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

1""" 

2Test simulator module. 

3""" 

4 

5import sys 

6import os 

7import pytest 

8 

9from dataclasses import replace 

10 

11# Add current path 

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

13 

14from tests.test_utils import temp_sys_path 

15 

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

17 from sim_types import WorkflowConfig 

18 from sim_types import Model 

19 from sim_types import Objective 

20 from sim_types import GPUType 

21 

22 from constants import SECONDS_IN_HOUR 

23 from constants import DEFAULT_WORKFLOW_CONFIG 

24 

25 from data_loading import load_latency_data 

26 from data_loading import load_power_data 

27 

28 from auto_model_allocator import AutoModelAllocator 

29 from model_provisioner.greedy import GreedyAllocator 

30 

31 from model_provisioner.policies import STREAMWISE_POLICY 

32 from model_provisioner.policies import NAIVE_POLICY 

33 

34 

35def test_estimate_total_time() -> None: 

36 """8 A100 + 8 H100.""" 

37 latency_data = load_latency_data("simulator/data/") 

38 

39 allocator = GreedyAllocator( 

40 workflow=DEFAULT_WORKFLOW_CONFIG, 

41 latency_data=latency_data, 

42 ) 

43 result = allocator.allocate( 

44 num_gpus={GPUType.A100: 8, GPUType.H100: 8}, 

45 verbose=True, 

46 ) 

47 assert result.gpus_used[GPUType.A100] == 8 

48 assert result.gpus_used[GPUType.H100] == 8 

49 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

50 # assert 4000 <= result.total_time_s <= 4400 

51 assert 3000 <= result.total_time_s <= 4000 

52 # assert 3500 <= result.ttff_s <= 3800 

53 assert 2500 <= result.ttff_s <= 3000 

54 assert 0.2 <= result.tbf_s <= 0.4 

55 

56 

57def test_naive() -> None: 

58 """Naive policy with 8 A100 + 8 H100.""" 

59 latency_data = load_latency_data("simulator/data/") 

60 allocator = AutoModelAllocator( 

61 workflow=DEFAULT_WORKFLOW_CONFIG, 

62 latency_data=latency_data, 

63 policy=NAIVE_POLICY, 

64 ) 

65 result = allocator.allocate( 

66 num_gpus={GPUType.A100: 8, GPUType.H100: 8}, 

67 verbose=True, 

68 ) 

69 assert result.gpus_used[GPUType.A100] <= 8 

70 assert result.gpus_used[GPUType.H100] <= 8 

71 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

72 

73 video_seconds = DEFAULT_WORKFLOW_CONFIG.total_video_seconds 

74 expected_length_upper = 8 * SECONDS_IN_HOUR # 8 hours 

75 expected_length_lower = 0.5 * SECONDS_IN_HOUR # 30 minutes 

76 total_frames = DEFAULT_WORKFLOW_CONFIG.total_frames[Model.FT] 

77 assert expected_length_lower <= result.total_time_s <= expected_length_upper 

78 assert expected_length_lower - video_seconds <= result.ttff_s <= expected_length_upper - video_seconds 

79 assert expected_length_lower / total_frames <= result.tbf_s <= expected_length_upper / total_frames 

80 

81 

82def test_workflow_config() -> None: 

83 latency_data = load_latency_data("simulator/data/") 

84 

85 workflow_config = WorkflowConfig( 

86 total_scenes=5, 

87 total_video_seconds=30, 

88 num_steps={ 

89 Model.FLUX: 25, 

90 Model.HF: 20, 

91 Model.FT: 20, 

92 }, 

93 hf_frames=[36, 72, 108, 144, 324], 

94 ft_frames=[9, 21, 41, 61, 77], 

95 frames_per_step_idx=2, 

96 total_frames={Model.HF: 20, Model.FT: 20}, 

97 per_subscene_frames={Model.HF: 10, Model.FT: 10}, 

98 total_subscenes=10, 

99 ) 

100 allocator = GreedyAllocator( 

101 workflow=workflow_config, 

102 latency_data=latency_data, 

103 ) 

104 result = allocator.allocate( 

105 num_gpus={GPUType.A100: 8, GPUType.H100: 8}, 

106 verbose=True, 

107 ) 

