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

1import sys 

2import os 

3import pytest 

4 

5# Add current path 

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

7 

8from tests.test_utils import temp_sys_path 

9 

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

11 from constants import DEFAULT_WORKFLOW_CONFIG 

12 

13 from provisioning import get_provisioning_results 

14 from provisioning import get_provisioning_adaptive_results 

15 from provisioning import Provision 

16 from provisioning import get_provisions 

17 from provisioning import GPU_PROVISIONS 

18 from provisioning import GPU_PROVISIONS_SHORT 

19 

20 from sim_types import GPUType 

21 from sim_types import QualityLevel 

22 from sim_types import Solver 

23 

24 from data_loading import load_latency_data 

25 

26 from model_provisioner.policies import NAIVE_POLICY 

27 from model_provisioner.policies import STREAMWISE_POLICY 

28 from model_provisioner.policies import HEXGEN_POLICY 

29 

30 

31@pytest.mark.parametrize("gpu_type", [gpu_type for gpu_type in GPUType]) 

32@pytest.mark.parametrize("num_gpus", [8, 16, 32, 64]) 

33def test_provisioning_streamwise_single_gpu_type( 

34 gpu_type: GPUType, 

35 num_gpus: int, 

36) -> None: 

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

38 provisions = [ 

39 Provision({gpu_type: num_gpus}), 

40 ] 

41 result = get_provisioning_results( 

42 workflow=DEFAULT_WORKFLOW_CONFIG, 

43 latency_data=latency_data, 

44 policy=STREAMWISE_POLICY, 

45 provisions=provisions, 

46 ) 

47 assert len(result.latencies) == 1 

48 assert len(result.costs) == 1 

49 assert len(result.ttffs) == 1 

50 assert len(result.tbfs) == 1 

51 assert len(result.actual_provision) == 1 

52 assert len(result.config_provision) == 1 

53 

54 assert result.config_provision == [ 

55 {gpu_type: num_gpus}, 

56 ] 

57 assert result.actual_provision == [ 

58 {gpu_type: num_gpus}, 

59 ] 

60 

61 

62def test_provisioning() -> None: 

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

64 result = get_provisioning_results( 

65 provisions=[ 

66 # A100, H100, H200 

67 # This now gets ordered differently 

68 Provision({GPUType.H100: 16}), 

69 Provision({GPUType.A100: 32}), 

70 Provision({GPUType.A100: 64, GPUType.H200: 128}), 

71 Provision({GPUType.A100: 64, GPUType.H100: 64}), 

72 ], 

73 workflow=DEFAULT_WORKFLOW_CONFIG, 

74 latency_data=latency_data, 

75 policy=STREAMWISE_POLICY, 

76 verbose=False, 

77 ) 

78 assert len(result.latencies) == 4 

79 assert len(result.costs) == 4 

80 assert len(result.ttffs) == 4 

81 assert len(result.tbfs) == 4 

82 assert len(result.actual_provision) == 4 

83 assert len(result.config_provision) == 4 

84 

85 assert result.config_provision == [ 

86 {GPUType.H100: 16}, 

87 {GPUType.A100: 32}, 

88 {GPUType.A100: 64, GPUType.H200: 128}, 

89 {GPUType.A100: 64, GPUType.H100: 64}, 

90 ] 

91 assert result.actual_provision[0] == {GPUType.H100: 16} 

92 assert result.actual_provision[1] == {GPUType.A100: 32} 

93 assert result.actual_provision[2] == {GPUType.A100: 64, GPUType.H200: 128} 

94 assert result.actual_provision[3] == {GPUType.A100: 64, GPUType.H100: 64} 

95 

96 

97def test_provisioning_not_gpus_error() -> None: 

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

99 with pytest.raises(AssertionError, match="No GPUs provisioned"): 

100 get_provisioning_results( 

101 workflow=DEFAULT_WORKFLOW_CONFIG, 

102 latency_data=latency_data, 

103 policy=NAIVE_POLICY, 

104 provisions=[ 

105 Provision({ # No GPU provision 

106 GPUType.A100: 0, 

107 GPUType.H100: 0, 

108 GPUType.H200: 0 

109 }), 

110 ], 

111 ) 

112 

113 

114def test_provisioning_3_gpus_error() -> None: 

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

