Coverage for simulator/model_allocator.py: 100%

1""" 

2Defines the ModelAllocator abstract base class and its interface for model allocation strategies. 

3""" 

4 

5from __future__ import annotations 

6 

7from typing import Optional 

8 

9from abc import ABC 

10from abc import abstractmethod 

11 

12from sim_types import GPUType 

13from sim_types import Model 

14from sim_types import ModelAllocation 

15from sim_types import Policy 

16from sim_types import WorkflowConfig 

17from sim_types import LatencyData 

18from sim_types import PowerData 

19from sim_types import Result 

20 

21from models import FluxModelAllocation 

22from models import GemmaModelAllocation 

23from models import HFModelAllocation 

24from models import HFVAEModelAllocation 

25from models import FTModelAllocation 

26from models import FTVAEModelAllocation 

27from models import UpscalerModelAllocation 

28from models import OthersModelAllocation 

29 

30from model_provisioner.policies import NAIVE_POLICY 

31 

32 

33class ModelAllocator(ABC): 

34 """ 

35 Abstract base class for model allocators. 

36 """ 

37 

38 def __init__( 

39 self, 

40 workflow: WorkflowConfig, 

41 latency_data: LatencyData, 

42 power_data: Optional[PowerData] = None, 

43 policy: Policy = NAIVE_POLICY, 

44 ) -> None: 

45 self.workflow = workflow 

46 self.latency_data = latency_data 

47 self.power_data = power_data 

48 self.policy = policy 

49 

50 @abstractmethod 

51 def allocate( 

52 self, 

53 num_gpus: dict[GPUType, int], 

54 verbose: bool = False, 

55 ) -> Result: 

56 """Allocate models to GPUs and return the provisioning result.""" 

57 ... 

58 

59 def _init_single_server_models( 

60 self, 

61 gpu_type: GPUType, 

62 ) -> dict[GPUType, dict[Model, list[ModelAllocation]]]: 

63 """ 

64 Initialize model allocations for a single server (8 GPUs or fewer). 

65 Each model gets a single allocation entry. 

66 """ 

67 models: dict[GPUType, dict[Model, list[ModelAllocation]]] = { 

68 gpu_type: { 

69 Model.GEMMA: [ 

70 GemmaModelAllocation( 

71 gpu_type=gpu_type, 

72 devices=1, replicas=1) 

73 ], 

74 Model.FLUX: [ 

75 FluxModelAllocation( 

76 gpu_type=gpu_type, 

77 devices=1, replicas=1) 

78 ], 

79 Model.HF: [ 

80 HFModelAllocation( 

81 gpu_type=gpu_type, 

82 devices=1, replicas=2) 

83 ], 

84 Model.HF_VAE: [ 

85 HFVAEModelAllocation( 

86 gpu_type=gpu_type, 

87 devices=1, replicas=1) 

88 ], 

89 Model.FT: [ 

90 FTModelAllocation( 

91 gpu_type=gpu_type, 

92 devices=1, replicas=1) 

93 ], 

94 Model.FT_VAE: [ 

95 FTVAEModelAllocation( 

96 gpu_type=gpu_type, 

97 devices=1, replicas=1) 

98 ], 

99 Model.UPSCALER: [ 

100 UpscalerModelAllocation( 

101 gpu_type=gpu_type) 

102 ], 

103 Model.OTHERS: [ 

104 OthersModelAllocation( 

105 gpu_type=gpu_type, 

106 devices=1, replicas=1) # + 1 for Kokoro/YOLO 

107 ], 

108 }, 

109 } 

110 

111 if self.policy.use_upscaler: 

112 # HF -> UPSCALER 

113 models[gpu_type][Model.HF][0].replicas -= 1 

114 models[gpu_type][Model.UPSCALER][0].replicas += 1 

115 

116 if not self.policy.is_disaggregated(Model.HF): 

117 # HF_VAE -> HF 

118 models[gpu_type][Model.HF_VAE][0].replicas -= 1 

119 models[gpu_type][Model.HF][0].replicas += 1 

120 if not self.policy.is_disaggregated(Model.FT): 

121 # FT_VAE -> FT 

122 models[gpu_type][Model.FT_VAE][0].replicas -= 1 

123 models[gpu_type][Model.FT][0].replicas += 1 

124 

125 self._zero_out_unused_models(models) 

126 return models 

127 

128 def _init_single_device_models( 

129 self, 

130 gpu_type: GPUType, 

131 ) -> dict[GPUType, dict[Model, list[ModelAllocation]]]: 

132 """ 

133 Initialize model allocations for a single GPU type with >8 GPUs. 

134 Each model gets two allocation entries (active and inactive). 

135 """ 