108 assert result.gpus_used[GPUType.A100] == 8 

109 assert result.gpus_used[GPUType.H100] == 8 

110 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

111 assert 0 < result.tbf_s < 10 

112 

113 # 8 A100 + 0 H100 

114 result = allocator.allocate( 

115 num_gpus={GPUType.A100: 8, GPUType.H100: 0}, 

116 verbose=True, 

117 ) 

118 assert result.gpus_used[GPUType.A100] == 8 

119 assert result.gpus_used.get(GPUType.H100, 0) == 0 

120 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

121 

122 # 0 A100 + 8 H100 

123 result = allocator.allocate( 

124 num_gpus={GPUType.H100: 8}, 

125 verbose=True, 

126 ) 

127 assert result.gpus_used.get(GPUType.A100, 0) == 0 

128 assert result.gpus_used[GPUType.H100] == 8 

129 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

130 

131 # 32 A100 + 128 H100 

132 result = allocator.allocate( 

133 num_gpus={GPUType.A100: 32, GPUType.H100: 128}, 

134 verbose=True, 

135 ) 

136 assert result.gpus_used[GPUType.A100] == 32 

137 # assert result.gpus_used[GPUType.H100] == 128 

138 assert 128 - 8 < result.gpus_used[GPUType.H100] <= 128 # TODO we should try to make it 128 no fragmenting 

139 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

140 

141 # 16 A100 + 8 H100 with verbose 

142 result = allocator.allocate( 

143 num_gpus={GPUType.A100: 16, GPUType.H100: 8}, 

144 verbose=True, 

145 ) 

146 assert result.gpus_used[GPUType.A100] == 16 

147 assert result.gpus_used[GPUType.H100] == 8 

148 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

149 

150 

151def test_estimate_total_time_baseline() -> None: 

152 """Naive baseline with 8 A100 + 8 H100.""" 

153 latency_data = load_latency_data("simulator/data/") 

154 allocator = AutoModelAllocator( 

155 workflow=DEFAULT_WORKFLOW_CONFIG, 

156 latency_data=latency_data, 

157 policy=NAIVE_POLICY, 

158 ) 

159 result = allocator.allocate( 

160 num_gpus={GPUType.A100: 8, GPUType.H100: 8}, 

161 ) 

162 assert 0 < result.gpus_used[GPUType.A100] <= 8 

163 assert 0 < result.gpus_used[GPUType.H100] <= 8 

164 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

165 

166 

167def test_estimate_total_energy() -> None: 

168 """Energy with 8 A100 + 8 H100.""" 

169 latency_data = load_latency_data("simulator/data/") 

170 power_data = load_power_data("simulator/data/") 

171 

172 allocator = GreedyAllocator( 

173 workflow=DEFAULT_WORKFLOW_CONFIG, 

174 latency_data=latency_data, 

175 policy=STREAMWISE_POLICY, 

176 power_data=power_data, 

177 ) 

178 result = allocator.allocate( 

179 num_gpus={GPUType.A100: 8, GPUType.H100: 8}, 

180 ) 

181 assert 0 < result.gpus_used[GPUType.A100] <= 8 

182 assert 0 < result.gpus_used[GPUType.H100] == 8 

183 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

184 assert 0 < result.total_energy 

185 assert 0 < result.cost 

186 

187 

188def test_AvH() -> None: 

189 """H100 should be better than A100.""" 

190 latency_data = load_latency_data("simulator/data/") 

191 

192 # 8 A100 

193 allocator = GreedyAllocator( 

194 workflow=DEFAULT_WORKFLOW_CONFIG, 

195 latency_data=latency_data, 

196 ) 

197 result_8a = allocator.allocate( 

198 num_gpus={GPUType.A100: 8, GPUType.H100: 0}, 

199 ) 

200 assert result_8a.gpus_used[GPUType.A100] == 8 

201 assert result_8a.gpus_used.get(GPUType.H100, 0) == 0 

202 

203 # 8 H100 

204 result_8h = allocator.allocate( 

205 num_gpus={GPUType.H100: 8}, 

206 ) 

207 assert result_8h.gpus_used.get(GPUType.A100, 0) == 0 

208 assert result_8h.gpus_used[GPUType.H100] == 8 

