Coverage for simulator/auto_model_allocator.py: 100%

1""" 

2Factory helpers for selecting the right model allocator implementation. 

3""" 

4 

5from __future__ import annotations 

6 

7import logging 

8 

9from dataclasses import replace 

10from typing import Optional 

11 

12from sim_types import Policy 

13from sim_types import WorkflowConfig 

14from sim_types import LatencyData 

15from sim_types import Model 

16from sim_types import PowerData 

17from sim_types import QualityLevel 

18from sim_types import Solver 

19from sim_types import GPUType 

20from sim_types import Result 

21 

22from model_provisioner.policies import STREAMWISE_POLICY 

23 

24from model_allocator import ModelAllocator 

25 

26 

27class AutoModelAllocator(ModelAllocator): 

28 """Allocator wrapper that routes to a concrete allocator by solver.""" 

29 

30 policy: Policy 

31 

32 def __init__( 

33 self, 

34 workflow: WorkflowConfig, 

35 latency_data: LatencyData, 

36 power_data: Optional[PowerData] = None, 

37 policy: Policy = STREAMWISE_POLICY, 

38 ) -> None: 

39 super().__init__( 

40 workflow=workflow, 

41 latency_data=latency_data, 

42 power_data=power_data, 

43 policy=policy, 

44 ) 

45 self._allocator = self._build_allocator() 

46 

47 def _build_allocator(self) -> ModelAllocator: 

48 """Create concrete allocator based on configured solver.""" 

49 if self.policy.solver == Solver.GREEDY: 

50 from model_provisioner.greedy import GreedyAllocator 

51 return GreedyAllocator( 

52 workflow=self.workflow, 

53 latency_data=self.latency_data, 

54 power_data=self.power_data, 

55 policy=self.policy, 

56 ) 

57 if self.policy.solver == Solver.NAIVE: 

58 from model_provisioner.naive_baseline import NaiveAllocator 

59 return NaiveAllocator( 

60 workflow=self.workflow, 

61 latency_data=self.latency_data, 

62 power_data=self.power_data, 

63 policy=self.policy, 

64 ) 

65 if self.policy.solver in {Solver.GUROBI, Solver.HIGHS}: 

66 from model_provisioner.milp import MILPAllocator 

67 return MILPAllocator( 

68 workflow=self.workflow, 

69 latency_data=self.latency_data, 

70 power_data=self.power_data, 

71 policy=self.policy, 

72 ) 

73 if self.policy.solver == Solver.HEXGEN: 

74 from model_provisioner.hexgen import HexGenAllocator 

75 return HexGenAllocator( 

76 workflow=self.workflow, 

77 latency_data=self.latency_data, 

78 power_data=self.power_data, 

79 policy=self.policy, 

80 ) 

81 if self.policy.solver == Solver.HELIX: 

82 from model_provisioner.helix import HelixAllocator 

83 return HelixAllocator( 

84 workflow=self.workflow, 

85 latency_data=self.latency_data, 

86 power_data=self.power_data, 

87 policy=self.policy, 

88 ) 

89 raise ValueError(f"Unsupported solver for allocator selection: {self.policy.solver}") 

90 

91 def allocate( 

92 self, 

93 num_gpus: dict[GPUType, int], 

94 verbose: bool = False, 

95 ) -> Result: 

96 if self.policy.use_upscaler and self.workflow.target_resolution == QualityLevel.LOW: 

97 logging.warning( 

98 f"Policy {self.policy.name} uses upscaler, but workflow target resolution is LOW. " 

99 f"Disabling upscaler for this allocation.") 

100 self.policy = replace(self.policy, use_upscaler=False) 

101 self._allocator.policy = self.policy 

102 # Remove upscaler from model work 

103 self.workflow.model_work.pop(Model.UPSCALER, None) 

104 self._allocator.workflow = self.workflow 

105 

106 return self._allocator.allocate( 

107 num_gpus=num_gpus, 

108 verbose=verbose, 

109 )