Coverage for simulator/models.py: 96%
356 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Contains the definition for each model.
3It includes the calculations for time, energy, and cost.
4"""
5from __future__ import annotations
7import math
9from typing import override
10from typing import Callable
11from typing import Optional
12from typing import Type
13from typing import ClassVar
15from sim_types import LatencyData
16from sim_types import PowerData
17from sim_types import ModelAllocation
18from sim_types import Model
19from sim_types import Policy
20from sim_types import QualityLevel
21from sim_types import WorkflowConfig
22from sim_types import GPUType
24from constants import TOTAL_INPUT_TOKENS
27# ModelAllocation Factory
28ModelAllocationCls = Type[ModelAllocation]
30_MODEL_ALLOCATION_REGISTRY: dict[Model, ModelAllocationCls] = {}
33def register_model(
34 model: Model
35) -> Callable[[ModelAllocationCls], ModelAllocationCls]:
36 """Register a ModelAllocation class for the factory."""
37 def decorator(cls: ModelAllocationCls) -> ModelAllocationCls:
38 _MODEL_ALLOCATION_REGISTRY[model] = cls
39 return cls
40 return decorator
43def get_model_allocation(
44 *,
45 model: Model,
46 gpu_type: GPUType,
47 devices: int = 1,
48 replicas: int = 0,
49) -> ModelAllocation:
50 """Factory to get the ModelAllocation instance for a specific model."""
51 if model not in _MODEL_ALLOCATION_REGISTRY:
52 raise ValueError(f"No ModelAllocation for model {model}")
53 cls = _MODEL_ALLOCATION_REGISTRY[model]
54 return cls(
55 gpu_type=gpu_type,
56 devices=devices,
57 replicas=replicas,
58 )
61def _calculate_total_time(
62 total_work: float,
63 num_replicas: int,
64 time_per_work: float,
65) -> float:
66 """Calculate total time given work, replicas, and time per work unit."""
67 if num_replicas <= 0:
68 return 0.0
69 total_time = (total_work / num_replicas) * time_per_work
70 if total_time < time_per_work: # We cannot go faster than single work unit time
71 total_time = time_per_work
72 return total_time
75def assert_pixel_config(
76 workflow: WorkflowConfig
77) -> None:
78 """Verify that the workflow's pixel configuration is valid for upscaling."""
79 from sim_types import RESOLUTION_PIXELS
80 assert 0 < RESOLUTION_PIXELS[QualityLevel.MEDIUM] < RESOLUTION_PIXELS[QualityLevel.HIGH]
83@register_model(Model.GEMMA)
84class GemmaModelAllocation(ModelAllocation):
85 """Gemma model allocation."""
86 model: ClassVar[Model] = Model.GEMMA
88 @override
89 def get_max_replicas(
90 self,
91 workflow: WorkflowConfig,
92 ) -> int:
93 return workflow.model_work.get(Model.GEMMA, 1)
95 @override
96 def calculate_time(
97 self,
98 policy: Policy,
99 workflow: WorkflowConfig,
100 latency_data: LatencyData,
101 work_pct: float = 1.0,
102 ) -> float:
103 if self.get_num_gpus() == 0:
104 self.time = 0.0
105 return self.time
106 latency_first = latency_data[self.gpu_type].gemma_first_scene[self.devices]
107 latency_per_scene = latency_data[self.gpu_type].gemma_per_scene[self.devices]
108 latency_first *= workflow.total_input_tokens / TOTAL_INPUT_TOKENS
109 latency_per_scene *= workflow.total_input_tokens / TOTAL_INPUT_TOKENS
110 total_work = workflow.model_work.get(Model.GEMMA, 1)
111 if total_work > 1:
112 num_scenes = math.ceil(work_pct * total_work)
113 total_time_per_scene = latency_first + latency_per_scene * (num_scenes - 1)
114 self.time = _calculate_total_time(
115 num_scenes,
116 self.replicas,
117 total_time_per_scene / num_scenes)
118 else:
119 self.time = latency_first + latency_per_scene * (workflow.total_scenes - 1)
120 return self.time
122 @override
123 def calculate_time_first(
124 self,
125 policy: Policy,
126 workflow: WorkflowConfig,
127 latency_data: LatencyData,
128 ) -> float:
129 if self.get_num_gpus() == 0:
130 self.time_first = 0.0
131 return self.time_first
132 latency_first = latency_data[self.gpu_type].gemma_first_scene[self.devices]
133 latency_first *= workflow.total_input_tokens / TOTAL_INPUT_TOKENS
134 self.time_first = latency_first
135 return self.time_first
137 @override
138 def calculate_energy(
139 self,
140 workflow: WorkflowConfig,
141 power_data: Optional[PowerData] = None,
142 total_time_s: float = 0.0,
143 ) -> float:
144 if self.get_num_gpus() == 0 or power_data is None:
145 self.energy = 0.0
146 return self.energy
147 # Gemma energy
148 latency_first = self.time_first
149 latency_per_scene = max(0.0, self.time - latency_first)
150 power_first = power_data[self.gpu_type].gemma_first_scene[self.devices]
151 power_per_scene = power_data[self.gpu_type].gemma_per_scene[self.devices]
152 self.energy = \
153 power_first * latency_first + \
154 power_per_scene * latency_per_scene * (workflow.total_scenes - 1)
155 # Idle energy
156 power_idle = power_data[self.gpu_type]["idle"] * self.get_num_gpus()
157 time_idle = total_time_s - self.time
158 if time_idle > 0:
159 self.energy += power_idle * time_idle
160 return self.energy
163@register_model(Model.FLUX)
164class FluxModelAllocation(ModelAllocation):
165 """Flux model allocation."""
166 model: ClassVar[Model] = Model.FLUX
168 def _calc_time_per_scene(
169 self,
170 policy: Policy,
171 workflow: WorkflowConfig,
172 latency_data: LatencyData,
173 ) -> float:
174 return (
175 latency_data[self.gpu_type][self.model, self.devices]
176 * workflow.num_steps[Model.FLUX]
177 )
179 @override
180 def get_max_replicas(
181 self,
182 workflow: WorkflowConfig,
183 ) -> int:
184 return workflow.model_work.get(Model.FLUX, 1)
186 @override
187 def calculate_time(
188 self,
189 policy: Policy,
190 workflow: WorkflowConfig,
191 latency_data: LatencyData,
192 work_pct: float = 1.0,
193 ) -> float:
194 if self.get_num_gpus() == 0:
195 self.time = 0.0
196 return self.time
197 time_per_scene = self._calc_time_per_scene(
198 policy,
199 workflow,
200 latency_data,
201 )
202 total_work = workflow.model_work.get(Model.FLUX, 1)
203 if total_work > 1:
204 num_scenes = math.ceil(work_pct * total_work)
205 self.time = _calculate_total_time(
206 num_scenes,
207 self.replicas,
208 time_per_scene)
209 else:
210 self.time = time_per_scene
211 return self.time
213 @override
214 def calculate_time_first(
215 self,
216 policy: Policy,
217 workflow: WorkflowConfig,
218 latency_data: LatencyData,
219 ) -> float:
220 if self.get_num_gpus() == 0:
221 self.time_first = 0.0
222 return self.time_first
223 time_per_scene = self._calc_time_per_scene(
224 policy,
225 workflow,
226 latency_data,
227 )
228 self.time_first = time_per_scene
229 return self.time_first
231 @override
232 def calculate_energy(
233 self,
234 workflow: WorkflowConfig,
235 power_data: Optional[PowerData] = None,
236 total_time_s: float = 0.0,
237 ) -> float:
238 if self.get_num_gpus() == 0 or power_data is None:
239 self.energy = 0.0
240 return self.energy
241 power_flux = power_data[self.gpu_type][Model.FLUX, self.devices]
242 self.energy = power_flux * self.time * self.replicas
243 # Idle energy
244 power_idle = power_data[self.gpu_type]["idle"] * self.get_num_gpus()
245 time_idle = total_time_s - self.time
246 if time_idle > 0:
247 self.energy += power_idle * time_idle
248 return self.energy
251@register_model(Model.HF)
252class HFModelAllocation(ModelAllocation):
253 """HunyuanFramePack model allocation."""
254 model: ClassVar[Model] = Model.HF
256 def _calc_time_per_frame(
257 self,
258 policy: Policy,
259 workflow: WorkflowConfig,
260 latency_data: LatencyData,
261 ) -> float:
262 return (
263 latency_data[self.gpu_type][self.model, self.devices]
264 * workflow.get_resolution_scale(policy.use_upscaler)
265 * workflow.num_steps[Model.HF]
266 )
268 def _calc_time_per_subscene(
269 self,
270 policy: Policy,
271 workflow: WorkflowConfig,
272 latency_data: LatencyData,
273 ) -> float:
274 return (
275 workflow.per_subscene_frames[Model.HF]
276 / workflow.hf_frames[workflow.frames_per_step_idx]
277 * latency_data[self.gpu_type][self.model, self.devices]
278 * workflow.get_resolution_scale(policy.use_upscaler) # latency_ratio
279 * workflow.num_steps[Model.HF]
280 )
282 @override
283 def calculate_time(
284 self,
285 policy: Policy,
286 workflow: WorkflowConfig,
287 latency_data: LatencyData,
288 work_pct: float = 1.0,
289 ) -> float:
290 if self.get_num_gpus() == 0:
291 self.time = 0.0
292 return self.time
294 hf_time_per_subscene = self._calc_time_per_subscene(
295 policy,
296 workflow,
297 latency_data,
298 )
299 self.time = _calculate_total_time(
300 math.ceil(work_pct * workflow.total_subscenes),
301 self.replicas,
302 hf_time_per_subscene)
304 if not policy.is_disaggregated(Model.HF):
305 # Include VAE time in the same GPU when disaggregation is disabled
306 hf_vae_time_per_frame = (
307 latency_data[self.gpu_type][Model.HF_VAE, 1] # VAE is single-device only in current policy
308 * workflow.get_resolution_scale(policy.use_upscaler)
309 / workflow.hf_frames[workflow.frames_per_step_idx]
310 )
311 self.time += _calculate_total_time(
312 math.ceil(work_pct * workflow.total_frames[Model.HF]),
313 self.replicas,
314 hf_vae_time_per_frame)
316 return self.time
318 @override
319 def calculate_time_first(
320 self,
321 policy: Policy,
322 workflow: WorkflowConfig,
323 latency_data: LatencyData,
324 ) -> float:
325 if self.get_num_gpus() == 0:
326 self.time_first = 0.0
327 return self.time_first
329 if policy.is_disaggregated(Model.HF):
330 # HF for the first chunk
331 self.time_first = min(
332 # Option 1: the first few frames until the first chunk is done
333 workflow.hf_frames[0]
334 / workflow.hf_frames[workflow.frames_per_step_idx]
335 * self._calc_time_per_frame(
336 policy,
337 workflow,
338 latency_data
339 ),
340 # Option 2: the full subscene
341 self._calc_time_per_subscene(
342 policy,
343 workflow,
344 latency_data
345 ),
346 )
347 else:
348 # HF + VAE for the full subscene
349 hf_time_per_subscene = self._calc_time_per_subscene(
350 policy,
351 workflow,
352 latency_data)
353 hf_vae_time_per_subscene = (
354 workflow.per_subscene_frames[Model.HF]
355 / workflow.hf_frames[workflow.frames_per_step_idx]
356 * latency_data[self.gpu_type][Model.HF_VAE, 1] # VAE is single-device only in current policy
357 * workflow.get_resolution_scale(policy.use_upscaler)
358 )
359 self.time_first = hf_time_per_subscene + hf_vae_time_per_subscene
361 return self.time_first
363 @override
364 def calculate_energy(
365 self,
366 workflow: WorkflowConfig,
367 power_data: Optional[PowerData] = None,
368 total_time_s: float = 0.0,
369 ) -> float:
370 if self.get_num_gpus() == 0 or power_data is None:
371 self.energy = 0.0
372 return self.energy
373 power_hf = power_data[self.gpu_type][Model.HF, self.devices]
374 self.energy = power_hf * self.time * self.replicas
375 # Idle energy
376 power_idle = power_data[self.gpu_type]["idle"] * self.get_num_gpus()
377 time_idle = total_time_s - self.time
378 if time_idle > 0:
379 self.energy += power_idle * time_idle
380 return self.energy
382 @override
383 def get_max_replicas(
384 self,
385 workflow: WorkflowConfig,
386 ) -> int:
387 return workflow.model_work.get(Model.HF, 1)
390@register_model(Model.HF_VAE)
391class HFVAEModelAllocation(ModelAllocation):
392 """HunyuanFramePack VAE model allocation."""
393 model: ClassVar[Model] = Model.HF_VAE
395 def _calc_time_per_frame(
396 self,
397 policy: Policy,
398 workflow: WorkflowConfig,
399 latency_data: LatencyData,
400 ) -> float:
401 return (
402 latency_data[self.gpu_type][Model.HF_VAE, self.devices]
403 * workflow.get_resolution_scale(policy.use_upscaler)
404 / workflow.hf_frames[workflow.frames_per_step_idx]
405 )
407 @override
408 def calculate_time(
409 self,
410 policy: Policy,
411 workflow: WorkflowConfig,
412 latency_data: LatencyData,
413 work_pct: float = 1.0,
414 ) -> float:
415 if not policy.is_disaggregated(Model.HF):
416 assert self.get_num_gpus() == 0
417 self.time = 0.0
418 return self.time
419 if self.get_num_gpus() == 0:
420 self.time = 0.0
421 return self.time
423 vae_time_per_frame = self._calc_time_per_frame(
424 policy,
425 workflow,
426 latency_data
427 )
428 self.time = _calculate_total_time(
429 math.ceil(workflow.total_frames[Model.HF] * work_pct),
430 self.replicas,
431 vae_time_per_frame)
432 return self.time
434 @override
435 def calculate_time_first(
436 self,
437 policy: Policy,
438 workflow: WorkflowConfig,
439 latency_data: LatencyData,
440 ) -> float:
441 if not policy.is_disaggregated(Model.HF):
442 assert self.get_num_gpus() == 0
443 self.time_first = 0.0
444 return self.time_first
445 if self.get_num_gpus() == 0:
446 self.time_first = 0.0
447 return self.time_first
449 vae_time_per_frame = self._calc_time_per_frame(
450 policy,
451 workflow,
452 latency_data,
453 )
454 num_frames = workflow.per_subscene_frames[Model.HF]
455 self.time_first = num_frames * vae_time_per_frame
456 return self.time_first
458 @override
459 def calculate_energy(
460 self,
461 workflow: WorkflowConfig,
462 power_data: Optional[PowerData] = None,
463 total_time_s: float = 0.0,
464 ) -> float:
465 if self.get_num_gpus() == 0 or power_data is None:
466 self.energy = 0.0
467 return self.energy
468 self.energy = power_data[self.gpu_type][Model.HF_VAE, self.devices] * self.time * self.replicas
469 # Idle energy
470 power_idle = power_data[self.gpu_type]["idle"] * self.get_num_gpus()
471 time_idle = total_time_s - self.time
472 if time_idle > 0:
473 self.energy += power_idle * time_idle
474 return self.energy
476 @override
477 def get_max_replicas(
478 self,
479 workflow: WorkflowConfig,
480 ) -> int:
481 return workflow.model_work.get(Model.HF_VAE, 1)
484@register_model(Model.FT)
485class FTModelAllocation(ModelAllocation):
486 """FantasyTalking model allocation."""
487 model: ClassVar[Model] = Model.FT
489 def _calc_time_per_subscene(
490 self,
491 policy: Policy,
492 workflow: WorkflowConfig,
493 latency_data: LatencyData,
494 ) -> float:
495 return (
496 workflow.per_subscene_frames[Model.FT]
497 / workflow.ft_frames[workflow.frames_per_step_idx]
498 * latency_data[self.gpu_type][Model.FT, self.devices]
499 * workflow.get_resolution_scale(policy.use_upscaler)
500 * workflow.num_steps[Model.FT]
501 )
503 @override
504 def calculate_time(
505 self,
506 policy: Policy,
507 workflow: WorkflowConfig,
508 latency_data: LatencyData,
509 work_pct: float = 1.0,
510 ) -> float:
511 if self.get_num_gpus() == 0:
512 self.time = 0.0
513 return self.time
515 ft_time_per_subscene = self._calc_time_per_subscene(
516 policy,
517 workflow,
518 latency_data,
519 )
520 self.time = _calculate_total_time(
521 math.ceil(work_pct * workflow.total_subscenes),
522 self.replicas,
523 ft_time_per_subscene)
525 if not policy.is_disaggregated(Model.FT):
526 # Include VAE time in the same GPU when disaggregation is disabled
527 # Note: VAE latency uses devices=1 as VAE processing is not parallelized
528 # across multiple devices in the same way as the main FT diffusion
529 ft_vae_time_per_frame = (
530 latency_data[self.gpu_type][Model.FT_VAE, 1]
531 * workflow.get_resolution_scale(policy.use_upscaler)
532 / workflow.ft_frames[workflow.frames_per_step_idx]
533 )
534 self.time += _calculate_total_time(
535 math.ceil(work_pct * workflow.total_frames[Model.FT]),
536 self.replicas,
537 ft_vae_time_per_frame)
539 return self.time
541 @override
542 def calculate_time_first(
543 self,
544 policy: Policy,
545 workflow: WorkflowConfig,
546 latency_data: LatencyData,
547 ) -> float:
548 if self.get_num_gpus() == 0:
549 self.time_first = 0.0
550 return self.time_first
552 ft_time_per_subscene = self._calc_time_per_subscene(
553 policy,
554 workflow,
555 latency_data,
556 )
557 self.time_first = ft_time_per_subscene
559 if not policy.is_disaggregated(Model.FT):
560 # Include VAE time_first when FT-VAE is not disaggregated
561 # Note: VAE latency uses devices=1 (see note in calculate_time)
562 ft_vae_time_per_subscene = (
563 workflow.per_subscene_frames[Model.FT]
564 / workflow.ft_frames[workflow.frames_per_step_idx]
565 * latency_data[self.gpu_type][Model.FT_VAE, 1]
566 * workflow.get_resolution_scale(policy.use_upscaler)
567 )
568 self.time_first += ft_vae_time_per_subscene
570 return self.time_first
572 @override
573 def calculate_energy(
574 self,
575 workflow: WorkflowConfig,
576 power_data: Optional[PowerData] = None,
577 total_time_s: float = 0.0,
578 ) -> float:
579 if self.get_num_gpus() == 0 or power_data is None:
580 self.energy = 0.0
581 return self.energy
582 power_ft = power_data[self.gpu_type][Model.FT, self.devices]
583 self.energy = power_ft * self.time * self.replicas
584 # Idle energy
585 power_idle = power_data[self.gpu_type]["idle"] * self.get_num_gpus()
586 time_idle = total_time_s - self.time
587 if time_idle > 0:
588 self.energy += power_idle * time_idle
589 return self.energy
591 @override
592 def get_max_replicas(
593 self,
594 workflow: WorkflowConfig,
595 ) -> int:
596 return workflow.model_work.get(Model.FT, 1)
599@register_model(Model.FT_VAE)
600class FTVAEModelAllocation(ModelAllocation):
601 """FantasyTalking VAE model allocation."""
602 model: ClassVar[Model] = Model.FT_VAE
604 def _calc_time_per_frame(
605 self,
606 policy: Policy,
607 workflow: WorkflowConfig,
608 latency_data: LatencyData,
609 ) -> float:
610 return (
611 latency_data[self.gpu_type][Model.FT_VAE, self.devices]
612 * workflow.get_resolution_scale(policy.use_upscaler)
613 / workflow.ft_frames[workflow.frames_per_step_idx]
614 )
616 @override
617 def calculate_time(
618 self,
619 policy: Policy,
620 workflow: WorkflowConfig,
621 latency_data: LatencyData,
622 work_pct: float = 1.0,
623 ) -> float:
624 if not policy.is_disaggregated(Model.FT):
625 assert self.get_num_gpus() == 0
626 self.time = 0.0
627 return self.time
628 if self.get_num_gpus() == 0:
629 self.time = 0.0
630 return self.time
632 vae_time_per_frame = self._calc_time_per_frame(
633 policy,
634 workflow,
635 latency_data,
636 )
637 self.time = _calculate_total_time(
638 math.ceil(workflow.total_frames[Model.FT] * work_pct),
639 self.replicas,
640 vae_time_per_frame)
641 return self.time
643 @override
644 def calculate_time_first(
645 self,
646 policy: Policy,
647 workflow: WorkflowConfig,
648 latency_data: LatencyData,
649 ) -> float:
650 if not policy.is_disaggregated(Model.FT):
651 assert self.get_num_gpus() == 0
652 self.time_first = 0.0
653 return self.time_first
654 if self.get_num_gpus() == 0:
655 self.time_first = 0.0
656 return self.time_first
658 vae_time_per_frame = self._calc_time_per_frame(
659 policy,
660 workflow,
661 latency_data,
662 )
663 num_frames = workflow.per_subscene_frames[Model.FT]
664 self.time_first = num_frames * vae_time_per_frame
665 return self.time_first
667 @override
668 def calculate_energy(
669 self,
670 workflow: WorkflowConfig,
671 power_data: Optional[PowerData] = None,
672 total_time_s: float = 0.0,
673 ) -> float:
674 if self.get_num_gpus() == 0 or power_data is None:
675 self.energy = 0.0
676 return self.energy
677 self.energy = power_data[self.gpu_type][Model.FT_VAE, self.devices] * self.time * self.replicas
678 # Idle energy
679 power_idle = power_data[self.gpu_type]["idle"] * self.get_num_gpus()
680 time_idle = total_time_s - self.time
681 if time_idle > 0:
682 self.energy += power_idle * time_idle
683 return self.energy
685 @override
686 def get_max_replicas(
687 self,
688 workflow: WorkflowConfig,
689 ) -> int:
690 return workflow.model_work.get(Model.FT_VAE, 1)
693@register_model(Model.UPSCALER)
694class UpscalerModelAllocation(ModelAllocation):
695 """Upscaler model allocation."""
696 model: ClassVar[Model] = Model.UPSCALER
698 @override
699 def calculate_time(
700 self,
701 policy: Policy,
702 workflow: WorkflowConfig,
703 latency_data: LatencyData,
704 work_pct: float = 1.0,
705 ) -> float:
706 if self.get_num_gpus() == 0:
707 self.time = 0.0
708 return self.time
709 self.time = _calculate_total_time(
710 math.ceil(work_pct * workflow.total_frames[Model.FT]),
711 self.replicas,
712 latency_data[self.gpu_type][self.model, self.devices])
713 return self.time
715 @override
716 def calculate_time_first(
717 self,
718 policy: Policy,
719 workflow: WorkflowConfig,
720 latency_data: LatencyData,
721 ) -> float:
722 if not policy.use_upscaler:
723 assert self.get_num_gpus() == 0
724 if self.get_num_gpus() == 0:
725 self.time_first = 0.0
726 return self.time_first
728 self.time_first = (
729 workflow.per_subscene_frames[Model.FT]
730 * latency_data[self.gpu_type][self.model, self.devices]
731 )
732 return self.time_first
734 @override
735 def calculate_energy(
736 self,
737 workflow: WorkflowConfig,
738 power_data: Optional[PowerData] = None,
739 total_time_s: float = 0.0,
740 ) -> float:
741 if self.get_num_gpus() == 0 or power_data is None:
742 self.energy = 0.0
743 return self.energy
744 # Assumes a single device and multiple replicas
745 self.energy = power_data[self.gpu_type][self.model, self.devices] * self.time * self.replicas
746 # Idle energy
747 power_idle = power_data[self.gpu_type]["idle"] * self.get_num_gpus()
748 time_idle = total_time_s - self.time
749 if time_idle > 0:
750 self.energy += power_idle * time_idle
751 return self.energy
753 @override
754 def get_max_replicas(
755 self,
756 workflow: WorkflowConfig,
757 ) -> int:
758 return workflow.model_work.get(Model.UPSCALER, 1)
761@register_model(Model.OTHERS)
762class OthersModelAllocation(ModelAllocation):
763 """Others: Kokoro + YOLO."""
764 model: ClassVar[Model] = Model.OTHERS
766 @override
767 def calculate_time(
768 self,
769 policy: Policy,
770 workflow: WorkflowConfig,
771 latency_data: LatencyData,
772 work_pct: float = 1.0,
773 ) -> float:
774 if self.get_num_gpus() == 0:
775 self.time = 0.0
776 return self.time
778 self.time = (
779 workflow.total_scenes
780 * latency_data[self.gpu_type][self.model, self.devices]
781 )
782 return self.time
784 @override
785 def calculate_time_first(
786 self,
787 policy: Policy,
788 workflow: WorkflowConfig,
789 latency_data: LatencyData,
790 ) -> float:
791 if self.get_num_gpus() == 0:
792 self.time_first = 0.0
793 return self.time_first
795 self.time_first = latency_data[self.gpu_type][self.model, self.devices]
796 return self.time_first
798 @override
799 def calculate_energy(
800 self,
801 workflow: WorkflowConfig,
802 power_data: Optional[PowerData] = None,
803 total_time_s: float = 0.0,
804 ) -> float:
805 if self.get_num_gpus() == 0 or power_data is None:
806 self.energy = 0.0
807 return self.energy
808 # Idle energy; not much GPU usage
809 power_idle = power_data[self.gpu_type]["idle"] * self.get_num_gpus()
810 self.energy = power_idle * self.time
811 return self.energy