File size: 6,511 Bytes
e479c46 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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-stage invariant for the TRT export-mode fan-out (CPU-only).
``build_trt_pipeline._MODE_MAP`` routes one ``export_mode`` selector to a
``(build, verify, benchmark)`` tuple. Each stage produces/consumes a *set* of
engines, but nothing else binds those sets — a valid-but-wrong row (e.g. a
build mode that requires ONNX the export never wrote, or a verify mode that
loads engines the build never produced) is only discovered when the row is run.
These tests bind the rows statically:
* every component a build produces is sourced from the shared
``FULL_PIPELINE_COMPONENTS`` / ``EXPORT_MODE_COMPONENTS`` tables, and
* the verify and benchmark consumers' required engine sets are the documented
contract from ``trt_model_forward.setup_tensorrt_engines`` (the ``_setup_*``
loaders), which both stages share.
A future row that pairs a producer with a consumer (verify or benchmark)
needing more than it produces fails here instead of at deploy time.
"""
from __future__ import annotations
from gr00t.deployment.modes import (
EXPORT_MODE_COMPONENTS,
FULL_PIPELINE_COMPONENTS,
PIPELINE_STAGE_MODES as _MODE_MAP,
ExportMode,
)
import pytest
_ENGINE_OF = {c.name: c.engine for c in FULL_PIPELINE_COMPONENTS}
_FULL = EXPORT_MODE_COMPONENTS[ExportMode.full_pipeline]
# Engines each *setup mode* REQUIRES (non-fallback). Verify and benchmark use
# the same mode strings, both resolved by the trt_model_forward.py setup loaders
# (``_setup_dit_only`` / ``_setup_action_head`` / ``_setup_n17_full_pipeline``),
# so one table backs both stages. Engines that fall back to PyTorch when absent
# (LLM / VL-SA in n17_full_pipeline) are NOT required. Keep in sync if a
# ``_setup_*`` loader changes which engines it hard-requires.
SETUP_REQUIRED_ENGINES: dict[str, frozenset[str]] = {
"dit_only": frozenset({"dit_bf16.engine"}),
"action_head": frozenset(
{
"state_encoder.engine",
"action_encoder.engine",
"dit_bf16.engine",
"action_decoder.engine",
}
),
"n17_full_pipeline": frozenset(
{
"vit.engine",
"state_encoder.engine",
"action_encoder.engine",
"dit_bf16.engine",
"action_decoder.engine",
}
),
}
def _engines_built_for(export_mode: str, build_mode: str) -> frozenset[str]:
"""Engines the build stage produces for a ``_MODE_MAP`` row.
Mirrors ``build_trt_pipeline._run_build``: the ``single`` builder emits just
the DiT engine; the ``full_pipeline`` builder emits one engine per component
that ``export_mode`` wrote ONNX for (``EXPORT_MODE_COMPONENTS``).
"""
if build_mode == "single":
return frozenset({_ENGINE_OF["DiT"]})
produced = EXPORT_MODE_COMPONENTS[ExportMode(export_mode)]
return frozenset(_ENGINE_OF[name] for name in produced)
def _required_for(mode: str) -> frozenset[str]:
assert mode in SETUP_REQUIRED_ENGINES, (
f"setup mode {mode!r} is routed by _MODE_MAP but has no documented "
"required-engine set in this test."
)
return SETUP_REQUIRED_ENGINES[mode]
@pytest.mark.parametrize("export_mode", sorted(_MODE_MAP))
def test_build_produces_every_engine_verify_requires(export_mode):
build_mode, verify_mode, _bench = _MODE_MAP[export_mode]
built = _engines_built_for(export_mode, build_mode)
missing = _required_for(verify_mode) - built
assert not missing, (
f"export_mode={export_mode!r}: build ({build_mode}) produces {sorted(built)}, "
f"but verify ({verify_mode}) requires {sorted(missing)} it never builds."
)
@pytest.mark.parametrize("export_mode", sorted(_MODE_MAP))
def test_build_produces_every_engine_benchmark_requires(export_mode):
# Same fan-out hazard on the third map: benchmark loads engines via the same
# setup_tensorrt_engines loaders, so a row whose benchmark mode needs an
# engine the build never produced must also fail here, not at deploy time.
build_mode, _verify, bench_mode = _MODE_MAP[export_mode]
built = _engines_built_for(export_mode, build_mode)
missing = _required_for(bench_mode) - built
assert not missing, (
f"export_mode={export_mode!r}: build ({build_mode}) produces {sorted(built)}, "
f"but benchmark ({bench_mode}) requires {sorted(missing)} it never builds."
)
@pytest.mark.parametrize("export_mode", sorted(_MODE_MAP))
def test_export_mode_component_names_are_valid(export_mode):
valid = {c.name for c in FULL_PIPELINE_COMPONENTS}
unknown = set(EXPORT_MODE_COMPONENTS[ExportMode(export_mode)]) - valid
assert not unknown, f"{export_mode}: unknown component names {sorted(unknown)}"
def test_action_head_build_excludes_pytorch_only_components():
# The regression guard for B-2026-06-29-003: action_head keeps ViT/LLM/VL-SA
# in PyTorch, so the build must NOT require their ONNX.
produced = EXPORT_MODE_COMPONENTS[ExportMode.action_head]
for name in ("ViT", "LLM", "VL Self-Attention"):
assert name not in produced
assert produced == frozenset({"State Encoder", "Action Encoder", "DiT", "Action Decoder"})
def test_setup_modes_have_known_required_sets():
# Every verify AND benchmark mode referenced by _MODE_MAP must have a
# documented required set above — otherwise a new row escapes the invariant.
for _build, verify_mode, bench_mode in _MODE_MAP.values():
assert verify_mode in SETUP_REQUIRED_ENGINES, (
f"verify mode {verify_mode!r} is routed by _MODE_MAP but has no "
"documented required-engine set in this test."
)
assert bench_mode in SETUP_REQUIRED_ENGINES, (
f"benchmark mode {bench_mode!r} is routed by _MODE_MAP but has no "
"documented required-engine set in this test."
)
|