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

1""" 

2Test baselines. 

3""" 

4 

5import sys 

6import os 

7import pytest 

8 

9# Add current path 

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

11 

12from tests.test_utils import temp_sys_path 

13 

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

15 from sim_types import GPUType 

16 from sim_types import Model 

17 

18 from constants import DEFAULT_WORKFLOW_CONFIG 

19 from constants import SECONDS_IN_HOUR 

20 from constants import POWER_GPU_IDLE 

21 from constants import POWER_GPU_TDP 

22 

23 from data_loading import load_latency_data 

24 from data_loading import load_power_data 

25 

26 from auto_model_allocator import AutoModelAllocator 

27 from model_provisioner.naive_baseline import NaiveAllocator 

28 from model_provisioner.greedy import GreedyAllocator 

29 

30 from model_provisioner.policies import NAIVE_POLICY 

31 from model_provisioner.policies import BASELINE_POLICIES 

32 from model_provisioner.policies import STREAMWISE_POLICY 

33 

34 from workflows import SHORTS_WORKFLOW 

35 from workflows import WORKFLOWS 

36 

37 

38def test_baseline() -> None: 

39 """8 A100 + 8 H100.""" 

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

41 allocator = NaiveAllocator( 

42 workflow=DEFAULT_WORKFLOW_CONFIG, 

43 latency_data=latency_data, 

44 ) 

45 result = allocator.allocate( 

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

47 ) 

48 assert 0 < result.gpus_used.get(GPUType.A100, 0) <= 8 

49 assert 0 < result.gpus_used.get(GPUType.H100, 0) <= 8 

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

51 

52 

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

54def test_baseline_A_options( 

55 num_a100s: int, 

56) -> None: 

57 """A100 combinations.""" 

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

59 allocator = NaiveAllocator( 

60 workflow=DEFAULT_WORKFLOW_CONFIG, 

61 latency_data=latency_data, 

62 ) 

63 result = allocator.allocate( 

64 num_gpus={GPUType.A100: num_a100s}, 

65 verbose=True, 

66 ) 

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

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

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

70 

71 

72def test_baseline_8A_upscaler() -> None: 

73 """8 A100 + 0 H100 with upscaler.""" 

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

75 allocator = NaiveAllocator( 

76 workflow=DEFAULT_WORKFLOW_CONFIG, 

77 latency_data=latency_data, 

78 policy=BASELINE_POLICIES["naive upscaler"], 

79 ) 

80 result = allocator.allocate( 

81 num_gpus={GPUType.A100: 8}, 

82 verbose=True, 

83 ) 

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

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

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

87 

88 

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

90def test_baseline_H( 

91 num_h100s: int, 

92) -> None: 

93 """H100 combinations.""" 

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

95 allocator = NaiveAllocator( 

96 workflow=DEFAULT_WORKFLOW_CONFIG, 

97 latency_data=latency_data, 

98 ) 

99 result = allocator.allocate( 

100 num_gpus={GPUType.H100: num_h100s}, 

101 verbose=True, 

102 ) 

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

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

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

106 

107 

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

109def test_baseline_H_upscaler( 

110 num_h100s: int, 

111) -> None: 

112 """H100 with upscaler.""" 

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

114 allocator = NaiveAllocator( 

115 workflow=DEFAULT_WORKFLOW_CONFIG, 

116 latency_data=latency_data, 

117 policy=BASELINE_POLICIES["naive upscaler"], 

118 ) 

119 result = allocator.allocate( 

120 num_gpus={GPUType.H100: num_h100s}, 

121 ) 

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

123 assert result.gpus_used.get(GPUType.H100, 0) <= num_h100s 

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

125 

126 

127def test_baseline_24H() -> None: 

128 """0 A100 + 24 H100""" 

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

130 

131 allocator = NaiveAllocator( 

132 workflow=DEFAULT_WORKFLOW_CONFIG, 

133 latency_data=latency_data, 

134 ) 

135 result = allocator.allocate( 

136 num_gpus={GPUType.A100: 0, GPUType.H100: 24}, 

137 verbose=True, 

138 ) 

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

140 assert 24 - 8 < result.gpus_used.get(GPUType.H100, 0) <= 24 

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

