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

1""" 

2Tests for simulator/auto_model_allocator.py. 

3 

4Covers: 

5- Routing to each concrete allocator by solver type. 

6- Upscaler-disabled warning when target_resolution is LOW. 

7- ValueError when an unsupported solver is requested. 

8""" 

9 

10from __future__ import annotations 

11 

12import sys 

13import os 

14import logging 

15 

16import pytest 

17from pytest import LogCaptureFixture 

18from dataclasses import replace 

19from unittest.mock import patch as _patch 

20 

21# Add current path 

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

23 

24from tests.test_utils import temp_sys_path 

25 

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

27 from sim_types import GPUType 

28 from sim_types import Model 

29 from sim_types import QualityLevel 

30 from sim_types import Solver 

31 

32 from constants import DEFAULT_WORKFLOW_CONFIG 

33 

34 from data_loading import load_latency_data 

35 

36 from model_provisioner.policies import STREAMWISE_POLICY 

37 from model_provisioner.policies import NAIVE_POLICY 

38 from model_provisioner.policies import HEXGEN_POLICY 

39 from model_provisioner.policies import HELIX_POLICY 

40 

41 from auto_model_allocator import AutoModelAllocator 

42 

43 from model_provisioner.greedy import GreedyAllocator 

44 from model_provisioner.naive_baseline import NaiveAllocator 

45 from model_provisioner.hexgen import HexGenAllocator 

46 from model_provisioner.helix import HelixAllocator 

47 from model_provisioner.milp import MILPAllocator 

48 

49 from workflows import PODCAST_WORKFLOW 

50 

51 

52# --------------------------------------------------------------------------- 

53# Solver routing 

54# --------------------------------------------------------------------------- 

55 

56def test_greedy_solver_routes_to_greedy_allocator() -> None: 

57 """AutoModelAllocator uses GreedyAllocator when solver=GREEDY.""" 

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

59 allocator = AutoModelAllocator( 

60 workflow=DEFAULT_WORKFLOW_CONFIG, 

61 latency_data=latency_data, 

62 policy=replace(STREAMWISE_POLICY, solver=Solver.GREEDY), 

63 ) 

64 assert isinstance(allocator._allocator, GreedyAllocator) 

65 

66 

67def test_naive_solver_routes_to_naive_allocator() -> None: 

68 """AutoModelAllocator uses NaiveAllocator when solver=NAIVE.""" 

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

70 allocator = AutoModelAllocator( 

71 workflow=DEFAULT_WORKFLOW_CONFIG, 

72 latency_data=latency_data, 

73 policy=replace(NAIVE_POLICY, solver=Solver.NAIVE), 

74 ) 

75 assert isinstance(allocator._allocator, NaiveAllocator) 

76 

77 

78def test_hexgen_solver_routes_to_hexgen_allocator() -> None: 

79 """AutoModelAllocator uses HexGenAllocator when solver=HEXGEN.""" 

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

81 policy = replace(HEXGEN_POLICY, solver=Solver.HEXGEN) 

82 allocator = AutoModelAllocator( 

83 workflow=DEFAULT_WORKFLOW_CONFIG, 

84 latency_data=latency_data, 

85 policy=policy, 

86 ) 

87 assert isinstance(allocator._allocator, HexGenAllocator) 

88 

89 

90def test_helix_solver_routes_to_helix_allocator() -> None: 

91 """AutoModelAllocator uses HelixAllocator when solver=HELIX.""" 

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

93 policy = replace(HELIX_POLICY, solver=Solver.HELIX) 

94 allocator = AutoModelAllocator( 

95 workflow=DEFAULT_WORKFLOW_CONFIG, 

96 latency_data=latency_data, 

97 policy=policy, 

98 ) 

99 assert isinstance(allocator._allocator, HelixAllocator) 

100 

101 

102def test_highs_solver_routes_to_milp_allocator() -> None: 

103 """AutoModelAllocator uses MILPAllocator when solver=HIGHS.""" 

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

105 policy = replace(STREAMWISE_POLICY, solver=Solver.HIGHS) 

