Coverage for simulator/sim_types.py: 90%

1from __future__ import annotations 

2 

3import pandas as pd 

4import numpy as np 

5 

6from typing import Optional 

7from typing import ClassVar 

8 

9from abc import ABC 

10from abc import abstractmethod 

11 

12from dataclasses import dataclass 

13from dataclasses import field 

14 

15from enum import Enum 

16 

17 

18class GPUType(Enum): 

19 A100 = "A100" 

20 H100 = "H100" 

21 H200 = "H200" 

22 GB200 = "GB200" 

23 

24 def __lt__(self, other: object) -> bool: 

25 if not isinstance(other, GPUType): 

26 return NotImplemented 

27 order = [GPUType.A100, GPUType.H100, GPUType.H200, GPUType.GB200] 

28 return order.index(self) < order.index(other) 

29 

30 

31class QualityLevel(Enum): 

32 ORIGINAL = "original" 

33 HIGH = "high" 

34 MEDIUM = "medium" 

35 LOW = "low" 

36 

37 

38# Pixel counts per quality level (16:10 aspect ratio). 

39# Latency data is profiled at MEDIUM resolution. 

40RESOLUTION_PIXELS: dict[QualityLevel, int] = { 

41 QualityLevel.HIGH: 1280 * 800, 

42 QualityLevel.MEDIUM: 640 * 400, 

43 QualityLevel.LOW: 320 * 200, 

44} 

45 

46 

47class Model(Enum): 

48 GEMMA = "gemma" 

49 FLUX = "flux" 

50 HF = "hf" # HunyuanFramePack 

51 HF_VAE = "hf_vae" # HunyuanFramePack VAE 

52 FT = "ft" # FantasyTalking 

53 FT_VAE = "ft_vae" # FantasyTalking VAE 

54 UPSCALER = "upscaler" 

55 OTHERS = "others" # YOLO + Kokoro 

56 

57 

58# Used for FIFO 

59MODEL_ORDER: dict[Model, int] = { 

60 Model.GEMMA: 0, 

61 Model.FLUX: 1, 

62 Model.OTHERS: 2, 

63 Model.HF: 3, 

64 Model.HF_VAE: 4, 

65 Model.FT: 5, 

66 Model.FT_VAE: 6, 

67 Model.UPSCALER: 7, 

68} 

69 

70 

71@dataclass 

72class ModelAllocation(ABC): 

73 model: ClassVar[Model] 

74 

75 # policy TODO 

76 # workflow TODO 

77 gpu_type: GPUType 

78 devices: int = 1 

79 replicas: int = 0 # No replicas by default 

80 work: int = 0 

81 time: float = 0.0 

82 time_first: float = 0.0 

83 energy: float = 0.0 

84 cost: float = 0.0 

85 

86 def __str__(self) -> str: 

87 if self.replicas <= 0: 

88 assert self.time == 0.0, f"time must be 0 when no replicas, got {self.time:.2f}" 

89 assert self.energy == 0.0, f"energy must be 0 when no replicas, got {self.energy:.2f}" 

90 return "--" 

91 return \ 

92 f"devices={self.devices:2d}, " \ 

93 f"replicas={self.replicas}, " \ 

94 f"work={self.work}, " \ 

95 f"time={self.time:.2f} secs, " \ 

96 f"time_first={self.time_first:.2f} secs, " \ 

97 f"energy={self.energy / 60.0 / 60.0:.2f} Wh, " \ 

98 f"cost=${self.cost:.2f}" 

99 

100 def __repr__(self) -> str: 

101 return self.__str__() 

102 

103 def __post_init__(self) -> None: 

104 if self.replicas > 0: 

105 return 

106 if self.time != 0.0 or self.energy != 0.0: 

107 raise ValueError( 

108 f"time and energy must be 0.0 when no replicas, got time={self.time:.2f}, energy={self.energy:.2f}") 

109 

110 def get_num_gpus(self) -> int: 

111 if self.replicas <= 0: 

112 return 0 

113 return self.devices * self.replicas 

114 

115 def disable(self) -> None: 

116 self.devices = 0 

117 self.replicas = 0 

118 self.time = 0.0 

119 self.time_first = 0.0 

120 self.energy = 0.0 

121 

122 @abstractmethod 

