Integrate descriptor runtime candidates into harness
This commit is contained in:
@@ -4,7 +4,10 @@ import unittest
|
||||
|
||||
from aituner.engine_adapters.vllm import default_vllm_descriptors
|
||||
from aituner.knob_descriptor import KnobConstraints, KnobDescriptor
|
||||
from aituner.mechanism_planner import coordinate_line_search_candidates
|
||||
from aituner.mechanism_planner import (
|
||||
CoordinateSearchPolicy,
|
||||
coordinate_line_search_candidates,
|
||||
)
|
||||
|
||||
|
||||
class MechanismPlannerTests(unittest.TestCase):
|
||||
@@ -56,12 +59,26 @@ class MechanismPlannerTests(unittest.TestCase):
|
||||
descriptor = default_vllm_descriptors(tunable_flags=("gpu-memory-utilization",))[0]
|
||||
|
||||
candidates = coordinate_line_search_candidates(
|
||||
current_config={"gpu-memory-utilization": 0.98},
|
||||
current_config={"gpu-memory-utilization": 0.96},
|
||||
descriptors=(descriptor,),
|
||||
evidence_weights={"kv_memory_capacity": 0.8},
|
||||
)
|
||||
|
||||
self.assertEqual(candidates[0].patch, {"gpu-memory-utilization": 1.0})
|
||||
self.assertEqual(candidates[0].patch, {"gpu-memory-utilization": 0.97})
|
||||
|
||||
def test_coordinate_search_can_emit_larger_same_operator_steps(self) -> None:
|
||||
descriptor = default_vllm_descriptors(tunable_flags=("max-num-seqs",))[0]
|
||||
|
||||
candidates = coordinate_line_search_candidates(
|
||||
current_config={"max-num-seqs": 8},
|
||||
descriptors=(descriptor,),
|
||||
evidence_weights={"admission_capacity": 0.9},
|
||||
policy=CoordinateSearchPolicy(step_multipliers=(1.0, 2.0)),
|
||||
)
|
||||
|
||||
patches = [candidate.patch for candidate in candidates]
|
||||
self.assertIn({"max-num-seqs": 16}, patches)
|
||||
self.assertIn({"max-num-seqs": 24}, patches)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user