tpu-inference 0.12.0.dev20251222__py3-none-any.whl

tests/test_envs.py

@@ -0,0 +1,280 @@
+# 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):
+    # 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("ENABLE_QUANTIZED_MATMUL_KERNEL", "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")
+    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
+
+    # Test ENABLE_QUANTIZED_MATMUL_KERNEL (default False)
+    assert envs.ENABLE_QUANTIZED_MATMUL_KERNEL is False
+    monkeypatch.setenv("ENABLE_QUANTIZED_MATMUL_KERNEL", "1")
+    assert envs.ENABLE_QUANTIZED_MATMUL_KERNEL is True
+
+
+def test_boolean_env_vars_string_values(monkeypatch: pytest.MonkeyPatch):
+    """Test that boolean env vars accept string values like 'True' and 'False'"""
+
+    # Test NEW_MODEL_DESIGN with string "True"
+    monkeypatch.setenv("NEW_MODEL_DESIGN", "True")
+    assert envs.NEW_MODEL_DESIGN is True
+
+    monkeypatch.setenv("NEW_MODEL_DESIGN", "true")
+    assert envs.NEW_MODEL_DESIGN is True
+
+    monkeypatch.setenv("NEW_MODEL_DESIGN", "False")
+    assert envs.NEW_MODEL_DESIGN is False
+
+    monkeypatch.setenv("NEW_MODEL_DESIGN", "false")
+    assert envs.NEW_MODEL_DESIGN is False
+
+    # Test SKIP_JAX_PRECOMPILE with string values
+    monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "True")
+    assert envs.SKIP_JAX_PRECOMPILE is True
+
+    monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "false")
+    assert envs.SKIP_JAX_PRECOMPILE is False
+
+    # Test VLLM_XLA_CHECK_RECOMPILATION with string values
+    monkeypatch.setenv("VLLM_XLA_CHECK_RECOMPILATION", "TRUE")
+    assert envs.VLLM_XLA_CHECK_RECOMPILATION is True
+
+    monkeypatch.setenv("VLLM_XLA_CHECK_RECOMPILATION", "FALSE")
+    assert envs.VLLM_XLA_CHECK_RECOMPILATION is False
+
+    # Test USE_MOE_EP_KERNEL with string values
+    monkeypatch.setenv("USE_MOE_EP_KERNEL", "true")
+    assert envs.USE_MOE_EP_KERNEL is True
+
+    monkeypatch.setenv("USE_MOE_EP_KERNEL", "False")
+    assert envs.USE_MOE_EP_KERNEL is False
+
+
+def test_boolean_env_vars_invalid_values(monkeypatch: pytest.MonkeyPatch):
+    """Test that boolean env vars raise errors for invalid values"""
+
+    # Test invalid value for NEW_MODEL_DESIGN
+    monkeypatch.setenv("NEW_MODEL_DESIGN", "yes")
+    with pytest.raises(
+            ValueError,
+            match="Invalid boolean value 'yes' for NEW_MODEL_DESIGN"):
+        _ = envs.NEW_MODEL_DESIGN
+
+    monkeypatch.setenv("NEW_MODEL_DESIGN", "2")
+    with pytest.raises(ValueError,
+                       match="Invalid boolean value '2' for NEW_MODEL_DESIGN"):
+        _ = envs.NEW_MODEL_DESIGN
+
+    # Test invalid value for SKIP_JAX_PRECOMPILE
+    monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "invalid")
+    with pytest.raises(
+            ValueError,
+            match="Invalid boolean value 'invalid' for SKIP_JAX_PRECOMPILE"):
+        _ = envs.SKIP_JAX_PRECOMPILE
+
+
+def test_boolean_env_vars_empty_string(monkeypatch: pytest.MonkeyPatch):
+    """Test that empty string returns default value"""
+
+    monkeypatch.setenv("NEW_MODEL_DESIGN", "")
+    assert envs.NEW_MODEL_DESIGN is False  # Should return default
+
+    monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "")
+    assert envs.SKIP_JAX_PRECOMPILE is False  # Should return default
+
+
+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_model_impl_type_choices(monkeypatch: pytest.MonkeyPatch):
+    # Test case sensitive choices
+    monkeypatch.setenv("MODEL_IMPL_TYPE", "flax_nnx")
+    assert envs.MODEL_IMPL_TYPE == "flax_nnx"
+
+    monkeypatch.setenv("MODEL_IMPL_TYPE", "vllm")
+    assert envs.MODEL_IMPL_TYPE == "vllm"
+
+
+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"
+
+
+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 "VLLM_XLA_CHECK_RECOMPILATION" 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", "ray")
+    assert envs.TPU_MULTIHOST_BACKEND == "ray"
+
+
+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 == "auto"
+
+
+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"

