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

1import sys 

2import os 

3import pytest 

4from dataclasses import replace 

5 

6# Add current path 

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

8 

9from tests.test_utils import temp_sys_path 

10 

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

12 from constants import DEFAULT_WORKFLOW_CONFIG 

13 from constants import SECONDS_IN_HOUR 

14 

15 from workflows import WORKFLOWS 

16 

17 from sim_types import GPUType 

18 from sim_types import QualityLevel 

19 from sim_types import WorkflowConfig 

20 

21 from data_loading import load_latency_data 

22 from data_loading import load_power_data 

23 

24 from model_provisioner.greedy import GreedyAllocator 

25 

26 from model_provisioner.policies import STREAMWISE_POLICY 

27 

28 

29def test_allocate_8A_8H() -> None: 

30 """8 A100 + 8 H100.""" 

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

32 allocator = GreedyAllocator( 

33 workflow=DEFAULT_WORKFLOW_CONFIG, 

34 latency_data=latency_data, 

35 ) 

36 result = allocator.allocate( 

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

38 ) 

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

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

41 assert 0 < result.ttff_s < result.total_time_s 

42 # assert 4100 <= result.total_time_s <= 4400 # TODO old values 

43 assert 50 * 60 <= result.total_time_s <= 60 * 60 # 50-60 minutes 

44 # assert 3500 <= result.ttff_s <= 3800 # TODO old values 

45 assert 45 * 60 <= result.ttff_s <= 55 * 60 # 45-55 minutes 

46 assert 0.2 <= result.tbf_s <= 0.4 

47 

48 

49def test_allocate_32A_32H() -> None: 

50 """32 A100 + 32 H100""" 

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

52 allocator = GreedyAllocator( 

53 workflow=DEFAULT_WORKFLOW_CONFIG, 

54 latency_data=latency_data, 

55 ) 

56 result = allocator.allocate( 

57 num_gpus={GPUType.A100: 32, GPUType.H100: 32}, 

58 verbose=True, 

59 ) 

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

61 assert result.gpus_used[GPUType.H100] == 32 

62 assert 0 < result.ttff_s < result.total_time_s 

63 assert 0 < result.tbf_s < 1 

64 

65 

66def test_allocate_16A() -> None: 

67 """16 A100 + 0 H100.""" 

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

69 allocator = GreedyAllocator( 

70 workflow=DEFAULT_WORKFLOW_CONFIG, 

71 latency_data=latency_data, 

72 ) 

73 result = allocator.allocate( 

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

75 verbose=True, 

76 ) 

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

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

79 assert 0 < result.ttff_s < result.total_time_s < 10 * 60 * 60 

80 assert 0 < result.tbf_s < 1 

81 

82 

83def test_allocate_64A() -> None: 

84 """64 A100 + 0 H100.""" 

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

86 allocator = GreedyAllocator( 

87 workflow=DEFAULT_WORKFLOW_CONFIG, 

88 latency_data=latency_data, 

89 ) 

90 result = allocator.allocate( 

91 num_gpus={GPUType.A100: 64, GPUType.H100: 0}, 

92 verbose=True, 

93 ) 

94 assert result.gpus_used[GPUType.A100] == 64 

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

96 assert 0 < result.ttff_s < result.total_time_s < 10 * 60 * 60 

97 assert 0 < result.tbf_s < 1 

98 

99 

100def test_allocate_64H() -> None: 

101 """0 A100 + 64 H100.""" 

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

103 allocator = GreedyAllocator( 

104 workflow=DEFAULT_WORKFLOW_CONFIG, 

105 latency_data=latency_data, 

106 ) 

107 result = allocator.allocate( 

108 num_gpus={GPUType.A100: 0, GPUType.H100: 64}, 

109 ) 

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

111 assert result.gpus_used[GPUType.H100] == 64 

112 assert 0 < result.ttff_s < result.total_time_s < 10 * 60 * 60 