142 

143 

144def test_baseline_96A_96H() -> None: 

145 """96 A100 + 96 H100.""" 

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

147 

148 allocator = NaiveAllocator( 

149 workflow=DEFAULT_WORKFLOW_CONFIG, 

150 latency_data=latency_data, 

151 policy=NAIVE_POLICY, 

152 ) 

153 result = allocator.allocate( 

154 num_gpus={GPUType.A100: 96, GPUType.H100: 96}, 

155 verbose=True, 

156 ) 

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

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

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

160 

161 

162def test_baseline_64A_16H() -> None: 

163 """64 A100 + 16 H100.""" 

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

165 

166 allocator = NaiveAllocator( 

167 workflow=DEFAULT_WORKFLOW_CONFIG, 

168 latency_data=latency_data, 

169 policy=BASELINE_POLICIES["naive upscaler"], 

170 ) 

171 result = allocator.allocate( 

172 num_gpus={GPUType.A100: 64, GPUType.H100: 16}, 

173 verbose=True, 

174 ) 

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

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

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

178 

179 

180def test_baseline_no_disaggregation() -> None: 

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

182 

183 # 8 A100 + 8 H100 

184 allocator = NaiveAllocator( 

185 workflow=DEFAULT_WORKFLOW_CONFIG, 

186 latency_data=latency_data, 

187 policy=NAIVE_POLICY, 

188 ) 

189 result = allocator.allocate( 

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

191 verbose=True, 

192 ) 

193 assert 0 < result.gpus_used.get(GPUType.A100, 0) <= 8 

194 assert 0 < result.gpus_used.get(GPUType.H100, 0) <= 8 

195 

196 

197def test_baseline_0() -> None: 

198 """No GPUs.""" 

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

200 with pytest.raises(AssertionError, match="Total number of GPUs must be at least 8"): 

201 allocator = NaiveAllocator( 

202 workflow=DEFAULT_WORKFLOW_CONFIG, 

203 latency_data=latency_data, 

204 ) 

205 allocator.allocate( 

206 num_gpus={GPUType.A100: 0, GPUType.H100: 0}, 

207 ) 

208 

209 

210def test_baseline_timexcost_1024A() -> None: 

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

212 policy = BASELINE_POLICIES["naive ttff*cost allocator"] 

213 allocator = AutoModelAllocator( 

214 workflow=DEFAULT_WORKFLOW_CONFIG, 

215 latency_data=latency_data, 

216 policy=policy, 

217 ) 

218 result = allocator.allocate( 

219 num_gpus={GPUType.A100: 1024}, 

220 ) 

221 assert result.gpus_used[GPUType.A100] == 1024 

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

223 

224 

225@pytest.mark.parametrize("policy_name", BASELINE_POLICIES.keys()) 

226def test_baseline_policies( 

227 policy_name: str 

228) -> None: 

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

230 policy = BASELINE_POLICIES[policy_name] 

231 allocator = AutoModelAllocator( 

232 workflow=DEFAULT_WORKFLOW_CONFIG, 

233 latency_data=latency_data, 

234 policy=policy, 

235 ) 

236 result = allocator.allocate( 

237 num_gpus={ 

238 GPUType.A100: 64, 

239 GPUType.H100: 64 

240 }, 

241 ) 

242 assert 0 < result.gpus_used[GPUType.A100] <= 64 

243 assert 0 < result.gpus_used[GPUType.H100] <= 64 

244 

245 

246def test_baseline_streamwise() -> None: 

247 """StreamWise policy.""" 

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

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

250 

251 policy_streamwise = STREAMWISE_POLICY 

252 assert policy_streamwise.objective.value == "ttff_cost" 

253 

254 allocator = GreedyAllocator( 

255 workflow=DEFAULT_WORKFLOW_CONFIG, 

256 latency_data=latency_data, 

257 power_data=power_data, 

258 policy=policy_streamwise, 

259 ) 

260 result = allocator.allocate( 

261 num_gpus={GPUType.A100: 8}, 

262 ) 

263 

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

265 assert GPUType.H100 not in result.gpus_used 

266 assert GPUType.H200 not in result.gpus_used 

267 assert GPUType.GB200 not in result.gpus_used 

