Coverage for simulator/multirequests.py: 97%

176 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-09 04:47 +0000

1from __future__ import annotations 

2 

3import math 

4import os 

5from dataclasses import replace 

6 

7from sim_types import GPUType 

8from sim_types import Model 

9from sim_types import QualityLevel 

10from sim_types import RESOLUTION_PIXELS 

11from sim_types import Result 

12from sim_types import WorkflowConfig 

13from sim_types import LatencyData 

14 

15from data_loading import load_latency_data 

16from data_loading import load_power_data 

17from data_loading import load_adaptive_quality_data 

18 

19from workflows import PODCAST_WORKFLOW 

20 

21from model_provisioner.policies import STREAMWISE_POLICY 

22 

23from auto_model_allocator import AutoModelAllocator 

24 

25 

26# Queries per minute 

27QPM_LIST = [0.1, 1, 2, 5, 10, 20, 30, 50, 100] 

28 

29# Resolve the data directory relative to this file so imports work from any cwd. 

30_DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data") 

31 

32 

33# --------------------------------------------------------------------------- 

34# Hardware budget — the Pareto-optimal operating point used in the paper. 

35# --------------------------------------------------------------------------- 

36HARDWARE_BUDGET: dict[GPUType, int] = { 

37 GPUType.A100: 256, 

38 GPUType.H100: 64, 

39} 

40 

41 

42# --------------------------------------------------------------------------- 

43# Derivation helpers 

44# --------------------------------------------------------------------------- 

45 

46def _extract_from_result( 

47 result: Result, 

48) -> tuple[dict[GPUType, dict[Model, int]], dict[GPUType, dict[Model, float]]]: 

49 """Extract init_replicas (GPU counts) and time_per_req from a simulation result. 

50 

51 Returns 

52 ------- 

53 init_replicas: 

54 ``{gpu_type: {model: total_gpus}}`` — total GPU count allocated to each 

55 model on each GPU type (i.e. ``devices × replicas`` summed across instances). 

56 time_per_req: 

57 ``{gpu_type: {model: seconds}}`` — wall-clock time for the model to process 

58 one full request (10-min video) given the allocated resources. When a model 

59 has multiple instances on the same GPU type, we take the *maximum* time 

60 (the bottleneck). 

61 """ 

62 init_replicas: dict[GPUType, dict[Model, int]] = {} 

63 time_per_req: dict[GPUType, dict[Model, float]] = {} 

64 

65 for gpu_type, model_allocs in result.models.items(): 

66 init_replicas[gpu_type] = {} 

67 time_per_req[gpu_type] = {} 

68 for model, allocs in model_allocs.items(): 

69 total_gpus = sum(a.get_num_gpus() for a in allocs) 

70 times = [a.time for a in allocs if a.get_num_gpus() > 0] 

71 if total_gpus > 0: 

72 init_replicas[gpu_type][model] = total_gpus 

73 time_per_req[gpu_type][model] = max(times) if times else 0.0 

74 

75 return init_replicas, time_per_req 

76 

77 

78def derive_multirequest_params( 

79 budget: dict[GPUType, int] | None = None, 

80 data_dir: str = _DATA_DIR, 

81) -> tuple[dict[GPUType, dict[Model, int]], dict[GPUType, dict[Model, float]]]: 

82 """Run the StreamWise simulator and derive multi-request parameters. 

83 

84 Runs the greedy allocator with ``STREAMWISE_POLICY`` on ``PODCAST_WORKFLOW`` 

85 at the given hardware *budget* and extracts: 

86 

87 * **init_replicas** — total GPU count per model per GPU type 

88 * **time_per_req** — total time (seconds) per request per model per GPU type 

89 

90 Parameters 

91 ---------- 

92 budget: 

93 ``{GPUType: num_gpus}`` hardware budget to allocate. 

94 Defaults to ``HARDWARE_BUDGET`` when ``None``. 

95 data_dir: 

96 Path to the latency/power CSV data directory. 

97 """ 

98 if budget is None: 

99 budget = dict(HARDWARE_BUDGET) 

100 latency_data = load_latency_data(data_dir=data_dir) 

101 power_data = load_power_data(data_dir=data_dir) 

102 

103 allocator = AutoModelAllocator( 

104 workflow=PODCAST_WORKFLOW, 

105 latency_data=latency_data, 

106 power_data=power_data, 

107 policy=STREAMWISE_POLICY, 

108 ) 

109 result = allocator.allocate( 

110 num_gpus=budget, 

111 verbose=False, 

112 ) 

113 

114 return _extract_from_result(result) 

115 

116 

