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

1""" 

2Tests for the Helix allocator. 

3 

4Helix optimizes models one-by-one following MODEL_ORDER using per-model MILP. 

5""" 

6 

7import sys 

8import os 

9 

10# Add current path 

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

12 

13from tests.test_utils import temp_sys_path 

14 

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

16 from constants import DEFAULT_WORKFLOW_CONFIG 

17 from sim_types import GPUType 

18 from sim_types import Model 

19 from sim_types import MODEL_ORDER 

20 from sim_types import Solver 

21 from data_loading import load_latency_data 

22 from data_loading import load_power_data 

23 from model_provisioner.helix import HelixAllocator 

24 from model_provisioner.policies import HELIX_POLICY 

25 

26 

27def test_get_model_order() -> None: 

28 """Test that get_model_order() returns models sorted by MODEL_ORDER.""" 

29 order = DEFAULT_WORKFLOW_CONFIG.get_model_order() 

30 assert len(order) > 0 

31 # Check ordering is consistent with MODEL_ORDER 

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

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

34 # All models in order should be in the workflow 

35 for m in order: 

36 assert m in DEFAULT_WORKFLOW_CONFIG.models 

37 

38 

39def test_helix_policy_solver() -> None: 

40 """Verify HELIX_POLICY uses Solver.HELIX.""" 

41 assert HELIX_POLICY.solver == Solver.HELIX 

42 

43 

44def test_produces_valid_result() -> None: 

45 """Verify the Helix result has all required fields populated correctly.""" 

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

47 allocator = HelixAllocator( 

48 workflow=DEFAULT_WORKFLOW_CONFIG, 

49 latency_data=latency_data, 

50 ) 

51 result = allocator.allocate( 

52 num_gpus={GPUType.A100: 32}, 

53 milp_solver=Solver.HIGHS, 

54 ) 

55 assert result.total_time_s > 0 

56 assert result.ttff_s > 0 

57 assert result.tbf_s >= 0 

58 assert result.cost >= 0 

59 assert result.total_energy >= 0 

60 assert result.models is not None 

61 assert len(result.models) > 0 

62 

63 

64def test_produces_valid_result_verbose() -> None: 

65 """Verify verbose mode works without errors.""" 

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

67 allocator = HelixAllocator( 

68 workflow=DEFAULT_WORKFLOW_CONFIG, 

69 latency_data=latency_data, 

70 ) 

71 result = allocator.allocate( 

72 num_gpus={GPUType.A100: 32}, 

73 verbose=True, 

74 milp_solver=Solver.HIGHS, 

75 ) 

76 assert result.total_time_s > 0 

77 assert result.models is not None 

78 

79 

80def test_gpu_budget_respected() -> None: 

81 """Verify total GPUs used does not exceed the budget.""" 

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

83 allocator = HelixAllocator( 

84 workflow=DEFAULT_WORKFLOW_CONFIG, 

85 latency_data=latency_data, 

86 ) 

87 budget = {GPUType.A100: 16} 

88 result = allocator.allocate( 

89 num_gpus=budget, 

90 milp_solver=Solver.HIGHS, 

91 ) 

92 for gpu_type, used in result.gpus_used.items(): 

93 assert used <= budget.get(gpu_type, 0), \ 

94 f"{gpu_type.value}: used {used} > budget {budget.get(gpu_type, 0)}" 

95 

96 

97def test_all_workflow_models_allocated() -> None: 

98 """Verify models are allocated following MODEL_ORDER until GPUs are exhausted. 

99 

100 Helix allocates sequentially, so later models may be skipped if earlier 

101 models consume the whole budget. With a large GPU pool the first several 

102 models should still be present. 

103 """ 

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

105 allocator = HelixAllocator( 

106 workflow=DEFAULT_WORKFLOW_CONFIG, 

107 latency_data=latency_data, 

108 ) 

109 result = allocator.allocate( 

110 num_gpus={GPUType.A100: 64}, 

111 milp_solver=Solver.HIGHS, 

112 ) 

113 # Collect all models with non-empty allocations 

114 allocated_models: set[Model] = set() 

115 for gpu_type, model_dict in result.models.items(): 

116 for model, allocs in model_dict.items(): 

117 if allocs: 

118 allocated_models.add(model) 

119 

120 # At least the first models in MODEL_ORDER (GEMMA, FLUX, OTHERS) should 

121 # always be allocated since they require few GPUs. 

122 model_order = DEFAULT_WORKFLOW_CONFIG.get_model_order() 

123 for model in model_order[:3]: 

124 if DEFAULT_WORKFLOW_CONFIG.model_work.get(model, 0) > 0: 

125 assert model in allocated_models, \ 

126 f"Model {model.value} has work but was not allocated" 

127 

128 # At least one model should be allocated 

129 assert len(allocated_models) >= 1 

130 

131 

132def test_with_power_data() -> None: 

133 """Verify allocation works with power data provided.""" 

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

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

136 allocator = HelixAllocator( 

137 workflow=DEFAULT_WORKFLOW_CONFIG, 

138 latency_data=latency_data, 

139 power_data=power_data, 

140 ) 

141 result = allocator.allocate( 

142 num_gpus={GPUType.A100: 32}, 

143 milp_solver=Solver.HIGHS, 

144 ) 

145 assert result.total_time_s > 0 

146 assert result.total_energy >= 0 

147 assert result.cost >= 0 

148 

149 

150def test_custom_per_model_time_limit() -> None: 

151 """Verify custom per-model time limit is accepted.""" 

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

153 allocator = HelixAllocator( 

154 workflow=DEFAULT_WORKFLOW_CONFIG, 

155 latency_data=latency_data, 

156 ) 

157 result = allocator.allocate( 

158 num_gpus={GPUType.A100: 16}, 

159 per_model_time_limit=10, 

160 milp_solver=Solver.HIGHS, 

161 ) 

162 assert result.total_time_s > 0