Coverage for tests/simulator/test_hexgen.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 from sim_types import GPUType 

13 from data_loading import load_latency_data 

14 from model_provisioner.hexgen import HexGenAllocator 

15 from model_provisioner.hexgen import _get_model_order 

16 from sim_types import MODEL_ORDER 

17 

18 

19def test_get_model_order() -> None: 

20 """Test that _get_model_order returns models sorted by MODEL_ORDER.""" 

21 order = _get_model_order(DEFAULT_WORKFLOW_CONFIG) 

22 assert len(order) > 0 

23 # Check ordering is consistent with MODEL_ORDER 

24 for i in range(len(order) - 1): 

25 assert MODEL_ORDER[order[i]] < MODEL_ORDER[order[i + 1]] 

26 # All models in order should be in the workflow 

27 for m in order: 

28 assert m in DEFAULT_WORKFLOW_CONFIG.models 

29 

30 

31def test_8A() -> None: 

32 """8 x A100 (single server).""" 

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

34 allocator = HexGenAllocator( 

35 workflow=DEFAULT_WORKFLOW_CONFIG, 

36 latency_data=latency_data, 

37 ) 

38 result = allocator.allocate( 

39 num_gpus={GPUType.A100: 8}, 

40 ) 

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

42 assert 0 < result.ttff_s < result.total_time_s 

43 assert 0 < result.tbf_s 

44 

45 

46def test_8H() -> None: 

47 """8 x H100 (single server).""" 

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

49 allocator = HexGenAllocator( 

50 workflow=DEFAULT_WORKFLOW_CONFIG, 

51 latency_data=latency_data, 

52 ) 

53 result = allocator.allocate( 

54 num_gpus={GPUType.H100: 8}, 

55 ) 

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

57 assert 0 < result.ttff_s < result.total_time_s 

58 assert 0 < result.tbf_s 

59 

60 

61def test_16A() -> None: 

62 """16 x A100.""" 

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

64 allocator = HexGenAllocator( 

65 workflow=DEFAULT_WORKFLOW_CONFIG, 

66 latency_data=latency_data, 

67 ) 

68 result = allocator.allocate( 

69 num_gpus={GPUType.A100: 16}, 

70 verbose=True, 

71 ) 

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

73 assert 0 < result.ttff_s < result.total_time_s < 24 * 60 * 60 

74 assert 0 < result.tbf_s < 5 

75 

76 

77def test_64A() -> None: 

78 """64 x A100.""" 

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

80 allocator = HexGenAllocator( 

81 workflow=DEFAULT_WORKFLOW_CONFIG, 

82 latency_data=latency_data, 

83 ) 

84 result = allocator.allocate( 

85 num_gpus={GPUType.A100: 64}, 

86 verbose=True, 

87 ) 

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

89 assert 0 < result.ttff_s < result.total_time_s < 24 * 60 * 60 

90 assert 0 < result.tbf_s < 5 

91 

92 

93def test_64H() -> None: 

94 """64 x H100.""" 

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

96 allocator = HexGenAllocator( 

97 workflow=DEFAULT_WORKFLOW_CONFIG, 

98 latency_data=latency_data, 

99 ) 

100 result = allocator.allocate( 

101 num_gpus={GPUType.H100: 64}, 

102 ) 

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

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

105 assert 0 < result.ttff_s < result.total_time_s < 24 * 60 * 60 

106 assert 0 < result.tbf_s < 3 

107 

108 

109def test_8A_8H() -> None: 

110 """8 x A100 + 8 x H100 (mixed GPU types).""" 

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

112 allocator = HexGenAllocator( 

113 workflow=DEFAULT_WORKFLOW_CONFIG, 

114 latency_data=latency_data, 

115 ) 

116 result = allocator.allocate( 

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

118 ) 

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

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

121 assert 0 < result.ttff_s < result.total_time_s 

122 assert 0 < result.tbf_s 

123 

124 

125def test_32A_32H() -> None: 

126 """32 x A100 + 32 x H100 (mixed, larger).""" 

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

128 allocator = HexGenAllocator( 

129 workflow=DEFAULT_WORKFLOW_CONFIG, 

130 latency_data=latency_data, 

131 ) 

132 result = allocator.allocate( 

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

134 verbose=True, 

135 ) 

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

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

138 assert 0 < result.ttff_s < result.total_time_s 

139 assert 0 < result.tbf_s < 3 

140 

141 

142def test_no_gpus_error() -> None: 

143 """No GPUs should raise an error.""" 

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

145 allocator = HexGenAllocator( 

146 workflow=DEFAULT_WORKFLOW_CONFIG, 

147 latency_data=latency_data, 

148 ) 

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

150 allocator.allocate( 

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

152 ) 

153 

154 

155def test_is_subclass_of_greedy() -> None: 

156 """HexGenAllocator should extend GreedyAllocator.""" 

157 from model_provisioner.greedy import GreedyAllocator 

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

159 allocator = HexGenAllocator( 

160 workflow=DEFAULT_WORKFLOW_CONFIG, 

161 latency_data=latency_data, 

162 ) 

163 assert isinstance(allocator, GreedyAllocator) 

164 

165 

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

167def test_single_gpu_type_parametrized(gpu_type: GPUType) -> None: 

168 """Test HexGen with various single GPU types at 16 GPUs.""" 

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

170 allocator = HexGenAllocator( 

171 workflow=DEFAULT_WORKFLOW_CONFIG, 

172 latency_data=latency_data, 

173 ) 

174 result = allocator.allocate( 

175 num_gpus={gpu_type: 16}, 

176 ) 

177 assert result.gpus_used[gpu_type] == 16 

178 assert 0 < result.ttff_s < result.total_time_s 

179 assert 0 < result.tbf_s 

180 

181 

182def test_produces_valid_result() -> None: 

183 """Verify the HexGen result has all required fields populated correctly.""" 

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

185 allocator = HexGenAllocator( 

186 workflow=DEFAULT_WORKFLOW_CONFIG, 

187 latency_data=latency_data, 

188 ) 

189 result = allocator.allocate( 

190 num_gpus={GPUType.A100: 32}, 

191 ) 

192 assert result.total_time_s > 0 

193 assert result.ttff_s > 0 

194 assert result.tbf_s >= 0 

195 assert result.cost >= 0 

196 assert result.total_energy >= 0 

197 assert result.models is not None 

198 assert len(result.models) > 0