117def derive_adaptive_params( 

118 budget: dict[GPUType, int] | None = None, 

119 data_dir: str = _DATA_DIR, 

120) -> tuple[ 

121 dict[GPUType, dict[Model, int]], 

122 dict[GPUType, dict[Model, dict[QualityLevel, float]]], 

123]: 

124 """Run the simulator at each quality level and derive adaptive parameters. 

125 

126 Returns 

127 ------- 

128 init_replicas_adaptive: 

129 ``{gpu_type: {model: total_gpus}}`` from the HIGH-quality simulation run 

130 (the worst-case / most-demanding quality level sets the base allocation). 

131 time_per_req_adaptive: 

132 ``{gpu_type: {model: {quality: seconds}}}`` — per-quality time per request, 

133 every ``(gpu_type, model)`` in ``init_replicas_adaptive`` has a timing 

134 entry for every quality level. 

135 """ 

136 if budget is None: 

137 budget = dict(HARDWARE_BUDGET) 

138 

139 power_data = load_power_data(data_dir=data_dir) 

140 

141 qualities = [QualityLevel.HIGH, QualityLevel.MEDIUM, QualityLevel.LOW] 

142 results_by_quality: dict[QualityLevel, Result] = {} 

143 for quality in qualities: 

144 policy = replace(STREAMWISE_POLICY) 

145 policy.name = f"{STREAMWISE_POLICY.name} {quality.value}" 

146 

147 latency_data = load_adaptive_quality_data( 

148 data_dir=data_dir, 

149 level=quality, 

150 ) 

151 

152 allocator = AutoModelAllocator( 

153 workflow=PODCAST_WORKFLOW, 

154 latency_data=latency_data, 

155 power_data=power_data, 

156 policy=policy, 

157 ) 

158 result = allocator.allocate( 

159 num_gpus=budget, 

160 verbose=False, 

161 ) 

162 results_by_quality[quality] = result 

163 

164 init_replicas_adaptive, time_per_req_high = _extract_from_result( 

165 results_by_quality[QualityLevel.HIGH], 

166 ) 

167 

168 time_per_req_by_quality: dict[QualityLevel, dict[GPUType, dict[Model, float]]] = {} 

169 for quality, result in results_by_quality.items(): 

170 _, time_per_req_q = _extract_from_result(result) 

171 time_per_req_by_quality[quality] = time_per_req_q 

172 

173 time_per_req_adaptive: dict[GPUType, dict[Model, dict[QualityLevel, float]]] = {} 

174 for gpu_type, models in init_replicas_adaptive.items(): 

175 time_per_req_adaptive[gpu_type] = {} 

176 for model in models: 

177 high_time = time_per_req_high[gpu_type][model] 

178 quality_times: dict[QualityLevel, float] = {} 

179 for quality in qualities: 

180 quality_times[quality] = ( 

181 time_per_req_by_quality 

182 .get(quality, {}) 

183 .get(gpu_type, {}) 

184 .get(model, high_time) 

185 ) 

186 time_per_req_adaptive[gpu_type][model] = quality_times 

187 

188 return init_replicas_adaptive, time_per_req_adaptive 

189 

190 

191# --------------------------------------------------------------------------- 

192# Derived constants — computed by running the simulator at HARDWARE_BUDGET. 

193# 

194# TIME_PER_REQ / INIT_REPLICAS: single (HIGH) quality operating point. 

195# TIME_PER_REQ_ADAPTIVE / INIT_REPLICAS_ADAPTIVE: per-quality-level values. 

196# --------------------------------------------------------------------------- 

197INIT_REPLICAS, TIME_PER_REQ = derive_multirequest_params(budget=dict(HARDWARE_BUDGET)) 

198INIT_REPLICAS_ADAPTIVE, TIME_PER_REQ_ADAPTIVE = derive_adaptive_params(budget=dict(HARDWARE_BUDGET)) 

199 

200# Allocation of video frames across quality levels for a 10-minute output 

201# to fulfill a TTFF SLO. These are configurable weights used in adaptive 

202# quality cost aggregation. 

203QUALITY_PORTIONS = { 

204 QualityLevel.LOW: 112, 

205 QualityLevel.MEDIUM: 305, 

206 QualityLevel.HIGH: 13383, 

207} 

208 

209 

210# Initial setup based on minimal 8 A100 configuration 

211# 1 for Kokoro, 1 for Gemma, 1 for Flux, 1 for HF+VAE (co-located), 4 for FT 

212INIT_REPLICAS_BASELINE: dict[GPUType, dict[Model, int]] = { 

213 GPUType.A100: { 

214 Model.OTHERS: 1, 

215 Model.GEMMA: 1, 

216 Model.FLUX: 1, 

217 Model.HF: 1, # HF and VAE co-located on same GPU 

218 Model.FT: 4, 

219 }, 

220 GPUType.H100: { 

221 # Empty for baseline 

222 } 

223} 

