Coverage for tests/simulator/test_models.py: 100%

1""" 

2Direct unit tests for simulator/models.py. 

3Tests the model allocation factory, helper functions, and per-model 

4calculate_time / calculate_time_first / calculate_energy methods. 

5""" 

6 

7from __future__ import annotations 

8 

9import sys 

10import os 

11import pytest 

12from unittest.mock import patch 

13 

14# Add current path 

15sys.path.append(os.getcwd()) 

16 

17from tests.test_utils import temp_sys_path 

18 

19with temp_sys_path("simulator", "streamwise"): 

20 from sim_types import GPUType 

21 from sim_types import Model 

22 from sim_types import ModelAllocation 

23 from sim_types import QualityLevel 

24 from sim_types import LatencyData 

25 from sim_types import PowerData 

26 

27 from constants import DEFAULT_WORKFLOW_CONFIG 

28 

29 from data_loading import load_latency_data 

30 from data_loading import load_power_data 

31 

32 from model_provisioner.policies import STREAMWISE_POLICY 

33 from model_provisioner.policies import NAIVE_POLICY 

34 

35 from models import get_model_allocation 

36 from models import _calculate_total_time 

37 from models import assert_pixel_config 

38 from models import _MODEL_ALLOCATION_REGISTRY 

39 from models import GemmaModelAllocation 

40 from models import FluxModelAllocation 

41 from models import HFModelAllocation 

42 from models import HFVAEModelAllocation 

43 from models import FTModelAllocation 

44 from models import FTVAEModelAllocation 

45 from models import UpscalerModelAllocation 

46 from models import OthersModelAllocation 

47 

48 

49# --------------------------------------------------------------------------- 

50# Helpers 

51# --------------------------------------------------------------------------- 

52 

53def _make_latency_data() -> LatencyData: 

54 """Return minimal LatencyData built from the real data files.""" 

55 return load_latency_data("simulator/data/") 

56 

57 

58def _make_power_data() -> PowerData: 

59 """Return minimal PowerData built from the real data files.""" 

60 return load_power_data("simulator/data/") 

61 

62 

63# --------------------------------------------------------------------------- 

64# get_model_allocation factory 

65# --------------------------------------------------------------------------- 

66 

67def test_get_model_allocation_returns_correct_types() -> None: 

68 """Factory returns the right ModelAllocation subclass for each Model.""" 

69 expected: list[tuple[Model, type[ModelAllocation]]] = [ 

70 (Model.GEMMA, GemmaModelAllocation), 

71 (Model.FLUX, FluxModelAllocation), 

72 (Model.HF, HFModelAllocation), 

73 (Model.HF_VAE, HFVAEModelAllocation), 

74 (Model.FT, FTModelAllocation), 

75 (Model.FT_VAE, FTVAEModelAllocation), 

76 (Model.UPSCALER, UpscalerModelAllocation), 

77 (Model.OTHERS, OthersModelAllocation), 

78 ] 

79 for model, cls in expected: 

80 alloc = get_model_allocation( 

81 model=model, 

82 gpu_type=GPUType.A100, 

83 devices=1, 

84 replicas=1, 

85 ) 

86 assert isinstance(alloc, cls), f"Expected {cls.__name__} for {model}" 

87 assert alloc.gpu_type == GPUType.A100 

88 assert alloc.devices == 1 

89 assert alloc.replicas == 1 

90 

91 

92def test_get_model_allocation_zero_replicas() -> None: 

93 """Factory creates an allocation with zero replicas (disabled) by default.""" 

94 alloc = get_model_allocation( 

95 model=Model.FLUX, 

96 gpu_type=GPUType.H100, 

97 ) 

98 assert alloc.replicas == 0 

99 assert alloc.get_num_gpus() == 0 

100 

101 

102def test_get_model_allocation_unknown_model_raises() -> None: 

103 """Factory raises ValueError for an unregistered Model value.""" 

104 # Temporarily remove GEMMA from the registry using patch.dict 

105 with patch.dict(_MODEL_ALLOCATION_REGISTRY, {}, clear=False) as patched: 

106 patched.pop(Model.GEMMA, None) 

107 with pytest.raises(ValueError, match="No ModelAllocation for model"): 

108 get_model_allocation( 

109 model=Model.GEMMA, 

110 gpu_type=GPUType.A100, 

111 ) 

112 

113 

114# --------------------------------------------------------------------------- 

115# _calculate_total_time 

116# --------------------------------------------------------------------------- 

117 

118def test_calculate_total_time_zero_replicas() -> None: 

119 """Zero replicas → zero time.""" 

120 assert _calculate_total_time(100.0, 0, 1.0) == 0.0 

121 

