Coverage for streamwise/model_provisioner/naive_baseline.py: 97%
212 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Naive baseline for the StreamWise workflow allocation problem.
3"""
5from __future__ import annotations
7from typing import Optional
9from constants import NUM_GPUS_PER_SERVER
10from constants import DEVICE_OPTIONS
12from sim_types import Result
13from sim_types import GPUType
14from sim_types import WorkflowConfig
15from sim_types import LatencyData
16from sim_types import PowerData
17from sim_types import Policy
18from sim_types import Solver
19from sim_types import Model
20from sim_types import ModelAllocation
21from sim_types import Objective
23from models import FluxModelAllocation
24from models import GemmaModelAllocation
25from models import HFModelAllocation
26from models import HFVAEModelAllocation
27from models import FTModelAllocation
28from models import FTVAEModelAllocation
29from models import UpscalerModelAllocation
30from models import OthersModelAllocation
32from evaluator import evaluate_model_allocation
34from .policies import NAIVE_POLICY
35from .policies import MAX_DEVICES
37from model_allocator import ModelAllocator
40class NaiveAllocator(ModelAllocator):
41 """
42 Naive allocator that implements a simple heuristic.
43 """
44 def __init__(
45 self,
46 workflow: WorkflowConfig,
47 latency_data: LatencyData,
48 power_data: Optional[PowerData] = None,
49 policy: Policy = NAIVE_POLICY,
50 ) -> None:
51 super().__init__(
52 workflow,
53 latency_data,
54 power_data,
55 policy,
56 )
57 assert self.policy.solver == Solver.NAIVE
58 assert self.policy.objective == Objective.TTFF
60 def allocate(
61 self,
62 num_gpus: dict[GPUType, int],
63 verbose: bool = False,
64 ) -> Result:
65 total_gpus = sum(num_gpus.values())
66 assert total_gpus >= 8, f"Total number of GPUs must be at least 8 ({num_gpus})"
68 gpu_types = [
69 gpu_type
70 for gpu_type, count in num_gpus.items()
71 if count > 0
72 ]
73 assert 1 <= len(gpu_types) <= 2, f"Only up to two GPU types are supported ({len(gpu_types)})"
74 gpu_type1 = gpu_types[0]
76 if len(gpu_types) == 1:
77 models = self._naive_single(
78 num_gpus.get(gpu_type1, 0),
79 gpu_type=gpu_type1,
80 )
81 else:
82 # Mixed setup of GPU types (e.g., A100 and H100)
83 models = self._naive_two(num_gpus)
85 result = evaluate_model_allocation(
86 models=models,
87 num_gpus=num_gpus,
88 workflow=self.workflow,
89 latency_data=self.latency_data,
90 power_data=self.power_data,
91 policy=self.policy,
92 round_up_cost_to_server=True,
93 )
94 return result
96 def _naive_single(
97 self,
98 num_gpus: int,
99 gpu_type: GPUType,
100 ) -> dict[GPUType, dict[Model, list[ModelAllocation]]]:
101 """Naive allocation for single GPU type."""
102 return self._naive_parallelism_allocation(gpu_type, num_gpus)
104 def _naive_two(
105 self,
106 num_gpus: dict[GPUType, int],
107 ) -> dict[GPUType, dict[Model, list[ModelAllocation]]]:
108 """Naive allocation for two GPU types."""
109 gpu_types = list(num_gpus.keys())
110 assert len(gpu_types) == 2
111 assert len(num_gpus) == 2
112 gpu_type1 = gpu_types[0]
113 gpu_type2 = gpu_types[1]
114 assert num_gpus[gpu_type1] >= NUM_GPUS_PER_SERVER[gpu_type1]
115 assert num_gpus[gpu_type2] >= NUM_GPUS_PER_SERVER[gpu_type2]
117 # Initialize allocations with minimal setup
118 models: dict[GPUType, dict[Model, list[ModelAllocation]]] = {
119 gpu_type1: { # 3 x A100s (type1)
120 Model.GEMMA: [GemmaModelAllocation(
121 gpu_type=gpu_type1,
122 devices=1, replicas=1)],
123 Model.FLUX: [FluxModelAllocation(
124 gpu_type=gpu_type1,
125 devices=1, replicas=1)],
126 Model.HF: [],
127 Model.HF_VAE: [],
128 Model.FT: [],
129 Model.FT_VAE: [],
130 Model.UPSCALER: [],
131 Model.OTHERS: [OthersModelAllocation(
132 gpu_type=gpu_type1,
133 devices=1, replicas=1)], # + 1 for Kokoro/YOLO
134 },
135 gpu_type2: { # 4 (+1) X H100 GPUs (type2)
136 Model.GEMMA: [],
137 Model.FLUX: [],
138 Model.HF: [HFModelAllocation(
139 gpu_type=gpu_type2,
140 devices=1, replicas=1)],
141 Model.HF_VAE: [HFVAEModelAllocation(
142 gpu_type=gpu_type2,
143 devices=1, replicas=1)],
144 Model.FT: [FTModelAllocation(
145 gpu_type=gpu_type2,
146 devices=2, replicas=1)],
147 Model.FT_VAE: [FTVAEModelAllocation(
148 gpu_type=gpu_type2,
149 devices=1, replicas=1)],
150 Model.UPSCALER: [UpscalerModelAllocation(
151 gpu_type=gpu_type2)],
152 Model.OTHERS: [],
153 },
154 }
156 # Calculate remaining: starting - assigned
157 if not self.policy.is_disaggregated(Model.HF):
158 models[gpu_type2][Model.HF][0].replicas = 2
159 models[gpu_type2][Model.HF_VAE][0].replicas = 0
160 if not self.policy.is_disaggregated(Model.FT):
161 models[gpu_type2][Model.FT_VAE][0].replicas = 0
163 if self.policy.use_upscaler:
164 models[gpu_type2][Model.UPSCALER][0].replicas = 1
166 models_gpu_type1 = self._naive_parallelism_allocation(
167 gpu_type1,
168 num_gpus.get(gpu_type1, 0),
169 )
170 models_gpu_type2 = self._naive_parallelism_allocation(
171 gpu_type2,
172 num_gpus.get(gpu_type2, 0),
173 # Already allocated in first GPU type
174 skip_non_paralelizable_models=True,
175 )
176 models[gpu_type1] = models_gpu_type1[gpu_type1]
177 models[gpu_type2] = models_gpu_type2[gpu_type2]
179 # Apply per-GPU-type overrides after allocation
180 if self.policy.use_upscaler:
181 models[gpu_type2][Model.UPSCALER][0].replicas = 1
183 return models
185 def _naive_parallelism_allocation(
186 self,
187 gpu_type: GPUType,
188 num_devices: int,
189 skip_non_paralelizable_models: bool = False,
190 ) -> dict[GPUType, dict[Model, list[ModelAllocation]]]:
191 """
192 Device allocation for naive parallelism.
193 Max devices for each model.
194 Allocate devices to each model proportional to their max devices.
195 """
196 models: dict[GPUType, dict[Model, list[ModelAllocation]]] = {
197 gpu_type: {
198 Model.GEMMA: [GemmaModelAllocation(
199 gpu_type=gpu_type,
200 replicas=1)],
201 Model.FLUX: [FluxModelAllocation(
202 gpu_type=gpu_type,
203 replicas=1)],
204 Model.HF: [HFModelAllocation(
205 gpu_type=gpu_type,
206 replicas=1)],
207 Model.HF_VAE: [HFVAEModelAllocation(
208 gpu_type=gpu_type,
209 replicas=1 if self.policy.is_disaggregated(Model.HF) else 0)],
210 Model.FT: [FTModelAllocation(
211 gpu_type=gpu_type,
212 replicas=4)],
213 Model.FT_VAE: [FTVAEModelAllocation(
214 gpu_type=gpu_type,
215 replicas=1 if self.policy.is_disaggregated(Model.FT) else 0)],
216 Model.OTHERS: [OthersModelAllocation(
217 gpu_type=gpu_type,
218 replicas=1)], # + 1 for Kokoro/YOLO
219 Model.UPSCALER: [UpscalerModelAllocation(
220 gpu_type=gpu_type,
221 replicas=1 if self.policy.use_upscaler else 0)],
222 },
223 }
225 # Zero out replicas for models not in workflow
226 for model in Model:
227 if model not in self.workflow.models:
228 for alloc in models[gpu_type][model]:
229 alloc.replicas = 0
231 # Zero out replicas for models that are not parallelizable when skip_non_paralelizable_models is True
232 if skip_non_paralelizable_models:
233 for model in Model:
234 if not self.workflow.is_parallelizable(model):
235 for alloc in models[gpu_type][model]:
236 alloc.replicas = 0
238 # Assert only 1 allocation instance per model for naive parallelism
239 for model in Model:
240 assert len(models[gpu_type][model]) == 1, \
241 f"Expected only 1 allocation instance for {model}, got {len(models[gpu_type][model])}"
243 alloc_id = 0
244 model_gemma = models[gpu_type][Model.GEMMA][alloc_id]
245 model_flux = models[gpu_type][Model.FLUX][alloc_id]
246 model_hf = models[gpu_type][Model.HF][alloc_id]
247 model_vae = models[gpu_type][Model.HF_VAE][alloc_id]
248 model_ft = models[gpu_type][Model.FT][alloc_id]
249 model_ft_vae = models[gpu_type][Model.FT_VAE][alloc_id]
250 model_upscaler = models[gpu_type][Model.UPSCALER][alloc_id]
252 # TODO do we need to do something for Model.OTHERS
254 if num_devices == 8:
255 # single server case, use fixed allocation
256 if Model.FT in self.workflow.models:
257 model_ft.replicas = 4
258 if self.policy.use_upscaler and Model.UPSCALER in self.workflow.models:
259 model_upscaler.replicas = 1
260 if Model.FT in self.workflow.models:
261 model_ft.replicas -= 1
262 if self.policy.is_disaggregated(Model.HF) and Model.HF_VAE in self.workflow.models:
263 model_vae.replicas = 1
264 if Model.FT in self.workflow.models:
265 model_ft.replicas -= 1
266 if self.policy.is_disaggregated(Model.FT) and Model.FT_VAE in self.workflow.models:
267 model_ft_vae.replicas = 1
268 if Model.FT in self.workflow.models:
269 model_ft.replicas -= 1
270 return models
272 init_num_devices = sum([
273 model[0].devices * model[0].replicas
274 for model in models[gpu_type].values()
275 ])
277 # Allocate devices proportional to each model's max devices
278 max_devices = MAX_DEVICES
279 models_in_workflow = [
280 model
281 for model in max_devices.keys()
282 if model in self.workflow.models
283 ]
284 if skip_non_paralelizable_models:
285 for model in max_devices.keys():
286 if not self.workflow.is_parallelizable(model):
287 models_in_workflow.remove(model)
289 total_max_devices = sum([
290 max_devices[model]
291 for model in models_in_workflow
292 ])
293 for model in models_in_workflow:
294 # Calculate the number of devices to allocate for the model, proportional to its max devices among others
295 alloc_devices = int((num_devices - init_num_devices) * max_devices[model] / total_max_devices)
296 if model == Model.GEMMA:
297 max_devices_gemma = max_devices[Model.GEMMA]
298 if self.latency_data:
299 max_devices_gemma = min(max_devices_gemma, self.latency_data[gpu_type].get_max_parallelism(model))
300 model_gemma.devices += min(alloc_devices, max_devices_gemma)
301 # Round down nearest in DEVICE_OPTIONS_GEMMA
302 num_gemma_devices = max([
303 d
304 for d in DEVICE_OPTIONS[Model.GEMMA]
305 if d <= model_gemma.devices
306 ])
307 model_gemma.devices = num_gemma_devices
308 elif model == Model.FLUX:
309 max_devices_flux = max_devices[Model.FLUX]
310 if self.latency_data:
311 max_devices_flux = min(max_devices_flux, self.latency_data[gpu_type].get_max_parallelism(model))
312 model_flux.devices += min(alloc_devices, max_devices_flux)
313 # Round down nearest in DEVICE_OPTIONS_FLUX
314 model_flux.devices = max([
315 d
316 for d in DEVICE_OPTIONS[Model.FLUX]
317 if d <= model_flux.devices
318 ])
319 elif model == Model.HF:
320 max_devices_hf = max_devices[Model.HF]
321 if self.latency_data:
322 max_devices_hf = min(max_devices_hf, self.latency_data[gpu_type].get_max_parallelism(model))
323 model_hf.replicas += min(alloc_devices, max_devices_hf)
324 elif model == Model.HF_VAE:
325 if self.policy.is_disaggregated(Model.HF):
326 max_devices_vae = max_devices[Model.HF_VAE]
327 if self.latency_data:
328 max_devices_vae = min(max_devices_vae, self.latency_data[gpu_type].get_max_parallelism(model))
329 model_vae.replicas += min(alloc_devices, max_devices_vae)
330 elif model == Model.FT:
331 max_devices_ft = max_devices[Model.FT]
332 if self.latency_data:
333 max_devices_ft = min(max_devices_ft, self.latency_data[gpu_type].get_max_parallelism(model))
334 model_ft.replicas += min(alloc_devices, max_devices_ft)
335 elif model == Model.FT_VAE:
336 if self.policy.is_disaggregated(Model.FT):
337 max_devices_ft_vae = max_devices[Model.FT_VAE]
338 if self.latency_data:
339 max_devices_ft_vae = min(
340 max_devices_ft_vae, self.latency_data[gpu_type].get_max_parallelism(model)
341 )
342 model_ft_vae.replicas += min(alloc_devices, max_devices_ft_vae)
343 else:
344 raise ValueError(f"Unrecognized model {model}")
346 remaining_devices = num_devices
347 for model_name in models[gpu_type].keys():
348 for model_alloc in models[gpu_type][model_name]:
349 remaining_devices -= model_alloc.get_num_gpus()
351 # Distribute remaining devices to parallelizable models
352 distribute_models = self.workflow.filter_parallelizable_models(
353 models_in_workflow,
354 disaggregation=self.policy.disaggregation,
355 )
356 # Prioritise models that already hold more GPUs
357 distribute_models.sort(
358 key=lambda m: models[gpu_type][m][alloc_id].get_num_gpus(),
359 reverse=True,
360 )
361 num_distribute = len(distribute_models)
362 if num_distribute > 0 and remaining_devices > 0:
363 made_progress = True
364 while remaining_devices > 0 and made_progress:
365 made_progress = False
366 for model_name in distribute_models:
367 gpus_per_replica = models[gpu_type][model_name][alloc_id].devices
368 if gpus_per_replica <= 0 or remaining_devices < gpus_per_replica:
369 continue
370 models[gpu_type][model_name][alloc_id].replicas += 1
371 remaining_devices -= gpus_per_replica
372 made_progress = True
373 if remaining_devices <= 0:
374 break
376 remaining_devices = num_devices
377 for model_name in models[gpu_type].keys():
378 for model_alloc in models[gpu_type][model_name]:
379 remaining_devices -= model_alloc.get_num_gpus()
381 # TODO we should try to assign all resources
382 # assert remaining_devices == 0, \
383 assert remaining_devices >= 0, \
384 f"remaining={remaining_devices} != 0: " \
385 f"gpu={gpu_type.value} total={num_devices} remaining={remaining_devices}"
387 # Update replicas based on total devices
388 # Gemma (when parallelizable)
389 if self.workflow.is_parallelizable(Model.GEMMA) and Model.GEMMA in models_in_workflow:
390 model_gemma.devices, model_gemma.replicas, remaining_devices = _calculate_naive_num_devices(
391 model_gemma.devices,
392 model_gemma.replicas,
393 remaining_devices,
394 device_options=DEVICE_OPTIONS[Model.GEMMA],
395 replica_upper_bound=self.workflow.total_scenes)
397 # Flux (when parallelizable)
398 if self.workflow.is_parallelizable(Model.FLUX) and Model.FLUX in models_in_workflow:
399 model_flux.devices, model_flux.replicas, remaining_devices = _calculate_naive_num_devices(
400 model_flux.devices,
401 model_flux.replicas,
402 remaining_devices,
403 device_options=DEVICE_OPTIONS[Model.FLUX],
404 replica_upper_bound=self.workflow.total_scenes)
406 # Hunyuan FramePack
407 if Model.HF in self.workflow.models:
408 model_hf.devices, model_hf.replicas, remaining_devices = _calculate_naive_num_devices(
409 model_hf.devices,
410 model_hf.replicas,
411 remaining_devices,
412 device_options=DEVICE_OPTIONS[Model.HF],
413 replica_upper_bound=self.workflow.total_scenes)
415 # Hunyuan FramePack VAE
416 if self.policy.is_disaggregated(Model.HF) and Model.HF_VAE in self.workflow.models:
417 model_vae.devices, model_vae.replicas, remaining_devices = _calculate_naive_num_devices(
418 model_vae.devices,
419 model_vae.replicas,
420 remaining_devices,
421 device_options=None,
422 replica_upper_bound=self.workflow.total_frames[Model.HF],
423 )
425 # Fantasy Talking
426 if Model.FT in self.workflow.models:
427 model_ft.devices, model_ft.replicas, remaining_devices = _calculate_naive_num_devices(
428 model_ft.devices,
429 model_ft.replicas,
430 remaining_devices,
431 device_options=DEVICE_OPTIONS[Model.FT],
432 replica_upper_bound=self.workflow.total_subscenes,
433 )
435 # Fantasy Talking VAE
436 if self.policy.is_disaggregated(Model.FT) and Model.FT_VAE in self.workflow.models:
437 model_ft_vae.devices, model_ft_vae.replicas, remaining_devices = _calculate_naive_num_devices(
438 model_ft_vae.devices,
439 model_ft_vae.replicas,
440 remaining_devices,
441 device_options=None,
442 replica_upper_bound=self.workflow.total_frames[Model.FT],
443 )
445 return models
448def _calculate_naive_num_devices(
449 num_devices: int,
450 num_replicas: int,
451 remaining_devices: int,
452 device_options: Optional[list[int]] = [1],
453 replica_upper_bound: Optional[int] = None,
454) -> tuple[int, int, int]:
455 """Find the parallelism that maximizes the device usage."""
456 assert remaining_devices >= 0
458 model_quota = num_devices * num_replicas
460 if device_options:
461 best_product = 0
462 best_devices_per_replica = 1
463 best_replicas = 1
464 for devices_per_replica in device_options:
465 if devices_per_replica > model_quota:
466 continue
467 max_replicas = model_quota // devices_per_replica
468 if replica_upper_bound and max_replicas > replica_upper_bound:
469 max_replicas = replica_upper_bound
470 product = devices_per_replica * max_replicas
471 if product > best_product:
472 best_product = product
473 best_devices_per_replica = devices_per_replica
474 best_replicas = max_replicas
475 else:
476 # start with parallelism=1 instead
477 best_devices_per_replica = 1
478 best_replicas = model_quota
480 num_devices = best_devices_per_replica
481 num_replicas = best_replicas
482 remaining_devices += model_quota - num_replicas * num_devices
484 return num_devices, num_replicas, remaining_devices