Coverage for streamwise/allocator_bridge.py: 92%
167 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Bridge between the model provisioner's allocator output and StreamWise pod deployment.
4Translates ModelAllocation results (abstract Model enum + GPU counts) into concrete
5container deployment parameters compatible with pod_manager.add_pod().
6"""
8from __future__ import annotations
10import sys
11import os
13# Ensure the directory containing this file is on sys.path so model_provisioner is importable.
14_HERE = os.path.dirname(os.path.abspath(__file__))
15if _HERE not in sys.path:
16 sys.path.insert(0, _HERE)
17_REPO_ROOT = os.path.dirname(_HERE)
18if _REPO_ROOT not in sys.path:
19 sys.path.insert(0, _REPO_ROOT)
20import model_provisioner # noqa: E402, F401 — adds simulator/ to sys.path
22from dataclasses import dataclass
23from typing import Optional
25from sim_types import GPUType
26from sim_types import Model
27from sim_types import Result
29from auto_model_allocator import AutoModelAllocator
30from container_config import COLOCATED_CONTAINERS
31from container_config import CONTAINER_RESOURCES
32from container_config import ContainerResourceSpec
33from container_config import GPU_TYPE_TO_POD_STR
34from container_config import MIG_AVAILABLE
35from container_config import MIG_CAPABLE_GPU_TYPES
36from container_config import MIG_CONTAINERS
37from container_config import MODEL_TO_CONTAINER_NAME
38from data_loading import load_latency_data
39from model_provisioner.policies import STREAMWISE_POLICY
40from streamwise_apps import STREAMWISE_APPS
41from workflows import WORKFLOWS
44# Mapping from simulator Model enum to concrete container names used by pod_manager.
45# Some Model entries map to multiple containers (e.g., OTHERS -> kokoro + yolo).
46MODEL_TO_CONTAINERS: dict[Model, list[str]] = {
47 Model.GEMMA: [MODEL_TO_CONTAINER_NAME[Model.GEMMA]],
48 Model.FLUX: [MODEL_TO_CONTAINER_NAME[Model.FLUX]],
49 Model.HF: [MODEL_TO_CONTAINER_NAME[Model.HF]],
50 Model.HF_VAE: [MODEL_TO_CONTAINER_NAME[Model.HF_VAE]],
51 Model.FT: [MODEL_TO_CONTAINER_NAME[Model.FT]],
52 Model.FT_VAE: [], # FT_VAE is handled within fantasytalking container
53 Model.UPSCALER: ["realesrgan"],
54 Model.OTHERS: ["kokoro", "yolo"],
55}
58def get_mig_profile(container_name: str, gpu_type: GPUType) -> Optional[str]:
59 """Return a MIG profile only when MIG is available and the GPU type supports it."""
60 if not MIG_AVAILABLE:
61 return None
62 if gpu_type not in MIG_CAPABLE_GPU_TYPES:
63 return None
64 return MIG_CONTAINERS.get(container_name)
67# Mapping from StreamWise app name to simulator workflow key
68APP_TO_WORKFLOW: dict[str, str] = {
69 "streamcast": "podcast",
70 "streampersona": "slide",
71 "streamchat": "chat",
72 "streamshort": "short",
73 "streammovie": "movie",
74 "streamanimate": "story",
75 "streamlecture": "lecture",
76 "streamdub": "dubbing",
77 "streamedit": "editing",
78}
80# Ensure allocator knows about all StreamWise apps (catch drift early).
81assert set(APP_TO_WORKFLOW.keys()) == set(STREAMWISE_APPS), (
82 f"APP_TO_WORKFLOW keys {set(APP_TO_WORKFLOW.keys())} != STREAMWISE_APPS {set(STREAMWISE_APPS)}"
83)
86@dataclass
87class DeploymentSpec:
88 """A single container deployment specification."""
89 container_name: str
90 cpu: int
91 memory_gib: int
92 ephemeral_storage_gib: int
93 gpu: int
94 gpu_type: Optional[str]
95 mig_profile: Optional[str]
98@dataclass
99class DeploymentPlan:
100 """Complete deployment plan produced by the auto-allocator."""
101 specs: list[DeploymentSpec]
102 result: Result
103 workflow_name: str
104 gpu_budget: dict[str, int]
107def _get_data_dir() -> str:
108 """Get the path to the simulator data directory."""
109 default_path = os.path.join(_REPO_ROOT, "simulator", "data")
110 return os.getenv("SIMULATOR_DATA_DIR", default_path)
113# Reverse mapping from pod gpu_type string to GPUType enum
114_POD_STR_TO_GPU_TYPE: dict[str, GPUType] = {v: k for k, v in GPU_TYPE_TO_POD_STR.items()}
117def _calc_actual_gpus_per_type(specs: list['DeploymentSpec']) -> dict[GPUType, int]:
118 """Calculate actual GPUs needed per GPUType from deployment specs."""
119 result: dict[GPUType, int] = {}
120 for spec in specs:
121 if spec.mig_profile:
122 continue
123 gpu_type = _POD_STR_TO_GPU_TYPE.get(spec.gpu_type or "")
124 if gpu_type is not None:
125 result[gpu_type] = result.get(gpu_type, 0) + spec.gpu
126 return result
129def _trim_specs_for_type(
130 specs: list['DeploymentSpec'], gpu_type_str: str, excess: int
131) -> list['DeploymentSpec']:
132 """
133 Remove replicas from specs to reduce GPU usage on a specific type by `excess` GPUs.
135 Prefers removing replicas of the most-replicated scalable container (typically
136 realesrgan/upscaler) to minimize impact on pipeline throughput.
137 """
138 # Count replicas per container on this GPU type (only scalable ones)
139 from collections import Counter
140 type_counts: Counter[str] = Counter()
141 for spec in specs:
142 if spec.gpu_type == gpu_type_str and spec.gpu > 0 and spec.container_name not in COLOCATED_CONTAINERS:
143 type_counts[spec.container_name] += 1
145 # Prefer trimming containers with most replicas (least impact per removal)
146 trimmed = 0
147 result_specs = list(specs)
148 for container_name, _count in type_counts.most_common():
149 if trimmed >= excess:
150 break
151 # Remove replicas from the end of the list
152 for i in range(len(result_specs) - 1, -1, -1):
153 if trimmed >= excess:
154 break
155 spec = result_specs[i]
156 if (spec.container_name == container_name
157 and spec.gpu_type == gpu_type_str
158 and spec.gpu > 0):
159 trimmed += spec.gpu
160 result_specs.pop(i)
161 return result_specs
164def get_available_workflows() -> list[str]:
165 """Return list of available workflow names for the UI."""
166 return list(APP_TO_WORKFLOW.keys())
169def get_available_gpu_types() -> list[str]:
170 """Return list of available GPU type strings for the UI."""
171 return [gpu_type.value for gpu_type in GPUType]
174def run_allocator(
175 gpu_budget: dict[str, int],
176 workflow_name: str,
177) -> DeploymentPlan:
178 """
179 Run the greedy model allocator and return a deployment plan.
181 Args:
182 gpu_budget: GPU counts keyed by GPU type string (e.g., {"A100": 8, "H100": 0}).
183 workflow_name: StreamWise app name (e.g., "streamcast").
185 Returns:
186 DeploymentPlan with concrete container deployment specs.
188 Raises:
189 ValueError: If workflow_name or GPU types are invalid.
190 """
191 # Validate workflow
192 workflow_key = APP_TO_WORKFLOW.get(workflow_name)
193 if workflow_key is None:
194 raise ValueError(
195 f"Unknown workflow '{workflow_name}'. "
196 f"Available: {list(APP_TO_WORKFLOW.keys())}")
198 workflow = WORKFLOWS[workflow_key]
200 # Parse GPU budget into GPUType enum
201 num_gpus: dict[GPUType, int] = {}
202 for gpu_str, count in gpu_budget.items():
203 try:
204 gpu_type = GPUType(gpu_str)
205 except ValueError:
206 raise ValueError(
207 f"Unknown GPU type '{gpu_str}'. "
208 f"Available: {[g.value for g in GPUType]}")
209 if count > 0:
210 num_gpus[gpu_type] = count
212 if not num_gpus or sum(num_gpus.values()) < 8:
213 raise ValueError("Total GPU budget must be at least 8 GPUs.")
215 # The allocator requires GPU counts to be multiples of NUM_GPUS_PER_SERVER (8).
216 # Round up for the allocator, then trim back to the real budget afterward.
217 import math
218 from constants import NUM_GPUS_PER_SERVER
219 allocator_gpus: dict[GPUType, int] = {}
220 for gpu_type, count in num_gpus.items():
221 server_size = NUM_GPUS_PER_SERVER[gpu_type]
222 allocator_gpus[gpu_type] = math.ceil(count / server_size) * server_size
224 # Load latency data and run allocator
225 data_dir = _get_data_dir()
226 latency_data = load_latency_data(data_dir=data_dir)
228 allocator = AutoModelAllocator(
229 workflow=workflow,
230 latency_data=latency_data,
231 policy=STREAMWISE_POLICY,
232 )
234 result = allocator.allocate(num_gpus=allocator_gpus, verbose=False)
236 # Convert result to deployment specs
237 specs = result_to_deployment_specs(result)
239 # Trim deployment specs back to the user's actual budget.
240 # Also handles MIG-unavailable overflow (e.g., OTHERS allocates 1 GPU
241 # but kokoro+yolo each need a full GPU = 2).
242 actual_per_type = _calc_actual_gpus_per_type(specs)
243 for gpu_type, budget_count in num_gpus.items():
244 actual = actual_per_type.get(gpu_type, 0)
245 if actual <= budget_count:
246 continue
247 excess = actual - budget_count
248 gpu_type_str = GPU_TYPE_TO_POD_STR[gpu_type]
249 specs = _trim_specs_for_type(specs, gpu_type_str, excess)
251 return DeploymentPlan(
252 specs=specs,
253 result=result,
254 workflow_name=workflow_name,
255 gpu_budget=gpu_budget,
256 )
259def result_to_deployment_specs(result: Result) -> list[DeploymentSpec]:
260 """
261 Convert an allocator Result into a list of DeploymentSpec objects.
263 Each ModelAllocation with replicas > 0 is mapped to one or more container deployments.
264 When MIG is unavailable, containers that would normally use MIG slices get 1 full GPU instead.
265 """
266 specs: list[DeploymentSpec] = []
268 for gpu_type, model_dict in result.models.items():
269 gpu_type_str = GPU_TYPE_TO_POD_STR[gpu_type]
271 for model, allocations in model_dict.items():
272 containers = MODEL_TO_CONTAINERS.get(model, [])
273 if not containers:
274 continue
276 for allocation in allocations:
277 if allocation.replicas <= 0:
278 continue
280 for container_name in containers:
281 resources = CONTAINER_RESOURCES.get(
282 container_name,
283 ContainerResourceSpec(cpu=4, memory_gib=16, ephemeral_storage_gib=16, gpu=0),
284 )
285 cpu = resources.cpu
286 memory_gib = resources.memory_gib
287 ephemeral_storage_gib = resources.ephemeral_storage_gib
289 mig_profile: Optional[str] = None
290 # Co-locate VAE only when disaggregation is disabled
291 # TODO: make disaggregation a configuration exposed to the users
292 is_colocated = (
293 container_name in COLOCATED_CONTAINERS
294 and not STREAMWISE_POLICY.disaggregation.get(Model.HF, False)
295 )
296 if is_colocated:
297 gpu_count = 0
298 elif MIG_AVAILABLE and container_name in MIG_CONTAINERS:
299 mig_profile = MIG_CONTAINERS[container_name]
300 gpu_count = 1
301 elif container_name in MIG_CONTAINERS:
302 # MIG not available: use 1 full GPU instead of a MIG slice
303 gpu_count = 1
304 else:
305 gpu_count = allocation.devices
307 for _ in range(allocation.replicas):
308 specs.append(DeploymentSpec(
309 container_name=container_name,
310 cpu=cpu,
311 memory_gib=memory_gib,
312 ephemeral_storage_gib=ephemeral_storage_gib,
313 gpu=gpu_count,
314 gpu_type=gpu_type_str,
315 mig_profile=mig_profile,
316 ))
318 return specs
321def deployment_plan_to_json(plan: DeploymentPlan) -> dict:
322 """Serialize a DeploymentPlan to a JSON-friendly dict."""
323 # Calculate actual GPUs used by the deployment specs (may differ from allocator
324 # when MIG is unavailable and services fall back to full GPUs).
325 actual_gpus: dict[str, int] = {}
326 for spec in plan.specs:
327 if spec.mig_profile:
328 continue # MIG slices don't count against full GPU budget
329 gpu_type_key = spec.gpu_type or "unknown"
330 actual_gpus[gpu_type_key] = actual_gpus.get(gpu_type_key, 0) + spec.gpu
332 total_budget = sum(plan.gpu_budget.values())
333 total_actual = sum(actual_gpus.values())
334 budget_exceeded = total_actual > total_budget
336 warnings: list[str] = []
337 if budget_exceeded:
338 mig_hint = (
339 "Enable MIG to fit lightweight services (kokoro, yolo, realesrgan) "
340 "on shared GPU slices."
341 ) if not MIG_AVAILABLE else ""
342 warnings.append(
343 f"Deployment requires {total_actual} full GPUs but budget is "
344 f"{total_budget}. {mig_hint}"
345 )
347 return {
348 "workflow_name": plan.workflow_name,
349 "gpu_budget": plan.gpu_budget,
350 "metrics": {
351 "total_time_s": round(plan.result.total_time_s, 2),
352 "ttff_s": round(plan.result.ttff_s, 2),
353 "cost": round(plan.result.cost, 4),
354 "gpus_used": {
355 gpu_type.value: count
356 for gpu_type, count in plan.result.gpus_used.items()
357 },
358 "actual_gpus_needed": actual_gpus,
359 "budget_exceeded": budget_exceeded,
360 },
361 "warnings": warnings,
362 "mig_available": MIG_AVAILABLE,
363 "specs": [
364 {
365 "container_name": spec.container_name,
366 "cpu": spec.cpu,
367 "memory_gib": spec.memory_gib,
368 "ephemeral_storage_gib": spec.ephemeral_storage_gib,
369 "gpu": spec.gpu,
370 "gpu_type": spec.gpu_type,
371 "mig_profile": spec.mig_profile,
372 }
373 for spec in plan.specs
374 ],
375 }