122 

123def test_calculate_total_time_negative_replicas() -> None: 

124 """Negative replicas → zero time.""" 

125 assert _calculate_total_time(100.0, -1, 1.0) == 0.0 

126 

127 

128def test_calculate_total_time_single_replica() -> None: 

129 """Single replica: total_work * time_per_work, clamped to time_per_work.""" 

130 # 10 work / 1 replica * 5.0 = 50.0 

131 assert _calculate_total_time(10.0, 1, 5.0) == 50.0 

132 

133 

134def test_calculate_total_time_floor_at_single_work_unit() -> None: 

135 """Time cannot be less than time_per_work (single unit floor).""" 

136 # 1 work, 10 replicas → 1/10 * 5.0 = 0.5, floored to 5.0 

137 assert _calculate_total_time(1.0, 10, 5.0) == 5.0 

138 

139 

140def test_calculate_total_time_many_replicas() -> None: 

141 """Many replicas reduce time proportionally.""" 

142 # 20/4 * 2.0 = 10.0; 10.0 > time_per_work(2.0) → 10.0 

143 assert _calculate_total_time(20.0, 4, 2.0) == 10.0 

144 

145 

146# --------------------------------------------------------------------------- 

147# assert_pixel_config 

148# --------------------------------------------------------------------------- 

149 

150def test_assert_pixel_config() -> None: 

151 """assert_pixel_config passes for valid config and raises for invalid.""" 

152 assert_pixel_config(DEFAULT_WORKFLOW_CONFIG) 

153 

154 # Patching MEDIUM > HIGH violates the ordering constraint → AssertionError. 

155 with patch.dict("sim_types.RESOLUTION_PIXELS", 

156 {QualityLevel.MEDIUM: 1000, QualityLevel.HIGH: 500}): 

157 with pytest.raises(AssertionError): 

158 assert_pixel_config(DEFAULT_WORKFLOW_CONFIG) 

159 

160 

161# --------------------------------------------------------------------------- 

162# Zero-GPU paths (replicas=0 → time=0, energy=0) 

163# All models must return 0.0 when no GPUs are allocated. 

164# --------------------------------------------------------------------------- 

165 

166def test_zero_gpu_all_models() -> None: 

167 latency_data = _make_latency_data() 

168 power_data = _make_power_data() 

169 

170 zero_allocs = [ 

171 GemmaModelAllocation(gpu_type=GPUType.A100), # replicas=0 by default 

172 FluxModelAllocation(gpu_type=GPUType.A100), 

173 HFModelAllocation(gpu_type=GPUType.A100), 

174 HFVAEModelAllocation(gpu_type=GPUType.A100), 

175 FTModelAllocation(gpu_type=GPUType.A100), 

176 FTVAEModelAllocation(gpu_type=GPUType.A100), 

177 UpscalerModelAllocation(gpu_type=GPUType.A100), 

178 OthersModelAllocation(gpu_type=GPUType.A100), 

179 ] 

180 workflow = DEFAULT_WORKFLOW_CONFIG 

181 for alloc in zero_allocs: 

182 assert alloc.get_num_gpus() == 0 

183 t = alloc.calculate_time(STREAMWISE_POLICY, workflow, latency_data) 

184 assert t == 0.0, f"{type(alloc).__name__}.calculate_time with 0 GPUs should be 0" 

185 tf = alloc.calculate_time_first(STREAMWISE_POLICY, workflow, latency_data) 

186 assert tf == 0.0, f"{type(alloc).__name__}.calculate_time_first with 0 GPUs should be 0" 

187 e = alloc.calculate_energy(workflow, power_data, total_time_s=1000.0) 

188 assert e == 0.0, f"{type(alloc).__name__}.calculate_energy with 0 GPUs should be 0" 

189 

190 

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

192# GemmaModelAllocation 

193# --------------------------------------------------------------------------- 

194 

195def test_gemma_calculate_time_single_gpu() -> None: 

196 """Gemma time with 1 GPU should be positive and scale with tokens.""" 

197 latency_data = _make_latency_data() 

198 alloc = GemmaModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

199 t = alloc.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

200 assert t > 0.0 

201 assert alloc.time == t 

202 

203 

204def test_gemma_calculate_time_first_single_gpu() -> None: 

205 """Gemma TTFF < total time.""" 

206 latency_data = _make_latency_data() 

207 alloc = GemmaModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

208 alloc.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

209 tf = alloc.calculate_time_first(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

210 assert 0.0 < tf <= alloc.time 

211 

212 

213def test_gemma_calculate_energy() -> None: 

214 """Gemma energy > 0 when power data is provided.""" 

215 latency_data = _make_latency_data() 

216 power_data = _make_power_data() 

217 alloc = GemmaModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

218 alloc.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

219 alloc.calculate_time_first(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

220 e = alloc.calculate_energy( 

221 DEFAULT_WORKFLOW_CONFIG, 

222 power_data=power_data, 

223 total_time_s=alloc.time * 2, 

224 ) 

225 assert e > 0.0 

226 

227 

228def test_gemma_calculate_energy_no_power_data() -> None: 

229 """Energy is 0 when no power data is provided.""" 

230 latency_data = _make_latency_data() 

231 alloc = GemmaModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

232 alloc.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

233 e = alloc.calculate_energy(DEFAULT_WORKFLOW_CONFIG, power_data=None) 

234 assert e == 0.0 

235 

236 

237def test_gemma_get_max_replicas() -> None: 

238 """Gemma max replicas equals model_work.get(GEMMA, 1).""" 

239 alloc = GemmaModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

240 max_r = alloc.get_max_replicas(DEFAULT_WORKFLOW_CONFIG) 

241 assert max_r == DEFAULT_WORKFLOW_CONFIG.model_work.get(Model.GEMMA, 1) 

242 

243 

244# --------------------------------------------------------------------------- 

245# FluxModelAllocation 

246# --------------------------------------------------------------------------- 

247 

248def test_flux_calculate_time_single_gpu() -> None: 

249 """Flux time with 1 GPU should be positive.""" 

250 latency_data = _make_latency_data() 

251 alloc = FluxModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

252 t = alloc.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

253 assert t > 0.0 

254 

255 

256def test_flux_time_equals_time_first() -> None: 

257 """For Flux, time and time_first are the same (single-scene model).""" 

258 latency_data = _make_latency_data() 

259 alloc = FluxModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

260 t = alloc.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

261 tf = alloc.calculate_time_first(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

262 assert abs(t - tf) < 1e-6 

263 

264 

265def test_flux_multi_device_faster() -> None: 

266 """Flux with more GPU devices should be faster (or equal).""" 

267 latency_data = _make_latency_data() 

268 alloc1 = FluxModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

269 alloc2 = FluxModelAllocation(gpu_type=GPUType.A100, devices=2, replicas=1) 

270 t1 = alloc1.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

271 t2 = alloc2.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

272 assert t2 <= t1 

273 

274 

275# --------------------------------------------------------------------------- 

276# HFModelAllocation 

277# --------------------------------------------------------------------------- 

278 

279def test_hf_calculate_time_single_gpu() -> None: 

280 """HF time with 1 GPU should be positive.""" 

281 latency_data = _make_latency_data() 

282 alloc = HFModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

283 t = alloc.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

284 assert t > 0.0 

285 

286 

287def test_hf_more_replicas_reduces_time() -> None: 

288 """More HF replicas should reduce total time.""" 

289 latency_data = _make_latency_data() 

290 alloc1 = HFModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

291 alloc4 = HFModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=4) 

292 t1 = alloc1.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

293 t4 = alloc4.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

294 assert t4 < t1 

295 

296 

297def test_hf_get_max_replicas() -> None: 

298 alloc = HFModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

299 max_r = alloc.get_max_replicas(DEFAULT_WORKFLOW_CONFIG) 

300 assert max_r == DEFAULT_WORKFLOW_CONFIG.model_work.get(Model.HF, 1) 

301 

302 

303# --------------------------------------------------------------------------- 

304# FTModelAllocation 

305# --------------------------------------------------------------------------- 

306 

307def test_ft_calculate_time_single_gpu() -> None: 

308 """FT time with 1 GPU should be positive.""" 

309 latency_data = _make_latency_data() 

310 alloc = FTModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

311 t = alloc.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

312 assert t > 0.0 

313 

314 

315def test_ft_time_first_equals_one_subscene() -> None: 

316 """FT time_first corresponds to a single subscene duration.""" 

317 latency_data = _make_latency_data() 

318 alloc = FTModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

319 alloc.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

320 tf = alloc.calculate_time_first(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

321 assert 0.0 < tf <= alloc.time 

322 

323 

324def test_ft_get_max_replicas() -> None: 

325 alloc = FTModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

326 max_r = alloc.get_max_replicas(DEFAULT_WORKFLOW_CONFIG) 

327 assert max_r == DEFAULT_WORKFLOW_CONFIG.model_work.get(Model.FT, 1) 

328 

329 

330# --------------------------------------------------------------------------- 

331# HFVAEModelAllocation (disaggregated) 

332# --------------------------------------------------------------------------- 

333 

334def test_hf_vae_disabled_when_not_disaggregated() -> None: 

335 """HF_VAE with no disaggregation (NAIVE_POLICY) should stay at zero.""" 

336 # NAIVE_POLICY has disaggregation={} → HF not disaggregated 

337 alloc = HFVAEModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=0) 

338 latency_data = _make_latency_data() 

339 t = alloc.calculate_time(NAIVE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

340 assert t == 0.0 

341 

342 

343def test_hf_vae_positive_when_disaggregated() -> None: 

344 """HF_VAE with disaggregation enabled should give positive time.""" 

345 alloc = HFVAEModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

346 latency_data = _make_latency_data() 

347 t = alloc.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

348 assert t > 0.0 

349 

350 

351# --------------------------------------------------------------------------- 

352# FTVAEModelAllocation (disaggregated) 

353# --------------------------------------------------------------------------- 

354 

355def test_ft_vae_disabled_when_not_disaggregated() -> None: 

356 """FT_VAE with no disaggregation should stay at zero.""" 

357 alloc = FTVAEModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=0) 

358 latency_data = _make_latency_data() 

359 t = alloc.calculate_time(NAIVE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

360 assert t == 0.0 

361 

362 

363# --------------------------------------------------------------------------- 

364# UpscalerModelAllocation 

365# --------------------------------------------------------------------------- 

366 

367def test_upscaler_calculate_time() -> None: 

368 """Upscaler time should be positive with 1 replica.""" 

369 latency_data = _make_latency_data() 

370 alloc = UpscalerModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

371 t = alloc.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

372 assert t > 0.0 

373 

374 

375def test_upscaler_time_first() -> None: 

376 """Upscaler time_first should be positive when use_upscaler=True.""" 

377 latency_data = _make_latency_data() 

378 alloc = UpscalerModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

379 tf = alloc.calculate_time_first(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

380 assert tf > 0.0 

381 

382 

383def test_upscaler_get_max_replicas() -> None: 

384 alloc = UpscalerModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

385 max_r = alloc.get_max_replicas(DEFAULT_WORKFLOW_CONFIG) 

386 assert max_r == DEFAULT_WORKFLOW_CONFIG.model_work.get(Model.UPSCALER, 1) 

387 

388 

389# --------------------------------------------------------------------------- 

390# OthersModelAllocation 

391# --------------------------------------------------------------------------- 

392 

393def test_others_calculate_time() -> None: 

394 """Others time is proportional to total_scenes.""" 

395 latency_data = _make_latency_data() 

396 alloc = OthersModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

397 t = alloc.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

398 assert t > 0.0 

399 

400 

401def test_others_time_first_less_than_total() -> None: 

402 """Others time_first equals latency for a single scene.""" 

403 latency_data = _make_latency_data() 

404 alloc = OthersModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

405 alloc.calculate_time(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

406 tf = alloc.calculate_time_first(STREAMWISE_POLICY, DEFAULT_WORKFLOW_CONFIG, latency_data) 

407 assert 0.0 < tf <= alloc.time 

408 

409 

410# --------------------------------------------------------------------------- 

411# ModelAllocation.calculate() convenience method 

412# --------------------------------------------------------------------------- 

413 

414def test_calculate_convenience_method_populates_all_fields() -> None: 

415 """ModelAllocation.calculate() should fill time, time_first, cost, energy.""" 

416 latency_data = _make_latency_data() 

417 power_data = _make_power_data() 

418 alloc = FluxModelAllocation(gpu_type=GPUType.A100, devices=1, replicas=1) 

419 alloc.calculate( 

420 policy=STREAMWISE_POLICY, 

421 workflow=DEFAULT_WORKFLOW_CONFIG, 

422 latency_data=latency_data, 

423 power_data=power_data, 

424 total_time_s=100.0, 

425 ) 

426 assert alloc.time > 0.0 

427 assert alloc.time_first > 0.0 

428 assert alloc.cost > 0.0 

429 assert alloc.energy > 0.0 

430 

431 

432# --------------------------------------------------------------------------- 

433# ModelAllocation.disable() 

434# --------------------------------------------------------------------------- 

435 

436def test_disable_zeroes_allocation() -> None: 

437 """disable() should zero all fields.""" 

438 alloc = FluxModelAllocation(gpu_type=GPUType.A100, devices=2, replicas=2) 

439 alloc.time = 5.0 

440 alloc.time_first = 1.0 

441 alloc.energy = 100.0 

442 alloc.disable() 

443 assert alloc.devices == 0 

444 assert alloc.replicas == 0 

445 assert alloc.time == 0.0 

446 assert alloc.time_first == 0.0 

447 assert alloc.energy == 0.0 

448 assert alloc.get_num_gpus() == 0