268 

269 assert len(result.models) == 1 

270 assert GPUType.A100 in result.models 

271 models = result.models[GPUType.A100] 

272 # 8 GPUs used 

273 assert models[Model.GEMMA][0].get_num_gpus() == 1 

274 assert models[Model.OTHERS][0].get_num_gpus() == 1 

275 assert models[Model.FLUX][0].get_num_gpus() == 1 

276 assert models[Model.HF][0].get_num_gpus() == 1 

277 assert models[Model.HF_VAE][0].get_num_gpus() == 1 

278 assert models[Model.FT][0].get_num_gpus() == 2 

279 assert models[Model.UPSCALER][0].get_num_gpus() == 1 

280 

281 assert 4 * SECONDS_IN_HOUR < result.total_time_s < 5 * SECONDS_IN_HOUR # 4-5 hours 

282 assert 4 * SECONDS_IN_HOUR < result.ttff_s < result.total_time_s # 4-5 hours 

283 assert 1 < result.tbf_s < 2 # 1-2 seconds 

284 

285 assert result.total_time_s * 8 * POWER_GPU_IDLE[GPUType.A100] < result.total_energy 

286 assert result.total_energy < result.total_time_s * 8 * POWER_GPU_TDP[GPUType.A100] 

287 

288 assert 37 < result.cost < 38 # $37-38 

289 

290 

291def test_baseline_naive() -> None: 

292 """Naive policy.""" 

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

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

295 

296 policy_naive = BASELINE_POLICIES["naive"] 

297 assert policy_naive.objective.value == "ttff" 

298 assert policy_naive.solver.value == "naive" 

299 

300 allocator = NaiveAllocator( 

301 workflow=DEFAULT_WORKFLOW_CONFIG, 

302 latency_data=latency_data, 

303 power_data=power_data, 

304 policy=policy_naive, 

305 ) 

306 result = allocator.allocate( 

307 num_gpus={GPUType.A100: 8}, 

308 ) 

309 

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

311 assert GPUType.H100 not in result.gpus_used 

312 assert GPUType.H200 not in result.gpus_used 

313 assert GPUType.GB200 not in result.gpus_used 

314 

315 assert len(result.models) == 1 

316 assert GPUType.A100 in result.models 

317 models = result.models[GPUType.A100] 

318 # 8 GPUs used 

319 assert models[Model.GEMMA][0].get_num_gpus() == 1 

320 assert models[Model.OTHERS][0].get_num_gpus() == 1 

321 assert models[Model.FLUX][0].get_num_gpus() == 1 

322 assert models[Model.HF][0].get_num_gpus() == 1 

323 assert models[Model.HF_VAE][0].get_num_gpus() == 0 # No dissaggregation 

324 assert models[Model.FT][0].get_num_gpus() == 4 

325 assert models[Model.UPSCALER][0].get_num_gpus() == 0 # no upscaler 

326 

327 # 8.3 hours from the paper 

328 assert 8 * SECONDS_IN_HOUR < result.total_time_s < 9 * SECONDS_IN_HOUR # 8-9 hours 

329 assert 7 * SECONDS_IN_HOUR < result.ttff_s < result.total_time_s # 7-8 hours 

330 assert 2 < result.tbf_s < 3 # 2-3 seconds 

331 

332 assert result.total_time_s * 8 * POWER_GPU_IDLE[GPUType.A100] < result.total_energy 

333 assert result.total_energy < result.total_time_s * 8 * POWER_GPU_TDP[GPUType.A100] 

334 

335 # With Spot is ~$71 and with Reserved is ~$225 

336 assert 225 < result.cost < 230 # $225-230 

337 

338 

339def test_baseline_upscaler() -> None: 

340 """Naive policy with upscaler.""" 

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

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

343 

344 policy_naive_upscaler = BASELINE_POLICIES["naive upscaler"] 

345 assert policy_naive_upscaler.objective.value == "ttff" 

346 assert policy_naive_upscaler.solver.value == "naive" 

347 

348 allocator = NaiveAllocator( 

349 workflow=DEFAULT_WORKFLOW_CONFIG, 

350 latency_data=latency_data, 

351 power_data=power_data, 

352 policy=policy_naive_upscaler, 

353 ) 