123 def calculate_time( 

124 self, 

125 policy: Policy, 

126 workflow: WorkflowConfig, 

127 latency_data: LatencyData, 

128 work_pct: float = 1.0, 

129 ) -> float: 

130 ... 

131 

132 @abstractmethod 

133 def calculate_time_first( 

134 self, 

135 policy: Policy, 

136 workflow: WorkflowConfig, 

137 latency_data: LatencyData, 

138 ) -> float: 

139 ... 

140 

141 @abstractmethod 

142 def calculate_energy( 

143 self, 

144 workflow: WorkflowConfig, 

145 power_data: Optional[PowerData] = None, 

146 total_time_s: float = 0.0, 

147 ) -> float: 

148 ... 

149 

150 def calculate_cost( 

151 self, 

152 policy: Policy, 

153 total_time_s: float = 0.0, 

154 ) -> float: 

155 """Calculate the cost for this model allocation.""" 

156 SECONDS_IN_HOUR = 60 * 60 

157 gpu_cost = policy.gpu_cost[self.gpu_type] 

158 self.cost = total_time_s * (self.get_num_gpus() * gpu_cost) / SECONDS_IN_HOUR 

159 return self.cost 

160 

161 def calculate( 

162 self, 

163 policy: Policy, 

164 workflow: WorkflowConfig, 

165 latency_data: LatencyData, 

166 power_data: Optional[PowerData] = None, 

167 total_time_s: float = 0.0, 

168 work_pct: float = 1.0, 

169 ) -> None: 

170 """Calculate all the values for this model allocation.""" 

171 self.calculate_time(policy, workflow, latency_data, work_pct) 

172 self.calculate_time_first(policy, workflow, latency_data) 

173 self.calculate_cost(policy, total_time_s) 

174 self.calculate_energy(workflow, power_data, total_time_s) 

175 

176 def get_max_replicas( 

177 self, 

178 workflow: WorkflowConfig, 

179 ) -> int: 

180 """Get the maximum number of replicas that can leverage parallelism.""" 

181 return 1 

182 

183 

184class Objective(Enum): 

185 FIFO = "fifo" 

186 TIME = "time" 

187 TTFF = "ttff" 

188 COST = "cost" 

189 ENERGY = "energy" 

190 TIME_COST = "time_cost" 

191 TTFF_COST = "ttff_cost" 

192 ENERGY_COST = "energy_cost" 

193 TIME_ENERGY = "time_energy" 

194 RANDOM = "random" 

195 NONE = "none" 

196 

197 TTFF_THEN_TIME = "ttff_then_time" # first minimize ttff, then minimize time 

198 

199 def is_monotonic(self) -> bool: 

200 return self not in {Objective.RANDOM, Objective.FIFO} 

201 

202 

203@dataclass 

204class WorkflowConfig: 

205 total_video_seconds: int 

206 total_scenes: int 

207 total_frames: dict[Model, int] 

208 total_subscenes: int 

209 per_subscene_frames: dict[Model, int] 

210 # default per-frame number of denoising steps 

211 num_steps: dict[Model, int] 

212 # supported number of generation frames 

213 hf_frames: list[int] 

214 ft_frames: list[int] 

215 frames_per_step_idx: int 

216 # target output resolution (default: HIGH) 

217 target_resolution: QualityLevel = QualityLevel.HIGH 

218 

219 # total input tokens 

220 total_input_tokens: int = 0 

221 

222 # work per model (determines parallelism; work > 1 means parallelizable across replicas) 

223 # models included in the workflow are derived from the keys of this dict 

224 model_work: dict[Model, int] = field(default_factory=dict) 

225 

226 @property 

227 def models(self) -> list[Model]: 

228 """Models included in the workflow (derived from model_work keys).""" 

229 return list(self.model_work.keys()) 

230 

231 @property 

232 def work(self) -> dict[Model, int]: 

233 """Units of work per model (0 for models not in the workflow).""" 

234 return { 

235 model_name: self.model_work.get(model_name, 0) 

236 for model_name in Model 

237 } 

238 

239 def get_model_order(self) -> list[Model]: 

240 """Get ordered list of models in the workflow, sorted by MODEL_ORDER.""" 

