Coverage for tests/simulator/test_evaluator.py: 100%
78 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1import sys
2import os
3import pytest
5# Add current path
6sys.path.append(os.getcwd())
8from tests.test_utils import assert_equals_approx
9from tests.test_utils import temp_sys_path
11with temp_sys_path("simulator", "streamwise"):
12 from constants import DEFAULT_WORKFLOW_CONFIG
13 from constants import SECONDS_IN_HOUR
15 from sim_types import GPUType
16 from sim_types import Model
18 from data_loading import load_latency_data
19 from data_loading import load_power_data
21 from evaluator import evaluate_model_allocation
23 from model_provisioner.policies import STREAMWISE_POLICY
25 from models import FluxModelAllocation
26 from models import GemmaModelAllocation
27 from models import HFModelAllocation
28 from models import HFVAEModelAllocation
29 from models import FTModelAllocation
30 from models import UpscalerModelAllocation
31 from models import OthersModelAllocation
33 from utils import to_models_df
36def test_empty() -> None:
37 """No models."""
38 latency_data = load_latency_data("simulator/data/")
39 power_data = load_power_data("simulator/data/")
41 with pytest.raises(AssertionError, match="Expected at least one instance of Model.GEMMA, but found 0"):
42 evaluate_model_allocation(
43 models={},
44 num_gpus={GPUType.A100: 8},
45 workflow=DEFAULT_WORKFLOW_CONFIG,
46 latency_data=latency_data,
47 power_data=power_data,
48 policy=STREAMWISE_POLICY,
49 )
52def test_8A() -> None:
53 """Test with 8 A100 GPUs."""
54 latency_data = load_latency_data("simulator/data/")
55 power_data = load_power_data("simulator/data/")
57 models = {
58 GPUType.A100: {
59 Model.GEMMA: [GemmaModelAllocation(
60 gpu_type=GPUType.A100,
61 devices=1, replicas=1)],
62 Model.FLUX: [FluxModelAllocation(
63 gpu_type=GPUType.A100,
64 devices=1, replicas=1)],
65 Model.HF: [HFModelAllocation(
66 gpu_type=GPUType.A100,
67 devices=1, replicas=1)],
68 Model.HF_VAE: [HFVAEModelAllocation(
69 gpu_type=GPUType.A100,
70 devices=1, replicas=1)],
71 Model.FT: [FTModelAllocation(
72 gpu_type=GPUType.A100,
73 devices=1, replicas=2)],
74 Model.UPSCALER: [UpscalerModelAllocation(
75 gpu_type=GPUType.A100,
76 devices=1, replicas=1)],
77 Model.OTHERS: [OthersModelAllocation(
78 gpu_type=GPUType.A100,
79 devices=1, replicas=1)], # + 1 for Kokoro/YOLO
80 }
81 }
82 result = evaluate_model_allocation(
83 models=models,
84 num_gpus={GPUType.A100: 8},
85 workflow=DEFAULT_WORKFLOW_CONFIG,
86 latency_data=latency_data,
87 power_data=power_data,
88 policy=STREAMWISE_POLICY,
89 )
90 assert to_models_df(models).to_string() == (
91 " Devices Replicas Work #GPUs Time (s) TTFF (s) Energy (kWh) Cost ($)\n"
92 "A100 gemma 1 1 0 1 33.20 5.94 0.38 4.74\n"
93 " flux 1 1 0 1 9.75 9.75 0.29 4.74\n"
94 " hf 1 1 0 1 1491.06 2.97 0.31 4.74\n"
95 " hf_vae 1 1 0 1 342.97 2.00 0.29 4.74\n"
96 " ft 1 2 0 2 11387.49 132.45 2.62 9.48\n"
97 " upscaler 1 1 0 1 2663.40 15.63 0.41 4.74\n"
98 " others 1 1 0 1 25.80 0.60 0.00 4.74\n"
99 "TOTAL 7 8 0 8 15953.67 169.34 4.30 37.92"
100 )
101 assert str(result) == (
102 "Time:15953.68 s TTFF:15353.68 s Cost:$37.93 TTFF*Cost:582431.29 Energy:4.31 kWh GPUS: 8xA100"
103 )
105 assert result.gpus_used == {GPUType.A100: 8}
106 assert result.gpus_total == {GPUType.A100: 8}
107 assert_equals_approx(result.total_time_s, 15953.68)
108 assert_equals_approx(result.ttff_s, 15353.68)
109 assert_equals_approx(result.first_chunk_time, 169.34)
110 assert_equals_approx(result.total_energy / SECONDS_IN_HOUR / 1000, 4.31)
111 assert_equals_approx(result.cost, 37.93)
114def test_16H() -> None:
115 """Test with 16 H200 GPUs."""
116 latency_data = load_latency_data("simulator/data/")
117 power_data = load_power_data("simulator/data/")
119 models = {
120 GPUType.H200: {
121 Model.GEMMA: [GemmaModelAllocation(
122 gpu_type=GPUType.H200,
123 devices=1, replicas=1)],
124 Model.FLUX: [FluxModelAllocation(
125 gpu_type=GPUType.H200,
126 devices=1, replicas=1)],
127 Model.HF: [HFModelAllocation(
128 gpu_type=GPUType.H200,
129 devices=2, replicas=2)],
130 Model.HF_VAE: [HFVAEModelAllocation(
131 gpu_type=GPUType.H200,
132 devices=1, replicas=1)],
133 Model.FT: [FTModelAllocation(
134 gpu_type=GPUType.H200,
135 devices=2, replicas=2)],
136 Model.UPSCALER: [UpscalerModelAllocation(
137 gpu_type=GPUType.H200,
138 devices=1, replicas=2)],
139 Model.OTHERS: [OthersModelAllocation(
140 gpu_type=GPUType.H200,
141 devices=1, replicas=1)], # + 1 for Kokoro/YOLO
142 }
143 }
144 result = evaluate_model_allocation(
145 models=models,
146 num_gpus={GPUType.H200: 16},
147 workflow=DEFAULT_WORKFLOW_CONFIG,
148 latency_data=latency_data,
149 power_data=power_data,
150 policy=STREAMWISE_POLICY,
151 )
152 assert result is not None
153 assert result.gpus_used == {GPUType.H200: 14}
154 assert result.gpus_total == {GPUType.H200: 16}
155 assert_equals_approx(result.total_time_s, 4062.7)
156 assert_equals_approx(result.ttff_s, 3462.7)
157 assert_equals_approx(result.first_chunk_time, 51.68)
158 assert_equals_approx(result.total_energy / SECONDS_IN_HOUR / 1000, 2.87)
159 assert_equals_approx(result.cost, 66.67)
162def test_cost_optimal() -> None:
163 """
164 Test with 256xA100 + 64xH200 GPUs.
165 This is the cost-optimal configuration for the default workflow used in the paper.
166 """
167 latency_data = load_latency_data("simulator/data/")
168 power_data = load_power_data("simulator/data/")
170 models = {
171 GPUType.A100: {
172 Model.GEMMA: [GemmaModelAllocation(
173 gpu_type=GPUType.A100,
174 devices=8, replicas=1)],
175 Model.FLUX: [FluxModelAllocation(
176 gpu_type=GPUType.A100,
177 devices=16, replicas=1)],
178 Model.HF: [
179 HFModelAllocation(
180 gpu_type=GPUType.A100,
181 devices=2, replicas=6),
182 HFModelAllocation(
183 gpu_type=GPUType.A100,
184 devices=1, replicas=29),
185 ],
186 Model.HF_VAE: [HFVAEModelAllocation(
187 gpu_type=GPUType.A100,
188 devices=1, replicas=20)],
189 Model.FT: [
190 FTModelAllocation(
191 gpu_type=GPUType.A100,
192 devices=4, replicas=18),
193 FTModelAllocation(
194 gpu_type=GPUType.A100,
195 devices=2, replicas=12),
196 ],
197 Model.UPSCALER: [
198 UpscalerModelAllocation(
199 gpu_type=GPUType.A100,
200 devices=4, replicas=8),
201 UpscalerModelAllocation(
202 gpu_type=GPUType.A100,
203 devices=8, replicas=3),
204 UpscalerModelAllocation(
205 gpu_type=GPUType.A100,
206 devices=2, replicas=6),
207 UpscalerModelAllocation(
208 gpu_type=GPUType.A100,
209 devices=1, replicas=6),
210 ],
211 Model.OTHERS: [OthersModelAllocation(
212 gpu_type=GPUType.A100,
213 devices=1, replicas=1)],
214 },
215 GPUType.H200: {
216 Model.HF: [HFModelAllocation(
217 gpu_type=GPUType.H200,
218 devices=2, replicas=4)],
219 Model.HF_VAE: [HFVAEModelAllocation(
220 gpu_type=GPUType.H200,
221 devices=1, replicas=4)],
222 Model.FT: [
223 FTModelAllocation(
224 gpu_type=GPUType.H200,
225 devices=2, replicas=13),
226 FTModelAllocation(
227 gpu_type=GPUType.H200,
228 devices=24, replicas=1),
229 ],
230 Model.UPSCALER: [UpscalerModelAllocation(
231 gpu_type=GPUType.H200,
232 devices=2, replicas=1)],
233 }
234 }
235 result = evaluate_model_allocation(
236 models=models,
237 num_gpus={GPUType.A100: 256, GPUType.H200: 64},
238 workflow=DEFAULT_WORKFLOW_CONFIG,
239 latency_data=latency_data,
240 power_data=power_data,
241 policy=STREAMWISE_POLICY,
242 )
244 assert to_models_df(models).to_string() == (
245 " Devices Replicas Work #GPUs Time (s) TTFF (s) Energy (kWh) Cost ($)\n"
246 "A100 gemma 8 1 0 8 8.57 1.48 0.24 0.72\n"
247 " flux 16 1 0 16 0.95 0.95 0.09 1.45\n"
248 " hf 2 6 0 12 29.56 2.03 0.08 1.09\n"
249 " hf 1 29 0 29 29.30 2.97 0.17 2.62\n"
250 " hf_vae 1 20 0 20 11.75 2.00 0.11 1.81\n"
251 " ft 4 18 0 72 192.62 58.82 1.61 6.51\n"
252 " ft 2 12 0 24 192.61 79.76 0.54 2.17\n"
253 " upscaler 4 8 0 32 34.75 3.97 0.23 2.90\n"
254 " upscaler 8 3 0 24 34.75 2.02 0.17 2.17\n"
255 " upscaler 2 6 0 12 34.76 7.86 0.08 1.09\n"
256 " upscaler 1 6 0 6 34.77 15.63 0.04 0.54\n"
257 " others 1 1 0 1 25.80 0.60 0.00 0.09\n"
258 "H200 hf 2 4 0 8 29.49 0.90 0.06 2.85\n"
259 " hf_vae 1 4 0 4 11.75 0.87 0.03 1.43\n"
260 " ft 2 13 0 26 193.00 34.87 0.81 9.28\n"
261 " ft 24 1 0 24 191.93 14.78 0.72 8.56\n"
262 " upscaler 2 1 0 2 34.75 3.87 0.02 0.71\n"
263 "TOTAL 81 134 0 320 304.40 21.60 5.00 45.99"
264 )
265 assert str(result) == (
266 "Time:304.40 s TTFF:21.60 s Cost:$46.00 TTFF*Cost:993.61 "
267 "Energy:5.01 kWh GPUS: 256xA100+64xH200"
268 )
270 assert result.gpus_used == {
271 GPUType.A100: 256,
272 GPUType.H200: 64,
273 }
274 assert result.gpus_total == {
275 GPUType.A100: 256,
276 GPUType.H200: 64,
277 }
278 assert_equals_approx(result.total_time_s, 304.40)
279 assert_equals_approx(result.ttff_s, 21.60)
280 assert_equals_approx(result.first_chunk_time, 21.60)
281 assert_equals_approx(result.total_energy / SECONDS_IN_HOUR / 1000, 5.01)
282 assert_equals_approx(result.cost, 46.00)
284 assert models[GPUType.A100][Model.OTHERS][0].devices == 1
285 assert models[GPUType.A100][Model.OTHERS][0].replicas == 1
286 assert models[GPUType.A100][Model.OTHERS][0].time_first == 0.60
288 assert models[GPUType.H200][Model.FT][1].devices == 24
289 assert models[GPUType.H200][Model.FT][1].replicas == 1
290 assert_equals_approx(models[GPUType.H200][Model.FT][1].time, 191.93)
291 assert_equals_approx(models[GPUType.H200][Model.FT][1].time_first, 14.78)
292 assert_equals_approx(models[GPUType.H200][Model.FT][1].energy / SECONDS_IN_HOUR / 1000, 0.72)
293 assert_equals_approx(models[GPUType.H200][Model.FT][1].cost, 8.56)