354 result = allocator.allocate( 

355 num_gpus={GPUType.A100: 8}, 

356 ) 

357 

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

359 assert GPUType.H100 not in result.gpus_used 

360 assert GPUType.H200 not in result.gpus_used 

361 assert GPUType.GB200 not in result.gpus_used 

362 

363 assert len(result.models) == 1 

364 assert GPUType.A100 in result.models 

365 models = result.models[GPUType.A100] 

366 # 8 GPUs used 

367 assert models[Model.GEMMA][0].get_num_gpus() == 1 

368 assert models[Model.OTHERS][0].get_num_gpus() == 1 

369 assert models[Model.FLUX][0].get_num_gpus() == 1 

370 assert models[Model.HF][0].get_num_gpus() == 1 

371 assert models[Model.HF_VAE][0].get_num_gpus() == 0 # No dissaggregation 

372 assert models[Model.FT][0].get_num_gpus() == 3 

373 assert models[Model.UPSCALER][0].get_num_gpus() == 1 

374 

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

376 assert 3 * SECONDS_IN_HOUR < result.total_time_s < 4 * SECONDS_IN_HOUR # 3-4 hours 

377 assert 3 * SECONDS_IN_HOUR < result.ttff_s < result.total_time_s # 3-4 hours 

378 assert 0.5 < result.tbf_s < 1 # 0.5-1 seconds 

379 

380 assert result.total_time_s * 8 * POWER_GPU_IDLE[GPUType.A100] < result.total_energy 

381 assert result.total_energy < result.total_time_s * 8 * POWER_GPU_TDP[GPUType.A100] 

382 

383 assert 91 < result.cost < 92 # $91-92 

384 

385 

386def test_baseline_time_cost() -> None: 

387 """Naive policy with upscaler.""" 

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

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

390 

391 policy_naive_sched = BASELINE_POLICIES["naive ttff*cost allocator"] 

392 assert policy_naive_sched.objective.value == "ttff_cost" 

393 assert policy_naive_sched.solver.value == "greedy" 

394 

395 allocator = AutoModelAllocator( 

396 workflow=DEFAULT_WORKFLOW_CONFIG, 

397 latency_data=latency_data, 

398 power_data=power_data, 

399 policy=policy_naive_sched, 

400 ) 

401 result = allocator.allocate( 

402 num_gpus={GPUType.A100: 8}, 

403 ) 

404 

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

406 assert GPUType.H100 not in result.gpus_used 

407 assert GPUType.H200 not in result.gpus_used 

408 assert GPUType.GB200 not in result.gpus_used 

409 

410 assert len(result.models) == 1 

411 assert GPUType.A100 in result.models 

412 models = result.models[GPUType.A100] 

413 # 8 GPUs used 

414 assert models[Model.GEMMA][0].get_num_gpus() == 1 

415 assert models[Model.OTHERS][0].get_num_gpus() == 1 

416 assert models[Model.FLUX][0].get_num_gpus() == 1 

417 assert models[Model.HF][0].get_num_gpus() == 3 

418 assert models[Model.HF_VAE][0].get_num_gpus() == 0 # No dissaggregation 

419 assert models[Model.FT][0].get_num_gpus() == 2 

420 assert models[Model.UPSCALER][0].get_num_gpus() == 0 

421 

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

423 # TODO figure if this is what is expected 

424 # assert 8 * SECONDS_IN_HOUR < result.total_time_s < 9 * SECONDS_IN_HOUR # 8-9 hours 

425 assert 8 * SECONDS_IN_HOUR < result.total_time_s < 14 * SECONDS_IN_HOUR # 8-14 hours 

426 assert 7 * SECONDS_IN_HOUR < result.ttff_s < result.total_time_s # 7-8 hours 

427 assert 1 < result.tbf_s < 4 # 1-4 seconds 

428 assert result.total_time_s * 8 * POWER_GPU_IDLE[GPUType.A100] < result.total_energy 

429 assert result.total_energy < result.total_time_s * 8 * POWER_GPU_TDP[GPUType.A100] 

430 

431 

432def test_baseline_hardware_error() -> None: 

