From 0d1c0fa0320daa14e797677657baf3fda41870ed Mon Sep 17 00:00:00 2001 From: qgai Date: Wed, 22 Jul 2026 19:56:58 -0700 Subject: [PATCH] [None][test] Enable session prefetch for all test stages Signed-off-by: qgai --- tests/test_common/session_prefetcher.py | 38 +++++-------------- .../llmapi/test_session_prefetcher.py | 38 +++++++++---------- 2 files changed, 28 insertions(+), 48 deletions(-) diff --git a/tests/test_common/session_prefetcher.py b/tests/test_common/session_prefetcher.py index 10f23d73cad9..2325df25b26b 100644 --- a/tests/test_common/session_prefetcher.py +++ b/tests/test_common/session_prefetcher.py @@ -40,12 +40,10 @@ — the two layers compose: reuse covers the steady state, prefetch covers the misses. -Rollout: currently in a CANARY phase (reviewer request) — with no explicit -setting the plugin is active only on the stage groups in -``_CANARY_STAGE_PREFIXES`` (CI exports the stage name as ``stageName``). -``TRTLLM_TEST_PREFETCH_SESSION=1``/``0`` overrides in either direction and -remains the permanent kill switch. Suites that never import tensorrt_llm's -executor modules pay nothing — not even the tensorrt_llm import. +Enabled by default; ``TRTLLM_TEST_PREFETCH_SESSION=0`` disables BOTH pool +prefetch and weight warming (one kill switch for the whole plugin). Suites +that never import tensorrt_llm's executor modules pay nothing — not even +the tensorrt_llm import. """ import glob @@ -72,17 +70,6 @@ "tensorrt_llm.llmapi.llm", ) -# Canary rollout stage GROUPS (one multi-GPU LLM-dense, one single-GPU -# LLM-dense). Prefixes, not exact names: the numbered shard suffix (-1/-2/…) -# is assigned by dynamic load balancing and cannot be targeted stably. -# Graduation plan: after a week of clean canary runs (zero failures -# attributed to the plugin, healthy session-summary counters) this gate is -# removed and the plugin returns to enabled-by-default. -_CANARY_STAGE_PREFIXES = ( - "DGX_H100-4_GPUs-PyTorch-DeepSeek", - "A10-PyTorch", -) - def _reuse_layer_active() -> bool: """True when the MPI session-reuse layer owns the pool-creation seams. @@ -280,17 +267,12 @@ def enabled(self) -> bool: # pool plus a spare. Disable in xdist workers. if os.environ.get("PYTEST_XDIST_WORKER"): return False - explicit = os.environ.get("TRTLLM_TEST_PREFETCH_SESSION") - if explicit is not None: - return explicit.lower() in ("1", "true", "yes", "on") - # Canary rollout (reviewer request): with no explicit setting, enable - # only on the canary stage GROUPS. CI exports the stage name as - # ``stageName``; prefix matching covers dynamically numbered shards - # (per-shard targeting is impossible — the split is rebalanced as the - # test list changes). Everywhere else (including runs without a - # stageName) the plugin stays off until the canary graduates. - stage = os.environ.get("stageName", "") - return any(stage.startswith(p) for p in _CANARY_STAGE_PREFIXES) + return os.environ.get("TRTLLM_TEST_PREFETCH_SESSION", "1").lower() in ( + "1", + "true", + "yes", + "on", + ) @staticmethod def _next_model_map(items): diff --git a/tests/unittest/llmapi/test_session_prefetcher.py b/tests/unittest/llmapi/test_session_prefetcher.py index 2df7dfe5d50c..f88735d3a020 100644 --- a/tests/unittest/llmapi/test_session_prefetcher.py +++ b/tests/unittest/llmapi/test_session_prefetcher.py @@ -83,34 +83,32 @@ def _arm(prefetcher, pool, spec=4): prefetcher._publish(spec, pool, session_prefetcher._spawn_snapshot(), prefetcher._build_gen) -def test_canary_enabled_on_canary_stage_groups(monkeypatch): - # Canary phase: with no explicit setting, active only on the canary - # stage groups; prefix match covers dynamically numbered shards. +@pytest.mark.parametrize( + "stage_name", + [ + None, + "A10-PyTorch-2", + "DGX_H100-4_GPUs-PyTorch-DeepSeek-1", + "DGX_B200-PyTorch-4", + "Any-Future-Stage", + ], +) +def test_enabled_by_default_in_all_stages(monkeypatch, stage_name): monkeypatch.delenv("TRTLLM_TEST_PREFETCH_SESSION", raising=False) monkeypatch.delenv("PYTEST_XDIST_WORKER", raising=False) - monkeypatch.setenv("stageName", "A10-PyTorch-2") - assert SessionPrefetcher().enabled - monkeypatch.setenv("stageName", "DGX_H100-4_GPUs-PyTorch-DeepSeek-1") + if stage_name is None: + monkeypatch.delenv("stageName", raising=False) + else: + monkeypatch.setenv("stageName", stage_name) assert SessionPrefetcher().enabled -def test_canary_disabled_elsewhere(monkeypatch): - monkeypatch.delenv("TRTLLM_TEST_PREFETCH_SESSION", raising=False) +def test_explicit_env_overrides_default(monkeypatch): monkeypatch.delenv("PYTEST_XDIST_WORKER", raising=False) - monkeypatch.setenv("stageName", "DGX_B200-PyTorch-4") - assert not SessionPrefetcher().enabled - monkeypatch.delenv("stageName", raising=False) # local run, no stage + monkeypatch.setenv("TRTLLM_TEST_PREFETCH_SESSION", "0") assert not SessionPrefetcher().enabled - - -def test_explicit_env_overrides_canary_gate(monkeypatch): - monkeypatch.delenv("PYTEST_XDIST_WORKER", raising=False) - monkeypatch.setenv("stageName", "DGX_B200-PyTorch-4") # not a canary stage monkeypatch.setenv("TRTLLM_TEST_PREFETCH_SESSION", "1") - assert SessionPrefetcher().enabled # manual opt-in anywhere - monkeypatch.setenv("stageName", "A10-PyTorch-1") # canary stage - monkeypatch.setenv("TRTLLM_TEST_PREFETCH_SESSION", "0") - assert not SessionPrefetcher().enabled # kill switch beats the canary + assert SessionPrefetcher().enabled def test_disabled_is_noop(monkeypatch):