Add test_envs.py (vllm-project#1079)

Signed-off-by: Xing Liu <[email protected]>

Author: xingliu14

tests/test_envs.py

@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the tpu-inference project
+
+import pytest
+
+import tpu_inference.envs as envs
+from tpu_inference.envs import enable_envs_cache, environment_variables
+
+
+def test_getattr_without_cache(monkeypatch: pytest.MonkeyPatch):
+    assert envs.JAX_PLATFORMS == ""
+    assert envs.PHASED_PROFILING_DIR == ""
+    monkeypatch.setenv("JAX_PLATFORMS", "tpu")
+    monkeypatch.setenv("PHASED_PROFILING_DIR", "/tmp/profiling")
+    assert envs.JAX_PLATFORMS == "tpu"
+    assert envs.PHASED_PROFILING_DIR == "/tmp/profiling"
+
+    assert envs.TPU_NAME is None
+    assert envs.TPU_ACCELERATOR_TYPE is None
+    monkeypatch.setenv("TPU_NAME", "my-tpu")
+    monkeypatch.setenv("TPU_ACCELERATOR_TYPE", "v5litepod-16")
+    assert envs.TPU_NAME == "my-tpu"
+    assert envs.TPU_ACCELERATOR_TYPE == "v5litepod-16"
+
+    # __getattr__ is not decorated with functools.cache
+    assert not hasattr(envs.__getattr__, "cache_info")
+
+
+def test_getattr_with_cache(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("JAX_PLATFORMS", "tpu")
+    monkeypatch.setenv("TPU_NAME", "my-tpu")
+
+    # __getattr__ is not decorated with functools.cache
+    assert not hasattr(envs.__getattr__, "cache_info")
+
+    enable_envs_cache()
+
+    # __getattr__ is decorated with functools.cache
+    assert hasattr(envs.__getattr__, "cache_info")
+    start_hits = envs.__getattr__.cache_info().hits
+
+    # 2 more hits due to JAX_PLATFORMS and TPU_NAME accesses
+    assert envs.JAX_PLATFORMS == "tpu"
+    assert envs.TPU_NAME == "my-tpu"
+    assert envs.__getattr__.cache_info().hits == start_hits + 2
+
+    # All environment variables are cached
+    for environment_variable in environment_variables:
+        envs.__getattr__(environment_variable)
+    assert envs.__getattr__.cache_info(
+    ).hits == start_hits + 2 + len(environment_variables)
+
+    # Reset envs.__getattr__ back to non-cached version to
+    # avoid affecting other tests
+    envs.__getattr__ = envs.__getattr__.__wrapped__
+
+
+def test_boolean_env_vars(monkeypatch: pytest.MonkeyPatch):
+    # 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 NEW_MODEL_DESIGN (default False)
+    assert envs.NEW_MODEL_DESIGN is False
+    monkeypatch.setenv("NEW_MODEL_DESIGN", "1")
+    assert envs.NEW_MODEL_DESIGN is True
+
+    # Test USE_MOE_EP_KERNEL (default False)
+    assert envs.USE_MOE_EP_KERNEL is False
+    monkeypatch.setenv("USE_MOE_EP_KERNEL", "1")
+    assert envs.USE_MOE_EP_KERNEL is True
+
+
+def test_integer_env_vars(monkeypatch: pytest.MonkeyPatch):
+    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
+
+
+def test_lowercase_conversion(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("TPU_MULTIHOST_BACKEND", "GRPC")
+    assert envs.TPU_MULTIHOST_BACKEND == "grpc"
+
+    monkeypatch.setenv("MODEL_IMPL_TYPE", "FLAX_NNX")
+    assert envs.MODEL_IMPL_TYPE == "flax_nnx"
+
+
+def test_string_env_vars_defaults(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.delenv("JAX_PLATFORMS", raising=False)
+    monkeypatch.delenv("PREFILL_SLICES", raising=False)
+    monkeypatch.delenv("DECODE_SLICES", raising=False)
+
+    assert envs.JAX_PLATFORMS == ""
+    assert envs.PREFILL_SLICES == ""
+    assert envs.DECODE_SLICES == ""
+    assert envs.PHASED_PROFILING_DIR == ""
+
+
+def test_none_default_env_vars(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.delenv("TPU_ACCELERATOR_TYPE", raising=False)
+    monkeypatch.delenv("TPU_NAME", raising=False)
+    monkeypatch.delenv("TPU_WORKER_ID", raising=False)
+
+    assert envs.TPU_ACCELERATOR_TYPE is None
+    assert envs.TPU_NAME is None
+    assert envs.TPU_WORKER_ID is None
+
+
+def test_ray_env_vars(monkeypatch: pytest.MonkeyPatch):
+    assert envs.RAY_USAGE_STATS_ENABLED == "0"
+    monkeypatch.setenv("RAY_USAGE_STATS_ENABLED", "1")
+    assert envs.RAY_USAGE_STATS_ENABLED == "1"
+
+    assert envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE == "shm"
+    monkeypatch.setenv("VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE", "nccl")
+    assert envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE == "nccl"
+
+
+def test_invalid_attribute_raises_error():
+    with pytest.raises(AttributeError,
+                       match="has no attribute 'NONEXISTENT_VAR'"):
+        _ = envs.NONEXISTENT_VAR
+
+
+def test_dir_returns_all_env_vars():
+    env_vars = envs.__dir__()
+    assert isinstance(env_vars, list)
+    assert len(env_vars) == len(environment_variables)
+    assert "JAX_PLATFORMS" in env_vars
+    assert "TPU_NAME" in env_vars
+    assert "SKIP_JAX_PRECOMPILE" in env_vars
+    assert "MODEL_IMPL_TYPE" in env_vars
+
+
+def test_tpu_multihost_env_vars(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("TPU_WORKER_ID", "0")
+    assert envs.TPU_WORKER_ID == "0"
+
+    monkeypatch.setenv("TPU_MULTIHOST_BACKEND", "grpc")
+    assert envs.TPU_MULTIHOST_BACKEND == "grpc"
+
+    monkeypatch.setenv("TPU_MULTIHOST_BACKEND", "xla")
+    assert envs.TPU_MULTIHOST_BACKEND == "xla"
+
+
+def test_disaggregated_serving_env_vars(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("PREFILL_SLICES", "0,1,2,3")
+    assert envs.PREFILL_SLICES == "0,1,2,3"
+
+    monkeypatch.setenv("DECODE_SLICES", "4,5,6,7")
+    assert envs.DECODE_SLICES == "4,5,6,7"
+
+
+def test_model_impl_type_default(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.delenv("MODEL_IMPL_TYPE", raising=False)
+    assert envs.MODEL_IMPL_TYPE == "flax_nnx"
+
+
+def test_cache_preserves_values_across_env_changes(
+        monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("JAX_PLATFORMS", "tpu")
+
+    enable_envs_cache()
+
+    assert envs.JAX_PLATFORMS == "tpu"
+
+    # Change environment variable
+    monkeypatch.setenv("JAX_PLATFORMS", "cpu")
+
+    # Cached value should still be "tpu"
+    assert envs.JAX_PLATFORMS == "tpu"
+
+    # Reset envs.__getattr__ back to non-cached version
+    envs.__getattr__ = envs.__getattr__.__wrapped__
+
+    # Now it should reflect the new value
+    assert envs.JAX_PLATFORMS == "cpu"