113 assert 0 < result.tbf_s < 1 

114 

115 

116def test_allocate_64GB() -> None: 

117 """64 GB200.""" 

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

119 allocator = GreedyAllocator( 

120 workflow=DEFAULT_WORKFLOW_CONFIG, 

121 latency_data=latency_data, 

122 ) 

123 result = allocator.allocate( 

124 num_gpus={GPUType.GB200: 64}, 

125 verbose=True, 

126 ) 

127 assert result.gpus_used[GPUType.GB200] == 64 

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

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

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

131 assert 0 < result.ttff_s < result.total_time_s < 10 * 60 * 60 

132 assert 0 < result.tbf_s < 1 

133 

134 

135def test_allocate_32A_8H() -> None: 

136 """32 A100 + 8 H100.""" 

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

138 allocator = GreedyAllocator( 

139 workflow=DEFAULT_WORKFLOW_CONFIG, 

140 latency_data=latency_data, 

141 ) 

142 result = allocator.allocate( 

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

144 verbose=True, 

145 ) 

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

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

148 assert 0 < result.ttff_s < result.total_time_s < 10 * 60 * 60 

149 assert 0 < result.tbf_s < 1 

150 

151 

152def test_allocate_16A_80H() -> None: 

153 """16 A100 + 80 H100.""" 

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

155 allocator = GreedyAllocator( 

156 workflow=DEFAULT_WORKFLOW_CONFIG, 

157 latency_data=latency_data, 

158 ) 

159 result = allocator.allocate( 

160 num_gpus={GPUType.A100: 16, GPUType.H100: 80}, 

161 verbose=True, 

162 ) 

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

164 assert result.gpus_used[GPUType.H100] == 80 

165 assert 0 < result.ttff_s < result.total_time_s < 10 * 60 * 60 

166 assert 0 < result.tbf_s < 1 

167 

168 

169def test_allocate_0H() -> None: 

170 """No GPUs.""" 

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

172 allocator = GreedyAllocator( 

173 workflow=DEFAULT_WORKFLOW_CONFIG, 

174 latency_data=latency_data, 

175 ) 

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

177 allocator.allocate( 

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

179 verbose=True, 

180 ) 

181 

182 

183def test_32A_32H_options() -> None: 

184 """32 A100 + 32 H100 with all options enabled.""" 

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

186 allocator = GreedyAllocator( 

187 workflow=DEFAULT_WORKFLOW_CONFIG, 

188 latency_data=latency_data, 

189 ) 

190 result = allocator.allocate( 

191 num_gpus={ 

192 GPUType.A100: 32, 

193 GPUType.H100: 32, 

194 }, 

195 verbose=True, 

196 allow_removal=True, 

197 allow_merging=True, 

198 look_ahead_replicas=4, 

199 ) 

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

201 assert result.gpus_used[GPUType.H100] == 32 

202 assert 0 < result.ttff_s < result.total_time_s < 10 * 60 * 60 

203 assert 0 < result.tbf_s < 1 

204 

205 

206@pytest.mark.parametrize("workflow", WORKFLOWS.values()) 

207def test_workflows(workflow: WorkflowConfig) -> None: 

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

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

210 policy = STREAMWISE_POLICY 

211 if workflow.target_resolution != QualityLevel.HIGH: 

212 policy = replace(policy, use_upscaler=False) 

213 allocator = GreedyAllocator( 

214 workflow=workflow, 

215 latency_data=latency_data, 

216 power_data=power_data, 

217 policy=policy, 

218 ) 

219 result = allocator.allocate( 

220 num_gpus={ 

221 GPUType.A100: 16, 

222 GPUType.H100: 16 

223 }, 

224 verbose=True, 

225 ) 

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

227 assert result.gpus_used[GPUType.H100] == 16 

228 assert 0 < result.ttff_s < result.total_time_s < 24 * SECONDS_IN_HOUR 

229 assert 0 < result.tbf_s < 1