Coverage for streamwise/model_provisioner/policies.py: 100%
23 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1from __future__ import annotations
3from sim_types import Objective
4from sim_types import Policy
5from sim_types import GPUType
6from sim_types import Model
7from sim_types import Solver
9from constants import GPU_RESERVED_COST
10from constants import GPU_SPOT_COST
13# Max devices for each model
14# the logic is to allocate devices to each model proportional to their max devices
15MAX_DEVICES = {
16 Model.GEMMA: 8,
17 Model.FLUX: 16,
18 Model.HF: 40,
19 Model.HF_VAE: 1,
20 Model.FT: 40,
21 Model.FT_VAE: 1,
22}
24# Max iterations for the optimization loop to prevent infinite loops in case of non-monotonic allocators or other issues
25MAX_ITERATIONS = 100
27# Set to True if we want to use up all GPUs if there's no further improvements in the greedy optimization loop
28USE_ALL_GPUS = True
30# Default StreamWise policy configuration
31# TODO: Add a meta policy that picks the best among disaggregation options for HF/FT
32STREAMWISE_POLICY = Policy(
33 name="streamwise",
34 gpu_cost=GPU_SPOT_COST,
35 objective=Objective.TTFF_COST,
36 disaggregation={
37 Model.HF: True,
38 Model.FT: False,
39 },
40 use_upscaler=True,
41 hardware=list(GPUType),
42)
44STREAMWISE_MILP_POLICY = Policy(
45 name="streamwise",
46 gpu_cost=GPU_SPOT_COST,
47 objective=Objective.TTFF_COST,
48 disaggregation={
49 Model.HF: True,
50 Model.FT: False,
51 },
52 use_upscaler=True,
53 hardware=list(GPUType),
54 solver=Solver.GUROBI,
55)
58"""
59HexGen policy configuration.
60"""
61HEXGEN_POLICY = Policy(
62 name="hexgen",
63 gpu_cost=GPU_RESERVED_COST,
64 objective=Objective.TTFF, # Does not account for cost
65 disaggregation={
66 Model.HF: True,
67 Model.FT: False,
68 }, # Dissagregation
69 use_upscaler=False,
70 hardware=[ # Multiple hardware
71 GPUType.A100,
72 GPUType.H100,
73 GPUType.H200,
74 GPUType.GB200,
75 ],
76 solver=Solver.HEXGEN,
77)
80"""
81Helix policy configuration.
82Reference: https://github.com/Thesys-lab/Helix-ASPLOS25
83Optimizes models one-by-one following MODEL_ORDER using MILP.
84"""
85HELIX_POLICY = Policy(
86 name="helix",
87 gpu_cost=GPU_RESERVED_COST,
88 objective=Objective.TTFF, # Does not account for cost
89 disaggregation={
90 Model.HF: True,
91 Model.FT: False,
92 },
93 use_upscaler=False,
94 hardware=list(GPUType),
95 solver=Solver.HELIX,
96)
99"""
100DDiT policy configuration.
101Reference: https://arxiv.org/html/2506.13497v1
102"""
103DDIT_POLICY = Policy(
104 name="ddit",
105 gpu_cost=GPU_RESERVED_COST,
106 objective=Objective.TTFF,
107 disaggregation={
108 Model.HF: True,
109 Model.FT: False,
110 },
111 use_upscaler=False,
112 hardware=list(GPUType),
113 solver=Solver.NAIVE,
114)
117STREAMWISE_ENERGY_POLICY = Policy(
118 name="streamwise energy",
119 gpu_cost=GPU_SPOT_COST,
120 objective=Objective.TIME_ENERGY,
121 disaggregation={
122 Model.HF: True,
123 Model.FT: False,
124 },
125 use_upscaler=True,
126 hardware=list(GPUType),
127)
129NAIVE_POLICY = Policy(
130 name="naive",
131 gpu_cost=GPU_RESERVED_COST,
132 objective=Objective.TTFF,
133 disaggregation={},
134 use_upscaler=False,
135 hardware=[GPUType.A100],
136 solver=Solver.NAIVE,
137)
140BASELINE_POLICIES = {
141 "naive": NAIVE_POLICY,
142 "naive disag": Policy(
143 "naive disag",
144 gpu_cost=GPU_RESERVED_COST,
145 objective=Objective.TTFF,
146 disaggregation={
147 Model.HF: True,
148 Model.FT: True,
149 },
150 use_upscaler=False,
151 hardware=[GPUType.A100],
152 solver=Solver.NAIVE,
153 ),
154 "naive upscaler": Policy(
155 "naive upscaler",
156 gpu_cost=GPU_RESERVED_COST,
157 objective=Objective.TTFF,
158 disaggregation={},
159 use_upscaler=True, # Changed to True
160 hardware=[GPUType.A100],
161 solver=Solver.NAIVE,
162 ),
163 "naive spot": Policy(
164 "naive spot",
165 gpu_cost=GPU_SPOT_COST, # Changed to SPOT_COST
166 objective=Objective.TTFF,
167 disaggregation={},
168 use_upscaler=False,
169 hardware=[GPUType.A100],
170 solver=Solver.NAIVE,
171 ),
172 "naive ttff*cost allocator": Policy(
173 "naive ttff*cost allocator",
174 GPU_RESERVED_COST,
175 objective=Objective.TTFF_COST, # Changed to TTFF_COST
176 disaggregation={},
177 use_upscaler=False,
178 hardware=[GPUType.A100],
179 solver=Solver.GREEDY,
180 ),
181 "naive hardware": Policy(
182 "naive hardware",
183 GPU_RESERVED_COST,
184 objective=Objective.TTFF,
185 disaggregation={},
186 use_upscaler=False,
187 hardware=list(GPUType), # Changed hardware
188 solver=Solver.NAIVE,
189 ),
190}
193STREAMWISE_POLICIES = {
194 "streamwise": STREAMWISE_POLICY,
195 "streamwise no disag": Policy(
196 name="streamwise no disag",
197 gpu_cost=GPU_SPOT_COST,
198 objective=Objective.TTFF_COST,
199 disaggregation={},
200 use_upscaler=True,
201 hardware=list(GPUType),
202 solver=Solver.GREEDY,
203 ),
204 "streamwise no upscaler": Policy(
205 name="streamwise no upscaler",
206 gpu_cost=GPU_SPOT_COST,
207 objective=Objective.TTFF_COST,
208 disaggregation={
209 Model.HF: True,
210 Model.FT: False,
211 },
212 use_upscaler=False,
213 hardware=list(GPUType),
214 solver=Solver.GREEDY,
215 ),
216 "streamwise no spot": Policy(
217 name="streamwise no spot",
218 gpu_cost=GPU_RESERVED_COST,
219 objective=Objective.TTFF_COST,
220 disaggregation={
221 Model.HF: True,
222 Model.FT: False,
223 },
224 use_upscaler=True,
225 hardware=list(GPUType),
226 solver=Solver.GREEDY,
227 ),
228 "streamwise naive allocator": Policy(
229 name="streamwise naive allocator",
230 gpu_cost=GPU_SPOT_COST,
231 objective=Objective.TTFF,
232 disaggregation={
233 Model.HF: True,
234 Model.FT: False,
235 },
236 use_upscaler=True,
237 hardware=list(GPUType),
238 solver=Solver.NAIVE,
239 ),
240 "streamwise A100": Policy(
241 name="streamwise single hardware",
242 gpu_cost=GPU_SPOT_COST,
243 objective=Objective.TTFF_COST,
244 disaggregation={
245 Model.HF: True,
246 Model.FT: False,
247 },
248 use_upscaler=True,
249 hardware=[GPUType.A100],
250 solver=Solver.NAIVE,
251 ),
252}