refactor: unify operator selection behind generic dispatch
Add astrai/extension/dispatch.py: per-family decision tables over composable Specs, with explicit-strict / implicit-loose resolution, ASTR_OPS env overrides, profile presets, and explain traces. Migrate attention (behavior-preserving facade) and rotary onto it; new tests cover spec algebra, resolution semantics, and spec-vs-supports_call consistency.
This commit is contained in:
@@ -0,0 +1,388 @@
|
||||
"""Operator dispatch: one selection mechanism for all op families.
|
||||
|
||||
A family registers an ordered list of ``ImplRecord`` rows (name, impl
|
||||
object, capability ``Spec``, machine-level ``available``). Resolution:
|
||||
explicit/context selection (strict — raises when incapable) > ``ASTR_OPS``
|
||||
env entry (soft — falls through) > first capable row > family fallback.
|
||||
``Spec`` is a composable predicate over ``CallContext``; the rows are the
|
||||
family's decision table, printable via ``explain``.
|
||||
|
||||
Records flagged ``faithful=False`` change numerics (e.g. fp8) and are only
|
||||
reachable through an explicit selection, never the implicit chain.
|
||||
"""
|
||||
|
||||
import contextvars
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field, replace
|
||||
from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CallContext:
|
||||
"""Axes a decision may depend on, snapshotted per call.
|
||||
|
||||
``raw``/``raw_kwargs`` keep the original call so unregistered handles
|
||||
can still be probed through their own capability method.
|
||||
"""
|
||||
|
||||
family: str
|
||||
fwd: Optional[str] = None
|
||||
dtype: Optional[torch.dtype] = None
|
||||
device_cuda: bool = False
|
||||
ndim: Optional[int] = None
|
||||
head_dim: Optional[int] = None
|
||||
has_cache: bool = False
|
||||
has_mask: bool = False
|
||||
grad_enabled: bool = True
|
||||
raw: Tuple[Any, ...] = ()
|
||||
raw_kwargs: Mapping[str, Any] = field(default_factory=dict)
|
||||
|
||||
def evolve(self, **changes) -> "CallContext":
|
||||
return replace(self, **changes)
|
||||
|
||||
|
||||
class Spec:
|
||||
"""Composable, self-describing predicate over ``CallContext``."""
|
||||
|
||||
__slots__ = ("_fn", "_desc")
|
||||
|
||||
def __init__(self, fn: Callable[[CallContext], bool], desc: str):
|
||||
self._fn = fn
|
||||
self._desc = desc
|
||||
|
||||
def matches(self, ctx: CallContext) -> bool:
|
||||
return bool(self._fn(ctx))
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return self._desc
|
||||
|
||||
def __and__(self, other: "Spec") -> "Spec":
|
||||
return Spec(
|
||||
lambda ctx: self._fn(ctx) and other._fn(ctx),
|
||||
f"({self._desc} and {other._desc})",
|
||||
)
|
||||
|
||||
def __or__(self, other: "Spec") -> "Spec":
|
||||
return Spec(
|
||||
lambda ctx: self._fn(ctx) or other._fn(ctx),
|
||||
f"({self._desc} or {other._desc})",
|
||||
)
|
||||
|
||||
def __invert__(self) -> "Spec":
|
||||
return Spec(lambda ctx: not self._fn(ctx), f"not({self._desc})")
|
||||
|
||||
@classmethod
|
||||
def always(cls) -> "Spec":
|
||||
return cls(lambda ctx: True, "always")
|
||||
|
||||
@classmethod
|
||||
def of(cls, fn: Callable[[CallContext], bool], desc: str) -> "Spec":
|
||||
return cls(fn, desc)
|
||||
|
||||
@classmethod
|
||||
def dtype_in(cls, *dtypes: torch.dtype) -> "Spec":
|
||||
names = [str(getattr(d, "name", d)) for d in dtypes]
|
||||
return cls(
|
||||
lambda ctx: ctx.dtype in dtypes, f"dtype∈{{{','.join(names)}}}"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def ndim(cls, n: int) -> "Spec":
|
||||
return cls(lambda ctx: ctx.ndim == n, f"ndim=={n}")
|
||||
|
||||
@classmethod
|
||||
def head_dim_in(cls, *dims: int) -> "Spec":
|
||||
return cls(
|
||||
lambda ctx: ctx.head_dim in dims,
|
||||
f"head_dim∈{{{','.join(map(str, dims))}}}",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def fwd_in(cls, *modes: Optional[str]) -> "Spec":
|
||||
names = [str(m) for m in modes]
|
||||
return cls(lambda ctx: ctx.fwd in modes, f"fwd∈{{{','.join(names)}}}")
|
||||
|
||||
@classmethod
|
||||
def inference(cls) -> "Spec":
|
||||
return cls(lambda ctx: ctx.fwd is not None, "inference")
|
||||
|
||||
@classmethod
|
||||
def training(cls) -> "Spec":
|
||||
return cls(lambda ctx: ctx.fwd is None, "training")
|
||||
|
||||
@classmethod
|
||||
def no_grad(cls) -> "Spec":
|
||||
return cls(lambda ctx: not ctx.grad_enabled, "no_grad")
|
||||
|
||||
@classmethod
|
||||
def has_cache(cls) -> "Spec":
|
||||
return cls(lambda ctx: ctx.has_cache, "has_cache")
|
||||
|
||||
@classmethod
|
||||
def mask_free(cls) -> "Spec":
|
||||
return cls(lambda ctx: not ctx.has_mask, "mask_free")
|
||||
|
||||
@classmethod
|
||||
def cuda_device(cls) -> "Spec":
|
||||
return cls(lambda ctx: ctx.device_cuda, "cuda_device")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImplRecord:
|
||||
"""One decision-table row: an implementation plus its capability."""
|
||||
|
||||
family: str
|
||||
name: str
|
||||
obj: Any
|
||||
spec: Spec
|
||||
available: Callable[[], bool] = lambda: True
|
||||
priority: int = 100
|
||||
faithful: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class OpFamily:
|
||||
name: str
|
||||
provider: Callable[[], List[ImplRecord]]
|
||||
fallback: Callable[[], ImplRecord]
|
||||
|
||||
|
||||
_FAMILIES: Dict[str, OpFamily] = {}
|
||||
_ENV_ALIASES: Dict[str, str] = {}
|
||||
|
||||
_current_overrides: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar(
|
||||
"astrai_op_overrides", default={}
|
||||
)
|
||||
|
||||
_env_lock = threading.Lock()
|
||||
_env_cache: Dict[tuple, Optional[Dict[str, str]]] = {}
|
||||
_warned: set = set()
|
||||
|
||||
|
||||
def register_family(
|
||||
name: str,
|
||||
provider: Callable[[], List[ImplRecord]],
|
||||
fallback: Callable[[], ImplRecord],
|
||||
) -> None:
|
||||
"""Register (or replace) a family; ``provider`` is re-evaluated per
|
||||
resolution so availability changes (tests, late imports) are honored."""
|
||||
_FAMILIES[name] = OpFamily(name, provider, fallback)
|
||||
|
||||
|
||||
def register_env_alias(family: str, varname: str) -> None:
|
||||
"""Legacy single-value env var for a family (e.g. attention →
|
||||
ASTR_BACKEND); an ASTR_OPS entry wins when both are set."""
|
||||
_ENV_ALIASES[family] = varname
|
||||
|
||||
|
||||
def _family(name: str) -> OpFamily:
|
||||
fam = _FAMILIES.get(name)
|
||||
if fam is None:
|
||||
raise KeyError(f"no operator family registered under {name!r}")
|
||||
return fam
|
||||
|
||||
|
||||
def _warn_once(message: str) -> None:
|
||||
if message not in _warned:
|
||||
_warned.add(message)
|
||||
logger.warning(message)
|
||||
|
||||
|
||||
def set_override(family: str, handle: Any) -> contextvars.Token:
|
||||
overrides = dict(_current_overrides.get())
|
||||
overrides[family] = handle
|
||||
return _current_overrides.set(overrides)
|
||||
|
||||
|
||||
def reset_override(token: contextvars.Token) -> None:
|
||||
_current_overrides.reset(token)
|
||||
|
||||
|
||||
def get_override(family: str) -> Optional[Any]:
|
||||
return _current_overrides.get().get(family)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def op_backend(**handles: Any):
|
||||
"""Select implementations per family for the enclosed scope::
|
||||
|
||||
with op_backend(attention="torch_native", rotary="torch"):
|
||||
engine.generate(...)
|
||||
|
||||
String handles are validated eagerly against the family's currently
|
||||
available implementations; object handles pass through unchecked.
|
||||
"""
|
||||
for family, handle in handles.items():
|
||||
if isinstance(handle, str):
|
||||
fam = _FAMILIES.get(family)
|
||||
if fam is None:
|
||||
raise ValueError(f"unknown operator family {family!r}")
|
||||
if _record_for_handle(fam, handle) is None:
|
||||
raise ValueError(f"Unknown {family} implementation: {handle!r}")
|
||||
tokens = [set_override(f, h) for f, h in handles.items()]
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
for token in reversed(tokens):
|
||||
reset_override(token)
|
||||
|
||||
|
||||
def env_overrides() -> Dict[str, str]:
|
||||
"""Merged ASTR_OPS + legacy-alias selections (family or "profile").
|
||||
|
||||
Cached per distinct env content; unknown families / malformed entries
|
||||
warn once and are dropped (soft override, never fatal).
|
||||
"""
|
||||
with _env_lock:
|
||||
merged: Dict[str, str] = {}
|
||||
raw = os.environ.get("ASTR_OPS", "").strip()
|
||||
if raw:
|
||||
key = ("ASTR_OPS", raw)
|
||||
if key not in _env_cache:
|
||||
parsed: Dict[str, str] = {}
|
||||
for item in raw.split(","):
|
||||
key_part, sep, value = item.strip().partition("=")
|
||||
key_part, value = key_part.strip(), value.strip()
|
||||
if not sep or not key_part or not value:
|
||||
_warn_once(f"ASTR_OPS: ignoring malformed entry {item!r}")
|
||||
continue
|
||||
parsed[key_part] = value
|
||||
_env_cache[key] = parsed or None
|
||||
merged.update(_env_cache[key] or {})
|
||||
for fam, varname in _ENV_ALIASES.items():
|
||||
raw = os.environ.get(varname, "").strip()
|
||||
if raw:
|
||||
key = (varname, raw)
|
||||
if key not in _env_cache:
|
||||
_env_cache[key] = {fam: raw.lower()}
|
||||
merged.setdefault(fam, _env_cache[key][fam])
|
||||
for fam in [f for f in merged if f not in _FAMILIES and f != "profile"]:
|
||||
_warn_once(f"ASTR_OPS: unknown operator family {fam!r}; dropping it")
|
||||
merged.pop(fam)
|
||||
return merged
|
||||
|
||||
|
||||
def env_selection(family: str) -> Optional[str]:
|
||||
return env_overrides().get(family)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Resolution:
|
||||
record: ImplRecord
|
||||
origin: str
|
||||
|
||||
|
||||
class ExplicitSelectionError(RuntimeError):
|
||||
"""An explicitly selected implementation cannot handle the call."""
|
||||
|
||||
|
||||
def _record_for_handle(fam: OpFamily, handle: Any) -> Optional[ImplRecord]:
|
||||
records = sorted(fam.provider(), key=lambda r: r.priority)
|
||||
if isinstance(handle, str):
|
||||
return next((r for r in records if r.name == handle), None)
|
||||
return next((r for r in records if r.obj is handle), None)
|
||||
|
||||
|
||||
def _adhoc_record(family: str, handle: Any) -> ImplRecord:
|
||||
"""Wrap an unregistered object; capability probes its own method."""
|
||||
supports = getattr(handle, "supports_call", None)
|
||||
if supports is not None:
|
||||
spec = Spec.of(
|
||||
lambda ctx: bool(supports(*ctx.raw, **ctx.raw_kwargs)),
|
||||
f"{type(handle).__name__}.supports_call",
|
||||
)
|
||||
else:
|
||||
spec = Spec.always()
|
||||
return ImplRecord(family, type(handle).__name__, handle, spec)
|
||||
|
||||
|
||||
def resolve(
|
||||
family: str, ctx: CallContext, explicit: Optional[Any] = None
|
||||
) -> Resolution:
|
||||
"""Resolve one family for one call (explicit-strict / implicit-loose)."""
|
||||
fam = _family(family)
|
||||
|
||||
handle: Optional[Any] = None
|
||||
origin = "chain"
|
||||
if explicit is not None:
|
||||
handle, origin = explicit, "explicit"
|
||||
elif get_override(family) is not None:
|
||||
handle, origin = get_override(family), "context"
|
||||
else:
|
||||
env_name = env_selection(family)
|
||||
if env_name is not None:
|
||||
handle, origin = env_name, "env"
|
||||
|
||||
if handle is not None:
|
||||
record = _record_for_handle(fam, handle)
|
||||
if record is None and not isinstance(handle, str):
|
||||
record = _adhoc_record(family, handle)
|
||||
if record is None:
|
||||
if origin in ("explicit", "context"):
|
||||
raise ValueError(f"Unknown {family} implementation: {handle!r}")
|
||||
_warn_once(f"ASTR_OPS: {family}={handle!r} is not registered; ignoring")
|
||||
else:
|
||||
if record.available() and record.spec.matches(ctx):
|
||||
return Resolution(record, origin)
|
||||
if origin in ("explicit", "context"):
|
||||
raise ExplicitSelectionError(
|
||||
f"Explicitly-set backend {type(record.obj).__name__} cannot "
|
||||
f"handle this {family} call; required: {record.spec.description}"
|
||||
)
|
||||
|
||||
if handle is None and env_overrides().get("profile") == "reference":
|
||||
return Resolution(fam.fallback(), "profile")
|
||||
|
||||
for record in sorted(fam.provider(), key=lambda r: r.priority):
|
||||
if record.available() and record.faithful and record.spec.matches(ctx):
|
||||
return Resolution(record, "chain")
|
||||
return Resolution(fam.fallback(), "fallback")
|
||||
|
||||
|
||||
def resolve_plan(ctxs: Mapping[str, CallContext]) -> Dict[str, Resolution]:
|
||||
"""Resolve several families at once (one decision snapshot)."""
|
||||
return {family: resolve(family, ctx) for family, ctx in ctxs.items()}
|
||||
|
||||
|
||||
def _describe_ctx(ctx: CallContext) -> str:
|
||||
parts = [f"fwd={ctx.fwd}", f"dtype={ctx.dtype}"]
|
||||
if ctx.ndim is not None:
|
||||
parts.append(f"ndim={ctx.ndim}")
|
||||
if ctx.head_dim is not None:
|
||||
parts.append(f"head_dim={ctx.head_dim}")
|
||||
parts += [f"cache={ctx.has_cache}", f"mask={ctx.has_mask}", f"grad={ctx.grad_enabled}"]
|
||||
return " ".join(parts)
|
||||
|
||||
|
||||
def explain(family: str, ctx: CallContext, explicit: Optional[Any] = None) -> str:
|
||||
"""Human-readable decision trace for one family."""
|
||||
fam = _family(family)
|
||||
records = sorted(fam.provider(), key=lambda r: r.priority)
|
||||
lines = [f"[{family}] {_describe_ctx(ctx)}"]
|
||||
for record in records:
|
||||
if not record.available():
|
||||
lines.append(f" {record.name}: SKIP unavailable")
|
||||
elif not record.faithful:
|
||||
lines.append(f" {record.name}: SKIP not faithful (explicit-only)")
|
||||
elif record.spec.matches(ctx):
|
||||
lines.append(f" {record.name}: MATCH ({record.spec.description})")
|
||||
else:
|
||||
lines.append(f" {record.name}: reject ({record.spec.description})")
|
||||
try:
|
||||
resolution = resolve(family, ctx, explicit)
|
||||
lines.append(f" => {resolution.record.name} (origin={resolution.origin})")
|
||||
except (ExplicitSelectionError, ValueError) as exc:
|
||||
lines.append(f" => ERROR: {exc}")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def explain_plan(ctxs: Mapping[str, CallContext]) -> str:
|
||||
return "\n".join(explain(family, ctx) for family, ctx in ctxs.items())
|
||||
Reference in New Issue
Block a user