209 

210 # A100 should be worse than H100 

211 assert result_8a.total_time_s > result_8h.total_time_s 

212 assert result_8a.ttff_s > result_8h.ttff_s 

213 assert result_8a.tbf_s > result_8h.tbf_s 

214 

215 

216def test_estimate_total_time_A() -> None: 

217 """More A100s should lead to better performance.""" 

218 latency_data = load_latency_data("simulator/data/") 

219 

220 allocator = GreedyAllocator( 

221 workflow=DEFAULT_WORKFLOW_CONFIG, 

222 latency_data=latency_data, 

223 ) 

224 

225 # 8 A100 

226 result_8a = allocator.allocate( 

227 num_gpus={GPUType.A100: 8, GPUType.H100: 0}, 

228 ) 

229 assert result_8a.gpus_used[GPUType.A100] == 8 

230 assert result_8a.gpus_used.get(GPUType.H100, 0) == 0 

231 assert 0 < result_8a.tbf_s < result_8a.ttff_s < result_8a.total_time_s 

232 

233 # 16 A100 

234 result_16a = allocator.allocate( 

235 num_gpus={GPUType.A100: 16, GPUType.H100: 0}, 

236 ) 

237 assert result_16a.gpus_used[GPUType.A100] == 16 

238 assert result_16a.gpus_used.get(GPUType.H100, 0) == 0 

239 assert 0 < result_16a.tbf_s < result_16a.ttff_s < result_16a.total_time_s 

240 # 16 A100s should be faster than 8 A100s 

241 assert result_8a.total_time_s > result_16a.total_time_s 

242 assert result_8a.ttff_s > result_16a.ttff_s 

243 assert result_8a.tbf_s > result_16a.tbf_s 

244 

245 # 24 A100 

246 result_24a = allocator.allocate( 

247 num_gpus={GPUType.A100: 24}, 

248 ) 

249 assert result_24a.gpus_used[GPUType.A100] == 24 

250 assert result_24a.gpus_used.get(GPUType.H100, 0) == 0 

251 assert result_24a.gpus_used.get(GPUType.H200, 0) == 0 

252 assert result_24a.gpus_used.get(GPUType.GB200, 0) == 0 

253 assert 0 < result_24a.tbf_s < result_24a.ttff_s < result_24a.total_time_s 

254 # 24 A100s should be faster than 16 A100s 

255 assert result_16a.total_time_s > result_24a.total_time_s 

256 assert result_16a.ttff_s > result_24a.ttff_s 

257 assert result_16a.tbf_s > result_24a.tbf_s 

258 

259 # 4096 A100 

260 result_4096a = allocator.allocate( 

261 num_gpus={GPUType.A100: 4096}, 

262 ) 

263 assert result_4096a.gpus_used[GPUType.A100] <= 4096 

264 assert result_4096a.gpus_used.get(GPUType.H100, 0) == 0 

265 assert 0 < result_4096a.tbf_s < result_4096a.ttff_s < result_4096a.total_time_s 

266 # 4096 A100s should be faster than 24 A100s 

267 assert result_24a.total_time_s > result_4096a.total_time_s 

268 assert result_24a.ttff_s > result_4096a.ttff_s 

269 assert result_24a.tbf_s > result_4096a.tbf_s 

270 

271 

272def test_estimate_total_time_H() -> None: 

273 """More H100s should lead to better performance.""" 

274 latency_data = load_latency_data("simulator/data/") 

275 

276 allocator = GreedyAllocator( 

277 workflow=DEFAULT_WORKFLOW_CONFIG, 

278 latency_data=latency_data, 

279 ) 

280 

281 # 8 H100 

282 result_8h = allocator.allocate( 

283 num_gpus={GPUType.H100: 8}, 

284 ) 

285 assert result_8h.gpus_used.get(GPUType.A100, 0) == 0 

286 assert result_8h.gpus_used[GPUType.H100] == 8 

287 assert 0 < result_8h.tbf_s < result_8h.ttff_s < result_8h.total_time_s 

288 

289 # 16 H100 

290 result_16h = allocator.allocate( 

291 num_gpus={GPUType.A100: 0, GPUType.H100: 16}, 

292 ) 