241 return sorted( 

242 [m for m in self.models if m in MODEL_ORDER], 

243 key=lambda m: MODEL_ORDER[m], 

244 ) 

245 

246 def get_resolution_scale(self, use_upscaler: bool) -> float: 

247 """Compute latency scaling factor based on target resolution. 

248 

249 Latency data is profiled at MEDIUM resolution. The scale factor 

250 adjusts for the actual generation resolution: 

251 

252 1. Upscaler used, HIGH → 1.0 (models generate at MEDIUM) 

253 2. Upscaler used, MEDIUM → LOW / MEDIUM (models generate at LOW) 

254 3. No upscaler, HIGH → HIGH / MEDIUM (scale up) 

255 4. No upscaler, MEDIUM → 1.0 

256 5. No upscaler, LOW → LOW / MEDIUM (scale down) 

257 """ 

258 if use_upscaler: 

259 assert self.target_resolution in (QualityLevel.HIGH, QualityLevel.MEDIUM), \ 

260 "Upscaler can only be used when target resolution is HIGH or MEDIUM" 

261 if self.target_resolution == QualityLevel.HIGH: 

262 return 1.0 

263 # MEDIUM target with upscaler: generate at LOW, upscale to MEDIUM 

264 return RESOLUTION_PIXELS[QualityLevel.LOW] / RESOLUTION_PIXELS[QualityLevel.MEDIUM] 

265 if self.target_resolution == QualityLevel.MEDIUM: 

266 return 1.0 

267 return RESOLUTION_PIXELS[self.target_resolution] / RESOLUTION_PIXELS[QualityLevel.MEDIUM] 

268 

269 def is_parallelizable(self, model: Model) -> bool: 

270 """Whether the given model can be parallelized across multiple replicas.""" 

271 return self.model_work.get(model, 0) > 1 

272 

273 def filter_parallelizable_models( 

274 self, 

275 models: list[Model], 

276 disaggregation: dict[Model, bool], 

277 ) -> list[Model]: 

278 filtered_models = [ 

279 model 

280 for model in models 

281 if self.is_parallelizable(model) 

282 ] 

283 # Remove VAE models when their parent model disaggregation is disabled 

284 if not disaggregation.get(Model.HF, False): 

285 filtered_models = [m for m in filtered_models if m != Model.HF_VAE] 

286 if not disaggregation.get(Model.FT, False): 

287 filtered_models = [m for m in filtered_models if m != Model.FT_VAE] 

288 return filtered_models 

289 

290 def __post_init__(self) -> None: 

291 assert self.total_frames[Model.HF] > self.per_subscene_frames[Model.HF] 

292 assert self.total_frames[Model.FT] > self.per_subscene_frames[Model.FT] 

293 

294 # If no models specified, populate defaults for all models 

295 if not self.model_work: 

296 defaults: dict[Model, int] = { 

297 Model.GEMMA: 1, 

298 Model.FLUX: 1, 

299 Model.HF: self.total_subscenes, 

300 Model.HF_VAE: self.total_frames[Model.HF], 

301 Model.FT: self.total_subscenes, 

302 Model.FT_VAE: self.total_frames[Model.FT], 

303 Model.UPSCALER: self.total_frames[Model.FT], 

304 Model.OTHERS: 1, 

305 } 

306 for model, work in defaults.items(): 

307 self.model_work[model] = work 

308 if self.target_resolution != QualityLevel.HIGH: 

309 if Model.UPSCALER in self.model_work: 

310 del self.model_work[Model.UPSCALER] 

311 

312 @property 

313 def num_frames(self) -> int: 

314 """Number of frames generated by the workflow.""" 

315 if Model.FT in self.total_frames: 

316 return self.total_frames[Model.FT] 

317 return 0 

318 

319 

320class ActionName(Enum): 

321 MERGE = "merge" 

322 ADD_DEVICE = "add device" 

323 ADD_REPLICA = "add replica" 

324 ADD_DEVICE_REPLICA = "add device replica" 

325 ADD_INSTANCE = "add instance" 

326 REMOVE_DEVICE = "remove device" 

327 REMOVE_REPLICA = "remove replica" 

328 

329 

330@dataclass 

331class Action: 

332 """ 

333 Optimization action to take. 

334 """ 

335 name: ActionName 

336 model: Model 

