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.
127 lines
6.3 KiB
Python
127 lines
6.3 KiB
Python
"""ModelInstance — a resident, loaded card: component instances, model state, caches,
|
|
compiled graphs, and a parallel plan. One request may run several of the card's loops
|
|
against one ``ModelInstance``.
|
|
|
|
This is what makes omni native: when two loops bind the same component
|
|
(``shared_weight_components``), ``component()`` returns the *exact same* live object —
|
|
no duplicated weights, shared live state. See v2/README.md.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from v2.platform import Platform
|
|
from v2.card.specs import ModelCard
|
|
|
|
|
|
class ModelInstance:
|
|
"""The live, resident form of a ModelCard. Also serves as the ``ModelState``
|
|
passed to ``Loop.init`` (the loop reads components through it)."""
|
|
|
|
def __init__(self,
|
|
card: ModelCard,
|
|
parallel_plan: Any = None,
|
|
cache_manager: Any = None,
|
|
weights_version: str = "v0",
|
|
platform: Any = None):
|
|
self.card = card
|
|
self.parallel_plan = parallel_plan
|
|
self.caches = cache_manager
|
|
self.weights_version = weights_version
|
|
# The detected (device, arch). Resolves component/kernel implementations through the two
|
|
# backend registries; defaults to CPU/numpy. Swapping this to a GPU platform changes only
|
|
# the backend, not the loops/policies/training.
|
|
self.platform = platform if platform is not None else Platform.cpu()
|
|
# Piecewise CUDA-graph cache for loops declaring graph_capture="breakable_cudagraph". Lazily
|
|
# created by the runtime (kept as a plain attr so card/ stays free of any runtime import);
|
|
# a captured graph is tied to this instance's resident weights.
|
|
self.graphs: Any = None
|
|
self.adapter_versions: dict[str, str] = {}
|
|
# Per-component weight versions: a component's version changes only when IT is synced,
|
|
# so a transformer-only RL sync never invalidates the frozen text-encoder's feature cache.
|
|
self.component_versions: dict[str, str] = {cid: weights_version for cid in card.components}
|
|
self._components: dict[str, Any] = {}
|
|
self._loops: dict[str, Any] = {}
|
|
self._asleep: set[str] = set()
|
|
|
|
# --- components: shared by reference (the MoT requirement) ---------------- #
|
|
def component(self, component_id: str) -> Any:
|
|
if component_id in self._asleep:
|
|
raise RuntimeError(f"component {component_id!r} is asleep; wake it first")
|
|
if component_id not in self._components:
|
|
spec = self.card.components.get(component_id)
|
|
if spec is None:
|
|
raise KeyError(f"component {component_id!r} not declared on card {self.card.model_id!r}")
|
|
# The single materialization seam: the platform resolves (kind, device, variant) through
|
|
# the COMPONENTS registry, falling back to spec.factory as the cpu/numpy terminal rung.
|
|
self._components[component_id] = self.platform.build_component(spec, self)
|
|
return self._components[component_id]
|
|
|
|
def has_component(self, component_id: str) -> bool:
|
|
return component_id in self.card.components
|
|
|
|
# --- loops: one stateless Loop object per (instance, loop_id) ------------- #
|
|
def loop(self, loop_id: str) -> Any:
|
|
if loop_id not in self._loops:
|
|
spec = self.card.loops.get(loop_id)
|
|
if spec is None:
|
|
raise KeyError(f"loop {loop_id!r} not declared on card {self.card.model_id!r}")
|
|
if spec.loop_factory is None:
|
|
raise RuntimeError(f"loop {loop_id!r} has no loop_factory")
|
|
self._loops[loop_id] = spec.loop_factory()
|
|
return self._loops[loop_id]
|
|
|
|
# --- sleep/wake by component (CuMem tags = component names) --------------- #
|
|
def sleep(self, component_ids: list[str]) -> None:
|
|
for cid in component_ids:
|
|
self._components.pop(cid, None)
|
|
self._asleep.add(cid)
|
|
|
|
def wake(self, component_ids: list[str]) -> None:
|
|
for cid in component_ids:
|
|
self._asleep.discard(cid)
|
|
|
|
# --- weight sync: bump version + invalidate caches ----------------------- #
|
|
def version_of(self, component_id: str) -> str:
|
|
"""The component's own weights version (defaults to the instance version)."""
|
|
return self.component_versions.get(component_id, self.weights_version)
|
|
|
|
def set_weights_version(self, version: str, components: list[str] | None = None) -> None:
|
|
"""Publish a new weights version. If ``components`` is given, only those components' versions
|
|
bump and only their caches are invalidated (partition, not flush) — so a transformer-only RL
|
|
weight sync leaves the frozen text-encoder's feature cache intact."""
|
|
self.weights_version = version
|
|
changed = components if components is not None else list(self.card.components.keys())
|
|
for c in changed:
|
|
self.component_versions[c] = version
|
|
if self.caches is not None and hasattr(self.caches, "invalidate_components"):
|
|
self.caches.invalidate_components(set(changed))
|
|
# Evict captured CUDA graphs for the synced components too (else they leak on a real box;
|
|
# version-in-key already makes them unreachable). Duck-typed → card/ imports no runtime.
|
|
if self.graphs is not None and hasattr(self.graphs, "invalidate"):
|
|
self.graphs.invalidate(set(changed))
|
|
|
|
def __repr__(self) -> str:
|
|
return (f"ModelInstance(card={self.card.model_id!r}, weights={self.weights_version!r}, "
|
|
f"resident={sorted(self._components)})")
|
|
|
|
|
|
def load_card(card: ModelCard,
|
|
parallel_plan: Any = None,
|
|
cache_manager: Any = None,
|
|
*,
|
|
validate: bool = True,
|
|
platform: Any = None) -> ModelInstance:
|
|
"""The card-as-factory entrypoint: the card is a runtime factory.
|
|
|
|
``platform`` selects the backend (device, arch); when omitted it is detected (CPU/numpy here,
|
|
CUDA on a torch+GPU box). The same card loads on any backend — only the resolved component/kernel
|
|
implementations differ.
|
|
"""
|
|
if validate:
|
|
card.validate()
|
|
return ModelInstance(card,
|
|
parallel_plan=parallel_plan,
|
|
cache_manager=cache_manager,
|
|
platform=platform if platform is not None else Platform.detect())
|