433 """Naive policy without naive parallelism.""" 

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

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

436 

437 policy_naive_hardware = BASELINE_POLICIES["naive hardware"] 

438 assert policy_naive_hardware.objective.value == "ttff" 

439 assert policy_naive_hardware.solver.value == "naive" 

440 

441 allocator = NaiveAllocator( 

442 workflow=DEFAULT_WORKFLOW_CONFIG, 

443 latency_data=latency_data, 

444 power_data=power_data, 

445 policy=policy_naive_hardware, 

446 ) 

447 result = allocator.allocate( 

448 num_gpus={GPUType.A100: 24, GPUType.H200: 1280}, 

449 ) 

450 

451 assert result.gpus_used[GPUType.A100] == 24 

452 assert result.gpus_used[GPUType.H200] == 1264 

453 assert GPUType.H100 not in result.gpus_used 

454 assert GPUType.GB200 not in result.gpus_used 

455 

456 assert len(result.models) == 2 

457 assert GPUType.A100 in result.models 

458 assert GPUType.H200 in result.models 

459 models = result.models[GPUType.A100] 

460 assert models[Model.GEMMA][0].get_num_gpus() == 2 

461 assert models[Model.OTHERS][0].get_num_gpus() == 1 

462 assert models[Model.FLUX][0].get_num_gpus() == 2 

463 assert models[Model.HF][0].get_num_gpus() == 8 

464 assert models[Model.HF_VAE][0].get_num_gpus() == 0 

465 assert models[Model.FT][0].get_num_gpus() == 11 

466 assert models[Model.FT_VAE][0].get_num_gpus() == 0 

467 assert models[Model.UPSCALER][0].get_num_gpus() == 0 

468 

469 assert 500 < result.total_time_s < 810 

470 assert 130 < result.ttff_s < result.total_time_s 

471 assert 0 < result.tbf_s < 1 

472 assert result.total_time_s * (1265 + 24) * POWER_GPU_IDLE[GPUType.A100] < result.total_energy 

473 assert result.total_energy < result.total_time_s * (1265 + 24) * POWER_GPU_TDP[GPUType.H200] 

474 

475 

476def test_workflow_short() -> None: 

477 """Test the shorts workflow with the naive policy.""" 

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

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

480 

481 allocator = NaiveAllocator( 

482 workflow=SHORTS_WORKFLOW, 

483 latency_data=latency_data, 

484 power_data=power_data, 

485 ) 

486 result = allocator.allocate( 

487 num_gpus={GPUType.A100: 56}, 

488 ) 

489 

490 assert result.gpus_used[GPUType.A100] == 49 

491 assert GPUType.H100 not in result.gpus_used 

492 assert GPUType.H200 not in result.gpus_used 

493 assert GPUType.GB200 not in result.gpus_used 

494 

495 assert len(result.models) == 1 

496 assert GPUType.A100 in result.models 

497 models = result.models[GPUType.A100] 

498 assert models[Model.GEMMA][0].get_num_gpus() == 48 

499 assert models[Model.OTHERS][0].get_num_gpus() == 1 

500 assert models[Model.FLUX][0].get_num_gpus() == 0 

501 assert models[Model.HF][0].get_num_gpus() == 0 

502 assert models[Model.HF_VAE][0].get_num_gpus() == 0 

503 assert models[Model.FT][0].get_num_gpus() == 0 

504 assert models[Model.UPSCALER][0].get_num_gpus() == 0 

505 

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

507 assert result.total_time_s * 8 * POWER_GPU_IDLE[GPUType.A100] < result.total_energy 

508 # assert result.total_energy < result.total_time_s * 8 * POWER_GPU_TDP[GPUType.A100] 

509 

510 

511@pytest.mark.parametrize("workflow_name", WORKFLOWS.keys()) 

512def test_workflows(workflow_name: str) -> None: 

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

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

515 

516 allocator = NaiveAllocator( 

517 workflow=WORKFLOWS[workflow_name], 

518 latency_data=latency_data, 

519 power_data=power_data, 

520 ) 

521 result = allocator.allocate( 

522 num_gpus={GPUType.A100: 32}, 

523 ) 

524 assert GPUType.A100 in result.models 

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