Coverage for src/lilbee/providers/fleet/placement.py: 100%
287 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""VRAM-aware placement planner for the multi-GPU llama-server fleet.
3Estimates each role-model's VRAM footprint from GGUF metadata and bin-packs
4instances across GPUs in ``placement_rank`` order, largest-first within a rank: a
5model that fits one GPU runs as a single pinned instance, small models co-locate
6on a GPU with spare VRAM, and a model too big for any single GPU is tensor-split
7across enough GPUs to fit. Roles that never run in the same phase (ingest OCR vs a
8query) share one swap group when they cannot all co-reside, so only the phase in
9use is charged. On GPUs the estimate advises but never refuses: a role that fits
10nowhere is placed tight (best-effort, with a warning). See docs/architecture.md.
11"""
13from __future__ import annotations
15import logging
16from collections.abc import Callable, Sequence
17from dataclasses import dataclass, field
18from functools import reduce
19from math import gcd
21from lilbee.providers.fleet.placement_spec import PlacementError, PlacementSpec, RolePlacement
22from lilbee.providers.fleet.vram import usable_vram_fraction
23from lilbee.providers.roles import ROLE_REGISTRY, WorkerRole
25log = logging.getLogger(__name__)
27# (role, per-device tensor-split ratio) -> the instance's per-device VRAM footprint
28# vector aligned to that ratio. A split is accepted only when every card's entry
29# fits its own headroom, and each card is charged its own entry, not the sum.
30PeakEstimator = Callable[[WorkerRole, tuple[int, ...]], tuple[int, ...]]
32# (per-device tensor-split ratio, chosen cards' snapshot free-VRAM bytes) -> the per-slot
33# context the launch would serve on that chat shard. Lets the planner widen a chat
34# split onto idle cards when a tighter shard would starve KV below the target.
35SplitCtxFitter = Callable[[tuple[int, ...], Sequence[int]], int]
38@dataclass(frozen=True)
39class ModelPlacementInput:
40 """A role's model, its estimated single-instance footprint, and replica count.
42 ``replicas`` > 1 requests N data-parallel instances (one per GPU) for the role,
43 each charged ``est_vram_bytes``; capped at runtime by the GPUs with room.
45 ``est_ram_bytes`` is what the same estimate puts in system memory, which is
46 non-zero only when something offloads. Placement divides GPUs and so does not
47 read it; admission does, because system memory is a bound too.
48 """
50 role: WorkerRole
51 est_vram_bytes: int
52 replicas: int = 1
53 est_ram_bytes: int = 0
56@dataclass(frozen=True)
57class InstancePlan:
58 """One planned llama-server instance.
60 ``devices`` >1 means the model is split across them; ``tensor_split`` is the
61 per-device proportion (free VRAM in GiB) so an unequal pair splits by capacity
62 rather than evenly. Empty for a single-device instance. ``replica`` is the
63 instance's index within its role's data-parallel pool (0 for a single server).
64 """
66 role: WorkerRole
67 devices: tuple[int, ...]
68 tensor_split: tuple[int, ...] = ()
69 replica: int = 0
72@dataclass(frozen=True)
73class Placement:
74 """Planner output: server instances, swap tenants, and best-effort placements.
76 ``unplaceable_roles`` get no server, so a call to them surfaces a
77 ``ProviderError`` (there is no in-process fallback); only the shared-memory
78 path produces them (an oversize load there OOM-livelocks the host). On GPUs
79 a role that fits nowhere is placed anyway and listed in ``tight_roles`` with
80 its estimated shortfall in bytes. ``co_tenants`` share one llama-swap group
81 and evict each other on demand, so only one is resident at a time; each runs
82 a single instance.
83 """
85 instances: tuple[InstancePlan, ...]
86 unplaceable_roles: tuple[WorkerRole, ...]
87 co_tenants: frozenset[WorkerRole] = frozenset()
88 tight_roles: dict[WorkerRole, int] = field(default_factory=dict)
91def plan_placement(
92 models: list[ModelPlacementInput],
93 devices: list[tuple[int, int]],
94 *,
95 estimate_peak: PeakEstimator,
96 unified_budget: int | None = None,
97 chat_ctx_fit: SplitCtxFitter | None = None,
98 chat_ctx_target: int = 0,
99 free_headroom: dict[int, int] | None = None,
100) -> Placement:
101 """Bin-pack *models* onto *devices* (``[(index, vram_bytes), ...]``).
103 Roles are charged in ``placement_rank`` order, largest-first within a rank, with
104 a 90% headroom per GPU. A model that fits one GPU takes a single instance; one
105 too big for any single GPU is tensor-split. Chat is charged last: when it fits
106 only if vision's VRAM is refunded, the two become ``co_tenants`` of one swap
107 group instead of either becoming unplaceable. A model that fits nowhere, even
108 alone beside the pinned tier, is still placed tight and reported in
109 ``tight_roles``.
111 A chat split widens past the fewest fitting cards when ``chat_ctx_fit`` shows a
112 tighter shard would starve its served context below ``chat_ctx_target``; the
113 fitter is sized against ``free_headroom`` (the plan snapshot's free VRAM per
114 device index; see ``planning.capture_plan_probe``). See
115 docs/architecture.md (Placement). Other splits keep the fewest-cards behavior.
117 A ``unified_budget`` (free system RAM, bytes) means every device this host has
118 shares the host's memory, or it has none: an integrated GPU, an Apple Silicon
119 Mac, a GPU-less box. Those go through the shared pool whether or not a device
120 enumerated, because the constraint is the same RAM either way, and only that
121 path can refuse a role. Bin-packing them per device instead reads one pool as
122 several and never refuses anything, so a role set that cannot fit is admitted,
123 loads, and swap-livelocks the machine, which is what the budget exists to
124 prevent. ``None`` means at least one device has memory of its own, and the
125 per-GPU packing below applies.
126 """
127 if unified_budget is not None:
128 return _place_shared_memory(models, unified_budget)
129 if not devices:
130 return _place_ungated(models)
131 usable = usable_vram_fraction()
132 remaining: dict[int, float] = {idx: vram * usable for idx, vram in devices}
134 # The persistent singles are every replicas<=1 role plus replica 0 of each
135 # replicated role. They are charged before the elastic replicas, and the chat
136 # model is charged last of all (``placement_rank``).
137 replicated = [m for m in models if m.replicas > 1]
138 persistent_singles = [m for m in models if m.replicas <= 1] + [
139 _persistent_single(m) for m in replicated
140 ]
142 def place(model: ModelPlacementInput) -> _Placed | None:
143 return _place_single(
144 model,
145 remaining,
146 estimate_peak,
147 chat_ctx_fit=chat_ctx_fit,
148 chat_ctx_target=chat_ctx_target,
149 free_headroom=free_headroom,
150 )
152 instances, co_tenants, tight = _place_persistent(persistent_singles, place, remaining)
153 instances.extend(_place_elastic(replicated, instances, remaining, co_tenants))
155 return Placement(
156 instances=tuple(instances),
157 unplaceable_roles=(),
158 co_tenants=co_tenants,
159 tight_roles=tight,
160 )
163def _place_ungated(models: list[ModelPlacementInput]) -> Placement:
164 """Every role as a single un-pinned instance: no GPU and no measured RAM budget."""
165 return Placement(
166 instances=tuple(
167 InstancePlan(role=m.role, devices=(), replica=r)
168 for m in models
169 for r in range(m.replicas)
170 ),
171 unplaceable_roles=(),
172 )
175def _place_persistent(
176 persistent_singles: list[ModelPlacementInput],
177 place: Callable[[ModelPlacementInput], _Placed | None],
178 remaining: dict[int, float],
179) -> tuple[list[InstancePlan], frozenset[WorkerRole], dict[WorkerRole, int]]:
180 """Charge every persistent single in ``placement_rank`` order.
182 A role that does not fit refunds every already-charged role it is phase-disjoint
183 from (never co-resident with) and retries; the roles that let it in become one
184 swap group. This restores the in-process pool's behavior, where a role loaded
185 only when its phase ran, so ingest never paid for the query models or vice versa.
186 A role that still fits nowhere is placed tight (:func:`_place_tight`) instead
187 of refused.
188 """
189 instances: list[InstancePlan] = []
190 charges: dict[WorkerRole, dict[int, float]] = {}
191 co_tenants: set[WorkerRole] = set()
192 tight: dict[WorkerRole, int] = {}
193 for model in sorted(persistent_singles, key=_shared_pool_order):
194 placed = place(model)
195 if placed is None:
196 placed, group = _place_beside_disjoint(model, place, remaining, charges)
197 co_tenants |= group
198 if placed is None:
199 placed, group, shortfall = _place_tight(model, remaining, charges)
200 co_tenants |= group
201 tight[model.role] = shortfall
202 instances.append(placed.plan)
203 charges[model.role] = placed.charges
204 return instances, frozenset(co_tenants), tight
207def _place_beside_disjoint(
208 model: ModelPlacementInput,
209 place: Callable[[ModelPlacementInput], _Placed | None],
210 remaining: dict[int, float],
211 charges: dict[WorkerRole, dict[int, float]],
212) -> tuple[_Placed | None, frozenset[WorkerRole]]:
213 """Retry *model* with every phase-disjoint charged role's VRAM refunded.
215 Phase-disjoint roles are never resident together, so they can share one swap
216 group: only the resident one is charged. On success *model* and the refunded
217 roles form the group and the refunded roles drop out of ``charges`` (their VRAM
218 is now the group's shared budget). On failure every refund is rolled back and
219 *model* is genuinely oversize.
220 """
221 refunds = {r: c for r, c in charges.items() if _phase_disjoint(r, model.role)}
222 if not refunds:
223 return None, frozenset()
224 for charged in refunds.values():
225 for idx, held in charged.items():
226 remaining[idx] += held
227 placed = place(model)
228 if placed is not None:
229 slot = _group_slot(placed.charges, refunds, remaining)
230 for role in refunds:
231 del charges[role]
232 return _Placed(plan=placed.plan, charges=slot), frozenset(refunds) | {model.role}
233 for charged in refunds.values():
234 for idx, held in charged.items():
235 remaining[idx] -= held
236 return None, frozenset()
239def _tight_device_group(needed: int, remaining: dict[int, float]) -> tuple[int, ...]:
240 """Cards to give a model that fits nowhere else.
242 The most-free card alone when that one card is enough, otherwise every card
243 with headroom, ordered most-free first so the split's main device is the
244 roomiest. Not a minimal subset: finding the smallest group that fits would
245 need an estimate per candidate group, and this path is reached only after
246 the estimating search has already failed.
247 """
248 by_room = sorted(remaining, key=lambda idx: remaining[idx], reverse=True)
249 if not by_room:
250 return ()
251 if remaining[by_room[0]] >= needed:
252 return (by_room[0],)
253 usable = [idx for idx in by_room if remaining[idx] > 0]
254 return tuple(usable or by_room[:1])
257def _place_tight(
258 model: ModelPlacementInput,
259 remaining: dict[int, float],
260 charges: dict[WorkerRole, dict[int, float]],
261) -> tuple[_Placed, frozenset[WorkerRole], int]:
262 """Place an oversize *model* best-effort instead of refusing it.
264 Refunds every phase-disjoint charged role (they become *model*'s swap group),
265 then gives *model* the widest set of cards that helps, and drains them so the
266 elastic tier places nothing there. Returns the estimated shortfall in bytes.
268 Widest rather than the single most-free card. Pinned to one card with the
269 others excluded by that pin, a model too big for it has nowhere to go and the
270 tight placement is only a slower refusal; given the group, llama-server can
271 split across them. One card is still used when one card is enough, since a
272 split it does not need costs it interconnect bandwidth.
273 """
274 refunds = {r: c for r, c in charges.items() if _phase_disjoint(r, model.role)}
275 for charged in refunds.values():
276 for idx, held in charged.items():
277 remaining[idx] += held
278 devices = _tight_device_group(model.est_vram_bytes, remaining)
279 claimed = {idx: remaining[idx] for idx in devices}
280 available = sum(claimed.values())
281 for idx in devices:
282 remaining[idx] = 0.0
283 slot = _group_slot(claimed, refunds, remaining)
284 for role in refunds:
285 del charges[role]
286 placed = _Placed(
287 plan=InstancePlan(role=model.role, devices=devices),
288 charges=slot,
289 )
290 group = frozenset(refunds) | {model.role} if refunds else frozenset()
291 return placed, group, int(model.est_vram_bytes - available)
294def _group_slot(
295 trigger: dict[int, float],
296 refunds: dict[WorkerRole, dict[int, float]],
297 remaining: dict[int, float],
298) -> dict[int, float]:
299 """Charge each device the swap group's largest member, not just the trigger.
301 Only one member is resident at a time, but every evicted member must be able to
302 swap back in, so the group's per-device slot is the max across members. The
303 shortfall beyond the trigger's own charge is deducted from *remaining* here so
304 the elastic replica tier cannot claim VRAM an evicted member needs to return to.
305 The returned slot is recorded as the trigger's charge, so a later trigger that
306 refunds this group reclaims the whole slot.
307 """
308 slot = dict(trigger)
309 for charged in refunds.values():
310 for idx, held in charged.items():
311 need = max(slot.get(idx, 0.0), held)
312 remaining[idx] -= need - slot.get(idx, 0.0)
313 slot[idx] = need
314 return slot
317def _phase_disjoint(a: WorkerRole, b: WorkerRole) -> bool:
318 """True when *a* and *b* share no run phase, so they are never co-resident."""
319 return ROLE_REGISTRY[a].phases.isdisjoint(ROLE_REGISTRY[b].phases)
322def _place_elastic(
323 replicated: list[ModelPlacementInput],
324 instances: list[InstancePlan],
325 remaining: dict[int, float],
326 co_tenants: frozenset[WorkerRole],
327) -> list[InstancePlan]:
328 """Fill the residual VRAM with replicas 1..N-1 of each placed, non-co-tenant role."""
329 placed_roles = {plan.role for plan in instances}
330 elastic: list[InstancePlan] = []
331 for model in replicated:
332 if model.role not in placed_roles or model.role in co_tenants:
333 continue # unplaceable, or a co-tenant capped to its single instance
334 elastic.extend(_place_replicas(model, remaining, start=1))
335 return elastic
338def _persistent_single(model: ModelPlacementInput) -> ModelPlacementInput:
339 """The replica-0 persistent instance of a replicated role, sized as one server."""
340 return ModelPlacementInput(role=model.role, est_vram_bytes=model.est_vram_bytes, replicas=1)
343@dataclass(frozen=True)
344class _Placed:
345 """One placed instance and the per-device VRAM it was charged."""
347 plan: InstancePlan
348 charges: dict[int, float]
351def _place_single(
352 model: ModelPlacementInput,
353 remaining: dict[int, float],
354 estimate_peak: PeakEstimator,
355 *,
356 chat_ctx_fit: SplitCtxFitter | None = None,
357 chat_ctx_target: int = 0,
358 free_headroom: dict[int, int] | None = None,
359) -> _Placed | None:
360 """Place one instance: a single GPU when it fits, else a tensor-split, else None."""
361 single = _best_single_device(model.est_vram_bytes, remaining)
362 if single is not None:
363 remaining[single] -= model.est_vram_bytes
364 return _Placed(
365 plan=InstancePlan(role=model.role, devices=(single,)),
366 charges={single: float(model.est_vram_bytes)},
367 )
368 return _place_split(
369 model,
370 remaining,
371 estimate_peak,
372 chat_ctx_fit=chat_ctx_fit,
373 chat_ctx_target=chat_ctx_target,
374 free_headroom=free_headroom,
375 )
378def _place_split(
379 model: ModelPlacementInput,
380 remaining: dict[int, float],
381 estimate_peak: PeakEstimator,
382 *,
383 chat_ctx_fit: SplitCtxFitter | None = None,
384 chat_ctx_target: int = 0,
385 free_headroom: dict[int, int] | None = None,
386) -> _Placed | None:
387 """Tensor-split across the most-free GPUs whose per-device share each fits.
389 Charges each chosen card its own entry from *estimate_peak*'s vector, so the
390 busiest card (which OOMs first) gates the split, not the summed pool. A chat
391 split widens past the fewest fitting cards via *chat_ctx_fit* (see
392 :func:`plan_placement`); every other split takes the fewest that fit.
394 Each card count is tried at several proportions (:func:`_split_ratio_candidates`),
395 because a footprint that does not scale with the shard can overflow the
396 smallest card at the proportional ratio and fit at a shifted one. The sweep
397 stops after ``_MAX_SPLIT_ESTIMATES`` estimator calls and says so, since each
398 call is a subprocess.
399 """
400 from lilbee.providers.base import ProviderError
402 by_free = sorted(remaining, key=lambda idx: remaining[idx], reverse=True)
403 best: tuple[int, list[int], tuple[int, ...], tuple[int, ...]] | None = None
404 spent = 0
405 for count in range(2, len(by_free) + 1):
406 chosen = by_free[:count]
407 for ratio in _split_ratio_candidates(chosen, remaining):
408 if spent >= _MAX_SPLIT_ESTIMATES:
409 log.info(
410 "Stopped looking for a tensor split for %s after %d estimates; "
411 "wider layouts and finer proportions were not tried.",
412 model.role.value,
413 spent,
414 )
415 return _best_or_none(best, model, remaining)
416 spent += 1
417 try:
418 per_device = estimate_peak(model.role, ratio)
419 except (ProviderError, OSError):
420 # An unsizable model cannot evaluate a split; the tight single-card
421 # path downstream still places it.
422 continue
423 if len(per_device) != count or not all(
424 peak <= remaining[idx] for idx, peak in zip(chosen, per_device, strict=True)
425 ):
426 continue
427 # Only chat is widened past the fewest fitting cards; everything else (and
428 # the no-fitter generic path) takes the first shard that fits.
429 if model.role is not WorkerRole.CHAT or chat_ctx_fit is None or free_headroom is None:
430 return _charge_split(model, chosen, ratio, per_device, remaining)
431 # The context fit bisects, and every probe is another gguf-parser
432 # run, so it costs far more than the estimate that preceded it and
433 # has to be charged against the same budget. Uncounted, a wide box
434 # ran roughly 190 subprocesses against a documented cap of 24, all
435 # while holding the cross-process build lock that every other lilbee
436 # start waits on for 90 seconds before failing.
437 spent += _CTX_FIT_ESTIMATE_COST
438 served = chat_ctx_fit(ratio, [free_headroom[idx] for idx in chosen])
439 if served >= chat_ctx_target:
440 return _charge_split(model, chosen, ratio, per_device, remaining)
441 if best is None or served > best[0]:
442 best = (served, chosen, ratio, per_device)
443 return _best_or_none(best, model, remaining)
446def _best_or_none(
447 best: tuple[int, list[int], tuple[int, ...], tuple[int, ...]] | None,
448 model: ModelPlacementInput,
449 remaining: dict[int, float],
450) -> _Placed | None:
451 """Charge the widest chat split found short of the target, if there was one."""
452 if best is not None:
453 _served, chosen, ratio, per_device = best
454 return _charge_split(model, chosen, ratio, per_device, remaining)
455 return None
458def _charge_split(
459 model: ModelPlacementInput,
460 chosen: list[int],
461 ratio: tuple[int, ...],
462 per_device: tuple[int, ...],
463 remaining: dict[int, float],
464) -> _Placed:
465 """Debit each chosen card its own per-device share and return the split plan."""
466 for idx, peak in zip(chosen, per_device, strict=True):
467 remaining[idx] -= peak
468 return _Placed(
469 plan=InstancePlan(role=model.role, devices=tuple(chosen), tensor_split=ratio),
470 charges={idx: float(peak) for idx, peak in zip(chosen, per_device, strict=True)},
471 )
474def _place_replicas(
475 model: ModelPlacementInput, remaining: dict[int, float], *, start: int = 0
476) -> list[InstancePlan]:
477 """Place the elastic replicas ``start..model.replicas-1``, one per distinct GPU
478 (most-free first).
480 Spreads for throughput: each replica lands on a card not yet hosting one of this
481 role's replicas, only co-locating a second round once every card has one. Stops
482 early when no card has room, so the pool shrinks to the residual VRAM. ``start``
483 skips the indices already placed as persistent singles (1 for the elastic batch).
484 """
485 plans: list[InstancePlan] = []
486 used: set[int] = set()
487 for replica in range(start, model.replicas):
488 candidates = [idx for idx, free in remaining.items() if free >= model.est_vram_bytes]
489 if not candidates:
490 break
491 fresh = [idx for idx in candidates if idx not in used]
492 pick = max(fresh or candidates, key=lambda idx: remaining[idx])
493 remaining[pick] -= model.est_vram_bytes
494 used.add(pick)
495 if len(used) == len(remaining):
496 used = set()
497 plans.append(InstancePlan(role=model.role, devices=(pick,), replica=replica))
498 return plans
501def _place_shared_memory(models: list[ModelPlacementInput], budget: int) -> Placement:
502 """Fit un-pinned roles into one shared RAM *budget*.
504 Roles pack in ``placement_rank`` order (the elastic chat model last). Replicas
505 run as N co-resident processes against the shared pool (no per-GPU spread without
506 GPUs). A role that does not fit refunds every already-charged phase-disjoint role
507 (never resident with it), which caps those roles to a single instance and makes
508 the set one swap group; a role that still fits nowhere is unplaceable.
509 """
510 remaining = budget
511 instances: list[InstancePlan] = []
512 unplaceable: list[WorkerRole] = []
513 charged: dict[WorkerRole, int] = {}
514 # A charged role's swap-back need: its single-instance footprint, or, for the
515 # role holding a swap group's budget, the whole group slot.
516 swap_need = {m.role: m.est_vram_bytes for m in models}
517 co_tenants: set[WorkerRole] = set()
518 for model in sorted(models, key=_shared_pool_order):
519 placed = 0
520 for _ in range(model.replicas):
521 if model.est_vram_bytes > remaining:
522 break
523 remaining -= model.est_vram_bytes
524 instances.append(InstancePlan(role=model.role, devices=(), replica=placed))
525 placed += 1
526 if placed:
527 charged[model.role] = placed * model.est_vram_bytes
528 continue
529 remaining, group = _shared_beside_disjoint(model, remaining, charged, swap_need, instances)
530 if group:
531 instances.append(InstancePlan(role=model.role, devices=()))
532 co_tenants |= group
533 else:
534 unplaceable.append(model.role)
535 return Placement(
536 instances=tuple(instances),
537 unplaceable_roles=tuple(unplaceable),
538 co_tenants=frozenset(co_tenants),
539 )
542def _shared_beside_disjoint(
543 model: ModelPlacementInput,
544 remaining: int,
545 charged: dict[WorkerRole, int],
546 swap_need: dict[WorkerRole, int],
547 instances: list[InstancePlan],
548) -> tuple[int, frozenset[WorkerRole]]:
549 """Refund the shared-pool roles phase-disjoint from *model* and cap them to one instance.
551 The refunded roles and *model* become one swap group: only one member is resident
552 at a time, but every member must be able to swap back in, so the group is charged
553 its largest member's footprint (the slot), recorded under *model*'s role so a
554 later trigger reclaims the whole slot. Returns the budget after reclaiming the
555 refunded VRAM and charging the slot, plus the group; when even the slot does not
556 fit the reclaimed budget the group is empty and the budget unchanged.
557 """
558 refunds = [r for r in list(charged) if _phase_disjoint(r, model.role)]
559 if not refunds:
560 return remaining, frozenset()
561 reclaim = sum(charged[r] for r in refunds)
562 slot = max(model.est_vram_bytes, *(swap_need[r] for r in refunds))
563 if slot > remaining + reclaim:
564 return remaining, frozenset()
565 for role in refunds:
566 del charged[role]
567 instances[:] = [plan for plan in instances if plan.role is not role]
568 instances.append(InstancePlan(role=role, devices=(), replica=0))
569 charged[model.role] = slot
570 swap_need[model.role] = slot
571 return remaining + reclaim - slot, frozenset(refunds) | {model.role}
574def _shared_pool_order(model: ModelPlacementInput) -> tuple[int, int, int]:
575 """Sort key: placement rank, then most-phases-first, then largest-first.
577 A role in more phases is co-resident with more of the fleet and can never make
578 room by swapping (nothing is phase-disjoint from it), so it charges before a
579 same-rank single-phase sibling: a large LLM reranker must not crowd out the
580 embedder that both ingest and query need.
581 """
582 info = ROLE_REGISTRY[model.role]
583 return (info.placement_rank, -len(info.phases), -model.est_vram_bytes)
586def _best_single_device(need: int, remaining: dict[int, float]) -> int | None:
587 """Index of the device with the most free VRAM that still fits *need*."""
588 candidates = [idx for idx, free in remaining.items() if free >= need]
589 if not candidates:
590 return None
591 return max(candidates, key=lambda idx: remaining[idx])
594def placement_from_spec(
595 spec: PlacementSpec,
596 active_roles: tuple[WorkerRole, ...],
597 device_capacity: dict[int, int],
598 *,
599 estimate_peak: PeakEstimator,
600) -> Placement:
601 """Build a Placement from a manual *spec*, charging each card and failing loud.
603 ``device_capacity`` is each card's total VRAM (not instantaneous free): the
604 plan defines the fleet's full intended residency, so charging it against live
605 free VRAM would double-count models already loaded. Every active role must
606 have an entry; every device must exist; each card must fit the sum of the
607 per-device peaks charged to it, within the cfg.usable_vram_fraction headroom.
608 """
609 usable = usable_vram_fraction()
610 remaining = {idx: total * usable for idx, total in device_capacity.items()}
611 instances: list[InstancePlan] = []
612 for role in active_roles:
613 rp = _required_entry(spec, role, device_capacity)
614 ratio = rp.tensor_split or _vram_proportional_split(rp.devices, remaining)
615 per_device = estimate_peak(role, ratio)
616 split = ratio if len(rp.devices) > 1 else ()
617 for replica in range(rp.replicas):
618 _charge_devices(role, rp.devices, per_device, remaining, device_capacity)
619 instances.append(
620 InstancePlan(
621 role=role, devices=tuple(rp.devices), tensor_split=split, replica=replica
622 )
623 )
624 return Placement(instances=tuple(instances), unplaceable_roles=())
627def _vram_proportional_split(
628 devices: Sequence[int], remaining: dict[int, float], *, divisor: int = 1
629) -> tuple[int, ...]:
630 """Tensor-split ratio proportional to each card's remaining usable VRAM.
632 *divisor* sets the resolution: 1 is whole GiB, 4 is quarter-GiB shares.
634 Each card's shard tracks its remaining VRAM (whole GiB, min 1), so a card
635 already carrying other roles takes a smaller share. This is the single source
636 of the proportion for both placement paths: the auto planner (:func:`_place_split`)
637 and a manual spec entry with no explicit ``tensor_split`` (:func:`placement_from_spec`).
638 Keeping it in one place is what guarantees a manually-applied layout is charged
639 the same way the planner charges the identical layout, instead of an even split
640 that would falsely reject a fit the planner itself serves.
641 """
642 return tuple(max(1, int(remaining[idx] * divisor / 1024**3)) for idx in devices)
645# How many proportions the sweep will try per card count, and how many estimator
646# calls the whole sweep may spend. The estimator shells out to gguf-parser, so an
647# unbounded ladder on a wide box turns a plan into a minute of subprocesses; the
648# cap is what keeps the search's cost linear in cards rather than in cards times
649# candidates.
650# Rungs on the ladder below. Not a cap applied to it: the ladder builds exactly
651# this many, and a slice pretending to enforce a bound it cannot reach would be
652# decoration. Named so the estimator memo can be sized against a whole plan.
653_MAX_RATIO_CANDIDATES = 3
654_MAX_SPLIT_ESTIMATES = 24
655# What one context fit costs in estimator runs. It bisects the servable window,
656# so it is not one call but a handful, and charging it as one understated the
657# sweep by roughly an order of magnitude.
658_CTX_FIT_ESTIMATE_COST = 8
659# Sub-GiB resolution for the shifted candidates. Whole GiB is coarse enough that
660# two cards 700 MiB apart quantize to the same share.
661_RATIO_QUANTUM_DIVISOR = 4
664def _split_ratio_candidates(
665 devices: Sequence[int], remaining: dict[int, float]
666) -> tuple[tuple[int, ...], ...]:
667 """Proportions worth trying for a split across *devices*, best-first.
669 The VRAM-proportional ratio leads, because it is right whenever the footprint
670 scales with the shard. It is not always right: KV, compute buffers and a
671 fixed per-device overhead do not scale, so the proportional shard can
672 overflow the smallest card while the group has room. The rest of the ladder
673 shifts load toward the roomiest card at finer resolution, which is the
674 direction that helps when it does not.
676 Deduplicated by proportion rather than by tuple, so equal cards cost one
677 estimate rather than three: (24, 24) and (96, 96) are the same split asked
678 twice, and each ask is a subprocess.
679 """
680 quantum = _vram_proportional_split(devices, remaining, divisor=_RATIO_QUANTUM_DIVISOR)
681 candidates = [
682 _vram_proportional_split(devices, remaining),
683 quantum,
684 _shifted_toward_roomiest(devices, remaining, quantum),
685 ]
686 seen: dict[tuple[int, ...], tuple[int, ...]] = {}
687 for candidate in candidates:
688 seen.setdefault(_normalized(candidate), candidate)
689 return tuple(seen.values())
692def _normalized(ratio: tuple[int, ...]) -> tuple[int, ...]:
693 """*ratio* in lowest terms, so the same proportion compares equal."""
694 divisor = reduce(gcd, ratio)
695 return tuple(part // divisor for part in ratio)
698def _shifted_toward_roomiest(
699 devices: Sequence[int], remaining: dict[int, float], base: tuple[int, ...]
700) -> tuple[int, ...]:
701 """*base* with a share moved from the tightest card to the roomiest.
703 A tenth of the tightest card's share, which is enough to clear a fixed
704 per-device overhead without distorting a proportion that was nearly right.
705 Cards with identical room have nowhere to shift to, so *base* is returned
706 unchanged rather than making one of them worse. A lone card takes the same
707 path, being trivially equal to itself.
708 """
709 order = sorted(range(len(devices)), key=lambda pos: remaining[devices[pos]])
710 tightest, roomiest = order[0], order[-1]
711 if remaining[devices[tightest]] == remaining[devices[roomiest]]:
712 return base
713 moved = max(1, base[tightest] // 10)
714 shifted = list(base)
715 shifted[tightest] = max(1, shifted[tightest] - moved)
716 shifted[roomiest] += moved
717 return tuple(shifted)
720def _required_entry(
721 spec: PlacementSpec, role: WorkerRole, device_capacity: dict[int, int]
722) -> RolePlacement:
723 """Return *role*'s placement entry, failing loud if absent or pinned off-hardware."""
724 rp = spec.roles.get(role)
725 if rp is None:
726 raise PlacementError(f"{role.value} has a model but no placement entry in placement")
727 for idx in rp.devices:
728 if idx not in device_capacity:
729 raise PlacementError(
730 f"{role.value} pinned to device {idx} but only "
731 f"{len(device_capacity)} GPU(s) detected"
732 )
733 return rp
736def _charge_devices(
737 role: WorkerRole,
738 devices: tuple[int, ...],
739 per_device: tuple[int, ...],
740 remaining: dict[int, float],
741 device_capacity: dict[int, int],
742) -> None:
743 """Subtract one instance's per-device peaks from *remaining*; fail loud if a card overflows.
745 An estimate that does not cover every pinned device is a PlacementError rather
746 than a zip mismatch: gguf-parser returns no per-device breakdown for some
747 models, and the auto planner skips such a candidate (see :func:`_place_split`),
748 so the manual path must refuse in the currency callers already handle.
749 """
750 if len(per_device) != len(devices):
751 raise PlacementError(
752 f"{role.value} is pinned to {len(devices)} device(s) but its memory estimate "
753 f"covers {len(per_device)}; clear the placement to place it automatically"
754 )
755 for idx, peak in zip(devices, per_device, strict=True):
756 if peak > remaining[idx]:
757 raise PlacementError(
758 f"{role.value} pinned to device {idx} needs {peak / 1024**3:.1f} GiB but "
759 f"device {idx} has {remaining[idx] / 1024**3:.1f} GiB usable "
760 f"({device_capacity[idx] / 1024**3:.1f} GiB total, 90% headroom)"
761 )
762 for idx, peak in zip(devices, per_device, strict=True):
763 remaining[idx] -= peak