Coverage for streamwise/model_provisioner/milp.py: 84%
369 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2MILP formulation for the StreamWise workflow allocation problem.
3"""
5from __future__ import annotations
7import json
8import logging
10from typing import Callable
11from typing import Optional
13from pyomo.environ import ConcreteModel
14from pyomo.environ import Var
15from pyomo.environ import Set
16from pyomo.environ import Objective as OptObjective
17from pyomo.environ import Binary
18from pyomo.environ import NonNegativeIntegers
19from pyomo.environ import NonNegativeReals
20from pyomo.environ import minimize
21from pyomo.environ import SolverFactory
22from pyomo.environ import ConstraintList
24from sim_types import GPUType
25from sim_types import Model
26from sim_types import WorkflowConfig
27from sim_types import LatencyData
28from sim_types import PowerData
29from sim_types import Result
30from sim_types import Policy
31from sim_types import ModelAllocation
32from sim_types import Objective
33from sim_types import Solver
35from models import get_model_allocation
37from model_allocator import ModelAllocator
39from constants import DEVICE_OPTIONS
40from constants import NUM_GPUS_PER_SERVER
41from constants import SECONDS_IN_HOUR
43from .policies import STREAMWISE_MILP_POLICY
46MAX_INSTANCES = 16
48# Maximum time it can take: 24 hours in seconds
49# Used for big-M constraints to link TTFF and makespan to instance variables
50MAX_TIME = 24 * SECONDS_IN_HOUR
53# Allocators that require quadratic (bilinear) objectives - need Gurobi
54QUADRATIC_OBJECTIVES = [
55 Objective.TTFF_COST,
56 Objective.TIME_ENERGY,
57 Objective.ENERGY_COST,
58]
61def idx(
62 gpu_type: GPUType,
63 model_name: Model,
64 instance_id: int
65) -> tuple[str, str, int]:
66 """Helper to convert enum to index key for instance variables."""
67 return (gpu_type.value, model_name.value, instance_id)
70def dev_idx(
71 gpu_type: GPUType,
72 model_name: Model,
73 instance_id: int,
74 num_devices: int
75) -> tuple[str, str, int, int]:
76 """Helper to convert enum to index key for device variables."""
77 return (gpu_type.value, model_name.value, instance_id, num_devices)
80class MILPAllocator(ModelAllocator):
81 """
82 MILP-based allocator that computes the optimal model allocation.
83 """
84 def __init__(
85 self,
86 workflow: WorkflowConfig,
87 latency_data: LatencyData,
88 power_data: Optional[PowerData] = None,
89 policy: Policy = STREAMWISE_MILP_POLICY,
90 ) -> None:
91 super().__init__(
92 workflow,
93 latency_data,
94 power_data,
95 policy,
96 )
97 assert self.policy.solver in [Solver.GUROBI, Solver.HIGHS]
99 def allocate(
100 self,
101 num_gpus: dict[GPUType, int],
102 verbose: bool = False,
103 running_cost: bool = False, # If True, cost = active time only; False = makespan x GPUs
104 max_cost: Optional[float] = None, # If set, adds a constraint to limit cost
105 max_ttff: Optional[float] = None, # If set, adds a constraint to limit TTFF
106 max_makespan: Optional[float] = None, # If set, adds a constraint to limit makespan
107 time_limit: Optional[int] = None, # Time limit for the solver in seconds
108 save_solution_path: Optional[str] = None, # If set, saves the solution to a JSON file
109 warm_start_path: Optional[str] = None, # If set, loads a warm start solution from a JSON file
110 force_num_gpus: bool = False, # If True, adds constraints to force the use of all available GPUs
111 skip_server_constraint: bool = False, # If True, skips the GPU-per-server constraint
112 ) -> Result:
113 """
114 Calculate the optimal model allocation and resulting metrics using MILP formulation.
115 """
116 m = ConcreteModel()
118 # Options: "gurobi", "highs"
119 solver_name = self.policy.solver.value
121 # Define index sets
122 gpu_types = list(num_gpus.keys())
124 model_names = [
125 Model.GEMMA,
126 Model.FLUX,
127 Model.HF,
128 # Model.HF_VAE,
129 Model.FT,
130 # Model.FT_VAE,
131 # Model.UPSCALER,
132 Model.OTHERS,
133 ]
134 if self.policy.use_upscaler:
135 model_names.append(Model.UPSCALER)
136 if self.policy.is_disaggregated(Model.HF):
137 model_names.append(Model.HF_VAE)
138 if self.policy.is_disaggregated(Model.FT):
139 model_names.append(Model.FT_VAE)
141 # Remove models not in the workflow
142 model_names = [
143 model_name
144 for model_name in model_names
145 if model_name in self.workflow.models
146 ]
148 instance_ids = list(range(MAX_INSTANCES))
150 # The units of work that each model has to do
151 work: dict[Model, int] = self.workflow.work
153 # Create Pyomo Sets
154 m.GPU_TYPES = Set(initialize=[g.value for g in gpu_types])
155 m.MODEL_NAMES = Set(initialize=[mn.value for mn in model_names])
156 m.INSTANCES = Set(initialize=instance_ids)
158 # Create index set for device choices: (gpu_type, model_name, instance_id, device_count)
159 device_index_set = [
160 (gpu_type.value, model_name.value, instance_id, num_devices)
161 for gpu_type in gpu_types
162 for model_name in model_names
163 for instance_id in instance_ids
164 for num_devices in [0] + DEVICE_OPTIONS[model_name]
165 ]
166 m.DEVICE_INDEX = Set(initialize=device_index_set)
168 # Create index set for instance variables: (gpu_type, model_name, instance_id)
169 instance_index_set = [
170 (gpu_type.value, model_name.value, instance_id)
171 for gpu_type in gpu_types
172 for model_name in model_names
173 for instance_id in instance_ids
174 ]
175 m.INSTANCE_INDEX = Set(initialize=instance_index_set)
177 # Define indexed variables
178 m.device_choice = Var(m.DEVICE_INDEX, domain=Binary)
179 m.work_device = Var(m.DEVICE_INDEX, domain=NonNegativeIntegers) # Linearization: work per device choice
180 m.gpus = Var(m.INSTANCE_INDEX, domain=NonNegativeIntegers)
181 m.is_active = Var(m.INSTANCE_INDEX, domain=Binary)
182 m.is_min = Var(m.INSTANCE_INDEX, domain=Binary)
183 m.work = Var(m.INSTANCE_INDEX, domain=NonNegativeIntegers)
184 m.time = Var(m.INSTANCE_INDEX, domain=NonNegativeReals)
185 m.ttff = Var(m.INSTANCE_INDEX, domain=NonNegativeReals)
187 # Objective variables
188 m.makespan = Var(domain=NonNegativeReals)
189 m.ttff_user = Var(domain=NonNegativeReals)
190 m.ttff_min = Var(m.MODEL_NAMES, domain=NonNegativeReals) # Per-model minimum TTFF
191 m.time_max = Var(m.MODEL_NAMES, domain=NonNegativeReals) # Per-model maximum time
192 m.cost = Var(domain=NonNegativeReals)
193 m.energy = Var(domain=NonNegativeReals)
195 # Constraint list for dynamic constraints
196 m.constraints = ConstraintList()
198 for gpu_type in gpu_types:
199 for model_name in model_names:
200 for instance_id in instance_ids:
201 key = idx(gpu_type, model_name, instance_id)
203 # GPUs used = sum of num_devices * device_choice[num_devices]
204 m.constraints.add(
205 m.gpus[key] == sum(
206 num_devices * m.device_choice[dev_idx(gpu_type, model_name, instance_id, num_devices)]
207 for num_devices in [0] + DEVICE_OPTIONS[model_name]
208 )
209 )
211 # Cannot select inactive instance as min
212 m.constraints.add(m.is_min[key] <= m.is_active[key])
213 # If active = 0 -> GPUs = 0
214 m.constraints.add(m.gpus[key] <= num_gpus[gpu_type] * m.is_active[key])
215 # If active = 1 -> GPUs ≥ 1
216 m.constraints.add(m.gpus[key] >= m.is_active[key])
217 # If work = 0 -> active = 0 -> GPUs = 0
218 m.constraints.add(m.is_active[key] <= m.work[key])
220 # If device = 0 -> work = 0
221 dev_idx_0 = dev_idx(gpu_type, model_name, instance_id, 0)
222 m.constraints.add(
223 m.work[key]
224 <= work[model_name] * (1 - m.device_choice[dev_idx_0])
225 )
227 # Linearization: work_device links device_choice and work
228 # work = sum(work_device[d] for d in devices) - excludes 0 GPUs since they can't do work
229 m.constraints.add(
230 m.work[key] == sum(
231 m.work_device[dev_idx(gpu_type, model_name, instance_id, num_devices)]
232 for num_devices in DEVICE_OPTIONS[model_name]
233 )
234 )
235 # If any non-zero device is selected, work must be >= 1
236 m.constraints.add(
237 m.work[key] >= sum(
238 m.device_choice[dev_idx(gpu_type, model_name, instance_id, num_devices)]
239 for num_devices in DEVICE_OPTIONS[model_name]
240 )
241 )
242 # work_device[d] <= TOTAL_WORK * device_choice[d]
243 for num_devices in [0] + DEVICE_OPTIONS[model_name]:
244 didx = dev_idx(gpu_type, model_name, instance_id, num_devices)
245 m.constraints.add(
246 m.work_device[didx] <= work[model_name] * m.device_choice[didx]
247 )
249 # Link instance time to per-model max time
250 m.constraints.add(m.time[key] <= m.time_max[model_name.value])
252 # Link TTFF to per-model TTFF min
253 # If selected → ttff_min[model] == ttff_var
254 m.constraints.add(m.ttff_min[model_name.value] >= m.ttff[key] - MAX_TIME * (1 - m.is_min[key]))
255 m.constraints.add(m.ttff_min[model_name.value] <= m.ttff[key] + MAX_TIME * (1 - m.is_active[key]))
257 # One device per instance
258 for instance_id in instance_ids:
259 m.constraints.add(
260 sum(
261 m.device_choice[dev_idx(gpu_type, model_name, instance_id, num_devices)]
262 for num_devices in [0] + DEVICE_OPTIONS[model_name]
263 ) == 1
264 )
266 # Symmetry breaking (fill earlier instances first)
267 for instance_id in range(MAX_INSTANCES - 1):
268 m.constraints.add(
269 m.gpus[idx(gpu_type, model_name, instance_id)]
270 >= m.gpus[idx(gpu_type, model_name, instance_id + 1)]
271 )
273 # Makespan is the sum of max times per model (models run sequentially)
274 m.constraints.add(m.makespan == sum(m.time_max[model_name.value] for model_name in model_names))
276 # User TTFF definition: sum of min TTFF per model
277 m.constraints.add(m.ttff_user >= sum(m.ttff_min[model_name.value] for model_name in model_names))
278 m.constraints.add(m.ttff_user >= m.makespan - self.workflow.total_video_seconds)
280 # Select exactly 1 instance as the min TTFF instance per model
281 for model_name in model_names:
282 m.constraints.add(
283 sum(
284 m.is_min[idx(gpu_type, model_name, instance_id)]
285 for gpu_type in gpu_types
286 for instance_id in instance_ids
287 ) == 1
288 )
290 # Resolution scaling factor for HF/VAE/FT
291 latency_ratio = self.workflow.get_resolution_scale(self.policy.use_upscaler)
293 # Time constraints
294 # Each model block is guarded by membership in model_names so that
295 # the MILP can be built for a subset of models (e.g. Helix per-model).
296 for gpu_type in gpu_types:
297 # Gemma
298 if Model.GEMMA in model_names and work[Model.GEMMA] > 0:
299 model_name = Model.GEMMA
300 for instance_id in instance_ids:
301 key = idx(gpu_type, model_name, instance_id)
302 # Makespan is the max time across all instances
303 # Linearized: use work_device instead of device_choice * work
304 if work[model_name] > 1:
305 # Parallel: each work unit = 1 scene
306 # Time for w scenes
307 # = gemma_first_scene + gemma_per_scene * (w - 1)
308 # = (gemma_first_scene - gemma_per_scene) * is_active + gemma_per_scene * work
309 # Using linearized variables:
310 # = (gemma_first_scene[d] - gemma_per_scene[d]) * \
311 # device_choice[d] + gemma_per_scene[d] * work_device[d]
312 m.constraints.add(
313 m.time[key] == sum(
314 (
315 self.latency_data[gpu_type].gemma_first_scene[num_devices]
316 - self.latency_data[gpu_type].gemma_per_scene[num_devices]
317 )
318 * m.device_choice[dev_idx(gpu_type, model_name, instance_id, num_devices)]
319 + self.latency_data[gpu_type].gemma_per_scene[num_devices]
320 * m.work_device[dev_idx(gpu_type, model_name, instance_id, num_devices)]
321 for num_devices in DEVICE_OPTIONS[model_name]
322 )
323 )
324 else:
325 m.constraints.add(
326 m.time[key] == sum(
327 (
328 self.latency_data[gpu_type].gemma_first_scene[num_devices]
329 + self.latency_data[gpu_type].gemma_per_scene[num_devices]
330 * (self.workflow.total_scenes - 1)
331 )
332 * m.work_device[dev_idx(gpu_type, model_name, instance_id, num_devices)]
333 for num_devices in DEVICE_OPTIONS[model_name]
334 )
335 )
336 # TTFF is for 1 work unit
337 m.constraints.add(
338 m.ttff[key] == sum(
339 m.device_choice[dev_idx(gpu_type, model_name, instance_id, num_devices)]
340 * self.latency_data[gpu_type].gemma_first_scene[num_devices]
341 * 1 # TTFF for tokens in first scene
342 for num_devices in DEVICE_OPTIONS[model_name]
343 )
344 )
346 # Flux
347 if Model.FLUX in model_names and work[Model.FLUX] > 0:
348 model_name = Model.FLUX
349 for instance_id in instance_ids:
350 key = idx(gpu_type, model_name, instance_id)
351 # Makespan is the max time across all instances
352 # Linearized: use work_device instead of device_choice * work
353 if work[model_name] > 1:
354 # Parallel: each work unit = 1 scene
355 # Time for w scenes = latency * num_steps_flux * w
356 m.constraints.add(
357 m.time[key] == sum(
358 self.latency_data[gpu_type][model_name, num_devices]
359 * self.workflow.num_steps[model_name]
360 * m.work_device[dev_idx(gpu_type, model_name, instance_id, num_devices)]
361 for num_devices in DEVICE_OPTIONS[model_name]
362 )
363 )
364 else:
365 # Non-parallel: single work unit covers all scenes
366 m.constraints.add(
367 m.time[key] == sum(
368 self.latency_data[gpu_type][model_name, num_devices]
369 * self.workflow.num_steps[model_name]
370 * m.work_device[dev_idx(gpu_type, model_name, instance_id, num_devices)]
371 for num_devices in DEVICE_OPTIONS[model_name]
372 )
373 )
374 # TTFF is for 1 work unit
375 m.constraints.add(
376 m.ttff[key] == sum(
377 m.device_choice[dev_idx(gpu_type, model_name, instance_id, num_devices)]
378 * self.latency_data[gpu_type][model_name, num_devices]
379 * self.workflow.num_steps[model_name]
380 * 1 # TTFF for first work unit
381 for num_devices in DEVICE_OPTIONS[model_name]
382 )
383 )
385 # Hunyuan FramePack
386 if Model.HF in model_names and work[Model.HF] > 0:
387 model_name = Model.HF
388 for instance_id in instance_ids:
389 key = idx(gpu_type, model_name, instance_id)
391 """
392 from models import HFModelAllocation
393 HFModelAllocation(
394 gpu_type,
395 num_devices,
396 replicas=1,
397 )._calc_time_per_subscene(
398 self.policy,
399 self.workflow,
400 self.latency_data[gpu_type]
401 )
402 """
404 # Makespan is the max time across all instances
405 # Linearized: use work_device instead of device_choice * work
406 hf_time_expr = sum(
407 self.workflow.per_subscene_frames[model_name]
408 / self.workflow.hf_frames[self.workflow.frames_per_step_idx]
409 * self.latency_data[gpu_type][model_name, num_devices]
410 * latency_ratio
411 * self.workflow.num_steps[model_name]
412 * m.work_device[dev_idx(gpu_type, model_name, instance_id, num_devices)]
413 for num_devices in DEVICE_OPTIONS[model_name]
414 )
415 # When not disaggregated, VAE runs on the same instance
416 if not self.policy.is_disaggregated(Model.HF):
417 hf_vae_time_per_work = (
418 self.latency_data[gpu_type][Model.HF_VAE, 1]
419 * latency_ratio
420 / self.workflow.hf_frames[self.workflow.frames_per_step_idx]
421 )
422 hf_time_expr += hf_vae_time_per_work * m.work[key]
423 m.constraints.add(m.time[key] == hf_time_expr)
424 # TTFF is for first chunk (can be smaller than subscene when disaggregated)
425 ttff_frames_hf = min(
426 self.workflow.hf_frames[0],
427 self.workflow.per_subscene_frames[model_name])
428 hf_ttff_expr = sum(
429 m.device_choice[dev_idx(gpu_type, model_name, instance_id, num_devices)]
430 * ttff_frames_hf
431 / self.workflow.hf_frames[self.workflow.frames_per_step_idx]
432 * self.latency_data[gpu_type][model_name, num_devices]
433 * latency_ratio
434 * self.workflow.num_steps[model_name]
435 * 1 # TTFF for first chunk
436 for num_devices in DEVICE_OPTIONS[model_name]
437 )
438 # When not disaggregated, add VAE decode time for first chunk
439 if not self.policy.is_disaggregated(Model.HF):
440 hf_vae_ttff = (
441 ttff_frames_hf
442 / self.workflow.hf_frames[self.workflow.frames_per_step_idx]
443 * self.latency_data[gpu_type][Model.HF_VAE, 1]
444 * latency_ratio
445 )
446 hf_ttff_expr += hf_vae_ttff * m.is_active[key]
447 m.constraints.add(m.ttff[key] == hf_ttff_expr)
449 # Hunyuan FramePack VAE
450 if Model.HF_VAE in model_names and work[Model.HF_VAE] > 0:
451 model_name = Model.HF_VAE
452 for instance_id in instance_ids:
453 key = idx(gpu_type, model_name, instance_id)
454 # Makespan is the max time across all instances
455 # Linearized: use work_device instead of device_choice * work
456 m.constraints.add(
457 m.time[key] == sum(
458 self.latency_data[gpu_type][model_name, num_devices]
459 * latency_ratio
460 / self.workflow.hf_frames[self.workflow.frames_per_step_idx]
461 * m.work_device[dev_idx(gpu_type, model_name, instance_id, num_devices)]
462 for num_devices in DEVICE_OPTIONS[model_name]
463 )
464 )
465 # TTFF is for 1 subscene
466 m.constraints.add(
467 m.ttff[key] == sum(
468 m.device_choice[dev_idx(gpu_type, model_name, instance_id, num_devices)]
469 * self.workflow.per_subscene_frames[Model.HF]
470 * self.latency_data[gpu_type][model_name, num_devices]
471 * latency_ratio
472 / self.workflow.hf_frames[self.workflow.frames_per_step_idx] # frames_per_step_hf
473 * 1 # TTFF for first subscene
474 for num_devices in DEVICE_OPTIONS[model_name]
475 )
476 )
478 # Fantasy Talking
479 if Model.FT in model_names and work[Model.FT] > 0:
480 model_name = Model.FT
481 for instance_id in instance_ids:
482 key = idx(gpu_type, model_name, instance_id)
483 # Makespan is the max time across all instances
484 # Linearized: use work_device instead of device_choice * work
485 ft_time_expr = sum(
486 self.workflow.per_subscene_frames[model_name]
487 / self.workflow.ft_frames[self.workflow.frames_per_step_idx]
488 * self.latency_data[gpu_type][model_name, num_devices]
489 * latency_ratio
490 * self.workflow.num_steps[model_name]
491 * m.work_device[dev_idx(gpu_type, model_name, instance_id, num_devices)]
492 for num_devices in DEVICE_OPTIONS[model_name]
493 )
494 # When not disaggregated, VAE runs on the same instance
495 if not self.policy.is_disaggregated(Model.FT):
496 ft_vae_time_per_work = (
497 self.latency_data[gpu_type][Model.FT_VAE, 1]
498 * latency_ratio
499 / self.workflow.ft_frames[self.workflow.frames_per_step_idx]
500 )
501 ft_time_expr += ft_vae_time_per_work * m.work[key]
502 m.constraints.add(m.time[key] == ft_time_expr)
503 # TTFF is for 1 work unit (e.g., subscene)
504 ft_ttff_expr = sum(
505 m.device_choice[dev_idx(gpu_type, model_name, instance_id, num_devices)]
506 * self.workflow.per_subscene_frames[model_name]
507 / self.workflow.ft_frames[self.workflow.frames_per_step_idx]
508 * self.latency_data[gpu_type][model_name, num_devices]
509 * latency_ratio
510 * self.workflow.num_steps[model_name]
511 * 1 # TTFF for first subscene
512 for num_devices in DEVICE_OPTIONS[model_name]
513 )
514 # When not disaggregated, add VAE decode time for first subscene
515 if not self.policy.is_disaggregated(Model.FT):
516 ft_vae_ttff = (
517 self.workflow.per_subscene_frames[Model.FT]
518 / self.workflow.ft_frames[self.workflow.frames_per_step_idx]
519 * self.latency_data[gpu_type][Model.FT_VAE, 1]
520 * latency_ratio
521 )
522 ft_ttff_expr += ft_vae_ttff * m.is_active[key]
523 m.constraints.add(m.ttff[key] == ft_ttff_expr)
525 # Fantasy Talking VAE
526 if Model.FT_VAE in model_names and work[Model.FT_VAE] > 0:
527 model_name = Model.FT_VAE
528 for instance_id in instance_ids:
529 key = idx(gpu_type, model_name, instance_id)
530 # Makespan is the max time across all instances
531 # Linearized: use work_device instead of device_choice * work
532 m.constraints.add(
533 m.time[key] == sum(
534 self.latency_data[gpu_type][model_name, num_devices]
535 * latency_ratio
536 / self.workflow.ft_frames[self.workflow.frames_per_step_idx]
537 * m.work_device[dev_idx(gpu_type, model_name, instance_id, num_devices)]
538 for num_devices in DEVICE_OPTIONS[model_name]
539 )
540 )
541 # TTFF is for 1 subscene
542 m.constraints.add(
543 m.ttff[key] == sum(
544 m.device_choice[dev_idx(gpu_type, model_name, instance_id, num_devices)]
545 * self.workflow.per_subscene_frames[Model.FT]
546 * self.latency_data[gpu_type][model_name, num_devices]
547 * latency_ratio
548 / self.workflow.ft_frames[self.workflow.frames_per_step_idx] # frames_per_step_ft
549 * 1 # TTFF for first subscene
550 for num_devices in DEVICE_OPTIONS[model_name]
551 )
552 )
554 # Upscaler
555 if Model.UPSCALER in model_names and work[Model.UPSCALER] > 0 and self.policy.use_upscaler:
556 model_name = Model.UPSCALER
557 for instance_id in instance_ids:
558 key = idx(gpu_type, model_name, instance_id)
559 # Linearized: use work_device instead of device_choice * work
560 m.constraints.add(
561 m.time[key] == sum(
562 self.latency_data[gpu_type][model_name, num_devices]
563 * m.work_device[dev_idx(gpu_type, model_name, instance_id, num_devices)]
564 for num_devices in DEVICE_OPTIONS[model_name]
565 )
566 )
567 # TTFF is for 1 work unit (e.g., subscene)
568 m.constraints.add(
569 m.ttff[key] == sum(
570 m.device_choice[dev_idx(gpu_type, model_name, instance_id, num_devices)]
571 * self.latency_data[gpu_type][model_name, num_devices]
572 * self.workflow.per_subscene_frames[Model.FT]
573 * 1 # TTFF is for first subscene
574 for num_devices in DEVICE_OPTIONS[model_name]
575 )
576 )
578 # Others
579 if Model.OTHERS in model_names and work[Model.OTHERS] > 0:
580 model_name = Model.OTHERS
581 for instance_id in instance_ids:
582 key = idx(gpu_type, model_name, instance_id)
583 # Makespan is the max time across all instances
584 m.constraints.add(
585 m.time[key] == sum(
586 m.device_choice[dev_idx(gpu_type, model_name, instance_id, num_devices)]
587 * self.latency_data[gpu_type][model_name, num_devices]
588 * self.workflow.total_scenes
589 for num_devices in DEVICE_OPTIONS[model_name]
590 )
591 )
592 # TTFF is for 1 work unit
593 m.constraints.add(
594 m.ttff[key] == sum(
595 m.device_choice[dev_idx(gpu_type, model_name, instance_id, num_devices)]
596 * self.latency_data[gpu_type][model_name, num_devices]
597 * 1 # TTFF is for first scene
598 for num_devices in DEVICE_OPTIONS[model_name]
599 )
600 )
602 # Total work to do for each model
603 for model_name in model_names:
604 m.constraints.add(
605 sum(
606 m.work[idx(gpu_type, model_name, instance_id)]
607 for gpu_type in gpu_types
608 for instance_id in instance_ids
609 ) == work[model_name]
610 )
612 # Number of GPUs per type
613 # Add a variable to represent the number of servers for each GPU type
614 m.num_servers = Var(m.GPU_TYPES, domain=NonNegativeIntegers)
616 for gpu_type in gpu_types:
617 total_gpus = sum(
618 m.gpus[idx(gpu_type, model_name, instance_id)]
619 for model_name in model_names
620 for instance_id in instance_ids
621 )
622 if force_num_gpus:
623 m.constraints.add(total_gpus == num_gpus[gpu_type])
624 else:
625 m.constraints.add(total_gpus <= num_gpus[gpu_type])
627 # GPUs used must be a multiple of NUM_GPUS_PER_SERVER
628 if not skip_server_constraint:
629 m.constraints.add(total_gpus == m.num_servers[gpu_type.value] * NUM_GPUS_PER_SERVER[gpu_type])
631 # Cost calculation
632 # running_cost=True: cost based only on active model running time
633 if running_cost:
634 cost_expr = sum(
635 self._get_latency_per_work(
636 gpu_type,
637 model_name,
638 num_devices,
639 )
640 * num_devices
641 * m.work_device[dev_idx(gpu_type, model_name, instance_id, num_devices)]
642 * self.policy.gpu_cost[gpu_type] / SECONDS_IN_HOUR
643 for gpu_type in gpu_types
644 for model_name in model_names
645 for instance_id in instance_ids
646 for num_devices in DEVICE_OPTIONS[model_name]
647 )
648 # running_cost=False: cost = makespan × total_GPUs_used (GPUs allocated for full job duration)
649 else:
650 cost_expr = m.makespan * sum(
651 m.gpus[idx(gpu_type, model_name, instance_id)]
652 * self.policy.gpu_cost[gpu_type] / SECONDS_IN_HOUR
653 for gpu_type in gpu_types
654 for model_name in model_names
655 for instance_id in instance_ids
656 )
657 m.constraints.add(m.cost == cost_expr)
659 # Energy: model-specific power * active time + idle power * (makespan - active time)
660 if self.power_data is None:
661 energy_expr = 0.0
662 else:
663 # Active energy: Use model-specific power values (not TDP)
664 energy_expr = sum(
665 self._get_latency_per_work(
666 gpu_type,
667 model_name,
668 num_devices,
669 )
670 * num_devices
671 * m.work_device[dev_idx(gpu_type, model_name, instance_id, num_devices)]
672 * (
673 self._get_power_per_work(
674 gpu_type,
675 model_name,
676 num_devices,
677 ) - self.power_data[gpu_type]["idle"]
678 )
679 for gpu_type in gpu_types
680 for model_name in model_names
681 for instance_id in instance_ids
682 for num_devices in DEVICE_OPTIONS[model_name]
683 )
684 # Idle energy: idle power * num_gpus * makespan
685 energy_expr += sum(
686 self.power_data[gpu_type]["idle"] * num_gpus[gpu_type] * m.makespan
687 for gpu_type in gpu_types
688 )
689 m.constraints.add(m.energy == energy_expr)
691 # Bounds
692 if max_cost is not None:
693 m.constraints.add(m.cost <= max_cost)
694 if max_ttff is not None:
695 m.constraints.add(m.ttff_user <= max_ttff)
696 if max_makespan is not None:
697 m.constraints.add(m.makespan <= max_makespan)
699 # Objective functions
700 obj = get_objective(
701 m=m,
702 allocator=self.policy.objective,
703 solver_name=solver_name,
704 )
705 if obj is not None:
706 m.objective = obj
708 # Solve
709 solver = SolverFactory(solver_name)
710 if solver_name == "gurobi" and time_limit:
711 solver.options["TimeLimit"] = time_limit
712 if solver_name == "highs" and time_limit:
713 solver.options["time_limit"] = time_limit
714 if self.policy.objective in QUADRATIC_OBJECTIVES and solver_name == "gurobi":
715 solver.options['NonConvex'] = 2 # Option for bilinear objectives
716 if solver_name == "highs":
717 solver.options["time_limit"] = 50 # seconds
719 if warm_start_path is not None:
720 _load_warm_start(m, warm_start_path)
722 if solver_name == "gurobi":
723 opt_result = solver.solve(
724 m,
725 tee=verbose,
726 warmstart=warm_start_path is not None,
727 )
728 else:
729 opt_result = solver.solve(m, tee=verbose)
731 if opt_result.solver.status != "ok":
732 logging.error(f"Solver failed with status: {opt_result.solver.status}")
734 if save_solution_path is not None:
735 _save_solution(m, save_solution_path)
737 models = milp_to_models_dict(
738 m=m,
739 gpu_types=gpu_types,
740 model_names=model_names,
741 instance_ids=instance_ids,
742 idx=idx,
743 workflow=self.workflow,
744 power_data=self.power_data,
745 policy=self.policy,
746 )
748 if not self._is_valid_result(m):
749 return Result()
751 tbf_s = 0.0
752 if m.makespan.value and self.workflow.num_frames > 0:
753 tbf_s = m.makespan.value / self.workflow.num_frames
754 return Result(
755 models=models,
756 gpus_used=self._get_num_gpus(m, gpu_types, model_names, instance_ids),
757 total_time_s=m.makespan.value,
758 ttff_s=m.ttff_user.value,
759 tbf_s=tbf_s,
760 cost=m.cost.value,
761 total_energy=m.energy.value,
762 )
764 def _is_valid_result(self, m: ConcreteModel) -> bool:
765 for gpu_type in m.GPU_TYPES:
766 for model_name in m.MODEL_NAMES:
767 for instance_id in m.INSTANCES:
768 if m.gpus[gpu_type, model_name, instance_id].value is None:
769 return False
770 return True
772 def _get_num_gpus(
773 self,
774 m: ConcreteModel,
775 gpu_types: list[GPUType],
776 model_names: list[Model],
777 instance_ids: list[int],
778 ) -> dict[GPUType, int]:
779 if not self._is_valid_result(m):
780 return {}
781 return {
782 gpu_type: sum(
783 # round() snaps solver float to nearest int (e.g. 1.9999 -> 2)
784 int(round(m.gpus[idx(gpu_type, model_name, instance_id)].value))
785 for model_name in model_names
786 for instance_id in instance_ids
787 if m.gpus[idx(gpu_type, model_name, instance_id)].value is not None
788 )
789 for gpu_type in gpu_types
790 }
792 def _get_latency_per_work(
793 self,
794 gpu_type: GPUType,
795 model_name: Model,
796 num_devices: int,
797 ) -> float:
798 """
799 Cost per unit of work for a given model and GPU type, based on latency data.
800 Cost: Linearized - sum of (latency * work_device * num_devices * ratio)
801 This replaces the bilinear makespan * GPUs.
802 """
803 # Resolution scaling factor for HF/VAE/FT
804 latency_ratio = self.workflow.get_resolution_scale(self.policy.use_upscaler)
806 if model_name == Model.GEMMA:
807 return (
808 self.latency_data[gpu_type].gemma_first_scene[num_devices]
809 + self.latency_data[gpu_type].gemma_per_scene[num_devices] * (self.workflow.total_scenes - 1)
810 )
812 if model_name == Model.FLUX:
813 return (
814 self.latency_data[gpu_type][model_name, num_devices]
815 * self.workflow.num_steps[Model.FLUX]
816 )
818 if model_name == Model.HF:
819 time_per_work = (
820 self.workflow.per_subscene_frames[Model.HF]
821 / self.workflow.hf_frames[self.workflow.frames_per_step_idx]
822 * self.latency_data[gpu_type][model_name, num_devices]
823 * latency_ratio
824 * self.workflow.num_steps[Model.HF]
825 )
826 if not self.policy.is_disaggregated(Model.HF):
827 time_per_work += self._get_latency_per_work(
828 gpu_type,
829 Model.HF_VAE,
830 1, # VAE is single-device only in current policy
831 )
832 return time_per_work
834 if model_name == Model.HF_VAE:
835 return (
836 self.latency_data[gpu_type][model_name, num_devices]
837 * latency_ratio
838 / self.workflow.hf_frames[self.workflow.frames_per_step_idx]
839 )
841 if model_name == Model.FT:
842 time_per_work = (
843 self.workflow.per_subscene_frames[Model.FT]
844 / self.workflow.ft_frames[self.workflow.frames_per_step_idx]
845 * self.latency_data[gpu_type][model_name, num_devices]
846 * latency_ratio
847 * self.workflow.num_steps[Model.FT]
848 )
849 if not self.policy.is_disaggregated(Model.FT):
850 time_per_work += self._get_latency_per_work(
851 gpu_type,
852 Model.FT_VAE,
853 1, # VAE is single-device only in current policy
854 )
855 return time_per_work
857 if model_name == Model.FT_VAE:
858 return (
859 self.latency_data[gpu_type][model_name, num_devices]
860 * latency_ratio
861 / self.workflow.ft_frames[self.workflow.frames_per_step_idx]
862 )
864 if model_name == Model.UPSCALER:
865 return self.latency_data[gpu_type][model_name, num_devices]
867 if model_name == Model.OTHERS:
868 return self.latency_data[gpu_type][model_name, num_devices] * self.workflow.total_scenes
870 raise ValueError(f"Unknown model_name {model_name}")
872 def _get_power_per_work(
873 self,
874 gpu_type: GPUType,
875 model_name: Model,
876 num_devices: int,
877 ) -> float:
878 """
879 Average power per unit of work for a given model and GPU type.
880 Returns the time-weighted average power consumption in watts.
881 For energy calculation:
882 energy = _get_latency_per_work(...) * _get_power_per_work(...) * num_devices * work
883 """
884 if self.power_data is None:
885 return 0.0
887 if model_name == Model.GEMMA:
888 # For Gemma, power varies between first scene and subsequent scenes
889 # Compute energy then divide by total time to get average power
890 power_first = self.power_data[gpu_type].gemma_first_scene[num_devices]
891 power_per_scene = self.power_data[gpu_type].gemma_per_scene[num_devices]
892 latency_first = self.latency_data[gpu_type].gemma_first_scene[num_devices]
893 latency_per_scene = self.latency_data[gpu_type].gemma_per_scene[num_devices]
895 total_energy = (
896 power_first * latency_first
897 + power_per_scene * latency_per_scene * (self.workflow.total_scenes - 1)
898 )
899 total_time = latency_first + latency_per_scene * (self.workflow.total_scenes - 1)
901 return total_energy / total_time if total_time > 0 else power_first
903 if model_name == Model.FLUX:
904 return self.power_data[gpu_type][model_name, num_devices]
906 if model_name == Model.HF:
907 return self.power_data[gpu_type][model_name, num_devices]
909 if model_name == Model.HF_VAE:
910 return self.power_data[gpu_type][model_name, num_devices]
912 if model_name == Model.FT:
913 return self.power_data[gpu_type][model_name, num_devices]
915 if model_name == Model.FT_VAE:
916 return self.power_data[gpu_type][model_name, num_devices]
918 if model_name == Model.UPSCALER:
919 return self.power_data[gpu_type][model_name, num_devices]
921 if model_name == Model.OTHERS:
922 # OTHERS model uses minimal GPU power (mostly idle)
923 # See models.py OthersModelAllocation.calculate_energy - only uses idle power
924 return self.power_data[gpu_type]["idle"]
926 raise ValueError(f"Unknown model_name {model_name}")
929def milp_to_models_dict(
930 m: ConcreteModel,
931 gpu_types: list[GPUType],
932 model_names: list[Model],
933 instance_ids: list[int],
934 idx: Callable[[GPUType, Model, int], tuple[str, str, int]],
935 workflow: WorkflowConfig,
936 power_data: Optional[PowerData],
937 policy: Policy,
938) -> dict[GPUType, dict[Model, list[ModelAllocation]]]:
939 """
940 MILP result to models dictionary.
941 """
942 if m is None:
943 return {}
945 models: dict[GPUType, dict[Model, list[ModelAllocation]]] = {}
946 for gpu_type in gpu_types:
947 models[gpu_type] = {}
948 for model_name in model_names:
949 models[gpu_type][model_name] = []
950 for instance_id in instance_ids:
951 key = idx(gpu_type, model_name, instance_id)
952 gpus_val = m.gpus[key].value
953 work_val = m.work[key].value
954 if gpus_val is None or work_val is None:
955 continue
956 # round() snaps solver floats to nearest int (e.g. 1.9999 -> 2);
957 # banker's rounding is irrelevant here since MILP values can be
958 # near-integer, like 1.999 and 2.001
959 gpus = int(round(gpus_val))
960 work = int(round(work_val))
961 if gpus > 0 and work > 0:
962 model_allocation = get_model_allocation(
963 model=model_name,
964 gpu_type=gpu_type,
965 devices=gpus,
966 replicas=1,
967 )
968 model_allocation.work = work
969 model_allocation.time = m.time[key].value
970 model_allocation.time_first = m.ttff[key].value
971 model_allocation.calculate_energy(
972 workflow=workflow,
973 power_data=power_data,
974 total_time_s=m.makespan.value
975 )
976 model_allocation.calculate_cost(
977 policy,
978 total_time_s=m.makespan.value
979 )
980 models[gpu_type][model_name].append(model_allocation)
981 merged_models = models # coalesce_models(models)
982 return merged_models
985def get_objective(
986 m: ConcreteModel,
987 allocator: Objective,
988 solver_name: str,
989) -> Optional[OptObjective]:
990 if allocator == Objective.TIME:
991 return OptObjective(expr=m.makespan, sense=minimize)
993 if allocator == Objective.TTFF:
994 return OptObjective(expr=m.ttff_user, sense=minimize)
996 if allocator == Objective.TTFF_COST:
997 # Note: This creates a bilinear (nonconvex) objective - requires Gurobi
998 if solver_name == "gurobi":
999 return OptObjective(expr=m.ttff_user * m.cost, sense=minimize)
1000 logging.warning("TTFF_COST using linear utility function.")
1001 a = 1.0
1002 b = 1.0
1003 return OptObjective(expr=a * m.ttff_user + b * m.cost, sense=minimize)
1005 if allocator == Objective.COST:
1006 return OptObjective(expr=m.cost, sense=minimize)
1008 if allocator == Objective.ENERGY:
1009 return OptObjective(expr=m.energy, sense=minimize)
1011 if allocator == Objective.TIME_ENERGY:
1012 # Note: This creates a bilinear objective - requires Gurobi
1013 if solver_name == "gurobi":
1014 return OptObjective(expr=m.makespan * m.energy, sense=minimize)
1015 logging.warning("TIME_ENERGY using linear utility function.")
1016 a = 1.0
1017 b = 1.0
1018 return OptObjective(expr=a * m.makespan + b * m.energy, sense=minimize)
1020 if allocator == Objective.ENERGY_COST:
1021 if solver_name == "gurobi":
1022 return OptObjective(expr=m.energy * m.cost, sense=minimize)
1023 logging.warning("ENERGY_COST using linear utility function.")
1024 a = 1.0
1025 b = 1.0
1026 return OptObjective(expr=a * m.energy + b * m.cost, sense=minimize)
1028 if allocator == Objective.FIFO:
1029 logging.error("FIFO not implemented in MILP")
1031 if allocator == Objective.RANDOM:
1032 return None # No objective, just find a feasible solution
1034 if allocator == Objective.NONE:
1035 return None
1037 return OptObjective(expr=m.makespan, sense=minimize)
1040def _save_solution(
1041 m: ConcreteModel,
1042 save_solution_path: str,
1043) -> None:
1044 solution = {
1045 var.name: var.value
1046 for var in m.component_data_objects(Var, active=True)
1047 if var.value is not None
1048 }
1049 with open(save_solution_path, "w", encoding="utf-8") as output_file:
1050 json.dump(solution, output_file, indent=2)
1053def _load_warm_start(
1054 m: ConcreteModel,
1055 warm_start_path: str,
1056) -> None:
1057 """Load warm start values from a JSON file and apply them to the model variables."""
1058 with open(warm_start_path, "r", encoding="utf-8") as input_file:
1059 warm_start_values = json.load(input_file)
1061 warm_start_applied = 0
1062 for var in m.component_data_objects(Var, active=True):
1063 if var.name in warm_start_values:
1064 var.set_value(warm_start_values[var.name])
1065 warm_start_applied += 1
1067 logging.info(
1068 f"Warm start loaded from {warm_start_path}. "
1069 f"Applied values to {warm_start_applied} variables."
1070 )