337 gpu_type: GPUType 

338 models: dict[GPUType, dict[Model, list[ModelAllocation]]] 

339 

340 action_result: Result = field(repr=False) 

341 

342 arrival_time_s: float = 0.0 # For FIFO scheduling 

343 

344 # Derived fields from action_result (not passed by caller) 

345 time: float = field(init=False) # Total execution time 

346 ttff: float = field(init=False) # Time to first frame 

347 cost: float = field(init=False) # Cost in $ 

348 energy: float = field(init=False) # Energy in W*s 

349 

350 def __post_init__(self) -> None: 

351 # ---- type checks ---- 

352 if not isinstance(self.model, Model): 

353 raise ValueError(f"Model {self.model} [{type(self.model)}] not supported") 

354 if not isinstance(self.name, ActionName): 

355 raise ValueError(f"Action name {self.name} [{type(self.name)}] not supported") 

356 if not isinstance(self.models, dict): 

357 raise ValueError(f"models must be a dict, got {type(self.models)}") 

358 if not isinstance(self.gpu_type, GPUType): 

359 raise ValueError(f"Device type {self.gpu_type} [{type(self.gpu_type)}] not supported") 

360 """ 

361 if not isinstance(self.allocation_id, int) or self.allocation_id < 0: 

362 raise ValueError(f"Allocation ID {self.allocation_id} must be a non-negative integer") 

363 if self.num_replicas <= 0: 

364 raise ValueError(f"num_replicas {self.num_replicas} must be > 0") 

365 if self.num_devices <= 0: 

366 raise ValueError(f"num_devices {self.num_devices} must be > 0") 

367 """ 

368 # ---- derive values ---- 

369 self.time = self.action_result.total_time_s 

370 self.ttff = self.action_result.ttff_s 

371 self.cost = self.action_result.cost 

372 self.energy = self.action_result.total_energy 

373 if self.cost < 0.0: 

374 raise ValueError("cost must be >= 0") 

375 

376 def __str__(self) -> str: 

377 return ( 

378 f"Action(" 

379 f"{self.name.value}, " 

380 f"model={self.model.value}, " 

381 f"gpu={self.gpu_type.value}, " 

382 f"time={self.time:.2f} s, " 

383 f"ttff={self.ttff:.2f} s, " 

384 f"cost=${self.cost:.2f}, " 

385 f"time*cost={self.time_cost():.2f}, " 

386 f"ttff*cost={self.ttff_cost():.2f}, " 

387 f"energy*cost={self.energy_cost():.2f}, " 

388 f"time*energy={self.time_energy():.2f}, " 

389 f"energy={self.energy:.2f} Ws, " 

390 f"models={self.models}" 

391 f")" 

392 ) 

393 

394 def time_cost(self) -> float: 

395 """We use improvement in time * $.""" 

396 if self.time <= 0: 

397 return self.cost 

398 if self.cost <= 0: 

399 return self.time 

400 return self.time * self.cost 

401 

402 def ttff_cost(self) -> float: 

403 """We use improvement in TTFF * $.""" 

404 if self.ttff <= 0: 

405 return self.cost 

406 if self.cost <= 0: 

407 return self.ttff 

408 return self.ttff * self.cost 

409 

410 def energy_cost(self) -> float: 

411 """We use improvement in Wh * $.""" 

412 if self.cost <= 0: 

413 return self.energy 

414 if self.energy <= 0: 

415 return self.cost 

416 return self.energy * self.cost 

417 

418 def time_energy(self) -> float: 

419 """We use improvement in TTFF * Wh.""" 

420 if self.energy <= 0: 

421 return self.time 

422 if self.time <= 0: 

423 return self.energy 

424 return self.time * self.energy 

425 

426 def get_order(self) -> int: 

427 " ""For FIFO scheduling."" " 

428 return MODEL_ORDER[self.model] 

429 

430 def get_metric( 

431 self, 

432 obj: Objective, 

433 switch_objective: bool = False, 

434 ) -> float: 

435 if obj == Objective.RANDOM: 

436 return 0.0 

437 if obj == Objective.TIME: 

438 return self.time 

439 if obj == Objective.TTFF: 

440 return self.ttff 

441 if obj == Objective.COST: 

442 return self.cost 

