SolitaryThinker
2f044c02dd
Add v2/platform/: the (recipe, runtime) backend membrane that lets CPU, GPU, and other devices coexist behind one dispatch substrate. - Two tuple-keyed registries: COMPONENTS(kind, device, variant) for weight-bearing components and KERNELS(op, device, arch, variant) for stateless primitives, each with an availability predicate and an enumerable manifest (declared-but-unavailable cells listed without importing torch). - Platform: detected (device, arch) owning the device + arch fallback chains and a per-platform cached KernelTable; detect() is honest (CPU/numpy unless torch+CUDA are actually present). Arch fallback is monotonic — only degrades to older/portable archs, never a newer binary-incompatible one. - Seams wired: ModelInstance.component() -> platform.build_component(spec,self) with spec.factory as the numpy terminal rung (existing cards untouched); the wan21 denoise thunk dispatches solver ops through model.platform.kernels. - Three backends: cpu (numpy terminal + parity oracle), accel (pure-python stand-in proving cross-device resolution, arch fallback, and the oracle), torch_cuda (declared-but-unavailable; no faked GPU). - test_platform.py: 16 tests — detection, terminal rung, arch-fallback walk + monotonicity, device precedence, variant fallback, component override + per-kind device fallback, the parity oracle (accel == cpu, bit_identical via the C1 ladder), and matrix enumeration without importing torch. Scope is honestly bounded in the README/docstrings: only wan21-denoise routes through the kernel table and only the dit kind is overridden today (each a one-line adoption); the torch/CUDA path and a cudagraph workspace-safety contract are declared/deferred, not implemented. 183 tests pass.
113 lines
5.7 KiB
Python
113 lines
5.7 KiB
Python
"""ModelInstance — a resident, loaded card (design_v3 §4.3).
|
|
|
|
> A ``ModelInstance`` is a resident, loaded card: component instances, model state,
|
|
> caches, compiled graphs, and a parallel plan. **A request may run several of the
|
|
> card's loops against one ``ModelInstance``.**
|
|
|
|
That single sentence is the difference between this design and a stage-only design,
|
|
and it 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.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from ..platform import Platform
|
|
from .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 is the whole
|
|
# "change only the backend, not the loops/policies/training" story (design_v3 §17).
|
|
self.platform = platform if platform is not None else Platform.cpu()
|
|
self.adapter_versions: dict[str, str] = {}
|
|
# per-component weight versions (§7.1): 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 (design_v3 §7.3; 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 (design_v3 §10): 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 (design_v3 §7.1 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))
|
|
|
|
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 (design_v3 §4.1: 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())
|