[None][fix] Isolate modeling_v2 tests and scope shared test state

jenkins/scripts/cbts/cbts/rules/modeling_v2.py

@@ -121,7 +121,9 @@
 # A family's `routing.py` imports each target it routes to by this line, from
 # the target directory directly under the family. The claims test reads the
 # same line with the same pattern.
-_TARGET_IMPORT = re.compile(r"^\s*from \.(\w+)\.modeling import (\w+)$", re.MULTILINE)
+_TARGET_IMPORT = re.compile(
+    r"^\s*from \.(\w+)\.modeling import (?:(\w+)|\(\s*(\w+),?\s*\))$", re.MULTILINE
+)
 
 
 def _is_mv2_claim(path: str) -> bool:
@@ -230,7 +232,8 @@ def _gates_under(self, node: Path) -> Optional[set[str]]:
         if not routing.is_file():
             return None
         routed, gates = set(), set()
-        for target, target_cls in _TARGET_IMPORT.findall(routing.read_text(encoding="utf-8")):
+        for target, plain, wrapped in _TARGET_IMPORT.findall(routing.read_text(encoding="utf-8")):
+            target_cls = plain or wrapped
             leaf = Path(family, target)
             routed.add(leaf)
             if leaf == node or node in leaf.parents:

tensorrt_llm/_torch/_experimental/modeling_v2/README.md

