Coverage for tests/simulator/test_milp.py: 100%
128 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Test simulator module.
3"""
5import sys
6import os
7import pytest
9from copy import deepcopy
11# Add current path
12sys.path.append(os.getcwd())
14from tests.test_utils import temp_sys_path
16with temp_sys_path("simulator", "streamwise"):
17 from sim_types import LatencyData
18 from sim_types import PowerData
19 from sim_types import GPUType
20 from sim_types import Objective
21 from sim_types import Solver
22 from sim_types import QualityLevel
24 from data_loading import load_latency_data
25 from data_loading import load_power_data
27 from constants import DEFAULT_WORKFLOW_CONFIG
28 from constants import SECONDS_IN_HOUR
30 from model_provisioner.policies import STREAMWISE_MILP_POLICY
32 from workflows import WORKFLOWS
34 from model_provisioner.milp import MILPAllocator
36 from evaluator import evaluate_model_allocation
38 from utils import to_models_df
41def test_base() -> None:
42 workflow = DEFAULT_WORKFLOW_CONFIG
44 data_dir = "simulator/data/"
45 latency_data: LatencyData = load_latency_data(data_dir=data_dir)
46 power_data: PowerData = load_power_data(data_dir=data_dir)
48 policy = deepcopy(STREAMWISE_MILP_POLICY)
49 policy.solver = Solver.HIGHS # Options: "gurobi", "highs"
50 policy.objective = Objective.TTFF
52 num_gpus = {
53 GPUType.H200: 16,
54 GPUType.GB200: 16,
55 }
57 allocator = MILPAllocator(
58 workflow=workflow,
59 latency_data=latency_data,
60 power_data=power_data,
61 policy=policy,
62 )
63 result = allocator.allocate(
64 num_gpus=num_gpus,
65 running_cost=True, # To make it work with "highs"
66 verbose=True,
67 )
68 assert result is not None
69 assert result.models is not None
70 df_models = to_models_df(result.models)
71 assert not df_models.empty
72 assert 0 < result.ttff_s < result.total_time_s
73 assert 0 < result.cost
74 assert 0 < result.total_energy
75 assert result.gpus_used[GPUType.H200] > 8
76 assert result.gpus_used[GPUType.GB200] > 8
78 # Validation that the model allocation is correct
79 result_2 = evaluate_model_allocation(
80 models=result.models,
81 num_gpus=num_gpus,
82 workflow=workflow,
83 latency_data=latency_data,
84 power_data=power_data,
85 policy=policy,
86 )
87 assert result_2 is not None
88 assert result_2.models is not None
89 df_models = to_models_df(result_2.models)
90 assert not df_models.empty
91 assert 0 < result_2.ttff_s < result_2.total_time_s
92 assert 0 < result_2.cost
93 assert 0 < result_2.total_energy
94 assert result_2.gpus_used[GPUType.H200] > 8
95 assert result_2.gpus_used[GPUType.GB200] > 8
98def test_objective_cost() -> None:
99 """Test that the MILP allocator optimizes for cost."""
100 data_dir = "simulator/data/"
101 latency_data = load_latency_data(data_dir=data_dir)
102 power_data = load_power_data(data_dir=data_dir)
104 policy_cost = deepcopy(STREAMWISE_MILP_POLICY)
105 policy_cost.solver = Solver.HIGHS
106 policy_cost.objective = Objective.COST
108 allocator = MILPAllocator(
109 workflow=DEFAULT_WORKFLOW_CONFIG,
110 latency_data=latency_data,
111 power_data=power_data,
112 policy=policy_cost,
113 )
114 result = allocator.allocate(
115 num_gpus={
116 GPUType.H200: 24,
117 GPUType.GB200: 16,
118 },
119 running_cost=True, # Avoid: Highs interface does not support expressions of degree None
120 verbose=True,
121 max_ttff=10 * SECONDS_IN_HOUR, # 10 hours
122 max_makespan=10 * SECONDS_IN_HOUR, # 10 hours
123 force_num_gpus=True,
124 )
125 assert 0 < result.ttff_s <= result.total_time_s <= 10 * SECONDS_IN_HOUR
126 assert 0 < result.cost
127 assert 0 < result.total_energy
128 assert result.gpus_used[GPUType.H200] == 24
129 assert result.gpus_used[GPUType.GB200] == 16
132def test_objective_ttff() -> None:
133 """Test that the MILP allocator optimizes for TTFF."""
134 data_dir = "simulator/data/"
135 latency_data = load_latency_data(data_dir=data_dir)
136 power_data = load_power_data(data_dir=data_dir)
138 policy_ttff = deepcopy(STREAMWISE_MILP_POLICY)
139 policy_ttff.solver = Solver.HIGHS
140 policy_ttff.objective = Objective.TTFF
142 allocator = MILPAllocator(
143 workflow=DEFAULT_WORKFLOW_CONFIG,
144 latency_data=latency_data,
145 power_data=power_data,
146 policy=policy_ttff,
147 )
148 result = allocator.allocate(
149 num_gpus={
150 GPUType.A100: 16,
151 GPUType.GB200: 16,
152 },
153 running_cost=True,
154 verbose=True,
155 max_cost=1000,
156 force_num_gpus=True,
157 )
158 assert 0 < result.ttff_s < result.total_time_s
159 assert 0 < result.cost <= 1000
160 assert 0 < result.total_energy
161 assert result.gpus_used[GPUType.A100] > 8
162 assert result.gpus_used[GPUType.GB200] > 8
165def test_unfeasible() -> None:
166 """Test that the MILP allocator raises an exception for unfeasible constraints."""
167 data_dir = "simulator/data/"
168 latency_data = load_latency_data(data_dir=data_dir)
169 power_data = load_power_data(data_dir=data_dir)
171 policy_cost = deepcopy(STREAMWISE_MILP_POLICY)
172 policy_cost.solver = Solver.HIGHS
173 policy_cost.objective = Objective.COST
175 allocator = MILPAllocator(
176 workflow=DEFAULT_WORKFLOW_CONFIG,
177 latency_data=latency_data,
178 power_data=power_data,
179 policy=policy_cost,
180 )
181 with pytest.raises(Exception, match="A feasible solution was not found"):
182 allocator.allocate(
183 num_gpus={GPUType.A100: 16},
184 running_cost=True,
185 force_num_gpus=True,
186 max_cost=1, # Unfeasible constraint
187 )
190def test_highs_exception() -> None:
191 """Test that the MILP allocator with Highs raises an exception for unsupported expressions."""
192 data_dir = "simulator/data/"
193 latency_data = load_latency_data(data_dir=data_dir)
194 power_data = load_power_data(data_dir=data_dir)
196 policy_cost = deepcopy(STREAMWISE_MILP_POLICY)
197 policy_cost.solver = Solver.HIGHS
198 policy_cost.objective = Objective.COST
200 allocator = MILPAllocator(
201 workflow=DEFAULT_WORKFLOW_CONFIG,
202 latency_data=latency_data,
203 power_data=power_data,
204 policy=policy_cost,
205 )
206 with pytest.raises(Exception, match="Highs interface does not support expressions of degree None"):
207 allocator.allocate(
208 num_gpus={GPUType.H200: 32},
209 running_cost=False, # Trigger: Highs interface does not support expressions of degree None
210 )
213def test_gurobi() -> None:
214 """Test that the MILP allocator with Gurobi raises an exception ."""
215 data_dir = "simulator/data/"
216 latency_data = load_latency_data(data_dir=data_dir)
217 power_data = load_power_data(data_dir=data_dir)
219 policy_cost = deepcopy(STREAMWISE_MILP_POLICY)
220 policy_cost.solver = Solver.GUROBI
221 policy_cost.objective = Objective.COST
223 allocator = MILPAllocator(
224 workflow=DEFAULT_WORKFLOW_CONFIG,
225 latency_data=latency_data,
226 power_data=power_data,
227 policy=policy_cost,
228 )
229 with pytest.raises(
230 Exception,
231 match=r"(Model too large for size-limited license|No executable found for solver 'gurobi')",
232 ):
233 allocator.allocate(
234 num_gpus={GPUType.A100: 24},
235 )
238@pytest.mark.parametrize("workflow_name", WORKFLOWS.keys())
239def test_workflows(workflow_name: str) -> None:
240 """Test that the MILP allocator can allocate for all workflows."""
241 data_dir = "simulator/data/"
242 latency_data = load_latency_data(data_dir=data_dir)
243 power_data = load_power_data(data_dir=data_dir)
245 policy = deepcopy(STREAMWISE_MILP_POLICY)
246 policy.solver = Solver.HIGHS
247 policy.objective = Objective.TTFF
249 if WORKFLOWS[workflow_name].target_resolution != QualityLevel.HIGH:
250 policy.use_upscaler = False
252 allocator = MILPAllocator(
253 workflow=WORKFLOWS[workflow_name],
254 latency_data=latency_data,
255 power_data=power_data,
256 policy=policy,
257 )
258 result = allocator.allocate(
259 num_gpus={GPUType.H200: 16},
260 running_cost=True, # Avoid: Highs interface does not support expressions of degree None
261 verbose=True,
262 force_num_gpus=True,
263 )
264 assert result is not None
265 assert result.models is not None
266 assert 0 < result.ttff_s < result.total_time_s <= 10 * SECONDS_IN_HOUR
267 assert 0 < result.cost
268 assert 0 < result.total_energy
269 # assert result.gpus_used[GPUType.H200] == 16
270 assert 8 < result.gpus_used[GPUType.H200] <= 16