Coverage for streamwise/model_provisioner/greedy.py: 94%
236 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Greedy algorithm for the StreamWise workflow allocation problem.
3"""
5from __future__ import annotations
7import logging
9from tabulate import tabulate
11from typing import Optional
13from operator import itemgetter
15from constants import NUM_GPUS_PER_SERVER
16from constants import SECONDS_IN_MINUTE
17from constants import SECONDS_IN_HOUR
19from sim_types import Result
20from sim_types import GPUType
21from sim_types import WorkflowConfig
22from sim_types import LatencyData
23from sim_types import PowerData
24from sim_types import Model
25from sim_types import ModelAllocation
26from sim_types import Policy
27from sim_types import Solver
29from utils import simplify_model_allocations
31from evaluator import calc_used_gpus
32from evaluator import evaluate_model_allocation
34from model_allocator import ModelAllocator
36from .policies import STREAMWISE_POLICY
37from .policies import MAX_ITERATIONS
38from .policies import USE_ALL_GPUS
40from actions import gen_actions
41from actions import choose_action
42from actions import apply_action
45class GreedyAllocator(ModelAllocator):
46 """
47 Greedy allocator that iteratively applies the best action.
48 """
49 def __init__(
50 self,
51 workflow: WorkflowConfig,
52 latency_data: LatencyData,
53 power_data: Optional[PowerData] = None,
54 policy: Policy = STREAMWISE_POLICY,
55 ) -> None:
56 super().__init__(
57 workflow,
58 latency_data,
59 power_data,
60 policy,
61 )
62 assert self.policy.solver in {Solver.GREEDY, Solver.HEXGEN}
64 def allocate(
65 self,
66 num_gpus: dict[GPUType, int],
67 verbose: bool = False,
68 # Greedy policy parameters
69 allow_removal: bool = False,
70 allow_merging: bool = False,
71 look_ahead_replicas: int = 3,
72 ) -> Result:
73 total_gpus = sum(num_gpus.values())
74 assert total_gpus >= 8, f"Total number of GPUs must be at least 8 ({num_gpus})"
76 gpu_types = [
77 gpu_type
78 for gpu_type, count in num_gpus.items()
79 if count > 0
80 ]
81 assert 1 <= len(gpu_types) <= 2, f"Only up to two GPU types are supported ({len(gpu_types)})"
82 gpu_type1 = gpu_types[0]
84 if len(gpu_types) == 1 and num_gpus[gpu_type1] == 8:
85 # 8 x GPUs
86 return self._pick_from_single_server(
87 gpu_type=gpu_type1,
88 verbose=verbose,
89 )
91 if len(gpu_types) == 1:
92 # More than 8 x GPUs
93 return self._pick_from_single_device_mapping(
94 num_gpus.get(gpu_type1, 0),
95 gpu_type=gpu_type1,
96 verbose=verbose,
97 allow_removal=allow_removal,
98 allow_merging=allow_merging,
99 look_ahead_replicas=look_ahead_replicas,
100 )
102 # Mixed setup of GPU types (e.g., A100 and H100)
103 return self._pick_from_both_devices_mapping(
104 num_gpus,
105 verbose=verbose,
106 allow_removal=allow_removal,
107 allow_merging=allow_merging,
108 look_ahead_replicas=look_ahead_replicas,
109 )
111 def _pick_from_both_devices_mapping(
112 self,
113 num_gpus: dict[GPUType, int],
114 verbose: bool = False,
115 allow_removal: bool = False,
116 allow_merging: bool = False,
117 look_ahead_replicas: int = 3,
118 ) -> Result:
119 """
120 Calculate based on two GPU types.
121 """
122 gpu_types = list(num_gpus.keys())
123 assert len(gpu_types) == 2
124 assert len(num_gpus) == 2
125 gpu_type1 = gpu_types[0]
126 gpu_type2 = gpu_types[1]
127 assert num_gpus[gpu_type1] >= NUM_GPUS_PER_SERVER[gpu_type1]
128 assert num_gpus[gpu_type2] >= NUM_GPUS_PER_SERVER[gpu_type2]
130 # Initialize allocations with minimal setup
131 models = self._init_both_devices_models(gpu_type1, gpu_type2)
133 remaining_gpus = {}
134 for gpu_type in num_gpus.keys():
135 remaining_gpus[gpu_type] = num_gpus[gpu_type] - calc_used_gpus({gpu_type: models[gpu_type]})
137 # Optimization loop
138 if verbose:
139 evaluate_model_allocation(
140 models=models,
141 num_gpus=num_gpus,
142 workflow=self.workflow,
143 latency_data=self.latency_data,
144 power_data=self.power_data,
145 policy=self.policy,
146 round_up_cost_to_server=True,
147 )
148 self._print_iteration(0, models, num_gpus)
150 it = 1
151 prev_metric = None
152 switch_objective = False
153 while sum(remaining_gpus.values()) > 0:
154 # Calculate current iteration times
155 evaluate_model_allocation(
156 models=models,
157 num_gpus=num_gpus,
158 workflow=self.workflow,
159 latency_data=self.latency_data,
160 power_data=self.power_data,
161 policy=self.policy,
162 round_up_cost_to_server=False,
163 )
165 # Calculate potential actions for each optimization option
166 actions = gen_actions(
167 workflow=self.workflow,
168 latency_data=self.latency_data,
169 power_data=self.power_data,
170 num_gpus=num_gpus,
171 models=models,
172 policy=self.policy,
173 allow_removal=allow_removal,
174 allow_merging=allow_merging,
175 look_ahead_replicas=look_ahead_replicas,
176 )
178 if not actions:
179 logging.debug(f"No more actions possible after {it} iterations for {self.policy}.")
180 break
182 best_action = choose_action(actions, self.policy.objective, switch_objective=switch_objective)
184 if not best_action:
185 logging.debug("No actions selected.")
186 break
188 new_metric = best_action.get_metric(self.policy.objective, switch_objective=switch_objective)
190 if self.policy.objective.is_monotonic() and prev_metric is not None and new_metric >= prev_metric:
191 msg = f"No improvement after {it} iterations for {self.policy}."
192 msg += f" Best action: {best_action}, metric: {new_metric:.2f} >= previous {prev_metric:.2f}."
193 if verbose:
194 print(msg)
195 logging.debug(msg)
196 if not USE_ALL_GPUS:
197 logging.debug("Not using all GPUs as USE_ALL_GPUS is False. Stopping optimization loop.")
198 break
199 switch_objective = True
201 prev_metric = new_metric
203 models = apply_action(best_action, models=models)
205 models = simplify_model_allocations(models)
207 remaining_gpus.clear()
208 for gpu_type in num_gpus.keys():
209 remaining_gpus[gpu_type] = num_gpus[gpu_type] - calc_used_gpus({gpu_type: models[gpu_type]})
211 if verbose:
212 self._print_iteration(it, models, num_gpus)
213 print(f"{len(actions)} actions:")
214 for action in actions:
215 if action == best_action:
216 print(f"* {action} (best)")
217 else:
218 print(f" {action}")
219 print(f"Metric: {new_metric:.2f}")
220 print("Remaining devices:")
221 for gpu_type in remaining_gpus.keys():
222 print(f" {remaining_gpus[gpu_type]} x {gpu_type.value}")
224 it += 1
225 if it > MAX_ITERATIONS:
226 logging.debug(f"Reached max iterations ({MAX_ITERATIONS}). Stopping optimization loop.")
227 break
229 # Adjust for no disaggregation
230 if not self.policy.is_disaggregated(Model.HF):
231 for models_gpu in models.values():
232 for instance_id in range(len(models_gpu[Model.HF_VAE])):
233 assert models_gpu[Model.HF_VAE][instance_id].get_num_gpus() == 0, \
234 "HF_VAE must have 0 GPUs when HF disaggregation is disabled"
235 if not self.policy.is_disaggregated(Model.FT):
236 for models_gpu in models.values():
237 for instance_id in range(len(models_gpu[Model.FT_VAE])):
238 assert models_gpu[Model.FT_VAE][instance_id].get_num_gpus() == 0, \
239 "FT_VAE must have 0 GPUs when FT disaggregation is disabled"
241 # Final calculations
242 result = evaluate_model_allocation(
243 models=models,
244 num_gpus=num_gpus,
245 workflow=self.workflow,
246 latency_data=self.latency_data,
247 power_data=self.power_data,
248 policy=self.policy,
249 round_up_cost_to_server=True,
250 )
252 if verbose:
253 self._print_final_allocation(
254 models=models,
255 used_devices=result.gpus_used,
256 total_devices={
257 gpu_type1: num_gpus.get(gpu_type1, 0),
258 gpu_type2: num_gpus.get(gpu_type2, 0),
259 },
260 power_data=self.power_data,
261 total_time_s=result.total_time_s,
262 ttff_s=result.ttff_s,
263 first_chunk_time=result.first_chunk_time,
264 tbf_s=result.tbf_s,
265 total_energy=result.total_energy if self.power_data else 0.0,
266 cost=result.cost,
267 )
269 assert result.gpus_used[gpu_type1] <= num_gpus.get(gpu_type1, 0), \
270 f"{gpu_type1.value}: {result.gpus_used[gpu_type1]} > {num_gpus.get(gpu_type1, 0)}"
271 assert result.gpus_used[gpu_type2] <= num_gpus.get(gpu_type2, 0), \
272 f"{gpu_type2.value}: {result.gpus_used[gpu_type2]} > {num_gpus.get(gpu_type2, 0)}"
274 return Result(
275 total_time_s=result.total_time_s,
276 models=models,
277 gpus_used=result.gpus_used,
278 ttff_s=result.ttff_s,
279 tbf_s=result.tbf_s,
280 total_energy=result.total_energy if self.power_data else 0.0,
281 cost=result.cost,
282 )
284 def _pick_from_single_server(
285 self,
286 gpu_type: GPUType,
287 verbose: bool = False,
288 ) -> Result:
289 """
290 The minimal setup with a servers with a single server (8 GPUs or 4 for GB200).
291 No parallelism across scenes/subscenes.
292 """
294 # Number of devices
295 num_gpus = NUM_GPUS_PER_SERVER[gpu_type]
296 models = self._init_single_server_models(gpu_type)
298 result = evaluate_model_allocation(
299 models=models,
300 num_gpus={gpu_type: 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 )
308 if verbose:
309 model_device = models[gpu_type]
310 print_data = [
311 [Model.GEMMA.value, round(model_device[Model.GEMMA][0].time, 2)],
312 [Model.FLUX.value, round(model_device[Model.FLUX][0].time, 2)],
313 [Model.HF.value, round(model_device[Model.HF][0].time, 2)],
314 [Model.HF_VAE.value, round(model_device[Model.HF_VAE][0].time, 2)],
315 [Model.FT.value, round(model_device[Model.FT][0].time, 2)],
316 [Model.FT_VAE.value, round(model_device[Model.FT_VAE][0].time, 2)],
317 ]
318 if self.policy.use_upscaler:
319 print_data.append([Model.UPSCALER.value, round(model_device[Model.UPSCALER][0].time, 2)])
320 print(f"Total time: {result.total_time_s:.2f} seconds")
321 print(tabulate(
322 print_data,
323 headers=["Model", "Time (seconds)"],
324 tablefmt="pretty",
325 colalign=["left", "right"]
326 ))
327 self._print_final_allocation(
328 models=models,
329 used_devices={gpu_type: num_gpus},
330 total_devices={gpu_type: num_gpus},
331 power_data=self.power_data,
332 total_time_s=result.total_time_s,
333 ttff_s=result.ttff_s,
334 first_chunk_time=result.first_chunk_time,
335 tbf_s=result.tbf_s,
336 total_energy=result.total_energy if self.power_data else 0.0,
337 cost=result.cost,
338 )
340 return Result(
341 total_time_s=result.total_time_s,
342 models=models,
343 gpus_used={gpu_type: num_gpus},
344 ttff_s=result.ttff_s,
345 tbf_s=result.tbf_s,
346 total_energy=result.total_energy if self.power_data else 0.0,
347 cost=result.cost,
348 )
350 def _pick_from_single_device_mapping(
351 self,
352 num_gpus: int,
353 gpu_type: GPUType,
354 verbose: bool = False,
355 allow_removal: bool = False,
356 allow_merging: bool = False,
357 look_ahead_replicas: int = 3,
358 ) -> Result:
359 """
360 Calculate time and energy based on a single GPU type.
361 """
362 assert num_gpus >= NUM_GPUS_PER_SERVER[gpu_type]
363 latency_gpu_data = self.latency_data[gpu_type]
364 assert gpu_type == latency_gpu_data.gpu_type
366 if self.power_data is not None:
367 power_gpu_data = self.power_data[gpu_type]
368 assert gpu_type == power_gpu_data.gpu_type
370 # Initialize allocations
371 models = self._init_single_device_models(gpu_type)
373 remaining_gpus = num_gpus - calc_used_gpus(models)
375 assert 0 <= remaining_gpus <= num_gpus
377 # Optimization loop
378 it = 0
379 prev_metric = None
380 switch_objective = False
381 while remaining_gpus > 0:
382 # Calculate current iteration times
383 evaluate_model_allocation(
384 models=models,
385 num_gpus={gpu_type: num_gpus},
386 workflow=self.workflow,
387 latency_data=self.latency_data,
388 power_data=self.power_data,
389 policy=self.policy,
390 round_up_cost_to_server=False,
391 )
393 # Calculate potential actions for each optimization option
394 actions = gen_actions(
395 num_gpus={gpu_type: num_gpus},
396 latency_data=self.latency_data,
397 power_data=self.power_data,
398 workflow=self.workflow,
399 models=models,
400 policy=self.policy,
401 allow_removal=allow_removal,
402 allow_merging=allow_merging,
403 look_ahead_replicas=look_ahead_replicas,
404 )
406 if not actions:
407 logging.debug(f"No more actions possible after {it} iterations for {self.policy}")
408 break
410 best_action = choose_action(
411 actions,
412 self.policy.objective,
413 switch_objective=switch_objective)
415 if not best_action:
416 logging.debug("No action selected.")
417 break
419 new_metric = best_action.get_metric(self.policy.objective, switch_objective=switch_objective)
420 if self.policy.objective.is_monotonic() and prev_metric is not None and new_metric >= prev_metric:
421 msg = f"No improvement from actions after {it} iterations for {self.policy}."
422 msg += f" Best action: {best_action}, metric: {new_metric:.2f} >= previous {prev_metric:.2f}."
423 if verbose:
424 print(msg)
425 logging.debug(msg)
426 if not USE_ALL_GPUS:
427 logging.debug("Not using all GPUs as USE_ALL_GPUS is False. Stopping optimization loop.")
428 break
429 switch_objective = True
431 models = apply_action(best_action, models)
433 models = simplify_model_allocations(models)
435 remaining_gpus = num_gpus - calc_used_gpus(models)
436 prev_metric = new_metric
438 if verbose:
439 self._print_iteration(it, models, {gpu_type: num_gpus})
440 print(f"Metric: {new_metric:.2f}")
441 print(f"{len(actions)} actions:")
442 for action in actions:
443 if action == best_action:
444 print(f" * {action} (best)")
445 else:
446 print(f" {action}")
447 print(f"Applied: {best_action}")
448 print(f"Remaining devices: {remaining_gpus}x{gpu_type}")
450 it += 1
451 if it > MAX_ITERATIONS:
452 logging.debug(f"Reached max iterations ({MAX_ITERATIONS}). Stopping optimization loop.")
453 break
455 result = evaluate_model_allocation(
456 models=models,
457 num_gpus={gpu_type: num_gpus},
458 workflow=self.workflow,
459 latency_data=self.latency_data,
460 power_data=self.power_data,
461 policy=self.policy,
462 round_up_cost_to_server=True,
463 )
465 if verbose:
466 self._print_final_allocation(
467 models=models,
468 used_devices=result.gpus_used,
469 total_devices={gpu_type: num_gpus},
470 power_data=self.power_data,
471 total_time_s=result.total_time_s,
472 ttff_s=result.ttff_s,
473 first_chunk_time=result.first_chunk_time,
474 tbf_s=result.tbf_s,
475 total_energy=result.total_energy if self.power_data else 0.0,
476 cost=result.cost,
477 )
479 if not self.policy.is_disaggregated(Model.HF):
480 if models[gpu_type][Model.HF_VAE]:
481 assert models[gpu_type][Model.HF_VAE][0].get_num_gpus() == 0, \
482 "HF_VAE must have 0 GPUs when HF disaggregation is disabled"
483 if not self.policy.is_disaggregated(Model.FT):
484 if models[gpu_type][Model.FT_VAE]:
485 assert models[gpu_type][Model.FT_VAE][0].get_num_gpus() == 0, \
486 "FT_VAE must have 0 GPUs when FT disaggregation is disabled"
487 num_gpus_used = result.gpus_used[gpu_type]
488 assert num_gpus_used <= num_gpus, f"{num_gpus_used}>{num_gpus} for {gpu_type.value}"
490 return Result(
491 total_time_s=result.total_time_s,
492 models=models,
493 gpus_used={gpu_type: num_gpus_used},
494 gpus_total={gpu_type: num_gpus},
495 ttff_s=result.ttff_s,
496 tbf_s=result.tbf_s,
497 total_energy=result.total_energy if self.power_data else 0.0,
498 cost=result.cost,
499 )
501 def _print_iteration(
502 self,
503 it: int,
504 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
505 num_gpus: dict[GPUType, int],
506 ) -> None:
507 print(f"--- Iteration {it} ---")
509 for gpu_type in models.keys():
510 total_gpus = calc_used_gpus({gpu_type: models[gpu_type]})
511 print(f"Current {gpu_type.value} allocation: {total_gpus}/{num_gpus[gpu_type]} GPUs")
512 for model in Model:
513 for model_instance in models[gpu_type][model]:
514 if model_instance.get_num_gpus() > 0:
515 print(f" {model.value:10s}:\t{model_instance}")
517 # Find the bottleneck stage
518 stage_times: dict[Model, float] = {}
519 ttff_times: dict[Model, float] = {}
520 for model_name in Model:
521 times = []
522 times_first = []
523 for gpu_type in models.keys():
524 for model_alloc in models[gpu_type][model_name]:
525 times.append(model_alloc.time)
526 times_first.append(model_alloc.time_first)
527 stage_times[model_name] = max(times) if times else 0.0
528 ttff_times[model_name] = max(times_first) if times_first else 0.0
530 bottleneck_stage, bottleneck_time = max(
531 stage_times.items(),
532 key=itemgetter(1)
533 )
534 bottleneck_ttff_stage, bottleneck_ttff_time = max(
535 ttff_times.items(),
536 key=itemgetter(1)
537 )
538 print(f"Bottleneck: {bottleneck_stage} ({bottleneck_time:.2f}s)")
539 print(f"Bottleneck TTFF: {bottleneck_ttff_stage} ({bottleneck_ttff_time:.2f}s)")
540 # bottleneck stage is not necessarily the stage with the
541 # highest potential gain from scaling up/out
543 def _print_final_allocation(
544 self,
545 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
546 used_devices: dict[GPUType, int],
547 total_devices: dict[GPUType, int],
548 power_data: Optional[PowerData],
549 total_time_s: float,
550 ttff_s: float,
551 first_chunk_time: float,
552 tbf_s: float,
553 total_energy: float,
554 cost: float,
555 ) -> None:
556 print("=== FINAL ALLOCATION ===")
557 print("Total devices used/available:")
558 for gpu_type, total_device in total_devices.items():
559 used_device = used_devices[gpu_type]
560 print(f" {gpu_type.value}: {used_device}/{total_device}")
561 print("Model allocations:")
562 for gpu_type in models.keys():
563 print(f" {gpu_type.value} ({used_devices[gpu_type]} used):")
564 for model in Model:
565 for model_alloc in models[gpu_type][model]:
566 print(f" {model.value:10s}:\t{model_alloc}")
567 print(f"Total time: {total_time_s:.2f} seconds ({total_time_s / SECONDS_IN_MINUTE:.2f} minutes)")
568 print(f"TTFF: {ttff_s:.2f} seconds")
569 print(f"First chunk time: {first_chunk_time:.2f} seconds")
570 print(f"TBF: {tbf_s:.2f} seconds")
571 print(f"Total cost: ${cost:.2f}")
572 if power_data is not None:
573 print(f"Total energy: {total_energy:.2f} Ws ({total_energy / SECONDS_IN_HOUR / 1000:.2f} kWh)")