443 if obj == Objective.ENERGY: 

444 return self.energy 

445 if obj == Objective.TIME_COST: 

446 return self.time_cost() 

447 if obj == Objective.TTFF_COST: 

448 return self.ttff_cost() 

449 if obj == Objective.ENERGY_COST: 

450 return self.energy_cost() 

451 if obj == Objective.TIME_ENERGY: 

452 return self.time_energy() 

453 if obj == Objective.FIFO: 

454 # return self.get_order() 

455 return 0 # TODO 

456 if obj == Objective.TTFF_THEN_TIME: 

457 if switch_objective: 

458 return self.time 

459 else: 

460 return self.ttff 

461 raise ValueError(f"Unknown objective {obj}") 

462 

463 

464@dataclass 

465class Result: 

466 total_time_s: float = 0.0 

467 first_chunk_time: float = 0.0 # Time to first chunk 

468 ttff_s: float = 0.0 # Time to first frame (accounts for total time and workflow length) 

469 tbf_s: float = 0.0 # Time between frames 

470 total_energy: float = 0.0 # Watts x second 

471 cost: float = 0.0 # Total $ cost 

472 gpus_used: dict[GPUType, int] = field(default_factory=dict) 

473 gpus_total: dict[GPUType, int] = field(default_factory=dict) 

474 models: dict[GPUType, dict[Model, list[ModelAllocation]]] = field(default_factory=dict) 

475 

476 def __post_init__(self) -> None: 

477 assert self.total_time_s >= 0.0, f"total_time_s={self.total_time_s} must be >= 0.0" 

478 assert self.first_chunk_time >= 0.0, f"first_chunk_time={self.first_chunk_time} must be >= 0.0" 

479 assert self.ttff_s >= 0.0, f"ttff_s={self.ttff_s} must be >= 0.0" 

480 assert self.tbf_s >= 0.0, f"tbf_s={self.tbf_s} must be >= 0.0" 

481 assert self.total_energy >= 0.0, f"total_energy={self.total_energy} must be >= 0.0" 

482 assert self.cost >= 0.0, f"cost={self.cost} must be >= 0.0" 

483 assert len(self.gpus_used) >= 0, f"gpus_used cannot be empty: {self.gpus_used}" 

484 for gpu_used in self.gpus_used.values(): 

485 assert gpu_used >= 0, f"all gpus_used value {self.gpus_used} must be >= 0" 

486 

487 def to_csv(self) -> str: 

488 num_a100 = self.gpus_used.get(GPUType.A100, 0) 

489 num_h100 = self.gpus_used.get(GPUType.H100, 0) 

490 num_h200 = self.gpus_used.get(GPUType.H200, 0) 

491 num_gb200 = self.gpus_used.get(GPUType.GB200, 0) 

492 return ( 

493 f"{num_a100},{num_h100},{num_h200},{num_gb200}," 

494 f"{self.ttff_s:.2f},{self.tbf_s:.2f},{self.cost:.2f}," 

495 f"{self.total_time_s:.2f},{self.total_energy:.2f}" 

496 ) 

497 

498 def __str__(self) -> str: 

499 SECONDS_IN_HOUR = 60 * 60 

500 return ( 

501 f"Time:{self.total_time_s:.2f} s TTFF:{self.ttff_s:.2f} s " 

502 f"Cost:${self.cost:.2f} TTFF*Cost:{self.ttff_s * self.cost:.2f} " 

503 f"Energy:{self.total_energy / SECONDS_IN_HOUR / 1000:.2f} kWh " 

504 f"GPUS: {num_gpus_to_str(self.gpus_used)}" 

505 ) 

506 

507 def __repr__(self) -> str: 

508 return self.__str__() 

509 

510 

511@dataclass 

512class LatencyGPUTypeData: 

513 gpu_type: GPUType 

514 # TP -> latency mappings 

515 flux: dict[int, float] = field(default_factory=dict) 

516 hf: dict[int, float] = field(default_factory=dict) 

517 hf_high: dict[int, float] = field(default_factory=dict) 

518 hf_vae: dict[int, float] = field(default_factory=dict) 

519 hf_vae_high: dict[int, float] = field(default_factory=dict) 

520 ft: dict[int, float] = field(default_factory=dict) 

