Coverage for streamwise/model_provisioner/helix.py: 86%
155 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Helix algorithm for the StreamWise workflow allocation problem.
4Reference: https://github.com/Thesys-lab/Helix-ASPLOS25
6Helix optimizes models one-by-one following MODEL_ORDER, using MILP
7for each model's resource allocation. After each model reaches convergence
8(solver optimality or per-model time limit), its allocation is fixed and the
9remaining GPU budget is passed to the next model.
11Design rationale:
12 HelixAllocator does NOT inherit from MILPAllocator because the parent's
13 allocate() builds a single joint MILP for all models simultaneously.
14 Instead, HelixAllocator extends ModelAllocator and *composes*
15 MILPAllocator instances — one per model in the workflow.
17 For each model, a per-model WorkflowConfig is created where only the
18 target model has non-zero work (all others set to 0). The existing MILP
19 constraints (is_active <= work, gpus <= num_gpus * is_active) naturally
20 force 0 GPU allocation for those 0-work models, so no changes to
21 milp.py are required.
22"""
24from __future__ import annotations
26import logging
28from dataclasses import replace
29from typing import Optional
31from sim_types import Result
32from sim_types import GPUType
33from sim_types import WorkflowConfig
34from sim_types import PowerData
35from sim_types import LatencyData
36from sim_types import Model
37from sim_types import ModelAllocation
38from sim_types import Policy
39from sim_types import Solver
40from sim_types import MODEL_ORDER
42from model_allocator import ModelAllocator
44from evaluator import evaluate_model_allocation
46from .milp import MILPAllocator
48from .policies import HELIX_POLICY
49from .policies import MAX_DEVICES
51from constants import DEVICE_OPTIONS
54# Default per-model MILP solver time limit in seconds.
55# Each model gets this long to converge before the solver moves on.
56DEFAULT_PER_MODEL_TIME_LIMIT = 30
59def _compute_per_model_gpu_budget(
60 model_order: list[Model],
61 num_gpus: dict[GPUType, int],
62 workflow: WorkflowConfig,
63) -> dict[Model, dict[GPUType, int]]:
64 """Compute a per-model GPU budget so every model gets a fair share.
66 Budget is proportional to each model's ``MAX_DEVICES`` weight (capped
67 by the model's actual maximum useful device count from ``DEVICE_OPTIONS``).
68 Models not in ``MAX_DEVICES`` (e.g. OTHERS, UPSCALER) receive a minimum
69 allocation of ``min(DEVICE_OPTIONS)`` GPUs.
71 The allocations are floored per model, and any remainder is distributed
72 round-robin starting from the first model.
74 Returns:
75 Mapping ``model -> {gpu_type -> max_gpus}`` that the model may use.
76 """
77 # Effective weight per model (max useful devices)
78 weights: dict[Model, int] = {}
79 for m in model_order:
80 if workflow.model_work.get(m, 0) == 0:
81 continue
82 if m in MAX_DEVICES:
83 weights[m] = MAX_DEVICES[m]
84 else:
85 # Models not in MAX_DEVICES (OTHERS, UPSCALER) get min allocation
86 weights[m] = min(DEVICE_OPTIONS.get(m, [1]))
88 total_weight = sum(weights.values())
89 if total_weight == 0:
90 # Fallback: equal split
91 total_weight = len(weights) or 1
92 weights = {m: 1 for m in weights}
94 budget: dict[Model, dict[GPUType, int]] = {}
95 for gpu_type, total in num_gpus.items():
96 # Floor allocation per model
97 allocated = 0
98 per_model: dict[Model, int] = {}
99 for m in model_order:
100 if m not in weights:
101 continue
102 share = int(total * weights[m] / total_weight)
103 # Ensure at least 1 GPU per model (if GPUs available)
104 share = max(share, 1) if total - allocated >= 1 else 0
105 per_model[m] = share
106 allocated += share
108 # Distribute remainder round-robin
109 remainder = total - allocated
110 idx = 0
111 models_list = [m for m in model_order if m in per_model]
112 while remainder > 0 and models_list:
113 m = models_list[idx % len(models_list)]
114 per_model[m] += 1
115 remainder -= 1
116 idx += 1
118 for m in model_order:
119 if m not in per_model:
120 continue
121 if m not in budget:
122 budget[m] = {}
123 budget[m][gpu_type] = per_model[m]
125 return budget
128class HelixAllocator(ModelAllocator):
129 """
130 Helix-style allocator that optimizes models one at a time
131 using MILP, sequentially following MODEL_ORDER.
133 Reference: https://github.com/Thesys-lab/Helix-ASPLOS25
135 Key approach:
136 1. For each model in MODEL_ORDER, create a per-model MILP sub-problem
137 where only the target model has non-zero work.
138 2. Solve the MILP with the remaining GPU budget and a per-model time limit.
139 3. Fix the allocation for that model and subtract used GPUs.
140 4. Move to the next model with the remaining GPU budget.
141 5. Combine all per-model allocations into the final result.
143 The HelixAllocator uses composition (not inheritance) with MILPAllocator,
144 creating a separate MILPAllocator instance for each model's sub-problem.
145 This avoids modifying the joint MILP formulation and allows per-model
146 solver configurations.
147 """
149 def __init__(
150 self,
151 workflow: WorkflowConfig,
152 latency_data: LatencyData,
153 power_data: Optional[PowerData] = None,
154 policy: Policy = HELIX_POLICY,
155 ) -> None:
156 super().__init__(
157 workflow,
158 latency_data,
159 power_data,
160 policy,
161 )
162 assert self.policy.solver == Solver.HELIX
164 def allocate(
165 self,
166 num_gpus: dict[GPUType, int],
167 verbose: bool = False,
168 per_model_time_limit: int = DEFAULT_PER_MODEL_TIME_LIMIT,
169 milp_solver: Solver = Solver.HIGHS,
170 ) -> Result:
171 """
172 Allocate resources model-by-model following MODEL_ORDER.
174 For each model, a MILPAllocator is created with a workflow where
175 only the target model has non-zero work. The MILP solver optimizes
176 the allocation for that model within the remaining GPU budget.
178 Args:
179 num_gpus: Available GPUs per type.
180 verbose: If True, print per-model allocation details.
181 per_model_time_limit: Time limit (seconds) for each per-model MILP solve.
182 milp_solver: MILP solver backend to use (GUROBI or HIGHS).
184 Returns:
185 Combined Result across all models.
186 """
187 assert milp_solver in (Solver.GUROBI, Solver.HIGHS), \
188 f"milp_solver must be GUROBI or HIGHS, got {milp_solver}"
190 model_order = self.workflow.get_model_order()
191 if not self.policy.use_upscaler and Model.UPSCALER in model_order:
192 # Remove UPSCALER from model_order if not using upscaler to avoid unnecessary MILP solve
193 model_order.remove(Model.UPSCALER)
194 remaining_gpus = dict(num_gpus)
196 # ---- GPU budget partitioning ----
197 # Pre-compute a per-model GPU budget proportional to MAX_DEVICES
198 # so that early models cannot starve later ones. Unused GPUs from
199 # one model roll over to subsequent models.
200 gpu_budget = _compute_per_model_gpu_budget(
201 model_order, num_gpus, self.workflow,
202 )
204 if verbose:
205 logging.info("Helix GPU budget per model:")
206 for m in model_order:
207 if m in gpu_budget:
208 logging.info(f" {m.value}: {gpu_budget[m]}")
210 # Accumulated per-model allocations and metrics
211 all_model_allocations: dict[GPUType, dict[Model, list[ModelAllocation]]] = {}
212 total_makespan = 0.0
213 total_ttff = 0.0
214 total_cost = 0.0
215 total_energy = 0.0
216 total_gpus_used: dict[GPUType, int] = {gt: 0 for gt in num_gpus}
218 for model in model_order:
219 work = self.workflow.model_work.get(model, 0)
220 if work == 0:
221 continue
223 # Skip VAE models when disaggregation is disabled for the parent.
224 # Their latency is folded into the parent model's time calculation.
225 if model == Model.HF_VAE and not self.policy.is_disaggregated(Model.HF):
226 continue
227 if model == Model.FT_VAE and not self.policy.is_disaggregated(Model.FT):
228 continue
230 # Check if any GPUs remain
231 if all(v <= 0 for v in remaining_gpus.values()):
232 logging.warning(
233 f"Helix: No GPUs remaining for {model.value}. Skipping.")
234 continue
236 # Filter out GPU types with 0 remaining.
237 # Cap per-model GPUs to the budget so later models are not starved.
238 model_budget = gpu_budget.get(model, {})
239 active_gpus = {
240 gt: min(count, model_budget.get(gt, count))
241 for gt, count in remaining_gpus.items()
242 if count > 0 and (gt not in model_budget or model_budget[gt] > 0)
243 }
245 if verbose:
246 logging.info(
247 f"--- Helix: Optimizing {model.value} "
248 f"(work={work}) with remaining GPUs: {active_gpus} ---"
249 )
251 # ---- build per-model workflow ----
252 # Only the target model has work; other models are excluded from
253 # model_work so the MILP only builds variables/constraints for it.
254 per_model_work = {model: self.workflow.model_work[model]}
255 per_model_workflow = replace(
256 self.workflow,
257 model_work=per_model_work,
258 )
260 # ---- build MILP-compatible policy ----
261 # The inner MILPAllocator requires solver ∈ {GUROBI, HIGHS}.
262 # Force disaggregation / use_upscaler flags so that the inner
263 # MILP's ``model_names`` list includes VAE / UPSCALER when those
264 # are the target model. Without this, the MILP would construct
265 # an empty model set and produce a trivial (infeasible) problem.
266 disag = {} # dict(self.policy.disaggregation)
267 if model == Model.HF_VAE and self.policy.is_disaggregated(Model.HF):
268 disag[Model.HF] = True
269 if model == Model.FT_VAE and self.policy.is_disaggregated(Model.FT):
270 disag[Model.FT] = True
271 milp_policy = Policy(
272 name=self.policy.name,
273 gpu_cost=self.policy.gpu_cost,
274 objective=self.policy.objective,
275 # disaggregation=self.policy.disaggregation or model == Model.HF_VAE,
276 disaggregation=disag,
277 use_upscaler=self.policy.use_upscaler or model == Model.UPSCALER,
278 hardware=self.policy.hardware,
279 solver=milp_solver,
280 )
282 # ---- solve per-model MILP ----
283 milp_allocator = MILPAllocator(
284 workflow=per_model_workflow,
285 latency_data=self.latency_data,
286 power_data=self.power_data,
287 policy=milp_policy,
288 )
290 result = milp_allocator.allocate(
291 num_gpus=active_gpus,
292 verbose=verbose,
293 time_limit=per_model_time_limit,
294 # Use running_cost=True for linear cost formulation (HiGHS-compatible)
295 running_cost=(milp_solver == Solver.HIGHS),
296 # Skip server constraint: per-model allocations don't need
297 # to be multiples of NUM_GPUS_PER_SERVER.
298 skip_server_constraint=True,
299 )
301 if result.total_time_s == 0.0 and not result.models:
302 logging.warning(
303 f"Helix: MILP failed for {model.value}. Skipping.")
304 continue
306 # ---- record allocations & snap devices to DEVICE_OPTIONS ----
307 # The MILP constrains devices to DEVICE_OPTIONS, but floating-point
308 # precision in the solver can occasionally produce off-by-one values
309 # (e.g. 31 instead of 32). Snap each replica to the nearest valid
310 # option, adjusting the GPU accounting so we don't exceed the total
311 # budget passed to evaluate_model_allocation at the end.
312 for gpu_type, model_dict in result.models.items():
313 if gpu_type not in all_model_allocations:
314 all_model_allocations[gpu_type] = {}
315 for m_name, allocs in model_dict.items():
316 for alloc in allocs:
317 valid_devices = DEVICE_OPTIONS.get(m_name, [1])
318 if alloc.devices not in valid_devices:
319 nearest = min(valid_devices, key=lambda d: abs(d - alloc.devices))
320 diff = nearest - alloc.devices # positive = round up
321 gpu_avail = remaining_gpus.get(gpu_type, 0) - result.gpus_used.get(gpu_type, 0)
322 if diff > 0 and gpu_avail < diff:
323 # Not enough spare GPUs to round up; round down instead
324 nearest = max(
325 (d for d in valid_devices if d <= alloc.devices),
326 default=valid_devices[0],
327 )
328 diff = nearest - alloc.devices
329 logging.info(
330 f"Helix: snapping {m_name.value} from "
331 f"{alloc.devices} to {nearest} devices "
332 f"(solver precision fix, diff={diff:+d})")
333 # Adjust GPU accounting for this model's result
334 result.gpus_used[gpu_type] = result.gpus_used.get(gpu_type, 0) + diff
335 alloc.devices = nearest
336 all_model_allocations[gpu_type][m_name] = allocs
338 # ---- accumulate metrics ----
339 total_makespan += result.total_time_s
340 total_ttff += result.ttff_s
341 total_cost += result.cost
342 total_energy += result.total_energy
343 if verbose:
344 print(f'Model {model.value} - Time: {result.total_time_s:.2f}s,'
345 f'TTFF: {result.ttff_s:.2f}s, Cost: ${result.cost:.2f}')
346 print(f'Total cost so far: ${total_cost:.2f}, Total time so far: {total_makespan:.2f}s,'
347 f'Total TTFF so far: {total_ttff:.2f}s')
348 print(f'GPUs allocated for {model.value}: {result.gpus_used}')
350 # ---- subtract used GPUs ----
351 for gpu_type, used in result.gpus_used.items():
352 remaining_gpus[gpu_type] = remaining_gpus.get(gpu_type, 0) - used
353 total_gpus_used[gpu_type] = total_gpus_used.get(gpu_type, 0) + used
355 # ---- roll over unused budget to next models ----
356 # If this model used fewer GPUs than its budget, the surplus
357 # is distributed evenly among the remaining models.
358 remaining_models = [
359 m for m in model_order
360 if m in gpu_budget and MODEL_ORDER.get(m, 0) > MODEL_ORDER.get(model, 0)
361 ]
362 if remaining_models:
363 for gpu_type in num_gpus:
364 budget_for_model = model_budget.get(gpu_type, 0)
365 used_by_model = result.gpus_used.get(gpu_type, 0)
366 surplus = budget_for_model - used_by_model
367 if surplus > 0:
368 per_model_extra = surplus // len(remaining_models)
369 leftover = surplus % len(remaining_models)
370 for i, rm in enumerate(remaining_models):
371 extra = per_model_extra + (1 if i < leftover else 0)
372 gpu_budget[rm][gpu_type] = gpu_budget[rm].get(gpu_type, 0) + extra
374 if verbose:
375 print(
376 f"Helix: {model.value} allocated. "
377 f"Time: {result.total_time_s:.2f}s, "
378 f"TTFF: {result.ttff_s:.2f}s, "
379 f"GPUs used: {result.gpus_used}, "
380 f"Remaining: {remaining_gpus}"
381 )
383 result = evaluate_model_allocation(
384 workflow=self.workflow,
385 latency_data=self.latency_data,
386 power_data=self.power_data,
387 policy=self.policy,
388 models=all_model_allocations,
389 num_gpus=num_gpus,
390 )
392 if verbose:
393 print(
394 f"=== Helix final: "
395 f"Makespan={result.total_time_s:.2f}s, "
396 f"TTFF={result.ttff_s:.2f}s, "
397 f"TBF={result.tbf_s:.4f}s, "
398 f"Cost=${result.cost:.2f}, "
399 f"Energy={result.total_energy:.2f}Ws, "
400 f"GPUs used={result.gpus_used} ==="
401 )
403 return result