116 with pytest.raises(AssertionError, match="Only support up to 2 GPU types in a provision"): 

117 get_provisioning_results( 

118 workflow=DEFAULT_WORKFLOW_CONFIG, 

119 latency_data=latency_data, 

120 policy=NAIVE_POLICY, 

121 provisions=[ 

122 Provision({ # All not supported 

123 GPUType.A100: 8, 

124 GPUType.H100: 8, 

125 GPUType.H200: 8, 

126 }), 

127 ], 

128 ) 

129 

130 

131def test_provisioning_adaptive_too_small() -> None: 

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

133 provisions = [ 

134 Provision({GPUType.A100: 8}), 

135 ] 

136 provisioning_result_high = get_provisioning_results( 

137 provisions=provisions, 

138 workflow=DEFAULT_WORKFLOW_CONFIG, 

139 latency_data=latency_data, 

140 policy=STREAMWISE_POLICY, 

141 ) 

142 provisioning_result_medium = get_provisioning_results( 

143 provisions=provisions, 

144 workflow=DEFAULT_WORKFLOW_CONFIG, 

145 latency_data=latency_data, 

146 policy=STREAMWISE_POLICY, 

147 ) 

148 provisioning_result_low = get_provisioning_results( 

149 provisions=provisions, 

150 workflow=DEFAULT_WORKFLOW_CONFIG, 

151 latency_data=latency_data, 

152 policy=STREAMWISE_POLICY, 

153 ) 

154 

155 with pytest.raises(ValueError, match="Low quality TTFF .+ exceeds video length .+"): 

156 get_provisioning_adaptive_results( 

157 workflow_config=DEFAULT_WORKFLOW_CONFIG, 

158 provisioning_qualities={ 

159 QualityLevel.HIGH: provisioning_result_high, 

160 QualityLevel.MEDIUM: provisioning_result_medium, 

161 QualityLevel.LOW: provisioning_result_low, 

162 }, 

163 ) 

164 

165 

166def test_provisioning_adaptive() -> None: 

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

168 provisions = [ 

169 Provision({GPUType.A100: 128}), 

170 Provision({GPUType.H100: 128}), 

171 Provision({GPUType.A100: 24, GPUType.H100: 32}), 

172 Provision({GPUType.A100: 128, GPUType.H100: 32}), # TODO This should work with fewer H100 

173 ] 

174 provisioning_result_high = get_provisioning_results( 

175 provisions=provisions, 

176 workflow=DEFAULT_WORKFLOW_CONFIG, 

177 latency_data=latency_data, 

178 policy=STREAMWISE_POLICY, 

179 ) 

180 provisioning_result_medium = get_provisioning_results( 

181 provisions=provisions, 

182 workflow=DEFAULT_WORKFLOW_CONFIG, 

183 latency_data=latency_data, 

184 policy=STREAMWISE_POLICY, 

185 ) 

186 provisioning_result_low = get_provisioning_results( 

187 provisions=provisions, 

188 workflow=DEFAULT_WORKFLOW_CONFIG, 

189 latency_data=latency_data, 

190 policy=STREAMWISE_POLICY, 

191 ) 

192 adaptive_results = get_provisioning_adaptive_results( 

193 workflow_config=DEFAULT_WORKFLOW_CONFIG, 

194 provisioning_qualities={ 

195 QualityLevel.HIGH: provisioning_result_high, 

196 QualityLevel.MEDIUM: provisioning_result_medium, 

197 QualityLevel.LOW: provisioning_result_low, 

198 }, 

199 ) 

200 assert adaptive_results is not None 

201 assert len(adaptive_results.costs) == len(provisions) 

202 assert len(adaptive_results.latencies) == len(provisions) 

203 assert len(adaptive_results.ttffs) == len(provisions) 

204 assert len(adaptive_results.tbfs) == len(provisions) 

205 assert len(adaptive_results.actual_provision) == len(provisions) 

206 assert len(adaptive_results.config_provision) == len(provisions) 

207 

208 

209def test_get_provisions() -> None: 

210 """Test get_provisions() function for various GPU type combinations.""" 

211 provisions_a100 = get_provisions([GPUType.A100]) 

212 assert 10 < len(provisions_a100) == len(GPU_PROVISIONS) - 1 

