From bb123f7837824c77075aba59192d8a58d5449d98 Mon Sep 17 00:00:00 2001 From: aryan5v Date: Fri, 31 Jul 2026 10:57:44 -0700 Subject: [PATCH 1/2] [feat]: dispatch trusted graph artifacts with native fallback Run a packaged optimization artifact in place of a repeated block's forward when, and only when, it provably matches -- otherwise run natively. Dispatch attaches to the same model-independent structure capture uses: children of an nn.ModuleList that share a class. No Wan-, LTX-, Cosmos- or Kandinsky-specific conditional exists anywhere in this path. Per stack and per observed input signature it runs the first call natively (which reveals the output signature), recomputes the module's graph fingerprint through the capture module so the value matches what the producer recorded, then selects a bundle whose fingerprint, tensor signatures and declared environment all match. The entry point is called as candidate(module, *args, **kwargs). Trust: executable code is loaded only from FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR, every declared file is re-hashed immediately before import, undeclared files in a bundle are a hard rejection, and bundles are imported under a private module namespace rather than sys.path. Fallback: a missing match, an unloadable bundle, an untraceable module or an exception from the candidate falls back to native execution and records a structured reason; a candidate that raised is demoted, not retried. The optional diagnostics report is metadata only. With FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR unset nothing is wrapped at all -- no forward is patched, no graph is traced, no artifact code is read. Tests: 28 new CPU tests with fake kernels and fake modules; 60 passed across the optimization suites. Verified end to end against a bundle packaged by MotionKernel's packager. Note: the pre-commit mypy hook fails with "FastVideo-v1-dispatch is not a valid Python package name" -- that is the checkout directory name and reproduces on untouched files. yapf, ruff and codespell pass. --- examples/inference/optimizations/README.md | 46 ++ fastvideo/envs.py | 41 ++ fastvideo/optimization/__init__.py | 9 +- fastvideo/optimization/artifact.py | 695 ++++++++++++++++++ fastvideo/optimization/dispatch.py | 494 +++++++++++++ fastvideo/optimization/identity.py | 79 ++ fastvideo/pipelines/composed_pipeline_base.py | 8 +- fastvideo/tests/optimization/test_dispatch.py | 690 +++++++++++++++++ 8 files changed, 2059 insertions(+), 3 deletions(-) create mode 100644 fastvideo/optimization/artifact.py create mode 100644 fastvideo/optimization/dispatch.py create mode 100644 fastvideo/optimization/identity.py create mode 100644 fastvideo/tests/optimization/test_dispatch.py diff --git a/examples/inference/optimizations/README.md b/examples/inference/optimizations/README.md index 70e3bc9fe2..cd139ac3a7 100644 --- a/examples/inference/optimizations/README.md +++ b/examples/inference/optimizations/README.md @@ -52,3 +52,49 @@ Result files use `schema_version: 1` and record wall time, generation time, peak memory, environment identity, optional frames path, and failure reasons. `candidate` mode writes `candidate_result.json`, so it cannot overwrite a previously validated `optimized_result.json`. + +## Generic graph dispatch + +Once a kernel has been packaged as an artifact bundle, FastVideo can run it +without any model-specific code. Point the runtime at a **trusted** directory +of bundles: + +```bash +FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR=/path/to/artifacts \ +FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID=Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ +FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS=/tmp/dispatch.json \ +python examples/inference/optimizations/generation_launcher.py \ + --workload /path/to/motionkernel/workloads/wan_t2v_1.3b_480p.yaml \ + --mode candidate \ + --output-dir /tmp/wan_candidate +``` + +Dispatch attaches to repeated block stacks -- children of an `nn.ModuleList` +that share a class -- exactly as capture does, so no architecture is named +anywhere. Per stack and per observed input signature it runs the first call +natively (which is what reveals the output signature), recomputes the module's +graph fingerprint with the capture tracer, and selects a bundle whose +fingerprint, tensor signatures and declared environment all match. The chosen +entry point is called as `candidate(module, *args, **kwargs)`. + +Everything is fail-safe. A missing match, an unloadable bundle, an untraceable +module or an exception raised by the candidate falls back to native execution +and records a structured reason; the candidate is not retried afterwards. + +| Variable | Default | Purpose | +| --- | --- | --- | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR` | `""` | Trusted artifact root. **Unset means nothing is wrapped at all** and generation is byte-for-byte identical to a build without this feature. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER` | `symbolic` | Tracer used to recompute the fingerprint. `symbolic` does not re-execute the module with real inputs; `export` and `dynamo` do. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID` | profile model id | Model identity matched against each bundle. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION` | `*` | Revision matched against each bundle. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE` | auto | Sharding mode. Auto-detection only resolves a single-rank run; a multi-rank run must declare its mode or no artifact is selected. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES` | `64` | Upper bound on dispatched block stacks. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES` | `8` | Upper bound on resolved input signatures per stack. | +| `FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS` | `""` | Optional path for the structured dispatch/fallback report. | + +The diagnostics report is metadata only: it records each scope, shape key, +decision reason, artifact id, rejection codes and call counts, plus the +registry and runtime identity. It never contains tensor or prompt data. + +The bundle format, its packager and the matching rules are documented in +MotionKernel's `docs/ARTIFACT_BUNDLE.md`. diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 3d68898a55..57ccb297c7 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -45,6 +45,14 @@ FASTVIDEO_OPTIMIZATION_PROFILE_FX_TRACER: str = "auto" FASTVIDEO_OPTIMIZATION_PROFILE_FX_MAX_SCOPES: int = 64 FASTVIDEO_OPTIMIZATION_PROFILE_FX_MAX_SHAPES: int = 8 + FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR: str = "" + FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER: str = "symbolic" + FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES: int = 64 + FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES: int = 8 + FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID: str = "" + FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION: str = "*" + FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE: str = "" + FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS: str = "" FASTVIDEO_TRACE_ACTIVATIONS: bool = False FASTVIDEO_TRACE_LAYERS: str = "" FASTVIDEO_TRACE_STATS: str = "abs_mean,sum" @@ -305,6 +313,39 @@ def maybe_convert_int(value: str | None) -> int | None: "FASTVIDEO_OPTIMIZATION_PROFILE_FX_MAX_SHAPES": lambda: int(os.getenv("FASTVIDEO_OPTIMIZATION_PROFILE_FX_MAX_SHAPES", "8")), + # Trusted directory holding packaged optimization artifacts. Executable + # candidate code is loaded only from here, and only after every declared + # file matches the hash recorded in its manifest. Leaving this unset + # disables graph dispatch entirely: no forward is wrapped and generation + # behaves exactly as it does without the feature. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR": + lambda: os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR", ""), + # Tracer used to recompute a module's graph fingerprint at dispatch time. + # ``symbolic`` is the default because, unlike ``export`` and ``dynamo``, it + # does not re-execute the module with real inputs. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER": + lambda: os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER", "symbolic"), + # Upper bound on dispatched block stacks. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES": + lambda: int(os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES", "64")), + # Upper bound on distinct input signatures resolved per stack. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES": + lambda: int(os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES", "8")), + # Model identity matched against each artifact's declared compatibility. + # Falls back to the optimization profile's model id when unset. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID": + lambda: os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID", ""), + "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION": + lambda: os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION", "*"), + # Sharding mode declared to the matcher. Empty means auto-detect, which + # only resolves a single-rank run; a multi-rank run must declare its mode + # explicitly or no artifact is selected. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE": + lambda: os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE", ""), + # Optional path for the structured dispatch/fallback report. + "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS": + lambda: os.getenv("FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS", ""), + # Enable activation trace hooks if set. "FASTVIDEO_TRACE_ACTIVATIONS": lambda: bool(os.getenv("FASTVIDEO_TRACE_ACTIVATIONS", "0") != "0"), diff --git a/fastvideo/optimization/__init__.py b/fastvideo/optimization/__init__.py index d8938c52aa..7b0cea5bfa 100644 --- a/fastvideo/optimization/__init__.py +++ b/fastvideo/optimization/__init__.py @@ -1,6 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -"""Model-independent optimization discovery helpers.""" +"""Model-independent optimization discovery and dispatch helpers.""" +from fastvideo.optimization.dispatch import (attach_graph_dispatch, detach_graph_dispatch) from fastvideo.optimization.profiler import optimization_profile -__all__ = ["optimization_profile"] +__all__ = [ + "attach_graph_dispatch", + "detach_graph_dispatch", + "optimization_profile", +] diff --git a/fastvideo/optimization/artifact.py b/fastvideo/optimization/artifact.py new file mode 100644 index 0000000000..8476434011 --- /dev/null +++ b/fastvideo/optimization/artifact.py @@ -0,0 +1,695 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Trusted artifact bundle loading for generic graph dispatch. + +A bundle is a directory holding an ``artifact.json`` manifest plus the files it +pins by SHA-256. The manifest is produced by MotionKernel, which owns the +normative schema; this module is the consumer half and deliberately re-derives +the checks rather than importing the producer, exactly as +:mod:`fastvideo.optimization.fx_capture` re-derives the graph fingerprint. + +Three rules hold everywhere in this module: + +* Executable code is imported only from inside the explicitly configured + trusted root. Nothing here reads a path from the manifest and follows it. +* Every declared file is hashed and compared before the entry point is + imported. +* Anything unexpected -- an unknown schema version, a missing field, a changed + byte -- is a hard rejection with a structured reason, never a best-effort + load. +""" + +from __future__ import annotations + +import hashlib +import importlib.util +import json +import re +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from collections.abc import Callable + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +#: Manifest format understood here. A bundle declaring anything else is +#: rejected rather than parsed optimistically. +SUPPORTED_ARTIFACT_SCHEMA_VERSION = 1 + +MANIFEST_FILENAME = "artifact.json" + +#: Wildcard accepted by the string-valued compatibility fields. +ANY = "*" + +#: Synthetic package the bundles are imported under, so they never shadow a +#: real module for the rest of the process. +_MODULE_NAMESPACE = "fastvideo._artifacts" + +_IGNORED_DIRECTORIES = frozenset({"__pycache__"}) +_READ_CHUNK_BYTES = 1 << 20 + +_FINGERPRINT_PATTERN = re.compile(r"^[0-9a-f]{32}$") +_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$") +_SYMBOL_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]{0,127}$") +_RELATIVE_FILE_PATTERN = re.compile(r"^[A-Za-z0-9_][A-Za-z0-9._/-]{0,255}$") + +# Structured rejection reasons. These are logged and exported verbatim, so they +# are part of the contract with the producer. +REASON_FINGERPRINT_MISMATCH = "fingerprint_mismatch" +REASON_INPUT_SIGNATURE_MISMATCH = "input_signature_mismatch" +REASON_OUTPUT_SIGNATURE_MISMATCH = "output_signature_mismatch" +REASON_MODEL_MISMATCH = "model_mismatch" +REASON_REVISION_MISMATCH = "model_revision_mismatch" +REASON_ARCHITECTURE_MISMATCH = "gpu_architecture_mismatch" +REASON_TORCH_VERSION = "torch_version_unsupported" +REASON_CUDA_VERSION = "cuda_version_unsupported" +REASON_TRITON_VERSION = "triton_version_unsupported" +REASON_EXECUTION_MODE = "execution_mode_unsupported" +REASON_DISTRIBUTED_MODE = "distributed_mode_unsupported" +REASON_NOT_PROMOTED = "not_promoted" +REASON_EVIDENCE_INCOMPLETE = "evidence_incomplete" + + +class ArtifactError(ValueError): + """Raised when a bundle is malformed, altered, or unsafe to load.""" + + +def _fail(source: object, location: str, message: str) -> ArtifactError: + return ArtifactError(f"artifact bundle {source!r}: {location}: {message}") + + +def _mapping(value: Any, source: object, location: str) -> dict[str, Any]: + if not isinstance(value, dict) or not value: + raise _fail(source, location, "must be a non-empty object") + for key in value: + if not isinstance(key, str) or not key: + raise _fail(source, location, "keys must be non-empty strings") + return value + + +def _sequence(value: Any, source: object, location: str) -> list[Any]: + if not isinstance(value, list): + raise _fail(source, location, "must be a list") + return value + + +def _text(value: Any, source: object, location: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise _fail(source, location, "must be a non-empty string") + return value + + +def _bool(value: Any, source: object, location: str) -> bool: + if not isinstance(value, bool): + raise _fail(source, location, "must be a bool") + return value + + +def _pattern(value: Any, pattern: re.Pattern[str], source: object, location: str, description: str) -> str: + text = _text(value, source, location) + if not pattern.fullmatch(text): + raise _fail(source, location, f"must be {description}") + return text + + +def _relative_path(value: Any, source: object, location: str) -> str: + text = _pattern(value, _RELATIVE_FILE_PATTERN, source, location, "a relative POSIX path") + if any(part in ("", ".", "..") for part in text.split("/")): + raise _fail(source, location, "must not contain empty or relative segments") + return text + + +def parse_version(text: str | None) -> tuple[int, ...] | None: + """Parse a leading dotted-numeric version, ignoring local/pre-release tags. + + ``"2.8.0+cu128"`` and ``"2.8.0a0"`` both parse to ``(2, 8, 0)``. ``None`` is + returned when there is no numeric prefix, which callers treat as "cannot + compare" rather than "compatible". + """ + if text is None: + return None + parts: list[int] = [] + for chunk in str(text).strip().split("."): + digits = "" + for character in chunk: + if not character.isdigit(): + break + digits += character + if not digits: + break + parts.append(int(digits)) + return tuple(parts) or None + + +def _compare_versions(left: tuple[int, ...], right: tuple[int, ...]) -> int: + width = max(len(left), len(right)) + padded_left = left + (0, ) * (width - len(left)) + padded_right = right + (0, ) * (width - len(right)) + if padded_left < padded_right: + return -1 + return 0 if padded_left == padded_right else 1 + + +@dataclass(frozen=True) +class VersionRange: + """An inclusive-minimum, exclusive-maximum version window.""" + + minimum: str | None = None + maximum_exclusive: str | None = None + + @classmethod + def from_dict(cls, raw: Any, *, source: object, location: str) -> VersionRange: + if raw is None: + return cls() + if not isinstance(raw, dict): + raise _fail(source, location, "must be an object") + unknown = sorted(set(raw) - {"min", "max_exclusive"}) + if unknown: + raise _fail(source, location, f"unknown field(s) {unknown}") + minimum = raw.get("min") + maximum = raw.get("max_exclusive") + if minimum is not None: + minimum = _text(minimum, source, f"{location}.min") + if maximum is not None: + maximum = _text(maximum, source, f"{location}.max_exclusive") + return cls(minimum=minimum, maximum_exclusive=maximum) + + @property + def unbounded(self) -> bool: + return self.minimum is None and self.maximum_exclusive is None + + def contains(self, version: str | None) -> bool: + """Whether ``version`` satisfies this range. + + A bound that cannot be evaluated -- because the runtime version is + missing or unparsable -- is never assumed to hold. + """ + if self.unbounded: + return True + observed = parse_version(version) + if observed is None: + return False + if self.minimum is not None: + low = parse_version(self.minimum) + if low is None or _compare_versions(observed, low) < 0: + return False + if self.maximum_exclusive is not None: + high = parse_version(self.maximum_exclusive) + if high is None or _compare_versions(observed, high) >= 0: + return False + return True + + def describe(self) -> str: + if self.unbounded: + return "any" + low = self.minimum if self.minimum is not None else "any" + high = self.maximum_exclusive if self.maximum_exclusive is not None else "any" + return f">={low},<{high}" + + +def signature_key(meta: dict[str, Any]) -> tuple[Any, ...]: + """The comparable identity of one tensor signature. + + Built from the capture module's tensor metadata so a live invocation and a + packaged artifact are compared on exactly the same fields. ``name`` is + excluded: argument naming is a runtime detail, not part of the layout. + """ + return ( + tuple(int(dim) for dim in meta.get("shape", ())), + tuple(int(step) for step in meta.get("stride", ())), + str(meta.get("dtype", "")), + str(meta.get("device_type", "")), + bool(meta.get("requires_grad", False)), + ) + + +def _signature_keys(raw: Any, source: object, location: str) -> tuple[tuple[Any, ...], ...]: + items = _sequence(raw, source, location) + if not items: + raise _fail(source, location, "must be a non-empty list") + keys = [] + for index, item in enumerate(items): + entry = _mapping(item, source, f"{location}[{index}]") + shape = _sequence(entry.get("shape"), source, f"{location}[{index}].shape") + stride = _sequence(entry.get("stride"), source, f"{location}[{index}].stride") + for values, name in ((shape, "shape"), (stride, "stride")): + for position, value in enumerate(values): + if isinstance(value, bool) or not isinstance(value, int): + raise _fail( + source, + f"{location}[{index}].{name}[{position}]", + "must be an integer", + ) + if len(shape) != len(stride): + raise _fail( + source, + f"{location}[{index}].stride", + "must have the same length as shape", + ) + _text(entry.get("dtype"), source, f"{location}[{index}].dtype") + _text(entry.get("device_type"), source, f"{location}[{index}].device_type") + keys.append(signature_key(entry)) + return tuple(keys) + + +@dataclass(frozen=True) +class ArtifactFile: + """One bundled file, pinned by size and content hash.""" + + path: str + sha256: str + size: int + + +@dataclass(frozen=True) +class ArtifactManifest: + """The subset of a bundle manifest this runtime acts on.""" + + artifact_id: str + graph_fingerprint: str + operation_name: str + parent_module: str + input_keys: tuple[tuple[Any, ...], ...] + output_keys: tuple[tuple[Any, ...], ...] + entry_file: str + entry_symbol: str + files: tuple[ArtifactFile, ...] + model_id: str + model_revision: str + gpu_architectures: tuple[str, ...] + torch_range: VersionRange + cuda_range: VersionRange + triton_range: VersionRange + execution_modes: tuple[str, ...] + distributed_modes: tuple[str, ...] + promotion_decision: str + evidence_passed: bool + speedup: float + directory: Path + + @classmethod + def from_dict(cls, raw_value: Any, *, directory: Path) -> ArtifactManifest: + source = str(directory) + raw = _mapping(raw_value, source, "top level") + version = raw.get("schema_version") + if isinstance(version, bool) or not isinstance(version, int): + raise _fail(source, "schema_version", "must be an integer") + if version != SUPPORTED_ARTIFACT_SCHEMA_VERSION: + raise _fail( + source, + "schema_version", + f"unsupported version {version}; expected " + f"{SUPPORTED_ARTIFACT_SCHEMA_VERSION}", + ) + + operation = _mapping(raw.get("operation"), source, "operation") + signature = _mapping(raw.get("signature"), source, "signature") + entry_point = _mapping(raw.get("entry_point"), source, "entry_point") + compatibility = _mapping(raw.get("compatibility"), source, "compatibility") + evidence = _mapping(raw.get("evidence"), source, "evidence") + promotion = _mapping(raw.get("promotion"), source, "promotion") + + files = [] + for index, item in enumerate(_sequence(raw.get("files"), source, "files")): + entry = _mapping(item, source, f"files[{index}]") + size = entry.get("bytes") + if isinstance(size, bool) or not isinstance(size, int) or size < 0: + raise _fail(source, f"files[{index}].bytes", "must be a non-negative integer") + files.append( + ArtifactFile( + path=_relative_path(entry.get("path"), source, f"files[{index}].path"), + sha256=_pattern( + entry.get("sha256"), + _SHA256_PATTERN, + source, + f"files[{index}].sha256", + "64 lowercase hex characters", + ), + size=size, + )) + if not files: + raise _fail(source, "files", "must be a non-empty list") + paths = [item.path for item in files] + if len(paths) != len(set(paths)): + raise _fail(source, "files", "contains duplicate paths") + + entry_file = _relative_path(entry_point.get("file"), source, "entry_point.file") + if entry_file not in paths: + raise _fail(source, "entry_point.file", f"{entry_file!r} is not a declared file") + if not entry_file.endswith(".py"): + raise _fail(source, "entry_point.file", "must be a .py file") + + benchmark = _mapping(evidence.get("benchmark"), source, "evidence.benchmark") + generation = _mapping(evidence.get("generation"), source, "evidence.generation") + speedup = benchmark.get("speedup") + if isinstance(speedup, bool) or not isinstance(speedup, int | float): + raise _fail(source, "evidence.benchmark.speedup", "must be a number") + + return cls( + artifact_id=_text(raw.get("artifact_id"), source, "artifact_id"), + graph_fingerprint=_pattern( + operation.get("graph_fingerprint"), + _FINGERPRINT_PATTERN, + source, + "operation.graph_fingerprint", + "32 lowercase hex characters", + ), + operation_name=_text(operation.get("name"), source, "operation.name"), + parent_module=_text(operation.get("parent_module"), source, "operation.parent_module"), + input_keys=_signature_keys(signature.get("inputs"), source, "signature.inputs"), + output_keys=_signature_keys(signature.get("outputs"), source, "signature.outputs"), + entry_file=entry_file, + entry_symbol=_pattern( + entry_point.get("symbol"), + _SYMBOL_PATTERN, + source, + "entry_point.symbol", + "a Python identifier", + ), + files=tuple(files), + model_id=_text(compatibility.get("model_id"), source, "compatibility.model_id"), + model_revision=_text( + compatibility.get("model_revision"), + source, + "compatibility.model_revision", + ), + gpu_architectures=tuple( + _text(item, source, f"compatibility.gpu_architectures[{index}]") for index, item in enumerate( + _sequence( + compatibility.get("gpu_architectures"), + source, + "compatibility.gpu_architectures", + ))), + torch_range=VersionRange.from_dict( + compatibility.get("torch"), + source=source, + location="compatibility.torch", + ), + cuda_range=VersionRange.from_dict( + compatibility.get("cuda"), + source=source, + location="compatibility.cuda", + ), + triton_range=VersionRange.from_dict( + compatibility.get("triton"), + source=source, + location="compatibility.triton", + ), + execution_modes=tuple( + _text(item, source, f"compatibility.execution_modes[{index}]") for index, item in enumerate( + _sequence( + compatibility.get("execution_modes"), + source, + "compatibility.execution_modes", + ))), + distributed_modes=tuple( + _text(item, source, f"compatibility.distributed_modes[{index}]") for index, item in enumerate( + _sequence( + compatibility.get("distributed_modes"), + source, + "compatibility.distributed_modes", + ))), + promotion_decision=_text(promotion.get("decision"), source, "promotion.decision"), + evidence_passed=(_bool(benchmark.get("passed"), source, "evidence.benchmark.passed") + and _bool(generation.get("passed"), source, "evidence.generation.passed")), + speedup=float(speedup), + directory=directory, + ) + + +@dataclass(frozen=True) +class RuntimeProfile: + """The environment this process is executing in.""" + + model_id: str + model_revision: str + gpu_architecture: str + torch_version: str + cuda_version: str | None = None + triton_version: str | None = None + execution_mode: str = "inference" + distributed_mode: str = "single" + + @classmethod + def detect( + cls, + *, + model_id: str, + model_revision: str = ANY, + execution_mode: str = "inference", + distributed_mode: str = "single", + ) -> RuntimeProfile: + """Read torch/CUDA/Triton identity from the running process.""" + import torch + + architecture = "cpu" + cuda_version = None + if torch.cuda.is_available(): + index = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(index) + architecture = f"sm{major}{minor}" + cuda_version = getattr(torch.version, "cuda", None) + try: + import triton + + triton_version = getattr(triton, "__version__", None) + except Exception: # noqa: BLE001 - Triton is optional on every platform + triton_version = None + return cls( + model_id=model_id, + model_revision=model_revision, + gpu_architecture=architecture, + torch_version=str(getattr(torch, "__version__", "")), + cuda_version=cuda_version, + triton_version=triton_version, + execution_mode=execution_mode, + distributed_mode=distributed_mode, + ) + + +def _wildcard_equal(declared: str, observed: str) -> bool: + return declared in (ANY, observed) + + +def check_compatibility( + manifest: ArtifactManifest, + *, + graph_fingerprint: str, + input_keys: tuple[tuple[Any, ...], ...], + output_keys: tuple[tuple[Any, ...], ...], + runtime: RuntimeProfile, +) -> str | None: + """Return the reason ``manifest`` cannot serve this call, or ``None``. + + Ordering is cheapest-first: graph identity, then tensor layout, then the + declared environment window. + """ + if manifest.graph_fingerprint != graph_fingerprint: + return REASON_FINGERPRINT_MISMATCH + if manifest.input_keys != input_keys: + return REASON_INPUT_SIGNATURE_MISMATCH + if manifest.output_keys != output_keys: + return REASON_OUTPUT_SIGNATURE_MISMATCH + if not _wildcard_equal(manifest.model_id, runtime.model_id): + return REASON_MODEL_MISMATCH + if not _wildcard_equal(manifest.model_revision, runtime.model_revision): + return REASON_REVISION_MISMATCH + if not any(_wildcard_equal(item, runtime.gpu_architecture) for item in manifest.gpu_architectures): + return REASON_ARCHITECTURE_MISMATCH + if not manifest.torch_range.contains(runtime.torch_version): + return REASON_TORCH_VERSION + if not manifest.cuda_range.contains(runtime.cuda_version): + return REASON_CUDA_VERSION + if not manifest.triton_range.contains(runtime.triton_version): + return REASON_TRITON_VERSION + if runtime.execution_mode not in manifest.execution_modes: + return REASON_EXECUTION_MODE + if runtime.distributed_mode not in manifest.distributed_modes: + return REASON_DISTRIBUTED_MODE + if manifest.promotion_decision != "promoted": + return REASON_NOT_PROMOTED + if not manifest.evidence_passed: + return REASON_EVIDENCE_INCOMPLETE + return None + + +def file_sha256(path: Path) -> str: + """Content hash of one file, streamed so large bundles stay cheap.""" + digest = hashlib.sha256() + with open(path, "rb") as handle: + while True: + chunk = handle.read(_READ_CHUNK_BYTES) + if not chunk: + break + digest.update(chunk) + return digest.hexdigest() + + +def _bundle_contents(directory: Path) -> set[str]: + present = set() + for path in directory.rglob("*"): + if path.is_dir(): + continue + relative = path.relative_to(directory) + if any(part in _IGNORED_DIRECTORIES for part in relative.parts): + continue + present.add(relative.as_posix()) + return present + + +def verify_bundle(directory: Path) -> ArtifactManifest: + """Parse a bundle and confirm every declared file is byte-for-byte intact. + + Undeclared files are refused as well: an attacker who can drop an extra + module next to a signed entry point could otherwise have it imported. + """ + manifest_path = directory / MANIFEST_FILENAME + if not manifest_path.is_file(): + raise _fail(str(directory), MANIFEST_FILENAME, "not found") + try: + raw = json.loads(manifest_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise _fail(str(directory), MANIFEST_FILENAME, f"invalid JSON: {exc}") from exc + manifest = ArtifactManifest.from_dict(raw, directory=directory) + + declared = {item.path for item in manifest.files} + undeclared = sorted(_bundle_contents(directory) - declared - {MANIFEST_FILENAME}) + if undeclared: + raise _fail(str(directory), "files", f"undeclared file(s) {undeclared}") + + for entry in sorted(manifest.files, key=lambda item: item.path): + path = directory / entry.path + if not path.is_file(): + raise _fail(str(directory), "files", f"{entry.path!r} is missing") + size = path.stat().st_size + if size != entry.size: + raise _fail( + str(directory), + "files", + f"{entry.path!r} is {size} bytes, manifest records {entry.size}", + ) + actual = file_sha256(path) + if actual != entry.sha256: + raise _fail( + str(directory), + "files", + f"{entry.path!r} hash {actual} does not match manifest {entry.sha256}", + ) + return manifest + + +def _resolve_inside(root: Path, candidate: Path) -> Path: + """Resolve ``candidate`` and require it to stay under ``root``.""" + resolved_root = root.resolve(strict=True) + resolved = candidate.resolve(strict=True) + if resolved != resolved_root and resolved_root not in resolved.parents: + raise ArtifactError(f"artifact bundle {str(candidate)!r}: resolves outside the trusted " + f"root {str(resolved_root)!r}") + return resolved + + +def load_entry_point(manifest: ArtifactManifest, *, trusted_root: Path) -> Callable[..., Any]: + """Re-verify a bundle and import its candidate callable. + + The bundle is verified again here even if the caller already validated it: + the gap between validation and import is precisely where a swapped file + would land. + """ + directory = _resolve_inside(trusted_root, manifest.directory) + verified = verify_bundle(directory) + if verified.artifact_id != manifest.artifact_id: + raise _fail( + str(directory), + "artifact_id", + f"changed from {manifest.artifact_id!r} to {verified.artifact_id!r} " + "since validation", + ) + + entry_file = _resolve_inside(directory, directory / verified.entry_file) + module_name = f"{_MODULE_NAMESPACE}.{re.sub(r'[^A-Za-z0-9_]', '_', verified.artifact_id)}" + spec = importlib.util.spec_from_file_location(module_name, entry_file) + if spec is None or spec.loader is None: + raise _fail(str(directory), "entry_point", f"cannot load {verified.entry_file!r}") + module = importlib.util.module_from_spec(spec) + # Registered before execution so the module can look itself up, and removed + # again on failure so a half-initialized module is never reachable. + sys.modules[module_name] = module + try: + spec.loader.exec_module(module) + except Exception as exc: # noqa: BLE001 - untrusted code, any failure is a rejection + sys.modules.pop(module_name, None) + raise _fail( + str(directory), + "entry_point", + f"importing {verified.entry_file!r} raised {type(exc).__name__}", + ) from exc + + candidate = getattr(module, verified.entry_symbol, None) + if candidate is None or not callable(candidate): + sys.modules.pop(module_name, None) + raise _fail( + str(directory), + "entry_point", + f"{verified.entry_symbol!r} is not a callable in {verified.entry_file!r}", + ) + return candidate + + +def discover_bundles(root: Path) -> list[Path]: + """List bundle directories under ``root``: the root itself plus children.""" + if not root.is_dir(): + return [] + found = [] + if (root / MANIFEST_FILENAME).is_file(): + found.append(root) + for child in sorted(root.iterdir()): + if child.is_dir() and (child / MANIFEST_FILENAME).is_file(): + found.append(child) + return found + + +class ArtifactRegistry: + """Every valid bundle found in the trusted artifact directory. + + Loading is eager and one-shot: bundles are parsed and hashed when the + registry is built, so a corrupt artifact surfaces before generation starts + rather than mid-run. Bundles that fail validation are recorded in + :attr:`errors` and skipped; they never disable the ones that passed. + """ + + def __init__(self, root: Path | None) -> None: + self.root = root + self.manifests: list[ArtifactManifest] = [] + self.errors: list[str] = [] + if root is None: + return + if not root.is_dir(): + self.errors.append(f"artifact directory {str(root)!r}: not a directory") + return + for path in discover_bundles(root): + try: + self.manifests.append(verify_bundle(path)) + except ArtifactError as exc: + self.errors.append(str(exc)) + + @property + def enabled(self) -> bool: + return bool(self.manifests) + + def candidates_for(self, input_keys: tuple[tuple[Any, ...], ...]) -> list[ArtifactManifest]: + """Bundles whose input layout matches, before any graph is traced. + + This is the cheap pre-filter that keeps dispatch from tracing a module + when the store holds nothing shaped like the live call. + """ + return [item for item in self.manifests if item.input_keys == input_keys] + + def summary(self) -> dict[str, Any]: + return { + "root": str(self.root) if self.root is not None else None, + "loaded": len(self.manifests), + "artifact_ids": sorted(item.artifact_id for item in self.manifests), + "errors": list(self.errors), + } diff --git a/fastvideo/optimization/dispatch.py b/fastvideo/optimization/dispatch.py new file mode 100644 index 0000000000..6f7816fa25 --- /dev/null +++ b/fastvideo/optimization/dispatch.py @@ -0,0 +1,494 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generic graph dispatch: run a packaged kernel, or fall back to native. + +Dispatch attaches to the same model-independent structure the capture pipeline +uses -- stacks of identically-typed children under an ``nn.ModuleList`` -- and +decides, per stack and per observed input signature, whether a trusted artifact +may replace the native forward. Nothing in this module knows the name of a +model, a block class, or an architecture; a new model is supported by +publishing an artifact, never by editing code here. + +The decision sequence for one scope and input signature is: + +1. The first call always runs natively. It is what reveals the output + signature, which is part of the artifact's identity. +2. The registry is pre-filtered on the input layout. If nothing in the store + is shaped like this call, no graph is ever traced. +3. The module is traced once to recompute its graph fingerprint, using the + capture module so the value matches what the producer recorded. +4. Compatibility is checked, the winning bundle is re-verified and imported, + and the callable is cached for every later call with that signature. + +Every step is guarded. A failure at any point demotes the signature to native +execution permanently and records a structured reason; it never propagates to +the caller. With no artifact directory configured, nothing is attached at all +and the model runs byte-for-byte as it does today. + +Candidate calling convention +---------------------------- +A bundle's entry point is called as ``candidate(module, *args, **kwargs)``: +the native module comes first, then the exact arguments its ``forward`` was +given. Passing the module is what lets one artifact serve every block in a +stack -- the kernel reads the parameters it needs from the module it was handed +instead of the loader having to know which parameters exist. +""" + +from __future__ import annotations + +import json +from collections import defaultdict +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from collections.abc import Callable + +from torch import nn + +from fastvideo import envs +from fastvideo.logger import init_logger +from fastvideo.optimization.artifact import ( + ArtifactManifest, + ArtifactRegistry, + RuntimeProfile, + check_compatibility, + load_entry_point, + signature_key, +) +from fastvideo.optimization.fx_capture import default_capture_targets +from fastvideo.optimization.identity import ( + graph_identity, + input_signatures, + output_signatures, + shape_key_for, +) + +logger = init_logger(__name__) + +#: Diagnostic report format. Bumped independently of the capture and profiler +#: schemas so a consumer can tell the three apart. +DISPATCH_SCHEMA_VERSION = 1 + +# Structured fallback reasons. Consumers group on these, so they are stable. +REASON_SELECTED = "artifact_selected" +REASON_NO_SIGNATURE_MATCH = "no_artifact_for_input_signature" +REASON_NO_COMPATIBLE_ARTIFACT = "no_compatible_artifact" +REASON_NO_TENSOR_INPUTS = "no_tensor_inputs" +_REASON_IDENTITY_PREFIX = "graph_identity_unavailable" +_REASON_LOAD_PREFIX = "artifact_load_failed" +_REASON_RUNTIME_PREFIX = "candidate_runtime_error" +_REASON_DECISION_PREFIX = "decision_failed" + + +@dataclass +class _Decision: + """What to do for one scope and input signature, and why.""" + + candidate: Callable[..., Any] | None + reason: str + artifact_id: str | None = None + rejections: tuple[str, ...] = () + #: Calls made *after* this decision was resolved. The one call that + #: produced it ran natively and is not counted here. + calls: int = 0 + candidate_calls: int = 0 + runtime_fallbacks: int = 0 + + def as_dict(self, scope: str, shape_key: str) -> dict[str, Any]: + return { + "scope": scope, + "shape_key": shape_key, + "reason": self.reason, + "artifact_id": self.artifact_id, + "rejections": list(self.rejections), + "calls": self.calls, + "candidate_calls": self.candidate_calls, + "runtime_fallbacks": self.runtime_fallbacks, + "active": self.candidate is not None, + } + + +@dataclass +class _Wrapper: + """One patched module, remembered so ``detach`` can restore it exactly.""" + + scope: str + module: nn.Module + native_forward: Callable[..., Any] + had_own_forward: bool = False + + +class GraphDispatchSession: + """Route repeated module stacks through trusted artifacts when possible.""" + + def __init__( + self, + registry: ArtifactRegistry, + runtime: RuntimeProfile, + *, + tracer: str = "symbolic", + max_scopes: int = 64, + max_shape_variants: int = 8, + ) -> None: + self.registry = registry + self.runtime = runtime + self.tracer = tracer + self.max_scopes = max_scopes + self.max_shape_variants = max_shape_variants + self._wrappers: list[_Wrapper] = [] + self._scopes: set[str] = set() + self._decisions: dict[tuple[str, str], _Decision] = {} + self._dropped_variants: dict[str, int] = defaultdict(int) + self._dropped_scopes = 0 + self._errors: list[str] = [] + + # -- attach --------------------------------------------------------- + + def attach_modules(self, modules: dict[str, Any] | None) -> int: + """Wrap every repeated block stack in a pipeline's module mapping.""" + if not modules or not self.registry.enabled: + return 0 + wrapped = 0 + for name, module in modules.items(): + if isinstance(module, nn.Module): + wrapped += self.attach(module, prefix=str(name)) + return wrapped + + def attach(self, root: nn.Module, *, prefix: str = "") -> int: + """Wrap every repeated block stack under ``root``.""" + try: + targets = default_capture_targets(root) + except Exception as exc: # noqa: BLE001 - dispatch never breaks generation + self._record_exception("target_selection_failed", exc) + return 0 + wrapped = 0 + for scope, module in targets: + qualified = f"{prefix}.{scope}" if prefix else scope + if qualified not in self._scopes and len(self._scopes) >= self.max_scopes: + self._dropped_scopes += 1 + continue + try: + self._install(qualified, module) + self._scopes.add(qualified) + wrapped += 1 + except Exception as exc: # noqa: BLE001 + self._record_exception("install_failed", exc, scope=qualified) + return wrapped + + def _install(self, scope: str, module: nn.Module) -> None: + """Replace ``module.forward`` with the dispatching wrapper.""" + native_forward = module.forward + wrapper = _Wrapper( + scope=scope, + module=module, + native_forward=native_forward, + had_own_forward="forward" in vars(module), + ) + + def dispatching_forward(*args: Any, **kwargs: Any) -> Any: + return self._dispatch(wrapper, args, kwargs) + + module.forward = dispatching_forward # type: ignore[method-assign] + self._wrappers.append(wrapper) + + def detach(self) -> None: + """Restore every native forward. Safe to call more than once.""" + for wrapper in reversed(self._wrappers): + try: + if wrapper.had_own_forward: + wrapper.module.forward = wrapper.native_forward # type: ignore[method-assign] + else: + # The wrapper shadowed the class-level method; deleting the + # instance attribute restores the original binding exactly. + del wrapper.module.forward # type: ignore[misc] + except Exception as exc: # noqa: BLE001 + self._record_exception("detach_failed", exc, scope=wrapper.scope) + self._wrappers.clear() + + # -- dispatch ------------------------------------------------------- + + def _dispatch(self, wrapper: _Wrapper, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: + native = wrapper.native_forward + try: + input_metas = input_signatures(args, kwargs) + shape_key = shape_key_for(input_metas) + except Exception as exc: # noqa: BLE001 + self._record_exception("shape_key_failed", exc, scope=wrapper.scope) + return native(*args, **kwargs) + + key = (wrapper.scope, shape_key) + decision = self._decisions.get(key) + if decision is None: + # First call for this signature: run native, learn the output + # layout, then decide once for every later call. + output = native(*args, **kwargs) + self._decisions[key] = self._decide(wrapper, args, kwargs, output, input_metas) + return output + + decision.calls += 1 + if decision.candidate is None: + return native(*args, **kwargs) + try: + result = decision.candidate(wrapper.module, *args, **kwargs) + except Exception as exc: # noqa: BLE001 - untrusted candidate code + # Demote permanently: a candidate that raised once is not trusted + # to be retried thousands of times over the rest of the run. + decision.runtime_fallbacks += 1 + decision.candidate = None + decision.reason = f"{_REASON_RUNTIME_PREFIX}:{type(exc).__name__}" + logger.warning( + "Artifact %s failed at runtime for %s; falling back to native " + "execution for the rest of this run", + decision.artifact_id, + wrapper.scope, + exc_info=True, + ) + return native(*args, **kwargs) + decision.candidate_calls += 1 + return result + + def _decide( + self, + wrapper: _Wrapper, + args: tuple[Any, ...], + kwargs: dict[str, Any], + output: Any, + input_metas: list[dict[str, Any]], + ) -> _Decision: + """Resolve one signature to an artifact or to a native fallback.""" + try: + return self._resolve(wrapper, args, kwargs, output, input_metas) + except Exception as exc: # noqa: BLE001 - resolution is best effort + self._record_exception("decide_failed", exc, scope=wrapper.scope) + return _Decision(None, f"{_REASON_DECISION_PREFIX}:{type(exc).__name__}") + + def _resolve( + self, + wrapper: _Wrapper, + args: tuple[Any, ...], + kwargs: dict[str, Any], + output: Any, + input_metas: list[dict[str, Any]], + ) -> _Decision: + scope = wrapper.scope + if self._variant_count(scope) >= self.max_shape_variants: + self._dropped_variants[scope] += 1 + return _Decision(None, "shape_variant_budget_exhausted") + + if not input_metas: + return _Decision(None, REASON_NO_TENSOR_INPUTS) + input_keys = tuple(signature_key(meta) for meta in input_metas) + candidates = self.registry.candidates_for(input_keys) + if not candidates: + return _Decision(None, REASON_NO_SIGNATURE_MATCH) + + try: + region = graph_identity( + wrapper.module, + args, + kwargs, + output, + scope=scope, + tracer=self.tracer, + ) + except Exception as exc: # noqa: BLE001 - untraceable modules stay native + reason = f"{_REASON_IDENTITY_PREFIX}:{type(exc).__name__}" + logger.info("Graph identity unavailable for %s: %s", scope, reason) + return _Decision(None, reason) + + fingerprint = str(region.get("fingerprint", "")) + output_keys = tuple(signature_key(meta) for meta in output_signatures(output)) + matched: list[ArtifactManifest] = [] + rejections: list[str] = [] + for manifest in candidates: + reason = check_compatibility( + manifest, + graph_fingerprint=fingerprint, + input_keys=input_keys, + output_keys=output_keys, + runtime=self.runtime, + ) + if reason is None: + matched.append(manifest) + else: + rejections.append(f"{manifest.artifact_id}:{reason}") + if not matched: + return _Decision( + None, + REASON_NO_COMPATIBLE_ARTIFACT, + rejections=tuple(sorted(rejections)), + ) + + best = max(matched, key=lambda item: (item.speedup, item.artifact_id)) + rejections.extend(f"{item.artifact_id}:not_selected" for item in matched if item is not best) + root = self.registry.root + if root is None: + return _Decision(None, f"{_REASON_LOAD_PREFIX}:NoTrustedRoot") + try: + candidate = load_entry_point(best, trusted_root=root) + except Exception as exc: # noqa: BLE001 - a bad bundle must not stop generation + logger.warning("Artifact %s could not be loaded; staying native", best.artifact_id, exc_info=True) + return _Decision( + None, + f"{_REASON_LOAD_PREFIX}:{type(exc).__name__}", + artifact_id=best.artifact_id, + rejections=tuple(sorted(rejections)), + ) + logger.info( + "Dispatching %s to artifact %s (fingerprint %s)", + scope, + best.artifact_id, + fingerprint, + ) + return _Decision( + candidate, + REASON_SELECTED, + artifact_id=best.artifact_id, + rejections=tuple(sorted(rejections)), + ) + + def _variant_count(self, scope: str) -> int: + return sum(1 for existing_scope, _ in self._decisions if existing_scope == scope) + + # -- diagnostics ---------------------------------------------------- + + def _record_error(self, message: str) -> None: + if len(self._errors) < 64: + self._errors.append(message) + + def _record_exception(self, code: str, exc: Exception, *, scope: str | None = None) -> None: + # Exception text can carry argument reprs; only the type is recorded. + location = f"[{scope}]" if scope else "" + self._record_error(f"{code}{location}:{type(exc).__name__}") + + def diagnostics(self) -> dict[str, Any]: + """A metadata-only report of every dispatch decision made.""" + decisions = [ + decision.as_dict(scope, shape_key) for (scope, shape_key), decision in sorted(self._decisions.items()) + ] + reason_counts: dict[str, int] = defaultdict(int) + for decision in decisions: + reason_counts[str(decision["reason"])] += 1 + return { + "dispatch": { + "dispatch_schema_version": DISPATCH_SCHEMA_VERSION, + "registry": self.registry.summary(), + "runtime": { + "model_id": self.runtime.model_id, + "model_revision": self.runtime.model_revision, + "gpu_architecture": self.runtime.gpu_architecture, + "torch_version": self.runtime.torch_version, + "cuda_version": self.runtime.cuda_version, + "triton_version": self.runtime.triton_version, + "execution_mode": self.runtime.execution_mode, + "distributed_mode": self.runtime.distributed_mode, + }, + "tracer": self.tracer, + "scopes": sorted(self._scopes), + "dropped_scopes": self._dropped_scopes, + "dropped_shape_variants": dict(sorted(self._dropped_variants.items())), + "errors": list(self._errors), + "reason_counts": dict(sorted(reason_counts.items())), + }, + "decisions": decisions, + } + + def write_diagnostics(self, path: str | Path) -> Path | None: + """Write the diagnostic report, returning ``None`` if that failed.""" + output = Path(path).expanduser() + try: + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_suffix(output.suffix + ".tmp") + temporary.write_text( + json.dumps(self.diagnostics(), indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary.replace(output) + except Exception: # noqa: BLE001 - diagnostics must never break a run + logger.warning("Could not write dispatch diagnostics to %s", output, exc_info=True) + return None + return output + + +def _execution_mode(modules: dict[str, Any] | None) -> str: + """Report ``training`` when any hooked module is in training mode.""" + for module in (modules or {}).values(): + if isinstance(module, nn.Module) and module.training: + return "training" + return "inference" + + +def _distributed_mode() -> str: + """Report the sharding mode this process runs in. + + A multi-rank run whose mode is not declared explicitly reports + ``unspecified``, which matches no artifact: an unsharded kernel silently + applied to a sharded module would be a correctness bug, so the ambiguous + case fails closed. + """ + declared = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE + if declared: + return declared + try: + import torch.distributed as distributed + + if distributed.is_available() and distributed.is_initialized(): + return "single" if distributed.get_world_size() == 1 else "unspecified" + except Exception: # noqa: BLE001 - absence of distributed means single process + return "single" + return "single" + + +def attach_graph_dispatch(modules: dict[str, Any] | None) -> GraphDispatchSession | None: + """Attach generic dispatch, or return ``None`` when it is not configured. + + Returning ``None`` is the zero-effect path: no forward is wrapped, no graph + is traced, and no artifact code is read. This is what makes an unset + ``FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR`` behaviorally identical to a build + without this feature. + """ + configured = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR + if not configured: + return None + try: + root = Path(configured).expanduser() + registry = ArtifactRegistry(root) + for error in registry.errors: + logger.warning("Skipping artifact bundle: %s", error) + if not registry.enabled: + logger.info("No usable artifacts in %s; running natively", root) + return None + runtime = RuntimeProfile.detect( + model_id=(envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID or envs.FASTVIDEO_OPTIMIZATION_PROFILE_MODEL_ID), + model_revision=envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION, + execution_mode=_execution_mode(modules), + distributed_mode=_distributed_mode(), + ) + session = GraphDispatchSession( + registry, + runtime, + tracer=envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER, + max_scopes=envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES, + max_shape_variants=envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES, + ) + wrapped = session.attach_modules(modules) + except Exception: # noqa: BLE001 - dispatch setup never breaks generation + logger.warning("Graph dispatch could not be started; running natively", exc_info=True) + return None + if not wrapped: + logger.info("Graph dispatch found no repeated module stacks; running natively") + return None + logger.info("Graph dispatch attached to %d module(s) from %s", wrapped, registry.root) + return session + + +def detach_graph_dispatch(session: GraphDispatchSession | None) -> None: + """Restore every native forward and write diagnostics if requested.""" + if session is None: + return + try: + session.detach() + except Exception: # noqa: BLE001 + logger.warning("Graph dispatch detach failed", exc_info=True) + output = envs.FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS + if output: + session.write_diagnostics(output) diff --git a/fastvideo/optimization/identity.py b/fastvideo/optimization/identity.py new file mode 100644 index 0000000000..882a1d3709 --- /dev/null +++ b/fastvideo/optimization/identity.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Graph identity for a single observed module invocation. + +Dispatch has to answer one question about a live call: *which captured region +is this?* The answer must be the same value the producer recorded, or no +artifact would ever match. + +That is why this module reuses :mod:`fastvideo.optimization.fx_capture`'s own +helpers rather than re-deriving tensor metadata, shape keys or fingerprints. +Those helpers are the canonical implementation; a second implementation here +would be one refactor away from silently disagreeing with the exports it is +supposed to match, and a fingerprint that disagrees is indistinguishable from +"no artifact available". +""" + +from __future__ import annotations + +from typing import Any + +from torch import nn + +from fastvideo.optimization.fx_capture import ( + FXCaptureSession, + _input_metas, + _output_metas, + _shape_key, + _Scope, +) + + +def input_signatures( + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> list[dict[str, Any]]: + """Layout metadata for every tensor reaching a module's forward.""" + return _input_metas(args, kwargs) + + +def output_signatures(output: Any) -> list[dict[str, Any]]: + """Layout metadata for every tensor a module's forward returned.""" + return _output_metas(output) + + +def shape_key_for(metas: list[dict[str, Any]]) -> str: + """The stable key identifying one observed shape variant.""" + return _shape_key(metas) + + +def graph_identity( + module: nn.Module, + args: tuple[Any, ...], + kwargs: dict[str, Any], + output: Any, + *, + scope: str, + tracer: str = "symbolic", +) -> dict[str, Any]: + """Trace one observed invocation and return its captured region record. + + This is the capture pipeline applied to a single call, so the resulting + ``fingerprint`` is produced by exactly the code that produced the + fingerprints in an exported profile. + + Unlike :meth:`FXCaptureSession.finalize`, a failure is raised rather than + recorded: the caller needs a decision, not a report. + """ + session = FXCaptureSession(tracer=tracer, max_scopes=1, max_shape_variants=1) + record = _Scope( + scope=scope, + class_name=type(module).__name__, + module=module, + ) + session._scopes[scope] = record + session._observe(record, args, kwargs, output) + regions = session._regions_for(record) + if not regions: + reasons = [str(item.get("reason", "")) for item in session._graph_breaks] + raise RuntimeError(reasons[0] if reasons else "trace_produced_no_region") + return regions[0] diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 4aa925ac10..78acd18537 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -19,7 +19,7 @@ from fastvideo.hooks.activation_trace import attach_activation_trace, detach_activation_trace from fastvideo.logger import init_logger from fastvideo.profiler import get_or_create_profiler -from fastvideo.optimization import optimization_profile +from fastvideo.optimization import (attach_graph_dispatch, detach_graph_dispatch, optimization_profile) from fastvideo.models.loader.component_loader import PipelineComponentLoader from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages import PipelineStage @@ -65,6 +65,7 @@ def __init__(self, self._stages: list[PipelineStage] = [] self._stage_name_mapping: dict[str, PipelineStage] = {} self._trace_mgr = None + self._dispatch_session = None self._optimization_profile_calls = 0 if required_config_modules is not None: @@ -237,6 +238,9 @@ def post_init(self) -> None: logger.info("Torch Compile enabled for audio VAE") self._trace_mgr = attach_activation_trace(self.modules.get("transformer")) + # Generic graph dispatch. Returns None unless a trusted artifact + # directory is configured, in which case nothing is wrapped at all. + self._dispatch_session = attach_graph_dispatch(self.modules) if not self.fastvideo_args.training_mode: logger.info("Creating pipeline stages...") @@ -526,6 +530,8 @@ def train(self) -> None: def close(self) -> None: detach_activation_trace(getattr(self, "_trace_mgr", None)) self._trace_mgr = None + detach_graph_dispatch(getattr(self, "_dispatch_session", None)) + self._dispatch_session = None def __del__(self): self.close() diff --git a/fastvideo/tests/optimization/test_dispatch.py b/fastvideo/tests/optimization/test_dispatch.py new file mode 100644 index 0000000000..8560ed1816 --- /dev/null +++ b/fastvideo/tests/optimization/test_dispatch.py @@ -0,0 +1,690 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for generic artifact dispatch and native fallback. + +Every kernel and module here is a CPU fake: no CUDA, no Triton, and no model +code. What is exercised is the contract -- identity matching, hash +verification, fallback behavior and diagnostics -- not any particular model. +""" + +from __future__ import annotations + +import copy +import hashlib +import json +from pathlib import Path + +import pytest +import torch +from torch import nn + +from fastvideo import envs +from fastvideo.optimization.artifact import (ANY, MANIFEST_FILENAME, ArtifactRegistry, RuntimeProfile, verify_bundle) +from fastvideo.optimization.dispatch import (REASON_NO_COMPATIBLE_ARTIFACT, REASON_NO_SIGNATURE_MATCH, REASON_SELECTED, + GraphDispatchSession, attach_graph_dispatch, detach_graph_dispatch) +from fastvideo.optimization.identity import (graph_identity, input_signatures, output_signatures) + +WIDTH = 4 +MARKER = 7.0 + +# A CPU fake kernel. It reads what it needs from the module it is handed, which +# is what lets one artifact serve every block in a stack. +FAKE_KERNEL = '''"""CPU fake kernel used by the dispatch tests.""" +import torch + + +def fused_block(module, hidden): + return torch.full_like(hidden, float(module.marker)) +''' + +RAISING_KERNEL = '''"""CPU fake kernel that fails at call time.""" + + +def fused_block(module, hidden): + raise RuntimeError("candidate exploded") +''' + +IMPORT_FAILURE_KERNEL = '''"""CPU fake kernel that fails at import time.""" +raise RuntimeError("import side effect") + + +def fused_block(module, hidden): + return hidden +''' + + +class _Block(nn.Module): + """One traceable block, standing in for a transformer block.""" + + def __init__(self, width: int) -> None: + super().__init__() + self.proj = nn.Linear(width, width) + self.marker = MARKER + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.silu(self.proj(hidden)) + hidden + + +class _UntraceableBlock(nn.Module): + """A block whose data-dependent branch defeats symbolic tracing.""" + + def __init__(self, width: int) -> None: + super().__init__() + self.proj = nn.Linear(width, width) + self.marker = MARKER + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + if bool(hidden.sum() > 0): + return self.proj(hidden) + return self.proj(hidden) * 2 + + +class _Transformer(nn.Module): + """A repeated block stack: the structure dispatch keys on.""" + + def __init__(self, width: int = WIDTH, depth: int = 2, block_cls=_Block) -> None: + super().__init__() + self.blocks = nn.ModuleList([block_cls(width) for _ in range(depth)]) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + for block in self.blocks: + hidden = block(hidden) + return hidden + + +def _pipeline_modules(block_cls=_Block) -> dict[str, nn.Module]: + torch.manual_seed(0) + return {"transformer": _Transformer(block_cls=block_cls).eval()} + + +def _runtime(**overrides) -> RuntimeProfile: + profile = { + "model_id": "fake/model", + "model_revision": "main", + "gpu_architecture": "cpu", + "torch_version": torch.__version__, + "cuda_version": None, + "triton_version": None, + "execution_mode": "inference", + "distributed_mode": "single", + } + profile.update(overrides) + return RuntimeProfile(**profile) + + +def _observed_identity(module: nn.Module, hidden: torch.Tensor, *, scope: str = "transformer.blocks"): + """Trace one call the way dispatch does, to learn the artifact identity.""" + with torch.no_grad(): + output = module(hidden) + region = graph_identity(module, (hidden, ), {}, output, scope=scope, tracer="symbolic") + return ( + region["fingerprint"], + input_signatures((hidden, ), {}), + output_signatures(output), + ) + + +def _sections(fingerprint, inputs, outputs, **overrides) -> dict: + sections = { + "schema_version": 1, + "artifact_id": "fake-fused-block", + "created_at": "2026-07-31T00:00:00+00:00", + "producer": { + "name": "motionkernel", + "version": "1.0.0" + }, + "operation": { + "name": "generated_blocks_fused", + "graph_fingerprint": fingerprint, + "parent_module": "transformer.blocks", + "operations": ["aten::linear", "aten::silu", "aten::add"], + }, + "signature": { + "inputs": copy.deepcopy(inputs), + "outputs": copy.deepcopy(outputs), + }, + "entry_point": { + "file": "kernel.py", + "symbol": "fused_block" + }, + "compatibility": { + "model_id": "fake/model", + "model_revision": ANY, + "gpu_architectures": ["cpu"], + "torch": {}, + "cuda": {}, + "triton": {}, + "execution_modes": ["inference"], + "distributed_modes": ["single"], + }, + "evidence": { + "benchmark": { + "harness": "motionkernel-bench", + "device": "cpu", + "samples": 20, + "baseline_us": 100.0, + "candidate_us": 50.0, + "speedup": 2.0, + "max_abs_error": 0.0, + "max_rel_error": 0.0, + "atol": 1e-5, + "rtol": 1e-5, + "passed": True, + "result_ref": "", + }, + "generation": { + "workload_id": "fake-workload", + "steps": 2, + "metric": "max_abs_latent_diff", + "value": 0.0, + "threshold": 1e-3, + "passed": True, + "baseline_ref": "", + "candidate_ref": "", + }, + }, + "promotion": { + "decision": "promoted", + "reason": "2x with full-generation parity", + "decided_at": "2026-07-31T00:00:00+00:00", + "campaign": { + "campaign_id": "campaign-1", + "source": "cpu-fake", + "target_name": "blocks_fused", + }, + }, + } + sections.update(overrides) + return sections + + +def _write_bundle(root: Path, sections: dict, *, kernel_source: str = FAKE_KERNEL, name: str = "fused") -> Path: + """Write a bundle the way the producer would, hashes included.""" + directory = root / name + directory.mkdir(parents=True, exist_ok=True) + kernel = directory / "kernel.py" + kernel.write_text(kernel_source, encoding="utf-8") + document = dict(sections) + document["files"] = [{ + "path": "kernel.py", + "sha256": hashlib.sha256(kernel.read_bytes()).hexdigest(), + "bytes": kernel.stat().st_size, + }] + (directory / MANIFEST_FILENAME).write_text(json.dumps(document, indent=2, sort_keys=True), encoding="utf-8") + return directory + + +@pytest.fixture() +def hidden() -> torch.Tensor: + torch.manual_seed(1) + return torch.randn(2, WIDTH) + + +@pytest.fixture() +def store(tmp_path: Path) -> Path: + directory = tmp_path / "artifacts" + directory.mkdir() + return directory + + +def _session(store: Path, **overrides) -> GraphDispatchSession: + return GraphDispatchSession(ArtifactRegistry(store), _runtime(**overrides), tracer="symbolic") + + +def _decisions(session: GraphDispatchSession) -> list[dict]: + return session.diagnostics()["decisions"] + + +def _reasons(session: GraphDispatchSession) -> list[str]: + return [item["reason"] for item in _decisions(session)] + + +# -- zero-effect baseline ----------------------------------------------------- + + +def test_no_artifact_directory_behaves_exactly_like_native(monkeypatch, hidden): + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR", "") + modules = _pipeline_modules() + transformer = modules["transformer"] + before = [type(block).forward for block in transformer.blocks] + with torch.no_grad(): + expected = transformer(hidden) + + session = attach_graph_dispatch(modules) + + assert session is None + # No instance attribute shadows the class method: the module graph is + # untouched, so behavior is identical to a build without this feature. + assert all("forward" not in vars(block) for block in transformer.blocks) + assert [type(block).forward for block in transformer.blocks] == before + with torch.no_grad(): + assert torch.equal(transformer(hidden), expected) + + +def test_empty_artifact_directory_stays_native(store, monkeypatch, hidden): + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR", str(store)) + modules = _pipeline_modules() + + session = attach_graph_dispatch(modules) + + assert session is None + + +def test_detach_restores_the_native_forward(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + with torch.no_grad(): + expected = transformer(hidden) + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + session = _session(store) + + assert session.attach_modules(modules) == 2 + assert all("forward" in vars(block) for block in transformer.blocks) + + session.detach() + + assert all("forward" not in vars(block) for block in transformer.blocks) + with torch.no_grad(): + assert torch.equal(transformer(hidden), expected) + + +# -- selection ---------------------------------------------------------------- + + +def test_compatible_artifact_is_selected(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + first = transformer(hidden) + second = transformer(hidden) + + session.detach() + # One decision covers the whole stack, so only the very first block call + # runs natively; every later block -- including the second block of the + # first pass -- goes through the candidate, which fills the tensor with the + # marker it reads off the module it was handed. + assert torch.allclose(first, torch.full_like(first, MARKER)) + assert torch.allclose(second, torch.full_like(second, MARKER)) + assert _reasons(session) == [REASON_SELECTED] + decision = _decisions(session)[0] + assert decision["artifact_id"] == "fake-fused-block" + assert decision["candidate_calls"] == 3 + assert decision["runtime_fallbacks"] == 0 + + +def test_fastest_compatible_artifact_wins(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs), name="slow") + fast = _sections(fingerprint, inputs, outputs) + fast["artifact_id"] = "fake-fused-block-fast" + fast["evidence"]["benchmark"]["speedup"] = 4.0 + _write_bundle(store, fast, name="fast") + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + transformer(hidden) + transformer(hidden) + + session.detach() + decision = _decisions(session)[0] + assert decision["artifact_id"] == "fake-fused-block-fast" + assert decision["rejections"] == ["fake-fused-block:not_selected"] + + +# -- rejections --------------------------------------------------------------- + + +def test_fingerprint_mismatch_falls_back_to_native(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + _, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections("0" * 32, inputs, outputs)) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + expected = _pipeline_modules()["transformer"](hidden) + transformer(hidden) + actual = transformer(hidden) + + session.detach() + assert torch.equal(actual, expected) + assert _reasons(session) == [REASON_NO_COMPATIBLE_ARTIFACT] + assert _decisions(session)[0]["rejections"] == ["fake-fused-block:fingerprint_mismatch"] + + +@pytest.mark.parametrize( + ("mutate", "field"), + [ + (lambda item: item.update({"dtype": "bfloat16"}), "dtype"), + (lambda item: item.update({"shape": [8, WIDTH]}), "shape"), + (lambda item: item.update({"device_type": "cuda"}), "device_type"), + ], +) +def test_signature_mismatches_never_reach_a_trace(store, hidden, mutate, field): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + mutated = copy.deepcopy(inputs) + mutate(mutated[0]) + _write_bundle(store, _sections(fingerprint, mutated, outputs)) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + expected = _pipeline_modules()["transformer"](hidden) + transformer(hidden) + actual = transformer(hidden) + + session.detach() + assert torch.equal(actual, expected) + # The input pre-filter rejects these before any graph is traced. + assert _reasons(session) == [REASON_NO_SIGNATURE_MATCH] + + +def test_output_signature_mismatch_falls_back_to_native(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + mutated = copy.deepcopy(outputs) + mutated[0]["shape"] = [2, WIDTH * 2] + mutated[0]["stride"] = [WIDTH * 2, 1] + _write_bundle(store, _sections(fingerprint, inputs, mutated)) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + transformer(hidden) + transformer(hidden) + + session.detach() + assert _decisions(session)[0]["rejections"] == ["fake-fused-block:output_signature_mismatch"] + + +@pytest.mark.parametrize( + ("override", "expected"), + [ + ({ + "gpu_architectures": ["sm90"] + }, "gpu_architecture_mismatch"), + ({ + "torch": { + "min": "99.0.0" + } + }, "torch_version_unsupported"), + ({ + "cuda": { + "min": "12.0" + } + }, "cuda_version_unsupported"), + ({ + "triton": { + "min": "3.0.0" + } + }, "triton_version_unsupported"), + ({ + "model_id": "other/model" + }, "model_mismatch"), + ({ + "execution_modes": ["training"] + }, "execution_mode_unsupported"), + ({ + "distributed_modes": ["tensor_parallel"] + }, "distributed_mode_unsupported"), + ], +) +def test_environment_mismatches_fall_back_to_native(store, hidden, override, expected): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + sections = _sections(fingerprint, inputs, outputs) + sections["compatibility"].update(override) + _write_bundle(store, sections) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + native = _pipeline_modules()["transformer"](hidden) + transformer(hidden) + actual = transformer(hidden) + + session.detach() + assert torch.equal(actual, native) + assert _decisions(session)[0]["rejections"] == [f"fake-fused-block:{expected}"] + + +def test_unpromoted_artifact_is_never_dispatched(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + sections = _sections(fingerprint, inputs, outputs) + sections["promotion"]["decision"] = "quarantined" + _write_bundle(store, sections) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + transformer(hidden) + transformer(hidden) + + session.detach() + assert _decisions(session)[0]["rejections"] == ["fake-fused-block:not_promoted"] + + +# -- tampering and load failures ---------------------------------------------- + + +def test_tampered_kernel_is_rejected_before_import(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + directory = _write_bundle(store, _sections(fingerprint, inputs, outputs)) + # Swap the kernel for one that would be obvious if it ever executed. + (directory / "kernel.py").write_text('raise SystemExit("payload executed")\n', encoding="utf-8") + + registry = ArtifactRegistry(store) + + assert registry.manifests == [] + assert len(registry.errors) == 1 + assert "does not match manifest" in registry.errors[0] or "bytes" in registry.errors[0] + + session = GraphDispatchSession(registry, _runtime()) + assert session.attach_modules(modules) == 0 + with torch.no_grad(): + assert torch.equal(transformer(hidden), _pipeline_modules()["transformer"](hidden)) + + +def test_undeclared_file_is_rejected(store, hidden): + modules = _pipeline_modules() + fingerprint, inputs, outputs = _observed_identity(modules["transformer"].blocks[0], hidden) + directory = _write_bundle(store, _sections(fingerprint, inputs, outputs)) + (directory / "extra.py").write_text("SECRET = 1\n", encoding="utf-8") + + registry = ArtifactRegistry(store) + + assert registry.manifests == [] + assert "undeclared file" in registry.errors[0] + + +def test_unknown_schema_version_is_rejected(store, hidden): + modules = _pipeline_modules() + fingerprint, inputs, outputs = _observed_identity(modules["transformer"].blocks[0], hidden) + sections = _sections(fingerprint, inputs, outputs) + sections["schema_version"] = 99 + _write_bundle(store, sections) + + registry = ArtifactRegistry(store) + + assert registry.manifests == [] + assert "unsupported version" in registry.errors[0] + + +def test_failing_import_falls_back_to_native(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs), kernel_source=IMPORT_FAILURE_KERNEL) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + native = _pipeline_modules()["transformer"](hidden) + transformer(hidden) + actual = transformer(hidden) + + session.detach() + assert torch.equal(actual, native) + assert _reasons(session) == ["artifact_load_failed:ArtifactError"] + + +def test_bundle_outside_the_trusted_root_is_never_loaded(store, tmp_path, hidden): + modules = _pipeline_modules() + fingerprint, inputs, outputs = _observed_identity(modules["transformer"].blocks[0], hidden) + outside = _write_bundle(tmp_path / "elsewhere", _sections(fingerprint, inputs, outputs)) + + # The registry only ever looks inside its own root. + assert ArtifactRegistry(store).manifests == [] + assert verify_bundle(outside).artifact_id == "fake-fused-block" + + +# -- runtime failure ---------------------------------------------------------- + + +def test_candidate_exception_falls_back_to_native_output(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs), kernel_source=RAISING_KERNEL) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + native = _pipeline_modules()["transformer"](hidden) + transformer(hidden) + actual = transformer(hidden) + again = transformer(hidden) + + session.detach() + assert torch.equal(actual, native) + assert torch.equal(again, native) + decision = _decisions(session)[0] + assert decision["reason"] == "candidate_runtime_error:RuntimeError" + assert decision["active"] is False + # Demoted after the first failure, not retried on every later call. + assert decision["runtime_fallbacks"] == 1 + + +# -- tracing failures --------------------------------------------------------- + + +def test_untraceable_module_stays_native(store, hidden): + modules = _pipeline_modules(block_cls=_UntraceableBlock) + transformer = modules["transformer"] + with torch.no_grad(): + reference = transformer(hidden) + # Build a bundle whose input signature matches, so the pre-filter passes + # and the tracer is genuinely attempted. + inputs = input_signatures((hidden, ), {}) + outputs = output_signatures(reference) + _write_bundle(store, _sections("1" * 32, inputs, outputs)) + session = _session(store) + session.attach_modules(modules) + + with torch.no_grad(): + transformer(hidden) + actual = transformer(hidden) + + session.detach() + assert torch.equal(actual, reference) + assert _reasons(session)[0].startswith("graph_identity_unavailable") + + +# -- diagnostics -------------------------------------------------------------- + + +def test_diagnostics_are_metadata_only(store, tmp_path, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + session = _session(store) + session.attach_modules(modules) + with torch.no_grad(): + transformer(hidden) + transformer(hidden) + session.detach() + + output = session.write_diagnostics(tmp_path / "dispatch.json") + + assert output is not None + report = json.loads(output.read_text(encoding="utf-8")) + assert report["dispatch"]["dispatch_schema_version"] == 1 + assert report["dispatch"]["registry"]["artifact_ids"] == ["fake-fused-block"] + assert report["dispatch"]["reason_counts"] == {REASON_SELECTED: 1} + # Nothing tensor-shaped or prompt-shaped may appear anywhere in the report. + serialized = json.dumps(report).lower() + for forbidden in ("prompt", "tensor_values", "weights", "password", "secret"): + assert forbidden not in serialized + + +def test_attach_graph_dispatch_uses_the_configured_directory(store, monkeypatch, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR", str(store)) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID", "fake/model") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION", ANY) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER", "symbolic") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES", 64) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES", 8) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE", "") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS", "") + + session = attach_graph_dispatch(modules) + + assert session is not None + try: + with torch.no_grad(): + transformer(hidden) + actual = transformer(hidden) + assert torch.allclose(actual, torch.full_like(actual, MARKER)) + finally: + detach_graph_dispatch(session) + assert all("forward" not in vars(block) for block in transformer.blocks) + + +def test_training_mode_is_detected_and_rejects_inference_only_artifacts(store, monkeypatch, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + transformer.train() + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR", str(store)) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID", "fake/model") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION", ANY) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER", "symbolic") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES", 64) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES", 8) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE", "") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS", "") + + session = attach_graph_dispatch(modules) + + assert session is not None + try: + with torch.no_grad(): + transformer(hidden) + actual = transformer(hidden) + assert not torch.allclose(actual, torch.full_like(actual, MARKER)) + assert _decisions(session)[0]["rejections"] == ["fake-fused-block:execution_mode_unsupported"] + finally: + detach_graph_dispatch(session) From df8cfee245227969f9ef9508cc62822318b18785 Mon Sep 17 00:00:00 2001 From: aryan5v Date: Fri, 31 Jul 2026 11:39:26 -0700 Subject: [PATCH 2/2] [bugfix]: harden generic artifact dispatch --- fastvideo/optimization/artifact.py | 91 ++++-- fastvideo/optimization/dispatch.py | 41 ++- fastvideo/optimization/fx_capture.py | 267 ++++++++---------- fastvideo/optimization/identity.py | 69 +---- fastvideo/tests/optimization/test_dispatch.py | 133 ++++++++- 5 files changed, 362 insertions(+), 239 deletions(-) diff --git a/fastvideo/optimization/artifact.py b/fastvideo/optimization/artifact.py index 8476434011..5236dbfe0f 100644 --- a/fastvideo/optimization/artifact.py +++ b/fastvideo/optimization/artifact.py @@ -23,6 +23,7 @@ import hashlib import importlib.util import json +import math import re import sys from dataclasses import dataclass @@ -174,6 +175,18 @@ def from_dict(cls, raw: Any, *, source: object, location: str) -> VersionRange: minimum = _text(minimum, source, f"{location}.min") if maximum is not None: maximum = _text(maximum, source, f"{location}.max_exclusive") + low = parse_version(minimum) + high = parse_version(maximum) + if minimum is not None and low is None: + raise _fail(source, f"{location}.min", "must be a dotted version") + if maximum is not None and high is None: + raise _fail( + source, + f"{location}.max_exclusive", + "must be a dotted version", + ) + if low is not None and high is not None and _compare_versions(low, high) >= 0: + raise _fail(source, location, "min must be lower than max_exclusive") return cls(minimum=minimum, maximum_exclusive=maximum) @property @@ -232,6 +245,13 @@ def _signature_keys(raw: Any, source: object, location: str) -> tuple[tuple[Any, keys = [] for index, item in enumerate(items): entry = _mapping(item, source, f"{location}[{index}]") + unknown = sorted(set(entry) - {"name", "shape", "stride", "dtype", "device_type", "requires_grad"}) + if unknown: + raise _fail( + source, + f"{location}[{index}]", + f"unknown field(s) {unknown}", + ) shape = _sequence(entry.get("shape"), source, f"{location}[{index}].shape") stride = _sequence(entry.get("stride"), source, f"{location}[{index}].stride") for values, name in ((shape, "shape"), (stride, "stride")): @@ -242,6 +262,12 @@ def _signature_keys(raw: Any, source: object, location: str) -> tuple[tuple[Any, f"{location}[{index}].{name}[{position}]", "must be an integer", ) + if name == "shape" and value < 0: + raise _fail( + source, + f"{location}[{index}].shape[{position}]", + "must be non-negative", + ) if len(shape) != len(stride): raise _fail( source, @@ -250,6 +276,11 @@ def _signature_keys(raw: Any, source: object, location: str) -> tuple[tuple[Any, ) _text(entry.get("dtype"), source, f"{location}[{index}].dtype") _text(entry.get("device_type"), source, f"{location}[{index}].device_type") + _bool( + entry.get("requires_grad", False), + source, + f"{location}[{index}].requires_grad", + ) keys.append(signature_key(entry)) return tuple(keys) @@ -344,8 +375,36 @@ def from_dict(cls, raw_value: Any, *, directory: Path) -> ArtifactManifest: benchmark = _mapping(evidence.get("benchmark"), source, "evidence.benchmark") generation = _mapping(evidence.get("generation"), source, "evidence.generation") speedup = benchmark.get("speedup") - if isinstance(speedup, bool) or not isinstance(speedup, int | float): - raise _fail(source, "evidence.benchmark.speedup", "must be a number") + if (isinstance(speedup, bool) or not isinstance(speedup, int | float) or not math.isfinite(float(speedup)) + or float(speedup) < 0): + raise _fail( + source, + "evidence.benchmark.speedup", + "must be a finite non-negative number", + ) + + architectures = _sequence( + compatibility.get("gpu_architectures"), + source, + "compatibility.gpu_architectures", + ) + execution_modes = _sequence( + compatibility.get("execution_modes"), + source, + "compatibility.execution_modes", + ) + distributed_modes = _sequence( + compatibility.get("distributed_modes"), + source, + "compatibility.distributed_modes", + ) + for items, location in ( + (architectures, "compatibility.gpu_architectures"), + (execution_modes, "compatibility.execution_modes"), + (distributed_modes, "compatibility.distributed_modes"), + ): + if not items: + raise _fail(source, location, "must be a non-empty list") return cls( artifact_id=_text(raw.get("artifact_id"), source, "artifact_id"), @@ -376,12 +435,8 @@ def from_dict(cls, raw_value: Any, *, directory: Path) -> ArtifactManifest: "compatibility.model_revision", ), gpu_architectures=tuple( - _text(item, source, f"compatibility.gpu_architectures[{index}]") for index, item in enumerate( - _sequence( - compatibility.get("gpu_architectures"), - source, - "compatibility.gpu_architectures", - ))), + _text(item, source, f"compatibility.gpu_architectures[{index}]") + for index, item in enumerate(architectures)), torch_range=VersionRange.from_dict( compatibility.get("torch"), source=source, @@ -398,19 +453,11 @@ def from_dict(cls, raw_value: Any, *, directory: Path) -> ArtifactManifest: location="compatibility.triton", ), execution_modes=tuple( - _text(item, source, f"compatibility.execution_modes[{index}]") for index, item in enumerate( - _sequence( - compatibility.get("execution_modes"), - source, - "compatibility.execution_modes", - ))), + _text(item, source, f"compatibility.execution_modes[{index}]") + for index, item in enumerate(execution_modes)), distributed_modes=tuple( - _text(item, source, f"compatibility.distributed_modes[{index}]") for index, item in enumerate( - _sequence( - compatibility.get("distributed_modes"), - source, - "compatibility.distributed_modes", - ))), + _text(item, source, f"compatibility.distributed_modes[{index}]") + for index, item in enumerate(distributed_modes)), promotion_decision=_text(promotion.get("decision"), source, "promotion.decision"), evidence_passed=(_bool(benchmark.get("passed"), source, "evidence.benchmark.passed") and _bool(generation.get("passed"), source, "evidence.generation.passed")), @@ -608,7 +655,9 @@ def load_entry_point(manifest: ArtifactManifest, *, trusted_root: Path) -> Calla ) entry_file = _resolve_inside(directory, directory / verified.entry_file) - module_name = f"{_MODULE_NAMESPACE}.{re.sub(r'[^A-Za-z0-9_]', '_', verified.artifact_id)}" + readable_id = re.sub(r"[^A-Za-z0-9_]", "_", verified.artifact_id) + unique_id = hashlib.sha256(verified.artifact_id.encode("utf-8")).hexdigest()[:16] + module_name = f"{_MODULE_NAMESPACE}.{readable_id}_{unique_id}" spec = importlib.util.spec_from_file_location(module_name, entry_file) if spec is None or spec.loader is None: raise _fail(str(directory), "entry_point", f"cannot load {verified.entry_file!r}") diff --git a/fastvideo/optimization/dispatch.py b/fastvideo/optimization/dispatch.py index 6f7816fa25..dd38eb84f8 100644 --- a/fastvideo/optimization/dispatch.py +++ b/fastvideo/optimization/dispatch.py @@ -37,6 +37,7 @@ import json from collections import defaultdict +from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from typing import Any @@ -282,14 +283,15 @@ def _resolve( return _Decision(None, REASON_NO_SIGNATURE_MATCH) try: - region = graph_identity( - wrapper.module, - args, - kwargs, - output, - scope=scope, - tracer=self.tracer, - ) + with self._native_forward_for_identity(wrapper): + region = graph_identity( + wrapper.module, + args, + kwargs, + output, + scope=scope, + tracer=self.tracer, + ) except Exception as exc: # noqa: BLE001 - untraceable modules stay native reason = f"{_REASON_IDENTITY_PREFIX}:{type(exc).__name__}" logger.info("Graph identity unavailable for %s: %s", scope, reason) @@ -300,17 +302,17 @@ def _resolve( matched: list[ArtifactManifest] = [] rejections: list[str] = [] for manifest in candidates: - reason = check_compatibility( + compatibility_reason = check_compatibility( manifest, graph_fingerprint=fingerprint, input_keys=input_keys, output_keys=output_keys, runtime=self.runtime, ) - if reason is None: + if compatibility_reason is None: matched.append(manifest) else: - rejections.append(f"{manifest.artifact_id}:{reason}") + rejections.append(f"{manifest.artifact_id}:{compatibility_reason}") if not matched: return _Decision( None, @@ -346,6 +348,21 @@ def _resolve( rejections=tuple(sorted(rejections)), ) + @contextmanager + def _native_forward_for_identity(self, wrapper: _Wrapper): + """Expose the real forward while export/Dynamo recompute identity. + + Symbolic FX traces the class method, but export and Dynamo call the + instance forward. Leaving the dispatch wrapper installed would recurse + back into this session during identity capture. + """ + installed_forward = wrapper.module.forward + wrapper.module.forward = wrapper.native_forward # type: ignore[method-assign] + try: + yield + finally: + wrapper.module.forward = installed_forward # type: ignore[method-assign] + def _variant_count(self, scope: str) -> int: return sum(1 for existing_scope, _ in self._decisions if existing_scope == scope) @@ -412,7 +429,7 @@ def write_diagnostics(self, path: str | Path) -> Path | None: def _execution_mode(modules: dict[str, Any] | None) -> str: """Report ``training`` when any hooked module is in training mode.""" for module in (modules or {}).values(): - if isinstance(module, nn.Module) and module.training: + if isinstance(module, nn.Module) and any(child.training for child in module.modules()): return "training" return "inference" diff --git a/fastvideo/optimization/fx_capture.py b/fastvideo/optimization/fx_capture.py index 3c619e92e0..8b09524783 100644 --- a/fastvideo/optimization/fx_capture.py +++ b/fastvideo/optimization/fx_capture.py @@ -159,22 +159,18 @@ def _tensor_leaves( leaves: list[tuple[str, torch.Tensor]] = [] if isinstance(value, tuple | list): for index, item in enumerate(value): - leaves.extend( - _tensor_leaves( - item, - prefix=f"{prefix}_{index}", - seen=seen, - ) - ) + leaves.extend(_tensor_leaves( + item, + prefix=f"{prefix}_{index}", + seen=seen, + )) elif isinstance(value, Mapping): for key, item in value.items(): - leaves.extend( - _tensor_leaves( - item, - prefix=f"{prefix}_{_safe_name(str(key))}", - seen=seen, - ) - ) + leaves.extend(_tensor_leaves( + item, + prefix=f"{prefix}_{_safe_name(str(key))}", + seen=seen, + )) elif is_dataclass(value) and not isinstance(value, type): for item in fields(value): leaves.extend( @@ -182,8 +178,7 @@ def _tensor_leaves( getattr(value, item.name), prefix=f"{prefix}_{_safe_name(item.name)}", seen=seen, - ) - ) + )) return leaves @@ -194,25 +189,18 @@ def _input_metas( leaves: list[tuple[str, torch.Tensor]] = [] seen: set[int] = set() for index, value in enumerate(args): - leaves.extend( - _tensor_leaves(value, prefix=f"input_{index}", seen=seen) - ) + leaves.extend(_tensor_leaves(value, prefix=f"input_{index}", seen=seen)) for key, value in kwargs.items(): - leaves.extend( - _tensor_leaves( - value, - prefix=f"kwarg_{_safe_name(str(key))}", - seen=seen, - ) - ) + leaves.extend(_tensor_leaves( + value, + prefix=f"kwarg_{_safe_name(str(key))}", + seen=seen, + )) return [_tensor_meta(name, tensor) for name, tensor in leaves] def _output_metas(output: Any) -> list[dict[str, Any]]: - return [ - _tensor_meta(name, tensor) - for name, tensor in _tensor_leaves(output, prefix="output") - ] + return [_tensor_meta(name, tensor) for name, tensor in _tensor_leaves(output, prefix="output")] def _input_shape_key(args: tuple[Any, ...], kwargs: dict[str, Any]) -> str: @@ -255,25 +243,19 @@ def _sanitize_constants(raw: dict[str, Any]) -> dict[str, Any]: def _capture_failure_reason(mode: str, exc: Exception) -> str: """Classify failures without exporting exception text or source snippets.""" text = str(exc).lower() - if any( - marker in text - for marker in ( + if any(marker in text for marker in ( "data-dependent", "data dependent", "guardondatadependentsymnode", ".item()", - ) - ): + )): code = "data_dependent_control_flow" - elif any( - marker in text - for marker in ( + elif any(marker in text for marker in ( "control flow", "proxy object", "symbolically traced variables", "cannot be iterated", - ) - ): + )): code = "dynamic_python_control_flow" elif "alias" in text: code = "unknown_aliasing" @@ -324,11 +306,7 @@ def _capture_ready_module( return forward_hooks = getattr(manager, "forward_hooks", {}) - unsupported_hooks = [ - str(name) - for name in forward_hooks - if str(name) != "LayerwiseOffloadHook" - ] + unsupported_hooks = [str(name) for name in forward_hooks if str(name) != "LayerwiseOffloadHook"] if unsupported_hooks: raise RuntimeError("unsupported_module_hooks") @@ -342,15 +320,8 @@ def _capture_ready_module( if offload_hook is not None and state is None: raise RuntimeError("missing_offload_state") - parameter_data = { - name: parameter.data - for name, parameter in module.named_parameters() - } - gpu_parameters = ( - dict(getattr(state, "gpu_named_parameters", {})) - if state is not None - else {} - ) + parameter_data = {name: parameter.data for name, parameter in module.named_parameters()} + gpu_parameters = (dict(getattr(state, "gpu_named_parameters", {})) if state is not None else {}) try: if state is not None: state.wait_and_replace_params() @@ -369,9 +340,7 @@ def _capture_ready_module( @contextmanager -def _capture_forward_context( - observed_context: tuple[Any, Any] | None, -) -> Iterator[None]: +def _capture_forward_context(observed_context: tuple[Any, Any] | None, ) -> Iterator[None]: """Re-establish the bounded runtime context needed by attention layers.""" if observed_context is None: yield @@ -381,9 +350,9 @@ def _capture_forward_context( current_timestep, attention_metadata = observed_context with set_forward_context( - current_timestep=current_timestep, - attn_metadata=attention_metadata, - forward_batch=None, + current_timestep=current_timestep, + attn_metadata=attention_metadata, + forward_batch=None, ): yield @@ -487,27 +456,21 @@ def record_constants(node: Any, key: str) -> None: if _is_safe_constant(value): constants[f"{node.name}.arg{arg_index}"] = value elif value is not None and not tuple(node_names(value)): - notes.append( - f"{key}: unsafe positional constant arg{arg_index}" - ) + notes.append(f"{key}: unsafe positional constant arg{arg_index}") for kwarg, value in (node.kwargs or {}).items(): if hasattr(value, "op"): continue if _is_safe_constant(value): constants[f"{node.name}.{kwarg}"] = value elif value is not None and not tuple(node_names(value)): - notes.append( - f"{key}: non-scalar constant {kwarg!r} dropped" - ) + notes.append(f"{key}: non-scalar constant {kwarg!r} dropped") for node in graph.nodes: if node.op in {"placeholder", "output"}: continue if node.op == "get_attr": target = str(node.target) - notes.append( - f"get_attr:{target}: lifted attribute value not exported" - ) + notes.append(f"get_attr:{target}: lifted attribute value not exported") continue if node.op == "call_function": key = _op_key(node.target) @@ -589,14 +552,12 @@ def _ir_argument( return {"device": "runtime"} if isinstance(value, tuple | list): kind = "tuple" if isinstance(value, tuple) else "list" - return { - kind: [_ir_argument(item, node_expressions) for item in value] - } + return {kind: [_ir_argument(item, node_expressions) for item in value]} # Static SymInts are shape metadata and safe to materialize. if type(value).__module__.startswith("torch") and type(value).__name__ in { - "SymInt", - "SymBool", - "SymFloat", + "SymInt", + "SymBool", + "SymFloat", }: try: scalar = int(value) if type(value).__name__ != "SymFloat" else float(value) @@ -625,9 +586,7 @@ def _extract_executable_ir(exported: Any, graph: Any) -> dict[str, Any]: inputs: list[dict[str, Any]] = [] runtime_index = 0 lifted_index = 0 - for index, (node, input_spec) in enumerate( - zip(placeholders, input_specs, strict=True) - ): + for index, (node, input_spec) in enumerate(zip(placeholders, input_specs, strict=True)): meta = _ir_tensor_meta((node.meta or {}).get("val")) if meta is None: constant = _ir_argument((node.meta or {}).get("val"), {}) @@ -661,9 +620,7 @@ def _extract_executable_ir(exported: Any, graph: Any) -> dict[str, Any]: item: dict[str, Any] = { "id": node_id, "target": _ir_target(node), - "args": [ - _ir_argument(value, node_expressions) for value in node.args - ], + "args": [_ir_argument(value, node_expressions) for value in node.args], "kwargs": { _safe_name(str(key)): _ir_argument(value, node_expressions) for key, value in (node.kwargs or {}).items() @@ -678,11 +635,7 @@ def _extract_executable_ir(exported: Any, graph: Any) -> dict[str, Any]: if len(output_nodes) != 1 or not output_nodes[0].args: raise RuntimeError("invalid_graph_output") encoded_output = _ir_argument(output_nodes[0].args[0], node_expressions) - outputs = ( - encoded_output["tuple"] - if set(encoded_output) == {"tuple"} - else [encoded_output] - ) + outputs = (encoded_output["tuple"] if set(encoded_output) == {"tuple"} else [encoded_output]) return { "schema_version": 1, "inputs": inputs, @@ -898,11 +851,9 @@ def _mode_order(self) -> tuple[str, ...]: if self.tracer in {"auto", "fallback"}: return _CAPTURE_MODES if self.tracer in _CAPTURE_MODES: - return (self.tracer,) - raise ValueError( - f"unsupported tracer {self.tracer!r}; " - "use auto, symbolic, export, or dynamo" - ) + return (self.tracer, ) + raise ValueError(f"unsupported tracer {self.tracer!r}; " + "use auto, symbolic, export, or dynamo") def _trace( self, @@ -915,11 +866,9 @@ def _trace( if args is None or kwargs is None: raise RuntimeError("example arguments unavailable") with _capture_ready_module(module, args, kwargs) as ( - capture_args, - capture_kwargs, - ), _capture_forward_context( - variant.observed_context - ), _traceable_module_forwards(module): + capture_args, + capture_kwargs, + ), _capture_forward_context(variant.observed_context), _traceable_module_forwards(module): if mode == "symbolic": return torch.fx.symbolic_trace(module) if mode == "export": @@ -937,11 +886,7 @@ def _trace( graph_module = getattr(exported, "graph_module", None) if graph_module is None and isinstance(exported, tuple): graph_module = exported[0] - return ( - graph_module - if graph_module is not None - else exported - ) + return (graph_module if graph_module is not None else exported) raise ValueError(f"unsupported capture mode {mode!r}") def _regions_for(self, record: _Scope) -> list[dict[str, Any]]: @@ -962,13 +907,9 @@ def _regions_for(self, record: _Scope) -> list[dict[str, Any]]: graph = getattr(traced, "graph", None) if graph is None: raise RuntimeError("trace result has no FX graph") - operations, dependencies, constants, notes = _extract_graph( - graph - ) + operations, dependencies, constants, notes = _extract_graph(graph) if not operations: - raise RuntimeError( - "trace result has no tensor operations" - ) + raise RuntimeError("trace result has no tensor operations") capture_mode = mode if mode == "export": try: @@ -979,20 +920,18 @@ def _regions_for(self, record: _Scope) -> list[dict[str, Any]]: except Exception as exc: # noqa: BLE001 known_reason = str(exc) if known_reason not in { - "graph_signature_input_mismatch", - "invalid_graph_output", - "missing_graph_signature", - "unlifted_graph_attribute", - "unsupported_graph_node_kind", - "unsupported_non_tensor_graph_input", + "graph_signature_input_mismatch", + "invalid_graph_output", + "missing_graph_signature", + "unlifted_graph_attribute", + "unsupported_graph_node_kind", + "unsupported_non_tensor_graph_input", }: known_reason = type(exc).__name__ self._graph_breaks.append({ "scope": record.scope, - "reason": ( - "executable_ir_unavailable:" - f"{known_reason}" - ), + "reason": ("executable_ir_unavailable:" + f"{known_reason}"), "count": max(variant.calls, 1), }) break @@ -1054,9 +993,7 @@ def _regions_for(self, record: _Scope) -> list[dict[str, Any]]: "module_class": record.class_name, "tracer": self.tracer, "capture_mode": capture_mode, - "capture_attempts": list(attempts)[ - : list(attempts).index(capture_mode) + 1 - ], + "capture_attempts": list(attempts)[:list(attempts).index(capture_mode) + 1], "capture_failures": failures, }, } @@ -1087,32 +1024,28 @@ def finalize(self) -> dict[str, Any]: unsupported = _coalesce(self._unsupported, ("op_name", "reason", "scope")) payload = { "capture": { - "capture_schema_version": CAPTURE_SCHEMA_VERSION, - "tracer": self.tracer, - "scopes": sorted(self._scopes), + "capture_schema_version": + CAPTURE_SCHEMA_VERSION, + "tracer": + self.tracer, + "scopes": + sorted(self._scopes), "scope_calls": { record.scope: record.calls for record in self._scopes.values() }, - "dropped_scopes": self._dropped_scopes, - "dropped_shape_variants": sum(record.dropped_variants for record in self._scopes.values()), - "errors": list(self._errors), - "capture_mode_breakdown": dict( - sorted( - ( - mode, - sum( - 1 - for region in regions - if region.get("attributes", {}).get( - "capture_mode" - ) - == mode - ), - ) - for mode in _CAPTURE_MODES - ) - ), + "dropped_scopes": + self._dropped_scopes, + "dropped_shape_variants": + sum(record.dropped_variants for record in self._scopes.values()), + "errors": + list(self._errors), + "capture_mode_breakdown": + dict( + sorted(( + mode, + sum(1 for region in regions if region.get("attributes", {}).get("capture_mode") == mode), + ) for mode in _CAPTURE_MODES)), }, "regions": regions, "graph_breaks": breaks, @@ -1122,6 +1055,56 @@ def finalize(self) -> dict[str, Any]: return payload +# Public dispatch-facing capture API. Keeping these operations here makes +# exported/runtime fingerprints share one implementation without forcing the +# dispatcher to import underscore-private capture internals. +def invocation_input_signatures( + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> list[dict[str, Any]]: + """Return metadata-only tensor signatures for a module invocation.""" + return _input_metas(args, kwargs) + + +def invocation_output_signatures(output: Any) -> list[dict[str, Any]]: + """Return metadata-only tensor signatures for a module output.""" + return _output_metas(output) + + +def invocation_shape_key(metas: list[dict[str, Any]]) -> str: + """Return the canonical capture shape key for an invocation.""" + return _shape_key(metas) + + +def capture_invocation_identity( + module: nn.Module, + args: tuple[Any, ...], + kwargs: dict[str, Any], + output: Any, + *, + scope: str, + tracer: str = "symbolic", +) -> dict[str, Any]: + """Capture one observed invocation and return its region identity. + + Failures are raised because dispatch needs one definitive decision rather + than a full discovery report containing graph-break records. + """ + session = FXCaptureSession(tracer=tracer, max_scopes=1, max_shape_variants=1) + record = _Scope( + scope=scope, + class_name=type(module).__name__, + module=module, + ) + session._scopes[scope] = record + session._observe(record, args, kwargs, output) + regions = session._regions_for(record) + if not regions: + reasons = [str(item.get("reason", "")) for item in session._graph_breaks] + raise RuntimeError(reasons[0] if reasons else "trace_produced_no_region") + return regions[0] + + def _coalesce(records: list[dict[str, Any]], keys: tuple[str, ...]) -> list[dict[str, Any]]: """Merge duplicate records, summing their ``count`` fields.""" merged: dict[tuple[Any, ...], dict[str, Any]] = {} diff --git a/fastvideo/optimization/identity.py b/fastvideo/optimization/identity.py index 882a1d3709..85c5545771 100644 --- a/fastvideo/optimization/identity.py +++ b/fastvideo/optimization/identity.py @@ -15,65 +15,16 @@ from __future__ import annotations -from typing import Any - -from torch import nn - from fastvideo.optimization.fx_capture import ( - FXCaptureSession, - _input_metas, - _output_metas, - _shape_key, - _Scope, + capture_invocation_identity as graph_identity, + invocation_input_signatures as input_signatures, + invocation_output_signatures as output_signatures, + invocation_shape_key as shape_key_for, ) - -def input_signatures( - args: tuple[Any, ...], - kwargs: dict[str, Any], -) -> list[dict[str, Any]]: - """Layout metadata for every tensor reaching a module's forward.""" - return _input_metas(args, kwargs) - - -def output_signatures(output: Any) -> list[dict[str, Any]]: - """Layout metadata for every tensor a module's forward returned.""" - return _output_metas(output) - - -def shape_key_for(metas: list[dict[str, Any]]) -> str: - """The stable key identifying one observed shape variant.""" - return _shape_key(metas) - - -def graph_identity( - module: nn.Module, - args: tuple[Any, ...], - kwargs: dict[str, Any], - output: Any, - *, - scope: str, - tracer: str = "symbolic", -) -> dict[str, Any]: - """Trace one observed invocation and return its captured region record. - - This is the capture pipeline applied to a single call, so the resulting - ``fingerprint`` is produced by exactly the code that produced the - fingerprints in an exported profile. - - Unlike :meth:`FXCaptureSession.finalize`, a failure is raised rather than - recorded: the caller needs a decision, not a report. - """ - session = FXCaptureSession(tracer=tracer, max_scopes=1, max_shape_variants=1) - record = _Scope( - scope=scope, - class_name=type(module).__name__, - module=module, - ) - session._scopes[scope] = record - session._observe(record, args, kwargs, output) - regions = session._regions_for(record) - if not regions: - reasons = [str(item.get("reason", "")) for item in session._graph_breaks] - raise RuntimeError(reasons[0] if reasons else "trace_produced_no_region") - return regions[0] +__all__ = [ + "graph_identity", + "input_signatures", + "output_signatures", + "shape_key_for", +] diff --git a/fastvideo/tests/optimization/test_dispatch.py b/fastvideo/tests/optimization/test_dispatch.py index 8560ed1816..3f14a28101 100644 --- a/fastvideo/tests/optimization/test_dispatch.py +++ b/fastvideo/tests/optimization/test_dispatch.py @@ -18,7 +18,8 @@ from torch import nn from fastvideo import envs -from fastvideo.optimization.artifact import (ANY, MANIFEST_FILENAME, ArtifactRegistry, RuntimeProfile, verify_bundle) +from fastvideo.optimization.artifact import (ANY, MANIFEST_FILENAME, ArtifactRegistry, RuntimeProfile, load_entry_point, + verify_bundle) from fastvideo.optimization.dispatch import (REASON_NO_COMPATIBLE_ARTIFACT, REASON_NO_SIGNATURE_MATCH, REASON_SELECTED, GraphDispatchSession, attach_graph_dispatch, detach_graph_dispatch) from fastvideo.optimization.identity import (graph_identity, input_signatures, output_signatures) @@ -111,11 +112,17 @@ def _runtime(**overrides) -> RuntimeProfile: return RuntimeProfile(**profile) -def _observed_identity(module: nn.Module, hidden: torch.Tensor, *, scope: str = "transformer.blocks"): +def _observed_identity( + module: nn.Module, + hidden: torch.Tensor, + *, + scope: str = "transformer.blocks", + tracer: str = "symbolic", +): """Trace one call the way dispatch does, to learn the artifact identity.""" with torch.no_grad(): output = module(hidden) - region = graph_identity(module, (hidden, ), {}, output, scope=scope, tracer="symbolic") + region = graph_identity(module, (hidden, ), {}, output, scope=scope, tracer=tracer) return ( region["fingerprint"], input_signatures((hidden, ), {}), @@ -226,8 +233,12 @@ def store(tmp_path: Path) -> Path: return directory -def _session(store: Path, **overrides) -> GraphDispatchSession: - return GraphDispatchSession(ArtifactRegistry(store), _runtime(**overrides), tracer="symbolic") +def _session(store: Path, *, tracer: str = "symbolic", **overrides) -> GraphDispatchSession: + return GraphDispatchSession( + ArtifactRegistry(store), + _runtime(**overrides), + tracer=tracer, + ) def _decisions(session: GraphDispatchSession) -> list[dict]: @@ -317,6 +328,26 @@ def test_compatible_artifact_is_selected(store, hidden): assert decision["runtime_fallbacks"] == 0 +def test_export_identity_traces_native_forward_without_dispatch_recursion(store, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity( + transformer.blocks[0], hidden, tracer="export" + ) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + session = _session(store, tracer="export") + session.attach_modules(modules) + + with torch.no_grad(): + first = transformer(hidden) + second = transformer(hidden) + + session.detach() + assert torch.allclose(first, torch.full_like(first, MARKER)) + assert torch.allclose(second, torch.full_like(second, MARKER)) + assert _reasons(session) == [REASON_SELECTED] + + def test_fastest_compatible_artifact_wins(store, hidden): modules = _pipeline_modules() transformer = modules["transformer"] @@ -527,6 +558,68 @@ def test_unknown_schema_version_is_rejected(store, hidden): assert "unsupported version" in registry.errors[0] +@pytest.mark.parametrize( + ("mutate", "match"), + [ + ( + lambda document: document["signature"]["inputs"][0].update( + requires_grad="false" + ), + "requires_grad", + ), + ( + lambda document: document["evidence"]["benchmark"].update( + speedup=float("nan") + ), + "finite non-negative", + ), + ( + lambda document: document["compatibility"].update( + torch={"min": "3.0", "max_exclusive": "2.0"} + ), + "min must be lower", + ), + ], +) +def test_malformed_acted_on_manifest_fields_fail_closed( + store, hidden, mutate, match +): + modules = _pipeline_modules() + fingerprint, inputs, outputs = _observed_identity( + modules["transformer"].blocks[0], hidden + ) + document = _sections(fingerprint, inputs, outputs) + mutate(document) + _write_bundle(store, document) + + registry = ArtifactRegistry(store) + + assert registry.manifests == [] + assert match in registry.errors[0] + + +def test_artifact_module_names_do_not_collide_after_id_sanitizing(store, hidden): + modules = _pipeline_modules() + fingerprint, inputs, outputs = _observed_identity( + modules["transformer"].blocks[0], hidden + ) + dashed = _sections(fingerprint, inputs, outputs, artifact_id="my-kernel-v1") + underscored = _sections( + fingerprint, inputs, outputs, artifact_id="my_kernel_v1" + ) + dashed_dir = _write_bundle(store, dashed, name="dashed") + underscored_dir = _write_bundle(store, underscored, name="underscored") + + dashed_candidate = load_entry_point( + verify_bundle(dashed_dir), trusted_root=store + ) + underscored_candidate = load_entry_point( + verify_bundle(underscored_dir), trusted_root=store + ) + + assert dashed_candidate.__module__ != underscored_candidate.__module__ + + def test_failing_import_falls_back_to_native(store, hidden): modules = _pipeline_modules() transformer = modules["transformer"] @@ -688,3 +781,33 @@ def test_training_mode_is_detected_and_rejects_inference_only_artifacts(store, m assert _decisions(session)[0]["rejections"] == ["fake-fused-block:execution_mode_unsupported"] finally: detach_graph_dispatch(session) + + +def test_nested_training_module_is_detected_fail_closed(store, monkeypatch, hidden): + modules = _pipeline_modules() + transformer = modules["transformer"] + fingerprint, inputs, outputs = _observed_identity(transformer.blocks[0], hidden) + _write_bundle(store, _sections(fingerprint, inputs, outputs)) + transformer.eval() + transformer.blocks[0].train() + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIR", str(store)) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_ID", "fake/model") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MODEL_REVISION", ANY) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_TRACER", "symbolic") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SCOPES", 64) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_MAX_SHAPES", 8) + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DISTRIBUTED_MODE", "") + monkeypatch.setattr(envs, "FASTVIDEO_OPTIMIZATION_ARTIFACT_DIAGNOSTICS", "") + + session = attach_graph_dispatch(modules) + + assert session is not None + try: + with torch.no_grad(): + transformer(hidden) + transformer(hidden) + assert _decisions(session)[0]["rejections"] == [ + "fake-fused-block:execution_mode_unsupported" + ] + finally: + detach_graph_dispatch(session)