Refactor environment variable handling to use EnvironmentVariables class

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

Author: xingliu14

examples/offline_inference.py

@@ -1,9 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import os
-
-import vllm.envs as envs
+import tpu_inference.envs as envs
+import vllm.envs as vllm_envs
 from vllm import LLM, EngineArgs
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
@@ -87,10 +86,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.
@@ -104,7 +103,7 @@ def main(args: dict):
 
 if __name__ == "__main__":
     # Skip long warmup for local simple test.
-    os.environ['SKIP_JAX_PRECOMPILE'] = '1'
+    envs.environment_variables['SKIP_JAX_PRECOMPILE'] = lambda: True
 
     parser = create_parser()
     args: dict = vars(parser.parse_args())

examples/offline_safety_model_inference.py

@@ -18,9 +18,8 @@
     --max-num_batched_tokens=4096
 """
 
-import os
-
-import vllm.envs as envs
+import tpu_inference.envs as envs
+import vllm.envs as vllm_envs
 from vllm import LLM, EngineArgs
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
@@ -170,7 +169,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 +178,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
@@ -220,7 +219,7 @@ def main(args: dict):
 
 if __name__ == "__main__":
     # Skip long warmup for local simple test.
-    os.environ['SKIP_JAX_PRECOMPILE'] = '1'
+    envs.environment_variables['SKIP_JAX_PRECOMPILE'] = lambda: True
 
     parser = create_parser()
     args: dict = vars(parser.parse_args())

tests/e2e/test_data_parallel.py

@@ -6,6 +6,7 @@
 from dataclasses import asdict
 
 import pytest
+import tpu_inference.envs as envs
 from vllm import LLM, EngineArgs, SamplingParams
 
 
@@ -173,8 +174,8 @@ def test_data_parallelism_correctness(
     This test compares outputs from a single-device run with data parallel runs
     to ensure correctness, including log probabilities.
     """
-    os.environ['SKIP_JAX_PRECOMPILE'] = '1'
-    os.environ['VLLM_XLA_CHECK_RECOMPILATION'] = '0'
+    envs.environment_variables['SKIP_JAX_PRECOMPILE'] = lambda: True
+    envs.environment_variables['VLLM_XLA_CHECK_RECOMPILATION'] = lambda: False
     model_name = "Qwen/Qwen2.5-1.5B-Instruct"
     # Use a smaller subset of prompts for correctness testing
     small_prompts = test_prompts[:10]

tests/e2e/test_multi_modal_inference.py

@@ -4,9 +4,9 @@
 # This script is a self-contained test that runs a single prompt and
 # compares the output to a known-good output.
 
-import os
 from dataclasses import asdict
 
+import tpu_inference.envs as envs
 from vllm import LLM, EngineArgs, SamplingParams
 from vllm.assets.image import ImageAsset
 from vllm.multimodal.image import convert_image_mode
@@ -24,9 +24,8 @@ def test_multi_modal_inference(monkeypatch):
     """
     Runs multi-modal inference and verifies the output.
     """
-    os.environ['SKIP_JAX_PRECOMPILE'] = '1'  # Skip warmup to save time.
-    os.environ[
-        'VLLM_XLA_CHECK_RECOMPILATION'] = '0'  # Allow compilation during execution.
+    envs.environment_variables['SKIP_JAX_PRECOMPILE'] = lambda: True  # Skip warmup to save time.
+    envs.environment_variables['VLLM_XLA_CHECK_RECOMPILATION'] = lambda: False  # Allow compilation during execution.
 
     monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
 

tests/runner/test_tpu_runner_mesh.py

@@ -4,6 +4,7 @@
 
 import pytest
 
+import tpu_inference.envs as envs
 from tpu_inference.runner.tpu_runner import TPUModelRunner
 
 
@@ -53,7 +54,7 @@ def runner_instance(self, mock_vllm_config, mock_devices):
     def test_init_mesh_2d_model_without_device_order(self, runner_instance,
                                                      mock_vllm_config):
         """Test 2d mesh creation without enforced device order."""
-        with patch.dict(os.environ, {'NEW_MODEL_DESIGN': ''}), \
+        with patch.dict(envs.environment_variables, {'NEW_MODEL_DESIGN': lambda: False}), \
              patch('tpu_inference.runner.tpu_runner.make_optimized_mesh') as mock_make_mesh, \
              patch('tpu_inference.runner.tpu_runner.logger'):
 
