Coverage for simulator/actions.py: 78%
293 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Actions for scaling models for the greedy allocator.
3"""
5from __future__ import annotations
7import random
9from collections import Counter
11from copy import deepcopy
13from typing import Optional
15from constants import DEVICE_OPTIONS
16from constants import SINGLE_INSTANCE_MODELS
17from constants import SINGLE_DEVICE_MODELS
19from sim_types import Action
20from sim_types import ActionName
21from sim_types import Model
22from sim_types import ModelAllocation
23from sim_types import GPUType
24from sim_types import WorkflowConfig
25from sim_types import LatencyData
26from sim_types import PowerData
27from sim_types import Objective
28from sim_types import Policy
30from model_provisioner.policies import STREAMWISE_POLICY
32from models import get_model_allocation
34from evaluator import evaluate_model_allocation
35from evaluator import calc_used_gpus
38def _is_single_instance(
39 model_name: Model,
40 workflow: Optional[WorkflowConfig] = None,
41) -> bool:
42 """Check if a model is single-instance, considering workflow parallelism settings."""
43 if model_name not in SINGLE_INSTANCE_MODELS:
44 return False
45 if workflow is not None and workflow.is_parallelizable(model_name):
46 return False
47 return True
50def find_next_devices(
51 device_options: list[int],
52 num_devices: int,
53 num_replicas: int,
54 remaining_devices: int,
55 max_num_devices: Optional[int] = None,
56) -> Optional[int]:
57 """
58 Find the next device combination.
59 For example, with device options [2, 4, 8, 16, 40], current devices 8, 1 replica, we get 16.
60 """
61 if num_replicas == 0:
62 # means we haven't allocated any replicas yet so start from smallest device option
63 return device_options[0] if device_options[0] <= remaining_devices else None
65 for device_option in device_options:
66 # if device_option > num_devices and device_option <= remaining_devices + num_devices:
67 if (
68 device_option > num_devices
69 and (device_option - num_devices) * num_replicas <= remaining_devices
70 and (max_num_devices is None or device_option <= max_num_devices)
71 ):
72 return device_option
73 return None
76def choose_action(
77 actions: list[Action],
78 objective: Objective,
79 switch_objective: bool = False,
80) -> Optional[Action]:
81 """Schedule requests."""
82 if not actions:
83 return None
85 if objective == Objective.TIME_COST:
86 # return min(actions, key=lambda a: a.time)
87 return min(
88 actions,
89 key=lambda a: (
90 a.time_cost(),
91 a.time,
92 ),
93 )
94 if objective == Objective.TIME_COST:
95 return min(
96 actions,
97 key=lambda a: (
98 a.time_cost(),
99 a.time,
100 ),
101 )
102 if objective == Objective.TTFF_COST:
103 return min(
104 actions,
105 key=lambda a: (
106 a.ttff_cost(),
107 a.ttff,
108 ),
109 )
110 if objective == Objective.FIFO:
111 # return min(actions, key=lambda a: a.arrival_time_s)
112 return min(actions, key=lambda a: a.get_order())
113 if objective == Objective.TIME:
114 return min(actions, key=lambda a: a.time)
115 if objective == Objective.TTFF:
116 return min(actions, key=lambda a: a.ttff)
117 if objective == Objective.COST:
118 return min(actions, key=lambda a: a.cost)
119 if objective == Objective.ENERGY:
120 return min(actions, key=lambda a: a.energy)
121 if objective == Objective.TIME_ENERGY:
122 return min(actions, key=lambda a: a.time_energy())
123 if objective == Objective.ENERGY_COST:
124 return min(actions, key=lambda a: a.energy_cost())
125 if objective == Objective.RANDOM:
126 # randomly pick an improvement to simulate naive allocation
127 return random.choice(actions)
128 if objective == Objective.TTFF_THEN_TIME:
129 if switch_objective:
130 return min(actions, key=lambda a: a.time)
131 else:
132 return min(actions, key=lambda a: a.ttff)
133 if objective == Objective.NONE:
134 return None
135 raise ValueError(f"Cannot recognize objective {objective}")
138def apply_action(
139 action: Action,
140 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
141) -> dict[GPUType, dict[Model, list[ModelAllocation]]]:
142 """Apply the chosen action to the models and update remaining devices."""
144 for gpu_type in action.models.keys():
145 if gpu_type not in models:
146 raise ValueError(f"Cannot find gpu type {gpu_type} in {models.keys()}")
147 for model in action.models[gpu_type].keys():
148 if model not in models[gpu_type]:
149 raise ValueError(f"Cannot find model {model} in {models[gpu_type].keys()}")
150 allocs_to_remove = []
151 for alloc_id in range(len(action.models[gpu_type][model])):
152 # check if devices and replicas are non-negative
153 num_devices = action.models[gpu_type][model][alloc_id].devices
154 if num_devices < 0:
155 raise ValueError(f"Action devices {num_devices} must be >= 0")
156 if action.models[gpu_type][model][alloc_id].replicas <= 0:
157 # remove that instance if replicas is 0 or negative
158 allocs_to_remove.append(alloc_id)
159 for alloc_id in reversed(allocs_to_remove):
160 del action.models[gpu_type][model][alloc_id]
162 return action.models
165def gen_actions(
166 workflow: WorkflowConfig,
167 num_gpus: dict[GPUType, int],
168 latency_data: LatencyData,
169 power_data: Optional[PowerData] = None,
170 models: dict[GPUType, dict[Model, list[ModelAllocation]]] = {},
171 policy: Policy = STREAMWISE_POLICY,
172 allow_removal: bool = False,
173 allow_merging: bool = False,
174 look_ahead_replicas: int = 3,
175) -> list[Action]:
176 actions: list[Action] = []
178 # Extract GPU types from models
179 gpu_types = list(models.keys())
180 assert len(gpu_types) == len(num_gpus), \
181 f"Number of GPU types in models {len(gpu_types)} must match num_gpus {len(num_gpus)}"
183 remaining_gpus = {}
184 for gpu_type in num_gpus.keys():
185 remaining_gpus[gpu_type] = num_gpus[gpu_type] - calc_used_gpus({gpu_type: models[gpu_type]})
187 # Option 1: Provision more by increasing <devices, replicas> for each model allocation
188 for model in Model:
189 if model not in workflow.models:
190 continue
191 for gpu_type in gpu_types:
192 for alloc_id in range(len(models[gpu_type][model])):
193 actions.extend(_gen_add_device_replica_actions(
194 models=models,
195 num_gpus=num_gpus,
196 remaining_gpus=remaining_gpus[gpu_type],
197 gpu_type=gpu_type,
198 model_name=model,
199 allocation_id=alloc_id,
200 workflow=workflow,
201 policy=policy,
202 latency_data=latency_data,
203 power_data=power_data,
204 look_ahead_replicas=look_ahead_replicas,
205 ))
207 # Option 2: Add a model instance of <devices, replicas>
208 for model in Model:
209 if model not in workflow.models:
210 continue
211 for gpu_type in gpu_types:
212 actions.extend(_gen_add_instance(
213 models=models,
214 num_gpus=num_gpus,
215 remaining_gpus=remaining_gpus[gpu_type],
216 gpu_type=gpu_type,
217 model_name=model,
218 workflow=workflow,
219 policy=policy,
220 latency_data=latency_data,
221 power_data=power_data,
222 look_ahead_replicas=look_ahead_replicas,
223 ))
225 if allow_removal:
226 # Option 3: Remove replicas for each model allocation
227 for model in Model:
228 if model not in workflow.models:
229 continue
230 for gpu_type in gpu_types:
231 model_instances = models[gpu_type][model]
232 for alloc_id in range(len(model_instances)):
233 action = _gen_remove_replica_action(
234 models=models,
235 num_gpus=num_gpus,
236 gpu_type=gpu_type,
237 model_name=model,
238 allocation_id=alloc_id,
239 workflow=workflow,
240 policy=policy,
241 latency_data=latency_data,
242 power_data=power_data,
243 )
244 if action:
245 actions.append(action)
247 if allow_merging:
248 # Option 4: Merge across model allocations
249 for model in Model:
250 if model not in workflow.models:
251 continue
252 for gpu_type in gpu_types:
253 actions.extend(_gen_merge_replicas_actions(
254 models=models,
255 num_gpus=num_gpus,
256 gpu_type=gpu_type,
257 model_name=model,
258 workflow=workflow,
259 policy=policy,
260 latency_data=latency_data,
261 power_data=power_data,
262 ))
264 return actions
267def _get_min_device_combinations(
268 num_gpus: int,
269 model: Model,
270) -> list[tuple[int, int]]:
271 """
272 Get the minimum device combinations for a given number of GPUs and model.
273 [(device_count, num_replicas), ...]
274 For example, for 64, it would return [(40, 1), (16, 1)].
275 """
276 remaining = num_gpus
277 result: list[int] = []
278 for size in sorted(DEVICE_OPTIONS[model], reverse=True):
279 while remaining >= size:
280 result.append(size)
281 remaining -= size
282 if remaining > 0:
283 raise ValueError(f"Cannot exactly decompose {num_gpus} with DEVICE_OPTIONS")
284 counts = Counter(result)
285 return sorted(counts.items(), reverse=True) # Sort by device count descending
288def _get_large_instance_many_small_combinations(
289 num_gpus: int,
290 model: Model,
291) -> list[tuple[int, int]]:
292 """
293 Get the largest instance possible and then split the rest into 1 GPU instances.
294 For example, for 64, it would return [(40, 1), (1, 16)].
295 """
296 assert num_gpus > 0
297 assert model in DEVICE_OPTIONS
298 assert DEVICE_OPTIONS[model][0] == 1 # must have 1 GPU option to use this function
300 remaining_gpus = num_gpus
301 result: list[tuple[int, int]] = []
302 for size in sorted(DEVICE_OPTIONS[model], reverse=True):
303 if remaining_gpus >= size:
304 result = [(size, 1)]
305 remaining_gpus -= size
306 break
307 if remaining_gpus > 0:
308 result.append((1, remaining_gpus))
309 return result
312def _gen_add_device_replica_actions(
313 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
314 num_gpus: dict[GPUType, int],
315 remaining_gpus: int,
316 gpu_type: GPUType,
317 model_name: Model,
318 allocation_id: int,
319 workflow: WorkflowConfig,
320 policy: Policy,
321 latency_data: LatencyData,
322 power_data: Optional[PowerData] = None,
323 look_ahead_replicas: int = 3,
324) -> list[Action]:
325 """
326 Generate actions that explore all valid (replicas, devices) provisioning
327 options for a given model allocation, using the remaining GPUs.
329 From the current replicas * devices, find the next options by distributing the remaining devices.
330 For example, if currently 2 replicas at parallelism 4 with 4 remaining devices, options include:
331 - 3 replicas, 4 devices (uses 12 total, 4 more than current 8)
332 - 1 replica, 10 devices (uses 10 total, 2 more than current 8)
333 - etc.
334 """
335 actions: list[Action] = []
337 if model_name in SINGLE_DEVICE_MODELS and _is_single_instance(model_name, workflow):
338 return actions # No scaling possible
340 alloc = models[gpu_type][model_name][allocation_id]
341 current_total = alloc.devices * max(alloc.replicas, 0)
342 current_replicas = alloc.replicas
343 total_available = current_total + remaining_gpus
345 max_num_devices = latency_data[gpu_type].get_max_parallelism(model_name)
346 max_replicas = alloc.get_max_replicas(workflow)
347 is_single_instance = _is_single_instance(model_name, workflow)
348 is_single_device = model_name in SINGLE_DEVICE_MODELS
350 seen: set[tuple[int, int]] = set()
351 seen.add((max(alloc.replicas, 0), alloc.devices)) # skip current config
353 for new_devices in DEVICE_OPTIONS[model_name]:
354 if new_devices > max_num_devices:
355 continue # Exceeds max parallelism from latency data
356 if is_single_device and new_devices > 1:
357 continue # Model only supports single device
358 if (model_name, new_devices) not in latency_data[gpu_type]:
359 continue # No latency data for this device count
361 # Determine the range of replicas possible with this device count
362 if is_single_instance:
363 replica_candidates = [1]
364 else:
365 max_r = min(max_replicas, total_available // new_devices) if new_devices > 0 else 0
366 # limit max replicas to original replicas + X to avoid too many combinations
367 max_r = min(max_r, current_replicas + look_ahead_replicas)
368 replica_candidates = list(range(1, max_r + 1))
370 for new_replicas in replica_candidates:
371 new_total = new_replicas * new_devices
372 if new_total <= current_total:
373 continue # Must be an increase
374 if new_total > total_available:
375 continue # Not enough GPUs
376 if (new_replicas, new_devices) in seen:
377 continue
378 seen.add((new_replicas, new_devices))
380 try:
381 new_models = deepcopy(models)
382 new_models[gpu_type][model_name][allocation_id] = get_model_allocation(
383 model=model_name,
384 gpu_type=gpu_type,
385 devices=new_devices,
386 replicas=new_replicas,
387 )
388 action_result = evaluate_model_allocation(
389 models=new_models,
390 num_gpus=num_gpus,
391 workflow=workflow,
392 latency_data=latency_data,
393 power_data=power_data,
394 policy=policy,
395 include_models=[model_name],
396 )
397 actions.append(Action(
398 name=ActionName.ADD_DEVICE_REPLICA,
399 model=model_name,
400 gpu_type=gpu_type,
401 models=new_models,
402 action_result=action_result,
403 arrival_time_s=alloc.time,
404 ))
405 except Exception:
406 pass # Invalid configuration, skip
408 return actions
411def _gen_add_device_action(
412 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
413 num_gpus: dict[GPUType, int],
414 remaining_gpus: int,
415 gpu_type: GPUType,
416 model_name: Model,
417 allocation_id: int,
418 workflow: WorkflowConfig,
419 policy: Policy,
420 latency_data: LatencyData,
421 power_data: Optional[PowerData] = None,
422) -> Optional[Action]:
423 """
424 Action to add devices (increase parallelism) for a specific model allocation.
425 """
426 action: Optional[Action] = None
428 if model_name in SINGLE_DEVICE_MODELS:
429 return action # These models only run on a single GPU, so we don't add more devices
431 alloc = models[gpu_type][model_name][allocation_id]
433 max_num_devices = latency_data[gpu_type].get_max_parallelism(model_name)
434 next_num_devices = find_next_devices(
435 DEVICE_OPTIONS[model_name],
436 num_devices=alloc.devices,
437 num_replicas=alloc.replicas,
438 remaining_devices=remaining_gpus,
439 max_num_devices=max_num_devices)
441 if not next_num_devices:
442 return action # No valid next device option, skip
443 if (model_name, next_num_devices) not in latency_data[gpu_type]:
444 return action # No latency data for this device option, skip
446 new_models = deepcopy(models)
447 new_models[gpu_type][model_name][allocation_id] = get_model_allocation(
448 model=model_name,
449 gpu_type=gpu_type,
450 devices=next_num_devices,
451 replicas=max(1, alloc.replicas),
452 )
453 try:
454 action_result = evaluate_model_allocation(
455 models=new_models,
456 num_gpus=num_gpus,
457 workflow=workflow,
458 latency_data=latency_data,
459 power_data=power_data,
460 policy=policy,
461 include_models=[model_name],
462 )
463 action = Action(
464 name=ActionName.ADD_DEVICE,
465 model=model_name,
466 gpu_type=gpu_type,
467 models=new_models,
468 action_result=action_result,
469 arrival_time_s=alloc.time,
470 )
471 except Exception:
472 pass # Invalid action
474 return action
477def _gen_merge_replicas_actions(
478 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
479 gpu_type: GPUType,
480 model_name: Model,
481 num_gpus: dict[GPUType, int],
482 workflow: WorkflowConfig,
483 policy: Policy,
484 latency_data: LatencyData,
485 power_data: Optional[PowerData] = None,
486) -> list[Action]:
487 actions: list[Action] = []
489 if _is_single_instance(model_name, workflow):
490 return actions # These models only support a single instance, so no need to merge
492 model_instances = models[gpu_type][model_name]
493 model_num_gpus = 0
494 for model_instance in model_instances:
495 model_num_gpus += model_instance.get_num_gpus()
496 if model_num_gpus <= 1:
497 return actions # No replicas to merge for this model and GPU type
499 for device_combos in [
500 _get_min_device_combinations(model_num_gpus, model_name),
501 _get_large_instance_many_small_combinations(model_num_gpus, model_name)
502 ]:
503 new_models = deepcopy(models)
504 new_models[gpu_type][model_name] = []
506 for new_num_devices, new_num_replicas in device_combos:
507 new_models[gpu_type][model_name].append(get_model_allocation(
508 model=model_name,
509 gpu_type=gpu_type,
510 devices=new_num_devices,
511 replicas=new_num_replicas,
512 ))
514 try:
515 action_result = evaluate_model_allocation(
516 models=new_models,
517 num_gpus=num_gpus,
518 workflow=workflow,
519 latency_data=latency_data,
520 power_data=power_data,
521 policy=policy,
522 include_models=[model_name],
523 )
525 instance_id = 0
526 actions.append(Action(
527 name=ActionName.MERGE,
528 model=model_name,
529 gpu_type=gpu_type,
530 models=new_models,
531 action_result=action_result,
532 arrival_time_s=new_models[gpu_type][model_name][instance_id].time,
533 ))
534 except Exception:
535 pass # Invalid action
537 return actions
540def _gen_add_instance(
541 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
542 num_gpus: dict[GPUType, int],
543 remaining_gpus: int,
544 gpu_type: GPUType,
545 model_name: Model,
546 workflow: WorkflowConfig,
547 policy: Policy,
548 latency_data: LatencyData,
549 power_data: Optional[PowerData] = None,
550 look_ahead_replicas: int = 3,
551) -> list[Action]:
552 actions: list[Action] = []
554 if _is_single_instance(model_name, workflow):
555 return actions # These models only support a single instance, so we don't add more
557 for new_num_devices in DEVICE_OPTIONS[model_name]:
558 for new_num_replicas in list(range(1, look_ahead_replicas + 1)):
559 new_instance = get_model_allocation(
560 model=model_name,
561 gpu_type=gpu_type,
562 devices=new_num_devices,
563 replicas=new_num_replicas,
564 )
565 if new_instance.get_num_gpus() > remaining_gpus:
566 continue # Not enough remaining GPUs for this new instance
568 new_models = deepcopy(models)
569 new_models[gpu_type][model_name].append(new_instance)
571 try:
572 action_result = evaluate_model_allocation(
573 models=new_models,
574 num_gpus=num_gpus,
575 workflow=workflow,
576 latency_data=latency_data,
577 power_data=power_data,
578 policy=policy,
579 include_models=[model_name],
580 )
581 action = Action(
582 name=ActionName.ADD_INSTANCE,
583 model=model_name,
584 gpu_type=gpu_type,
585 models=new_models,
586 action_result=action_result,
587 arrival_time_s=new_instance.time,
588 )
589 actions.append(action)
590 except Exception:
591 pass # Invalid action
593 return actions
596def _gen_remove_replica_action(
597 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
598 num_gpus: dict[GPUType, int],
599 gpu_type: GPUType,
600 model_name: Model,
601 allocation_id: int,
602 workflow: WorkflowConfig,
603 policy: Policy,
604 latency_data: LatencyData,
605 power_data: Optional[PowerData] = None,
606) -> Optional[Action]:
607 action: Optional[Action] = None
609 model = models[gpu_type][model_name][allocation_id]
611 if model.replicas == 0:
612 return action # No replicas to remove for this model and GPU type
614 new_models = deepcopy(models)
615 new_models[gpu_type][model_name][allocation_id] = get_model_allocation(
616 model=model_name,
617 gpu_type=gpu_type,
618 devices=model.devices,
619 replicas=model.replicas - 1,
620 )
622 if len(num_gpus) == 2:
623 # For dual GPU setting, initialize removed replica on the other GPU type to see if it improves performance
624 gpu_types = list(num_gpus.keys())
625 other_gpu_type = gpu_types[0] if gpu_type == gpu_types[1] else gpu_types[1]
626 if _is_single_instance(model_name, workflow):
627 if new_models[gpu_type][model_name][allocation_id].replicas == 0:
628 # If this is a single instance model and we're removing the only replica, add it to the other GPU type
629 new_models[other_gpu_type][model_name].append(get_model_allocation(
630 model=model_name,
631 gpu_type=other_gpu_type,
632 devices=model.devices,
633 replicas=1,
634 ))
636 try:
637 action_result = evaluate_model_allocation(
638 models=new_models,
639 num_gpus=num_gpus,
640 workflow=workflow,
641 latency_data=latency_data,
642 power_data=power_data,
643 policy=policy,
644 include_models=[model_name],
645 )
646 action = Action(
647 name=ActionName.REMOVE_REPLICA,
648 model=model_name,
649 gpu_type=gpu_type,
650 models=new_models,
651 action_result=action_result,
652 arrival_time_s=new_models[gpu_type][model_name][allocation_id].time,
653 )
654 except Exception:
655 pass # Ignore not possible action
656 return action
659def _gen_add_replica_action(
660 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
661 num_gpus: dict[GPUType, int],
662 remaining_gpus: int,
663 gpu_type: GPUType,
664 model_name: Model,
665 allocation_id: int,
666 workflow: WorkflowConfig,
667 policy: Policy,
668 latency_data: LatencyData,
669 power_data: Optional[PowerData] = None,
670) -> Optional[Action]:
671 """
672 Action to add replicas for a specific model allocation.
673 """
674 action: Optional[Action] = None
676 if _is_single_instance(model_name, workflow):
677 return action # These models don't support replication, so we skip
679 model = models[gpu_type][model_name][allocation_id]
681 if remaining_gpus < model.devices:
682 return action # Not enough remaining GPUs to add another replica
684 max_replicas = model.get_max_replicas(workflow)
685 if model.replicas >= max_replicas:
686 return action # Already at max replicas, skip
688 new_num_replicas = min(
689 model.replicas + 1,
690 max_replicas, # - models[other_gpu_type][Model.HF].replicas
691 model.replicas + remaining_gpus // model.devices
692 )
693 if new_num_replicas == model.replicas:
694 return action # No changes, skip
696 new_models = deepcopy(models)
697 new_models[gpu_type][model_name][allocation_id] = get_model_allocation(
698 model=model_name,
699 gpu_type=gpu_type,
700 devices=model.devices,
701 replicas=new_num_replicas,
702 )
704 try:
705 action_result = evaluate_model_allocation(
706 models=new_models,
707 num_gpus=num_gpus,
708 workflow=workflow,
709 latency_data=latency_data,
710 power_data=power_data,
711 policy=policy,
712 include_models=[model_name],
713 )
714 action = Action(
715 name=ActionName.ADD_REPLICA,
716 model=model_name,
717 gpu_type=gpu_type,
718 models=new_models,
719 action_result=action_result,
720 arrival_time_s=model.time,
721 )
722 except Exception:
723 pass # Invalid action
725 return action
728def max_time(
729 models: dict[GPUType, dict[Model, list[ModelAllocation]]],
730 model_name: Model,
731) -> float:
732 values = []
733 for models_gpu in models.values():
734 if model_name in models_gpu:
735 for alloc in models_gpu[model_name]:
736 values.append(alloc.time)
737 return max(values)