521 ft_high: dict[int, float] = field(default_factory=dict) 

522 ft_vae: dict[int, float] = field(default_factory=dict) 

523 ft_vae_high: dict[int, float] = field(default_factory=dict) 

524 upscaler: dict[int, float] = field(default_factory=dict) 

525 gemma_first_scene: dict[int, float] = field(default_factory=dict) 

526 gemma_per_scene: dict[int, float] = field(default_factory=dict) 

527 others: dict[int, float] = field(default_factory=dict) 

528 

529 def __getitem__( 

530 self, 

531 key: Model | tuple[Model, int] 

532 ) -> float: 

533 if isinstance(key, tuple): 

534 assert isinstance(key[0], Model) 

535 assert isinstance(key[1], int) 

536 model, num_devices = key 

537 if model == Model.FLUX: 

538 return self.flux[num_devices] 

539 if model == Model.HF: 

540 return self.hf[num_devices] 

541 if model == Model.HF_VAE: 

542 return self.hf_vae[num_devices] 

543 if model == Model.FT: 

544 return self.ft[num_devices] 

545 if model == Model.FT_VAE: 

546 return self.ft_vae[num_devices] 

547 if model == Model.GEMMA: 

548 return self.gemma_first_scene[num_devices] 

549 if model == Model.UPSCALER: 

550 return self.upscaler[num_devices] 

551 if model == Model.OTHERS: 

552 return self.others[num_devices] 

553 raise KeyError(f"Latency for model {key} not found") 

554 

555 def __contains__(self, key: Model | tuple[Model, int]) -> bool: 

556 if isinstance(key, tuple): 

557 assert isinstance(key[0], Model) 

558 assert isinstance(key[1], int) 

559 model, num_devices = key 

560 if model == Model.GEMMA: 

561 return num_devices in self.gemma_first_scene 

562 if model == Model.FLUX: 

563 return num_devices in self.flux 

564 if model == Model.HF: 

565 return num_devices in self.hf 

566 if model == Model.HF_VAE: 

567 return num_devices in self.hf_vae 

568 if model == Model.FT: 

569 return num_devices in self.ft 

570 if model == Model.FT_VAE: 

571 return num_devices in self.ft_vae 

572 if model == Model.UPSCALER: 

573 return num_devices in self.upscaler 

574 if model == Model.HF_VAE: 

575 return num_devices in self.hf_vae 

576 if model == Model.OTHERS: 

577 return num_devices in self.others 

578 return False 

579 

580 def get_max_parallelism(self, model: Model) -> int: 

581 """Max number of devices supported for the given model.""" 

582 if model == Model.FLUX: 

583 return max(self.flux.keys()) 

584 if model == Model.HF: 

585 return max(self.hf.keys()) 

586 if model == Model.FT: 

587 return max(self.ft.keys()) 

588 if model == Model.FT_VAE: 

589 return max(self.ft_vae.keys()) 

590 if model == Model.GEMMA: 

591 return max(self.gemma_first_scene.keys()) 

592 if model == Model.UPSCALER: 

593 return max(self.upscaler.keys()) 

594 if model == Model.HF_VAE: 

595 return max(self.hf_vae.keys()) 

596 if model == Model.OTHERS: 

597 return max(self.others.keys()) 

598 raise KeyError(f"Model {model} not found in latency data") 

599 

600 

601@dataclass 

602class PowerGPUTypeData: 

603 gpu_type: GPUType 

604 # TP -> power mappings 

605 flux: dict[int, float] = field(default_factory=dict) 

606 hf: dict[int, float] = field(default_factory=dict) 

607 hf_high: dict[int, float] = field(default_factory=dict) 

608 hf_vae: dict[int, float] = field(default_factory=dict) 

609 hf_vae_high: dict[int, float] = field(default_factory=dict) 

610 ft: dict[int, float] = field(default_factory=dict) 

611 ft_high: dict[int, float] = field(default_factory=dict) 

612 ft_vae: dict[int, float] = field(default_factory=dict) 

613 ft_vae_high: dict[int, float] = field(default_factory=dict) 

614 upscaler: dict[int, float] = field(default_factory=dict) 

615 gemma_first_scene: dict[int, float] = field(default_factory=dict) 

