Coverage for tests/streamwise/test_allocator_bridge.py: 100%
100 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Tests for streamwise/allocator_bridge.py.
4Covers:
5- Model-to-container name mapping.
6- Result to deployment specs conversion.
7- run_allocator end-to-end (with real latency data).
8- Error handling for invalid inputs.
9"""
11from __future__ import annotations
13import sys
14import os
16import pytest
18# Add current path and simulator/ permanently so lazy imports
19# (e.g. GreedyAllocator via auto_model_allocator) resolve at test time.
20sys.path.append(os.getcwd())
21sys.path[:0] = [os.path.join(os.getcwd(), "simulator")]
23from tests.test_utils import temp_sys_path
25with temp_sys_path("streamwise", "simulator"):
26 from allocator_bridge import (
27 MODEL_TO_CONTAINERS,
28 CONTAINER_RESOURCES,
29 GPU_TYPE_TO_POD_STR,
30 APP_TO_WORKFLOW,
31 DeploymentSpec,
32 DeploymentPlan,
33 get_available_workflows,
34 get_available_gpu_types,
35 result_to_deployment_specs,
36 deployment_plan_to_json,
37 run_allocator,
38 )
39 from sim_types import GPUType, Model, Result
40 from models import (
41 GemmaModelAllocation,
42 FluxModelAllocation,
43 HFModelAllocation,
44 HFVAEModelAllocation,
45 FTModelAllocation,
46 OthersModelAllocation,
47 UpscalerModelAllocation,
48 )
51# ---------------------------------------------------------------------------
52# Mapping correctness
53# ---------------------------------------------------------------------------
55def test_model_to_containers_covers_all_models() -> None:
56 """Every Model enum value must have a mapping entry."""
57 for model in Model:
58 assert model in MODEL_TO_CONTAINERS, f"Missing mapping for {model}"
61def test_container_resources_covers_all_mapped_containers() -> None:
62 """Every container referenced in MODEL_TO_CONTAINERS must have resource defaults."""
63 for model, containers in MODEL_TO_CONTAINERS.items():
64 for container in containers:
65 assert container in CONTAINER_RESOURCES, (
66 f"Missing CONTAINER_RESOURCES for '{container}' (from {model})")
69def test_gpu_type_to_pod_str_covers_all_gpu_types() -> None:
70 """Every GPUType enum value must have a pod string mapping."""
71 for gpu_type in GPUType:
72 assert gpu_type in GPU_TYPE_TO_POD_STR
75def test_app_to_workflow_has_expected_entries() -> None:
76 """Key StreamWise apps should map to workflows."""
77 assert "streamcast" in APP_TO_WORKFLOW
78 assert "streampersona" in APP_TO_WORKFLOW
79 assert "streamchat" in APP_TO_WORKFLOW
82# ---------------------------------------------------------------------------
83# Utility functions
84# ---------------------------------------------------------------------------
86def test_get_available_workflows() -> None:
87 workflows = get_available_workflows()
88 assert isinstance(workflows, list)
89 assert "streamcast" in workflows
90 assert len(workflows) >= 5
93def test_get_available_gpu_types() -> None:
94 gpu_types = get_available_gpu_types()
95 assert isinstance(gpu_types, list)
96 assert "A100" in gpu_types
97 assert "H100" in gpu_types
100# ---------------------------------------------------------------------------
101# result_to_deployment_specs
102# ---------------------------------------------------------------------------
104def test_result_to_deployment_specs_basic() -> None:
105 """A simple result with one active allocation maps to the right container."""
106 models = {
107 GPUType.A100: {
108 Model.GEMMA: [GemmaModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1)],
109 Model.FLUX: [FluxModelAllocation(gpu_type=GPUType.A100, devices=2, replicas=1)],
110 Model.HF: [HFModelAllocation(gpu_type=GPUType.A100, devices=2, replicas=2)],
111 Model.HF_VAE: [HFVAEModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1)],
112 Model.FT: [FTModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=0)],
113 Model.FT_VAE: [],
114 Model.UPSCALER: [UpscalerModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=0)],
115 Model.OTHERS: [OthersModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1)],
116 }
117 }
118 result = Result(
119 total_time_s=100.0,
120 ttff_s=10.0,
121 cost=1.0,
122 gpus_used={GPUType.A100: 8},
123 gpus_total={GPUType.A100: 8},
124 models=models,
125 )
127 specs = result_to_deployment_specs(result)
128 assert isinstance(specs, list)
129 assert len(specs) > 0
131 container_names = [s.container_name for s in specs]
132 assert "gemma" in container_names
133 assert "flux" in container_names
134 assert "hunyuanframepackf1" in container_names # HF model
135 assert "hunyuanframepackvae" in container_names # HF_VAE model
137 # OTHERS maps to kokoro + yolo
138 assert "kokoro" in container_names
139 assert "yolo" in container_names
141 # Check GPU type mapping
142 gemma_spec = next(s for s in specs if s.container_name == "gemma")
143 assert gemma_spec.gpu_type == "a100"
144 assert gemma_spec.gpu == 1
146 # Without MIG, kokoro gets no mig_profile (full GPU)
147 kokoro_spec = next(s for s in specs if s.container_name == "kokoro")
148 assert kokoro_spec.mig_profile is None
149 assert kokoro_spec.gpu == 1
151 # With disaggregation=True for HF, VAE runs on its own GPU
152 vae_spec = next(s for s in specs if s.container_name == "hunyuanframepackvae")
153 assert vae_spec.gpu == 1
156def test_result_to_deployment_specs_skips_zero_replicas() -> None:
157 """Allocations with zero replicas should not produce deployment specs."""
158 models = {
159 GPUType.A100: {
160 Model.GEMMA: [GemmaModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=0)],
161 Model.FLUX: [FluxModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=0)],
162 Model.HF: [HFModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=0)],
163 Model.HF_VAE: [],
164 Model.FT: [],
165 Model.FT_VAE: [],
166 Model.UPSCALER: [],
167 Model.OTHERS: [],
168 }
169 }
170 result = Result(
171 total_time_s=0.0,
172 ttff_s=0.0,
173 cost=0.0,
174 gpus_used={GPUType.A100: 0},
175 gpus_total={GPUType.A100: 8},
176 models=models,
177 )
178 specs = result_to_deployment_specs(result)
179 assert specs == []
182def test_result_to_deployment_specs_multiple_replicas() -> None:
183 """Multiple replicas should produce multiple deployment specs for same container."""
184 models = {
185 GPUType.H100: {
186 Model.GEMMA: [GemmaModelAllocation(gpu_type=GPUType.H100, devices=1, replicas=1)],
187 Model.FLUX: [FluxModelAllocation(gpu_type=GPUType.H100, devices=1, replicas=1)],
188 Model.HF: [HFModelAllocation(gpu_type=GPUType.H100, devices=2, replicas=3)],
189 Model.HF_VAE: [],
190 Model.FT: [],
191 Model.FT_VAE: [],
192 Model.UPSCALER: [],
193 Model.OTHERS: [],
194 }
195 }
196 result = Result(
197 total_time_s=50.0,
198 ttff_s=5.0,
199 cost=0.5,
200 gpus_used={GPUType.H100: 8},
201 gpus_total={GPUType.H100: 16},
202 models=models,
203 )
204 specs = result_to_deployment_specs(result)
205 hf_specs = [s for s in specs if s.container_name == "hunyuanframepackf1"]
206 assert len(hf_specs) == 3 # 3 replicas
207 for spec in hf_specs:
208 assert spec.gpu == 2
209 assert spec.gpu_type == "h100"
212# ---------------------------------------------------------------------------
213# deployment_plan_to_json
214# ---------------------------------------------------------------------------
216def test_deployment_plan_to_json() -> None:
217 """Serialization should produce all expected keys."""
218 result = Result(
219 total_time_s=100.0,
220 ttff_s=10.0,
221 cost=1.5,
222 gpus_used={GPUType.A100: 8},
223 gpus_total={GPUType.A100: 8},
224 models={},
225 )
226 plan = DeploymentPlan(
227 specs=[
228 DeploymentSpec(
229 container_name="gemma", cpu=16, memory_gib=192,
230 ephemeral_storage_gib=64, gpu=2, gpu_type="a100", mig_profile=None)
231 ],
232 result=result,
233 workflow_name="streamcast",
234 gpu_budget={"A100": 8},
235 )
236 data = deployment_plan_to_json(plan)
237 assert data["workflow_name"] == "streamcast"
238 assert data["gpu_budget"] == {"A100": 8}
239 assert data["metrics"]["total_time_s"] == 100.0
240 assert data["metrics"]["ttff_s"] == 10.0
241 assert len(data["specs"]) == 1
242 assert data["specs"][0]["container_name"] == "gemma"
245# ---------------------------------------------------------------------------
246# run_allocator (integration with real data)
247# ---------------------------------------------------------------------------
249def test_run_allocator_streamcast_8_a100() -> None:
250 """Run allocator for StreamCast with 8 A100s — should produce a valid plan."""
251 plan = run_allocator(
252 gpu_budget={"A100": 8},
253 workflow_name="streamcast",
254 )
255 assert isinstance(plan, DeploymentPlan)
256 assert len(plan.specs) > 0
257 assert plan.result.total_time_s > 0
258 assert plan.result.ttff_s > 0
259 assert plan.workflow_name == "streamcast"
262def test_run_allocator_streamchat_8_h100() -> None:
263 """Run allocator for StreamChat with 8 H100s."""
264 plan = run_allocator(
265 gpu_budget={"H100": 8},
266 workflow_name="streamchat",
267 )
268 assert isinstance(plan, DeploymentPlan)
269 assert len(plan.specs) > 0
272def test_run_allocator_invalid_workflow() -> None:
273 """Unknown workflow name raises ValueError."""
274 with pytest.raises(ValueError, match="Unknown workflow"):
275 run_allocator(gpu_budget={"A100": 8}, workflow_name="nonexistent")
278def test_run_allocator_invalid_gpu_type() -> None:
279 """Unknown GPU type raises ValueError."""
280 with pytest.raises(ValueError, match="Unknown GPU type"):
281 run_allocator(gpu_budget={"RTX4090": 8}, workflow_name="streamcast")
284def test_run_allocator_insufficient_gpus() -> None:
285 """Too few GPUs raises ValueError."""
286 with pytest.raises(ValueError, match="at least 8"):
287 run_allocator(gpu_budget={"A100": 4}, workflow_name="streamcast")