@@ -83,10 +83,6 @@ LLM(model=..., modeling_v2=...)  ->  ModelLoader
                                      returns that class; the registry is not consulted
 ```
 
-`route` imports the target class on the branch that returns it, so a routing
-module that declines never loads modeling code, and the class a reader sees
-returned is the class that runs -- no name to look up elsewhere.
-
 ## What may decide a target
 
 A quantity may be a routing criterion only if it is known when

tests/integration/defs/accuracy/modeling_v2/_env.py

@@ -0,0 +1,29 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Engine-scoped environment overrides for accuracy tests."""
+
+import os
+from collections.abc import Iterator, Mapping
+from contextlib import contextmanager
+
+
+@contextmanager
+def scoped_env_overrides(overrides: Mapping[str, str]) -> Iterator[dict[str, str]]:
+    """Apply rank overrides locally and restore their prior values on exit.
+
+    Pass the yielded mapping to ``LLM(env_overrides=...)`` inside this scope,
+    so rank workers receive the same environment and engine shutdown completes
+    before the parent environment is restored. Only the overridden keys are
+    restored; unrelated environment changes are preserved.
+    """
+    values = dict(overrides)
+    previous = {key: os.environ.get(key) for key in values}
+    try:
+        os.environ.update(values)
+        yield values
+    finally:
+        for key, value in previous.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = value

tests/integration/defs/accuracy/modeling_v2/test_deepseek_v3.py

@@ -45,6 +45,7 @@
     assert_acceptance_length,
     compute_acceptance_length,
 )
+from ._env import scoped_env_overrides
 
 # The targets assert their own SM at construction: certification is per GPU
 # architecture, and a receipt from another one says nothing here.
@@ -152,12 +153,15 @@ def test_gsm8k(self, mocker, parallel):
         # about metadata that engine prepared, so the only place it can be
         # checked honestly is a run that prepared it; a synthesized metadata
         # object would be checking the synthesizer.
-        with LLM(
-            self.MODEL_PATH,
-            modeling_v2="require",
-            env_overrides={MODELING_V2_VALIDATE_ENV: "1"},
-            **parallel,
-        ) as llm:
+        with (
+            scoped_env_overrides({MODELING_V2_VALIDATE_ENV: "1"}) as rank_env,
+            LLM(
+                self.MODEL_PATH,
+                modeling_v2="require",
+                env_overrides=rank_env,
+                **parallel,
+            ) as llm,
+        ):
             task = GSM8K(self.MODEL_NAME)
             task.evaluate(llm)
 
@@ -186,16 +190,22 @@ def test_gsm8k_identity_vs_mtp3(self, mocker, parallel):
         # engines were built under the same mode.
         modeling_v2 = dict(modeling_v2="require", env_overrides={MODELING_V2_VALIDATE_ENV: "1"})
 
-        with LLM(self.MODEL_PATH, **modeling_v2, **parallel) as llm:
+        with (
+            scoped_env_overrides(modeling_v2["env_overrides"]),
+            LLM(self.MODEL_PATH, **modeling_v2, **parallel) as llm,
+        ):
             identity = task.evaluate(llm)
 
-        with LLM(
-            self.MODEL_PATH,
-            speculative_config=self.MTP3,
-            kv_cache_config=self.MTP3_KV,
-            **modeling_v2,
-            **parallel,
-        ) as llm:
+        with (
+            scoped_env_overrides(modeling_v2["env_overrides"]),
+            LLM(
+                self.MODEL_PATH,
+                speculative_config=self.MTP3,
+                kv_cache_config=self.MTP3_KV,
+                **modeling_v2,
+                **parallel,
+            ) as llm,
+        ):
             mtp3 = task.evaluate(llm)
 
         delta = mtp3 - identity
@@ -242,16 +252,19 @@ def test_mtp3_acceptance(self, mode, parallel, mocker):
 
         mocker.patch.dict(GSM8K.EVALUATE_KWARGS, _SCORES_FILTER)
 
-        with LLM(
-            self.MODEL_PATH,
-            speculative_config=self.MTP3,
-            kv_cache_config=self.MTP3_KV,
-            cuda_graph_config=CudaGraphConfig(),
-            enable_iter_perf_stats=True,
-            modeling_v2=mode,
-            env_overrides=rank_env,
-            **parallel,
-        ) as llm:
+        with (
+            scoped_env_overrides(rank_env),
+            LLM(
+                self.MODEL_PATH,
+                speculative_config=self.MTP3,
+                kv_cache_config=self.MTP3_KV,
+                cuda_graph_config=CudaGraphConfig(),
+                enable_iter_perf_stats=True,
+                modeling_v2=mode,
+                env_overrides=rank_env,
+                **parallel,
+            ) as llm,
+        ):
             task = GSM8K(self.MODEL_NAME)
             task.evaluate(llm)
             acceptance_length = compute_acceptance_length(llm)
@@ -276,13 +289,16 @@ def test_gsm8k_chunked_prefill(self, mocker, parallel):
         """
         mocker.patch.dict(GSM8K.EVALUATE_KWARGS, _SCORES_FILTER)
 
-        with LLM(
-            self.MODEL_PATH,
-            enable_chunked_prefill=True,
-            max_num_tokens=256,
-            modeling_v2="require",
-            env_overrides={MODELING_V2_VALIDATE_ENV: "1"},
-            **parallel,
-        ) as llm:
+        with (
+            scoped_env_overrides({MODELING_V2_VALIDATE_ENV: "1"}) as rank_env,
+            LLM(
+                self.MODEL_PATH,
+                enable_chunked_prefill=True,
+                max_num_tokens=256,
+                modeling_v2="require",
+                env_overrides=rank_env,
+                **parallel,
+            ) as llm,
+        ):
             task = GSM8K(self.MODEL_NAME)
             task.evaluate(llm)

tests/integration/defs/accuracy/modeling_v2/test_glm_moe_dsa.py

@@ -61,6 +61,7 @@
     assert_acceptance_length,
     compute_acceptance_length,
 )
+from ._env import scoped_env_overrides
 
 _TP_SIZE = 16
 _NUM_SAMPLES = 1319
@@ -197,10 +198,7 @@ def test_mtp7_gsm8k_and_acceptance(self, monkeypatch: pytest.MonkeyPatch, tmp_pa
         assert os.environ.get("INTEGRATION_TEST") != "1", "This guard requires all 1319 samples"
         monkeypatch.setenv("TLLM_EVAL_MAX_IN_FLIGHT", "1")
         monkeypatch.setenv("TLLM_EVAL_PARTIAL_SCORES_EVERY", "100")
-        for key, value in _RANK_ENV.items():
-            monkeypatch.setenv(key, value)
-
-        with LLM(_MODEL_PATH, **_engine_kwargs()) as llm:
+        with scoped_env_overrides(_RANK_ENV), LLM(_MODEL_PATH, **_engine_kwargs()) as llm:
             _assert_optimized_paths(llm, tmp_path / "optimized-path-traces")
             _request_stats(llm, 1)  # Exclude the coverage probe from corpus AL.
             evaluator = GSM8KInferenceX(
@@ -359,24 +357,30 @@ def test_gsm8k_identity_vs_mtp1_and_acceptance(self):
         # would be checking the synthesizer.
         modeling_v2 = dict(modeling_v2="require", env_overrides={MODELING_V2_VALIDATE_ENV: "1"})
 
-        with LLM(
-            self.MODEL_PATH,
-            kv_cache_config=self.IDENTITY_KV,
-            **modeling_v2,
-            **self.TEP4,
-            **self.BOUNDS,
-        ) as llm:
+        with (
+            scoped_env_overrides(modeling_v2["env_overrides"]),
+            LLM(
+                self.MODEL_PATH,
+                kv_cache_config=self.IDENTITY_KV,
+                **modeling_v2,
+                **self.TEP4,
+                **self.BOUNDS,
+            ) as llm,
+        ):
             identity = task.evaluate(llm)
 
-        with LLM(
-            self.MODEL_PATH,
-            speculative_config=self.MTP1,
-            kv_cache_config=self.MTP1_KV,
-            enable_iter_perf_stats=True,
-            **modeling_v2,
-            **self.TEP4,
-            **self.BOUNDS,
-        ) as llm:
+        with (
+            scoped_env_overrides(modeling_v2["env_overrides"]),
+            LLM(
+                self.MODEL_PATH,
+                speculative_config=self.MTP1,
+                kv_cache_config=self.MTP1_KV,
+                enable_iter_perf_stats=True,
+                **modeling_v2,
+                **self.TEP4,
+                **self.BOUNDS,
+            ) as llm,
+        ):
             mtp1 = task.evaluate(llm)
             acceptance_length = compute_acceptance_length(llm)
 

tests/integration/defs/accuracy/modeling_v2/test_gpt_oss.py

@@ -35,6 +35,7 @@
 
 from ...conftest import llm_models_root
 from ..accuracy_core import GSM8K, LlmapiAccuracyTestHarness
+from ._env import scoped_env_overrides
 
 # The targets assert their own SM at construction: certification is per GPU
 # architecture, and a receipt from another one says nothing here.
@@ -114,14 +115,17 @@ def test_gsm8k(self, mocker, world_size, moe_tp_size, moe_ep_size, attention_dp)
         # checking the synthesizer. One forward of the run below pays for it.
         # It travels as ``env_overrides``: the check runs where the metadata
         # does, which is the rank rather than this process.
-        with LLM(
-            self.MODEL_PATH,
-            tensor_parallel_size=world_size,
-            moe_tensor_parallel_size=moe_tp_size,
-            moe_expert_parallel_size=moe_ep_size,
-            enable_attention_dp=attention_dp,
-            modeling_v2="require",
-            env_overrides={MODELING_V2_VALIDATE_ENV: "1"},
-        ) as llm:
+        with (
+            scoped_env_overrides({MODELING_V2_VALIDATE_ENV: "1"}) as rank_env,
+            LLM(
+                self.MODEL_PATH,
+                tensor_parallel_size=world_size,
+                moe_tensor_parallel_size=moe_tp_size,
+                moe_expert_parallel_size=moe_ep_size,
+                enable_attention_dp=attention_dp,
+                modeling_v2="require",
+                env_overrides=rank_env,
+            ) as llm,
+        ):
             task = GSM8K(self.MODEL_NAME)
             task.evaluate(llm, extra_evaluator_kwargs=self.extra_evaluator_kwargs)

tests/integration/defs/accuracy/modeling_v2/test_kimi_k3_vl.py

@@ -27,6 +27,7 @@
 
 from ...conftest import check_device_contain, llm_models_root
 from ..accuracy_core import GSM8K, LlmapiAccuracyTestHarness
+from ._env import scoped_env_overrides
 
 _MODEL_NAME = "moonshotai/Kimi-K3"
 _MODEL_PATH = f"{llm_models_root()}/Kimi-K3"
@@ -38,32 +39,35 @@
 
 
 def _run_gsm8k(model_path: str, moe_tp_size: int, moe_ep_size: int) -> None:
-    with LLM(
-        model_path,
-        dtype="bfloat16",
-        tensor_parallel_size=16,
-        pipeline_parallel_size=1,
-        moe_tensor_parallel_size=moe_tp_size,
-        moe_expert_parallel_size=moe_ep_size,
-        enable_attention_dp=False,
-        modeling_v2="require",
-        trust_remote_code=True,
-        max_batch_size=1,
-        max_num_tokens=GSM8K.MAX_INPUT_LEN,
-        max_seq_len=8192,
-        enable_chunked_prefill=True,
-        kv_cache_config=KvCacheConfig(
-            dtype="auto",
-            tokens_per_block=64,
-            enable_block_reuse=False,
-            use_kv_cache_manager_v2=True,
-            mamba_ssm_cache_dtype="float32",
-            max_tokens=8192,
-        ),
-        cuda_graph_config=CudaGraphConfig(batch_sizes=[1], enable_padding=True),
-        moe_config=MoeConfig(backend="TRTLLM", use_low_precision_moe_combine=True),
-        env_overrides={MODELING_V2_VALIDATE_ENV: "1"},
-    ) as llm:
+    with (
+        scoped_env_overrides({MODELING_V2_VALIDATE_ENV: "1"}) as rank_env,
+        LLM(
+            model_path,
+            dtype="bfloat16",
+            tensor_parallel_size=16,
+            pipeline_parallel_size=1,
+            moe_tensor_parallel_size=moe_tp_size,
+            moe_expert_parallel_size=moe_ep_size,
+            enable_attention_dp=False,
+            modeling_v2="require",
+            trust_remote_code=True,
+            max_batch_size=1,
+            max_num_tokens=GSM8K.MAX_INPUT_LEN,
+            max_seq_len=8192,
+            enable_chunked_prefill=True,
+            kv_cache_config=KvCacheConfig(
+                dtype="auto",
+                tokens_per_block=64,
+                enable_block_reuse=False,
+                use_kv_cache_manager_v2=True,
+                mamba_ssm_cache_dtype="float32",
+                max_tokens=8192,
+            ),
+            cuda_graph_config=CudaGraphConfig(batch_sizes=[1], enable_padding=True),
+            moe_config=MoeConfig(backend="TRTLLM", use_low_precision_moe_combine=True),
+            env_overrides=rank_env,
+        ) as llm,
+    ):
         GSM8K(_MODEL_NAME).evaluate(llm)
 
 
@@ -76,8 +80,7 @@ class TestModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4(LlmapiAccuracyTestHarness):
     MODEL_NAME = _MODEL_NAME
     MODEL_PATH = _MODEL_PATH
 
-    def test_gsm8k(self, monkeypatch: pytest.MonkeyPatch) -> None:
-        monkeypatch.setenv(MODELING_V2_VALIDATE_ENV, "1")
+    def test_gsm8k(self) -> None:
         _run_gsm8k(self.MODEL_PATH, moe_tp_size=4, moe_ep_size=4)
 
 
@@ -90,6 +93,5 @@ class TestModelingV2KimiK3Mxfp4Sm100Tp16Moetp16ep1(LlmapiAccuracyTestHarness):
     MODEL_NAME = _MODEL_NAME
     MODEL_PATH = _MODEL_PATH
 
-    def test_gsm8k(self, monkeypatch: pytest.MonkeyPatch) -> None:
-        monkeypatch.setenv(MODELING_V2_VALIDATE_ENV, "1")
+    def test_gsm8k(self) -> None:
         _run_gsm8k(self.MODEL_PATH, moe_tp_size=16, moe_ep_size=1)