@@ -79,7 +80,7 @@ def test_init_mesh_2d_model_with_device_order(self, runner_instance,
         """Test 2d mesh creation with enforced device order."""
         mock_vllm_config.sharding_config.device_indexes = [0, 1, 2, 3]
 
-        with patch.dict(os.environ, {'NEW_MODEL_DESIGN': ''}), \
+        with patch.dict(envs.environment_variables, {'NEW_MODEL_DESIGN': lambda: False}), \
              patch('jax.make_mesh') as mock_jax_mesh, \
              patch('tpu_inference.runner.tpu_runner.logger'):
 
@@ -103,7 +104,7 @@ def test_init_mesh_2d_model_with_device_order(self, runner_instance,
     def test_init_mesh_new_model_single_slice(self, runner_instance,
                                               mock_vllm_config):
         """Test new model mesh creation with single slice."""
-        with patch.dict(os.environ, {'NEW_MODEL_DESIGN': '1', 'NUM_SLICES': '1'}), \
+        with patch.dict(envs.environment_variables, {'NEW_MODEL_DESIGN': lambda: True, 'NUM_SLICES': lambda: 1}), \
              patch('tpu_inference.runner.tpu_runner.mesh_utils') as mock_mesh_utils, \
              patch('jax.sharding.Mesh') as mock_jax_mesh, \
              patch('tpu_inference.runner.tpu_runner.logger'):
@@ -134,7 +135,7 @@ def test_init_mesh_new_model_multi_slice(self, runner_instance,
                                              mock_vllm_config):
         """Test new model mesh creation with multiple slices."""
         num_slices = 2
-        with patch.dict(os.environ, {'NEW_MODEL_DESIGN': '1', 'NUM_SLICES': str(num_slices)}), \
+        with patch.dict(envs.environment_variables, {'NEW_MODEL_DESIGN': lambda: True, 'NUM_SLICES': lambda: num_slices}), \
              patch('tpu_inference.runner.tpu_runner.mesh_utils') as mock_mesh_utils, \
              patch('jax.sharding.Mesh') as mock_jax_mesh, \
              patch('tpu_inference.runner.tpu_runner.logger'):

tests/test_envs.py

@@ -63,6 +63,13 @@ def test_boolean_env_vars(monkeypatch: pytest.MonkeyPatch):
     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")
@@ -81,6 +88,13 @@ def test_integer_env_vars(monkeypatch: pytest.MonkeyPatch):
     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 +148,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

@@ -6,6 +6,7 @@
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.outputs import DraftTokenIds
 
+import tpu_inference.envs as envs
 # The class we are testing
 from tpu_inference.worker.tpu_worker import TPUWorker
 
@@ -280,7 +281,7 @@ def test_add_lora_not_implemented_lora_request(self, mock_vllm_config):
     #
 
     @patch('tpu_inference.worker.tpu_worker.jax')
-    @patch.dict('os.environ', {"PYTHON_TRACER_LEVEL": "1"}, clear=True)
+    @patch.dict(envs.environment_variables, {"PYTHON_TRACER_LEVEL": lambda: 1})
     def test_profile_start(self, mock_jax, mock_vllm_config):
         """Tests starting the JAX profiler."""
         worker = TPUWorker(vllm_config=mock_vllm_config,
@@ -296,7 +297,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"
 
@@ -48,6 +50,9 @@
     # Skip JAX precompilation step during initialization
     "SKIP_JAX_PRECOMPILE":
     lambda: bool(int(os.getenv("SKIP_JAX_PRECOMPILE", "0"))),
+    # Check for XLA recompilation during execution
+    "VLLM_XLA_CHECK_RECOMPILATION":
+    lambda: bool(int(os.getenv("VLLM_XLA_CHECK_RECOMPILATION", "0"))),
     # Model implementation type (e.g., "flax_nnx")
     "MODEL_IMPL_TYPE":
     lambda: os.getenv("MODEL_IMPL_TYPE", "flax_nnx").lower(),
@@ -63,6 +68,9 @@
     # Use custom expert-parallel kernel for MoE (Mixture of Experts)
     "USE_MOE_EP_KERNEL":
     lambda: bool(int(os.getenv("USE_MOE_EP_KERNEL", "0"))),
+    # Number of TPU slices for multi-slice mesh
+    "NUM_SLICES":
+    lambda: int(os.getenv("NUM_SLICES", "1")),
     # Enable/disable Ray usage statistics collection
     "RAY_USAGE_STATS_ENABLED":
     lambda: os.getenv("RAY_USAGE_STATS_ENABLED", "0"),

tpu_inference/executors/ray_distributed_executor.py

@@ -18,6 +18,7 @@
 from vllm.v1.executor.ray_executor import RayWorkerMetaData
 from vllm.v1.executor.ray_utils import RayWorkerWrapper, _wait_until_pg_ready
 
+import tpu_inference.envs as tpu_envs
 from tpu_inference.logger import init_logger
 
 try:
@@ -72,7 +73,8 @@ class RayDistributedExecutor(RayDistributedExecutorV1):
     def _init_executor(self) -> None:
         self.forward_dag: Optional[ray.dag.CompiledDAG] = None
 
-        os.environ["VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE"] = "shm"
+        # Ensure Ray compiled DAG channel type is set for vLLM
+        os.environ["VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE"] = tpu_envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE
 
         # Currently, this requires USE_RAY_SPMD_WORKER=True.
         self.use_ray_compiled_dag = True
@@ -86,10 +88,10 @@ def _init_executor(self) -> None:
         self._initialize_ray_cluster()
         placement_group = self.parallel_config.placement_group
 
-        # Disable Ray usage stats collection.
-        ray_usage = os.environ.get("RAY_USAGE_STATS_ENABLED", "0")
-        if ray_usage != "1":
-            os.environ["RAY_USAGE_STATS_ENABLED"] = "0"
+        # Ensure Ray usage stats collection setting is propagated to Ray workers.
+        # Ray workers inherit environment variables, so we explicitly set this
+        # based on our configuration (defaults to "0" to disable stats).
+        os.environ["RAY_USAGE_STATS_ENABLED"] = tpu_envs.RAY_USAGE_STATS_ENABLED
 
         # Create the parallel GPU workers.
         self._init_workers_ray(placement_group)

tpu_inference/layers/common/sharding.py

@@ -8,7 +8,7 @@
 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 +48,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:
@@ -166,7 +166,7 @@ def validate(cls, vllm_config, sharding_strategy):
                     f"LoRA is not supported with data parallelism "
                     f"(DP size: {total_dp_size}). Please disable LoRA or "
                     f"set data parallelism to 1.")
-            if not os.environ.get("NEW_MODEL_DESIGN", False):
+            if not envs.NEW_MODEL_DESIGN:
                 raise ValueError(
                     "Must run DP with NEW_MODEL_DESIGN enabled. Please set the "
                     "NEW_MODEL_DESIGN=True.")