Coverage for streamwise/model_provisioner/hexgen.py: 59%
259 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2HexGen algorithm for the StreamWise workflow allocation problem.
4Reference: https://arxiv.org/abs/2311.11514
6HexGen treats each model in the workflow as an independent component for optimization.
7It tracks metrics per model and optimizes models sequentially according to MODEL_ORDER.
8When a model's metric converges (stops dropping), it moves to the next model.
9After the last model converges, it cycles back to the first model and allocates
10remaining GPUs until exhausted.
11"""
13from __future__ import annotations
14import logging
15from typing import Optional
17from sim_types import Result
18from sim_types import GPUType
19from sim_types import WorkflowConfig
20from sim_types import PowerData
21from sim_types import LatencyData
22from sim_types import Model
23from sim_types import ModelAllocation
24from sim_types import Policy
25from sim_types import Solver
26from sim_types import MODEL_ORDER
28from utils import simplify_model_allocations
30from evaluator import calc_used_gpus
31from evaluator import evaluate_model_allocation
33from .greedy import GreedyAllocator
35from actions import gen_actions
36from actions import choose_action
37from actions import apply_action
39from .policies import HEXGEN_POLICY
40from .policies import MAX_ITERATIONS
41from .policies import USE_ALL_GPUS
44def _get_model_order(workflow: WorkflowConfig) -> list[Model]:
45 """Get ordered list of models in the workflow, sorted by MODEL_ORDER."""
46 return sorted(
47 [m for m in workflow.models if m in MODEL_ORDER],
48 key=lambda m: MODEL_ORDER[m],
49 )
52class HexGenAllocator(GreedyAllocator):
53 """
54 HexGen-style allocator that optimizes models one at a time,
55 sequentially following MODEL_ORDER.
57 Reference: https://arxiv.org/abs/2311.11514
59 Key differences from GreedyAllocator:
60 1. Each model is treated as an independent optimization target.
61 2. Per-model metrics are tracked separately.
62 3. Models are optimized in MODEL_ORDER sequence. When a model's metric
63 converges, it moves to the next model. After the last model converges,
64 it cycles back to the first and allocates remaining GPUs.
65 """
67 def __init__(
68 self,
69 workflow: WorkflowConfig,
70 latency_data: LatencyData,
71 power_data: Optional[PowerData] = None,
72 policy: Policy = HEXGEN_POLICY,
73 ) -> None:
74 super().__init__(
75 workflow,
76 latency_data,
77 power_data,
78 policy,
79 )
80 assert self.policy.solver == Solver.HEXGEN
82 def _pick_from_single_device_mapping(
83 self,
84 num_gpus: int,
85 gpu_type: GPUType,
86 verbose: bool = False,
87 allow_removal: bool = False,
88 allow_merging: bool = False,
89 look_ahead_replicas: int = 3,
90 ) -> Result:
91 """
92 HexGen-style allocation for a single GPU type (>8 GPUs).
93 Optimizes models one at a time following MODEL_ORDER.
94 """
95 from constants import NUM_GPUS_PER_SERVER
97 assert num_gpus >= NUM_GPUS_PER_SERVER[gpu_type]
99 # Initialize allocations (same as GreedyAllocator)
100 models = self._init_single_device_models(gpu_type)
102 remaining_gpus = num_gpus - calc_used_gpus(models)
103 assert 0 <= remaining_gpus <= num_gpus
105 # --- HexGen per-model sequential optimization ---
106 model_order = _get_model_order(self.workflow)
107 per_model_metrics: dict[Model, Optional[float]] = {m: None for m in model_order}
109 it = 0
110 current_model_idx = 0
111 cycles_without_progress = 0 # track full cycles without any improvement
112 total_models = len(model_order)
114 while remaining_gpus > 0:
115 if current_model_idx >= total_models:
116 # Completed a full cycle, wrap around
117 current_model_idx = 0
118 cycles_without_progress += 1
119 if cycles_without_progress >= 1:
120 logging.debug(
121 f"HexGen: No progress after {cycles_without_progress} full cycles.")
122 break
124 current_model = model_order[current_model_idx]
126 if verbose:
127 print(f"--- HexGen: Optimizing {current_model.value} "
128 f"(model {current_model_idx + 1}/{total_models}) ---")
130 # Inner loop: keep optimizing current model until convergence
131 inner_it = 0
132 while remaining_gpus > 0:
133 # Evaluate current state
134 evaluate_model_allocation(
135 models=models,
136 num_gpus={gpu_type: num_gpus},
137 workflow=self.workflow,
138 latency_data=self.latency_data,
139 power_data=self.power_data,
140 policy=self.policy,
141 round_up_cost_to_server=False,
142 )
144 # Generate actions only for the current model
145 all_actions = gen_actions(
146 num_gpus={gpu_type: num_gpus},
147 latency_data=self.latency_data,
148 power_data=self.power_data,
149 workflow=self.workflow,
150 models=models,
151 policy=self.policy,
152 )
154 # Filter to actions targeting the current model only
155 model_actions = [a for a in all_actions if a.model == current_model]
157 if not model_actions:
158 logging.debug(
159 f"HexGen: No actions for {current_model.value} after {inner_it} inner iterations.")
160 break
162 best_action = choose_action(model_actions, self.policy.objective)
164 if not best_action:
165 logging.debug(f"HexGen: No action selected for {current_model.value}.")
166 break
168 new_metric = best_action.get_metric(self.policy.objective)
169 prev_metric = per_model_metrics[current_model]
171 if self.policy.objective.is_monotonic() and prev_metric is not None and new_metric >= prev_metric:
172 msg = (
173 f"HexGen: {current_model.value} converged after {inner_it} inner iterations. "
174 f"Metric: {new_metric:.2f} >= previous {prev_metric:.2f}."
175 )
176 if verbose:
177 print(msg)
178 logging.debug(msg)
179 break
181 per_model_metrics[current_model] = new_metric
183 models = apply_action(best_action, models=models)
184 models = simplify_model_allocations(models)
186 remaining_gpus = num_gpus - calc_used_gpus(models)
188 if verbose:
189 self._print_iteration(it, models, {gpu_type: num_gpus})
190 print(f"HexGen: Applied action for {current_model.value}, "
191 f"metric: {new_metric:.2f}, remaining: {remaining_gpus}")
193 it += 1
194 inner_it += 1
196 if it > MAX_ITERATIONS:
197 logging.debug(f"HexGen: Reached max iterations ({MAX_ITERATIONS}). Stopping.")
198 break
200 if it > MAX_ITERATIONS:
201 break
203 current_model_idx += 1
205 # --- USE_ALL_GPUS: fill remaining GPUs by cycling through MODEL_ORDER ---
206 remaining_gpus = num_gpus - calc_used_gpus(models)
207 if USE_ALL_GPUS and remaining_gpus > 0:
208 models = self._fill_remaining_gpus_single(
209 models=models,
210 num_gpus=num_gpus,
211 gpu_type=gpu_type,
212 model_order=model_order,
213 it=it,
214 verbose=verbose,
215 )
217 # Final evaluation
218 result = evaluate_model_allocation(
219 models=models,
220 num_gpus={gpu_type: num_gpus},
221 workflow=self.workflow,
222 latency_data=self.latency_data,
223 power_data=self.power_data,
224 policy=self.policy,
225 round_up_cost_to_server=True,
226 )
228 if verbose:
229 self._print_final_allocation(
230 models=models,
231 used_devices=result.gpus_used,
232 total_devices={gpu_type: num_gpus},
233 power_data=self.power_data,
234 total_time_s=result.total_time_s,
235 ttff_s=result.ttff_s,
236 first_chunk_time=result.first_chunk_time,
237 tbf_s=result.tbf_s,
238 total_energy=result.total_energy if self.power_data else 0.0,
239 cost=result.cost,
240 )
242 if not self.policy.is_disaggregated(Model.HF):
243 if models[gpu_type][Model.HF_VAE]:
244 assert models[gpu_type][Model.HF_VAE][0].get_num_gpus() == 0, \
245 "HF_VAE must have 0 GPUs when HF disaggregation is disabled"
246 if not self.policy.is_disaggregated(Model.FT):
247 if models[gpu_type][Model.FT_VAE]:
248 assert models[gpu_type][Model.FT_VAE][0].get_num_gpus() == 0, \
249 "FT_VAE must have 0 GPUs when FT disaggregation is disabled"
251 num_gpus_used = result.gpus_used[gpu_type]
252 assert num_gpus_used <= num_gpus, f"{num_gpus_used}>{num_gpus} for {gpu_type.value}"
254 return Result(
255 total_time_s=result.total_time_s,
256 models=models,
257 gpus_used={gpu_type: num_gpus_used},
258 gpus_total={gpu_type: num_gpus},
259 ttff_s=result.ttff_s,
260 tbf_s=result.tbf_s,
261 total_energy=result.total_energy if self.power_data else 0.0,
262 cost=result.cost,
263 )
265 def _pick_from_both_devices_mapping(
266 self,
267 num_gpus: dict[GPUType, int],
268 verbose: bool = False,
269 allow_removal: bool = False,
270 allow_merging: bool = False,
271 look_ahead_replicas: int = 3,
272 ) -> Result:
273 """
274 HexGen-style allocation for two GPU types.
275 Optimizes models one at a time following MODEL_ORDER.
276 """
277 from constants import NUM_GPUS_PER_SERVER
279 gpu_types = list(num_gpus.keys())
280 assert len(gpu_types) == 2
281 gpu_type1 = gpu_types[0]
282 gpu_type2 = gpu_types[1]
283 assert num_gpus[gpu_type1] >= NUM_GPUS_PER_SERVER[gpu_type1]
284 assert num_gpus[gpu_type2] >= NUM_GPUS_PER_SERVER[gpu_type2]
286 # Initialize allocations (same as GreedyAllocator)
287 models = self._init_both_devices_models(gpu_type1, gpu_type2)
289 remaining_gpus: dict[GPUType, int] = {}
290 for gpu_type in num_gpus.keys():
291 remaining_gpus[gpu_type] = num_gpus[gpu_type] - calc_used_gpus({gpu_type: models[gpu_type]})
293 # --- HexGen per-model sequential optimization ---
294 model_order = _get_model_order(self.workflow)
295 per_model_metrics: dict[Model, Optional[float]] = {m: None for m in model_order}
297 if verbose:
298 evaluate_model_allocation(
299 models=models,
300 num_gpus=num_gpus,
301 workflow=self.workflow,
302 latency_data=self.latency_data,
303 power_data=self.power_data,
304 policy=self.policy,
305 round_up_cost_to_server=True,
306 )
307 self._print_iteration(0, models, num_gpus)
309 it = 1
310 current_model_idx = 0
311 cycles_without_progress = 0
312 total_models = len(model_order)
314 while sum(remaining_gpus.values()) > 0:
315 if current_model_idx >= total_models:
316 current_model_idx = 0
317 cycles_without_progress += 1
318 if cycles_without_progress >= 1:
319 logging.debug(
320 f"HexGen: No progress after {cycles_without_progress} full cycles.")
321 break
323 current_model = model_order[current_model_idx]
325 if verbose:
326 print(f"--- HexGen: Optimizing {current_model.value} "
327 f"(model {current_model_idx + 1}/{total_models}) ---")
329 inner_it = 0
331 while sum(remaining_gpus.values()) > 0:
332 evaluate_model_allocation(
333 models=models,
334 num_gpus=num_gpus,
335 workflow=self.workflow,
336 latency_data=self.latency_data,
337 power_data=self.power_data,
338 policy=self.policy,
339 round_up_cost_to_server=False,
340 )
342 all_actions = gen_actions(
343 workflow=self.workflow,
344 latency_data=self.latency_data,
345 power_data=self.power_data,
346 num_gpus=num_gpus,
347 models=models,
348 policy=self.policy,
349 )
351 # Filter to current model
352 model_actions = [a for a in all_actions if a.model == current_model]
354 if not model_actions:
355 logging.debug(
356 f"HexGen: No actions for {current_model.value} after {inner_it} inner iterations.")
357 break
359 best_action = choose_action(model_actions, self.policy.objective)
361 if not best_action:
362 logging.debug(f"HexGen: No action selected for {current_model.value}.")
363 break
365 new_metric = best_action.get_metric(self.policy.objective)
366 prev_metric = per_model_metrics[current_model]
368 if self.policy.objective.is_monotonic() and prev_metric is not None and new_metric >= prev_metric:
369 msg = (
370 f"HexGen: {current_model.value} converged. "
371 f"Metric: {new_metric:.2f} >= previous {prev_metric:.2f}."
372 )
373 if verbose:
374 print(msg)
375 logging.debug(msg)
376 break
378 per_model_metrics[current_model] = new_metric
380 models = apply_action(best_action, models=models)
381 models = simplify_model_allocations(models)
383 remaining_gpus.clear()
384 for gpu_type in num_gpus.keys():
385 remaining_gpus[gpu_type] = num_gpus[gpu_type] - calc_used_gpus({gpu_type: models[gpu_type]})
387 if verbose:
388 self._print_iteration(it, models, num_gpus)
389 print(f"HexGen: Applied action for {current_model.value}, "
390 f"metric: {new_metric:.2f}")
391 print("Remaining devices:")
392 for gt in remaining_gpus:
393 print(f" {remaining_gpus[gt]} x {gt.value}")
395 it += 1
396 inner_it += 1
398 if it > MAX_ITERATIONS:
399 logging.debug(f"HexGen: Reached max iterations ({MAX_ITERATIONS}). Stopping.")
400 break
402 if it > MAX_ITERATIONS:
403 break
405 current_model_idx += 1
407 # --- USE_ALL_GPUS: fill remaining GPUs by cycling through MODEL_ORDER ---
408 remaining_gpus_total = sum(
409 num_gpus[gt] - calc_used_gpus({gt: models[gt]})
410 for gt in num_gpus
411 )
412 if USE_ALL_GPUS and remaining_gpus_total > 0:
413 models = self._fill_remaining_gpus_multi(
414 models=models,
415 num_gpus=num_gpus,
416 model_order=model_order,
417 it=it,
418 verbose=verbose,
419 )
421 # Adjust for no disaggregation
422 if not self.policy.is_disaggregated(Model.HF):
423 for models_gpu in models.values():
424 for instance_id in range(len(models_gpu[Model.HF_VAE])):
425 assert models_gpu[Model.HF_VAE][instance_id].get_num_gpus() == 0, \
426 "HF_VAE must have 0 GPUs when HF disaggregation is disabled"
427 if not self.policy.is_disaggregated(Model.FT):
428 for models_gpu in models.values():
429 for instance_id in range(len(models_gpu[Model.FT_VAE])):
430 assert models_gpu[Model.FT_VAE][instance_id].get_num_gpus() == 0, \
431 "FT_VAE must have 0 GPUs when FT disaggregation is disabled"
433 # Final evaluation
434 result = evaluate_model_allocation(
435 models=models,
436 num_gpus=num_gpus,
437 workflow=self.workflow,
438 latency_data=self.latency_data,
439 power_data=self.power_data,
440 policy=self.policy,
441 round_up_cost_to_server=True,
442 )
444 if verbose:
445 self._print_final_allocation(
446 models=models,
447 used_devices=result.gpus_used,
448 total_devices={
449 gpu_type1: num_gpus.get(gpu_type1, 0),
450 gpu_type2: num_gpus.get(gpu_type2, 0),
451 },
452 power_data=self.power_data,
453 total_time_s=result.total_time_s,
454 ttff_s=result.ttff_s,
455 first_chunk_time=result.first_chunk_time,
456 tbf_s=result.tbf_s,
457 total_energy=result.total_energy if self.power_data else 0.0,
458 cost=result.cost,
459 )
461 assert result.gpus_used[gpu_type1] <= num_gpus.get(gpu_type1, 0), \
462 f"{gpu_type1.value}: {result.gpus_used[gpu_type1]} > {num_gpus.get(gpu_type1, 0)}"
463 assert result.gpus_used[gpu_type2] <= num_gpus.get(gpu_type2, 0), \
464 f"{gpu_type2.value}: {result.gpus_used[gpu_type2]} > {num_gpus.get(gpu_type2, 0)}"
466 return Result(
467 total_time_s=result.total_time_s,
468 models=models,
469 gpus_used=result.gpus_used,
470 ttff_s=result.ttff_s,
471 tbf_s=result.tbf_s,
472 total_energy=result.total_energy if self.power_data else 0.0,
473 cost=result.cost,
474 )
476 def _fill_remaining_gpus_single(
477 self,
478 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
479 num_gpus: int,
480 gpu_type: GPUType,
481 model_order: list[Model],
482 it: int = 0,
483 verbose: bool = False,
484 ) -> dict[GPUType, dict[Model, list[ModelAllocation]]]:
485 """
486 Fill remaining GPUs by cycling through MODEL_ORDER (single GPU type).
487 Applies any available action per model, ignoring metric convergence.
488 Stops when all GPUs are used or no model can accept more.
489 """
490 remaining_gpus = num_gpus - calc_used_gpus(models)
491 total_models = len(model_order)
492 model_idx = 0
493 models_exhausted: set[Model] = set()
495 if verbose:
496 print(f"--- HexGen: USE_ALL_GPUS fill phase, {remaining_gpus} remaining ---")
498 while remaining_gpus > 0 and len(models_exhausted) < total_models:
499 current_model = model_order[model_idx % total_models]
500 model_idx += 1
502 if current_model in models_exhausted:
503 continue
505 evaluate_model_allocation(
506 models=models,
507 num_gpus={gpu_type: num_gpus},
508 workflow=self.workflow,
509 latency_data=self.latency_data,
510 power_data=self.power_data,
511 policy=self.policy,
512 round_up_cost_to_server=False,
513 )
515 all_actions = gen_actions(
516 num_gpus={gpu_type: num_gpus},
517 latency_data=self.latency_data,
518 power_data=self.power_data,
519 workflow=self.workflow,
520 models=models,
521 policy=self.policy,
522 )
523 model_actions = [a for a in all_actions if a.model == current_model]
525 if not model_actions:
526 models_exhausted.add(current_model)
527 logging.debug(f"HexGen fill: {current_model.value} exhausted (no actions).")
528 continue
530 best_action = choose_action(model_actions, self.policy.objective)
531 if not best_action:
532 models_exhausted.add(current_model)
533 logging.debug(f"HexGen fill: {current_model.value} exhausted (no action selected).")
534 continue
536 models = apply_action(best_action, models=models)
537 models = simplify_model_allocations(models)
538 remaining_gpus = num_gpus - calc_used_gpus(models)
540 if verbose:
541 self._print_iteration(it, models, {gpu_type: num_gpus})
542 print(f"HexGen fill: Allocated to {current_model.value}, remaining: {remaining_gpus}")
544 it += 1
545 if it > MAX_ITERATIONS:
546 logging.debug(f"HexGen fill: Reached max iterations ({MAX_ITERATIONS}). Stopping.")
547 break
549 return models
551 def _fill_remaining_gpus_multi(
552 self,
553 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
554 num_gpus: dict[GPUType, int],
555 model_order: list[Model],
556 it: int = 0,
557 verbose: bool = False,
558 ) -> dict[GPUType, dict[Model, list[ModelAllocation]]]:
559 """
560 Fill remaining GPUs by cycling through MODEL_ORDER (multi GPU type).
561 Applies any available action per model, ignoring metric convergence.
562 Stops when all GPUs are used or no model can accept more.
563 """
564 total_remaining = sum(
565 num_gpus[gt] - calc_used_gpus({gt: models[gt]})
566 for gt in num_gpus
567 )
568 total_models = len(model_order)
569 model_idx = 0
570 models_exhausted: set[Model] = set()
572 if verbose:
573 print(f"--- HexGen: USE_ALL_GPUS fill phase, {total_remaining} remaining ---")
575 while total_remaining > 0 and len(models_exhausted) < total_models:
576 current_model = model_order[model_idx % total_models]
577 model_idx += 1
579 if current_model in models_exhausted:
580 continue
582 evaluate_model_allocation(
583 models=models,
584 num_gpus=num_gpus,
585 workflow=self.workflow,
586 latency_data=self.latency_data,
587 power_data=self.power_data,
588 policy=self.policy,
589 round_up_cost_to_server=False,
590 )
592 all_actions = gen_actions(
593 workflow=self.workflow,
594 latency_data=self.latency_data,
595 power_data=self.power_data,
596 num_gpus=num_gpus,
597 models=models,
598 policy=self.policy,
599 )
600 model_actions = [a for a in all_actions if a.model == current_model]
602 if not model_actions:
603 models_exhausted.add(current_model)
604 logging.debug(f"HexGen fill: {current_model.value} exhausted (no actions).")
605 continue
607 best_action = choose_action(model_actions, self.policy.objective)
608 if not best_action:
609 models_exhausted.add(current_model)
610 logging.debug(f"HexGen fill: {current_model.value} exhausted (no action selected).")
611 continue
613 models = apply_action(best_action, models=models)
614 models = simplify_model_allocations(models)
615 total_remaining = sum(
616 num_gpus[gt] - calc_used_gpus({gt: models[gt]})
617 for gt in num_gpus
618 )
620 if verbose:
621 self._print_iteration(it, models, num_gpus)
622 print(f"HexGen fill: Allocated to {current_model.value}, remaining: {total_remaining}")
624 it += 1
625 if it > MAX_ITERATIONS:
626 logging.debug(f"HexGen fill: Reached max iterations ({MAX_ITERATIONS}). Stopping.")
627 break
629 return models