Centralizes environment variable access by routing variables reads through the envs.py module (#1147)

Signed-off-by: Xing Liu <xingliu14@gmail.com>

Author: xingliu14

examples/offline_inference.py

@@ -3,7 +3,7 @@
 
 import os
 
-import vllm.envs as envs
+import vllm.envs as vllm_envs
 from vllm import LLM, EngineArgs
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
@@ -87,10 +87,10 @@ def main(args: dict):
         'Who wrote the novel "Pride and Prejudice"?',
     ]
 
-    if envs.VLLM_TORCH_PROFILER_DIR is not None:
+    if vllm_envs.VLLM_TORCH_PROFILER_DIR is not None:
         llm.start_profile()
     outputs = llm.generate(prompts, sampling_params)
-    if envs.VLLM_TORCH_PROFILER_DIR is not None:
+    if vllm_envs.VLLM_TORCH_PROFILER_DIR is not None:
         llm.stop_profile()
 
     # Print the outputs.

examples/offline_lora_inference.py

@@ -4,7 +4,7 @@
 import os
 import time
 
-import vllm.envs as envs
+import vllm.envs as vllm_envs
 from vllm import LLM, EngineArgs
 from vllm.lora.request import LoRARequest
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -55,13 +55,13 @@ def main(args: dict):
         "lora_adapter_3", 3,
         "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_3_adapter")
 
-    if envs.VLLM_TORCH_PROFILER_DIR is not None:
+    if vllm_envs.VLLM_TORCH_PROFILER_DIR is not None:
         llm.start_profile()
     start = time.perf_counter()
     outputs = llm.generate(prompt,
                            sampling_params=sampling_params,
                            lora_request=lora_request)
-    if envs.VLLM_TORCH_PROFILER_DIR is not None:
+    if vllm_envs.VLLM_TORCH_PROFILER_DIR is not None:
         llm.stop_profile()
 
     # Print the outputs.

examples/offline_safety_model_inference.py

@@ -20,7 +20,7 @@
 
 import os
 
-import vllm.envs as envs
+import vllm.envs as vllm_envs
 from vllm import LLM, EngineArgs
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
@@ -170,7 +170,7 @@ def main(args: dict):
 
         prompts.append(TokensPrompt(prompt_token_ids=tokenized_prompt))
 
-    if envs.VLLM_TORCH_PROFILER_DIR is not None:
+    if vllm_envs.VLLM_TORCH_PROFILER_DIR is not None:
         llm.start_profile()
 
     outputs = llm.generate(
@@ -179,7 +179,7 @@ def main(args: dict):
         use_tqdm=True,
     )
 
-    if envs.VLLM_TORCH_PROFILER_DIR is not None:
+    if vllm_envs.VLLM_TORCH_PROFILER_DIR is not None:
         llm.stop_profile()
 
     passed_tests = 0

tests/test_envs.py

@@ -56,13 +56,26 @@ def test_getattr_with_cache(monkeypatch: pytest.MonkeyPatch):
 
 
 def test_boolean_env_vars(monkeypatch: pytest.MonkeyPatch):
+    # Ensure clean environment for boolean vars by setting to default "0"
+    monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "0")
+    monkeypatch.setenv("VLLM_XLA_CHECK_RECOMPILATION", "0")
+    monkeypatch.setenv("NEW_MODEL_DESIGN", "0")
+    monkeypatch.setenv("USE_MOE_EP_KERNEL", "0")
+
     # Test SKIP_JAX_PRECOMPILE (default False)
     assert envs.SKIP_JAX_PRECOMPILE is False
     monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "1")
     assert envs.SKIP_JAX_PRECOMPILE is True
     monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "0")
     assert envs.SKIP_JAX_PRECOMPILE is False
 