136 models: dict[GPUType, dict[Model, list[ModelAllocation]]] = { 

137 gpu_type: { 

138 Model.GEMMA: [ 

139 GemmaModelAllocation( 

140 gpu_type=gpu_type, 

141 devices=1, replicas=1), 

142 GemmaModelAllocation( 

143 gpu_type=gpu_type), 

144 ], 

145 Model.FLUX: [ 

146 FluxModelAllocation( 

147 gpu_type=gpu_type, 

148 devices=1, replicas=1), 

149 FluxModelAllocation( 

150 gpu_type=gpu_type), 

151 ], 

152 Model.HF: [ 

153 HFModelAllocation( 

154 gpu_type=gpu_type, 

155 devices=1, replicas=1), 

156 HFModelAllocation( 

157 gpu_type=gpu_type), 

158 ], 

159 Model.HF_VAE: [ 

160 HFVAEModelAllocation( 

161 gpu_type=gpu_type, 

162 devices=1, replicas=1), 

163 HFVAEModelAllocation( 

164 gpu_type=gpu_type), 

165 ], 

166 Model.FT: [ 

167 FTModelAllocation( 

168 gpu_type=gpu_type, 

169 devices=2, replicas=1), 

170 FTModelAllocation( 

171 gpu_type=gpu_type), 

172 ], 

173 Model.FT_VAE: [ 

174 FTVAEModelAllocation( 

175 gpu_type=gpu_type, 

176 devices=1, replicas=1), 

177 FTVAEModelAllocation( 

178 gpu_type=gpu_type), 

179 ], 

180 Model.UPSCALER: [ 

181 UpscalerModelAllocation( 

182 gpu_type=gpu_type), 

183 UpscalerModelAllocation( 

184 gpu_type=gpu_type), 

185 ], 

186 Model.OTHERS: [ 

187 OthersModelAllocation( 

188 gpu_type=gpu_type, 

189 devices=1, replicas=1), 

190 OthersModelAllocation( 

191 gpu_type=gpu_type), 

192 ], 

193 }, 

194 } 

195 

196 if self.policy.use_upscaler: 

197 models[gpu_type][Model.UPSCALER][0].replicas = 1 

198 

199 if not self.policy.is_disaggregated(Model.HF): 

200 # HF_VAE -> HF 

201 models[gpu_type][Model.HF_VAE][0].replicas -= 1 

202 models[gpu_type][Model.HF][0].replicas += 1 

203 if not self.policy.is_disaggregated(Model.FT): 

204 # FT_VAE -> FT 

205 models[gpu_type][Model.FT_VAE][0].replicas -= 1 

206 models[gpu_type][Model.FT][0].replicas += 1 

207 

208 self._zero_out_unused_models(models) 

209 return models 

210 

211 def _init_both_devices_models( 

212 self, 

213 gpu_type1: GPUType, 

214 gpu_type2: GPUType, 

215 ) -> dict[GPUType, dict[Model, list[ModelAllocation]]]: 

216 """ 

217 Initialize model allocations for two GPU types. 

218 gpu_type1 gets GEMMA, FLUX, OTHERS; gpu_type2 gets HF, VAE, FT, UPSCALER. 

219 """ 

220 models: dict[GPUType, dict[Model, list[ModelAllocation]]] = { 

221 gpu_type1: { 

222 Model.GEMMA: [GemmaModelAllocation( 

223 gpu_type=gpu_type1, 

224 devices=1, replicas=1)], 

225 Model.FLUX: [FluxModelAllocation( 

226 gpu_type=gpu_type1, 

227 devices=1, replicas=1)], 

228 Model.HF: [], 

229 Model.HF_VAE: [], 

230 Model.FT: [], 

231 Model.FT_VAE: [], 

232 Model.UPSCALER: [], 

233 Model.OTHERS: [OthersModelAllocation( 

234 gpu_type=gpu_type1, 

235 devices=1, replicas=1)], # + 1 for Kokoro/YOLO 

236 }, 

237 gpu_type2: { 

238 Model.GEMMA: [], 

239 Model.FLUX: [], 

240 Model.HF: [HFModelAllocation( 

241 gpu_type=gpu_type2, 

242 devices=1, replicas=1)], 

243 Model.HF_VAE: [HFVAEModelAllocation( 

244 gpu_type=gpu_type2, 

245 devices=1, replicas=1)], 

246 Model.FT: [FTModelAllocation( 

247 gpu_type=gpu_type2, 

248 devices=2, replicas=1)], 

249 Model.FT_VAE: [FTVAEModelAllocation( 

250 gpu_type=gpu_type2, 

251 devices=1, replicas=1)], 

252 Model.UPSCALER: [UpscalerModelAllocation( 

253 gpu_type=gpu_type2)], 

254 Model.OTHERS: [], 

255 }, 

256 } 

257 

258 if not self.policy.is_disaggregated(Model.HF): 

259 # HF_VAE -> HF 

260 models[gpu_type2][Model.HF_VAE][0].replicas -= 1 

261 models[gpu_type2][Model.HF][0].replicas += 1 

262 if not self.policy.is_disaggregated(Model.FT): 

263 # FT_VAE -> FT 

264 models[gpu_type2][Model.FT_VAE][0].replicas -= 1 

265 models[gpu_type2][Model.FT][0].replicas += 1 

266 

267 if self.policy.use_upscaler: 

268 models[gpu_type2][Model.UPSCALER][0].replicas = 1 

269 

270 self._zero_out_unused_models(models) 

271 return models 

272 

273 def _zero_out_unused_models( 

274 self, 

275 models: dict[GPUType, dict[Model, list[ModelAllocation]]], 

276 ) -> None: 

277 """Zero out replicas for models not in the workflow.""" 

278 for gpu_type in models: 

279 for model in Model: 

280 if model not in self.workflow.models: 

281 for alloc in models[gpu_type][model]: 

282 alloc.replicas = 0