213 for provision in provisions_a100: 

214 assert provision.num_gpus.get(GPUType.A100, 0) > 0 

215 assert provision.num_gpus.get(GPUType.H100, 0) == 0 

216 assert provision.num_gpus.get(GPUType.H200, 0) == 0 

217 

218 provisions_h100 = get_provisions([GPUType.H100]) 

219 assert 10 < len(provisions_h100) == len(GPU_PROVISIONS) - 1 

220 for provision in provisions_h100: 

221 assert provision.num_gpus.get(GPUType.A100, 0) == 0 

222 assert provision.num_gpus.get(GPUType.H100, 0) > 0 

223 assert provision.num_gpus.get(GPUType.H200, 0) == 0 

224 

225 provisions_h200 = get_provisions([GPUType.H200]) 

226 assert 10 < len(provisions_h200) == len(GPU_PROVISIONS) - 1 

227 for provision in provisions_h200: 

228 assert provision.num_gpus.get(GPUType.A100, 0) == 0 

229 assert provision.num_gpus.get(GPUType.H100, 0) == 0 

230 assert provision.num_gpus.get(GPUType.H200, 0) > 0 

231 

232 # Mixed mode (A100 + H100) 

233 provisions_mixed = get_provisions([GPUType.A100, GPUType.H100]) 

234 assert len(provisions_mixed) == len(GPU_PROVISIONS) ** 2 - 1 # num_gpus ^ num_types - 1 (empty) 

235 for provision in provisions_mixed: 

236 assert ( 

237 provision.num_gpus.get(GPUType.A100, 0) > 0 

238 or provision.num_gpus.get(GPUType.H100, 0) > 0 

239 ) 

240 

241 # Mixed mode (all types) 

242 provisions_mixed = get_provisions([ 

243 GPUType.A100, 

244 GPUType.H100, 

245 GPUType.H200, 

246 ], limits_pairs=False) 

247 assert len(provisions_mixed) == len(GPU_PROVISIONS) ** 3 - 1 # num_gpus ^ num_types - 1 (empty) 

248 for provision in provisions_mixed: 

249 assert ( 

250 provision.num_gpus.get(GPUType.A100, 0) > 0 

251 or provision.num_gpus.get(GPUType.H100, 0) > 0 

252 or provision.num_gpus.get(GPUType.H200, 0) > 0 

253 ) 

254 # Make sure specific cases exist 

255 assert any( 

256 provision.num_gpus.get(GPUType.A100, 0) == 8 

257 and provision.num_gpus.get(GPUType.H100, 0) == 8 

258 and provision.num_gpus.get(GPUType.H200, 0) == 0 

259 for provision in provisions_mixed 

260 ) 

261 assert any( 

262 provision.num_gpus.get(GPUType.A100, 0) == 8 

263 and provision.num_gpus.get(GPUType.H100, 0) == 0 

264 and provision.num_gpus.get(GPUType.H200, 0) == 0 

265 for provision in provisions_mixed 

266 ) 

267 assert any( 

268 provision.num_gpus.get(GPUType.A100, 0) == 8 

269 and provision.num_gpus.get(GPUType.H100, 0) == 8 

270 and provision.num_gpus.get(GPUType.H200, 0) == 8 

271 for provision in provisions_mixed 

272 ) 

273 

274 

275def test_get_provisions_short() -> None: 

276 """Test get_provisions() function for various GPU type combinations.""" 

277 provisions_a100 = get_provisions([GPUType.A100], short_list=True) 

278 assert 10 < len(provisions_a100) == len(GPU_PROVISIONS_SHORT) - 1 

279 for provision in provisions_a100: 

280 assert provision.num_gpus.get(GPUType.A100, 0) > 0 

281 assert provision.num_gpus.get(GPUType.H100, 0) == 0 

282 assert provision.num_gpus.get(GPUType.H200, 0) == 0 

283 

284 provisions_mixed = get_provisions([ 

285 GPUType.A100, 

286 GPUType.H100, 

287 GPUType.H200, 

288 ], limits_pairs=False, short_list=True) 

289 assert len(provisions_mixed) == len(GPU_PROVISIONS_SHORT) ** 3 - 1 # num_gpus ^ num_types - 1 (empty) 

290 for provision in provisions_mixed: 