224 

225 

226def get_time_per_request_baseline( 

227 workflow_config: WorkflowConfig, 

228 latency_data: LatencyData, 

229 init_replicas: dict[GPUType, dict[Model, int]] = INIT_REPLICAS_BASELINE, 

230) -> dict[GPUType, dict[Model, float]]: 

231 """Get time per request for baseline (single quality).""" 

232 # Calculate time per request for each component (using baseline latencies) 

233 # Using A100 latencies from the Naive Baseline section 

234 # NOTE: In naive baseline, HF and VAE are co-located and run sequentially (not concurrently) 

235 

236 total_scenes = workflow_config.total_scenes 

237 num_steps_flux = workflow_config.num_steps[Model.FLUX] 

238 

239 total_frames_hf = workflow_config.total_frames[Model.HF] 

240 num_steps_hf = workflow_config.num_steps[Model.HF] 

241 hf_frames = workflow_config.hf_frames 

242 frames_per_step_idx = workflow_config.frames_per_step_idx 

243 total_frames_ft = workflow_config.total_frames[Model.FT] 

244 num_steps_ft = workflow_config.num_steps[Model.FT] 

245 ft_frames = workflow_config.ft_frames 

246 

247 num_pixels_high = RESOLUTION_PIXELS[QualityLevel.HIGH] 

248 num_pixels_medium = RESOLUTION_PIXELS[QualityLevel.MEDIUM] 

249 

250 # Latencies 

251 latency_hf_mapping_a100 = { 

252 k: v * num_pixels_high / num_pixels_medium 

253 for k, v in latency_data.gpus[GPUType.A100].hf.items() 

254 } 

255 latency_hf_vae_a100 = latency_data.gpus[GPUType.A100][Model.HF_VAE, 1] * num_pixels_high / num_pixels_medium 

256 latency_ft_mapping_a100 = { 

257 k: v * num_pixels_high / num_pixels_medium 

258 for k, v in latency_data.gpus[GPUType.A100].ft.items() 

259 } 

260 latency_ft_vae_a100 = latency_data.gpus[GPUType.A100][Model.FT_VAE, 1] * num_pixels_high / num_pixels_medium 

261 

262 num_gemma_gpus = 1 

263 num_flux_gpus = 1 

264 num_hf_gpus = 1 

265 num_ft_gpus = 1 

266 num_ft_replicas = init_replicas[GPUType.A100][Model.FT] 

267 

268 latency_gemma_first = latency_data.gpus[GPUType.A100].gemma_first_scene[num_gemma_gpus] 

269 latency_gemma_per = latency_data.gpus[GPUType.A100].gemma_per_scene[num_gemma_gpus] 

270 latency_flux = latency_data.gpus[GPUType.A100][Model.FLUX, num_flux_gpus] 

271 latency_hf = latency_hf_mapping_a100[num_hf_gpus] 

272 time_hf = ( 

273 (total_frames_hf / hf_frames[frames_per_step_idx] * latency_hf * num_steps_hf) 

274 + (total_frames_hf / hf_frames[frames_per_step_idx] * latency_hf_vae_a100) 

275 ) 

276 latency_ft = latency_ft_mapping_a100[num_ft_gpus] / num_ft_replicas 

277 time_ft = ( 

278 (total_frames_ft / ft_frames[frames_per_step_idx] * latency_ft * num_steps_ft) 

279 + (total_frames_ft / ft_frames[frames_per_step_idx] * latency_ft_vae_a100) 

280 ) 

281 

282 return { 

283 GPUType.A100: { 

284 Model.OTHERS: total_scenes * 0.6, 

285 Model.GEMMA: latency_gemma_first + latency_gemma_per * (total_scenes - 1), 

286 Model.FLUX: latency_flux * num_steps_flux, 

287 Model.HF: time_hf, 

288 Model.FT: time_ft, 

289 }, 

290 GPUType.H100: { 

291 # None in baseline 

292 }, 

293 } 

294 

295 

296def aggregate_time_per_request_by_quality( 

297 time_per_req: dict[GPUType, dict[Model, float | dict[QualityLevel, float]]], 

298 quality_portions: dict[QualityLevel, int], 

299) -> dict[GPUType, dict[Model, float]]: 

300 """Aggregate time per request metrics.""" 

301 ret: dict[GPUType, dict[Model, float]] = {} 

302 

303 total_portions = sum(quality_portions.values()) 

304 

305 for gpu_type in time_per_req.keys(): 

306 ret[gpu_type] = {} 

307 for model in time_per_req[gpu_type].keys(): 