106 allocator = AutoModelAllocator( 

107 workflow=DEFAULT_WORKFLOW_CONFIG, 

108 latency_data=latency_data, 

109 policy=policy, 

110 ) 

111 assert isinstance(allocator._allocator, MILPAllocator) 

112 

113 

114def test_gurobi_solver_routes_to_milp_allocator() -> None: 

115 """AutoModelAllocator uses MILPAllocator when solver=GUROBI.""" 

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

117 policy = replace(STREAMWISE_POLICY, solver=Solver.GUROBI) 

118 allocator = AutoModelAllocator( 

119 workflow=DEFAULT_WORKFLOW_CONFIG, 

120 latency_data=latency_data, 

121 policy=policy, 

122 ) 

123 assert isinstance(allocator._allocator, MILPAllocator) 

124 

125 

126def test_unsupported_solver_raises() -> None: 

127 """Building AutoModelAllocator with an unrecognised solver raises ValueError.""" 

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

129 policy = replace(STREAMWISE_POLICY) 

130 

131 allocator = AutoModelAllocator( 

132 workflow=DEFAULT_WORKFLOW_CONFIG, 

133 latency_data=latency_data, 

134 policy=policy, 

135 ) 

136 

137 # Patch the allocator's policy.solver to an unsupported sentinel value so 

138 # that _build_allocator falls through all known solver branches. 

139 bad_solver = object() 

140 with _patch.object(allocator.policy, 'solver', new=bad_solver, create=True): 

141 with pytest.raises(ValueError): 

142 allocator._build_allocator() 

143 

144 

145# --------------------------------------------------------------------------- 

146# Upscaler auto-disable when target_resolution is LOW 

147# --------------------------------------------------------------------------- 

148 

149def test_upscaler_disabled_for_low_resolution(caplog: LogCaptureFixture) -> None: 

150 """ 

151 When use_upscaler=True but target_resolution=LOW, allocate() should log a 

152 warning, disable the upscaler flag, and still return a valid Result. 

153 """ 

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

155 

156 # Build a workflow with LOW target resolution (no upscaler work). 

157 # Use a fresh copy of model_work to avoid mutating the global PODCAST_WORKFLOW. 

158 low_workflow = replace( 

159 PODCAST_WORKFLOW, 

160 target_resolution=QualityLevel.LOW, 

161 model_work=dict(PODCAST_WORKFLOW.model_work), 

162 ) 

163 # __post_init__ strips UPSCALER from model_work when resolution is not HIGH. 

164 assert Model.UPSCALER not in low_workflow.model_work 

165 

166 # Policy says use_upscaler=True but workflow cannot support it. 

167 policy = replace(STREAMWISE_POLICY, use_upscaler=True) 

168 allocator = AutoModelAllocator( 

169 workflow=low_workflow, 

170 latency_data=latency_data, 

171 policy=policy, 

172 ) 

173 

174 with caplog.at_level(logging.WARNING, logger="root"): 

175 result = allocator.allocate(num_gpus={GPUType.A100: 8}) 

176 

177 # The upscaler flag should have been cleared. 

178 assert allocator.policy.use_upscaler is False 

179 # The allocation should still succeed. 

180 assert result is not None 

181 assert result.total_time_s > 0.0 

182 

183 

184# --------------------------------------------------------------------------- 

185# Basic end-to-end allocation through AutoModelAllocator 

186# --------------------------------------------------------------------------- 

187 

188def test_greedy_allocation_produces_valid_result() -> None: 

189 """AutoModelAllocator with GREEDY solver produces a valid Result.""" 

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

191 allocator = AutoModelAllocator( 

192 workflow=DEFAULT_WORKFLOW_CONFIG, 

193 latency_data=latency_data, 

194 policy=replace(STREAMWISE_POLICY, solver=Solver.GREEDY), 

195 ) 

196 result = allocator.allocate(num_gpus={GPUType.A100: 8}) 

197 assert result.total_time_s > 0.0 

198 assert result.ttff_s > 0.0 

199 assert result.cost > 0.0 

200 assert result.gpus_used.get(GPUType.A100, 0) > 0