Revert test file changes

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

Author: xingliu14

tests/e2e/test_data_parallel.py

@@ -6,14 +6,13 @@
 from dataclasses import asdict
 
 import pytest
-import tpu_inference.envs as envs
 from vllm import LLM, EngineArgs, SamplingParams
 
 
 @pytest.fixture(autouse=True)
 def setup_new_model_design():
     """Automatically set NEW_MODEL_DESIGN=True for all tests."""
-    os.environ['NEW_MODEL_DESIGN'] = '1'
+    os.environ['NEW_MODEL_DESIGN'] = 'True'
 
 
 @pytest.fixture
@@ -174,8 +173,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.
     """
-    envs.environment_variables['SKIP_JAX_PRECOMPILE'] = lambda: True
-    envs.environment_variables['VLLM_XLA_CHECK_RECOMPILATION'] = lambda: False
+    os.environ['SKIP_JAX_PRECOMPILE'] = '1'
+    os.environ['VLLM_XLA_CHECK_RECOMPILATION'] = '0'
     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,8 +24,9 @@ def test_multi_modal_inference(monkeypatch):
     """
     Runs multi-modal inference and verifies the output.
     """
-    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.
+    os.environ['SKIP_JAX_PRECOMPILE'] = '1'  # Skip warmup to save time.
+    os.environ[
+        'VLLM_XLA_CHECK_RECOMPILATION'] = '0'  # Allow compilation during execution.
 
     monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")