Coverage for simulator/evaluator.py: 92%
193 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Evaluate the performance of a given model allocation in terms of time, energy, and cost.
3It includes some assertions (e.g., only one instance of Gemma and Flux).
4"""
5from __future__ import annotations
7import math
8import logging
10from typing import Optional
12from constants import NUM_GPUS_PER_SERVER
13from constants import TOTAL_INPUT_TOKENS
14from constants import SECONDS_IN_HOUR
16from sim_types import Result
17from sim_types import GPUType
18from sim_types import WorkflowConfig
19from sim_types import PowerData
20from sim_types import LatencyData
21from sim_types import Model
22from sim_types import ModelAllocation
23from sim_types import Policy
25from sim_types_json import models_to_json
26from sim_types_json import workflow_to_json
27from sim_types_json import policy_to_json
30def _count_instances(
31 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
32 model: Model,
33) -> int:
34 num_instances = 0
35 for model_gpus in models.values():
36 if model in model_gpus:
37 for model_allocation in model_gpus[model]:
38 if model_allocation.get_num_gpus() > 0:
39 num_instances += 1
40 return num_instances
43def _assert_single_instance(
44 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
45 model: Model,
46) -> None:
47 num_instances = _count_instances(models, model)
48 assert num_instances == 1, f"Expected exactly one instance of {model}, but found {num_instances}"
51def _assert_at_least_one_instance(
52 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
53 model: Model,
54) -> None:
55 num_instances = _count_instances(models, model)
56 assert num_instances > 0, f"Expected at least one instance of {model}, but found {num_instances}"
59def _assert_no_instances(
60 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
61 model: Model,
62) -> None:
63 num_instances = _count_instances(models, model)
64 assert num_instances == 0, f"Expected no instances of {model}, but found {num_instances}"
67def evaluate_times(
68 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
69 latency_data: LatencyData,
70 workflow: WorkflowConfig,
71 policy: Policy,
72 include_models: Optional[list[Model]] = None,
73) -> None:
74 """
75 Compute the total time for the given model allocation and workflow, using the latency data.
76 It only evaluates the models specified in "include_models" if provided.
77 """
78 gpu_types = list(models.keys())
80 upscaler_gpus = sum(
81 model_alloc.get_num_gpus()
82 for gpu_type in gpu_types
83 for model_alloc in models.get(gpu_type, {}).get(Model.UPSCALER, [])
84 )
85 if not policy.use_upscaler:
86 assert upscaler_gpus == 0
88 for model_name in workflow.models:
89 if include_models is not None and model_name not in include_models:
90 continue
92 # Special conditions: models that require a policy flag
93 if model_name == Model.HF_VAE and not policy.is_disaggregated(Model.HF):
94 _assert_no_instances(models, Model.HF_VAE)
95 continue
96 if model_name == Model.FT_VAE and not policy.is_disaggregated(Model.FT):
97 _assert_no_instances(models, Model.FT_VAE)
98 continue
99 if model_name == Model.UPSCALER and not policy.use_upscaler:
100 _assert_no_instances(models, Model.UPSCALER)
101 continue
103 _assert_at_least_one_instance(models, model_name)
105 if not workflow.is_parallelizable(model_name):
106 # Single-instance: no work splitting
107 for gpu_type in gpu_types:
108 if model_name in models[gpu_type]:
109 for model_alloc in models[gpu_type][model_name]:
110 model_alloc.calculate_time(
111 policy, workflow, latency_data)
112 model_alloc.calculate_time_first(
113 policy, workflow, latency_data)
114 continue
116 # Parallel: capacity-based work splitting (throughput-weighted)
117 capacities: dict[GPUType, list[float]] = {}
118 for gpu_type in gpu_types:
119 capacities[gpu_type] = []
120 if model_name not in models[gpu_type]:
121 continue
122 for model_alloc in models[gpu_type][model_name]:
123 if model_alloc.get_num_gpus() > 0:
124 latency = latency_data[gpu_type][model_name, model_alloc.devices]
125 # When not disaggregated, include VAE overhead in capacity
126 if model_name == Model.FT and not policy.is_disaggregated(Model.FT):
127 latency += latency_data[gpu_type][Model.FT_VAE, 1] / workflow.num_steps[Model.FT]
128 if model_name == Model.HF and not policy.is_disaggregated(Model.HF):
129 latency += latency_data[gpu_type][Model.HF_VAE, 1] / workflow.num_steps[Model.HF]
130 if model_name in (Model.HF, Model.HF_VAE, Model.FT, Model.FT_VAE):
131 latency *= workflow.get_resolution_scale(policy.use_upscaler)
132 if model_name == Model.GEMMA:
133 latency *= workflow.total_input_tokens / TOTAL_INPUT_TOKENS
134 if latency == 0:
135 capacities[gpu_type].append(0.0)
136 else:
137 capacities[gpu_type].append(model_alloc.replicas / latency)
139 total_capacity = sum(sum(c) for c in capacities.values())
140 for gpu_type in gpu_types:
141 if model_name not in models[gpu_type]:
142 continue
143 cap_idx = 0
144 for model_alloc in models[gpu_type][model_name]:
145 if model_alloc.get_num_gpus() > 0:
146 work_pct = capacities[gpu_type][cap_idx] / total_capacity if total_capacity > 0 else 0.0
147 model_alloc.calculate_time(
148 policy, workflow, latency_data,
149 work_pct=work_pct)
150 model_alloc.calculate_time_first(
151 policy, workflow, latency_data)
152 cap_idx += 1
155def evaluate_energy(
156 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
157 power_data: PowerData,
158 workflow: WorkflowConfig,
159 total_time_s: float = 0.0,
160) -> None:
161 """
162 Calculate total energy (power * time * replicas for each model).
163 Need to run after evaluate_times since energy calculation depends on time.
164 """
165 for gpu_type_allocs in models.values():
166 for model_allocation_list in gpu_type_allocs.values():
167 for model_allocation in model_allocation_list:
168 model_allocation.calculate_energy(
169 workflow,
170 power_data,
171 total_time_s)
174def evaluate_cost(
175 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
176 total_time_s: float,
177 policy: Policy,
178) -> None:
179 """
180 Calculate total cost based on GPU hours used.
181 Need to run after evaluate_times since cost calculation depends on time.
182 """
183 for gpu_type_allocs in models.values():
184 for model_allocation_list in gpu_type_allocs.values():
185 for model in model_allocation_list:
186 model.calculate_cost(policy, total_time_s)
189_EVALUATOR_CACHE: dict[str, Result] = {}
192def evaluate_model_allocation(
193 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
194 num_gpus: dict[GPUType, int],
195 workflow: WorkflowConfig,
196 latency_data: LatencyData,
197 power_data: Optional[PowerData],
198 policy: Policy,
199 include_models: Optional[list[Model]] = None,
200 cache_results: bool = False,
201 round_up_cost_to_server: bool = False,
202) -> Result:
203 """
204 Evaluate the metrics for a given allocation of models to GPUs.
205 It only evaluates the models in "include_models" if specified.
206 """
207 cache_key = None
208 if cache_results:
209 cache_key = models_to_json(models) + \
210 workflow_to_json(workflow) + \
211 str(latency_data) + \
212 str(power_data) + \
213 policy_to_json(policy) + \
214 str(include_models)
215 if cache_key in _EVALUATOR_CACHE:
216 return _EVALUATOR_CACHE[cache_key]
218 # Check if setup is possible
219 gpus_used = {}
220 for gpu_type, model_gpu in models.items():
221 gpus_used[gpu_type] = calc_used_gpus({gpu_type: model_gpu})
222 assert num_gpus[gpu_type] % NUM_GPUS_PER_SERVER[gpu_type] == 0, \
223 f"{gpu_type.value}: {num_gpus[gpu_type]} % {NUM_GPUS_PER_SERVER[gpu_type]}"
224 assert gpus_used[gpu_type] <= num_gpus[gpu_type], \
225 f"{gpu_type.value}: {gpus_used[gpu_type]} > {num_gpus[gpu_type]}"
227 # Assert input models are built correctly
228 for gpu_type in models.keys():
229 for model_name in models[gpu_type].keys():
230 for instance_id in range(len(models[gpu_type][model_name])):
231 assert models[gpu_type][model_name][instance_id].model == model_name
232 assert models[gpu_type][model_name][instance_id].gpu_type == gpu_type
234 # Actual evaluation
235 evaluate_times(
236 models, latency_data, workflow, policy,
237 include_models=include_models,
238 )
239 time_s = calc_total_time(models)
241 first_chunk_time = calc_ttff(models)
242 ttff_s = max(
243 first_chunk_time,
244 time_s - workflow.total_video_seconds
245 )
247 num_frames = (workflow.total_frames[Model.FT] - workflow.per_subscene_frames[Model.FT])
248 tbf_s = (time_s - first_chunk_time) / num_frames
249 if tbf_s < 0:
250 logging.debug(
251 f"Negative TBF: "
252 F"{tbf_s:.2f} = ({time_s:.2f} - {first_chunk_time:.2f}) / {num_frames}")
253 tbf_s = 0.0
255 # Calculate total energy (power * time * replicas for each model)
256 energy = 0.0
257 if power_data is not None:
258 evaluate_energy(models, power_data, workflow, time_s)
259 energy = calc_energy(models=models)
261 evaluate_cost(models, time_s, policy)
262 cost = calc_cost(
263 models, time_s, policy,
264 round_up_to_server=round_up_cost_to_server)
266 ret = Result(
267 models=models,
268 gpus_used=gpus_used,
269 gpus_total=num_gpus,
270 total_time_s=time_s,
271 first_chunk_time=first_chunk_time,
272 ttff_s=ttff_s,
273 tbf_s=tbf_s,
274 total_energy=energy if power_data else 0.0,
275 cost=cost,
276 )
278 if cache_key is not None:
279 _EVALUATOR_CACHE[cache_key] = ret
281 return ret
284def calc_energy(
285 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
286) -> float:
287 """
288 Calculate total energy (power * time * replicas for each model).
289 Energy in Watt x seconds (Joules).
290 This assumes that evaluate_energy() has been called already.
291 """
292 energy = 0.0 # Total energy in Watt-seconds (Joules = Watt x second)
293 for model_dict in models.values():
294 for model_allocations in model_dict.values():
295 for model_allocation in model_allocations:
296 energy += model_allocation.energy
297 return energy
300def calc_model_cost(
301 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
302) -> float:
303 """
304 Calculate total cost based on GPU hours used.
305 This assumes that evaluate_cost() has been called already.
306 """
307 costs = {}
308 for gpu_type, model_dict in models.items():
309 costs[gpu_type] = 0.0
310 for model_allocations in model_dict.values():
311 for model_allocation in model_allocations:
312 costs[gpu_type] += model_allocation.cost
313 return sum(costs.values())
316def calc_cost(
317 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
318 time_s: float,
319 policy: Policy,
320 round_up_to_server: bool = True,
321) -> float:
322 """
323 Calculate total cost based on GPU hours used.
324 """
325 used_gpus = calc_used_gpus_per_type(models)
327 # Round up to the nearest server (pack of GPUs) since we pay for whole servers
328 if round_up_to_server:
329 for gpu_type, used in used_gpus.items():
330 used_pack = math.ceil(used / NUM_GPUS_PER_SERVER[gpu_type]) * NUM_GPUS_PER_SERVER[gpu_type]
331 used_gpus[gpu_type] = used_pack
333 return calc_cost_total(used_gpus, time_s, policy)
336def calc_cost_total(
337 num_gpus: dict[GPUType, int],
338 time_s: float,
339 policy: Policy,
340) -> float:
341 """
342 Calculate total cost based on GPU hours used.
343 It includes the idle GPUs not assigned to a model.
344 """
345 cost = 0.0
346 for gpu_type, num in num_gpus.items():
347 cost += num * (time_s / SECONDS_IN_HOUR) * policy.gpu_cost[gpu_type]
348 return cost
351def calc_used_gpus_per_type(
352 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
353) -> dict[GPUType, int]:
354 """
355 Calculate number of GPUs used per GPU type across all models.
356 """
357 gpus_used = {}
358 for gpu_type, model_gpu in models.items():
359 gpus_used[gpu_type] = 0
360 for model_allocations in model_gpu.values():
361 for model_allocation in model_allocations:
362 gpus_used[gpu_type] += model_allocation.get_num_gpus()
363 return gpus_used
366def calc_used_gpus(
367 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
368) -> int:
369 """
370 Calculate total number of GPUs used across all models and GPU types.
371 """
372 gpus_used = calc_used_gpus_per_type(models)
373 return sum(gpus_used.values())
376def calc_total_time(
377 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
378) -> float:
379 """
380 Calculate total time considering all stages and dependencies.
381 This assumes that evaluate_time() has been called already.
382 """
383 total_time_secs = 0.0
384 for model_name in Model:
385 model_alloc_times = [
386 model_alloc.time
387 for gpu_type in GPUType
388 if gpu_type in models and model_name in models[gpu_type]
389 for model_alloc in models[gpu_type][model_name]
390 ]
391 model_time = max(model_alloc_times) if model_alloc_times else 0.0
392 total_time_secs += model_time
393 return total_time_secs
396def calc_ttff(
397 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
398) -> float:
399 """
400 Calculate time to first frame (chunk).
401 It takes the time to first frame (TTFF) for each model.
402 This assumes that evaluate_time() has been called already.
403 """
404 models_time_first: dict[Model, float] = {}
405 for model_name in Model:
406 times_first = []
407 for gpu_type in models.keys():
408 if model_name in models[gpu_type]:
409 for model_alloc in models[gpu_type][model_name]:
410 if model_alloc.get_num_gpus() > 0:
411 times_first.append(model_alloc.time_first)
412 if len(times_first) > 0:
413 models_time_first[model_name] = min(times_first) # The fastest model determines TTFF
414 return sum(models_time_first.values())