tests/test_tpu_info.py

@@ -0,0 +1,134 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from unittest.mock import MagicMock, patch
+
+import pytest
+import requests
+
+from tpu_inference.tpu_info import (get_node_name, get_node_worker_id,
+                                    get_num_chips, get_num_cores_per_chip,
+                                    get_tpu_metadata, get_tpu_type)
+
+
+# Mock requests.get for get_tpu_metadata tests
+@patch("tpu_inference.tpu_info.requests.get")
+def test_get_tpu_metadata_success(mock_get):
+    """Test get_tpu_metadata when the request is successful."""
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.text = "test_metadata_value"
+    mock_get.return_value = mock_response
+    assert get_tpu_metadata(key="test-key") == "test_metadata_value"
+
+
+@patch("tpu_inference.tpu_info.requests.get")
+def test_get_tpu_metadata_request_error(mock_get):
+    """Test get_tpu_metadata when a RequestException is raised."""
+    mock_get.side_effect = requests.RequestException("Test RequestException")
+    assert get_tpu_metadata(key="test-key") is None
+
+
+# Test get_tpu_type
+@patch("tpu_inference.tpu_info.get_tpu_metadata")
+@patch.dict(os.environ, {"TPU_ACCELERATOR_TYPE": "env_tpu_type"})
+def test_get_tpu_type_from_env(mock_get_tpu_metadata):
+    """Test get_tpu_type when TPU_ACCELERATOR_TYPE is set in environment."""
+    # The function should return the env var value and not call get_tpu_metadata
+    assert get_tpu_type() == "env_tpu_type"
+    mock_get_tpu_metadata.assert_not_called()
+
+
+@patch.dict(os.environ, {}, clear=True)
+@patch("tpu_inference.tpu_info.get_tpu_metadata",
+       return_value="metadata_tpu_type")
+def test_get_tpu_type_from_metadata(mock_get_tpu_metadata):
+    """Test get_tpu_type when environment variable is not set."""
+    assert get_tpu_type() == "metadata_tpu_type"
+    mock_get_tpu_metadata.assert_called_once_with(key="accelerator-type")
+
+
+# Test get_node_name
+@patch("tpu_inference.tpu_info.get_tpu_metadata")
+@patch.dict(os.environ, {"TPU_NAME": "env_tpu_name"})
+def test_get_node_name_from_env(mock_get_tpu_metadata):
+    """Test get_node_name when TPU_NAME is set in environment."""
+    assert get_node_name() == "env_tpu_name"
+    mock_get_tpu_metadata.assert_not_called()
+
+
+@patch.dict(os.environ, {}, clear=True)
+@patch("tpu_inference.tpu_info.get_tpu_metadata",
+       return_value="metadata_tpu_name")
+def test_get_node_name_from_metadata(mock_get_tpu_metadata):
+    """Test get_node_name when environment variable is not set."""
+    assert get_node_name() == "metadata_tpu_name"
+    mock_get_tpu_metadata.assert_called_once_with(key="instance-id")
+
+
+# Test get_node_worker_id
+@patch("tpu_inference.tpu_info.get_tpu_metadata")
+@patch.dict(os.environ, {"TPU_WORKER_ID": "5"})
+def test_get_node_worker_id_from_env(mock_get_tpu_metadata):
+    """Test get_node_worker_id when TPU_WORKER_ID is set in environment."""
+    assert get_node_worker_id() == 5
+    mock_get_tpu_metadata.assert_not_called()
+
+
+@patch.dict(os.environ, {}, clear=True)
+@patch("tpu_inference.tpu_info.get_tpu_metadata", return_value="10")
+def test_get_node_worker_id_from_metadata(mock_get_tpu_metadata):
+    """Test get_node_worker_id when environment variable is not set."""
+    assert get_node_worker_id() == 10
+    mock_get_tpu_metadata.assert_called_once_with(key="agent-worker-number")
+
+
+# Test get_num_cores_per_chip
+@pytest.mark.parametrize(
+    "tpu_type, expected",
+    [
+        ("v5litepod-4", 1),
+        ("v6e-8", 1),
+        ("v4-8", 2),
+        ("v5p-16", 2),
+        ("unknown-type", 2)  # Default case
+    ])
+@patch("tpu_inference.tpu_info.get_tpu_type")
+def test_get_num_cores_per_chip(mock_get_tpu_type, tpu_type, expected):
+    """Test get_num_cores_per_chip with different TPU types."""
+    mock_get_tpu_type.return_value = tpu_type
+    assert get_num_cores_per_chip() == expected
+
+
+# Test get_num_chips
+@patch("tpu_inference.tpu_info.glob.glob",
+       return_value=["/dev/accel0", "/dev/accel1"])
+def test_get_num_chips_from_accel(mock_glob):
+    """Test get_num_chips when /dev/accel* files exist."""
+    assert get_num_chips() == 2
+
+
+@patch("tpu_inference.tpu_info.glob.glob", return_value=[])
+@patch("tpu_inference.tpu_info.os.listdir", return_value=["0", "1", "2"])
+def test_get_num_chips_from_vfio(mock_listdir, mock_glob):
+    """Test get_num_chips when /dev/accel* files don't exist but /dev/vfio entries do."""
+    assert get_num_chips() == 3
+
+
+@patch("tpu_inference.tpu_info.glob.glob", return_value=[])
+@patch("tpu_inference.tpu_info.os.listdir", side_effect=FileNotFoundError)
+def test_get_num_chips_not_found(mock_listdir, mock_glob, caplog):
+    """Test get_num_chips when neither files nor directory are found."""
+    assert get_num_chips() == 0