308 val = time_per_req[gpu_type][model] 

309 if isinstance(val, float): 

310 time_val: float = val 

311 ret[gpu_type][model] = time_val 

312 elif isinstance(val, dict): 

313 dict_quality: dict[QualityLevel, float] = val 

314 agg_val = 0.0 

315 for quality_level in dict_quality.keys(): 

316 fraction = quality_portions[quality_level] / total_portions 

317 agg_val += dict_quality[quality_level] * fraction 

318 ret[gpu_type][model] = agg_val 

319 else: 

320 raise ValueError("Invalid time_per_req format") 

321 return ret 

322 

323 

324def required_replicas( 

325 name: str, 

326 video_seconds: float, 

327 ttff: float, 

328 per_sec: float, 

329 partition: str, 

330 req_per_min: float, 

331) -> float: 

332 """Calculate required replicas for a model.""" 

333 ttff_total = 0.0 

334 per_sec_total = 0.0 

335 

336 if partition == "scenes": 

337 ttff_total = ttff 

338 per_sec_total = (video_seconds * per_sec) 

339 elif partition == "frames": 

340 if name == "hf_vae": 

341 ttff_total = ttff 

342 per_sec_total = (video_seconds * per_sec) 

343 if name == "upscaler": 

344 ttff_total = ttff 

345 per_sec_total = (video_seconds * per_sec) 

346 elif partition == "subscenes": 

347 ttff_total = ttff 

348 per_sec_total = (video_seconds * per_sec) 

349 else: 

350 ttff_total = ttff 

351 per_sec_total = video_seconds * per_sec 

352 

353 total_time_per_request = ttff_total + per_sec_total 

354 total_time_per_minute = total_time_per_request * req_per_min 

355 return total_time_per_minute / 60.0 

356 

357 

358def get_replicas( 

359 video_seconds: float = 10 * 60, # 10 minutes video 

360 requests_per_minute: float = 0.5, 

361 time_per_req: dict[GPUType, dict[Model, float]] = TIME_PER_REQ, 

362 init_replicas: dict[GPUType, dict[Model, int]] = INIT_REPLICAS, 

363 qpms: list[float] = QPM_LIST, 

364) -> dict[GPUType, dict[Model, list[int]]]: 

365 """Get required replicas for different QPM levels.""" 

366 assert video_seconds > 0 

367 assert requests_per_minute > 0 

368 assert 0 < len(time_per_req) == len(init_replicas) 

369 assert len(qpms) > 0 

370 

371 video_minutes = video_seconds / 60 

372 

373 capacity: dict[GPUType, dict[Model, float]] = {} 

374 for gpu_type in time_per_req.keys(): 

375 capacity[gpu_type] = {} 

376 for model in time_per_req[gpu_type].keys(): 

377 capacity[gpu_type][model] = video_seconds / time_per_req[gpu_type][model] 

378 

379 replicas: dict[GPUType, dict[Model, list[int]]] = {} 

380 for qpm in qpms: 

381 arrival = qpm * video_minutes 

382 for gpu_type in init_replicas.keys(): 

383 if gpu_type not in replicas: 

384 replicas[gpu_type] = {} 

385 for model in init_replicas[gpu_type].keys(): 

386 if model not in replicas[gpu_type]: 

387 replicas[gpu_type][model] = [] 

388 # num_replicas = arrival * time_per_req[gpu_type][model] 

389 scale_factor = max(1, arrival / capacity[gpu_type][model]) 

390 num_replicas = math.ceil(init_replicas[gpu_type][model] * scale_factor) 

391 replicas[gpu_type][model].append(num_replicas) 

392 

393 return replicas 

394 

395 

396def get_costs( 

397 replicas: dict[GPUType, dict[Model, list[int]]], 

398 gpu_costs: dict[GPUType, float], 

399) -> dict[GPUType, dict[Model, list[float]]]: 

400 costs: dict[GPUType, dict[Model, list[float]]] = {} 

401 for gpu_type in replicas.keys(): 

402 costs[gpu_type] = {} 

403 for model in replicas[gpu_type].keys(): 

404 costs[gpu_type][model] = [ 

405 replica * gpu_costs[gpu_type] 

406 for replica in replicas[gpu_type][model] 

407 ] 

408 return costs 

409 

410 

411def get_total_costs( 

412 costs: dict[GPUType, dict[Model, list[float]]], 

413 qpms: list[float] = QPM_LIST, 

414) -> list[float]: 

415 total_costs = [ 

416 sum( 

417 costs[gpu_type][model][i] 

418 for gpu_type in costs.keys() 

419 for model in costs[gpu_type].keys() 

420 ) 

421 for i in range(len(qpms)) 

422 ] 

423 return total_costs