SolitaryThinker
db866dfb04
Sweep all 296 v2 modules: simplify verbose docstrings/comments and remove 376 dangling "(design_vN §X)" citations to the now-deleted design docs (v2/README.md is the source of truth). Comment/docstring-only — AST-verified code-identical; the CPU suite holds at 237 passed / 5 pre-existing. Also applies yapf + ruff --fix auto-fixes (import ordering, forward-ref annotation de-quoting under `from __future__ import annotations`; behavior- neutral, suite-confirmed) and adds the legitimate domain terms mot/clen/te to the codespell ignore-list. Remaining ruff (24) + mypy (68 no-untyped-call) findings are pre-existing v2 debt, untouched here.
125 lines
4.5 KiB
Python
125 lines
4.5 KiB
Python
"""Role / stage pools (port of sglang multimodal_gen disagg).
|
|
|
|
A ``RolePool`` is a single-node worker pool for one role (encoder / denoiser / decoder / …) holding
|
|
its OWN resident ModelInstance (a component subset), its OWN cache manager, an outbound connector,
|
|
and a capacity (max concurrent requests → capacity-aware dispatch). ``DeployConfig`` maps program
|
|
nodes onto pools. Pools are single-node; multi-node scale is *multiple pools* fronted by the fleet —
|
|
never a cross-node mesh inside one pool.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from v2.cache import CacheManager
|
|
from v2.card import load_card
|
|
from v2.transport import Connector, make_connector
|
|
|
|
|
|
@dataclass
|
|
class RolePoolSpec:
|
|
role: str # "encoder" | "denoiser" | "decoder" | ...
|
|
pool_id: str
|
|
components: tuple[str, ...] = () # components this pool is responsible for
|
|
capacity: int = 2 # max concurrent requests occupying this pool
|
|
replicas: int = 1
|
|
connector: str = "in_proc" # outbound transport backend
|
|
slo_class: str = "throughput" # "latency" | "throughput" | "cost"
|
|
|
|
|
|
@dataclass
|
|
class DeployConfig:
|
|
"""Maps a Program's nodes onto role pools (the vllm-omni deploy-YAML analog)."""
|
|
pools: list[RolePoolSpec] = field(default_factory=list)
|
|
placement: dict[str, str] = field(default_factory=dict) # node_id -> pool_id
|
|
|
|
def pool_id_for(self, node_id: str, default: str | None = None) -> str:
|
|
return self.placement.get(node_id, default or (self.pools[0].pool_id if self.pools else ""))
|
|
|
|
|
|
class RolePool:
|
|
|
|
def __init__(self, spec: RolePoolSpec, card: Any):
|
|
self.spec = spec
|
|
self.role = spec.role
|
|
self.pool_id = spec.pool_id
|
|
self.components = set(spec.components)
|
|
self.capacity = spec.capacity
|
|
self.instance = load_card(card, cache_manager=CacheManager.from_card(card), validate=False)
|
|
self.connector: Connector = make_connector(spec.connector, name=spec.pool_id)
|
|
self.in_flight = 0
|
|
self.served = 0
|
|
self.transfers_out = 0
|
|
|
|
def can_admit(self) -> bool:
|
|
return self.in_flight < self.capacity
|
|
|
|
def enter(self) -> None:
|
|
self.in_flight += 1
|
|
|
|
def leave(self) -> None:
|
|
if self.in_flight > 0:
|
|
self.in_flight -= 1
|
|
self.served += 1
|
|
|
|
@property
|
|
def utilization(self) -> float:
|
|
return self.in_flight / max(1, self.capacity)
|
|
|
|
def stats(self) -> dict[str, Any]:
|
|
return {
|
|
"role": self.role,
|
|
"in_flight": self.in_flight,
|
|
"capacity": self.capacity,
|
|
"served": self.served,
|
|
"transfers_out": self.connector.transfers,
|
|
"credits_available": self.connector.credits.available
|
|
}
|
|
|
|
|
|
class PoolSet:
|
|
"""The collection of role pools for one deployment + the node→pool placement."""
|
|
|
|
def __init__(self, deploy: DeployConfig, card: Any):
|
|
self.deploy = deploy
|
|
self.card = card
|
|
self.by_id: dict[str, RolePool] = {s.pool_id: RolePool(s, card) for s in deploy.pools}
|
|
if not self.by_id:
|
|
raise ValueError("DeployConfig declares no pools")
|
|
|
|
def pool_for(self, node_id: str) -> RolePool:
|
|
return self.by_id[self.deploy.pool_id_for(node_id)]
|
|
|
|
def warmup(self) -> dict[str, int]:
|
|
"""Instantiate each pool's components ahead of first request."""
|
|
built: dict[str, int] = {}
|
|
for pool in self.by_id.values():
|
|
for cid in pool.components:
|
|
if pool.instance.has_component(cid):
|
|
pool.instance.component(cid)
|
|
built[pool.pool_id] = len(pool.components)
|
|
return built
|
|
|
|
def stats(self) -> dict[str, Any]:
|
|
return {pid: p.stats() for pid, p in self.by_id.items()}
|
|
|
|
|
|
def wan_t2v_disaggregated() -> DeployConfig:
|
|
"""The canonical N:M:K split for Wan T2V (encoder → denoiser → decoder)."""
|
|
return DeployConfig(
|
|
pools=[
|
|
RolePoolSpec(role="encoder", pool_id="enc", components=("text_encoder", ), capacity=4, slo_class="latency"),
|
|
RolePoolSpec(role="denoiser",
|
|
pool_id="den",
|
|
components=("transformer", ),
|
|
capacity=1,
|
|
slo_class="throughput"), # jumbo video step stays batch-of-1
|
|
RolePoolSpec(role="decoder", pool_id="dec", components=("vae", ), capacity=2),
|
|
],
|
|
placement={
|
|
"text_encode": "enc",
|
|
"denoise": "den",
|
|
"vae_decode": "dec"
|
|
},
|
|
)
|