tests/test_utils.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+from unittest.mock import MagicMock, patch
+
+import jax.numpy as jnp
+import pytest
+
+# Import the functions to be tested
+from tpu_inference.utils import (GBYTES, enable_megacore,
+                                 get_jax_dtype_from_str_dtype, get_megacore,
+                                 get_padded_head_dim, hbm_usage_bytes,
+                                 hbm_usage_gb)
+
+
+def test_enable_and_get_megacore():
+    """Tests the enable_megacore and get_megacore functions."""
+    assert not get_megacore()
+    enable_megacore()
+    assert get_megacore()
+
+
+@patch.dict(os.environ, {"TPU_MULTIHOST_BACKEND": "ray"})
+def test_hbm_usage_bytes_ray_backend():
+    """Tests hbm_usage_bytes when TPU_MULTIHOST_BACKEND is ray."""
+    mock_device1 = MagicMock()
+    mock_device1.memory_stats.return_value = {
+        "bytes_in_use": 100 * GBYTES,
+        "bytes_limit": 128 * GBYTES
+    }
+    mock_device2 = MagicMock()
+    mock_device2.memory_stats.side_effect = Exception("Memory stats failed")
+
+    devices = [mock_device1, mock_device2]
+    usage = hbm_usage_bytes(devices)
+
+    expected_usage = [(100 * GBYTES, 128 * GBYTES),
+                      (100 * GBYTES, 128 * GBYTES)]
+    assert usage == expected_usage
+
+
+@patch("vllm.envs.VLLM_TPU_USING_PATHWAYS", False)
+def test_hbm_usage_bytes_pathways_disabled():
+    """Tests hbm_usage_bytes when VLLM_TPU_USING_PATHWAYS is False."""
+    mock_device1 = MagicMock()
+    mock_device1.memory_stats.return_value = {
+        "bytes_in_use": 100 * GBYTES,
+        "bytes_limit": 128 * GBYTES
+    }
+    mock_device2 = MagicMock()
+    mock_device2.memory_stats.return_value = {
+        "bytes_in_use": 50 * GBYTES,
+        "bytes_limit": 128 * GBYTES
+    }
+
+    devices = [mock_device1, mock_device2]
+    usage = hbm_usage_bytes(devices)
+
+    expected_usage = [(100 * GBYTES, 128 * GBYTES),
+                      (50 * GBYTES, 128 * GBYTES)]
+    assert usage == expected_usage
+
+
+@patch("vllm.envs.VLLM_TPU_USING_PATHWAYS", True)
+@patch("jax.live_arrays")
+@patch("jax.devices")
+def test_hbm_usage_bytes_pathways_enabled(mock_devices, mock_live_arrays):
+    """Tests hbm_usage_bytes when VLLM_TPU_USING_PATHWAYS is True."""
+    # Mock TPU v5p devices
+    mock_jax_device = MagicMock()
+    mock_jax_device.device_kind = "TPU v5p"
+    mock_devices.return_value = [mock_jax_device]
+
+    # Create mock devices
+    mock_device1 = MagicMock()
+    mock_device2 = MagicMock()
+    devices = [mock_device1, mock_device2]
+
+    # Create mock addressable shards with data property
+    mock_data1_dev1 = MagicMock()
+    mock_data1_dev1.device = mock_device1
+    mock_data1_dev1.nbytes = 2000  # 2000 bytes on device1
+
+    mock_data1_dev2 = MagicMock()
+    mock_data1_dev2.device = mock_device2
+    mock_data1_dev2.nbytes = 2000  # 2000 bytes on device2
+
+    mock_data2_dev1 = MagicMock()
+    mock_data2_dev1.device = mock_device1
+    mock_data2_dev1.nbytes = 1000  # 1000 bytes on device1
+
+    mock_shard1_dev1 = MagicMock()
+    mock_shard1_dev1.data = mock_data1_dev1
+
+    mock_shard1_dev2 = MagicMock()
+    mock_shard1_dev2.data = mock_data1_dev2
+
+    mock_shard2_dev1 = MagicMock()
+    mock_shard2_dev1.data = mock_data2_dev1
+
+    # Create mock arrays with addressable_shards
+    mock_array1 = MagicMock()
+    mock_array1.addressable_shards = [mock_shard1_dev1, mock_shard1_dev2]
+
+    mock_array2 = MagicMock()
+    mock_array2.addressable_shards = [mock_shard2_dev1]
+
+    mock_live_arrays.return_value = [mock_array1, mock_array2]
+
+    usage = hbm_usage_bytes(devices)
+
+    # Expected calculations:
+    # Array1: 2000 bytes on device1, 2000 bytes on device2
+    # Array2: 1000 bytes on device1
+    # Device1 total: 2000 + 1000 = 3000 bytes
+    # Device2 total: 2000 + 0 = 2000 bytes
+    # hbm_limit = 95 * GBYTES for TPU v5p
+    expected_usage = [(3000, 95 * GBYTES), (2000, 95 * GBYTES)]
+    assert usage == expected_usage
+
+
+@patch("vllm.envs.VLLM_TPU_USING_PATHWAYS", False)
+def test_hbm_usage_gb_pathways_disabled():
+    """Tests hbm_usage_gb when VLLM_TPU_USING_PATHWAYS is False."""
+    mock_device1 = MagicMock()
+    mock_device1.memory_stats.return_value = {
+        "bytes_in_use": 100 * GBYTES,
+        "bytes_limit": 128 * GBYTES
+    }
+    mock_device2 = MagicMock()
+    mock_device2.memory_stats.return_value = {
+        "bytes_in_use": 50.5 * GBYTES,
+        "bytes_limit": 128.0 * GBYTES
+    }
+
+    devices = [mock_device1, mock_device2]
+    usage = hbm_usage_gb(devices)
+
+    expected_usage = [(100.0, 128.0), (50.5, 128.0)]
+    assert usage == expected_usage
+
+
+@patch("vllm.envs.VLLM_TPU_USING_PATHWAYS", True)
+@patch("jax.live_arrays")
+@patch("jax.devices")
+def test_hbm_usage_bytes_pathways_no_arrays(mock_devices, mock_live_arrays):
+    """Tests hbm_usage_bytes when VLLM_TPU_USING_PATHWAYS is True but no live arrays."""
+    # Mock TPU v6e devices
+    mock_jax_device = MagicMock()
+    mock_jax_device.device_kind = "TPU v6e"
+    mock_devices.return_value = [mock_jax_device]
+
+    mock_device1 = MagicMock()
+    mock_device2 = MagicMock()
+    devices = [mock_device1, mock_device2]
+
+    # No live arrays
+    mock_live_arrays.return_value = []
+
+    usage = hbm_usage_bytes(devices)
+
+    # No arrays means no memory usage, defaultdict returns 0 for missing keys
+    # HBM limit for TPU v6e is 32 GB
+    expected_usage = [(0, 32 * GBYTES), (0, 32 * GBYTES)]
+    assert usage == expected_usage
+
+
+@pytest.mark.parametrize(
+    "head_dim, expected_padded_head_dim",
+    [
+        (1, 128),
+        (64, 64),
+        (127, 128),
+        (128, 128),
+        (129, 256),
+        (255, 256),
+        (256, 256),
+        (0, 0),  # Although head_dim is usually positive, testing boundary
+    ],
+)
+def test_get_padded_head_dim(head_dim, expected_padded_head_dim):
+    """Tests the get_padded_head_dim function."""
+    assert get_padded_head_dim(head_dim) == expected_padded_head_dim
+
+
+def test_get_jax_dtype_from_str_dtype():
+    """
+    Test the get_jax_dtype_from_str_dtype function
+    """
+    assert get_jax_dtype_from_str_dtype("int8") == jnp.int8
+    assert get_jax_dtype_from_str_dtype("bfloat16") == jnp.bfloat16
+    assert get_jax_dtype_from_str_dtype("fp8") == jnp.float8_e4m3fn
+    assert get_jax_dtype_from_str_dtype("fp8_e4m3") == jnp.float8_e4m3fn
+    assert get_jax_dtype_from_str_dtype("fp8_e5m2") == jnp.float8_e5m2

tests/worker/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.