291 assert ( 

292 provision.num_gpus.get(GPUType.A100, 0) > 0 

293 or provision.num_gpus.get(GPUType.H100, 0) > 0 

294 or provision.num_gpus.get(GPUType.H200, 0) > 0 

295 ) 

296 

297 

298# --- HexGenAllocator provisioning tests --- 

299@pytest.mark.parametrize("gpu_type", [GPUType.A100, GPUType.H100]) 

300@pytest.mark.parametrize("num_gpus", [8, 16, 32]) 

301def test_provisioning_hexgen_single_gpu_type( 

302 gpu_type: GPUType, 

303 num_gpus: int, 

304) -> None: 

305 """Test provisioning with HexGenAllocator for single GPU types.""" 

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

307 provisions = [ 

308 Provision({gpu_type: num_gpus}), 

309 ] 

310 result = get_provisioning_results( 

311 workflow=DEFAULT_WORKFLOW_CONFIG, 

312 latency_data=latency_data, 

313 policy=HEXGEN_POLICY, 

314 provisions=provisions, 

315 ) 

316 assert len(result.latencies) == 1 

317 assert len(result.costs) == 1 

318 assert len(result.ttffs) == 1 

319 assert len(result.tbfs) == 1 

320 assert len(result.actual_provision) == 1 

321 assert len(result.config_provision) == 1 

322 

323 assert result.config_provision == [ 

324 {gpu_type: num_gpus}, 

325 ] 

326 assert result.actual_provision == [ 

327 {gpu_type: num_gpus}, 

328 ] 

329 

330 

331def test_hexgen_mixed() -> None: 

332 """Test provisioning with HexGenAllocator for mixed GPU types.""" 

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

334 result = get_provisioning_results( 

335 provisions=[ 

336 Provision({GPUType.A100: 8, GPUType.H100: 8}), 

337 Provision({GPUType.A100: 32}), 

338 ], 

339 workflow=DEFAULT_WORKFLOW_CONFIG, 

340 latency_data=latency_data, 

341 policy=HEXGEN_POLICY, 

342 verbose=False, 

343 ) 

344 assert len(result.latencies) == 2 

345 assert len(result.costs) == 2 

346 assert len(result.ttffs) == 2 

347 assert len(result.tbfs) == 2 

348 assert len(result.actual_provision) == 2 

349 assert len(result.config_provision) == 2 

350 

351 

352def test_default_allocator_is_greedy() -> None: 

353 """Verify default policy allocator is greedy (backward compat).""" 

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

355 provisions = [ 

356 Provision({GPUType.A100: 8}), 

357 ] 

358 # STREAMWISE_POLICY uses default allocator="greedy" 

359 result = get_provisioning_results( 

360 workflow=DEFAULT_WORKFLOW_CONFIG, 

361 latency_data=latency_data, 

362 policy=STREAMWISE_POLICY, 

363 provisions=provisions, 

364 ) 

365 assert STREAMWISE_POLICY.solver == Solver.GREEDY 

366 assert len(result.latencies) == 1 

367 

368 

369def test_hexgen_vs_greedy() -> None: 

370 """Both allocators should produce valid results for the same provision.""" 

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

372 provisions = [ 

373 Provision({GPUType.A100: 32}), 

374 ] 

375 result_greedy = get_provisioning_results( 

376 workflow=DEFAULT_WORKFLOW_CONFIG, 

377 latency_data=latency_data, 

378 policy=STREAMWISE_POLICY, 

379 provisions=provisions, 

380 ) 

381 result_hexgen = get_provisioning_results( 

382 workflow=DEFAULT_WORKFLOW_CONFIG, 

383 latency_data=latency_data, 

384 policy=HEXGEN_POLICY, 

385 provisions=provisions, 

386 ) 

387 # Both should produce 1 result 

388 assert len(result_greedy.latencies) == 1 

389 assert len(result_hexgen.latencies) == 1 

390 # Both should have positive metrics 

391 assert result_greedy.latencies[0] > 0 

392 assert result_hexgen.latencies[0] > 0 

393 assert result_greedy.costs[0] > 0 

394 assert result_hexgen.costs[0] > 0 

395 assert result_greedy.ttffs[0] > 0 

396 assert result_hexgen.ttffs[0] > 0