293 assert result_16h.gpus_used.get(GPUType.A100, 0) == 0 

294 assert result_16h.gpus_used[GPUType.H100] == 16 

295 assert 0 < result_16h.tbf_s < result_16h.ttff_s < result_16h.total_time_s 

296 # 16 H100s should be faster than 8 H100s 

297 assert result_8h.total_time_s > result_16h.total_time_s 

298 assert result_8h.ttff_s > result_16h.ttff_s 

299 assert result_8h.tbf_s > result_16h.tbf_s 

300 

301 # 24 H100 

302 result_24h = allocator.allocate( 

303 num_gpus={GPUType.H100: 24}, 

304 ) 

305 assert result_24h.gpus_used.get(GPUType.A100, 0) == 0 

306 assert result_24h.gpus_used[GPUType.H100] == 24 

307 assert 0 < result_24h.tbf_s < result_24h.ttff_s < result_24h.total_time_s 

308 # 24 H100s should be faster than 16 H100s 

309 assert result_16h.total_time_s > result_24h.total_time_s 

310 assert result_16h.ttff_s > result_24h.ttff_s 

311 assert result_16h.tbf_s > result_24h.tbf_s 

312 

313 

314def test_time_vs_power() -> None: 

315 """Compare time and power estimation for the same setup.""" 

316 latency_data = load_latency_data("simulator/data/") 

317 power_data = load_power_data("simulator/data/") 

318 

319 # Time only 

320 allocator_time = GreedyAllocator( 

321 workflow=DEFAULT_WORKFLOW_CONFIG, 

322 latency_data=latency_data, 

323 policy=STREAMWISE_POLICY, 

324 ) 

325 result_time = allocator_time.allocate( 

326 num_gpus={GPUType.A100: 0, GPUType.H100: 8}, 

327 verbose=True, 

328 ) 

329 assert 0 < result_time.tbf_s < result_time.ttff_s < result_time.total_time_s 

330 

331 # Time+Energy only 

332 allocator_energy = GreedyAllocator( 

333 workflow=DEFAULT_WORKFLOW_CONFIG, 

334 latency_data=latency_data, 

335 power_data=power_data, 

336 policy=STREAMWISE_POLICY, 

337 ) 

338 result_energy = allocator_energy.allocate( 

339 num_gpus={GPUType.H100: 8}, 

340 verbose=True, 

341 ) 

342 assert 0 < result_energy.tbf_s < result_energy.ttff_s < result_energy.total_time_s 

343 assert result_energy.total_time_s == result_time.total_time_s 

344 assert result_energy.ttff_s == result_time.ttff_s 

345 assert result_energy.tbf_s == result_time.tbf_s 

346 

347 assert 0 == result_time.total_energy 

348 assert 0 < result_energy.total_energy 

349 

350 

351@pytest.mark.parametrize("objective", list(Objective)) 

352def test_scheduler( 

353 objective: Objective, 

354) -> None: 

355 latency_data = load_latency_data("simulator/data/") 

356 power_data = load_power_data("simulator/data/") 

357 

358 policy = replace(STREAMWISE_POLICY) 

359 policy.objective = objective 

360 

361 allocator = AutoModelAllocator( 

362 workflow=DEFAULT_WORKFLOW_CONFIG, 

363 latency_data=latency_data, 

364 power_data=power_data, 

365 policy=policy, 

366 ) 

367 

368 result = allocator.allocate( 

369 num_gpus={GPUType.A100: 8, GPUType.H100: 8}, 

370 verbose=True, 

371 ) 

372 assert result.cost > 0 

373 if objective == Objective.NONE: 

374 assert result.gpus_used[GPUType.A100] == 3 

375 else: 

376 assert result.gpus_used[GPUType.A100] == 8 

377 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

378 assert 0 < result.total_energy 

379 

380 # Single server type 

381 result = allocator.allocate( 

382 num_gpus={GPUType.H100: 8}, 

383 ) 

384 assert result.gpus_used.get(GPUType.A100, 0) == 0 

385 assert 0 < result.gpus_used[GPUType.H100] <= 8 

386 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

387 assert 0 < result.total_energy 

388 

389 

390def test_scheduler_wrong() -> None: 

391 latency_data = load_latency_data("simulator/data/") 

