Coverage for simulator/multirequests.py: 97%
176 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1from __future__ import annotations
3import math
4import os
5from dataclasses import replace
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
15from data_loading import load_latency_data
16from data_loading import load_power_data
17from data_loading import load_adaptive_quality_data
19from workflows import PODCAST_WORKFLOW
21from model_provisioner.policies import STREAMWISE_POLICY
23from auto_model_allocator import AutoModelAllocator
26# Queries per minute
27QPM_LIST = [0.1, 1, 2, 5, 10, 20, 30, 50, 100]
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")
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}
42# ---------------------------------------------------------------------------
43# Derivation helpers
44# ---------------------------------------------------------------------------
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.
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]] = {}
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
75 return init_replicas, time_per_req
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.
84 Runs the greedy allocator with ``STREAMWISE_POLICY`` on ``PODCAST_WORKFLOW``
85 at the given hardware *budget* and extracts:
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
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)
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 )
114 return _extract_from_result(result)
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.
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)
139 power_data = load_power_data(data_dir=data_dir)
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}"
147 latency_data = load_adaptive_quality_data(
148 data_dir=data_dir,
149 level=quality,
150 )
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
164 init_replicas_adaptive, time_per_req_high = _extract_from_result(
165 results_by_quality[QualityLevel.HIGH],
166 )
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
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
188 return init_replicas_adaptive, time_per_req_adaptive
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))
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}
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}
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)
236 total_scenes = workflow_config.total_scenes
237 num_steps_flux = workflow_config.num_steps[Model.FLUX]
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
247 num_pixels_high = RESOLUTION_PIXELS[QualityLevel.HIGH]
248 num_pixels_medium = RESOLUTION_PIXELS[QualityLevel.MEDIUM]
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
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]
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 )
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 }
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]] = {}
303 total_portions = sum(quality_portions.values())
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
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
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
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
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
371 video_minutes = video_seconds / 60
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]
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)
393 return replicas
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
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