616 gemma_per_scene: dict[int, float] = field(default_factory=dict) 

617 # Other values 

618 idle: float = 0.0 # Idle power in Watts 

619 tdp: float = 0.0 # TDP power in Watts 

620 

621 def __getitem__( 

622 self, 

623 key: Model | tuple[Model, int] | str 

624 ) -> float: 

625 if isinstance(key, tuple): 

626 assert isinstance(key[0], Model) 

627 assert isinstance(key[1], int) 

628 model, devices = key 

629 if model == Model.FLUX: 

630 return self.flux[devices] 

631 if model == Model.HF: 

632 return self.hf[devices] 

633 if model == Model.HF_VAE: 

634 return self.hf_vae[devices] 

635 if model == Model.FT: 

636 return self.ft[devices] 

637 if model == Model.FT_VAE: 

638 return self.ft_vae[devices] 

639 if model == Model.UPSCALER: 

640 return self.upscaler[devices] 

641 if isinstance(key, str): 

642 if key == "idle": 

643 return self.idle 

644 if key == "tdp": 

645 return self.tdp 

646 raise KeyError(f"Power for {key} not found") 

647 

648 

649@dataclass 

650class LatencyData: 

651 gpus: dict[GPUType, LatencyGPUTypeData] 

652 

653 def __getitem__(self, gpu_type: GPUType) -> LatencyGPUTypeData: 

654 return self.gpus[gpu_type] 

655 

656 def __setitem__( 

657 self, 

658 gpu_type: GPUType, 

659 latency_data: LatencyGPUTypeData 

660 ) -> None: 

661 self.gpus[gpu_type] = latency_data 

662 

663 

664@dataclass 

665class PowerData: 

666 gpus: dict[GPUType, PowerGPUTypeData] 

667 

668 def __getitem__(self, gpu_type: GPUType) -> PowerGPUTypeData: 

669 return self.gpus[gpu_type] 

670 

671 def __setitem__( 

672 self, 

673 gpu_type: GPUType, 

674 power_data: PowerGPUTypeData 

675 ) -> None: 

676 self.gpus[gpu_type] = power_data 

677 

678 

679def num_gpus_to_str( 

680 provision: dict[GPUType, int] 

681) -> str: 

682 return "+".join([ 

683 f"{num_gpus}x{gpu_type.name}" 

684 for gpu_type, num_gpus in provision.items() 

685 if num_gpus > 0 

686 ]) 

687 

688 

689@dataclass 

690class Provision: 

691 num_gpus: dict[GPUType, int] = field(default_factory=dict) 

692 

693 def __getitem__(self, gpu_type: GPUType) -> int: 

694 return self.num_gpus[gpu_type] 

695 

696 def __str__(self) -> str: 

697 return num_gpus_to_str(self.num_gpus) 

698 

699 

700@dataclass 

701class ProvisioningResult: 

702 latencies: list[float] 

703 costs: list[float] 

704 ttffs: list[float] 

705 tbfs: list[float] 

706 actual_provision: list[dict[GPUType, int]] 

707 config_provision: list[dict[GPUType, int]] 

708 model_provision: list[dict[GPUType, dict[Model, list[ModelAllocation]]]] 

709 qualities: list[float] = field(default_factory=list) 

710 energies: list[float] = field(default_factory=list) 

711 

712 def save( 

713 self, 

714 policy_name: str, 

715 results_dir: str, 

716 ) -> None: 

717 """Save the provisioning results to a CSV file.""" 

718 num_a100: list[int] = [] 

719 num_h100: list[int] = [] 

720 num_h200: list[int] = [] 

721 num_gb200: list[int] = [] 

722 for provision in self.actual_provision: 

723 num_a100.append(provision.get(GPUType.A100, 0)) 

724 num_h100.append(provision.get(GPUType.H100, 0)) 

725 num_h200.append(provision.get(GPUType.H200, 0)) 

726 num_gb200.append(provision.get(GPUType.GB200, 0)) 

727 df_latency = pd.DataFrame({ 

728 'num_a100': num_a100, 

729 'num_h100': num_h100, 

730 'num_h200': num_h200, 

731 'num_gb200': num_gb200, 

732 'ttff_s': self.ttffs, 

733 'tbf_s': self.tbfs, 

734 'cost': self.costs,