392 power_data = load_power_data("simulator/data/") 

393 

394 policy = replace(STREAMWISE_POLICY) 

395 policy.objective = "wrong" 

396 

397 allocator = GreedyAllocator( 

398 workflow=DEFAULT_WORKFLOW_CONFIG, 

399 latency_data=latency_data, 

400 power_data=power_data, 

401 policy=policy, 

402 ) 

403 

404 with pytest.raises(ValueError, match="Cannot recognize objective wrong"): 

405 allocator.allocate( 

406 num_gpus={ 

407 GPUType.A100: 8, 

408 GPUType.H100: 8 

409 }, 

410 ) 

411 

412 

413@pytest.mark.parametrize("num_a100s", [num_a100s for num_a100s in range(8, 128 + 1, 8)]) 

414def test_estimate_A100( 

415 num_a100s: int, 

416) -> None: 

417 latency_data = load_latency_data("simulator/data/") 

418 

419 allocator = GreedyAllocator( 

420 workflow=DEFAULT_WORKFLOW_CONFIG, 

421 latency_data=latency_data, 

422 policy=STREAMWISE_POLICY, 

423 ) 

424 result = allocator.allocate( 

425 num_gpus={GPUType.A100: num_a100s, GPUType.H100: 0}, 

426 verbose=True, 

427 ) 

428 assert result.gpus_used[GPUType.A100] == num_a100s 

429 assert result.gpus_used.get(GPUType.H100, 0) == 0 

430 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

431 

432 

433@pytest.mark.parametrize("num_h100s", [num_h100s for num_h100s in range(8, 128 + 1, 8)]) 

434def test_estimate_H100( 

435 num_h100s: int, 

436) -> None: 

437 latency_data = load_latency_data("simulator/data/") 

438 

439 allocator = GreedyAllocator( 

440 workflow=DEFAULT_WORKFLOW_CONFIG, 

441 latency_data=latency_data, 

442 ) 

443 result = allocator.allocate( 

444 num_gpus={GPUType.H100: num_h100s}, 

445 verbose=True, 

446 ) 

447 assert result.gpus_used.get(GPUType.A100, 0) == 0 

448 assert result.gpus_used[GPUType.H100] == num_h100s 

449 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

450 

451 

452@pytest.mark.parametrize("num_h200s", [num_h200s for num_h200s in range(8, 128 + 1, 8)]) 

453def test_estimate_H200( 

454 num_h200s: int, 

455) -> None: 

456 latency_data = load_latency_data("simulator/data/") 

457 

458 allocator = GreedyAllocator( 

459 workflow=DEFAULT_WORKFLOW_CONFIG, 

460 latency_data=latency_data, 

461 policy=STREAMWISE_POLICY, 

462 ) 

463 result = allocator.allocate( 

464 num_gpus={GPUType.H200: num_h200s}, 

465 verbose=True, 

466 ) 

467 assert result.gpus_used.get(GPUType.A100, 0) == 0 

468 assert result.gpus_used.get(GPUType.H100, 0) == 0 

469 assert result.gpus_used[GPUType.H200] == num_h200s 

470 assert result.gpus_used.get(GPUType.GB200, 0) == 0 

471 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s 

472 

473 

474@pytest.mark.parametrize("num_gb200s", [num_gb200s for num_gb200s in range(8, 128 + 1, 8)]) 

475def test_estimate_GB200( 

476 num_gb200s: int, 

477) -> None: 

478 latency_data = load_latency_data("simulator/data/") 

479 

480 allocator = GreedyAllocator( 

481 workflow=DEFAULT_WORKFLOW_CONFIG, 

482 latency_data=latency_data, 

483 policy=STREAMWISE_POLICY, 

484 ) 

485 result = allocator.allocate( 

486 num_gpus={GPUType.GB200: num_gb200s}, 

487 verbose=True, 

488 ) 

489 assert result.gpus_used.get(GPUType.A100, 0) == 0 

490 assert result.gpus_used.get(GPUType.H100, 0) == 0 

491 assert result.gpus_used.get(GPUType.H200, 0) == 0 

492 assert result.gpus_used[GPUType.GB200] == num_gb200s 

493 assert 0 < result.tbf_s < result.ttff_s < result.total_time_s