+    # Test VLLM_XLA_CHECK_RECOMPILATION (default False)
+    assert envs.VLLM_XLA_CHECK_RECOMPILATION is False
+    monkeypatch.setenv("VLLM_XLA_CHECK_RECOMPILATION", "1")
+    assert envs.VLLM_XLA_CHECK_RECOMPILATION is True
+    monkeypatch.setenv("VLLM_XLA_CHECK_RECOMPILATION", "0")
+    assert envs.VLLM_XLA_CHECK_RECOMPILATION is False
+
     # Test NEW_MODEL_DESIGN (default False)
     assert envs.NEW_MODEL_DESIGN is False
     monkeypatch.setenv("NEW_MODEL_DESIGN", "1")
@@ -75,12 +88,23 @@ def test_boolean_env_vars(monkeypatch: pytest.MonkeyPatch):
 
 
 def test_integer_env_vars(monkeypatch: pytest.MonkeyPatch):
+    # Ensure clean environment for integer vars by setting to defaults
+    monkeypatch.setenv("PYTHON_TRACER_LEVEL", "1")
+    monkeypatch.setenv("NUM_SLICES", "1")
+
     assert envs.PYTHON_TRACER_LEVEL == 1
     monkeypatch.setenv("PYTHON_TRACER_LEVEL", "3")
     assert envs.PYTHON_TRACER_LEVEL == 3
     monkeypatch.setenv("PYTHON_TRACER_LEVEL", "0")
     assert envs.PYTHON_TRACER_LEVEL == 0
 
+    # Test NUM_SLICES (default 1)
+    assert envs.NUM_SLICES == 1
+    monkeypatch.setenv("NUM_SLICES", "2")
+    assert envs.NUM_SLICES == 2
+    monkeypatch.setenv("NUM_SLICES", "4")
+    assert envs.NUM_SLICES == 4
+
 
 def test_lowercase_conversion(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setenv("TPU_MULTIHOST_BACKEND", "GRPC")
@@ -134,6 +158,7 @@ def test_dir_returns_all_env_vars():
     assert "JAX_PLATFORMS" in env_vars
     assert "TPU_NAME" in env_vars
     assert "SKIP_JAX_PRECOMPILE" in env_vars
+    assert "VLLM_XLA_CHECK_RECOMPILATION" in env_vars
     assert "MODEL_IMPL_TYPE" in env_vars
 
 

tests/worker/tpu_worker_test.py

@@ -294,7 +294,7 @@ def test_profile_start(self, mock_jax, mock_vllm_config):
         args, kwargs = mock_jax.profiler.start_trace.call_args
         assert args[0] == "/tmp/profile_dir"
         # Verify options from env var were used
-        assert kwargs['profiler_options'].python_tracer_level == '1'
+        assert kwargs['profiler_options'].python_tracer_level == 1
 
     @patch('tpu_inference.worker.tpu_worker.jax')
     def test_profile_stop(self, mock_jax, mock_vllm_config):

tpu_inference/envs.py

@@ -15,11 +15,13 @@
     PREFILL_SLICES: str = ""
     DECODE_SLICES: str = ""
     SKIP_JAX_PRECOMPILE: bool = False
+    VLLM_XLA_CHECK_RECOMPILATION: bool = False
     MODEL_IMPL_TYPE: str = "flax_nnx"
     NEW_MODEL_DESIGN: bool = False
     PHASED_PROFILING_DIR: str = ""
     PYTHON_TRACER_LEVEL: int = 1
     USE_MOE_EP_KERNEL: bool = False
+    NUM_SLICES: int = 1
     RAY_USAGE_STATS_ENABLED: str = "0"
     VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: str = "shm"
 
@@ -47,22 +49,28 @@
     lambda: os.getenv("DECODE_SLICES", ""),
     # Skip JAX precompilation step during initialization
     "SKIP_JAX_PRECOMPILE":
-    lambda: bool(int(os.getenv("SKIP_JAX_PRECOMPILE", "0"))),
+    lambda: bool(int(os.getenv("SKIP_JAX_PRECOMPILE") or "0")),
+    # Check for XLA recompilation during execution
+    "VLLM_XLA_CHECK_RECOMPILATION":
+    lambda: bool(int(os.getenv("VLLM_XLA_CHECK_RECOMPILATION") or "0")),
     # Model implementation type (e.g., "flax_nnx")
     "MODEL_IMPL_TYPE":
     lambda: os.getenv("MODEL_IMPL_TYPE", "flax_nnx").lower(),
     # Enable new experimental model design
     "NEW_MODEL_DESIGN":
-    lambda: bool(int(os.getenv("NEW_MODEL_DESIGN", "0"))),
+    lambda: bool(int(os.getenv("NEW_MODEL_DESIGN") or "0")),
     # Directory to store phased profiling output
     "PHASED_PROFILING_DIR":
     lambda: os.getenv("PHASED_PROFILING_DIR", ""),
     # Python tracer level for profiling
     "PYTHON_TRACER_LEVEL":
-    lambda: int(os.getenv("PYTHON_TRACER_LEVEL", "1")),
+    lambda: int(os.getenv("PYTHON_TRACER_LEVEL") or "1"),
     # Use custom expert-parallel kernel for MoE (Mixture of Experts)
     "USE_MOE_EP_KERNEL":
-    lambda: bool(int(os.getenv("USE_MOE_EP_KERNEL", "0"))),
+    lambda: bool(int(os.getenv("USE_MOE_EP_KERNEL") or "0")),
+    # Number of TPU slices for multi-slice mesh
+    "NUM_SLICES":
+    lambda: int(os.getenv("NUM_SLICES") or "1"),
     # Enable/disable Ray usage statistics collection
     "RAY_USAGE_STATS_ENABLED":
     lambda: os.getenv("RAY_USAGE_STATS_ENABLED", "0"),

tpu_inference/layers/common/sharding.py

@@ -1,14 +1,13 @@
 import json
 import math
-import os
 from dataclasses import asdict, dataclass
 from typing import TYPE_CHECKING, List, Optional
 
 import jax.numpy as jnp
 import numpy as np
 from jax.sharding import Mesh
 
-from tpu_inference import utils
+from tpu_inference import envs, utils
 
 if TYPE_CHECKING:
     from vllm.v1.configs.vllm_config import VllmConfig
@@ -48,7 +47,7 @@ class ShardingAxisName2D:
 
 
 try:
-    _use_base_sharding = os.getenv("NEW_MODEL_DESIGN", False)
+    _use_base_sharding = envs.NEW_MODEL_DESIGN
     if _use_base_sharding:
         ShardingAxisName = ShardingAxisNameBase
     else:
@@ -167,7 +166,7 @@ def validate(cls, vllm_config, sharding_strategy):
                     f"(DP size: {total_dp_size}). Please disable LoRA or "
                     f"set data parallelism to 1.")
         if sharding_strategy.attention_data_parallelism > 1:
-            if not os.environ.get("NEW_MODEL_DESIGN", False):
+            if not envs.NEW_MODEL_DESIGN:
                 raise ValueError(
                     "Must run Attention DP with NEW_MODEL_DESIGN enabled. Please set the "
                     "NEW_MODEL_DESIGN=True.")

tpu_inference/runner/compilation_manager.py

@@ -1,13 +1,13 @@
-import os
 import time
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
 
 import jax
 import jax.numpy as jnp
 import numpy as np
-import vllm.envs as envs
+import vllm.envs as vllm_envs
 from jax.sharding import NamedSharding, PartitionSpec
 
+import tpu_inference.envs as envs
 from tpu_inference.core.disagg_utils import is_disagg_enabled
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
 from tpu_inference.layers.common.sharding import ShardingAxisName
@@ -30,10 +30,10 @@ class CompilationManager:
 
     def __init__(self, runner: "TPUModelRunner"):
         self.runner = runner
-        if not envs.VLLM_DISABLE_COMPILE_CACHE:
+        if not vllm_envs.VLLM_DISABLE_COMPILE_CACHE:
             logger.info("Enabling JAX compile cache.")
             jax.config.update("jax_compilation_cache_dir",
-                              envs.VLLM_XLA_CACHE_PATH)
+                              vllm_envs.VLLM_XLA_CACHE_PATH)
 
     def _create_dummy_tensor(self,
                              shape: Tuple[int, ...],
@@ -67,8 +67,7 @@ def _run_compilation(self, name: str, fn: Callable, *args,
         logger.info("Compilation finished in %.2f [secs].", end - start)
 
     def capture_model(self) -> None:
-        if os.getenv("SKIP_JAX_PRECOMPILE",
-                     False) or self.runner.model_config.enforce_eager:
+        if envs.SKIP_JAX_PRECOMPILE or self.runner.model_config.enforce_eager:
             return
         logger.info("Precompile all the subgraphs with possible input shapes.")
 

tpu_inference/runner/tpu_runner.py

@@ -1,6 +1,5 @@
 import copy
 import functools
-import os
 import random
 from contextlib import nullcontext
 from dataclasses import dataclass
@@ -11,7 +10,7 @@
 import jaxtyping
 import numpy as np
 import torch
-import vllm.envs as envs
+import vllm.envs as vllm_envs
 from flax import nnx
 from jax.experimental import mesh_utils
 from jax.sharding import NamedSharding, PartitionSpec
@@ -35,6 +34,7 @@
     KVConnectorModelRunnerMixin
 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
 
+import tpu_inference.envs as envs
 from tpu_inference import utils as common_utils
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
 from tpu_inference.layers.common.sharding import (MESH_AXIS_NAMES,
@@ -291,7 +291,7 @@ def _init_random(self):
         self.rng_key = jax.random.key(self.model_config.seed)
 
     def _init_mesh(self) -> None:
-        if os.getenv("NEW_MODEL_DESIGN", False):
+        if envs.NEW_MODEL_DESIGN:
             self.mesh = self._create_new_model_mesh()
         else:
             # NOTE(wenxindongwork): The new MoE kernel expects a 2D mesh, so we need
@@ -302,7 +302,7 @@ def _init_mesh(self) -> None:
         logger.info(f"Init mesh | mesh={self.mesh}")
 
     def _create_new_model_mesh(self) -> jax.sharding.Mesh:
-        num_slices = int(os.environ.get('NUM_SLICES', 1))
+        num_slices = envs.NUM_SLICES
 
         logger.info(f"Creating new model mesh | devices={len(self.devices)}, "
                     f"num_slices={num_slices}")
@@ -371,7 +371,7 @@ def _create_2d_mesh(self) -> jax.sharding.Mesh:
                                        devices=self.devices)
 
     def _init_phased_profiling(self) -> None:
-        self.phased_profiling_dir = os.getenv("PHASED_PROFILING_DIR", "")
+        self.phased_profiling_dir = envs.PHASED_PROFILING_DIR
         self.phase_based_profiler = None
         if self.phased_profiling_dir:
             self.phase_based_profiler = runner_utils.PhasedBasedProfiler(
@@ -413,7 +413,7 @@ def _init_inputs(self) -> None:
             min_token_size=max(16, self.dp_size),
             max_token_size=scheduler_config.max_num_batched_tokens *
             self.dp_size,
-            padding_gap=envs.VLLM_TPU_BUCKET_PADDING_GAP)
+            padding_gap=vllm_envs.VLLM_TPU_BUCKET_PADDING_GAP)
         self.num_tokens_paddings_per_dp = [
             padding // self.dp_size for padding in self.num_tokens_paddings
         ]

tpu_inference/runner/utils.py

@@ -15,6 +15,7 @@
 from jax._src.interpreters import pxla
 from vllm.v1.core.sched.output import SchedulerOutput as VllmSchedulerOutput
 
+from tpu_inference import envs
 from tpu_inference.logger import init_logger
 from tpu_inference.runner.input_batch import InputBatch
 
@@ -306,8 +307,7 @@ def __init__(self, profile_dir: str):
             InferencePhase.BALANCED: False
         }
         self.default_profiling_options = jax.profiler.ProfileOptions()
-        self.default_profiling_options.python_tracer_level = os.getenv(
-            "PYTHON_TRACER_LEVEL", 0)
+        self.default_profiling_options.python_tracer_level = envs.PYTHON_TRACER_LEVEL
 
         self.current_phase: str = ""