diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e47fb0c98..770009920 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -220,6 +220,14 @@ jobs: - name: pgw#1049 guard - torch-settings writes outside the authority run: uv run python scripts/lint_settings_writers.py + # GATING (pgw#1143, §1.33): `Slot(layouts=...)` is the per-slot layout + # DEMAND the hub's gate reads. It must be readable by the AST sweep as + # well as by the import, so a computed declaration is refused where it + # is written — otherwise the published manifest becomes the only place + # the demand can be read, which is the dual-declaration hazard. + - name: pgw#1143 guard - Slot(layouts=) declarations are literals + run: uv run python scripts/lint_layout_declarations.py + # GATING (pgw#1122): the compute child holds no worker credential BY # CONSTRUCTION, so a gate that reads one to answer "who am I?" or "is # there a hub?" refuses on every real serving pod (pgw#1108, then diff --git a/changelog.d/pgw1143.md b/changelog.d/pgw1143.md new file mode 100644 index 000000000..5d3e1888f --- /dev/null +++ b/changelog.d/pgw1143.md @@ -0,0 +1 @@ +- **pgw#1143 steps 1-3 (DESIGN-RULINGS §1.33, with th#1809 T1-T2): a slot DECLARES the tensor layouts its code can execute.** `Slot(layouts={component_path: (handle, ...)})` is the code-side DEMAND half of the layout contract — an ordered tuple per component path (`"*"` is the whole tree), ordered by preference, published into the release manifest as `functions[].slots[].layouts` so the hub can gate a rebind against what the AUTHOR says the slot needs instead of against what the installed wheel happens to decode (2 distinct decoder signatures across 64 releases). **Absent is UNDECLARED** — a tri-state, never "accepts everything": an empty mapping and an empty tuple are both decoration-time errors, because collapsing that tri-state is th#1580's fail-open defect wearing a new name. It lives on `Slot` and deliberately not on `Compile`, whose fields feed `contract_axes()` and would re-key every cell in the fleet for a fact §1.33 point 5 says must never enter the key. Handles are validated against `KNOWN_CONTRACTS` at the constructor and component keys against the DERIVED component tree at decoration (a key matching nothing is a build error naming the tree); the SDK emits HANDLES only and the hub resolves each to its descriptor digest at manifest ingest, so the transcription is allowed to be stale and is CHECKED, never authoritative. The decoder census stays a LOWER bound: a declared handle no `@implements_contract` decoder backs lands on the manifest as `layouts_census_unbacked` and does not refuse, since the blockwise conditioner is decoded natively by `transformers` with zero cozy markers. New `fast gates` guard `scripts/lint_layout_declarations.py` refuses a declaration the AST sweep cannot read (a comprehension, an f-string, a helper-built dict), because a computed declaration would make the published manifest the only place the demand can be read — the dual-declaration hazard pgw#1107 spent a program deleting. diff --git a/docs/endpoint-authoring.md b/docs/endpoint-authoring.md index f348d8752..77dfc142c 100644 --- a/docs/endpoint-authoring.md +++ b/docs/endpoint-authoring.md @@ -368,6 +368,34 @@ class Generate: components (the qwen text encoder, a shared VAE) load once and are refcounted across checkpoint picks; there is nothing to declare (`share_components=` is deleted). +- **`layouts=` declares WHAT BYTES THIS SLOT'S CODE CAN EXECUTE** (§1.33, + pgw#1143) — the DEMAND half of the tensor-layout contract, per component + path, ordered by preference: + + ```python + from gen_worker.models.tensor_layout_contract import ( + CONTRACT_HF_FP8_BLOCKWISE, CONTRACT_PLAIN_BF16) + + Slot(StableDiffusionXLPipeline, selected_by="model", layouts={ + "*": (CONTRACT_PLAIN_BF16,), + "text_encoder": (CONTRACT_HF_FP8_BLOCKWISE, CONTRACT_PLAIN_BF16), + }) + ``` + + `"*"` is the whole-tree default; a component key overrides it for that + component. The hub compares an artifact's PROVEN layout against this set at + rebind and refuses a mismatch with both sides named, before any pod is + bought. **Omitting `layouts=` leaves the slot UNDECLARED** — the gate then + falls back to the image-wide decoder census, and absence is never read as + "accepts everything"; an empty mapping or an empty tuple is a + decoration-time error. Handles must be registered + (`KNOWN_CONTRACTS`, transcribed from tensorhub's `internal/tensorlayout`) + and written as LITERALS or as constants imported from that module — + `scripts/lint_layout_declarations.py` refuses anything the AST sweep cannot + read. Declaring a handle no decoder in the image backs is NOT an error: it + lands in the build log as `layouts_census_unbacked`, because plenty of + layouts are decoded natively by `transformers`/`diffusers` with no cozy + marker. **Per-family defaults vocabulary**: a typed, versioned, JSON-Schema-exportable struct per architecture — the shape tensorhub validates diff --git a/scripts/lint_layout_declarations.py b/scripts/lint_layout_declarations.py new file mode 100644 index 000000000..bd8ec21ac --- /dev/null +++ b/scripts/lint_layout_declarations.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python3 +"""pgw#1143 step 3 (§1.33): `Slot(layouts=...)` is READABLE WITHOUT IMPORTING. + +The per-slot DEMAND is the one fact the hub's layout gate reads out of the +manifest, so it has to be legible to BOTH paths that ever look at it: + +* the IMPORT path — `gen_worker.discovery` imports endpoint modules with heavy + deps stubbed (`discovery/heavy_deps.py`: "that metadata is torch-free by + design"), and +* the AST path — the sweep pgw#1107 used to find the eight class-less + `Compile(...)` declarations, and this script. + +A computed declaration (a comprehension, an f-string, a dict built by a +helper, a value read from config) is invisible to the second. It would make +the published manifest the only place the demand can be read, which is exactly +the dual-declaration hazard pgw#1107 spent a program deleting — and it is +unreviewable in a diff, which is where a layout demand is actually judged. + +So: `layouts=` must be a DICT LITERAL whose keys are string literals and whose +values are tuple/list literals of string literals or of names imported from +`gen_worker.models.tensor_layout_contract`. Nothing else. + +This is a fence, not a taste rule. It carries no allowlist by design: there is +no declaration this refuses that could not be written literally instead. + +Usage: + + python scripts/lint_layout_declarations.py [PATH ...] + +Defaults to this repository's `src/`. An ENDPOINT repo runs the same script +against its own tree — the constraint is on the declaration, not on where it +lives. +""" + +from __future__ import annotations + +import ast +import sys +from pathlib import Path +from typing import Iterator, List, Set, Tuple + +REPO = Path(__file__).resolve().parents[1] + +#: `src/` only, deliberately. The subject is a declaration that gets +#: PUBLISHED — an endpoint's `Slot(layouts=...)` reaching a release manifest. +#: `tests/` holds the opposite by construction: a test proving the constructor +#: refuses a computed declaration has to WRITE one, and sweeping it would make +#: this fence and its own negative test mutually exclusive. An endpoint repo +#: runs this against its own tree, where every declaration is a real one. +DEFAULT_ROOTS = (REPO / "src",) + +#: The module whose module-level constants may stand in for a handle literal. +#: It is the SDK's transcription of tensorhub's registry, so a name imported +#: from it resolves to a handle string the lint can still classify by NAME. +VOCABULARY_MODULE = "gen_worker.models.tensor_layout_contract" + + +def _iter_python_files(roots: Tuple[Path, ...]) -> Iterator[Path]: + for root in roots: + if root.is_file() and root.suffix == ".py": + yield root + continue + if not root.is_dir(): + continue + for path in sorted(root.rglob("*.py")): + yield path + + +def _vocabulary_names(tree: ast.Module) -> Set[str]: + """Local names bound to constants of the vocabulary module.""" + names: Set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module == VOCABULARY_MODULE: + for alias in node.names: + names.add(alias.asname or alias.name) + elif isinstance(node, ast.Import): + for alias in node.names: + if alias.name == VOCABULARY_MODULE: + names.add(alias.asname or alias.name.split(".")[0]) + return names + + +def _handle_is_readable(node: ast.expr, vocabulary: Set[str]) -> bool: + if isinstance(node, ast.Constant) and isinstance(node.value, str): + return True + if isinstance(node, ast.Name): + return node.id in vocabulary + if isinstance(node, ast.Attribute): + # `tensor_layout_contract.CONTRACT_PLAIN_BF16` — the module itself was + # imported under a name the sweep can resolve. + root: ast.expr = node + while isinstance(root, ast.Attribute): + root = root.value + return isinstance(root, ast.Name) and root.id in vocabulary + return False + + +def _check_layouts( + path: Path, call: ast.Call, value: ast.expr, vocabulary: Set[str], +) -> List[str]: + where = f"{path}:{value.lineno}" + if not isinstance(value, ast.Dict): + return [ + f"{where}: layouts= is a {type(value).__name__}, not a dict " + "literal — the AST sweep cannot read it, so the hub's copy would " + "be the only place this demand can be read" + ] + problems: List[str] = [] + for key, item in zip(value.keys, value.values): + if key is None: + problems.append( + f"{where}: layouts= splats another mapping (**); a splat hides " + "the component paths from the sweep") + continue + if not (isinstance(key, ast.Constant) and isinstance(key.value, str)): + problems.append( + f"{path}:{key.lineno}: layouts= key is not a string literal") + continue + if not isinstance(item, (ast.Tuple, ast.List)): + problems.append( + f"{path}:{item.lineno}: layouts[{key.value!r}] is a " + f"{type(item).__name__}, not a tuple literal — order is " + "preference and a computed sequence has no reviewable order") + continue + for element in item.elts: + if not _handle_is_readable(element, vocabulary): + problems.append( + f"{path}:{element.lineno}: layouts[{key.value!r}] holds a " + f"{type(element).__name__}; every handle must be a string " + f"literal or a constant imported from {VOCABULARY_MODULE}") + return problems + + +def _call_name(call: ast.Call) -> str: + func = call.func + if isinstance(func, ast.Name): + return func.id + if isinstance(func, ast.Attribute): + return func.attr + return "" + + +def scan(path: Path) -> List[str]: + try: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + except (SyntaxError, UnicodeDecodeError) as exc: + return [f"{path}: could not parse ({exc})"] + vocabulary = _vocabulary_names(tree) + problems: List[str] = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call) or _call_name(node) != "Slot": + continue + for kw in node.keywords: + if kw.arg == "layouts": + problems.extend(_check_layouts(path, node, kw.value, vocabulary)) + elif kw.arg is None: + # `Slot(cls, **kwargs)` — a layouts= could be hiding in there. + problems.append( + f"{path}:{node.lineno}: Slot(**kwargs) may carry a " + "layouts= the sweep cannot see; pass the declaration " + "explicitly") + return problems + + +def main(argv: List[str]) -> int: + roots = tuple(Path(a).resolve() for a in argv[1:]) or DEFAULT_ROOTS + problems: List[str] = [] + for path in _iter_python_files(roots): + problems.extend(scan(path)) + if problems: + print("pgw#1143: Slot(layouts=...) declarations that the AST sweep " + "cannot read:\n", file=sys.stderr) + for problem in problems: + print(f" {problem}", file=sys.stderr) + print( + "\nWrite the declaration literally — a dict literal of string " + "literals mapping to tuple literals of handles. The demand is " + "reviewed in the diff and read by the hub; both need it visible " + "without running the module.", + file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv)) diff --git a/src/gen_worker/api/slot.py b/src/gen_worker/api/slot.py index 5ec3393fb..8de258623 100644 --- a/src/gen_worker/api/slot.py +++ b/src/gen_worker/api/slot.py @@ -35,12 +35,13 @@ def generate(self, ctx: RequestContext[SdxlDefaults], p: In) -> Out: from __future__ import annotations -from typing import Any, Dict, Generic, Optional, Sequence, Type, TypeVar +from typing import Any, Dict, Generic, Mapping, Optional, Sequence, Type, TypeVar import msgspec from .binding import ModelRef from ..families.base import KIND_LORA, GenerationDefaults, family_for +from ..models.tensor_layout_contract import normalize_layout_demand D = TypeVar("D", bound=GenerationDefaults) @@ -125,6 +126,32 @@ class (diffusers-style: exposes ``_get_signature_keys`` or a components multi-slot shape must mark exactly one root — ambiguity is a decoration-time error, never a silent fallback. + ``layouts`` is the slot's per-component DEMAND (§1.33, pgw#1143): a + mapping from component path to the ORDERED tuple of tensor-layout + contract handles this slot's code can execute. ``"*"`` is the whole-tree + default; a component key overrides it for that component only:: + + Slot(StableDiffusionXLPipeline, selected_by="model", layouts={ + "*": (CONTRACT_PLAIN_BF16,), + "text_encoder": (CONTRACT_HF_FP8_BLOCKWISE, CONTRACT_PLAIN_BF16), + }) + + Order IS preference. The handles are validated here against the SDK's + transcribed ``KNOWN_CONTRACTS``; the KEYS are validated at decoration + against the slot's DERIVED component tree, so a key that would silently + match nothing is a build error naming the tree. **Absent is UNDECLARED** + — a tri-state, neither "accepts everything" nor "accepts nothing": the + hub's gate has no teeth on a slot that declares nothing and falls back to + the image-wide decoder census. An empty mapping or an empty tuple is a + decoration-time error, because collapsing the tri-state is th#1580's + fail-open defect wearing a new name. + + This lives on ``Slot`` and NOT on ``Compile``: ``Compile``'s fields feed + ``contract_axes()``, a cell-key input, and §1.33 point 5 is that + conversion is upstream of compute and invisible to cell identity. A + layout declaration there would either re-key every cell in the fleet or + sit inside the key struct while deliberately not participating. + ``optional`` is DERIVED, never passed: a slot is optional exactly when its ``setup()`` parameter carries a default (``edit: Pipe | None = None``). The signature is the single source of truth, so the two can @@ -135,7 +162,7 @@ class (diffusers-style: exposes ``_get_signature_keys`` or a components __slots__ = ( "pipeline_cls", "selected_by", "family", "default_checkpoint", "root", - "optional", + "optional", "layouts", ) def __init__( @@ -146,6 +173,7 @@ def __init__( family: Optional[str] = None, default_checkpoint: Optional[ModelRef] = None, root: bool = False, + layouts: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: if not isinstance(pipeline_cls, type): raise TypeError( @@ -163,12 +191,21 @@ def __init__( self.default_checkpoint = default_checkpoint self.root = bool(root) self.optional = False # derived at decoration from setup()'s default + # None is UNDECLARED and stays None; anything else normalizes now, at + # the declaration site, so the traceback names the Slot the author + # wrote rather than a manifest key. + self.layouts: Optional[Dict[str, tuple[str, ...]]] = ( + None if layouts is None + else normalize_layout_demand( + layouts, where=f"Slot({pipeline_cls.__name__})") + ) def __repr__(self) -> str: # pragma: no cover - debug aid return ( f"Slot({self.pipeline_cls.__name__}, selected_by={self.selected_by!r}, " f"family={self.family!r}, default_checkpoint={self.default_checkpoint!r}, " - f"root={self.root!r}, optional={self.optional!r})" + f"root={self.root!r}, optional={self.optional!r}, " + f"layouts={self.layouts!r})" ) diff --git a/src/gen_worker/discovery/discover.py b/src/gen_worker/discovery/discover.py index a8efea5f9..2b8e30726 100644 --- a/src/gen_worker/discovery/discover.py +++ b/src/gen_worker/discovery/discover.py @@ -31,6 +31,7 @@ VideoAsset, ) from gen_worker.discovery.execution_lanes import ( + DerivedExecutionLanes, derive_execution_lanes, execution_lanes_for_function, manifest_block, @@ -451,6 +452,16 @@ def _slot_to_manifest( {"name": part, "kind": kind} for part, kind in sorted(components.items()) ] + # §1.33 / pgw#1143: the per-component DEMAND. Absent means UNDECLARED — + # the hub's gate falls back to the image-wide decoder census for this + # slot, and never reads absence as "accepts everything". Handles only: + # the hub resolves each to its descriptor DIGEST at ingest, against its + # own registry, which is the only moment one wheel and one hub are both + # pinned. + if slot.layouts: + out["layouts"] = { + path: list(handles) for path, handles in slot.layouts.items() + } return out @@ -939,6 +950,9 @@ def discover_manifest(root: Optional[Path] = None) -> Dict[str, Any]: {"execution_lane": e.execution_lane, "reason": e.reason} for e in exclusions ] + unbacked = _census_unbacked_layouts(fn, derived) + if unbacked: + fn["layouts_census_unbacked"] = list(unbacked) manifest: Dict[str, Any] = { "functions": functions, @@ -950,6 +964,28 @@ def discover_manifest(root: Optional[Path] = None) -> Dict[str, Any]: return manifest +def _census_unbacked_layouts( + fn: Dict[str, Any], derived: DerivedExecutionLanes, +) -> List[str]: + """Declared handles no `@implements_contract` decoder in this image backs. + + **NOT a refusal**, deliberately (§1.33: the census is a LOWER-bound sanity + check, and a lower bound that refuses is an upper bound). The ruling's own + worked example is decoded natively by `transformers` via + `quantization_config` with zero cozy markers, so refusing here would make + the design's motivating case illegal. It lands on the manifest and in the + build log so an author can see the gap and judge it. + """ + backed = {c.contract for c in derived.contracts} + unbacked: List[str] = [] + for slot in fn.get("slots") or []: + for handles in (slot.get("layouts") or {}).values(): + for handle in handles: + if handle not in backed and handle not in unbacked: + unbacked.append(handle) + return sorted(unbacked) + + def _strip_none(obj: Any) -> Any: """Recursively remove None values from dicts/lists (TOML has no null type).""" if isinstance(obj, dict): diff --git a/src/gen_worker/models/tensor_layout_contract.py b/src/gen_worker/models/tensor_layout_contract.py index 85dfbcf67..3d6cdf561 100644 --- a/src/gen_worker/models/tensor_layout_contract.py +++ b/src/gen_worker/models/tensor_layout_contract.py @@ -147,3 +147,120 @@ def contract_decoders_of(obj: Any) -> tuple[ContractDecoder, ...]: if not isinstance(marked, tuple): return () return tuple(d for d in marked if isinstance(d, ContractDecoder)) + + +# ── §1.33 / pgw#1143: the DEMAND side of the same vocabulary ────────────────── +# +# `@implements_contract` above is the SUPPLY-adjacent census — "which decoders +# does this IMAGE contain". The DEMAND is what `Slot(layouts=...)` declares: +# "what does this slot's code need in order to run". Two facts, one vocabulary, +# so the handle grammar and the registration refusal are shared verbatim rather +# than transcribed twice. +# +# The SDK emits HANDLES and never digests: descriptors are Go, in tensorhub +# (th#1580 A2). The hub resolves handle -> `Contract.Digest()` at MANIFEST +# INGEST against its own registry and stores both; a handle its registry does +# not know fails the manifest there, not at rebind. So `KNOWN_CONTRACTS` is +# honestly what it is — a transcription that is allowed to be stale and is +# CHECKED, never authoritative. + +#: The whole-tree key: this slot's demand for every component that has no +#: more specific declaration. +LAYOUT_KEY_ANY_COMPONENT = "*" + + +class LayoutDeclarationError(ValueError): + """A `Slot(layouts=...)` declaration the SDK refuses where it is written.""" + + +def validate_layout_handle(handle: object, *, where: str) -> str: + """One declared handle, normalized, or a decoration-time refusal. + + The quant axis only. §1.33's rendered `"+"` pair needs a + topology REGISTRY to compare against (th#1809 T3); until that exists a + composite is a handle naming a vocabulary the platform cannot resolve, and + exact-or-refused means refusing it rather than storing half a pair. + """ + if not isinstance(handle, str): + raise LayoutDeclarationError( + f"{where}: layout handle must be a string, got " + f"{type(handle).__name__}" + ) + text = handle.strip() + if "+" in text: + raise LayoutDeclarationError( + f"{where}: {text!r} names a + pair. The topology " + "axis has no registry yet (th#1809 T3) — declare the quant handle " + "alone until it does; a pair the hub cannot resolve field-wise is " + "not exact." + ) + if not _HANDLE_RE.match(text): + raise LayoutDeclarationError( + f"{where}: {text!r} is not a contract handle (want ns.name@N)" + ) + if text not in KNOWN_CONTRACTS: + raise LayoutDeclarationError( + f"{where}: contract {text!r} is not registered. " + "Contracts are CODE (th#1580 A2): register it in tensorhub's " + "internal/tensorlayout with a descriptor and a probe set " + f"before a slot may demand it. Known: {', '.join(KNOWN_CONTRACTS)}" + ) + return text + + +def normalize_layout_demand( + layouts: object, *, where: str, +) -> dict[str, tuple[str, ...]]: + """`Slot(layouts=...)` -> `{component_path: ordered handles}`. + + Ordering IS preference (§1.33 point 2), so the tuple is kept as written + and never sorted. Component-path keys are validated against the DERIVED + component tree separately, at decoration time, by the registry — this + function owns only the shape and the vocabulary. + """ + if not isinstance(layouts, dict): + raise LayoutDeclarationError( + f"{where}: layouts= must be a mapping of component path -> ordered " + f"handles, got {type(layouts).__name__}" + ) + if not layouts: + raise LayoutDeclarationError( + f"{where}: layouts={{}} declares nothing. Omit layouts= to leave " + "this slot UNDECLARED; an empty declaration is neither 'accepts " + "everything' nor 'accepts nothing' and the platform will not " + "guess which." + ) + out: dict[str, tuple[str, ...]] = {} + for raw_key, raw_value in layouts.items(): + if not isinstance(raw_key, str) or not raw_key.strip(): + raise LayoutDeclarationError( + f"{where}: layouts= key {raw_key!r} must be a non-empty " + f"component path or {LAYOUT_KEY_ANY_COMPONENT!r}" + ) + key = raw_key.strip() + if isinstance(raw_value, (str, bytes)) or not isinstance( + raw_value, (tuple, list)): + raise LayoutDeclarationError( + f"{where}: layouts[{key!r}] must be an ORDERED tuple of " + f"handles (order is preference), got " + f"{type(raw_value).__name__}" + ) + if not raw_value: + raise LayoutDeclarationError( + f"{where}: layouts[{key!r}] is empty. A component that " + "accepts no layout cannot be bound at all; omit the key to " + "fall back to the whole-tree declaration, or omit layouts= " + "entirely to leave the slot UNDECLARED." + ) + handles: list[str] = [] + for item in raw_value: + handle = validate_layout_handle( + item, where=f"{where}: layouts[{key!r}]") + if handle in handles: + raise LayoutDeclarationError( + f"{where}: layouts[{key!r}] repeats {handle!r}; the " + "second position is unreachable" + ) + handles.append(handle) + out[key] = tuple(handles) + return out diff --git a/src/gen_worker/registry.py b/src/gen_worker/registry.py index c5ef01e50..8bbbd2667 100644 --- a/src/gen_worker/registry.py +++ b/src/gen_worker/registry.py @@ -28,6 +28,7 @@ from .discovery.names import slugify_name from .discovery.walk import find_endpoints from .families.base import GenerationDefaults +from .models.tensor_layout_contract import LAYOUT_KEY_ANY_COMPONENT from .warmup import validate_class_warmup import dataclasses from .api.compile_axis import warm_guidance_values @@ -411,6 +412,44 @@ def _validate_compile_arms( ) +def _validate_slot_layout_keys( + owner: str, + slots: Dict[str, Slot], + slot_components: Dict[str, Dict[str, str]], +) -> None: + """Every `Slot(layouts=...)` key is `"*"` or a real component path. + + §1.33 pins the DEMAND per (slot, component path), and the path vocabulary + is the DERIVED tree the manifest already publishes — so a key that matches + no component is not a stricter declaration, it is a silently absent one. + The handles themselves were validated at the Slot constructor; only the + keys need the tree, and the tree only exists here. + """ + for name, slot in slots.items(): + declared = getattr(slot, "layouts", None) + if not declared: + continue + tree = slot_components.get(name) or {} + for key in declared: + if key == LAYOUT_KEY_ANY_COMPONENT: + continue + if key in tree: + continue + if not tree: + raise ValueError( + f"{owner}: slot {name!r} declares layouts[{key!r}] but " + f"{slot.pipeline_cls.__name__} is not introspectable, so " + "no component tree is derived — a self-loading slot can " + f"only declare {LAYOUT_KEY_ANY_COMPONENT!r}." + ) + raise ValueError( + f"{owner}: slot {name!r} declares layouts[{key!r}], which is " + f"not a component of {slot.pipeline_cls.__name__}. Its derived " + f"tree is: {', '.join(sorted(tree))} (or " + f"{LAYOUT_KEY_ANY_COMPONENT!r} for the whole tree)." + ) + + def _slot_is_family_agnostic( name: str, slot: Slot, slots: Dict[str, Slot], ) -> bool: @@ -515,6 +554,10 @@ def _spec_for_handler( for name, slot in slots.items() ) if tree } + # §1.33 / pgw#1143: the DEMAND's component keys are checked against the + # tree that was just derived. A key matching nothing is the failure mode + # to prevent — it reads as a declaration and gates nothing. + _validate_slot_layout_keys(owner, slots, slot_components) payload_axes = extract_payload_axes(owner, payload_type) ret = hints.get("return") if ret is None: diff --git a/tests/test_slot_layout_demand_pgw1143.py b/tests/test_slot_layout_demand_pgw1143.py new file mode 100644 index 000000000..96965bdd4 --- /dev/null +++ b/tests/test_slot_layout_demand_pgw1143.py @@ -0,0 +1,322 @@ +"""pgw#1143 steps 1-3 (§1.33): the per-slot layout DEMAND, end to end on the +SDK side. + +What is proven here, in the order the design puts it: + +1. a declared slot's demand reaches the PUBLISHED MANIFEST, per component + path, in declaration order — that is the whole point, since the manifest is + the only thing the hub reads; +2. absence is UNDECLARED — no key at all, not an empty one, so the hub cannot + read it as "accepts everything"; +3. an invalid declaration is refused WHERE IT IS WRITTEN — an unknown handle + at the Slot constructor, a component key that matches nothing at + decoration, a non-literal declaration at the lint; +4. the census cross reports an unbacked handle and does NOT refuse it. + +Run: uv run pytest tests/test_slot_layout_demand_pgw1143.py +""" + +from __future__ import annotations + +import subprocess +import sys +import textwrap +from pathlib import Path + +import msgspec +import pytest + +from gen_worker import RequestContext, Slot, endpoint +from gen_worker.discovery.discover import ( + _census_unbacked_layouts, + _slot_to_manifest, +) +from gen_worker.discovery.execution_lanes import ( + DerivedContract, + DerivedExecutionLanes, +) +from gen_worker.families.base import GenerationDefaults +from gen_worker.models.tensor_layout_contract import ( + CONTRACT_HF_FP8_BLOCKWISE, + CONTRACT_PLAIN_BF16, + LayoutDeclarationError, +) +from gen_worker.registry import extract_specs + +REPO = Path(__file__).resolve().parents[1] + + +class _Vae: + pass + + +class _TextEncoder: + pass + + +class _Unet: + pass + + +class FakePipeline: + """An introspectable pipeline: the derived tree is {vae, text_encoder, + unet}, through the same `_get_signature_keys` hook diffusers exposes.""" + + def __init__(self, vae: _Vae, text_encoder: _TextEncoder, unet: _Unet): + self.vae = vae + self.text_encoder = text_encoder + self.unet = unet + + @classmethod + def _get_signature_keys(cls, _obj: object) -> tuple: + return {"vae", "text_encoder", "unet"}, set() + + +class _In(msgspec.Struct): + prompt: str = "" + + +class _Out(msgspec.Struct): + ok: bool = True + + +class _Defaults(GenerationDefaults): + pass + + +def _components() -> dict: + return {"vae": "weights", "text_encoder": "weights", "unet": "weights"} + + +# ── 1. the declaration reaches the manifest, ordered, per component ────────── + + +def test_a_declared_slot_publishes_its_demand_per_component_in_order() -> None: + slot = Slot( + FakePipeline, + selected_by="model", + layouts={ + "*": (CONTRACT_PLAIN_BF16,), + "text_encoder": (CONTRACT_HF_FP8_BLOCKWISE, CONTRACT_PLAIN_BF16), + }, + ) + block = _slot_to_manifest( + "pipeline", slot, family="", components=_components()) + assert block["layouts"] == { + "*": ["plain.bf16@1"], + "text_encoder": ["hf.fp8-blockwise@1", "plain.bf16@1"], + } + # Order IS preference (§1.33 pt 2) — the fp8 encoder is declared FIRST and + # must not be sorted into second place by anything on the way out. + assert block["layouts"]["text_encoder"][0] == "hf.fp8-blockwise@1" + + +def test_an_undeclared_slot_emits_no_layouts_key_at_all() -> None: + block = _slot_to_manifest( + "pipeline", Slot(FakePipeline), family="", components=_components()) + assert "layouts" not in block, ( + "UNDECLARED must be ABSENT. An empty mapping on the wire is a " + "declaration of nothing, and the hub cannot tell it from a slot whose " + "author has not spoken." + ) + + +# ── 2. refused where written ───────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "layouts, fragment", + [ + ({}, "declares nothing"), + ({"*": ()}, "is empty"), + ({"*": "plain.bf16@1"}, "ORDERED tuple"), + ({"*": ("plain.bf16",)}, "not a contract handle"), + ({"*": ("cozy.not-registered@1",)}, "is not registered"), + ({"*": ("plain.bf16@1", "plain.bf16@1")}, "repeats"), + ({"": ("plain.bf16@1",)}, "must be a non-empty component path"), + ( + {"*": ("diffusers.multifile@1+plain.bf16@1",)}, + "topology axis has no registry yet", + ), + ], +) +def test_an_invalid_declaration_is_refused_at_the_declaration_site( + layouts: object, fragment: str, +) -> None: + with pytest.raises(LayoutDeclarationError) as excinfo: + Slot(FakePipeline, layouts=layouts) # type: ignore[arg-type] + assert fragment in str(excinfo.value) + + +def test_a_component_key_matching_nothing_is_a_decoration_error() -> None: + """The key vocabulary is the DERIVED tree. A key that matches no component + reads as a declaration and gates nothing, which is worse than silence.""" + + @endpoint(models={ + "pipeline": Slot( + FakePipeline, + layouts={"txt_encoder": (CONTRACT_PLAIN_BF16,)}, + ), + }) + class Bad: + def setup(self, pipeline: FakePipeline) -> None: + self.pipeline = pipeline + + def generate(self, ctx: RequestContext[_Defaults], p: _In) -> _Out: + return _Out() + + with pytest.raises(ValueError) as excinfo: + extract_specs(Bad) + message = str(excinfo.value) + assert "txt_encoder" in message + # The refusal names the tree, because the author's next action is to pick + # a real path out of it. + assert "text_encoder" in message and "vae" in message + + +def test_a_real_component_key_is_accepted_at_decoration() -> None: + @endpoint(models={ + "pipeline": Slot( + FakePipeline, + layouts={ + "*": (CONTRACT_PLAIN_BF16,), + "text_encoder": (CONTRACT_HF_FP8_BLOCKWISE,), + }, + ), + }) + class Good: + def setup(self, pipeline: FakePipeline) -> None: + self.pipeline = pipeline + + def generate(self, ctx: RequestContext[_Defaults], p: _In) -> _Out: + return _Out() + + specs = extract_specs(Good) + slot = specs[0].slots["pipeline"] + assert slot.layouts == { + "*": ("plain.bf16@1",), + "text_encoder": ("hf.fp8-blockwise@1",), + } + + +# ── 3. the lint refuses what the AST sweep cannot read ─────────────────────── + + +def _run_lint(target: Path) -> subprocess.CompletedProcess: + return subprocess.run( + [sys.executable, str(REPO / "scripts" / "lint_layout_declarations.py"), + str(target)], + capture_output=True, text=True, timeout=120, + ) + + +def test_the_lint_refuses_a_computed_declaration(tmp_path: Path) -> None: + module = tmp_path / "computed.py" + module.write_text(textwrap.dedent( + """ + from gen_worker import Slot + + HANDLES = ["plain.bf16@1"] + + def build(): + return Slot(object, layouts={"*": tuple(h for h in HANDLES)}) + """ + ), encoding="utf-8") + result = _run_lint(module) + assert result.returncode == 1, result.stdout + result.stderr + assert "layouts['*']" in result.stderr + + +def test_the_lint_accepts_literals_and_vocabulary_constants( + tmp_path: Path, +) -> None: + module = tmp_path / "literal.py" + module.write_text(textwrap.dedent( + """ + from gen_worker import Slot + from gen_worker.models.tensor_layout_contract import CONTRACT_PLAIN_BF16 + + SLOT = Slot(object, layouts={ + "*": (CONTRACT_PLAIN_BF16,), + "text_encoder": ("hf.fp8-blockwise@1", CONTRACT_PLAIN_BF16), + }) + """ + ), encoding="utf-8") + result = _run_lint(module) + assert result.returncode == 0, result.stdout + result.stderr + + +def test_the_shipped_tree_passes_the_lint() -> None: + result = _run_lint(REPO / "src") + assert result.returncode == 0, result.stderr + + +# ── 4. the census is a LOWER bound: it reports, it does not refuse ─────────── + + +def test_a_declared_handle_no_decoder_backs_is_reported_not_refused() -> None: + derived = DerivedExecutionLanes( + derivation="gen_worker.discovery.execution_lanes@1", + execution_lanes=("bf16-w16a16+eager",), + contracts=(DerivedContract( + contract=CONTRACT_PLAIN_BF16, + decoder="gen_worker.models.loading:load_from_pretrained", + execution_lanes=("bf16-w16a16+eager",), + composes_lora=True, + ),), + excluded_modules=(), + ) + fn = { + "name": "generate", + "slots": [_slot_to_manifest( + "pipeline", + Slot(FakePipeline, layouts={ + "*": (CONTRACT_PLAIN_BF16,), + "text_encoder": (CONTRACT_HF_FP8_BLOCKWISE,), + }), + family="", components=_components(), + )], + } + # `hf.fp8-blockwise@1` is decoded natively by transformers through + # `quantization_config`, with zero cozy markers. Refusing it would make + # §1.33's own worked example illegal. + assert _census_unbacked_layouts(fn, derived) == ["hf.fp8-blockwise@1"] + + +def test_the_census_is_silent_for_an_undeclared_slot() -> None: + derived = DerivedExecutionLanes( + derivation="gen_worker.discovery.execution_lanes@1", + execution_lanes=(), contracts=(), excluded_modules=(), + ) + fn = { + "name": "generate", + "slots": [_slot_to_manifest( + "pipeline", Slot(FakePipeline), family="", + components=_components())], + } + assert _census_unbacked_layouts(fn, derived) == [] + + +# ── 5. the wire the hub actually reads ─────────────────────────────────────── + + +def test_the_demand_survives_the_endpoint_lock_toml_round_trip() -> None: + """endpoint.lock is TOML and the hub parses JSON, so the declaration + crosses two encoders before any gate sees it. A nested table inside an + array-of-tables is exactly where a TOML encoder's key ordering can bite — + and the hub's `manifestSlotLayoutDoc` decodes `map[string][]string`, so a + shape change here is a silent UNDECLARED on the other side.""" + block = _slot_to_manifest( + "pipeline", + Slot(FakePipeline, selected_by="model", layouts={ + "*": (CONTRACT_PLAIN_BF16,), + "text_encoder": (CONTRACT_HF_FP8_BLOCKWISE, CONTRACT_PLAIN_BF16), + }), + family="sdxl", components=_components(), + ) + doc = {"functions": [{"name": "generate", "slots": [block]}]} + decoded = msgspec.toml.decode(msgspec.toml.encode(doc)) + assert decoded["functions"][0]["slots"][0]["layouts"] == { + "*": ["plain.bf16@1"], + "text_encoder": ["hf.fp8-blockwise@1", "plain.bf16@1"], + }