diff --git a/autoware_ml/deployment/backends/__init__.py b/autoware_ml/deployment/backends/__init__.py new file mode 100644 index 00000000..173404f0 --- /dev/null +++ b/autoware_ml/deployment/backends/__init__.py @@ -0,0 +1,20 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Runtime module runners for exported deployment artifacts (ONNX / TensorRT).""" + +from autoware_ml.deployment.backends.onnx_runner import OnnxModuleRunner +from autoware_ml.deployment.backends.tensorrt_runner import TensorRTModuleRunner + +__all__ = ["OnnxModuleRunner", "TensorRTModuleRunner"] diff --git a/autoware_ml/deployment/backends/onnx_runner.py b/autoware_ml/deployment/backends/onnx_runner.py new file mode 100644 index 00000000..48b933ef --- /dev/null +++ b/autoware_ml/deployment/backends/onnx_runner.py @@ -0,0 +1,118 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""ONNX Runtime module runner. + +One implementation of the ONNX Runtime session plumbing (provider selection, +name discovery, tensor conversion, wall-clock timing) shared by every +per-model deployment pipeline, so the run loop cannot drift between backends. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +import time + +import torch + +logger = logging.getLogger(__name__) + +# ONNX Runtime element-type string -> torch dtype, for casting feeds to the +# graph's declared input types (feeding fp16 into a float32 graph is an error). +_ORT_TYPE_TO_TORCH_DTYPE = { + "tensor(float)": torch.float32, + "tensor(float16)": torch.float16, + "tensor(double)": torch.float64, + "tensor(int64)": torch.int64, + "tensor(int32)": torch.int32, + "tensor(int8)": torch.int8, + "tensor(uint8)": torch.uint8, + "tensor(bool)": torch.bool, +} + + +class OnnxModuleRunner: + """Run one exported ONNX module through ONNX Runtime. + + Args: + onnx_path: Path to the exported ``.onnx`` file. + device: Torch device the module should execute on. ``cuda`` requires the + CUDA execution provider. + """ + + def __init__(self, onnx_path: str | Path, device: torch.device) -> None: + import onnxruntime as ort + + onnx_path = Path(onnx_path) + if not onnx_path.exists(): + raise FileNotFoundError(f"ONNX module not found: {onnx_path}") + + self.device = torch.device(device) + if self.device.type == "cuda": + providers = [ + ("CUDAExecutionProvider", {"device_id": self.device.index or 0}), + "CPUExecutionProvider", + ] + else: + providers = ["CPUExecutionProvider"] + + self.session = ort.InferenceSession(str(onnx_path), providers=providers) + if ( + self.device.type == "cuda" + and "CUDAExecutionProvider" not in self.session.get_providers() + ): + raise RuntimeError( + f"CUDA execution provider unavailable for ONNX module {onnx_path.name} — " + "ONNX Runtime silently fell back to CPU, which would corrupt latency " + "comparisons. Install onnxruntime-gpu / check the CUDA setup, or request " + "device=cpu explicitly." + ) + self.input_names = [node.name for node in self.session.get_inputs()] + self.output_names = [node.name for node in self.session.get_outputs()] + self._input_torch_dtypes = { + node.name: _ORT_TYPE_TO_TORCH_DTYPE.get(node.type) for node in self.session.get_inputs() + } + logger.info( + "Loaded ONNX module %s (inputs=%s, outputs=%s, providers=%s)", + onnx_path.name, + self.input_names, + self.output_names, + self.session.get_providers(), + ) + + def run(self, inputs: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], float]: + """Run the module once. + + Args: + inputs: Input tensor per ONNX input name. Tensors may live on any device. + + Returns: + Tuple of (outputs by ONNX output name on :attr:`device`, wall-clock time in ms + for ``session.run`` only — host/device transfers excluded). + """ + feed = {} + for name, tensor in inputs.items(): + expected_dtype = self._input_torch_dtypes.get(name) + if expected_dtype is not None and tensor.dtype != expected_dtype: + tensor = tensor.to(expected_dtype) + feed[name] = tensor.detach().cpu().numpy() + start = time.perf_counter() + raw_outputs = self.session.run(self.output_names, feed) + elapsed_ms = (time.perf_counter() - start) * 1000.0 + outputs = { + name: torch.from_numpy(array).to(self.device) + for name, array in zip(self.output_names, raw_outputs) + } + return outputs, elapsed_ms diff --git a/autoware_ml/deployment/backends/tensorrt_builder.py b/autoware_ml/deployment/backends/tensorrt_builder.py new file mode 100644 index 00000000..cddcbcb6 --- /dev/null +++ b/autoware_ml/deployment/backends/tensorrt_builder.py @@ -0,0 +1,158 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TensorRT engine builder — the build half of the TensorRT backend. + +Owns everything between an ONNX file and a serialized ``.engine``: builder/network +creation, the workspace pool, optimization profiles, plugin loading, and +serialization. The runtime half lives next door in :mod:`.tensorrt_runner`. + +Every network is built STRONGLY TYPED: the ONNX graph's own tensor types are +binding, and precision therefore lives in the ONNX, not in builder flags — +quantized precisions come from Quantize/DequantizeLinear nodes (explicit +quantization), FP16 from the exported graph's tensor types +(:func:`autoware_ml.deployment.onnx.autocast.autocast_to_fp16`). This matches the +TensorRT direction: the weak-typing precision flags (``BuilderFlag.FP16`` & co) +were deprecated in TensorRT 10.12 and removed in TensorRT 11, where all networks +are strongly typed. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any, Mapping, Sequence + +from autoware_ml.deployment.config import ShapeProfile + +logger = logging.getLogger(__name__) + + +def load_tensorrt_plugin_libraries(plugin_libraries: Sequence[str] | None) -> None: + """Load custom TensorRT plugin shared libraries before plugin registry init. + + Args: + plugin_libraries: Paths to plugin ``.so`` files, or None/empty for none. + + Raises: + FileNotFoundError: If a configured plugin library does not exist. + """ + if not plugin_libraries: + return + import ctypes + + for library in plugin_libraries: + library_path = Path(library) + if not library_path.exists(): + raise FileNotFoundError(f"TensorRT plugin library not found: {library_path}") + ctypes.CDLL(str(library_path), mode=ctypes.RTLD_GLOBAL) + logger.info("Loaded TensorRT plugin library: %s", library_path) + + +def _create_builder(workspace_size: int, plugin_libraries: Sequence[str]): + """Create ``(builder, network, parser, config)`` for one strongly typed engine build.""" + import tensorrt as trt + + # Custom plugins must be loadable before plugin registry initialization. + load_tensorrt_plugin_libraries(plugin_libraries) + + trt_logger = trt.Logger(trt.Logger.WARNING) + trt.init_libnvinfer_plugins(trt_logger, "") + builder = trt.Builder(trt_logger) + + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + parser = trt.OnnxParser(network, trt_logger) + config = builder.create_builder_config() + + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, int(workspace_size)) + logger.info("Workspace size: %.2f GB", workspace_size / (1024**3)) + return builder, network, parser, config + + +def _parse_onnx_file(parser: Any, onnx_path: Path) -> None: + with open(onnx_path, "rb") as f: + onnx_data = f.read() + + if not parser.parse(onnx_data): + errors = [parser.get_error(i) for i in range(parser.num_errors)] + error_msg = "\n".join(f"TensorRT parser error {i}: {err}" for i, err in enumerate(errors)) + if "plugin" in error_msg.lower() or "INVALID_NODE" in error_msg: + error_msg += ( + "\nHint: if this graph carries custom ops (e.g. autoware::*), make sure " + "deploy.tensorrt.plugin_libraries lists the plugin .so for this environment." + ) + raise RuntimeError(f"Failed to parse ONNX file:\n{error_msg}") + + logger.info("Successfully parsed ONNX file") + + +def _create_optimization_profile(builder: Any, input_shapes: Mapping[str, ShapeProfile]): + profile = builder.create_optimization_profile() + for input_name, shapes in input_shapes.items(): + profile.set_shape( + input_name, + min=list(shapes.min_shape), + opt=list(shapes.opt_shape), + max=list(shapes.max_shape), + ) + logger.info( + "Optimization profile for '%s': min=%s, opt=%s, max=%s", + input_name, + list(shapes.min_shape), + list(shapes.opt_shape), + list(shapes.max_shape), + ) + return profile + + +def build_engine( + onnx_path: Path, + output_path: Path, + *, + workspace_size: int = 1 << 32, + plugin_libraries: Sequence[str] = (), + input_shapes: Mapping[str, ShapeProfile] | None = None, +) -> None: + """Build and serialize one strongly typed TensorRT engine from an ONNX file. + + Precision is read from the ONNX graph (module docstring); there are no + precision knobs here by design. + + Args: + onnx_path: Exported ONNX model. + output_path: Destination ``.engine`` path. + workspace_size: Workspace memory-pool limit in bytes. + plugin_libraries: Custom plugin ``.so`` paths to load before parsing. + input_shapes: Optimization-profile shapes per dynamic input; ``None``/empty + builds without an explicit profile (static-shape graphs). + + Raises: + RuntimeError: When ONNX parsing or the engine build fails. + """ + logger.info("Building TensorRT engine (strongly typed)...") + builder, network, parser, config = _create_builder(workspace_size, plugin_libraries) + _parse_onnx_file(parser, onnx_path) + + if input_shapes: + config.add_optimization_profile(_create_optimization_profile(builder, input_shapes)) + + logger.info("Building TensorRT engine (this may take a while)...") + serialized_engine = builder.build_serialized_network(network, config) + if serialized_engine is None: + raise RuntimeError("Failed to build TensorRT engine.") + + with open(output_path, "wb") as f: + f.write(serialized_engine) + + logger.info("Successfully built TensorRT engine: %s", output_path) diff --git a/autoware_ml/deployment/backends/tensorrt_runner.py b/autoware_ml/deployment/backends/tensorrt_runner.py new file mode 100644 index 00000000..6b00255d --- /dev/null +++ b/autoware_ml/deployment/backends/tensorrt_runner.py @@ -0,0 +1,205 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared TensorRT engine runner (torch-native I/O). + +One battle-tested implementation of the TensorRT run loop, reused by every +per-model deployment pipeline so the GPU plumbing cannot drift between +backends. Unlike the classic pycuda variant this runner uses torch CUDA +tensors as device buffers: inputs that already live on the GPU are bound +in place (no host round-trip) and outputs are returned as CUDA tensors, +which is exactly what the downstream torch stages (scatter, decode) want. + +Timing brackets only ``execute_async_v3`` with CUDA events on the current +torch stream, so the reported time is the engine's pure GPU compute. +""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import numpy as np +import torch + +logger = logging.getLogger(__name__) + + +def load_trt_engine(engine_path: str | Path, *, component_name: str | None = None): + """Deserialize a TensorRT engine and create its execution context, failing loud. + + Args: + engine_path: Path to the serialized ``.engine`` file. + component_name: Optional component label for error messages. + + Returns: + Tuple of ``(engine, execution_context)``. + + Raises: + RuntimeError: If deserialization or context creation fails + (context failure is usually GPU out-of-memory). + """ + import tensorrt as trt + + engine_path = Path(engine_path) + label = component_name or engine_path.name + if not engine_path.exists(): + raise FileNotFoundError(f"TensorRT engine not found: {engine_path}") + + trt_logger = trt.Logger(trt.Logger.WARNING) + trt.init_libnvinfer_plugins(trt_logger, "") + runtime = trt.Runtime(trt_logger) + with open(engine_path, "rb") as f: + engine = runtime.deserialize_cuda_engine(f.read()) + if engine is None: + raise RuntimeError(f"Failed to deserialize TensorRT engine: {engine_path}") + + context = engine.create_execution_context() + if context is None: + raise RuntimeError( + f"Failed to create TensorRT execution context for {label} (likely GPU out-of-memory)." + ) + return engine, context + + +def list_trt_io_names(engine) -> tuple[list[str], list[str]]: + """Return ``(input_names, output_names)`` in TensorRT tensor-index order.""" + import tensorrt as trt + + inputs: list[str] = [] + outputs: list[str] = [] + for i in range(engine.num_io_tensors): + name = engine.get_tensor_name(i) + if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + inputs.append(name) + else: + outputs.append(name) + return inputs, outputs + + +def _trt_dtype_to_torch(trt_dtype) -> torch.dtype: + """Map a TensorRT dtype to the matching torch dtype, failing loud on unknowns. + + Guessing a size (e.g. defaulting to float32) would mis-size the GPU buffer and + silently corrupt the data, so unknown dtypes raise via TensorRT's own ``nptype``. + """ + import tensorrt as trt + + numpy_dtype = np.dtype(trt.nptype(trt_dtype)) + return torch.from_numpy(np.zeros(0, dtype=numpy_dtype)).dtype + + +class TensorRTModuleRunner: + """Run one serialized TensorRT engine with torch tensors as device buffers. + + Args: + engine_path: Path to the serialized ``.engine`` file. + device: CUDA device the engine executes on. + """ + + def __init__(self, engine_path: str | Path, device: torch.device) -> None: + self.device = torch.device(device) + if self.device.type != "cuda": + raise ValueError(f"TensorRT requires a CUDA device, got {self.device}.") + self.engine, self.context = load_trt_engine(engine_path) + self.input_names, self.output_names = list_trt_io_names(self.engine) + # Bindings persist across calls: input shapes are re-declared and output buffers + # re-allocated only when a shape actually changes. Re-binding every call was + # measured to inflate the reported per-engine time by ~0.5 ms/engine on + # BEVFusion (the deployment report showed 6.08 ms for a chain whose paired + # single-window measurement is 5.16 ms). + self._bound_input_shapes: dict[str, tuple[int, ...]] = {} + self._output_buffers: dict[str, torch.Tensor] = {} + self._start_event = torch.cuda.Event(enable_timing=True) + self._end_event = torch.cuda.Event(enable_timing=True) + logger.info( + "Loaded TensorRT engine %s (inputs=%s, outputs=%s)", + Path(engine_path).name, + self.input_names, + self.output_names, + ) + + def _cast_to_binding_dtype(self, tensor_name: str, tensor: torch.Tensor) -> torch.Tensor: + """Return ``tensor`` on :attr:`device`, contiguous, in the engine binding's dtype. + + Matching the binding dtype is critical for FP16 engines: a graph traced with + FP32 inputs may bind as ``HALF``, and feeding float32 bytes into a HALF binding + misaligns the GPU buffer and silently corrupts the activations. + """ + target_dtype = _trt_dtype_to_torch(self.engine.get_tensor_dtype(tensor_name)) + if tensor.dtype != target_dtype: + logger.debug( + "[trt-io] casting tensor %r: %s -> %s", tensor_name, tensor.dtype, target_dtype + ) + tensor = tensor.to(target_dtype) + return tensor.to(self.device).contiguous() + + def run(self, inputs: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], float]: + """Run the engine once and return ``(outputs_by_name, pure_gpu_time_ms)``. + + Args: + inputs: Engine input tensor name -> torch tensor (any device / dtype; + cast and moved as needed). + + Returns: + Tuple of (outputs by engine output name as CUDA tensors, pure-GPU time in + ms measured with CUDA events around ``execute_async_v3`` only). + + NOTE: output tensors are owned by the runner and REUSED on the next ``run`` + with the same shapes — consume (or copy) them before calling ``run`` again. + Every current caller is strictly sequential per runner (evaluation processes a + frame to completion; verification compares per batch, and its reference and + test pipelines hold separate runners). + + Raises: + RuntimeError: If ``execute_async_v3`` reports a failure status. + """ + device_inputs = { + name: self._cast_to_binding_dtype(name, tensor) for name, tensor in inputs.items() + } + shapes_changed = False + for name, tensor in device_inputs.items(): + shape = tuple(tensor.shape) + if self._bound_input_shapes.get(name) != shape: + self.context.set_input_shape(name, shape) + self._bound_input_shapes[name] = shape + shapes_changed = True + # Input tensors arrive from the caller, so their addresses change per call. + self.context.set_tensor_address(name, int(tensor.data_ptr())) + + # Output shapes can depend on the input shapes, so re-derive (and re-allocate + # only what actually changed) when any input shape did. + if shapes_changed or not self._output_buffers: + for name in self.output_names: + shape = tuple(self.context.get_tensor_shape(name)) + buffer = self._output_buffers.get(name) + if buffer is None or tuple(buffer.shape) != shape: + buffer = torch.empty( + shape, + dtype=_trt_dtype_to_torch(self.engine.get_tensor_dtype(name)), + device=self.device, + ) + self._output_buffers[name] = buffer + self.context.set_tensor_address(name, int(buffer.data_ptr())) + + stream = torch.cuda.current_stream(self.device) + self._start_event.record(stream) + succeeded = self.context.execute_async_v3(stream_handle=stream.cuda_stream) + if not succeeded: + raise RuntimeError("TensorRT execute_async_v3 returned failure status.") + self._end_event.record(stream) + self._end_event.synchronize() + gpu_time_ms = float(self._start_event.elapsed_time(self._end_event)) + + return dict(self._output_buffers), gpu_time_ms diff --git a/autoware_ml/deployment/config.py b/autoware_ml/deployment/config.py new file mode 100644 index 00000000..c2c4d6e0 --- /dev/null +++ b/autoware_ml/deployment/config.py @@ -0,0 +1,372 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Typed view of the Hydra ``deploy`` config section — parsed once, typo-guarded. + +Layout (mirrors the stage graph): + +.. code-block:: yaml + + deploy: + onnx: { enabled, dynamo, opset_version, do_constant_folding, precision, modify_graph } # global + tensorrt: { enabled, workspace_size, plugin_libraries } # global + stages: # per GraphStage, keyed by stage name + : + onnx: { dynamic_axes | dynamic_shapes } + tensorrt: { input_shapes: { : { min_shape, opt_shape, max_shape } } } + verification: { enabled, tolerance, num_verify_batches, scenarios } + evaluation: { enabled, num_samples, num_warmup, backends: { : { enabled, device } } } + +Every mapping rejects unknown keys: a misspelled option would otherwise silently fall +back to a default (``opset_versoin`` exports with the wrong opset; a stage name that +does not match the model's declaration silently drops its shape profile). +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from enum import Enum +from typing import Any + +from autoware_ml.deployment.verification.backend_verifier import VerificationScenario +from autoware_ml.types.backend import Backend +from autoware_ml.utils.config_parsing import reject_unknown_keys as _reject_unknown + + +def _mapping(raw: Any, where: str) -> Mapping[str, Any]: + if raw is None: + return {} + if not isinstance(raw, Mapping): + raise TypeError(f"{where} must be a mapping, got {type(raw).__name__}.") + return raw + + +class OnnxPrecision(str, Enum): + """Precision of the exported graphs. + + Engines build strongly typed, so this is where FP16 is decided: ``FP16`` runs + ModelOpt AutoCast on every exported stage without Q/DQ nodes (quantized stages + keep the precision their checkpoint bakes in); ``FP32`` exports as traced. + """ + + FP32 = "fp32" + FP16 = "fp16" + + +@dataclass(frozen=True) +class OnnxConfig: + """Global ONNX export options (``deploy.onnx``).""" + + enabled: bool = True + dynamo: bool = True + opset_version: int = 21 + do_constant_folding: bool = True + precision: OnnxPrecision = OnnxPrecision.FP32 + #: Optional Hydra-instantiable graph modifier applied to every exported ONNX. + modify_graph: Any = None + + KNOWN_KEYS = frozenset( + {"enabled", "dynamo", "opset_version", "do_constant_folding", "precision", "modify_graph"} + ) + + @classmethod + def from_dict(cls, raw: Any) -> OnnxConfig: + raw = _mapping(raw, "deploy.onnx") + _reject_unknown(raw, cls.KNOWN_KEYS, "deploy.onnx") + raw_precision = str(raw.get("precision", OnnxPrecision.FP32.value)).lower() + try: + precision = OnnxPrecision(raw_precision) + except ValueError: + raise ValueError( + f"Unknown deploy.onnx.precision {raw_precision!r}. " + f"Valid: {[p.value for p in OnnxPrecision]} " + "(quantized precisions come from the checkpoint's Q/DQ nodes, not from here)." + ) from None + return cls( + enabled=bool(raw.get("enabled", True)), + dynamo=bool(raw.get("dynamo", True)), + opset_version=int(raw.get("opset_version", 21)), + do_constant_folding=bool(raw.get("do_constant_folding", True)), + precision=precision, + modify_graph=raw.get("modify_graph"), + ) + + +@dataclass(frozen=True) +class TensorRTConfig: + """Global TensorRT build options (``deploy.tensorrt``). + + There is no precision knob: engines build strongly typed and read precision from + the ONNX graph (Q/DQ for quantized precisions, tensor types for FP16 — see + ``deploy.onnx.precision``). TensorRT deprecated the weak-typing precision flags in + 10.12 and removed them in 11. + """ + + enabled: bool = True + workspace_size: int = 1 << 32 + plugin_libraries: tuple[str, ...] = () + + KNOWN_KEYS = frozenset({"enabled", "workspace_size", "plugin_libraries"}) + + @classmethod + def from_dict(cls, raw: Any) -> TensorRTConfig: + raw = _mapping(raw, "deploy.tensorrt") + _reject_unknown(raw, cls.KNOWN_KEYS, "deploy.tensorrt") + return cls( + enabled=bool(raw.get("enabled", True)), + workspace_size=int(raw.get("workspace_size", 1 << 32)), + plugin_libraries=tuple(str(p) for p in (raw.get("plugin_libraries") or ())), + ) + + +@dataclass(frozen=True) +class StageOnnxConfig: + """Per-stage ONNX options (``deploy.stages..onnx``). + + Shape declarations plus an optional per-stage precision; input/output *names* are + never configured — they come from the stage declaration. + """ + + #: Legacy exporter (``dynamo=false``): ``{tensor_name: {dim_index: dim_name}}``. + dynamic_axes: Mapping[str, Mapping[int, str]] | None = None + #: Dynamo exporter: ``{input_name: {dim_index: dim_name | {name, min, max}}}``. + dynamic_shapes: Mapping[str, Mapping[int, Any]] | None = None + #: Overrides ``deploy.onnx.precision`` for this stage only — for a pipeline whose + #: stages need different precisions (one numerically fragile head kept FP32, say). + #: ``None`` inherits the global setting. + precision: OnnxPrecision | None = None + + KNOWN_KEYS = frozenset({"dynamic_axes", "dynamic_shapes", "precision"}) + + @classmethod + def from_dict(cls, raw: Any, stage: str) -> StageOnnxConfig: + raw = _mapping(raw, f"deploy.stages.{stage}.onnx") + _reject_unknown(raw, cls.KNOWN_KEYS, f"deploy.stages.{stage}.onnx") + raw_precision = raw.get("precision") + try: + precision = OnnxPrecision(str(raw_precision).lower()) if raw_precision else None + except ValueError: + raise ValueError( + f"deploy.stages.{stage}.onnx.precision={raw_precision!r} — valid values: " + f"{[p.value for p in OnnxPrecision]}." + ) from None + return cls( + dynamic_axes=raw.get("dynamic_axes"), + dynamic_shapes=raw.get("dynamic_shapes"), + precision=precision, + ) + + +@dataclass(frozen=True) +class ShapeProfile: + """One TensorRT optimization-profile entry.""" + + min_shape: tuple[int, ...] + opt_shape: tuple[int, ...] + max_shape: tuple[int, ...] + + KNOWN_KEYS = frozenset({"min_shape", "opt_shape", "max_shape"}) + + @classmethod + def from_dict(cls, raw: Any, where: str) -> ShapeProfile: + raw = _mapping(raw, where) + _reject_unknown(raw, cls.KNOWN_KEYS, where) + missing = cls.KNOWN_KEYS - set(raw) + if missing: + raise ValueError(f"{where} is incomplete: missing {sorted(missing)}.") + return cls( + min_shape=tuple(int(x) for x in raw["min_shape"]), + opt_shape=tuple(int(x) for x in raw["opt_shape"]), + max_shape=tuple(int(x) for x in raw["max_shape"]), + ) + + +@dataclass(frozen=True) +class StageTensorRTConfig: + """Per-stage TensorRT options (``deploy.stages..tensorrt``).""" + + input_shapes: Mapping[str, ShapeProfile] = field(default_factory=dict) + + KNOWN_KEYS = frozenset({"input_shapes"}) + + @classmethod + def from_dict(cls, raw: Any, stage: str) -> StageTensorRTConfig: + where = f"deploy.stages.{stage}.tensorrt" + raw = _mapping(raw, where) + _reject_unknown(raw, cls.KNOWN_KEYS, where) + shapes = _mapping(raw.get("input_shapes"), f"{where}.input_shapes") + return cls( + input_shapes={ + str(name): ShapeProfile.from_dict(profile, f"{where}.input_shapes.{name}") + for name, profile in shapes.items() + } + ) + + +@dataclass(frozen=True) +class StageConfig: + """Everything configured for one exportable stage.""" + + onnx: StageOnnxConfig = StageOnnxConfig() + tensorrt: StageTensorRTConfig = StageTensorRTConfig() + + KNOWN_KEYS = frozenset({"onnx", "tensorrt"}) + + @classmethod + def from_dict(cls, raw: Any, stage: str) -> StageConfig: + raw = _mapping(raw, f"deploy.stages.{stage}") + _reject_unknown(raw, cls.KNOWN_KEYS, f"deploy.stages.{stage}") + return cls( + onnx=StageOnnxConfig.from_dict(raw.get("onnx"), stage), + tensorrt=StageTensorRTConfig.from_dict(raw.get("tensorrt"), stage), + ) + + +@dataclass(frozen=True) +class VerificationConfig: + """Cross-backend parity stage (``deploy.verification``).""" + + enabled: bool = False + #: Default absolute tolerance on raw graph outputs. Lossy backends (fp16 / int8) + #: set an explicit per-scenario ``tolerance`` instead of loosening this. + tolerance: float = 0.01 + num_verify_batches: int = 1 + scenarios: tuple[VerificationScenario, ...] = () + + KNOWN_KEYS = frozenset({"enabled", "tolerance", "num_verify_batches", "scenarios"}) + + @classmethod + def from_dict(cls, raw: Any) -> VerificationConfig: + raw = _mapping(raw, "deploy.verification") + _reject_unknown(raw, cls.KNOWN_KEYS, "deploy.verification") + return cls( + enabled=bool(raw.get("enabled", False)), + tolerance=float(raw.get("tolerance", 0.01)), + num_verify_batches=int(raw.get("num_verify_batches", 1)), + scenarios=tuple( + VerificationScenario.from_dict(s) for s in (raw.get("scenarios") or ()) + ), + ) + + +@dataclass(frozen=True) +class BackendEvaluationConfig: + """One entry of ``deploy.evaluation.backends``.""" + + enabled: bool = True + device: str = "cuda" + + KNOWN_KEYS = frozenset({"enabled", "device"}) + + @classmethod + def from_dict(cls, raw: Any, backend: str) -> BackendEvaluationConfig: + where = f"deploy.evaluation.backends.{backend}" + raw = _mapping(raw, where) + _reject_unknown(raw, cls.KNOWN_KEYS, where) + return cls(enabled=bool(raw.get("enabled", True)), device=str(raw.get("device", "cuda"))) + + +@dataclass(frozen=True) +class EvaluationConfig: + """Per-backend ground-truth evaluation stage (``deploy.evaluation``).""" + + enabled: bool = False + #: Split the backends are scored on: ``test`` (default, the predict dataloader) or + #: ``val`` — for when the test split is unavailable or held back. Metric keys carry + #: the split, so a val evaluation reports under ``val/{backend}/...``. + split: str = "test" + #: Samples per backend; -1 = the whole split. + num_samples: int = -1 + #: Extra re-runs of the first batch that prime the GPU / TensorRT (discarded). + num_warmup: int = 2 + backends: Mapping[Backend, BackendEvaluationConfig] = field(default_factory=dict) + + KNOWN_KEYS = frozenset({"enabled", "split", "num_samples", "num_warmup", "backends"}) + + @classmethod + def from_dict(cls, raw: Any) -> EvaluationConfig: + raw = _mapping(raw, "deploy.evaluation") + _reject_unknown(raw, cls.KNOWN_KEYS, "deploy.evaluation") + backends = _mapping(raw.get("backends"), "deploy.evaluation.backends") + split = str(raw.get("split", "test")) + if split not in ("test", "val"): + raise ValueError( + f"deploy.evaluation.split={split!r} — valid values: 'test', 'val'." + ) + return cls( + enabled=bool(raw.get("enabled", False)), + split=split, + num_samples=int(raw.get("num_samples", -1)), + num_warmup=int(raw.get("num_warmup", 2)), + backends={ + Backend.parse(name): BackendEvaluationConfig.from_dict(cfg, name) + for name, cfg in backends.items() + }, + ) + + def enabled_backends(self) -> list[tuple[Backend, BackendEvaluationConfig]]: + """Backends with ``enabled: true``, in configuration order.""" + return [(backend, cfg) for backend, cfg in self.backends.items() if cfg.enabled] + + +@dataclass(frozen=True) +class DeployConfig: + """Typed view of the whole ``deploy`` section.""" + + onnx: OnnxConfig = OnnxConfig() + tensorrt: TensorRTConfig = TensorRTConfig() + stages: Mapping[str, StageConfig] = field(default_factory=dict) + verification: VerificationConfig = VerificationConfig() + evaluation: EvaluationConfig = EvaluationConfig() + + KNOWN_KEYS = frozenset({"onnx", "tensorrt", "stages", "verification", "evaluation"}) + + @classmethod + def from_dict(cls, raw: Any) -> DeployConfig: + """Parse the resolved ``deploy`` mapping (``OmegaConf.to_container(..., resolve=True)``). + + Raises: + ValueError: On unknown keys anywhere in the section or an invalid value. + TypeError: When a sub-section is not a mapping. + """ + raw = _mapping(raw, "deploy") + _reject_unknown(raw, cls.KNOWN_KEYS, "deploy") + stages = _mapping(raw.get("stages"), "deploy.stages") + return cls( + onnx=OnnxConfig.from_dict(raw.get("onnx")), + tensorrt=TensorRTConfig.from_dict(raw.get("tensorrt")), + stages={ + str(name): StageConfig.from_dict(cfg, str(name)) for name, cfg in stages.items() + }, + verification=VerificationConfig.from_dict(raw.get("verification")), + evaluation=EvaluationConfig.from_dict(raw.get("evaluation")), + ) + + def stage(self, name: str) -> StageConfig: + """Per-stage options, defaulting to empty when the stage has no entry.""" + return self.stages.get(name, StageConfig()) + + def check_stage_names(self, declared: Sequence[str]) -> None: + """Raise when ``deploy.stages`` names a stage the model does not declare. + + A stage name typo would otherwise silently drop that stage's dynamic axes and + TensorRT shape profile. + """ + unknown = sorted(set(self.stages) - set(declared)) + if unknown: + raise ValueError( + f"deploy.stages configures unknown stage(s) {unknown}; the model declares " + f"exportable stages {list(declared)}." + ) diff --git a/autoware_ml/deployment/pipeline.py b/autoware_ml/deployment/pipeline.py new file mode 100644 index 00000000..9a0a5015 --- /dev/null +++ b/autoware_ml/deployment/pipeline.py @@ -0,0 +1,278 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run a model's stage graph on one backend. + +:class:`StagedPipeline` is the generic executor of a :mod:`~autoware_ml.deployment.stages` +declaration. Glue stages always run their PyTorch callable; exportable stages run either +their module (``pytorch`` backend) or the runner of their exported artifact +(``onnx`` / ``tensorrt``). Every backend therefore consumes and produces exactly the +same named tensors, so cross-backend differences are pure backend differences. + +Preprocessing (``model.preprocess_batch``) and decoding/metrics stay outside the +pipeline: it starts from preprocessed inputs and ends at the final graph stage's raw +outputs plus a per-stage latency breakdown. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from contextlib import contextmanager +import copy +from dataclasses import dataclass, field +from pathlib import Path +import time +from typing import Any, Iterator, Sequence + +import torch +from torch import nn + +from autoware_ml.dataclasses.multi_task_batch_inputs import MultiTaskBatchInputs +from autoware_ml.dataclasses.multi_task_predictions import MultiTaskPredictions +from autoware_ml.deployment.stages import ( + GraphStage, + Stage, + StageContext, + TorchStage, + artifact_path, + final_stage, + graph_stages, + validate_stages, +) +from autoware_ml.types.backend import Backend + +AssembleFn = Callable[[Mapping[str, torch.Tensor]], MultiTaskPredictions] + + +@dataclass +class PipelineResult: + """Raw outputs of one pipeline run. + + Attributes: + outputs: Final graph stage's output tensors keyed by their ONNX output name. + output_names: Output names in the declared (frozen ABI) order. + stage_times_ms: Per-stage latency in milliseconds. + graph_stage_names: Names of the exportable stages; :attr:`model_ms` sums these + (pure GPU time for TensorRT). + """ + + outputs: dict[str, torch.Tensor] + output_names: list[str] + stage_times_ms: dict[str, float] = field(default_factory=dict) + graph_stage_names: tuple[str, ...] = () + + @property + def model_ms(self) -> float: + """Summed latency of the exportable stages.""" + return sum(self.stage_times_ms.get(name, 0.0) for name in self.graph_stage_names) + + def ordered_outputs(self) -> list[torch.Tensor]: + """Return the outputs in the frozen ABI order.""" + return [self.outputs[name] for name in self.output_names] + + +@contextmanager +def cuda_synced_timer( + times_ms: dict[str, float], stage: str, device: torch.device +) -> Iterator[None]: + """Time a block with wall clock, synchronizing CUDA so async kernels are included.""" + if device.type == "cuda": + torch.cuda.synchronize(device) + start = time.perf_counter() + yield + if device.type == "cuda": + torch.cuda.synchronize(device) + times_ms[stage] = times_ms.get(stage, 0.0) + (time.perf_counter() - start) * 1000.0 + + +class _ModuleRunner: + """Run a GraphStage's own module (the pytorch backend), timed like an artifact runner.""" + + def __init__(self, stage: GraphStage, device: torch.device) -> None: + module = stage.module + module_device = next(module.parameters(), torch.empty(0)).device + if module_device != device: + # Never move the shared model's modules: this pipeline gets its own copy. + module = copy.deepcopy(module).to(device) + self.module: nn.Module = module.eval() + self.device = device + self.stage = stage + + def run(self, inputs: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], float]: + args = tuple(inputs[name].to(self.device) for name in self.stage.inputs) + times: dict[str, float] = {} + with cuda_synced_timer(times, "run", self.device): + raw = self.module(*args) + if isinstance(raw, torch.Tensor): + raw = (raw,) + if len(raw) != len(self.stage.outputs): + raise ValueError( + f"GraphStage {self.stage.name!r} returned {len(raw)} tensor(s) but declares " + f"outputs {list(self.stage.outputs)}." + ) + return dict(zip(self.stage.outputs, raw)), times["run"] + + +def _artifact_runner( + stage: GraphStage, backend: Backend, device: torch.device, artifacts_dir: Path +): + path = artifact_path(artifacts_dir, stage.name, backend) + if backend is Backend.ONNX: + from autoware_ml.deployment.backends.onnx_runner import OnnxModuleRunner + + return OnnxModuleRunner(path, device) + if backend is Backend.TENSORRT: + from autoware_ml.deployment.backends.tensorrt_runner import TensorRTModuleRunner + + return TensorRTModuleRunner(path, device) + raise ValueError(f"No artifact runner for backend {backend}.") + + +class StagedPipeline: + """Execute a stage graph on one backend. + + Args: + stages: The model's stage declaration (``model.build_stages()``). + backend: Backend that runs the exportable stages. + device: Device the exportable stages execute on; glue stages hand their results + over on this device. + artifacts_dir: Directory holding ``.onnx`` / ``.engine`` (non-pytorch backends). + assemble: The model's ``assemble_predictions`` hook, used by :meth:`assemble`. + """ + + def __init__( + self, + stages: Sequence[Stage], + backend: str | Backend, + device: torch.device, + artifacts_dir: str | Path | None = None, + assemble: AssembleFn | None = None, + ) -> None: + self.stages = validate_stages(stages) + self.backend = Backend.parse(backend) + self.device = torch.device(device) + self._assemble = assemble + self.final_stage = final_stage(self.stages) + self.output_names = list(self.final_stage.outputs) + self.graph_stage_names = tuple(stage.name for stage in graph_stages(self.stages)) + + self._runners: dict[str, Any] = {} + #: Graph stages that run their PyTorch module on this (non-pytorch) backend + #: because they declare it as a fallback — reports must say so, or a backend + #: column can silently be the pytorch numbers under another name. + self.fallback_stage_names: tuple[str, ...] = tuple( + stage.name + for stage in graph_stages(self.stages) + if self.backend is not Backend.PYTORCH and self.backend in stage.torch_fallback_backends + ) + for stage in graph_stages(self.stages): + if self.backend is Backend.PYTORCH or self.backend in stage.torch_fallback_backends: + self._runners[stage.name] = _ModuleRunner(stage, self.device) + else: + if artifacts_dir is None: + raise ValueError( + f"artifacts_dir is required for the {self.backend.value} backend." + ) + self._runners[stage.name] = _artifact_runner( + stage, self.backend, self.device, Path(artifacts_dir) + ) + + def run(self, batch_inputs: MultiTaskBatchInputs) -> tuple[PipelineResult, StageContext]: + """Run every stage and return the result together with the full context. + + The context (every named tensor produced along the way) is what export uses + as trace inputs for each graph stage. + """ + context = StageContext(batch_inputs=batch_inputs, device=self.device) + times: dict[str, float] = {} + with torch.no_grad(): + for stage in self.stages: + if isinstance(stage, TorchStage): + with cuda_synced_timer(times, stage.name, self.device): + produced = stage.run(context) + context.tensors.update(produced) + else: + inputs = {name: context[name] for name in stage.inputs} + outputs, elapsed_ms = self._runners[stage.name].run(inputs) + times[stage.name] = times.get(stage.name, 0.0) + elapsed_ms + context.tensors.update({name: outputs[name] for name in stage.outputs}) + result = PipelineResult( + outputs={name: context[name] for name in self.output_names}, + output_names=list(self.output_names), + stage_times_ms=times, + graph_stage_names=self.graph_stage_names, + ) + return result, context + + def infer(self, batch_inputs: MultiTaskBatchInputs) -> PipelineResult: + """Run every stage and return the final raw outputs plus timing.""" + result, _ = self.run(batch_inputs) + return result + + def assemble( + self, result: PipelineResult, device: torch.device | None = None + ) -> MultiTaskPredictions: + """Turn a result into the model's predictions (what metrics consume). + + Args: + result: Output of :meth:`infer`. + device: Optional device to move the tensors to first (e.g. the metrics device). + """ + if self._assemble is None: + raise RuntimeError( + "StagedPipeline was built without the model's assemble_predictions hook." + ) + fields: dict[str, torch.Tensor] = {} + for onnx_name, field_name in self.final_stage.output_fields: + tensor = result.outputs[onnx_name] + if device is not None: + tensor = tensor.to(device) + fields[field_name] = tensor.float().contiguous() + return self._assemble(fields) + + +class PipelineCache: + """Build each ``(backend, device)`` pipeline once and share it across stages. + + Verification and evaluation both need pipelines for the same backends; loading + an ONNX session or deserializing a TensorRT engine twice per deploy run is waste. + One instance lives for one deploy run over one set of artifacts — there is no + invalidation: if an artifact on disk changes, use a fresh cache. + """ + + def __init__( + self, + stages: Sequence[Stage], + artifacts_dir: str | Path, + assemble: AssembleFn, + ) -> None: + self.stages = validate_stages(stages) + self.artifacts_dir = Path(artifacts_dir) + self._assemble = assemble + self._pipelines: dict[tuple[Backend, str], StagedPipeline] = {} + + def get(self, backend: str | Backend, device: str | torch.device) -> StagedPipeline: + """Return the cached pipeline for ``(backend, device)``, building it on first use.""" + backend = Backend.parse(backend) + device = torch.device(device) + key = (backend, str(device)) + if key not in self._pipelines: + self._pipelines[key] = StagedPipeline( + self.stages, + backend=backend, + device=device, + artifacts_dir=self.artifacts_dir, + assemble=self._assemble, + ) + return self._pipelines[key] diff --git a/autoware_ml/deployment/verification/__init__.py b/autoware_ml/deployment/verification/__init__.py new file mode 100644 index 00000000..158363df --- /dev/null +++ b/autoware_ml/deployment/verification/__init__.py @@ -0,0 +1,33 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Cross-backend numerical verification of exported deployment artifacts.""" + +from autoware_ml.deployment.verification.backend_verifier import ( + BackendVerifier, + VerificationScenario, +) +from autoware_ml.deployment.verification.output_comparator import ( + OutputComparator, + OutputDiffSummary, + TensorDiffDetail, +) + +__all__ = [ + "BackendVerifier", + "VerificationScenario", + "OutputComparator", + "OutputDiffSummary", + "TensorDiffDetail", +] diff --git a/autoware_ml/deployment/verification/backend_verifier.py b/autoware_ml/deployment/verification/backend_verifier.py new file mode 100644 index 00000000..dac84d26 --- /dev/null +++ b/autoware_ml/deployment/verification/backend_verifier.py @@ -0,0 +1,197 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Scenario-driven cross-backend verification. + +Each scenario names a reference and a test backend (with devices); the verifier +runs both pipelines on the same preprocessed batches and compares the final raw +graph outputs element-wise against an absolute tolerance (the verifier default, +or the scenario's own ``tolerance`` override). No ground truth is involved: this +is numerical parity, a peer of — not a form of — evaluation. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import logging +from typing import TYPE_CHECKING, Any, Mapping, Sequence + +from autoware_ml.dataclasses.multi_task_batch_inputs import MultiTaskBatchInputs +from autoware_ml.deployment.verification.output_comparator import OutputComparator +from autoware_ml.types.backend import Backend + +if TYPE_CHECKING: + from autoware_ml.deployment.pipeline import PipelineCache + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class VerificationScenario: + """One reference-vs-test backend comparison. + + Attributes: + ref_backend: Reference backend name (``pytorch`` / ``onnx`` / ``tensorrt``). + ref_device: Device string for the reference pipeline (e.g. ``cuda`` / ``cpu``). + test_backend: Test backend name. + test_device: Device string for the test pipeline. + tolerance: Optional per-scenario absolute tolerance overriding the + verifier default (e.g. relaxed for int8, tight for fp32-vs-onnx). + """ + + ref_backend: str + ref_device: str + test_backend: str + test_device: str + tolerance: float | None = None + + @classmethod + def from_dict(cls, raw: Mapping[str, Any]) -> VerificationScenario: + """Build a scenario from a ``{ref: {backend, device}, test: {backend, device}}`` mapping. + + An optional top-level ``tolerance`` key overrides the verifier default + for this scenario only. + """ + try: + raw_tolerance = raw.get("tolerance") + return cls( + ref_backend=Backend.parse(raw["ref"]["backend"]).value, + ref_device=str(raw["ref"].get("device", "cuda")), + test_backend=Backend.parse(raw["test"]["backend"]).value, + test_device=str(raw["test"].get("device", "cuda")), + tolerance=float(raw_tolerance) if raw_tolerance is not None else None, + ) + except (AttributeError, KeyError, TypeError) as error: + raise ValueError( + "A verification scenario must look like " + "{ref: {backend: ..., device: ...}, test: {backend: ..., device: ...}, " + "tolerance: }, " + f"got: {raw!r}" + ) from error + + def describe(self) -> str: + """Human-readable scenario label.""" + return f"{self.ref_backend}({self.ref_device}) vs {self.test_backend}({self.test_device})" + + +class BackendVerifier: + """Run verification scenarios over a set of preprocessed batches. + + Args: + pipelines: Shared pipeline cache (one pipeline per backend/device, reused by + evaluation). + tolerance: Default absolute element-wise tolerance on raw graph outputs; + a scenario's own ``tolerance`` overrides it for that scenario. + """ + + def __init__(self, pipelines: PipelineCache, tolerance: float) -> None: + self.pipelines = pipelines + self.tolerance = float(tolerance) + + def run( + self, + batches: Sequence[MultiTaskBatchInputs], + scenarios: Sequence[VerificationScenario], + available_backends: set[Backend], + ) -> bool: + """Run every applicable scenario; log a per-scenario report. + + Args: + batches: Preprocessed sample batches shared by all scenarios. + scenarios: Configured reference-vs-test comparisons. + available_backends: Backends whose artifacts exist for this run. + Scenarios touching an unavailable backend are skipped with a warning. + + Returns: + True when every executed scenario passed on every batch. + + Raises: + ValueError: If no scenario was executable (misconfiguration). + """ + if not scenarios: + logger.warning("Verification enabled but no scenarios configured; nothing to verify.") + return True + + executed = 0 + all_passed = True + for scenario in scenarios: + required = {Backend.parse(scenario.ref_backend), Backend.parse(scenario.test_backend)} + missing = required - available_backends + if missing: + logger.warning( + "Skipping verification scenario %s: backend(s) %s not exported in this run.", + scenario.describe(), + sorted(b.value for b in missing), + ) + continue + executed += 1 + all_passed &= self._run_scenario(scenario, batches) + + if executed == 0: + raise ValueError( + "No verification scenario was executable — every configured scenario " + f"references unavailable backends (available: {sorted(b.value for b in available_backends)})." + ) + return all_passed + + def _run_scenario( + self, scenario: VerificationScenario, batches: Sequence[MultiTaskBatchInputs] + ) -> bool: + ref_pipeline = self.pipelines.get(scenario.ref_backend, scenario.ref_device) + test_pipeline = self.pipelines.get(scenario.test_backend, scenario.test_device) + comparator = OutputComparator(output_names=test_pipeline.output_names) + + tolerance = scenario.tolerance if scenario.tolerance is not None else self.tolerance + tolerance_source = ( + "scenario override" if scenario.tolerance is not None else "verifier default" + ) + logger.info("=" * 70) + logger.info( + "Verification scenario: %s (tolerance=%s [%s], %d batch(es))", + scenario.describe(), + tolerance, + tolerance_source, + len(batches), + ) + passed_all = True + for index, batch_inputs in enumerate(batches): + ref_result = ref_pipeline.infer(batch_inputs) + test_result = test_pipeline.infer(batch_inputs) + summary, details = comparator.compare( + ref_result.ordered_outputs(), + test_result.ordered_outputs(), + tolerance=tolerance, + ) + status = "PASS" if summary.passed else "FAIL" + logger.info( + " sample %d: %s (max_diff=%.6f, mean_diff=%.6f)", + index, + status, + summary.max_diff, + summary.mean_diff, + ) + for detail in details: + logger.info( + " %-32s shape=%s max=%.6f mean=%.6f", + detail.path, + detail.shape, + detail.max_diff, + detail.mean_diff, + ) + if not summary.passed: + logger.error(" sample %d failed: %s", index, summary.reason) + passed_all = False + + logger.info("Scenario %s: %s", scenario.describe(), "PASSED" if passed_all else "FAILED") + return passed_all diff --git a/autoware_ml/deployment/verification/output_comparator.py b/autoware_ml/deployment/verification/output_comparator.py new file mode 100644 index 00000000..df2f6eb1 --- /dev/null +++ b/autoware_ml/deployment/verification/output_comparator.py @@ -0,0 +1,271 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Pure output comparison for model verification. + +This module contains `OutputComparator`, a stateless recursive comparator +for structured model outputs. Deployment verification expects pipeline raw +outputs that are **sequences (list/tuple) of tensors/arrays** and/or tensor +leaves; dict and bare scalar outputs are not handled (they fail with a type +mismatch). + +Naming: + - **OutputDiffSummary**: one object for the **whole output** — whether it + passed, overall max/mean diff (aggregated), and first failure reason. + - **TensorDiffDetail**: one row per **tensor** in the structure — path, + shape, and that tensor's max/mean diff (for per-head logging). + +Design notes: + - No logging here; callers (e.g. ``BackendVerifier``) render logs. + - :meth:`OutputComparator.compare` returns ``(OutputDiffSummary, list of + TensorDiffDetail)`` in a single traversal. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Iterable, Sequence + +import numpy as np +import torch + +#: Headroom multiplier for the gate suggested on a verification failure. Observed +#: max_diff varies a little run to run (kernel/tactic nondeterminism), so the suggested +#: gate leaves margin above one observation without becoming a rubber stamp. +SUGGESTED_GATE_HEADROOM = 1.25 + + +@dataclass(frozen=True) +class OutputDiffSummary: + """Rolled-up comparison for an entire structured output (or subtree). + + Use this for pass/fail and **global** max/mean diff. For each tensor's + own stats, see :class:`TensorDiffDetail`. + + Attributes: + passed: True if the full structure is within ``tolerance``. + max_diff: Largest per-tensor max diff anywhere in the tree. + mean_diff: Element-weighted mean of absolute differences over the tree. + num_elements: Total tensor elements compared (for weighted mean). + reason: First failing tensor's message, or ``None`` if passed. + """ + + passed: bool + max_diff: float + mean_diff: float + num_elements: int = 0 + reason: str | None = None + + +@dataclass(frozen=True) +class TensorDiffDetail: + """Stats for **one tensor** at a path (used for verbose per-output logs). + + Attributes: + path: Dot/bracket path (e.g. ``output[heatmap]``). + shape: NumPy shape of this tensor. + max_diff: Max absolute difference on this tensor. + mean_diff: Mean absolute difference on this tensor. + """ + + path: str + shape: tuple[int, ...] + max_diff: float + mean_diff: float + + +class OutputComparator: + """Recursively compare structured outputs within an absolute tolerance. + + Optional ``output_names`` label sequence slots (e.g. head names) in paths. + + Args: + output_names: Names aligned with sequence indices; children become + ``output[name]`` instead of ``output_0``, ``output_1``, ... + """ + + def __init__(self, output_names: Sequence[str] | None = None) -> None: + self._output_names: tuple[str, ...] | None = tuple(output_names) if output_names else None + + def compare( + self, + reference: Any, + test: Any, + tolerance: float, + path: str = "output", + ) -> tuple[OutputDiffSummary, list[TensorDiffDetail]]: + """Compare two structured outputs; collect per-tensor rows and a summary.""" + tensor_details: list[TensorDiffDetail] = [] + summary = self._compare_nested(reference, test, tolerance, path, tensor_details) + return summary, tensor_details + + def _compare_nested( + self, + reference: Any, + test: Any, + tolerance: float, + path: str, + tensor_details: list[TensorDiffDetail], + ) -> OutputDiffSummary: + """Recursive compare; appends one :class:`TensorDiffDetail` per tensor leaf.""" + if reference is None and test is None: + return OutputDiffSummary(passed=True, max_diff=0.0, mean_diff=0.0) + + if reference is None or test is None: + return _fail(path, "one side is None while the other is not") + + if isinstance(reference, (list, tuple)) and isinstance(test, (list, tuple)): + return self._compare_sequences(reference, test, tolerance, path, tensor_details) + + if self._is_array_like(reference) and self._is_array_like(test): + return self._compare_arrays(reference, test, tolerance, path, tensor_details) + + return _fail( + path, + f"type mismatch {type(reference).__name__} vs {type(test).__name__}", + ) + + def _compare_sequences( + self, + reference: list | tuple, + test: list | tuple, + tolerance: float, + path: str, + tensor_details: list[TensorDiffDetail], + ) -> OutputDiffSummary: + """Compare list/tuple outputs element-wise using ``output_names`` when provided.""" + if len(reference) != len(test): + return _fail(path, f"length mismatch {len(reference)} vs {len(test)}") + + names = self._output_names + + def _child_summaries(): + for idx, (ref_item, test_item) in enumerate(zip(reference, test)): + name = names[idx] if names and idx < len(names) else f"output_{idx}" + yield self._compare_nested( + ref_item, test_item, tolerance, f"{path}[{name}]", tensor_details + ) + + return self._merge_summaries(_child_summaries()) + + def _compare_arrays( + self, + reference: Any, + test: Any, + tolerance: float, + path: str, + tensor_details: list[TensorDiffDetail], + ) -> OutputDiffSummary: + """Compare tensor/ndarray leaves (same shape required).""" + ref_np = self._to_numpy(reference) + test_np = self._to_numpy(test) + + if ref_np.shape != test_np.shape: + tensor_details.append( + TensorDiffDetail( + path=path, + shape=tuple(int(x) for x in ref_np.shape), + max_diff=float("inf"), + mean_diff=float("inf"), + ) + ) + return _fail(path, f"shape mismatch {ref_np.shape} vs {test_np.shape}") + + diff = np.abs(ref_np.astype(np.float64) - test_np.astype(np.float64)) + max_diff = float(np.max(diff)) if diff.size else 0.0 + mean_diff = float(np.mean(diff)) if diff.size else 0.0 + num_elements = int(diff.size) + + passed = max_diff <= tolerance + reason = ( + None + if passed + else ( + f"{path}: max_diff={max_diff:.6f} > tolerance={tolerance:.6f} (shape={ref_np.shape}). " + "Raw-logit diffs between backends are expected for quantized/FP16 stages; if the " + f"per-backend metrics stay equal, recalibrate this scenario's gate to " + f"~{max_diff * SUGGESTED_GATE_HEADROOM:.1f} " + "and record the observed value in the config comment." + ) + ) + tensor_details.append( + TensorDiffDetail( + path=path, + shape=tuple(int(x) for x in ref_np.shape), + max_diff=max_diff, + mean_diff=mean_diff, + ) + ) + return OutputDiffSummary( + passed=passed, + max_diff=max_diff, + mean_diff=mean_diff, + num_elements=num_elements, + reason=reason, + ) + + @staticmethod + def _merge_summaries(results: Iterable[OutputDiffSummary]) -> OutputDiffSummary: + """Combine child :class:`OutputDiffSummary` values into one rollup.""" + max_diff = 0.0 + total_diff = 0.0 + total_elements = 0 + all_passed = True + first_reason: str | None = None + + for result in results: + max_diff = max(max_diff, result.max_diff) + # Skip zero-element children (shape/type mismatches carry mean_diff=inf, + # num_elements=0): inf * 0 is nan and would poison the whole mean. The + # mismatch still surfaces via all_passed=False and max_diff=inf. + if result.num_elements: + total_diff += result.mean_diff * result.num_elements + total_elements += result.num_elements + if not result.passed and all_passed: + all_passed = False + first_reason = result.reason + + mean_diff = total_diff / total_elements if total_elements > 0 else 0.0 + return OutputDiffSummary( + passed=all_passed, + max_diff=max_diff, + mean_diff=mean_diff, + num_elements=total_elements, + reason=first_reason, + ) + + @staticmethod + def _is_array_like(obj: Any) -> bool: + """Return True when ``obj`` is a tensor or ndarray (leaf comparison path).""" + return isinstance(obj, (torch.Tensor, np.ndarray)) + + @staticmethod + def _to_numpy(tensor: Any) -> np.ndarray: + """Convert tensors to CPU NumPy arrays; pass through ``ndarray``.""" + if isinstance(tensor, torch.Tensor): + return tensor.detach().cpu().numpy() + if isinstance(tensor, np.ndarray): + return tensor + return np.array(tensor) + + +def _fail(path: str, reason: str) -> OutputDiffSummary: + """Build a failing summary with infinite diffs and a short reason.""" + return OutputDiffSummary( + passed=False, + max_diff=float("inf"), + mean_diff=float("inf"), + reason=f"{path}: {reason}", + ) diff --git a/autoware_ml/models/base.py b/autoware_ml/models/base.py index 8ea6f26c..120b51d5 100644 --- a/autoware_ml/models/base.py +++ b/autoware_ml/models/base.py @@ -355,6 +355,7 @@ def predict_step(self, batch_inputs_dict: Mapping[str, Any], batch_idx: int) -> outputs = self(**forward_inputs) return self.predict_outputs(batch_inputs_dict, outputs) + # TODO(vividf): legacy export contract — new models must implement MultiTaskBaseModel.build_stages() (stage-graph export) instead; remove this path once the last BaseModel migrates. def build_export_spec(self, batch_inputs_dict: Mapping[str, Any]) -> ExportSpec: """Build the default deployment export specification for the model. diff --git a/autoware_ml/tests/deployment/test_backend_verifier_scenarios.py b/autoware_ml/tests/deployment/test_backend_verifier_scenarios.py new file mode 100644 index 00000000..3abda31a --- /dev/null +++ b/autoware_ml/tests/deployment/test_backend_verifier_scenarios.py @@ -0,0 +1,95 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for VerificationScenario.from_dict parsing.""" + +from __future__ import annotations + +import dataclasses + +import pytest + +backend_verifier = pytest.importorskip( + "autoware_ml.deployment.verification.backend_verifier", + reason="backend_verifier transitively needs the full model stack", +) +VerificationScenario = backend_verifier.VerificationScenario + + +class TestVerificationScenarioFromDict: + def test_valid_mapping(self): + scenario = VerificationScenario.from_dict( + { + "ref": {"backend": "pytorch", "device": "cpu"}, + "test": {"backend": "onnx", "device": "cuda"}, + } + ) + assert scenario.ref_backend == "pytorch" + assert scenario.ref_device == "cpu" + assert scenario.test_backend == "onnx" + assert scenario.test_device == "cuda" + assert scenario.tolerance is None + + def test_device_defaults_to_cuda(self): + scenario = VerificationScenario.from_dict( + {"ref": {"backend": "pytorch"}, "test": {"backend": "tensorrt"}} + ) + assert scenario.ref_device == "cuda" + assert scenario.test_device == "cuda" + + def test_per_scenario_tolerance_parsed_as_float(self): + scenario = VerificationScenario.from_dict( + { + "ref": {"backend": "pytorch"}, + "test": {"backend": "onnx"}, + "tolerance": "0.001", + } + ) + assert isinstance(scenario.tolerance, float) + assert scenario.tolerance == pytest.approx(0.001) + + def test_missing_tolerance_stays_none(self): + scenario = VerificationScenario.from_dict( + {"ref": {"backend": "pytorch"}, "test": {"backend": "onnx"}} + ) + assert scenario.tolerance is None + + @pytest.mark.parametrize( + "raw", + [ + {}, + {"ref": {"backend": "pytorch"}}, + {"ref": {"device": "cuda"}, "test": {"backend": "onnx"}}, + {"ref": "pytorch", "test": "onnx"}, + ], + ) + def test_malformed_input_raises_value_error(self, raw): + with pytest.raises(ValueError, match="verification scenario"): + VerificationScenario.from_dict(raw) + + def test_scenario_is_frozen(self): + scenario = VerificationScenario.from_dict( + {"ref": {"backend": "pytorch"}, "test": {"backend": "onnx"}} + ) + with pytest.raises(dataclasses.FrozenInstanceError): + scenario.tolerance = 0.5 + + def test_describe_names_backends_and_devices(self): + scenario = VerificationScenario.from_dict( + { + "ref": {"backend": "pytorch", "device": "cpu"}, + "test": {"backend": "onnx", "device": "cuda"}, + } + ) + assert scenario.describe() == "pytorch(cpu) vs onnx(cuda)" diff --git a/autoware_ml/tests/deployment/test_deploy_config.py b/autoware_ml/tests/deployment/test_deploy_config.py new file mode 100644 index 00000000..e73a5b12 --- /dev/null +++ b/autoware_ml/tests/deployment/test_deploy_config.py @@ -0,0 +1,162 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DeployConfig: typo guard, per-stage layout, stage-name check.""" + +from __future__ import annotations + +import pytest + +from autoware_ml.deployment.config import DeployConfig, OnnxPrecision +from autoware_ml.types.backend import Backend + +_RAW = { + "onnx": {"dynamo": False, "opset_version": 17, "precision": "fp16"}, + "tensorrt": {"enabled": False}, + "stages": { + "pts_voxel_encoder": { + "onnx": {"dynamic_axes": {"input_features": {0: "num_voxels"}}}, + "tensorrt": { + "input_shapes": { + "input_features": { + "min_shape": [1, 32, 11], + "opt_shape": [2, 32, 11], + "max_shape": [3, 32, 11], + } + } + }, + } + }, + "verification": { + "enabled": True, + "scenarios": [ + {"ref": {"backend": "pytorch"}, "test": {"backend": "onnx"}, "tolerance": 2.0} + ], + }, + "evaluation": { + "enabled": True, + "num_samples": 10, + "backends": { + "pytorch": {"enabled": True}, + "tensorrt": {"enabled": False, "device": "cuda:1"}, + }, + }, +} + + +class TestDeployConfig: + def test_round_trip(self): + cfg = DeployConfig.from_dict(_RAW) + assert cfg.onnx.dynamo is False and cfg.onnx.opset_version == 17 + assert cfg.onnx.precision is OnnxPrecision.FP16 + assert cfg.tensorrt.enabled is False + stage = cfg.stage("pts_voxel_encoder") + assert stage.onnx.dynamic_axes == {"input_features": {0: "num_voxels"}} + assert stage.tensorrt.input_shapes["input_features"].opt_shape == (2, 32, 11) + assert cfg.verification.scenarios[0].tolerance == 2.0 + assert cfg.evaluation.num_samples == 10 + assert [b for b, _ in cfg.evaluation.enabled_backends()] == [Backend.PYTORCH] + assert cfg.evaluation.backends[Backend.TENSORRT].device == "cuda:1" + + def test_absent_section_is_all_defaults(self): + cfg = DeployConfig.from_dict(None) + assert cfg.onnx.enabled and cfg.tensorrt.enabled + assert not cfg.verification.enabled and not cfg.evaluation.enabled + assert cfg.stage("anything").tensorrt.input_shapes == {} + + @pytest.mark.parametrize( + "raw, where", + [ + ({"onnx": {"opset_versoin": 17}}, "deploy.onnx"), + ({"tensorrt": {"workspace": 1}}, "deploy.tensorrt"), + ({"stages": {"s": {"onnx": {"input_names": ["x"]}}}}, "deploy.stages.s.onnx"), + ({"stages": {"s": {"trt": {}}}}, "deploy.stages.s"), + ( + {"evaluation": {"backends": {"onnx": {"devcie": "cuda"}}}}, + "deploy.evaluation.backends.onnx", + ), + ({"unknown_top": 1}, "deploy"), + ], + ) + def test_unknown_keys_rejected(self, raw, where): + with pytest.raises(ValueError, match=where): + DeployConfig.from_dict(raw) + + def test_incomplete_shape_profile_rejected(self): + raw = {"stages": {"s": {"tensorrt": {"input_shapes": {"x": {"min_shape": [1]}}}}}} + with pytest.raises(ValueError, match="incomplete"): + DeployConfig.from_dict(raw) + + def test_unknown_onnx_precision_rejected(self): + with pytest.raises(ValueError, match="deploy.onnx.precision"): + DeployConfig.from_dict({"onnx": {"precision": "int8"}}) + + def test_removed_precision_policy_rejected_as_unknown_key(self): + with pytest.raises(ValueError, match="precision_policy"): + DeployConfig.from_dict({"tensorrt": {"precision_policy": "fp16"}}) + + def test_unknown_evaluation_backend_rejected(self): + with pytest.raises(ValueError, match="Unknown backend"): + DeployConfig.from_dict({"evaluation": {"backends": {"tflite": {}}}}) + + def test_stage_name_typo_is_caught_against_the_declaration(self): + cfg = DeployConfig.from_dict({"stages": {"pts_voxel_encodr": {}}}) + with pytest.raises(ValueError, match="pts_voxel_encodr"): + cfg.check_stage_names(["pts_voxel_encoder", "pts_backbone_neck_head"]) + cfg = DeployConfig.from_dict(_RAW) + cfg.check_stage_names(["pts_voxel_encoder"]) + + +def test_stage_onnx_precision_overrides_the_global_setting() -> None: + """A stage may pin its own precision; unset stages inherit ``deploy.onnx.precision``.""" + from autoware_ml.deployment.config import DeployConfig, OnnxPrecision + + cfg = DeployConfig.from_dict( + { + "onnx": {"enabled": True, "precision": "fp16"}, + "tensorrt": {"enabled": False}, + "stages": {"fragile_head": {"onnx": {"precision": "fp32"}}}, + } + ) + assert cfg.onnx.precision is OnnxPrecision.FP16 + assert cfg.stage("fragile_head").onnx.precision is OnnxPrecision.FP32 + assert cfg.stage("other").onnx.precision is None + + +def test_stage_onnx_precision_rejects_unknown_values() -> None: + import pytest + + from autoware_ml.deployment.config import DeployConfig + + with pytest.raises(ValueError, match="fragile_head.onnx.precision"): + DeployConfig.from_dict( + { + "onnx": {"enabled": True}, + "tensorrt": {"enabled": False}, + "stages": {"fragile_head": {"onnx": {"precision": "fp42"}}}, + } + ) + + +def test_evaluation_split_parses_and_rejects_unknown_values() -> None: + import pytest + + from autoware_ml.deployment.config import DeployConfig + + base = {"onnx": {"enabled": True}, "tensorrt": {"enabled": False}, "stages": {}} + cfg = DeployConfig.from_dict({**base, "evaluation": {"enabled": True, "split": "val"}}) + assert cfg.evaluation.split == "val" + assert DeployConfig.from_dict(base).evaluation.split == "test" + with pytest.raises(ValueError, match="evaluation.split"): + DeployConfig.from_dict({**base, "evaluation": {"split": "train"}}) diff --git a/autoware_ml/tests/deployment/test_output_comparator.py b/autoware_ml/tests/deployment/test_output_comparator.py new file mode 100644 index 00000000..ca56a6fd --- /dev/null +++ b/autoware_ml/tests/deployment/test_output_comparator.py @@ -0,0 +1,75 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for OutputComparator, focused on the shape-mismatch path.""" + +from __future__ import annotations + +import numpy as np +import torch + +from autoware_ml.deployment.verification.output_comparator import OutputComparator + + +class TestOutputComparator: + def test_identical_arrays_pass(self): + a = [np.zeros((2, 3), dtype=np.float32)] + b = [np.zeros((2, 3), dtype=np.float32)] + summary, details = OutputComparator().compare(a, b, tolerance=1e-6) + assert summary.passed + assert summary.max_diff == 0.0 + assert len(details) == 1 + + def test_within_tolerance_passes(self): + a = [np.zeros((4,), dtype=np.float32)] + b = [np.full((4,), 0.05, dtype=np.float32)] + summary, _ = OutputComparator().compare(a, b, tolerance=0.1) + assert summary.passed + assert summary.max_diff <= 0.1 + + def test_exceeds_tolerance_fails(self): + a = [np.zeros((4,), dtype=np.float32)] + b = [np.full((4,), 1.0, dtype=np.float32)] + summary, _ = OutputComparator().compare(a, b, tolerance=0.1) + assert not summary.passed + assert "tolerance" in (summary.reason or "") + + def test_shape_mismatch_fails_with_reason(self): + a = [np.zeros((2, 3), dtype=np.float32)] + b = [np.zeros((2, 4), dtype=np.float32)] + summary, details = OutputComparator().compare(a, b, tolerance=1.0) + assert not summary.passed + assert "shape mismatch" in (summary.reason or "") + # The mismatched tensor is recorded with infinite diffs, not silently dropped. + assert details and details[0].max_diff == float("inf") + + def test_length_mismatch_fails(self): + a = [np.zeros((2,), dtype=np.float32), np.zeros((2,), dtype=np.float32)] + b = [np.zeros((2,), dtype=np.float32)] + summary, _ = OutputComparator().compare(a, b, tolerance=1.0) + assert not summary.passed + assert "length mismatch" in (summary.reason or "") + + def test_named_outputs_label_paths(self): + a = [np.zeros((2,), dtype=np.float32)] + b = [np.ones((2,), dtype=np.float32)] + summary, details = OutputComparator(output_names=["heatmap"]).compare(a, b, tolerance=0.1) + assert not summary.passed + assert "heatmap" in details[0].path + + def test_torch_and_numpy_mix(self): + a = [torch.zeros(2, 3)] + b = [np.zeros((2, 3), dtype=np.float32)] + summary, _ = OutputComparator().compare(a, b, tolerance=1e-6) + assert summary.passed diff --git a/autoware_ml/tests/deployment/test_pipeline_result.py b/autoware_ml/tests/deployment/test_pipeline_result.py new file mode 100644 index 00000000..b201ad4b --- /dev/null +++ b/autoware_ml/tests/deployment/test_pipeline_result.py @@ -0,0 +1,71 @@ +# Copyright 2026 TIER IV, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for PipelineResult timing rollup and output ordering.""" + +from __future__ import annotations + +import pytest +import torch + +pipeline = pytest.importorskip( + "autoware_ml.deployment.pipeline", + reason="deployment.pipeline transitively needs the project batch dataclasses", +) +PipelineResult = pipeline.PipelineResult + + +class TestPipelineResult: + def test_model_ms_sums_only_graph_stages(self): + result = PipelineResult( + outputs={}, + output_names=[], + stage_times_ms={"pillar_decorate": 1.0, "graph_a": 2.0, "graph_b": 3.5}, + graph_stage_names=("graph_a", "graph_b"), + ) + assert result.model_ms == pytest.approx(5.5) + + def test_model_ms_treats_missing_graph_stage_as_zero(self): + result = PipelineResult( + outputs={}, + output_names=[], + stage_times_ms={"graph_a": 2.0}, + graph_stage_names=("graph_a", "graph_b"), + ) + assert result.model_ms == pytest.approx(2.0) + + def test_model_ms_is_zero_without_graph_stages(self): + result = PipelineResult( + outputs={}, + output_names=[], + stage_times_ms={"anything": 4.0}, + ) + assert result.model_ms == 0.0 + + def test_ordered_outputs_respects_output_names_order(self): + heatmap = torch.zeros(1) + reg = torch.ones(1) + # Insertion order deliberately differs from the frozen ABI order. + result = PipelineResult( + outputs={"reg": reg, "heatmap": heatmap}, + output_names=["heatmap", "reg"], + ) + ordered = result.ordered_outputs() + assert ordered[0] is heatmap + assert ordered[1] is reg + + def test_ordered_outputs_raises_on_missing_name(self): + result = PipelineResult(outputs={"reg": torch.ones(1)}, output_names=["heatmap", "reg"]) + with pytest.raises(KeyError): + result.ordered_outputs() diff --git a/autoware_ml/utils/deploy.py b/autoware_ml/utils/deploy.py index 505b1093..6cc5a2cd 100644 --- a/autoware_ml/utils/deploy.py +++ b/autoware_ml/utils/deploy.py @@ -12,11 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Deployment utility types and helpers. - -This module defines the canonical per-module export contract used by -deployment code. Models expose ``build_export_specs(batch)`` and return a -mapping from module names to :class:`ExportSpec` objects. +"""LEGACY single-module export contract (``ExportSpec``) and its DictConfig adapters. + +TODO(vividf): delete this whole module once every legacy ``BaseModel`` (ptv3 / frnet / +transfusion / bevfusion / calibration_status) migrates to +``MultiTaskBaseModel.build_stages()`` (design doc Q5). Everything current lives in +:mod:`autoware_ml.deployment.onnx` (ONNX primitives) and +:mod:`autoware_ml.deployment.backends.tensorrt_builder` (engine build); the wrappers +here only adapt the legacy ``deploy.onnx.modules`` DictConfig schema onto them. """ from __future__ import annotations @@ -30,16 +33,21 @@ import lightning as L from omegaconf import DictConfig, OmegaConf import torch -from torch.export import Dim -from autoware_ml.ops.segment.scatter_reduce import register_scatter_reduce_onnx_symbolic +from autoware_ml.deployment.backends.tensorrt_builder import build_engine +from autoware_ml.deployment.config import ShapeProfile +from autoware_ml.deployment.onnx.export import export_to_onnx as _export_to_onnx +from autoware_ml.deployment.onnx.modify import ( # noqa: F401 (legacy re-exports) + modify_onnx_graph, + should_modify_graph, +) logger = logging.getLogger(__name__) @dataclass(frozen=True) class ExportSpec: - """Describe the module and tensor inputs used for model export. + """Describe the module and tensor inputs used for legacy single-module export. Attributes: module: Module instance exported to ONNX. @@ -59,30 +67,6 @@ class ExportSpec: supported_stages: frozenset[str] = frozenset({"onnx", "tensorrt"}) -def validate_cuda_available() -> None: - """Ensure CUDA is available for deployment export.""" - if not torch.cuda.is_available(): - raise RuntimeError( - "CUDA is not available. TensorRT requires CUDA. " - "Please run on a machine with CUDA support." - ) - - -def resolve_output_paths( - checkpoint_path: Path, - output_name: str | None, - output_dir: str | None, -) -> tuple[Path, Path, Path]: - """Resolve the output directory and export artifact paths.""" - base_name = output_name if output_name else checkpoint_path.stem - output_directory = Path(output_dir) if output_dir else checkpoint_path.parent - output_directory.mkdir(parents=True, exist_ok=True) - - onnx_path = output_directory / f"{base_name}.onnx" - engine_path = output_directory / f"{base_name}.engine" - return output_directory, onnx_path, engine_path - - def get_forward_signature(model: L.LightningModule) -> inspect.Signature: """Return the cached forward signature from BaseModel, or compute it.""" return getattr(model, "forward_signature", inspect.signature(model.forward)) @@ -111,20 +95,6 @@ def extract_input_from_batch(batch: dict[str, Any], param_name: str) -> Any: return input_value -def get_predict_batch( - datamodule: L.LightningDataModule, - model: L.LightningModule, - device: torch.device, -) -> dict[str, Any]: - """Load one prediction batch and apply transfer-time preprocessing.""" - datamodule.setup("predict") - predict_dataloader = datamodule.predict_dataloader() - batch = next(iter(predict_dataloader)) - batch = batch.to_device(device) - # batch = move_data_to_device(batch, device) - return model.on_after_batch_transfer(batch, dataloader_idx=0) - - def infer_export_spec(model: L.LightningModule, batch: dict[str, Any]) -> ExportSpec: """Infer an export specification directly from the model forward signature.""" forward_params = get_export_parameter_names(model) @@ -142,37 +112,31 @@ def infer_export_spec(model: L.LightningModule, batch: dict[str, Any]) -> Export return ExportSpec(module=model, args=input_args, input_param_names=forward_params) -def resolve_export_specs( +def get_predict_batch( datamodule: L.LightningDataModule, model: L.LightningModule, device: torch.device, -) -> dict[str, ExportSpec]: - """Resolve per-module export specifications for a model. +) -> dict[str, Any]: + """Load one prediction batch and apply transfer-time preprocessing.""" + datamodule.setup("predict") + predict_dataloader = datamodule.predict_dataloader() + batch = next(iter(predict_dataloader)) + batch = batch.to_device(device) + return model.on_after_batch_transfer(batch, dataloader_idx=0) - Args: - datamodule: Data module used to generate one prediction batch. - model: Model instance to export. - device: Device for tensor operations during export preparation. - Returns: - Ordered mapping of module name to export specification. - """ +def resolve_export_specs( + datamodule: L.LightningDataModule, + model: L.LightningModule, + device: torch.device, +) -> dict[str, ExportSpec]: + """Resolve per-module export specifications for a legacy ``BaseModel``.""" batch = get_predict_batch(datamodule, model, device) return model.build_export_specs(batch) def merge_module_onnx_cfg(onnx_cfg: DictConfig, module_name: str) -> DictConfig: - """Merge shared ONNX settings with per-module settings. - - Module-level settings override shared settings. The ``modules`` key itself - is excluded from the merged result. - - Args: - onnx_cfg: Top-level ONNX deploy config containing a ``modules`` mapping. - module_name: Key of the module to resolve within ``modules``. - - Returns: - Merged config with shared settings and module-specific overrides. + """Merge shared ONNX settings with per-module overrides (legacy ``onnx.modules`` schema). Raises: KeyError: If ``module_name`` is not found in ``onnx_cfg.modules``. @@ -186,131 +150,7 @@ def merge_module_onnx_cfg(onnx_cfg: DictConfig, module_name: str) -> DictConfig: k: v for k, v in OmegaConf.to_container(onnx_cfg, resolve=True).items() if k != "modules" } module_overrides = OmegaConf.to_container(onnx_cfg.modules[module_name], resolve=True) - return OmegaConf.create({**shared, **module_overrides}) - - -def log_export_inputs(input_args: tuple[Any, ...], input_names: list[str]) -> None: - """Log export input metadata for debugging.""" - for input_name, input_value in zip(input_names, input_args): - if isinstance(input_value, torch.Tensor): - logger.info( - "Input '%s': shape=%s, dtype=%s", - input_name, - tuple(input_value.shape), - input_value.dtype, - ) - else: - logger.info("Input '%s': type=%s", input_name, type(input_value).__name__) - - -def build_dynamic_shapes( - onnx_cfg: DictConfig, - forward_params: list[str], -) -> tuple[dict[int, Dim] | None, ...] | None: - """Build the ONNX dynamic-shape mapping from config.""" - if "dynamic_shapes" not in onnx_cfg or onnx_cfg.dynamic_shapes is None: - return None - - raw_dynamic_shapes = onnx_cfg.dynamic_shapes - unknown_params = [ - param_name for param_name in raw_dynamic_shapes if param_name not in forward_params - ] - if unknown_params: - raise ValueError( - f"Dynamic shape parameters {unknown_params} not found in export inputs. " - f"Available inputs: {forward_params}." - ) - - dynamic_shapes: list[dict[int, Dim] | None] = [] - for param_name in forward_params: - dim_mapping = raw_dynamic_shapes.get(param_name) - if dim_mapping is None: - dynamic_shapes.append(None) - continue - - param_dynamic_shapes: dict[int, Dim] = {} - for dim_idx, dim_spec in dim_mapping.items(): - if isinstance(dim_spec, str): - param_dynamic_shapes[int(dim_idx)] = Dim(dim_spec) - continue - - dim_name = dim_spec.get("name") - if dim_name is None: - raise ValueError( - f"Dynamic shape spec for '{param_name}[{dim_idx}]' must define 'name'." - ) - dim_kwargs = {key: dim_spec[key] for key in ("min", "max") if key in dim_spec} - param_dynamic_shapes[int(dim_idx)] = Dim(dim_name, **dim_kwargs) - - dynamic_shapes.append(param_dynamic_shapes or None) - - if all(param_dynamic_shapes is None for param_dynamic_shapes in dynamic_shapes): - return None - return tuple(dynamic_shapes) - - -def normalize_dynamic_shapes_for_model( - model: torch.nn.Module, - dynamic_shapes: tuple[dict[int, Dim] | None, ...] | None, -) -> tuple[Any, ...] | None: - """Adapt dynamic-shape structure to the model forward signature. - - ``torch.export`` requires ``dynamic_shapes`` to mirror the positional input - pytree passed to the model. Wrappers that expose ``forward(*args)`` receive - one tuple-valued positional argument, so their dynamic-shape structure must - be wrapped one level deeper. - """ - if dynamic_shapes is None: - return None - - signature = inspect.signature(model.forward) - parameters = [parameter for parameter in signature.parameters.values()] - if len(parameters) == 1 and parameters[0].kind == inspect.Parameter.VAR_POSITIONAL: - return (dynamic_shapes,) - return dynamic_shapes - - -def build_dynamic_axes(onnx_cfg: DictConfig) -> dict[str, dict[int, str]] | None: - """Build legacy ONNX dynamic-axes mapping from config. - - This path is used with ``torch.onnx.export(..., dynamo=False)`` to support - exports that still rely on the legacy exporter behavior. - """ - dynamic_axes_cfg = onnx_cfg.get("dynamic_axes") - if dynamic_axes_cfg is None: - dynamic_axes_cfg = onnx_cfg.get("dynamic_shapes") - if dynamic_axes_cfg is None: - return None - - dynamic_axes: dict[str, dict[int, str]] = {} - for tensor_name, dim_mapping in dynamic_axes_cfg.items(): - tensor_dynamic_axes: dict[int, str] = {} - for dim_idx, dim_spec in dim_mapping.items(): - if isinstance(dim_spec, str): - tensor_dynamic_axes[int(dim_idx)] = dim_spec - continue - - dim_name = dim_spec.get("name") - if dim_name is None: - raise ValueError( - f"Dynamic axis/shape spec for '{tensor_name}[{dim_idx}]' must define 'name'." - ) - tensor_dynamic_axes[int(dim_idx)] = dim_name - - if tensor_dynamic_axes: - dynamic_axes[tensor_name] = tensor_dynamic_axes - - return dynamic_axes or None - - -def merge_onnx_external_data(onnx_path: Path) -> None: - """Merge ONNX external data shards back into a single file.""" - import onnx - from onnx.external_data_helper import convert_model_from_external_data - - onnx_model = onnx.load(str(onnx_path), load_external_data=True) - convert_model_from_external_data(onnx_model) - onnx.save_model(onnx_model, str(onnx_path)) + return OmegaConf.create({**shared, **(module_overrides or {})}) def export_to_onnx( @@ -322,153 +162,21 @@ def export_to_onnx( dynamic_axes_override: dict[str, dict[int, str]] | None, output_path: Path, ) -> None: - """Export a model to ONNX.""" - logger.info("Exporting model to ONNX...") - + """Legacy DictConfig adapter over :func:`autoware_ml.deployment.onnx.export.export_to_onnx`.""" if not input_param_names: raise ValueError("Model forward signature has no parameters.") - - dynamo = onnx_cfg.get("dynamo", True) - dynamic_shapes = build_dynamic_shapes(onnx_cfg, input_param_names) if dynamo else None - dynamic_shapes = normalize_dynamic_shapes_for_model(model, dynamic_shapes) if dynamo else None - dynamic_axes = None - if not dynamo: - dynamic_axes = dynamic_axes_override or build_dynamic_axes(onnx_cfg) - input_names = list(onnx_cfg.get("input_names", input_param_names)) - output_names = list(output_names_override or onnx_cfg.get("output_names", ["output"])) - - logger.info("Dynamic shapes: %s", dynamic_shapes) - logger.info("Dynamic axes: %s", dynamic_axes) - logger.info("ONNX opset version: %s", onnx_cfg.opset_version) - logger.info("Input names: %s", input_names) - logger.info("Output names: %s", output_names) - log_export_inputs(input_sample, input_param_names) - - # Register shared ONNX symbolics needed by export-aware ops packages. - register_scatter_reduce_onnx_symbolic(opset_version=int(onnx_cfg.opset_version)) - - export_kwargs = { - "model": model, - "args": input_sample, - "f": str(output_path), - "input_names": input_names, - "output_names": output_names, - "opset_version": onnx_cfg.opset_version, - "dynamo": dynamo, - "do_constant_folding": onnx_cfg.get("do_constant_folding", True), - } - if dynamo: - export_kwargs["dynamic_shapes"] = dynamic_shapes - else: - export_kwargs["dynamic_axes"] = dynamic_axes - - torch.onnx.export(**export_kwargs) - - logger.info("Successfully exported ONNX model to %s", output_path) - - data_path = output_path.with_suffix(output_path.suffix + ".data") - if data_path.exists(): - logger.info("Found external data file %s. Merging into the ONNX file...", data_path) - merge_onnx_external_data(output_path) - data_path.unlink() - logger.info("Successfully merged external data into the ONNX file") - - -def instantiate_modifier(modify_graph_cfg: DictConfig) -> Any: - """Instantiate an ONNX graph modifier from config.""" - import hydra - - modifier = hydra.utils.instantiate(modify_graph_cfg) - if callable(modifier): - return modifier - if hasattr(modifier, "modify"): - return modifier - raise ValueError(f"Modifier {modifier} must be callable or have a 'modify' method.") - - -def apply_modifier(modifier: Any, onnx_path: Path) -> Path: - """Apply a configured ONNX graph modifier.""" - modified_path = modifier(onnx_path) if callable(modifier) else modifier.modify(onnx_path) - if modified_path is None: - raise ValueError("Modifier returned None. Must return Path or str.") - return Path(modified_path) - - -def should_modify_graph(modify_graph_cfg: DictConfig | None) -> bool: - """Return whether graph modification is enabled.""" - if modify_graph_cfg is None: - return False - if isinstance(modify_graph_cfg, DictConfig): - return OmegaConf.to_container(modify_graph_cfg, resolve=False) is not None - return True - - -def modify_onnx_graph(onnx_path: Path, modify_graph_cfg: DictConfig) -> Path: - """Modify an ONNX graph using the configured modifier.""" - logger.info("Modifying ONNX graph...") - modifier = instantiate_modifier(modify_graph_cfg) - modified_path = apply_modifier(modifier, onnx_path) - logger.info("Successfully modified ONNX graph: %s", modified_path) - return modified_path - - -def create_tensorrt_builder_config(tensorrt_cfg: DictConfig) -> tuple[Any, Any, Any, Any]: - """Create TensorRT builder objects for engine generation.""" - import tensorrt as trt - - trt_logger = trt.Logger(trt.Logger.WARNING) - trt.init_libnvinfer_plugins(trt_logger, "") - builder = trt.Builder(trt_logger) - # Always strongly typed: deploy.onnx.precision decides which dtypes the ONNX carries, and the - # engine has to use them as exported rather than let the builder reassign precisions. - network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) - parser = trt.OnnxParser(network, trt_logger) - config = builder.create_builder_config() - - workspace_size = tensorrt_cfg.get("workspace_size", 1 << 30) - config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_size) - logger.info("Workspace size: %.2f GB", workspace_size / (1024**3)) - return builder, network, parser, config - - -def parse_onnx_file(parser: Any, onnx_path: Path) -> None: - """Parse an ONNX file with a TensorRT parser.""" - with open(onnx_path, "rb") as f: - onnx_data = f.read() - - if not parser.parse(onnx_data): - errors = [parser.get_error(i) for i in range(parser.num_errors)] - error_msg = "\n".join(f"TensorRT parser error {i}: {err}" for i, err in enumerate(errors)) - raise RuntimeError(f"Failed to parse ONNX file:\n{error_msg}") - - logger.info("Successfully parsed ONNX file") - - -def create_optimization_profile(builder: Any, tensorrt_cfg: DictConfig) -> Any | None: - """Create a TensorRT optimization profile from config.""" - if "input_shapes" not in tensorrt_cfg: - return None - - profile = builder.create_optimization_profile() - for input_name, shapes in tensorrt_cfg.input_shapes.items(): - min_shape = shapes.get("min_shape") - opt_shape = shapes.get("opt_shape") - max_shape = shapes.get("max_shape") - if not (min_shape and opt_shape and max_shape): - raise ValueError( - f"TensorRT optimization profile for input '{input_name}' is incomplete. " - "All of min_shape, opt_shape, and max_shape must be specified." - ) - - profile.set_shape(input_name, min=min_shape, opt=opt_shape, max=max_shape) - logger.info( - "Optimization profile for '%s': min=%s, opt=%s, max=%s", - input_name, - min_shape, - opt_shape, - max_shape, - ) - return profile + _export_to_onnx( + model, + tuple(input_sample), + output_path, + input_names=list(onnx_cfg.get("input_names", input_param_names)), + output_names=list(output_names_override or onnx_cfg.get("output_names", ["output"])), + opset_version=int(onnx_cfg.opset_version), + dynamo=bool(onnx_cfg.get("dynamo", True)), + do_constant_folding=bool(onnx_cfg.get("do_constant_folding", True)), + dynamic_shapes=onnx_cfg.get("dynamic_shapes"), + dynamic_axes=dynamic_axes_override or onnx_cfg.get("dynamic_axes"), + ) def build_tensorrt_engine( @@ -476,32 +184,28 @@ def build_tensorrt_engine( deploy_cfg: DictConfig, output_path: Path, ) -> None: - """Build a TensorRT engine from an ONNX model.""" - logger.info("Building TensorRT engine...") + """Legacy DictConfig adapter over :func:`...backends.tensorrt_builder.build_engine`.""" tensorrt_cfg = deploy_cfg.tensorrt - builder, network, parser, config = create_tensorrt_builder_config(tensorrt_cfg) - parse_onnx_file(parser, onnx_path) - - profile = create_optimization_profile(builder, tensorrt_cfg) - if profile is not None: - config.add_optimization_profile(profile) - - logger.info("Building TensorRT engine (this may take a while)...") - serialized_engine = builder.build_serialized_network(network, config) - if serialized_engine is None: - raise RuntimeError("Failed to build TensorRT engine.") - - with open(output_path, "wb") as f: - f.write(serialized_engine) - - logger.info("Successfully built TensorRT engine: %s", output_path) - - -def should_export_stage(stage_cfg: DictConfig | None) -> bool: - """Return whether an export stage is enabled.""" - if stage_cfg is None: - return False - return bool(stage_cfg.get("enabled", True)) + policy = tensorrt_cfg.get("precision_policy") + if policy is not None and str(policy).lower() != "strongly_typed": + logger.warning( + "Ignoring legacy deploy.tensorrt.precision_policy=%r: engines always build " + "strongly typed now (precision lives in the ONNX graph).", + policy, + ) + raw_shapes = tensorrt_cfg.get("input_shapes") or {} + input_shapes = { + str(name): ShapeProfile.from_dict(profile, f"deploy.tensorrt.input_shapes.{name}") + for name, profile in raw_shapes.items() + } + plugin_libraries = tensorrt_cfg.get("plugin_libraries", None) + build_engine( + onnx_path, + output_path, + workspace_size=int(tensorrt_cfg.get("workspace_size", 1 << 30)), + plugin_libraries=list(plugin_libraries) if plugin_libraries is not None else (), + input_shapes=input_shapes, + ) def supports_export_stage(export_spec: ExportSpec, stage_name: str) -> bool: