chore: import upstream snapshot with attribution

python/ray/tests/accelerators/mock_acl.py

@@ -0,0 +1,8 @@
+def get_soc_name():
+    return "Ascend910B"
+
+
+class rt:
+    @staticmethod
+    def get_device_count():
+        return (4, 0)

python/ray/tests/accelerators/mock_dpctl_1.py

@@ -0,0 +1,16 @@
+class SyclContext:
+    def __init__(self, info):
+        pass
+
+    @property
+    def device_count(self):
+        return 6
+
+
+class SyclDevice:
+    def __init__(self, info):
+        pass
+
+    @property
+    def name(self):
+        return "Intel(R) Data Center GPU Max 1550"

python/ray/tests/accelerators/mock_dpctl_2.py

@@ -0,0 +1,16 @@
+class SyclContext:
+    def __init__(self, info):
+        pass
+
+    @property
+    def device_count(self):
+        return 4
+
+
+class SyclDevice:
+    def __init__(self, info):
+        pass
+
+    @property
+    def name(self):
+        return "Intel(R) Data Center GPU Max 1100"

python/ray/tests/accelerators/mock_furiosa_smi_py.py

@@ -0,0 +1,42 @@
+class _MockArch:
+    """Mock for the PyO3-generated Arch enum."""
+
+    def __init__(self, name: str = "Rngd"):
+        self._name = name
+
+    def __str__(self) -> str:
+        # PyO3 unit enums commonly stringify like "Arch.Rngd".
+        return f"Arch.{self._name}"
+
+
+class _MockDeviceInfo:
+    def __init__(self, arch_name: str = "Rngd", index: int = 0):
+        self._arch_name = arch_name
+        self._index = index
+
+    def arch(self) -> _MockArch:
+        return _MockArch(self._arch_name)
+
+    def name(self) -> str:
+        return f"npu{self._index}"
+
+    def index(self) -> int:
+        return self._index
+
+
+class _MockDevice:
+    def __init__(self, idx: int, arch_name: str = "Rngd"):
+        self._idx = idx
+        self._arch_name = arch_name
+
+    def device_info(self) -> _MockDeviceInfo:
+        return _MockDeviceInfo(arch_name=self._arch_name, index=self._idx)
+
+
+def init():
+    """Mock for ``furiosa_smi_py.init``. Real SMI library requires init() before use."""
+    return None
+
+
+def list_devices():
+    return [_MockDevice(i) for i in range(8)]

python/ray/tests/accelerators/mock_pynvml.py

@@ -0,0 +1,101 @@
+from typing import List
+from unittest.mock import patch
+
+import pytest
+
+import ray._private.thirdparty.pynvml as pynvml
+
+
+class DeviceHandleMock(dict):
+    def __init__(
+        self,
+        name: str,
+        uuid: str,
+        mig_devices: List["DeviceHandleMock"] = None,
+        **kwargs
+    ):
+        super().__init__()
+        self["name"] = name
+        self["uuid"] = uuid
+        if mig_devices is not None:
+            self["mig_devices"] = mig_devices
+        self.update(kwargs)
+
+
+# pnvml mock for gpu resources
+class PyNVMLMock:
+    def __init__(self, mock_data, driver_version="535.104.12"):
+        self._mock_data = mock_data
+        self.driver_version = driver_version
+
+    def nvmlInit(self):
+        return
+
+    def nvmlShutdown(self):
+        return
+
+    def nvmlSystemGetDriverVersion(self):
+        return self.driver_version
+
+    def nvmlDeviceGetCount(self):
+        return len(self._mock_data)
+
+    def nvmlDeviceGetHandleByIndex(self, index):
+        return self._mock_data[index]
+
+    def nvmlDeviceGetName(self, handle):
+        return handle.get("name", "")
+
+    def nvmlDeviceGetMaxMigDeviceCount(self, handle):
+        if "mig_devices" in handle:
+            return max(7, len(handle["mig_devices"]))
+        else:
+            raise pynvml.NVMLError_NotSupported
+
+    def nvmlDeviceGetMigDeviceHandleByIndex(self, handle, mig_index):
+        try:
+            return handle["mig_devices"][mig_index]
+        except IndexError:
+            raise pynvml.NVMLError_NotFound
+
+    def nvmlDeviceGetUUID(self, handle):
+        return handle.get("uuid", "")
+
+    def nvmlDeviceGetComputeInstanceId(self, mig_handle):
+        return mig_handle["ci_id"]
+
+    def nvmlDeviceGetGpuInstanceId(self, mig_handle):
+        return mig_handle["gi_id"]
+
+
+@pytest.fixture
+def patch_mock_pynvml(mock_nvml):
+    with patch("ray._private.thirdparty.pynvml.nvmlInit", mock_nvml.nvmlInit), patch(
+        "ray._private.thirdparty.pynvml.nvmlShutdown", mock_nvml.nvmlShutdown
+    ), patch(
+        "ray._private.thirdparty.pynvml.nvmlSystemGetDriverVersion",
+        mock_nvml.nvmlSystemGetDriverVersion,
+    ), patch(
+        "ray._private.thirdparty.pynvml.nvmlDeviceGetCount",
+        mock_nvml.nvmlDeviceGetCount,
+    ), patch(
+        "ray._private.thirdparty.pynvml.nvmlDeviceGetHandleByIndex",
+        mock_nvml.nvmlDeviceGetHandleByIndex,
+    ), patch(
+        "ray._private.thirdparty.pynvml.nvmlDeviceGetName", mock_nvml.nvmlDeviceGetName
+    ), patch(
+        "ray._private.thirdparty.pynvml.nvmlDeviceGetMaxMigDeviceCount",
+        mock_nvml.nvmlDeviceGetMaxMigDeviceCount,
+    ), patch(
+        "ray._private.thirdparty.pynvml.nvmlDeviceGetMigDeviceHandleByIndex",
+        mock_nvml.nvmlDeviceGetMigDeviceHandleByIndex,
+    ), patch(
+        "ray._private.thirdparty.pynvml.nvmlDeviceGetUUID", mock_nvml.nvmlDeviceGetUUID
+    ), patch(
+        "ray._private.thirdparty.pynvml.nvmlDeviceGetComputeInstanceId",
+        mock_nvml.nvmlDeviceGetComputeInstanceId,
+    ), patch(
+        "ray._private.thirdparty.pynvml.nvmlDeviceGetGpuInstanceId",
+        mock_nvml.nvmlDeviceGetGpuInstanceId,
+    ):
+        yield

python/ray/tests/accelerators/mock_rebel.py

@@ -0,0 +1,6 @@
+def device_count():
+    return 4
+
+
+def get_npu_name():
+    return "RBLN-CA02"

python/ray/tests/accelerators/test_accelerators.py

@@ -0,0 +1,22 @@
+import sys
+
+import pytest
+
+from ray.util import accelerators
+from ray.util.annotations import RayDeprecationWarning
+
+
+def test_accelerators():
+    assert accelerators.NVIDIA_TESLA_K80 == "K80"
+    assert accelerators.NVIDIA_A100 == "A100"
+    with pytest.raises(
+        AttributeError,
+        match="module 'ray.util.accelerators' has no attribute 'NVIDIA_INVALID'",
+    ):
+        _ = accelerators.NVIDIA_INVALID
+    with pytest.warns(RayDeprecationWarning):
+        assert accelerators.NVIDIA_TESLA_A100 == "A100"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/accelerators/test_amd_gpu.py

@@ -0,0 +1,92 @@
+import os
+import sys
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray._private.accelerators import (
+    AMDGPUAcceleratorManager,
+    get_accelerator_manager_for_resource,
+)
+
+
+@patch(
+    "ray._private.accelerators.AMDGPUAcceleratorManager.get_current_node_num_accelerators",  # noqa: E501
+    return_value=4,
+)
+def test_visible_amd_gpu_ids(mock_get_num_accelerators, monkeypatch, shutdown_only):
+    monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0,1,2")
+    # Delete the cache so it can be re-populated the next time
+    # we call get_accelerator_manager_for_resource
+    del get_accelerator_manager_for_resource._resource_name_to_accelerator_manager
+    ray.init()
+    _ = mock_get_num_accelerators.called
+    assert ray.available_resources()["GPU"] == 3
+
+
+@patch(
+    "ray._private.accelerators.AMDGPUAcceleratorManager._get_amd_device_ids",
+    return_value=["0x74a1", "0x74a1", "0x74a1", "0x74a1"],
+)
+def test_visible_amd_gpu_type(mock_get_amd_device_ids, shutdown_only):
+    ray.init()
+    _ = mock_get_amd_device_ids.called
+    assert (
+        AMDGPUAcceleratorManager.get_current_node_accelerator_type()
+        == "AMD-Instinct-MI300X-OAM"
+    )
+
+
+@patch(
+    "ray._private.accelerators.AMDGPUAcceleratorManager._get_amd_device_ids",
+    return_value=["0x640f", "0x640f", "0x640f", "0x640f"],
+)
+def test_visible_amd_gpu_type_bad_device_id(mock_get_num_accelerators, shutdown_only):
+    ray.init()
+    _ = mock_get_num_accelerators.called
+    assert AMDGPUAcceleratorManager.get_current_node_accelerator_type() is None
+
+
+def test_get_current_process_visible_accelerator_ids(monkeypatch):
+    monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0,1,2")
+    assert AMDGPUAcceleratorManager.get_current_process_visible_accelerator_ids() == [
+        "0",
+        "1",
+        "2",
+    ]
+
+    monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0,2,7")
+    assert AMDGPUAcceleratorManager.get_current_process_visible_accelerator_ids() == [
+        "0",
+        "2",
+        "7",
+    ]
+
+    monkeypatch.setenv("HIP_VISIBLE_DEVICES", "")
+    assert AMDGPUAcceleratorManager.get_current_process_visible_accelerator_ids() == []
+
+    del os.environ["HIP_VISIBLE_DEVICES"]
+    assert (
+        AMDGPUAcceleratorManager.get_current_process_visible_accelerator_ids() is None
+    )
+
+
+def test_set_current_process_visible_accelerator_ids():
+    AMDGPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0"])
+    env_var = AMDGPUAcceleratorManager.get_visible_accelerator_ids_env_var()
+    assert os.environ[env_var] == "0"
+
+    AMDGPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0", "1"])
+    assert os.environ[env_var] == "0,1"
+
+    AMDGPUAcceleratorManager.set_current_process_visible_accelerator_ids(
+        ["0", "1", "7"]
+    )
+    assert os.environ[env_var] == "0,1,7"
+
+    del os.environ[env_var]
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/accelerators/test_furiosa.py

@@ -0,0 +1,159 @@
+import os
+import sys
+
+import pytest
+
+from ray._private.accelerators.furiosa import (
+    FURIOSA_VISIBLE_DEVICES_ENV_VAR,
+    NOSET_FURIOSA_VISIBLE_DEVICES_ENV_VAR,
+    FuriosaAcceleratorManager,
+)
+
+
+@pytest.fixture(autouse=True)
+def mock_furiosa_smi_py(monkeypatch):
+    from ray.tests.accelerators import mock_furiosa_smi_py
+
+    monkeypatch.setitem(sys.modules, "furiosa_smi_py", mock_furiosa_smi_py)
+
+
+@pytest.fixture
+def clear_furiosa_environment():
+    original_env = os.environ.get(FURIOSA_VISIBLE_DEVICES_ENV_VAR)
+    original_no_set_env = os.environ.get(NOSET_FURIOSA_VISIBLE_DEVICES_ENV_VAR)
+
+    os.environ.pop(FURIOSA_VISIBLE_DEVICES_ENV_VAR, None)
+    os.environ.pop(NOSET_FURIOSA_VISIBLE_DEVICES_ENV_VAR, None)
+
+    yield
+
+    if original_env is not None:
+        os.environ[FURIOSA_VISIBLE_DEVICES_ENV_VAR] = original_env
+    if original_no_set_env is not None:
+        os.environ[NOSET_FURIOSA_VISIBLE_DEVICES_ENV_VAR] = original_no_set_env
+
+
+@pytest.mark.usefixtures("clear_furiosa_environment")
+class TestFuriosaAcceleratorManager:
+    def test_get_resource_name(self):
+        assert FuriosaAcceleratorManager.get_resource_name() == "FURIOSA"
+
+    def test_get_visible_accelerator_ids_env_var(self):
+        assert (
+            FuriosaAcceleratorManager.get_visible_accelerator_ids_env_var()
+            == FURIOSA_VISIBLE_DEVICES_ENV_VAR
+        )
+
+    @pytest.mark.parametrize(
+        "env_value,expected",
+        [
+            # furiosa-llm --devices form (preferred).
+            ("npu:0,npu:1,npu:2,npu:3", ["0", "1", "2", "3"]),
+            # Bare integer form is also accepted for convenience.
+            ("0,1,2,3", ["0", "1", "2", "3"]),
+            # Core range notation: only the device index is returned.
+            ("npu:0:0-3,npu:1:0-3", ["0", "1"]),
+            # Empty string yields an empty list.
+            ("", []),
+            # Sentinel ``None`` means the env var is unset.
+            (None, None),
+        ],
+    )
+    def test_get_current_process_visible_accelerator_ids(self, env_value, expected):
+        if env_value is None:
+            os.environ.pop(FURIOSA_VISIBLE_DEVICES_ENV_VAR, None)
+        else:
+            os.environ[FURIOSA_VISIBLE_DEVICES_ENV_VAR] = env_value
+        assert (
+            FuriosaAcceleratorManager.get_current_process_visible_accelerator_ids()
+            == expected
+        )
+
+    def test_get_current_node_num_accelerators(self):
+        assert FuriosaAcceleratorManager.get_current_node_num_accelerators() == 8
+
+    def test_get_current_node_accelerator_type(self):
+        assert (
+            FuriosaAcceleratorManager.get_current_node_accelerator_type()
+            == "FURIOSA_RNGD"
+        )
+
+    @pytest.mark.parametrize(
+        "arch_name,expected",
+        [
+            # PyO3 enum form (CamelCase).
+            ("Rngd", "FURIOSA_RNGD"),
+            ("RngdMax", "FURIOSA_RNGDMAX"),
+            ("RngdS", "FURIOSA_RNGDS"),
+            ("RngdPlus", "FURIOSA_RNGDPLUS"),
+            # ``Arch::ToString`` form is also accepted, and both forms must
+            # resolve to the same label.
+            ("rngd-max", "FURIOSA_RNGDMAX"),
+            # ``+`` must NOT be silently stripped, since that would collapse
+            # ``rngd+`` into ``rngd`` and collide with the base RNGD SKU; it
+            # is mapped to ``plus`` so the label matches ``RngdPlus``.
+            ("rngd+", "FURIOSA_RNGDPLUS"),
+        ],
+    )
+    def test_get_current_node_accelerator_type_dynamic(
+        self, monkeypatch, arch_name, expected
+    ):
+        from ray.tests.accelerators import mock_furiosa_smi_py
+
+        def mocked_list_devices():
+            return [mock_furiosa_smi_py._MockDevice(0, arch_name=arch_name)]
+
+        monkeypatch.setattr(mock_furiosa_smi_py, "list_devices", mocked_list_devices)
+        assert FuriosaAcceleratorManager.get_current_node_accelerator_type() == expected
+
+    def test_get_current_node_accelerator_type_no_devices(self, monkeypatch):
+        from ray.tests.accelerators import mock_furiosa_smi_py
+
+        monkeypatch.setattr(mock_furiosa_smi_py, "list_devices", lambda: [])
+        assert FuriosaAcceleratorManager.get_current_node_accelerator_type() is None
+
+    def test_get_current_node_accelerator_type_arch_is_none(self, monkeypatch):
+        """Regression: arch() returning None must not produce 'FURIOSA_NONE'."""
+        from ray.tests.accelerators import mock_furiosa_smi_py
+
+        class _NullArchDeviceInfo:
+            def arch(self):
+                return None
+
+        class _NullArchDevice:
+            def device_info(self):
+                return _NullArchDeviceInfo()
+
+        monkeypatch.setattr(
+            mock_furiosa_smi_py, "list_devices", lambda: [_NullArchDevice()]
+        )
+        assert FuriosaAcceleratorManager.get_current_node_accelerator_type() is None
+
+    def test_set_current_process_visible_accelerator_ids(self):
+        # Ray's scheduler hands us bare integer IDs; we serialize them in
+        # the ``npu:<id>`` form expected by ``furiosa-llm --devices``.
+        FuriosaAcceleratorManager.set_current_process_visible_accelerator_ids(
+            ["0", "1"]
+        )
+        assert os.environ[FURIOSA_VISIBLE_DEVICES_ENV_VAR] == "npu:0,npu:1"
+
+        os.environ[NOSET_FURIOSA_VISIBLE_DEVICES_ENV_VAR] = "1"
+        FuriosaAcceleratorManager.set_current_process_visible_accelerator_ids(
+            ["2", "3"]
+        )
+        assert os.environ[FURIOSA_VISIBLE_DEVICES_ENV_VAR] == "npu:0,npu:1"
+
+    def test_validate_resource_request_quantity(self):
+        valid, _ = FuriosaAcceleratorManager.validate_resource_request_quantity(1)
+        assert valid
+
+        valid, _ = FuriosaAcceleratorManager.validate_resource_request_quantity(2.0)
+        assert valid
+
+        valid, msg = FuriosaAcceleratorManager.validate_resource_request_quantity(0.5)
+        assert not valid
+        assert "whole number" in msg
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/accelerators/test_hpu.py

@@ -0,0 +1,296 @@
+import os
+import sys
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray._private.accelerators import HPUAcceleratorManager, hpu
+from ray.util.placement_group import placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+
+def test_user_configured_more_than_visible(monkeypatch, call_ray_stop_only):
+    # Test more hpus are configured than visible.
+    monkeypatch.setenv("HABANA_VISIBLE_MODULES", "0,1,2")
+    with pytest.raises(ValueError):
+        ray.init(resources={"HPU": 4})
+
+
+@patch(
+    "ray._private.accelerators.HPUAcceleratorManager.get_current_node_num_accelerators",  # noqa: E501
+    return_value=4,
+)
+def test_auto_detected_more_than_visible(
+    mock_get_num_accelerators, monkeypatch, shutdown_only
+):
+    # Test more hpus are detected than visible.
+    monkeypatch.setenv("HABANA_VISIBLE_MODULES", "0,1,2")
+    ray.init()
+    _ = mock_get_num_accelerators.called
+    assert ray.available_resources()["HPU"] == 3
+
+
+@patch(
+    "ray._private.accelerators.HPUAcceleratorManager.get_current_node_num_accelerators",  # noqa: E501
+    return_value=2,
+)
+def test_auto_detect_resources(mock_get_num_accelerators, shutdown_only):
+    # Test that ray node resources are filled with auto detected count.
+    ray.init()
+    _ = mock_get_num_accelerators.called
+    assert ray.available_resources()["HPU"] == 2
+
+
+def test_get_current_process_visible_accelerator_ids():
+    os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] = "0,1,2"
+    assert HPUAcceleratorManager.get_current_process_visible_accelerator_ids() == [
+        "0",
+        "1",
+        "2",
+    ]  # noqa: E501
+
+    del os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR]
+    assert HPUAcceleratorManager.get_current_process_visible_accelerator_ids() is None
+
+    os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] = ""
+    assert HPUAcceleratorManager.get_current_process_visible_accelerator_ids() == []
+
+    del os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR]
+
+
+def test_set_current_process_visible_accelerator_ids():
+    HPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0"])
+    assert os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] == "0"
+
+    HPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0", "1"])
+    assert os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] == "0,1"
+
+    HPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0", "1", "2"])
+    assert os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] == "0,1,2"
+
+    del os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR]
+
+
+@pytest.mark.parametrize(
+    "test_config",
+    [
+        (1, False),
+        (0.5, True),
+        (3, False),
+    ],
+)
+def test_validate_resource_request_quantity(test_config):
+    num_hpus, expect_error = test_config
+
+    if expect_error:
+        assert (
+            HPUAcceleratorManager.validate_resource_request_quantity(num_hpus)[0]
+            is False
+        )
+        assert (
+            HPUAcceleratorManager.validate_resource_request_quantity(num_hpus)[1]
+            is not None
+        )
+    else:
+        assert (
+            HPUAcceleratorManager.validate_resource_request_quantity(num_hpus)[0]
+            is True
+        )
+        assert (
+            HPUAcceleratorManager.validate_resource_request_quantity(num_hpus)[1]
+            is None
+        )
+
+
+def test_check_accelerator_info():
+
+    if HPUAcceleratorManager.is_initialized():
+        assert (
+            "Intel-GAUDI" in HPUAcceleratorManager.get_current_node_accelerator_type()
+        )
+    else:
+        assert HPUAcceleratorManager.get_current_node_accelerator_type() is None
+
+    assert HPUAcceleratorManager.get_resource_name() == "HPU"
+
+
+def test_decorator_args():
+
+    # This is a valid way of using the decorator.
+    @ray.remote(resources={"HPU": 1})  # noqa: F811
+    class Actor:  # noqa: F811
+        def __init__(self):
+            pass
+
+    # This is a valid way of using the decorator.
+    @ray.remote(num_cpus=1, resources={"HPU": 1})  # noqa: F811
+    class Actor:  # noqa: F811
+        def __init__(self):
+            pass
+
+
+def test_actor_deletion_with_hpus(shutdown_only):
+    ray.init(num_cpus=1, resources={"HPU": 1})
+
+    # When an actor that uses an HPU exits, make sure that the HPU resources
+    # are released.
+
+    @ray.remote(resources={"HPU": 1})
+    class Actor:
+        def getpid(self):
+            return os.getpid()
+
+    for _ in range(5):
+        # If we can successfully create an actor, that means that enough
+        # HPU resources are available.
+        a = Actor.remote()
+        ray.get(a.getpid.remote())
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Failing on Windows.")
+def test_actor_hpus(ray_start_cluster):
+    cluster = ray_start_cluster
+    num_nodes = 2
+    num_hpus_per_raylet = 2
+    for i in range(num_nodes):
+        cluster.add_node(num_cpus=10 * 2, resources={"HPU": num_hpus_per_raylet})
+    ray.init(address=cluster.address)
+
+    @ray.remote(resources={"HPU": 1})
+    class Actor1:
+        def __init__(self):
+            resource_ids = ray.get_runtime_context().get_accelerator_ids()
+            self.hpu_ids = resource_ids.get("HPU")
+
+        def get_location_and_ids(self):
+            return (
+                ray.get_runtime_context().get_node_id(),
+                tuple(self.hpu_ids),
+            )
+
+    # Create one actor per HPU.
+    actors = [Actor1.remote() for _ in range(num_nodes * num_hpus_per_raylet)]
+    # Make sure that no two actors are assigned to the same HPU.
+    locations_and_ids = ray.get(
+        [actor.get_location_and_ids.remote() for actor in actors]
+    )
+    node_names = {location for location, hpu_id in locations_and_ids}
+    assert len(node_names) == num_nodes
+    location_actor_combinations = []
+    for node_name in node_names:
+        for hpu_id in range(num_hpus_per_raylet):
+            location_actor_combinations.append((node_name, (f"{hpu_id}",)))
+
+    assert set(locations_and_ids) == set(location_actor_combinations)
+
+    # Creating a new actor should fail because all of the HPUs are being
+    # used.
+    a = Actor1.remote()
+    ready_ids, _ = ray.wait([a.get_location_and_ids.remote()], timeout=0.01)
+    assert ready_ids == []
+
+
+def test_actor_habana_visible_devices(shutdown_only):
+    """Test user can overwrite HABANA_VISIBLE_MODULES
+    after the actor is created."""
+    ray.init(resources={"HPU": 1})
+
+    @ray.remote(resources={"HPU": 1})
+    class Actor:
+        def set_habana_visible_devices(self, habana_visible_devices):
+            os.environ["HABANA_VISIBLE_MODULES"] = habana_visible_devices
+
+        def get_habana_visible_devices(self):
+            return os.environ["HABANA_VISIBLE_MODULES"]
+
+    actor = Actor.remote()
+    assert ray.get(actor.get_habana_visible_devices.remote()) == "0"
+    ray.get(actor.set_habana_visible_devices.remote("0,1"))
+    assert ray.get(actor.get_habana_visible_devices.remote()) == "0,1"
+
+
+def test_hpu_ids(shutdown_only):
+    num_hpus = 3
+    ray.init(num_cpus=num_hpus, resources={"HPU": num_hpus})
+
+    def get_hpu_ids(hpus_per_worker):
+        hpu_ids = ray.get_runtime_context().get_accelerator_ids()["HPU"]
+        assert len(hpu_ids) == hpus_per_worker
+        modules = os.environ.get("HABANA_VISIBLE_MODULES")
+        if modules is not None:
+            assert modules == ",".join([str(i) for i in hpu_ids])  # noqa
+        for hpu_id in hpu_ids:
+            assert hpu_id in [str(i) for i in range(num_hpus)]
+        return hpu_ids
+
+    f0 = ray.remote(resources={"HPU": 0})(lambda: get_hpu_ids(0))
+    f1 = ray.remote(resources={"HPU": 1})(lambda: get_hpu_ids(1))
+
+    list_of_ids = ray.get([f0.remote() for _ in range(10)])
+    assert list_of_ids == 10 * [[]]
+    ray.get([f1.remote() for _ in range(10)])
+
+    # Test that actors have HABANA_VISIBLE_MODULES set properly.
+
+    def _check_hpu_env(expected_num_hpus):
+        hpu_ids = ray.get_runtime_context().get_accelerator_ids()["HPU"]
+        assert len(hpu_ids) == expected_num_hpus
+        if expected_num_hpus > 0:
+            assert os.environ["HABANA_VISIBLE_MODULES"] == ",".join(
+                [str(i) for i in hpu_ids]  # noqa
+            )
+        else:
+            assert os.environ.get("HABANA_VISIBLE_MODULES") is None
+
+    @ray.remote
+    class Actor:
+        def __init__(self, num_hpus):
+            self.num_hpus = num_hpus
+            _check_hpu_env(num_hpus)
+            self.x = num_hpus
+
+        def test(self):
+            _check_hpu_env(self.num_hpus)
+            return self.x
+
+    a0 = Actor.remote(0)
+    assert ray.get(a0.test.remote()) == 0
+
+    a1 = Actor.options(resources={"HPU": 1}).remote(1)
+    assert ray.get(a1.test.remote()) == 1
+
+
+def test_hpu_with_placement_group(shutdown_only):
+    num_hpus = 2
+    ray.init(num_cpus=1, resources={"HPU": num_hpus})
+
+    @ray.remote(resources={"HPU": num_hpus})
+    class HPUActor:
+        def __init__(self):
+            pass
+
+        def ready(self):
+            hpu_ids = ray.get_runtime_context().get_accelerator_ids()["HPU"]
+            assert len(hpu_ids) == num_hpus
+            assert os.environ["HABANA_VISIBLE_MODULES"] == ",".join(
+                [str(i) for i in hpu_ids]  # noqa
+            )
+
+    # Reserve a placement group of 1 bundle that reserves 1 CPU and 2 HPU.
+    pg = placement_group([{"CPU": 1, "HPU": num_hpus}])
+
+    # Wait until placement group is created.
+    ray.get(pg.ready(), timeout=10)
+
+    actor = HPUActor.options(
+        scheduling_strategy=PlacementGroupSchedulingStrategy(
+            placement_group=pg,
+        )
+    ).remote()
+
+    ray.get(actor.ready.remote(), timeout=10)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/accelerators/test_intel_gpu.py

@@ -0,0 +1,117 @@
+import os
+import sys
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray._private.accelerators import (
+    IntelGPUAcceleratorManager as Accelerator,
+    get_accelerator_manager_for_resource,
+)
+from ray.util.accelerators import INTEL_MAX_1100, INTEL_MAX_1550
+
+
+def test_visible_intel_gpu_ids(shutdown_only):
+    with patch.object(Accelerator, "get_current_node_num_accelerators", return_value=4):
+        os.environ["ONEAPI_DEVICE_SELECTOR"] = "level_zero:0,1,2"
+        # Delete the cache so it can be re-populated the next time
+        # we call get_accelerator_manager_for_resource
+        del get_accelerator_manager_for_resource._resource_name_to_accelerator_manager
+        ray.init()
+        manager = get_accelerator_manager_for_resource("GPU")
+        assert manager.get_current_node_num_accelerators() == 4
+        assert manager.__name__ == "IntelGPUAcceleratorManager"
+        assert ray.available_resources()["GPU"] == 3
+
+
+def test_visible_intel_gpu_type(shutdown_only):
+    with patch.object(
+        Accelerator, "get_current_node_num_accelerators", return_value=4
+    ), patch.object(
+        Accelerator, "get_current_node_accelerator_type", return_value=INTEL_MAX_1550
+    ):
+        os.environ["ONEAPI_DEVICE_SELECTOR"] = "level_zero:0,1,2"
+        del get_accelerator_manager_for_resource._resource_name_to_accelerator_manager
+        ray.init()
+        manager = get_accelerator_manager_for_resource("GPU")
+        assert manager.get_current_node_accelerator_type() == INTEL_MAX_1550
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Not supported mock on Windows")
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12),
+    reason="Not passing on Python 3.12. Being followed up by external contributors.",
+)
+def test_get_current_node_num_accelerators():
+    old_dpctl = None
+    if "dpctl" in sys.modules:
+        old_dpctl = sys.modules["dpctl"]
+
+    sys.modules["dpctl"] = __import__("mock_dpctl_1")
+    assert Accelerator.get_current_node_num_accelerators() == 6
+
+    sys.modules["dpctl"] = __import__("mock_dpctl_2")
+    assert Accelerator.get_current_node_num_accelerators() == 4
+
+    if old_dpctl is not None:
+        sys.modules["dpctl"] = old_dpctl
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Not supported mock on Windows")
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12),
+    reason="Not passing on Python 3.12. Being followed up by external contributors.",
+)
+def test_get_current_node_accelerator_type():
+    old_dpctl = None
+    if "dpctl" in sys.modules:
+        old_dpctl = sys.modules["dpctl"]
+
+    sys.modules["dpctl"] = __import__("mock_dpctl_1")
+    assert Accelerator.get_current_node_accelerator_type() == INTEL_MAX_1550
+
+    sys.modules["dpctl"] = __import__("mock_dpctl_2")
+    assert Accelerator.get_current_node_accelerator_type() == INTEL_MAX_1100
+
+    if old_dpctl is not None:
+        sys.modules["dpctl"] = old_dpctl
+
+
+def test_intel_gpu_accelerator_manager_api():
+    assert Accelerator.get_resource_name() == "GPU"
+    assert Accelerator.get_visible_accelerator_ids_env_var() == "ONEAPI_DEVICE_SELECTOR"
+    assert Accelerator.validate_resource_request_quantity(0.1) == (True, None)
+
+
+def test_get_current_process_visible_accelerator_ids():
+    os.environ["ONEAPI_DEVICE_SELECTOR"] = "level_zero:0,1,2"
+    assert Accelerator.get_current_process_visible_accelerator_ids() == ["0", "1", "2"]
+
+    del os.environ["ONEAPI_DEVICE_SELECTOR"]
+    assert Accelerator.get_current_process_visible_accelerator_ids() is None
+
+    os.environ["ONEAPI_DEVICE_SELECTOR"] = ""
+    assert Accelerator.get_current_process_visible_accelerator_ids() == []
+
+    os.environ["ONEAPI_DEVICE_SELECTOR"] = "NoDevFiles"
+    assert Accelerator.get_current_process_visible_accelerator_ids() == []
+
+    del os.environ["ONEAPI_DEVICE_SELECTOR"]
+
+
+def test_set_current_process_visible_accelerator_ids():
+    Accelerator.set_current_process_visible_accelerator_ids(["0"])
+    assert os.environ["ONEAPI_DEVICE_SELECTOR"] == "level_zero:0"
+
+    Accelerator.set_current_process_visible_accelerator_ids(["0", "1"])
+    assert os.environ["ONEAPI_DEVICE_SELECTOR"] == "level_zero:0,1"
+
+    Accelerator.set_current_process_visible_accelerator_ids(["0", "1", "2"])
+    assert os.environ["ONEAPI_DEVICE_SELECTOR"] == "level_zero:0,1,2"
+
+    del os.environ["ONEAPI_DEVICE_SELECTOR"]
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/accelerators/test_intel_gpu_e2e.py

@@ -0,0 +1,228 @@
+"""
+Manual Intel GPU validation tests, not executed in automated runs.
+
+These tests are basic acceptance tests to validate Intel GPU support in Ray. They
+require a suitable Intel GPU environment with dpctl installed. They are intended to
+serve as an approved method to verify Intel GPU-based Ray deployments.
+"""
+
+import os
+import re
+from typing import Any, Dict, List
+
+import pytest
+
+import ray
+
+try:
+    import dpctl
+except ImportError:
+    pytest.skip(
+        "dpctl is not installed, skipping Intel GPU tests.", allow_module_level=True
+    )
+
+DEFAULT_SCALE_OUT_NODES = 2
+DEFAULT_SCALE_UP_DEVICES = 2
+
+USE_GPU = bool(os.environ.get("RAY_PYTEST_USE_GPU", 0))
+
+if not USE_GPU:
+    pytest.skip("Skipping, these tests require GPUs.", allow_module_level=True)
+
+
+@pytest.fixture
+def ray_gpu_session():
+    """Start a Ray session with caller-provided init kwargs."""
+
+    def _start_session(**init_kwargs):
+        if ray.is_initialized():
+            ray.shutdown()
+
+        ray.init(**init_kwargs)
+
+    try:
+        yield _start_session
+    finally:
+        if ray.is_initialized():
+            ray.shutdown()
+
+
+def _is_cluster_configured(address: str = "auto") -> bool:
+    try:
+        ray.init(
+            address=address,
+        )
+        return True
+    except (ray.exceptions.RaySystemError, ConnectionError, TimeoutError):
+        return False
+    finally:
+        if ray.is_initialized():
+            ray.shutdown()
+
+
+def _detect_available_gpu_count() -> int:
+    """Return the number of GPU devices detected via dpctl."""
+    try:
+        return dpctl.SyclContext("level_zero:gpu").device_count
+    except Exception:
+        # If dpctl cannot enumerate devices, assume no additional GPUs.
+        return 0
+
+
+def _require_min_gpus(required: int, context: str) -> None:
+    available = _detect_available_gpu_count()
+    if available < required:
+        pytest.skip(
+            f"Skipping {context}: requires {required} GPUs, detected {available} via dpctl."
+        )
+
+
+def _require_min_cluster_nodes(required_nodes: int, context: str) -> None:
+    alive_nodes = [node for node in ray.nodes() if node.get("Alive")]
+    unique_node_ids = {node.get("NodeID") for node in alive_nodes if node.get("NodeID")}
+
+    if len(unique_node_ids) < required_nodes:
+        pytest.skip(
+            f"Skipping {context}: requires {required_nodes} alive Ray nodes, detected {len(unique_node_ids)}."
+        )
+
+
+@ray.remote(num_gpus=1)
+def gpu_task() -> Dict[str, Any]:
+    context = ray.get_runtime_context()
+    gpu_ids = context.get_accelerator_ids().get("GPU", [])
+    return {
+        "gpu_ids": gpu_ids,
+        "pid": os.getpid(),
+        "oneapi_selector": os.environ.get("ONEAPI_DEVICE_SELECTOR"),
+    }
+
+
+@ray.remote(num_gpus=1)
+def cluster_probe_task() -> Dict[str, Any]:
+    context = ray.get_runtime_context()
+    return {
+        "node_id": context.get_node_id(),
+        "node_ip": ray.util.get_node_ip_address(),
+        "worker_id": context.get_worker_id(),
+        "gpu_ids": context.get_accelerator_ids().get("GPU", []),
+        "selector": os.environ.get("ONEAPI_DEVICE_SELECTOR"),
+    }
+
+
+def assert_valid_gpu_binding(result: Dict[str, Any], label: str) -> None:
+    primary_gpu_id = _validate_gpu_binding_common(result, label)
+    assert (
+        primary_gpu_id >= 0
+    ), f"Expected {label} to bind to a valid GPU, got {result.get('gpu_ids')}"
+
+
+def _validate_gpu_binding_common(
+    result: Dict[str, Any], label: str, selector_key: str = "oneapi_selector"
+) -> int:
+    """Validate basic GPU binding properties shared by single- and multi-GPU tests."""
+
+    gpu_ids = result.get("gpu_ids")
+    assert gpu_ids, f"No GPU IDs assigned for {label}."
+
+    primary_gpu_id = int(gpu_ids[0])
+
+    selector = result.get(selector_key)
+    assert selector, f"ONEAPI_DEVICE_SELECTOR not set in environment for {label}."
+    selector_lower = selector.lower()
+    assert (
+        "level_zero:" in selector_lower
+    ), f"ONEAPI_DEVICE_SELECTOR should target GPU devices for {label}, got: {selector}."
+
+    selector_gpu_ids = {int(match) for match in re.findall(r"\b\d+\b", selector_lower)}
+    assert (
+        primary_gpu_id in selector_gpu_ids
+    ), f"ONEAPI_DEVICE_SELECTOR does not reference bound GPU id for {label}: {selector}."
+
+    return primary_gpu_id
+
+
+def assert_valid_multi_gpu_binding(
+    results: List[Dict[str, Any]], num_gpus: int, label: str
+) -> None:
+    """Assert that multiple GPU tasks bind to different GPUs correctly."""
+    assert (
+        len(results) == num_gpus
+    ), f"Expected {num_gpus} results for {label}, got {len(results)}."
+
+    gpu_ids = []
+    for i, result in enumerate(results):
+        primary_gpu_id = _validate_gpu_binding_common(result, f"{label} instance {i}")
+        gpu_ids.append(primary_gpu_id)
+    assert (
+        len(set(gpu_ids)) == num_gpus
+    ), f"Expected {label} to bind to {num_gpus} distinct GPUs, got bindings to GPU IDs: {gpu_ids}."
+
+
+@pytest.mark.skipif(
+    _is_cluster_configured(),
+    reason="Environment setup for scale-out, skipping single-node test.",
+)
+def test_gpu_task_binding(ray_gpu_session) -> None:
+    _require_min_gpus(1, "single GPU task binding test")
+    ray_gpu_session(num_gpus=1)
+
+    task_result = ray.get(gpu_task.remote())
+    assert_valid_gpu_binding(task_result, "GPU task")
+
+
+@pytest.mark.skipif(
+    _is_cluster_configured(),
+    reason="Environment setup for scale-out, skipping single-node test.",
+)
+@pytest.mark.parametrize(
+    "num_gpus", [DEFAULT_SCALE_UP_DEVICES]
+)  # To be extended to required configurations
+def test_multi_gpu_task_binding(ray_gpu_session, num_gpus) -> None:
+    """Test that multiple GPU tasks bind to different GPUs correctly."""
+    _require_min_gpus(num_gpus, "multi-GPU task binding test")
+
+    ray_gpu_session(num_gpus=num_gpus)
+
+    task_futures = [gpu_task.remote() for _ in range(num_gpus)]
+    task_results = ray.get(task_futures)
+
+    assert_valid_multi_gpu_binding(task_results, num_gpus, f"GPU tasks (n={num_gpus})")
+
+
+@pytest.mark.skipif(
+    not _is_cluster_configured(), reason="Environment not setup for scale-out test."
+)
+@pytest.mark.parametrize(
+    "num_nodes", [DEFAULT_SCALE_OUT_NODES]
+)  # To be extended to required configurations
+def test_scale_out_task_distribution(ray_gpu_session, num_nodes) -> None:
+    """Ensure tasks can be scheduled across multiple nodes in the cluster."""
+    ray_gpu_session(address="auto")
+    _require_min_cluster_nodes(num_nodes, "scale-out task distribution test")
+
+    probe_handles = [
+        cluster_probe_task.options(scheduling_strategy="SPREAD").remote()
+        for _ in range(num_nodes)
+    ]
+    probe_results = ray.get(probe_handles)
+
+    node_ids = {
+        result.get("node_id") for result in probe_results if result.get("node_id")
+    }
+    node_ips = {
+        result.get("node_ip") for result in probe_results if result.get("node_ip")
+    }
+
+    for result in probe_results:
+        _validate_gpu_binding_common(result, "scale-out probe task", "selector")
+
+    assert len(node_ids) == num_nodes or len(node_ips) == num_nodes, (
+        f"Expected probe tasks to execute on {num_nodes} distinct nodes, "
+        f"got node_ids={node_ids} node_ips={node_ips}."
+    )
+
+    gpu_capable_results = [result for result in probe_results if result.get("gpu_ids")]
+    assert (
+        len(gpu_capable_results) == num_nodes
+    ), "Not all probe tasks reported GPU accelerator bindings in the cluster."

python/ray/tests/accelerators/test_metax_gpu.py

@@ -0,0 +1,69 @@
+import os
+import sys
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray._private.accelerators import (
+    MetaxGPUAcceleratorManager,
+    get_accelerator_manager_for_resource,
+)
+
+
+@patch(
+    "ray._private.accelerators.MetaxGPUAcceleratorManager.get_current_node_num_accelerators",
+    return_value=4,
+)
+def test_visible_metax_gpu_ids(mock_get_num_accelerators, monkeypatch, shutdown_only):
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1,2")
+    del get_accelerator_manager_for_resource._resource_name_to_accelerator_manager
+    ray.init()
+    assert mock_get_num_accelerators.called
+    assert ray.available_resources()["GPU"] == 3
+
+
+def test_get_current_process_visible_accelerator_ids(monkeypatch):
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0")
+    assert MetaxGPUAcceleratorManager.get_current_process_visible_accelerator_ids() == [
+        "0"
+    ]
+
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,4,7")
+    assert MetaxGPUAcceleratorManager.get_current_process_visible_accelerator_ids() == [
+        "0",
+        "4",
+        "7",
+    ]
+
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "")
+    assert (
+        MetaxGPUAcceleratorManager.get_current_process_visible_accelerator_ids() == []
+    )
+
+    monkeypatch.delenv("CUDA_VISIBLE_DEVICES")
+    assert (
+        MetaxGPUAcceleratorManager.get_current_process_visible_accelerator_ids() is None
+    )
+
+
+def test_set_current_process_visible_accelerator_ids():
+    MetaxGPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0"])
+    assert os.environ["CUDA_VISIBLE_DEVICES"] == "0"
+
+    MetaxGPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0", "1"])
+    assert os.environ["CUDA_VISIBLE_DEVICES"] == "0,1"
+
+    MetaxGPUAcceleratorManager.set_current_process_visible_accelerator_ids(
+        ["0", "1", "7"]
+    )
+    assert os.environ["CUDA_VISIBLE_DEVICES"] == "0,1,7"
+
+    del os.environ["CUDA_VISIBLE_DEVICES"]
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/accelerators/test_neuron.py

@@ -0,0 +1,104 @@
+import subprocess
+import sys
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray._private.accelerators import NeuronAcceleratorManager
+
+
+def test_user_configured_more_than_visible(monkeypatch, call_ray_stop_only):
+    # Test more neuron_cores are configured than visible.
+    monkeypatch.setenv("NEURON_RT_VISIBLE_CORES", "0,1,2")
+    with pytest.raises(ValueError):
+        ray.init(resources={"neuron_cores": 4})
+
+
+@patch(
+    "ray._private.accelerators.NeuronAcceleratorManager.get_current_node_num_accelerators",  # noqa: E501
+    return_value=4,
+)
+def test_auto_detected_more_than_visible(
+    mock_get_num_accelerators, monkeypatch, shutdown_only
+):
+    # Test more neuron_cores are detected than visible.
+    monkeypatch.setenv("NEURON_RT_VISIBLE_CORES", "0,1,2")
+    ray.init()
+    _ = mock_get_num_accelerators.called
+    assert ray.available_resources()["neuron_cores"] == 3
+
+
+@patch(
+    "ray._private.accelerators.NeuronAcceleratorManager.get_current_node_num_accelerators",  # noqa: E501
+    return_value=2,
+)
+def test_auto_detect_resources(mock_get_num_accelerators, shutdown_only):
+    # Test that ray node resources are filled with auto detected count.
+    ray.init()
+    _ = mock_get_num_accelerators.called
+    assert ray.available_resources()["neuron_cores"] == 2
+
+
+@patch(
+    "subprocess.run",
+    return_value=subprocess.CompletedProcess(
+        args=[],
+        returncode=0,
+        stdout=(
+            b'[{"neuron_device":0,"bdf":"00:1e.0",'
+            b'"connected_to":null,"nc_count":2,'
+            b'"memory_size":34359738368,"neuron_processes":[]}]'
+        ),
+    ),
+)
+@patch("os.path.isdir", return_value=True)
+@patch("sys.platform", "linux")
+def test_get_neuron_core_count_single_device(mock_isdir, mock_subprocess):
+    assert NeuronAcceleratorManager.get_current_node_num_accelerators() == 2
+
+
+@patch(
+    "subprocess.run",
+    return_value=subprocess.CompletedProcess(
+        args=[],
+        returncode=0,
+        stdout=(
+            b'[{"neuron_device":0,"bdf":"00:1e.0",'
+            b'"connected_to":null,"nc_count":2,'
+            b'"memory_size":34359738368,"neuron_processes":[]},'
+            b'{"neuron_device":1,"bdf":"00:1f.0","connected_to":null,'
+            b'"nc_count":2,"memory_size":34359738368,"neuron_processes":[]}]'
+        ),
+    ),
+)
+@patch("os.path.isdir", return_value=True)
+@patch("sys.platform", "linux")
+def test_get_neuron_core_count_multiple_devices(mock_isdir, mock_subprocess):
+    assert NeuronAcceleratorManager.get_current_node_num_accelerators() == 4
+
+
+@patch(
+    "subprocess.run",
+    return_value=subprocess.CompletedProcess(
+        args=[], returncode=1, stdout=b"AccessDenied"
+    ),
+)
+@patch("os.path.isdir", return_value=True)
+@patch("sys.platform", "linux")
+def test_get_neuron_core_count_failure_with_error(mock_isdir, mock_subprocess):
+    assert NeuronAcceleratorManager.get_current_node_num_accelerators() == 0
+
+
+@patch(
+    "subprocess.run",
+    return_value=subprocess.CompletedProcess(args=[], returncode=0, stdout=b"[{}]"),
+)
+@patch("os.path.isdir", return_value=True)
+@patch("sys.platform", "linux")
+def test_get_neuron_core_count_failure_with_empty_results(mock_isdir, mock_subprocess):
+    assert NeuronAcceleratorManager.get_current_node_num_accelerators() == 0
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/accelerators/test_npu.py

@@ -0,0 +1,131 @@
+import os
+import sys
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray._private.accelerators import NPUAcceleratorManager as Accelerator
+
+
+@patch("glob.glob")
+def test_autodetect_num_npus(mock_glob):
+    with patch.dict(sys.modules):
+        sys.modules["acl"] = None
+        mock_glob.return_value = [f"/dev/davinci{i}" for i in range(64)]
+        assert Accelerator.get_current_node_num_accelerators() == 64
+
+
+@patch("glob.glob")
+def test_autodetect_num_npus_without_devices(mock_glob):
+    with patch.dict(sys.modules):
+        sys.modules["acl"] = None
+        mock_glob.side_effect = Exception
+        assert Accelerator.get_current_node_num_accelerators() == 0
+
+
+def test_ascend_npu_accelerator_manager_api():
+    assert Accelerator.get_resource_name() == "NPU"
+    assert (
+        Accelerator.get_visible_accelerator_ids_env_var() == "ASCEND_RT_VISIBLE_DEVICES"
+    )
+    assert Accelerator.validate_resource_request_quantity(0.5) == (True, None)
+    assert Accelerator.validate_resource_request_quantity(1) == (True, None)
+
+
+def test_visible_ascend_npu_type(monkeypatch, shutdown_only):
+    with patch.object(
+        Accelerator, "get_current_node_num_accelerators", return_value=4
+    ), patch.object(
+        Accelerator, "get_current_node_accelerator_type", return_value="Ascend910B"
+    ):
+        monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "0,1,2")
+        manager = ray._private.accelerators.get_accelerator_manager_for_resource("NPU")
+        assert manager.get_current_node_accelerator_type() == "Ascend910B"
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Not supported mock on Windows")
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12),
+    reason="Not passing on Python 3.12. Being followed up by external contributors.",
+)
+def test_visible_ascend_npu_ids(monkeypatch, shutdown_only):
+    with patch.dict(sys.modules):
+        sys.modules["acl"] = __import__("mock_acl")
+
+        monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "0,1,2")
+        with patch.object(
+            Accelerator, "get_current_node_num_accelerators", return_value=4
+        ):
+
+            ray.init()
+            manager = ray._private.accelerators.get_accelerator_manager_for_resource(
+                "NPU"
+            )
+            assert manager.get_current_node_num_accelerators() == 4
+            assert manager.__name__ == "NPUAcceleratorManager"
+            assert ray.available_resources()["NPU"] == 3
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Not supported mock on Windows")
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12),
+    reason="Not passing on Python 3.12. Being followed up by external contributors.",
+)
+def test_acl_api_function(shutdown_only):
+    with patch.dict(sys.modules):
+        sys.modules["acl"] = __import__("mock_acl")
+
+        ray.init()
+        manager = ray._private.accelerators.get_accelerator_manager_for_resource("NPU")
+        assert manager.get_current_node_num_accelerators() == 4
+        assert manager.__name__ == "NPUAcceleratorManager"
+        assert manager.get_current_node_accelerator_type() == "Ascend910B"
+
+
+def test_get_current_process_visible_accelerator_ids(monkeypatch, shutdown_only):
+    monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "0,1,2")
+    assert Accelerator.get_current_process_visible_accelerator_ids() == ["0", "1", "2"]
+
+    monkeypatch.delenv("ASCEND_RT_VISIBLE_DEVICES")
+    assert Accelerator.get_current_process_visible_accelerator_ids() is None
+
+    monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "")
+    assert Accelerator.get_current_process_visible_accelerator_ids() == []
+
+    monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "NoDevFiles")
+    assert Accelerator.get_current_process_visible_accelerator_ids() == []
+
+
+def test_set_current_process_visible_accelerator_ids(shutdown_only):
+    Accelerator.set_current_process_visible_accelerator_ids(["0"])
+    assert os.environ["ASCEND_RT_VISIBLE_DEVICES"] == "0"
+
+    Accelerator.set_current_process_visible_accelerator_ids(["0", "1"])
+    assert os.environ["ASCEND_RT_VISIBLE_DEVICES"] == "0,1"
+
+    Accelerator.set_current_process_visible_accelerator_ids(["0", "1", "2"])
+    assert os.environ["ASCEND_RT_VISIBLE_DEVICES"] == "0,1,2"
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Not supported mock on Windows")
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12),
+    reason="Not passing on Python 3.12. Being followed up by external contributors.",
+)
+def test_auto_detected_more_than_visible(monkeypatch, shutdown_only):
+    with patch.dict(sys.modules):
+        sys.modules["acl"] = __import__("mock_acl")
+
+        with patch.object(
+            Accelerator, "get_current_node_num_accelerators", return_value=4
+        ):
+            # If more NPUs are detected than visible.
+            monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "0,1,2")
+
+            ray.init()
+            assert ray.available_resources()["NPU"] == 3
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/accelerators/test_nvidia_gpu.py

@@ -0,0 +1,103 @@
+import sys
+
+import pytest
+
+from ray._private.accelerators import NvidiaGPUAcceleratorManager
+from ray.tests.accelerators.mock_pynvml import (
+    DeviceHandleMock,
+    PyNVMLMock,
+    patch_mock_pynvml,
+)
+
+GPU_MOCK_DATA = [
+    DeviceHandleMock(
+        "Ampere A100-SXM4-40GB",
+        "GPU-8eaaebb8-bb64-8489-fda2-62256e821983",
+        mig_devices=[
+            DeviceHandleMock(
+                "Ampere A100-SXM4-40GB MIG 1g.5gb",
+                "MIG-c6d4f1ef-42e4-5de3-91c7-45d71c87eb3f",
+                gi_id=0,
+                ci_instance=0,
+            ),
+            DeviceHandleMock(
+                "Ampere A100-SXM4-40GB MIG 1g.5gb",
+                "MIG-0c757cd7-e942-5726-a0b8-0e8fb7067135",
+                gi_id=1,
+                ci_instance=0,
+            ),
+        ],
+    ),
+    DeviceHandleMock(
+        "Ampere A100-SXM4-40GB",
+        "GPU-8eaaebb8-bb64-8489-fda2-62256e821983",
+        mig_devices=[
+            DeviceHandleMock(
+                "Ampere A100-SXM4-40GB MIG 1g.5gb",
+                "MIG-a28ad590-3fda-56dd-84fc-0a0b96edc58d",
+                gi_id=0,
+                ci_instance=0,
+            )
+        ],
+    ),
+    DeviceHandleMock(
+        "Tesla V100-SXM2-16GB", "GPU-8eaaebb8-bb64-8489-fda2-62256e821983"
+    ),
+]
+
+mock_nvml = PyNVMLMock(GPU_MOCK_DATA)
+
+patch_mock_pynvml = patch_mock_pynvml  # avoid format error
+
+
+@pytest.mark.parametrize("mock_nvml", [mock_nvml])
+def test_num_gpus_parsing(patch_mock_pynvml):
+    # without mig instance
+    assert NvidiaGPUAcceleratorManager.get_current_node_num_accelerators() == len(
+        GPU_MOCK_DATA
+    )
+
+
+@pytest.mark.parametrize("mock_nvml", [mock_nvml])
+def test_gpu_info_parsing(patch_mock_pynvml):
+    assert NvidiaGPUAcceleratorManager.get_current_node_accelerator_type() == "A100"
+
+
+@pytest.mark.parametrize(
+    "name,expected",
+    [
+        # Legacy datacenter GPU names: keep labels produced by the previous
+        # parser stable.
+        ("Tesla V100-SXM2-16GB", "V100"),
+        ("Tesla P100-PCIE-16GB", "P100"),
+        ("Tesla T4", "T4"),
+        ("Tesla P4", "P4"),
+        ("Tesla K80", "K80"),
+        ("NVIDIA A10G", "A10G"),
+        ("NVIDIA L4", "L4"),
+        ("NVIDIA L40S", "L40S"),
+        ("NVIDIA A100-SXM4-40GB", "A100"),
+        ("NVIDIA H100 80GB HBM3", "H100"),
+        ("NVIDIA H200", "H200"),
+        ("NVIDIA H20", "H20"),
+        ("NVIDIA B200", "B200"),
+        ("NVIDIA B300", "B300"),
+        # Consumer GPUs: the regex does not match the mixed-case product line,
+        # so we fall back to a hyphen-joined product name.
+        ("NVIDIA GeForce RTX 5090", "GeForce-RTX-5090"),
+        ("NVIDIA GeForce RTX 4090", "GeForce-RTX-4090"),
+        # RTX PRO cards: "RTX" alone is just a brand prefix, so the model is
+        # captured through the first digit-containing token instead of
+        # collapsing to the ambiguous "RTX".
+        ("NVIDIA RTX PRO 6000 Blackwell Server Edition", "RTX-PRO-6000"),
+        # Edge cases.
+        (None, None),
+        ("", None),
+    ],
+)
+def test_gpu_name_to_accelerator_type(name, expected):
+    assert NvidiaGPUAcceleratorManager._gpu_name_to_accelerator_type(name) == expected
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/accelerators/test_rbln.py

@@ -0,0 +1,82 @@
+import os
+import sys
+
+import pytest
+
+from ray._private.accelerators.rbln import (
+    NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR,
+    RBLN_RT_VISIBLE_DEVICES_ENV_VAR,
+    RBLNAcceleratorManager,
+)
+
+
+@pytest.fixture(autouse=True)
+def mock_rebel_module(monkeypatch):
+    from ray.tests.accelerators import mock_rebel
+
+    monkeypatch.setitem(sys.modules, "rebel", mock_rebel)
+
+
+@pytest.fixture
+def clear_rbln_environment():
+    original_env = os.environ.get(RBLN_RT_VISIBLE_DEVICES_ENV_VAR)
+    original_no_set_env = os.environ.get(NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR)
+
+    os.environ.pop(RBLN_RT_VISIBLE_DEVICES_ENV_VAR, None)
+    os.environ.pop(NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR, None)
+
+    yield
+
+    if original_env is not None:
+        os.environ[RBLN_RT_VISIBLE_DEVICES_ENV_VAR] = original_env
+    if original_no_set_env is not None:
+        os.environ[NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR] = original_no_set_env
+
+
+@pytest.mark.usefixtures("clear_rbln_environment")
+class TestRBLNAcceleratorManager:
+    def test_get_resource_name(self):
+        assert RBLNAcceleratorManager.get_resource_name() == "RBLN"
+
+    def test_get_visible_accelerator_ids_env_var(self):
+        assert (
+            RBLNAcceleratorManager.get_visible_accelerator_ids_env_var()
+            == RBLN_RT_VISIBLE_DEVICES_ENV_VAR
+        )
+
+    def test_get_current_process_visible_accelerator_ids(self):
+        os.environ[RBLN_RT_VISIBLE_DEVICES_ENV_VAR] = "0,1,2,3"
+        assert RBLNAcceleratorManager.get_current_process_visible_accelerator_ids() == [
+            "0",
+            "1",
+            "2",
+            "3",
+        ]
+
+        os.environ[RBLN_RT_VISIBLE_DEVICES_ENV_VAR] = ""
+        assert (
+            RBLNAcceleratorManager.get_current_process_visible_accelerator_ids() == []
+        )
+
+        os.environ.pop(RBLN_RT_VISIBLE_DEVICES_ENV_VAR)
+        assert (
+            RBLNAcceleratorManager.get_current_process_visible_accelerator_ids() is None
+        )
+
+    def test_get_current_node_num_accelerators(self):
+        assert RBLNAcceleratorManager.get_current_node_num_accelerators() == 4
+
+    def test_get_current_node_accelerator_type(self):
+        assert RBLNAcceleratorManager.get_current_node_accelerator_type() == "RBLN-CA02"
+
+    def test_set_current_process_visible_accelerator_ids(self):
+        RBLNAcceleratorManager.set_current_process_visible_accelerator_ids(["0", "1"])
+        assert os.environ[RBLN_RT_VISIBLE_DEVICES_ENV_VAR] == "0,1"
+
+        os.environ[NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR] = "1"
+        RBLNAcceleratorManager.set_current_process_visible_accelerator_ids(["2", "3"])
+        assert os.environ[RBLN_RT_VISIBLE_DEVICES_ENV_VAR] == "0,1"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/accelerators/test_tpu.py

@@ -0,0 +1,346 @@
+import os
+import sys
+from unittest import mock
+from unittest.mock import patch
+
+import pytest
+import requests
+
+from ray._private.accelerators import TPUAcceleratorManager, tpu
+
+
+@patch("glob.glob")
+def test_autodetect_num_tpus_accel(mock_glob):
+    mock_glob.return_value = [
+        "/dev/accel0",
+        "/dev/accel1",
+        "/dev/accel2",
+        "/dev/accel3",
+    ]
+    TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()
+    assert TPUAcceleratorManager.get_current_node_num_accelerators() == 4
+
+
+@patch("os.path.isdir")
+@patch("glob.glob")
+@patch("os.listdir")
+def test_autodetect_num_tpus_accel_ignores_blackwell_directory(
+    mock_list, mock_glob, mock_isdir
+):
+    # NVIDIA drivers 570.x (Blackwell-class GPUs, e.g. RTX 5090) create
+    # /dev/accel as a directory containing /dev/accel/accel0. The non-recursive
+    # glob matches the directory entry; filtering directories out keeps real
+    # TPU chips (character devices at /dev/accel0..N) while rejecting the
+    # NVIDIA false positive.
+    mock_glob.return_value = ["/dev/accel"]
+    mock_isdir.side_effect = lambda p: p == "/dev/accel"
+    mock_list.side_effect = FileNotFoundError
+    TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()
+    assert TPUAcceleratorManager.get_current_node_num_accelerators() == 0
+
+
+@patch("glob.glob")
+@patch("os.listdir")
+def test_autodetect_num_tpus_vfio(mock_list, mock_glob):
+    mock_glob.return_value = []
+    mock_list.return_value = [f"{i}" for i in range(4)]
+    TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()
+    assert TPUAcceleratorManager.get_current_node_num_accelerators() == 4
+
+
+@patch("glob.glob")
+@patch("os.listdir")
+def test_autodetect_num_tpus_without_devices(mock_list, mock_glob):
+    mock_list.side_effect = FileNotFoundError
+    mock_glob.return_value = []
+    TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()
+    assert TPUAcceleratorManager.get_current_node_num_accelerators() == 0
+
+
+@pytest.mark.parametrize(
+    "accelerator_type_version_tuple",
+    [
+        ("gce", "v2-8", "TPU-V2"),
+        ("gce", "v2-32", "TPU-V2"),
+        ("gce", "v3-8", "TPU-V3"),
+        ("gce", "v3-128", "TPU-V3"),
+        ("gce", "v4-8", "TPU-V4"),
+        ("gce", "v4-2048", "TPU-V4"),
+        ("gce", "v5p-8", "TPU-V5P"),
+        ("gce", "v5litepod-8", "TPU-V5LITEPOD"),
+        ("gce", "v6e-8", "TPU-V6E"),
+        ("gke", "v2-8", "TPU-V2"),
+        ("gke", "v2-32", "TPU-V2"),
+        ("gke", "v3-8", "TPU-V3"),
+        ("gke", "v3-128", "TPU-V3"),
+        ("gke", "v4-8", "TPU-V4"),
+        ("gke", "v4-2048", "TPU-V4"),
+        ("gke", "v5p-8", "TPU-V5P"),
+        ("gke", "v5litepod-8", "TPU-V5LITEPOD"),
+        ("gke", "v6e-8", "TPU-V6E"),
+        ("gke", "tpu7x-16", "TPU-V7X"),
+    ],
+)
+@patch("requests.get")
+@patch("os.getenv")
+def test_autodetect_tpu_accelerator_type(
+    mock_os, mock_request, accelerator_type_version_tuple
+):
+    gce_or_gke, accelerator_type, expected_version = accelerator_type_version_tuple
+    if gce_or_gke == "gce":
+        mock_response = mock.MagicMock()
+        mock_response.status_code = 200
+        mock_response.text = accelerator_type
+        mock_request.return_value = mock_response
+        mock_os.return_value = None
+    else:
+        mock_os.return_value = accelerator_type
+    assert TPUAcceleratorManager.get_current_node_accelerator_type() == expected_version
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        ("gce", "0", 0),
+        ("gke", "0", 0),
+    ],
+)
+@patch("requests.get")
+@patch("os.getenv")
+def test_get_current_node_tpu_worker_id(mock_os, mock_request, test_case):
+    gce_or_gke, worker_id, expected_value = test_case
+    if gce_or_gke == "gce":
+        mock_response = mock.MagicMock()
+        mock_response.status_code = 200
+        mock_response.text = worker_id
+        mock_request.return_value = mock_response
+        mock_os.return_value = None
+    else:
+        mock_os.return_value = worker_id
+    assert TPUAcceleratorManager.get_current_node_tpu_worker_id() == expected_value
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        ("gce", "my-tpu"),
+        ("gke", "my-tpu"),
+    ],
+)
+@patch("requests.get")
+@patch("os.getenv")
+def test_get_tpu_unique_id(mock_os, mock_request, test_case):
+    gce_or_gke, worker_id = test_case
+    if gce_or_gke == "gce":
+        mock_response = mock.MagicMock()
+        mock_response.status_code = 200
+        mock_response.text = worker_id
+        mock_request.return_value = mock_response
+        mock_os.return_value = None
+    else:
+        mock_os.return_value = worker_id
+    assert TPUAcceleratorManager.get_current_node_tpu_name() == worker_id
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        ("gce", "not-a-valid-version"),
+        ("gce", "vNOTVALID-8"),
+        ("gce", "230498230948230948"),
+        # From issue #39913
+        ("gce", ""),
+        ("gke", "not-a-valid-version"),
+        ("gke", "vNOTVALID-8"),
+        ("gke", "230498230948230948"),
+    ],
+)
+@patch("requests.get")
+@patch("os.getenv")
+def test_autodetect_invalid_type(mock_os, mock_request, test_case):
+    gce_or_gke, accelerator_type = test_case
+    if gce_or_gke == "gce":
+        mock_response = mock.MagicMock()
+        mock_response.status_code = 200
+        mock_response.text = accelerator_type
+        mock_request.return_value = mock_response
+        mock_os.return_value = None
+    else:
+        mock_os.return_value = accelerator_type
+    assert TPUAcceleratorManager.get_current_node_accelerator_type() is None
+
+
+def test_autodetect_tpu_accelerator_type_fails_gracefully():
+    with patch("requests.get") as mock_get:
+        mock_get.side_effect = requests.exceptions.RequestException
+        assert TPUAcceleratorManager.get_current_node_accelerator_type() is None
+
+
+@pytest.mark.parametrize(
+    "test_config",
+    [
+        (1, False),
+        (0.5, True),
+        (3, True),
+    ],
+)
+def test_validate_resource_request_quantity(test_config):
+    num_tpus, expect_error = test_config
+
+    if expect_error:
+        assert (
+            TPUAcceleratorManager.validate_resource_request_quantity(num_tpus)[0]
+            is False
+        )
+        assert (
+            TPUAcceleratorManager.validate_resource_request_quantity(num_tpus)[1]
+            is not None
+        )
+    else:
+        assert (
+            TPUAcceleratorManager.validate_resource_request_quantity(num_tpus)[0]
+            is True
+        )
+        assert (
+            TPUAcceleratorManager.validate_resource_request_quantity(num_tpus)[1]
+            is None
+        )
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        (4, ["0"]),
+        (4, ["0", "1"]),
+        (4, ["0", "1", "2", "3"]),
+        (8, ["0", "1", "2", "3", "4", "5", "6", "7"]),
+    ],
+)
+@patch("glob.glob")
+def test_set_tpu_visible_ids_and_bounds(mock_glob, test_case):
+    num_devices, tpu_chips = test_case
+    mock_glob.return_value = ["/dev/accel" + str(x) for x in range(num_devices)]
+    with patch.dict("os.environ", {}, clear=True):
+        TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()
+        TPUAcceleratorManager.set_current_process_visible_accelerator_ids(tpu_chips)
+        if len(tpu_chips) == 1:
+            assert (
+                os.environ[tpu.TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR]
+                == tpu.TPU_CHIPS_PER_HOST_BOUNDS_1_CHIP_CONFIG
+            )
+            assert os.environ[tpu.TPU_HOST_BOUNDS_ENV_VAR] == tpu.TPU_SINGLE_HOST_BOUNDS
+            assert os.environ[tpu.TPU_VISIBLE_CHIPS_ENV_VAR] == ",".join(tpu_chips)
+        elif len(tpu_chips) == 2:
+            assert (
+                os.environ[tpu.TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR]
+                == tpu.TPU_CHIPS_PER_HOST_BOUNDS_2_CHIP_CONFIG
+            )
+            assert os.environ[tpu.TPU_HOST_BOUNDS_ENV_VAR] == tpu.TPU_SINGLE_HOST_BOUNDS
+            assert os.environ[tpu.TPU_VISIBLE_CHIPS_ENV_VAR] == ",".join(tpu_chips)
+        elif len(tpu_chips) == 4:
+            # Check that nothing is set, let the ML framework use the defaults.
+            assert os.environ.get(tpu.TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR, None) is None
+            assert os.environ.get(tpu.TPU_SINGLE_HOST_BOUNDS, None) is None
+            assert os.environ.get(tpu.TPU_VISIBLE_CHIPS_ENV_VAR, None) is None
+        else:  # len(tpu_chips) == 8
+            assert os.environ.get(tpu.TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR, None) is None
+            assert os.environ.get(tpu.TPU_SINGLE_HOST_BOUNDS, None) is None
+            assert os.environ.get(tpu.TPU_VISIBLE_CHIPS_ENV_VAR, None) is None
+
+
+@pytest.mark.parametrize(
+    "test_config",
+    [
+        (0, "v4-16", {"TPU-v4-16-head": 1, "my-tpu": 1}),
+        (1, "v4-16", {"my-tpu": 1}),
+        (0, "tpu7x-16", {"TPU-v7x-16-head": 1, "my-tpu": 1}),
+    ],
+)
+def test_tpu_pod_detect_and_configure_worker(test_config):
+    worker_id, pod_type, expected_value = test_config
+    final_resources = {}
+    with patch(
+        "ray._private.accelerators.tpu.TPUAcceleratorManager.get_current_node_tpu_name",
+        return_value="my-tpu",
+    ):
+        with patch(
+            "ray._private.accelerators.tpu.TPUAcceleratorManager.get_current_node_tpu_worker_id",
+            return_value=worker_id,
+        ):
+            with patch.dict(os.environ, {"TPU_ACCELERATOR_TYPE": pod_type}):
+                final_resources = (
+                    TPUAcceleratorManager.get_current_node_additional_resources()
+                )
+
+    assert final_resources == expected_value
+
+
+@pytest.mark.parametrize(
+    "accelerator_type, expected",
+    [
+        ("v2-8", True),
+        ("v3-32", True),
+        ("v4-8", True),
+        ("v5p-8", True),
+        ("v5litepod-8", True),
+        ("v6e-8", True),
+        ("tpu7x-16", True),
+        ("v7x-16", True),
+        ("v-8", False),
+        ("8", False),
+        ("tpu-8", False),
+        ("v2", False),
+        ("v2-", False),
+        ("random-string", False),
+    ],
+)
+def test_is_valid_tpu_accelerator_type(accelerator_type, expected):
+    assert (
+        TPUAcceleratorManager.is_valid_tpu_accelerator_type(accelerator_type)
+        == expected
+    )
+
+
+def test_get_total_chips_from_accelerator_type():
+    assert tpu.get_total_chips_from_accelerator_type("v6e-16") == 16
+    assert tpu.get_total_chips_from_accelerator_type("v6e-8") == 8
+    assert (
+        tpu.get_total_chips_from_accelerator_type("v7x-16") == 8
+    )  # v7x has 2 cores per chip
+    assert (
+        tpu.get_total_chips_from_accelerator_type("v4-8") == 4
+    )  # v4 has 2 cores per chip
+
+    # Test invalid cases
+    with pytest.raises(ValueError, match="Accelerator type must include size"):
+        tpu.get_total_chips_from_accelerator_type("v6e")
+
+    with pytest.raises(ValueError, match="Invalid accelerator type"):
+        tpu.get_total_chips_from_accelerator_type("invalid-8")
+
+
+def test_get_num_tpu_visible_chips_per_host():
+    # v6e multi-host (4 chips per VM)
+    assert tpu.get_num_tpu_visible_chips_per_host("v6e-16") == 4
+    assert tpu.get_num_tpu_visible_chips_per_host("v6e-32") == 4
+
+    # v6e single-host/sub-host (exact chip count)
+    assert tpu.get_num_tpu_visible_chips_per_host("v6e-8") == 8
+    assert tpu.get_num_tpu_visible_chips_per_host("v6e-4") == 4
+    assert tpu.get_num_tpu_visible_chips_per_host("v6e-1") == 1
+
+    # v5litepod multi-host defaults to 4, single-host is 8 chips
+    assert tpu.get_num_tpu_visible_chips_per_host("v5litepod-16") == 4
+    assert tpu.get_num_tpu_visible_chips_per_host("v5litepod-8") == 8
+
+    # v5litepod sub-host
+    assert tpu.get_num_tpu_visible_chips_per_host("v5litepod-4") == 4
+    assert tpu.get_num_tpu_visible_chips_per_host("v5litepod-1") == 1
+
+    # Other TPU generations default to 4
+    assert tpu.get_num_tpu_visible_chips_per_host("v4-8") == 4
+    assert tpu.get_num_tpu_visible_chips_per_host("v5p-8") == 4
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/additional_property.yaml

@@ -0,0 +1,29 @@
+# An unique identifier for the head node and workers of this cluster.
+cluster_name: minimal
+hi: 1
+# The maximum number of workers nodes to launch in addition to the head
+# node. This takes precedence over min_workers. min_workers default to 0.
+min_workers: 0
+max_workers: 0
+
+provider:
+    type: aws
+    region: us-west-2
+    availability_zone: us-west-2b
+
+# How Ray will authenticate with newly launched nodes.
+auth:
+    ssh_user: ubuntu
+
+head_node:
+    InstanceType: m4.10xlarge
+    ImageId: ami-3b6bce43  # Amazon Deep Learning AMI (Ubuntu)
+
+setup_commands:
+    - error me
+    #     - echo 'export PATH="$HOME/anaconda3/envs/tensorflow_p36/bin:$PATH"' >> ~/.bashrc
+
+# # Command to start ray on the head node. You don't need to change this.
+head_start_ray_commands:
+    - ray stop
+    - ray start --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml

python/ray/tests/authentication/BUILD.bazel

@@ -0,0 +1,23 @@
+load("@rules_python//python:defs.bzl", "py_library")
+load("//bazel:python.bzl", "py_test_run_all_subdirectory")
+
+py_library(
+    name = "conftest",
+    srcs = ["conftest.py"],
+)
+
+py_test_run_all_subdirectory(
+    size = "medium",
+    include = glob(["test_*.py"]),
+    exclude = [],
+    extra_srcs = [],
+    tags = [
+        "exclusive",
+        "medium_size_python_tests_shard_0",
+        "team:core",
+    ],
+    deps = [
+        ":conftest",
+        "//:ray_lib",
+    ],
+)

python/ray/tests/authentication/conftest.py

@@ -0,0 +1,144 @@
+import grpc
+import pytest
+from grpc import aio as aiogrpc
+
+from ray._private.authentication.authentication_token_generator import (
+    generate_new_authentication_token,
+)
+from ray._private.authentication_test_utils import (
+    authentication_env_guard,
+    reset_auth_token_state,
+    set_auth_mode,
+    set_env_auth_token,
+)
+from ray._private.grpc_utils import create_grpc_server_with_interceptors
+from ray.core.generated import reporter_pb2, reporter_pb2_grpc
+
+
+class SyncReporterService(reporter_pb2_grpc.ReporterServiceServicer):
+    """Simple synchronous test service for testing auth interceptors."""
+
+    def HealthCheck(self, request, context):
+        """Simple health check endpoint."""
+        return reporter_pb2.HealthCheckReply()
+
+
+class AsyncReporterService(reporter_pb2_grpc.ReporterServiceServicer):
+    """Simple asynchronous test service for testing auth interceptors."""
+
+    async def HealthCheck(self, request, context):
+        """Simple health check endpoint (async version)."""
+        return reporter_pb2.HealthCheckReply()
+
+
+class SyncLogService(reporter_pb2_grpc.LogServiceServicer):
+    """Simple synchronous log service for testing streaming auth interceptors."""
+
+    def StreamLog(self, request, context):
+        """Streaming log endpoint - yields test data chunks."""
+        for i in range(3):
+            yield reporter_pb2.StreamLogReply(data=f"chunk{i}".encode())
+
+
+class AsyncLogService(reporter_pb2_grpc.LogServiceServicer):
+    """Simple asynchronous log service for testing streaming auth interceptors."""
+
+    async def StreamLog(self, request, context):
+        """Streaming log endpoint (async version) - yields test data chunks."""
+        for i in range(3):
+            yield reporter_pb2.StreamLogReply(data=f"chunk{i}".encode())
+
+
+def _create_test_server_base(
+    *,
+    asynchronous: bool,
+    with_auth: bool,
+    reporter_servicer_cls,
+    log_servicer_cls,
+):
+    """Internal helper to create sync or async test server with optional auth."""
+
+    if with_auth:
+        # Auth is enabled - server will use interceptor
+        server = create_grpc_server_with_interceptors(
+            max_workers=None if asynchronous else 10,
+            thread_name_prefix="test_server",
+            options=None,
+            asynchronous=asynchronous,
+        )
+    else:
+        # Auth is disabled - create server without helper (no interceptor)
+        if asynchronous:
+            server = aiogrpc.server(options=None)
+        else:
+            from concurrent import futures
+
+            server = grpc.server(
+                futures.ThreadPoolExecutor(max_workers=10),
+                options=None,
+            )
+
+    # Add test services
+    reporter_servicer = reporter_servicer_cls()
+    reporter_pb2_grpc.add_ReporterServiceServicer_to_server(reporter_servicer, server)
+
+    log_servicer = log_servicer_cls()
+    reporter_pb2_grpc.add_LogServiceServicer_to_server(log_servicer, server)
+
+    # Bind to ephemeral port
+    port = server.add_insecure_port("[::]:0")
+
+    return server, port
+
+
+@pytest.fixture
+def create_sync_test_server():
+    """Factory to create synchronous gRPC test server.
+
+    Returns a function that creates a test server and returns (server, port).
+    The server must be stopped by the caller.
+    """
+
+    def _create(with_auth=True):
+        server, port = _create_test_server_base(
+            asynchronous=False,
+            with_auth=with_auth,
+            reporter_servicer_cls=SyncReporterService,
+            log_servicer_cls=SyncLogService,
+        )
+        server.start()
+        return server, port
+
+    return _create
+
+
+@pytest.fixture
+def create_async_test_server():
+    """Factory to create asynchronous gRPC test server.
+
+    Returns an async function that creates a test server and returns (server, port).
+    The server must be stopped by the caller.
+    """
+
+    async def _create(with_auth=True):
+        server, port = _create_test_server_base(
+            asynchronous=True,
+            with_auth=with_auth,
+            reporter_servicer_cls=AsyncReporterService,
+            log_servicer_cls=AsyncLogService,
+        )
+        await server.start()
+        return server, port
+
+    return _create
+
+
+@pytest.fixture
+def setup_auth_environment():
+    """Set up authentication environment with test token."""
+    test_token = generate_new_authentication_token()
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(test_token)
+        reset_auth_token_state()
+        yield test_token

python/ray/tests/authentication/test_async_grpc_interceptors.py

@@ -0,0 +1,221 @@
+import grpc
+import pytest
+from grpc import aio as aiogrpc
+
+from ray._private.authentication.authentication_token_generator import (
+    generate_new_authentication_token,
+)
+from ray._private.authentication_test_utils import (
+    authentication_env_guard,
+    reset_auth_token_state,
+    set_auth_mode,
+    set_env_auth_token,
+)
+from ray._private.grpc_utils import init_grpc_channel
+from ray.core.generated import reporter_pb2, reporter_pb2_grpc
+
+
+@pytest.mark.asyncio
+async def test_async_server_and_client_with_valid_token(create_async_test_server):
+    """Test async server + client with matching token succeeds."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        # Create server with auth enabled
+        server, port = await create_async_test_server(with_auth=True)
+
+        try:
+            # Client with auth interceptor via init_grpc_channel
+            channel = init_grpc_channel(
+                f"localhost:{port}",
+                options=None,
+                asynchronous=True,
+            )
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+            response = await stub.HealthCheck(request, timeout=5)
+            assert response is not None
+        finally:
+            await server.stop(grace=1)
+
+
+@pytest.mark.asyncio
+async def test_async_server_and_client_with_invalid_token(create_async_test_server):
+    """Test async server + client with mismatched token fails."""
+    server_token = generate_new_authentication_token()
+    wrong_token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        # Set up server with server_token
+        set_auth_mode("token")
+        set_env_auth_token(server_token)
+        reset_auth_token_state()
+
+        server, port = await create_async_test_server(with_auth=True)
+
+        try:
+            # Create client channel and manually add wrong token to metadata
+            channel = aiogrpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+
+            # Add invalid token to metadata (not using client interceptor)
+            metadata = (("authorization", f"Bearer {wrong_token}"),)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should fail with UNAUTHENTICATED
+            with pytest.raises(grpc.RpcError) as exc_info:
+                await stub.HealthCheck(request, metadata=metadata, timeout=5)
+
+            assert exc_info.value.code() == grpc.StatusCode.UNAUTHENTICATED
+        finally:
+            await server.stop(grace=1)
+
+
+@pytest.mark.asyncio
+async def test_async_server_with_auth_client_without_token(create_async_test_server):
+    """Test async server with auth, client without token fails."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        # Set up server with auth enabled
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        server, port = await create_async_test_server(with_auth=True)
+
+        try:
+            # Create channel without auth metadata
+            channel = aiogrpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should fail with UNAUTHENTICATED (no metadata provided)
+            with pytest.raises(grpc.RpcError) as exc_info:
+                await stub.HealthCheck(request, timeout=5)
+
+            assert exc_info.value.code() == grpc.StatusCode.UNAUTHENTICATED
+        finally:
+            await server.stop(grace=1)
+
+
+@pytest.mark.asyncio
+async def test_async_server_without_auth(create_async_test_server):
+    """Test async server without auth allows unauthenticated requests."""
+    with authentication_env_guard():
+        # Disable auth mode
+        set_auth_mode("disabled")
+        reset_auth_token_state()
+
+        # Create server without auth
+        server, port = await create_async_test_server(with_auth=False)
+
+        try:
+            # Client without auth
+            channel = aiogrpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should succeed without auth
+            response = await stub.HealthCheck(request, timeout=5)
+            assert response is not None
+        finally:
+            await server.stop(grace=1)
+
+
+@pytest.mark.asyncio
+async def test_async_server_with_auth_disabled_allows_all(create_async_test_server):
+    """Test async server allows requests when auth mode is disabled."""
+    with authentication_env_guard():
+        # Disable auth mode globally
+        set_auth_mode("disabled")
+        reset_auth_token_state()
+
+        # Even though we call create_async_test_server with with_auth=True,
+        # the server won't enforce auth because auth mode is disabled
+        server, port = await create_async_test_server(with_auth=True)
+
+        try:
+            # Client without token
+            channel = aiogrpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should succeed because auth is disabled
+            response = await stub.HealthCheck(request, timeout=5)
+            assert response is not None
+        finally:
+            await server.stop(grace=1)
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_response_with_valid_token(create_async_test_server):
+    """Test async server streaming response (unary_stream) works with valid token."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        # Create server with auth enabled
+        server, port = await create_async_test_server(with_auth=True)
+
+        try:
+            # Client with auth interceptor via init_grpc_channel
+            channel = init_grpc_channel(
+                f"localhost:{port}",
+                options=None,
+                asynchronous=True,
+            )
+            stub = reporter_pb2_grpc.LogServiceStub(channel)
+            request = reporter_pb2.StreamLogRequest(log_file_name="test.log")
+
+            # Stream the response - this tests the unary_stream RPC path
+            chunks = []
+            async for response in stub.StreamLog(request, timeout=5):
+                chunks.append(response.data)
+
+            # Verify we got all 3 chunks from the test service
+            assert len(chunks) == 3
+            assert chunks == [b"chunk0", b"chunk1", b"chunk2"]
+        finally:
+            await server.stop(grace=1)
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_response_without_token_fails(create_async_test_server):
+    """Test async server streaming response fails without token."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        server, port = await create_async_test_server(with_auth=True)
+
+        try:
+            # Client without auth token
+            channel = aiogrpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.LogServiceStub(channel)
+            request = reporter_pb2.StreamLogRequest(log_file_name="test.log")
+
+            # Should fail with UNAUTHENTICATED when trying to iterate
+            with pytest.raises(grpc.RpcError) as exc_info:
+                async for _ in stub.StreamLog(request, timeout=5):
+                    pass
+
+            assert exc_info.value.code() == grpc.StatusCode.UNAUTHENTICATED
+        finally:
+            await server.stop(grace=1)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-vv", __file__]))

python/ray/tests/authentication/test_sync_grpc_interceptors.py

@@ -0,0 +1,213 @@
+import grpc
+import pytest
+
+from ray._private.authentication.authentication_token_generator import (
+    generate_new_authentication_token,
+)
+from ray._private.authentication_test_utils import (
+    authentication_env_guard,
+    reset_auth_token_state,
+    set_auth_mode,
+    set_env_auth_token,
+)
+from ray._private.grpc_utils import init_grpc_channel
+from ray.core.generated import reporter_pb2, reporter_pb2_grpc
+
+
+def test_sync_server_and_client_with_valid_token(create_sync_test_server):
+    """Test sync server + client with matching token succeeds."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        # Create server with auth enabled
+        server, port = create_sync_test_server(with_auth=True)
+
+        try:
+            # Client with auth interceptor via init_grpc_channel
+            channel = init_grpc_channel(
+                f"localhost:{port}",
+                options=None,
+                asynchronous=False,
+            )
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+            response = stub.HealthCheck(request, timeout=5)
+            assert response is not None
+        finally:
+            server.stop(grace=1)
+
+
+def test_sync_server_and_client_with_invalid_token(create_sync_test_server):
+    """Test sync server + client with mismatched token fails."""
+    server_token = generate_new_authentication_token()
+    wrong_token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        # Set up server with server_token
+        set_auth_mode("token")
+        set_env_auth_token(server_token)
+        reset_auth_token_state()
+
+        server, port = create_sync_test_server(with_auth=True)
+
+        try:
+            # Create client channel and manually add wrong token to metadata
+            channel = grpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+
+            # Add invalid token to metadata (not using client interceptor)
+            metadata = (("authorization", f"Bearer {wrong_token}"),)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should fail with UNAUTHENTICATED
+            with pytest.raises(grpc.RpcError) as exc_info:
+                stub.HealthCheck(request, metadata=metadata, timeout=5)
+
+            assert exc_info.value.code() == grpc.StatusCode.UNAUTHENTICATED
+        finally:
+            server.stop(grace=1)
+
+
+def test_sync_server_with_auth_client_without_token(create_sync_test_server):
+    """Test server with auth, client without token fails."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        # Set up server with auth enabled
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        server, port = create_sync_test_server(with_auth=True)
+
+        try:
+            # Create channel without auth metadata
+            channel = grpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should fail with UNAUTHENTICATED (no metadata provided)
+            with pytest.raises(grpc.RpcError) as exc_info:
+                stub.HealthCheck(request, timeout=5)
+
+            assert exc_info.value.code() == grpc.StatusCode.UNAUTHENTICATED
+        finally:
+            server.stop(grace=1)
+
+
+def test_sync_server_without_auth(create_sync_test_server):
+    """Test server without auth allows unauthenticated requests."""
+    with authentication_env_guard():
+        # Disable auth mode
+        set_auth_mode("disabled")
+        reset_auth_token_state()
+
+        # Create server without auth
+        server, port = create_sync_test_server(with_auth=False)
+
+        try:
+            # Client without auth
+            channel = grpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should succeed without auth
+            response = stub.HealthCheck(request, timeout=5)
+            assert response is not None
+        finally:
+            server.stop(grace=1)
+
+
+def test_sync_server_with_auth_disabled_allows_all(create_sync_test_server):
+    """Test server allows requests when auth mode is disabled."""
+    with authentication_env_guard():
+        # Disable auth mode globally
+        set_auth_mode("disabled")
+        reset_auth_token_state()
+
+        # Even though we call create_sync_test_server with with_auth=True,
+        # the server won't enforce auth because auth mode is disabled
+        server, port = create_sync_test_server(with_auth=True)
+
+        try:
+            # Client without token
+            channel = grpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should succeed because auth is disabled
+            response = stub.HealthCheck(request, timeout=5)
+            assert response is not None
+        finally:
+            server.stop(grace=1)
+
+
+def test_sync_streaming_response_with_valid_token(create_sync_test_server):
+    """Test sync server streaming response (unary_stream) works with valid token."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        # Create server with auth enabled
+        server, port = create_sync_test_server(with_auth=True)
+
+        try:
+            # Client with auth interceptor via init_grpc_channel
+            channel = init_grpc_channel(
+                f"localhost:{port}",
+                options=None,
+                asynchronous=False,
+            )
+            stub = reporter_pb2_grpc.LogServiceStub(channel)
+            request = reporter_pb2.StreamLogRequest(log_file_name="test.log")
+
+            # Stream the response - this tests the unary_stream RPC path
+            chunks = []
+            for response in stub.StreamLog(request, timeout=5):
+                chunks.append(response.data)
+
+            # Verify we got all 3 chunks from the test service
+            assert len(chunks) == 3
+            assert chunks == [b"chunk0", b"chunk1", b"chunk2"]
+        finally:
+            server.stop(grace=1)
+
+
+def test_sync_streaming_response_without_token_fails(create_sync_test_server):
+    """Test sync server streaming response fails without token."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        server, port = create_sync_test_server(with_auth=True)
+
+        try:
+            # Client without auth token
+            channel = grpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.LogServiceStub(channel)
+            request = reporter_pb2.StreamLogRequest(log_file_name="test.log")
+
+            # Should fail with UNAUTHENTICATED when trying to iterate
+            with pytest.raises(grpc.RpcError) as exc_info:
+                for _ in stub.StreamLog(request, timeout=5):
+                    pass
+
+            assert exc_info.value.code() == grpc.StatusCode.UNAUTHENTICATED
+        finally:
+            server.stop(grace=1)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-vv", __file__]))

python/ray/tests/autoscaler/test_providers.py

@@ -0,0 +1,38 @@
+import unittest
+
+import yaml
+
+from ray.autoscaler._private.providers import (
+    _DEFAULT_CONFIGS,
+    _NODE_PROVIDERS,
+    _PROVIDER_PRETTY_NAMES,
+)
+
+
+class TestProviders(unittest.TestCase):
+    def test_node_providers(self):
+        for provider_name, provider_cls in _NODE_PROVIDERS.items():
+            config = {"module": "ray.autoscaler._private"}
+
+            try:
+                provider_cls(config)
+            except ImportError as e:
+                if f"ray.autoscaler.{provider_name}" in str(e):
+                    self.fail(
+                        f"Unexpected import error for provider {provider_name}: {e}"
+                    )
+
+    def test_provider_pretty_names(self):
+        self.assertEqual(
+            set(_NODE_PROVIDERS.keys()), set(_PROVIDER_PRETTY_NAMES.keys())
+        )
+
+    def test_default_configs(self):
+        for config_loader in _DEFAULT_CONFIGS.values():
+            config_path = config_loader()
+            with open(config_path) as f:
+                yaml.safe_load(f)
+
+
+if __name__ == "__main__":
+    unittest.main()

python/ray/tests/autoscaler/util.py

@@ -0,0 +1,68 @@
+import unittest
+from unittest.mock import Mock
+
+from ray.autoscaler._private.util import get_per_node_breakdown_as_dict
+
+
+class TestGetPerNodeBreakdown(unittest.TestCase):
+    def setUp(self):
+        # Create a mock LoadMetricsSummary object with the required attributes
+        lm_summary_mock_data = {
+            "e9919752e5e8d757765d97d8bec910a2e78e8826f20bce46fd58f92e": {
+                "node:172.31.6.57": [0.0, 1.0],
+                "object_store_memory": [0.0, 13984228147.0],
+                "memory": [0.0, 27968456295.0],
+                "node:__internal_head__": [0.0, 1.0],
+                "CPU": [0.0, 8.0],
+            }
+        }
+        self.lm_summary_mock = Mock()
+        self.lm_summary_mock.usage_by_node = lm_summary_mock_data
+
+    def test_get_per_node_breakdown_as_dict(self):
+        result = get_per_node_breakdown_as_dict(self.lm_summary_mock)
+
+        expected_output = {
+            "e9919752e5e8d757765d97d8bec910a2e78e8826f20bce46fd58f92e": (
+                "0.0/8.0 CPU\n0B/26.05GiB memory\n0B/13.02GiB object_store_memory"
+            )
+        }
+
+        self.assertEqual(result, expected_output)
+
+    def test_get_per_node_breakdown_as_dict_empty_summary(self):
+        # Test with an empty lm_summary
+        lm_summary_mock_data = {}
+        self.lm_summary_mock.usage_by_node = lm_summary_mock_data
+
+        result = get_per_node_breakdown_as_dict(self.lm_summary_mock)
+
+        expected_output = {}
+
+        self.assertEqual(result, expected_output)
+
+    def test_get_per_node_breakdown_as_dict_missing_usage(self):
+        # Test with missing usage data for a node
+        lm_summary_mock_data = {
+            "e9919752e5e8d757765d97d8bec910a2e78e8826f20bce46fd58f92e": {
+                "node:172.31.6.57": [0.0, 1.0],
+                "object_store_memory": [0.0, 13984228147.0],
+                # 'memory': [0.0, 27968456295.0],  # Missing memory data
+                "node:__internal_head__": [0.0, 1.0],
+                "CPU": [0.0, 8.0],
+            }
+        }
+        self.lm_summary_mock.usage_by_node = lm_summary_mock_data
+
+        result = get_per_node_breakdown_as_dict(self.lm_summary_mock)
+
+        expected_output = {
+            "e9919752e5e8d757765d97d8bec910a2e78e8826f20bce46fd58f92e": "0.0/8.0 CPU\n"
+            "0B/13.02GiB object_store_memory"
+        }
+
+        self.assertEqual(result, expected_output)
+
+
+if __name__ == "__main__":
+    unittest.main()

python/ray/tests/autoscaler_test_utils.py

@@ -0,0 +1,312 @@
+import re
+import threading
+from subprocess import CalledProcessError
+from typing import Any, Dict, List, Optional
+
+from ray.autoscaler.node_provider import NodeProvider
+
+
+class MockNode:
+    def __init__(
+        self,
+        node_id,
+        tags,
+        node_config,
+        node_type,
+        unique_ips=False,
+        resources=None,
+        labels=None,
+    ):
+        self.node_id = str(node_id)
+        self.state = "pending"
+        self.tags = tags
+        self.external_ip = "1.2.3.4"
+        self.internal_ip = "172.0.0.{}".format(self.node_id)
+        if unique_ips:
+            self.external_ip = f"1.2.3.{self.node_id}"
+
+        self.node_config = node_config
+        self.node_type = node_type
+        self.created_in_main_thread = (
+            threading.current_thread() is threading.main_thread()
+        )
+        self.resources = resources or {}
+        self.labels = labels or {}
+
+    def matches(self, tags):
+        for k, v in tags.items():
+            if k not in self.tags or self.tags[k] != v:
+                return False
+        return True
+
+
+class MockProcessRunner:
+    def __init__(self, fail_cmds=None, cmd_to_callback=None, print_out=False):
+        self.calls = []
+        self.cmd_to_callback = cmd_to_callback or {}  # type: Dict[str, Callable]
+        self.print_out = print_out
+        self.fail_cmds = fail_cmds or []
+        self.call_response = {}
+        self.ready_to_run = threading.Event()
+        self.ready_to_run.set()
+
+        self.lock = threading.RLock()
+
+    def check_call(self, cmd, *args, **kwargs):
+        with self.lock:
+            self.ready_to_run.wait()
+            self.calls.append(cmd)
+            if self.print_out:
+                print(f">>>Process runner: Executing \n {str(cmd)}")
+            for token in self.cmd_to_callback:
+                if token in str(cmd):
+                    # Trigger a callback if token is in cmd.
+                    # Can be used to simulate background events during a node
+                    # update (e.g. node disconnected).
+                    callback = self.cmd_to_callback[token]
+                    callback()
+
+            for token in self.fail_cmds:
+                if token in str(cmd):
+                    raise CalledProcessError(1, token, "Failing command on purpose")
+
+    def check_output(self, cmd):
+        with self.lock:
+            self.check_call(cmd)
+            return_string = "command-output"
+            key_to_shrink = None
+            for pattern, response_list in self.call_response.items():
+                if pattern in str(cmd):
+                    return_string = response_list[0]
+                    key_to_shrink = pattern
+                    break
+            if key_to_shrink:
+                self.call_response[key_to_shrink] = self.call_response[key_to_shrink][
+                    1:
+                ]
+                if len(self.call_response[key_to_shrink]) == 0:
+                    del self.call_response[key_to_shrink]
+
+            return return_string.encode()
+
+    def assert_has_call(
+        self, ip: str, pattern: Optional[str] = None, exact: Optional[List[str]] = None
+    ) -> bool:
+        """Checks if the given value was called by this process runner.
+
+        NOTE: Either pattern or exact must be specified, not both!
+
+        Args:
+            ip: IP address of the node that the given call was executed on.
+            pattern: RegEx that matches one specific call.
+            exact: List of strings that when joined exactly match one call.
+
+        Returns:
+            ``True`` when a matching call is found. Raises ``Exception`` when
+            no matching call is recorded.
+        """
+        with self.lock:
+            assert bool(pattern) ^ bool(
+                exact
+            ), "Must specify either a pattern or exact match."
+            debug_output = ""
+            if pattern is not None:
+                for cmd in self.command_history():
+                    if ip in cmd:
+                        debug_output += cmd
+                        debug_output += "\n"
+                        if re.search(pattern, cmd):
+                            return True
+                else:
+                    raise Exception(
+                        f"Did not find [{pattern}] in [{debug_output}] for "
+                        f"ip={ip}.\n\nFull output: {self.command_history()}"
+                    )
+            elif exact is not None:
+                exact_cmd = " ".join(exact)
+                for cmd in self.command_history():
+                    if ip in cmd:
+                        debug_output += cmd
+                        debug_output += "\n"
+                    if cmd == exact_cmd:
+                        return True
+                raise Exception(
+                    f"Did not find [{exact_cmd}] in [{debug_output}] for "
+                    f"ip={ip}.\n\nFull output: {self.command_history()}"
+                )
+
+    def assert_not_has_call(self, ip: str, pattern: str):
+        """Ensure that the given regex pattern was never called."""
+        with self.lock:
+            out = ""
+            for cmd in self.command_history():
+                if ip in cmd:
+                    out += cmd
+                    out += "\n"
+            if re.search(pattern, out):
+                raise Exception("Found [{}] in [{}] for {}".format(pattern, out, ip))
+            else:
+                return True
+
+    def clear_history(self):
+        with self.lock:
+            self.calls = []
+
+    def command_history(self):
+        with self.lock:
+            return [" ".join(cmd) for cmd in self.calls]
+
+    def respond_to_call(self, pattern, response_list):
+        with self.lock:
+            self.call_response[pattern] = response_list
+
+
+class MockProvider(NodeProvider):
+    def __init__(self, cache_stopped=False, unique_ips=False):
+        self.mock_nodes = {}
+        self.next_id = 0
+        self.throw = False
+        self.creation_error = None
+        self.termination_errors = None
+        self.fail_creates = False
+        self.ready_to_create = threading.Event()
+        self.ready_to_create.set()
+        self.cache_stopped = cache_stopped
+        self.unique_ips = unique_ips
+        self.fail_to_fetch_ip = False
+        self.safe_to_scale_flag = True
+        self.partical_success_count = None
+        # Many of these functions are called by node_launcher or updater in
+        # different threads. This can be treated as a global lock for
+        # everything.
+        self.lock = threading.Lock()
+        self.num_non_terminated_nodes_calls = 0
+        super().__init__(None, None)
+
+    def non_terminated_nodes(self, tag_filters):
+        self.num_non_terminated_nodes_calls += 1
+        with self.lock:
+            if self.throw:
+                raise Exception("oops")
+            return [
+                n.node_id
+                for n in self.mock_nodes.values()
+                if n.matches(tag_filters) and n.state not in ["stopped", "terminated"]
+            ]
+
+    def non_terminated_node_ips(self, tag_filters):
+        with self.lock:
+            if self.throw:
+                raise Exception("oops")
+            return [
+                n.internal_ip
+                for n in self.mock_nodes.values()
+                if n.matches(tag_filters) and n.state not in ["stopped", "terminated"]
+            ]
+
+    def is_running(self, node_id):
+        with self.lock:
+            return self.mock_nodes[node_id].state == "running"
+
+    def is_terminated(self, node_id):
+        if node_id is None:
+            # Circumvent test-cases where there's no head node.
+            return True
+        with self.lock:
+            return self.mock_nodes[node_id].state in ["stopped", "terminated"]
+
+    def node_tags(self, node_id):
+        if node_id is None:
+            # Circumvent test cases where there's no head node.
+            return {}
+        # Don't assume that node providers can retrieve tags from
+        # terminated nodes.
+        if self.is_terminated(node_id):
+            raise Exception(f"The node with id {node_id} has been terminated!")
+        with self.lock:
+            return self.mock_nodes[node_id].tags
+
+    def internal_ip(self, node_id):
+        if self.fail_to_fetch_ip:
+            raise Exception("Failed to fetch ip on purpose.")
+        if node_id is None:
+            # Circumvent test-cases where there's no head node.
+            return "mock"
+        with self.lock:
+            return self.mock_nodes[node_id].internal_ip
+
+    def external_ip(self, node_id):
+        with self.lock:
+            return self.mock_nodes[node_id].external_ip
+
+    def create_node(
+        self,
+        node_config: Dict[str, Any],
+        tags: Dict[str, str],
+        count: int,
+        _skip_wait=False,
+    ) -> Dict[str, Any]:
+        return self.create_node_with_resources_and_labels(
+            node_config, tags, count, {}, {}, _skip_wait=_skip_wait
+        )
+
+    def create_node_with_resources_and_labels(
+        self, node_config, tags, count, resources, labels, _skip_wait=False
+    ):
+        from ray.autoscaler.tags import TAG_RAY_USER_NODE_TYPE
+
+        if self.creation_error is not None:
+            raise self.creation_error
+        if not _skip_wait:
+            self.ready_to_create.wait()
+        if self.fail_creates:
+            return
+
+        created_nodes = {}
+        if self.partical_success_count is not None:
+            count = min(count, self.partical_success_count)
+        with self.lock:
+            if self.cache_stopped:
+                for node in self.mock_nodes.values():
+                    if node.state == "stopped" and count > 0:
+                        count -= 1
+                        node.state = "pending"
+                        node.tags.update(tags)
+                        created_nodes[node.node_id] = node
+            for _ in range(count):
+                new_node = MockNode(
+                    str(self.next_id),
+                    tags.copy(),
+                    node_config,
+                    tags.get(TAG_RAY_USER_NODE_TYPE),
+                    resources=resources,
+                    labels=labels,
+                    unique_ips=self.unique_ips,
+                )
+                self.mock_nodes[new_node.node_id] = new_node
+                created_nodes[new_node.node_id] = new_node
+                self.next_id += 1
+        return created_nodes
+
+    def set_node_tags(self, node_id, tags):
+        with self.lock:
+            self.mock_nodes[node_id].tags.update(tags)
+
+    def terminate_node(self, node_id):
+        with self.lock:
+            if self.termination_errors is not None:
+                raise self.termination_errors
+
+            if self.cache_stopped:
+                self.mock_nodes[node_id].state = "stopped"
+            else:
+                self.mock_nodes[node_id].state = "terminated"
+
+    def finish_starting_nodes(self):
+        with self.lock:
+            for node in self.mock_nodes.values():
+                if node.state == "pending":
+                    node.state = "running"
+
+    def safe_to_scale(self):
+        return self.safe_to_scale_flag

python/ray/tests/aws/conftest.py

@@ -0,0 +1,55 @@
+import pytest
+from botocore.stub import Stubber
+
+from ray.autoscaler._private.aws.utils import client_cache, resource_cache
+from ray.autoscaler._private.constants import BOTO_MAX_RETRIES
+
+
+@pytest.fixture()
+def iam_client_stub(request):
+    region = getattr(request, "param", "us-west-2")
+    resource = resource_cache("iam", region)
+    with Stubber(resource.meta.client) as stubber:
+        yield stubber
+        stubber.assert_no_pending_responses()
+
+
+@pytest.fixture()
+def ec2_client_stub(request):
+    region = getattr(request, "param", "us-west-2")
+    resource = resource_cache("ec2", region)
+    with Stubber(resource.meta.client) as stubber:
+        yield stubber
+        stubber.assert_no_pending_responses()
+
+
+@pytest.fixture()
+def ec2_client_stub_fail_fast():
+    resource = resource_cache("ec2", "us-west-2", 0)
+    with Stubber(resource.meta.client) as stubber:
+        yield stubber
+        stubber.assert_no_pending_responses()
+
+
+@pytest.fixture()
+def ec2_client_stub_max_retries():
+    resource = resource_cache("ec2", "us-west-2", BOTO_MAX_RETRIES)
+    with Stubber(resource.meta.client) as stubber:
+        yield stubber
+        stubber.assert_no_pending_responses()
+
+
+@pytest.fixture()
+def cloudwatch_client_stub():
+    resource = resource_cache("cloudwatch", "us-west-2")
+    with Stubber(resource.meta.client) as stubber:
+        yield stubber
+        stubber.assert_no_pending_responses()
+
+
+@pytest.fixture()
+def ssm_client_stub():
+    client = client_cache("ssm", "us-west-2")
+    with Stubber(client) as stubber:
+        yield stubber
+        stubber.assert_no_pending_responses()

python/ray/tests/aws/test_aws_batch_tag_update.py

@@ -0,0 +1,68 @@
+import sys
+import threading
+import time
+import unittest
+from unittest import mock
+
+import pytest
+
+from ray.autoscaler._private.aws.node_provider import TAG_BATCH_DELAY, AWSNodeProvider
+
+
+def mock_create_tags(provider, batch_updates):
+    # Increment batches sent.
+    provider.batch_counter += 1
+    # Increment tags updated.
+    provider.tag_update_counter += sum(len(batch_updates[x]) for x in batch_updates)
+
+
+def batch_test(num_threads, delay):
+    """Run AWSNodeProvider.set_node_tags in several threads, with a
+    specified delay between thread launches.
+
+    Return the number of batches of tag updates and the number of tags
+    updated.
+    """
+    with mock.patch(
+        "ray.autoscaler._private.aws.node_provider.make_ec2_resource"
+    ), mock.patch.object(AWSNodeProvider, "_create_tags", mock_create_tags):
+        provider = AWSNodeProvider(
+            provider_config={"region": "nowhere"}, cluster_name="default"
+        )
+        provider.batch_counter = 0
+        provider.tag_update_counter = 0
+        provider.tag_cache = {str(x): {} for x in range(num_threads)}
+
+        threads = []
+        for x in range(num_threads):
+            thread = threading.Thread(
+                target=provider.set_node_tags, args=(str(x), {"foo": "bar"})
+            )
+            threads.append(thread)
+
+        for thread in threads:
+            thread.start()
+            time.sleep(delay)
+        for thread in threads:
+            thread.join()
+
+        return provider.batch_counter, provider.tag_update_counter
+
+
+class TagBatchTest(unittest.TestCase):
+    def test_concurrent(self):
+        num_threads = 100
+        batches_sent, tags_updated = batch_test(num_threads, delay=0)
+        self.assertLess(batches_sent, num_threads / 10)
+        self.assertEqual(tags_updated, num_threads)
+
+    def test_serial(self):
+        num_threads = 5
+        long_delay = TAG_BATCH_DELAY * 1.2
+        batches_sent, tags_updated = batch_test(num_threads, delay=long_delay)
+        self.assertEqual(batches_sent, num_threads)
+        self.assertEqual(tags_updated, num_threads)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tests/aws/utils/constants.py

@@ -0,0 +1,228 @@
+import copy
+from datetime import datetime
+
+import ray
+from ray.autoscaler.tags import (
+    NODE_KIND_HEAD,
+    TAG_RAY_LAUNCH_CONFIG,
+    TAG_RAY_NODE_KIND,
+    TAG_RAY_USER_NODE_TYPE,
+)
+
+# Override global constants used in AWS autoscaler config artifact names.
+# This helps ensure that any unmocked test doesn't alter non-test artifacts.
+ray.autoscaler._private.aws.config.RAY = "ray-autoscaler-aws-test"
+ray.autoscaler._private.aws.config.DEFAULT_RAY_INSTANCE_PROFILE = (
+    ray.autoscaler._private.aws.config.RAY + "-v1"
+)
+ray.autoscaler._private.aws.config.DEFAULT_RAY_IAM_ROLE = (
+    ray.autoscaler._private.aws.config.RAY + "-v1"
+)
+ray.autoscaler._private.aws.config.SECURITY_GROUP_TEMPLATE = (
+    ray.autoscaler._private.aws.config.RAY + "-{}"
+)
+
+# Default IAM instance profile to expose to tests.
+DEFAULT_INSTANCE_PROFILE = {
+    "Arn": "arn:aws:iam::336924118301:instance-profile/ExampleInstanceProfile",
+    "CreateDate": datetime(2013, 6, 12, 23, 52, 2, 2),
+    "InstanceProfileId": "AIPA0000000000EXAMPLE",
+    "InstanceProfileName": "ExampleInstanceProfile",
+    "Path": "/",
+    "Roles": [
+        {
+            "Arn": "arn:aws:iam::123456789012:role/Test-Role",
+            "AssumeRolePolicyDocument": "ExampleAssumeRolePolicyDocument",
+            "CreateDate": datetime(2013, 1, 9, 6, 33, 26, 2),
+            "Path": "/",
+            "RoleId": "AROA0000000000EXAMPLE",
+            "RoleName": "Test-Role",
+        },
+    ],
+}
+
+# Default EC2 key pair to expose to tests.
+DEFAULT_KEY_PAIR = {
+    "KeyFingerprint": "00:11:22:33:44:55:66:77:88:99:AA:BB:CC:DD:EE:FF:00",
+    "KeyName": ray.autoscaler._private.aws.config.RAY + "_us-west-2",
+}
+
+# Primary EC2 subnet to expose to tests.
+DEFAULT_SUBNET = {
+    "AvailabilityZone": "us-west-2a",
+    "AvailableIpAddressCount": 251,
+    "CidrBlock": "10.0.1.0/24",
+    "DefaultForAz": False,
+    "MapPublicIpOnLaunch": True,
+    "State": "available",
+    "SubnetId": "subnet-0000000",
+    "VpcId": "vpc-0000000",
+}
+
+
+def subnet_in_vpc(vpc_num):
+    """Returns a copy of DEFAULT_SUBNET whose VpcId ends with the digits
+    of vpc_num."""
+    subnet = copy.copy(DEFAULT_SUBNET)
+    subnet["VpcId"] = f"vpc-{vpc_num:07d}"
+    return subnet
+
+
+A_THOUSAND_SUBNETS_IN_DIFFERENT_VPCS = [
+    subnet_in_vpc(vpc_num) for vpc_num in range(1, 1000)
+] + [DEFAULT_SUBNET]
+
+
+def subnet_in_az(idx):
+    azs = ["a", "b", "c", "d"]
+    subnet = copy.copy(DEFAULT_SUBNET)
+    subnet["AvailabilityZone"] = "us-west-2" + azs[idx % 4]
+    subnet["SubnetId"] = f"subnet-{idx:07d}"
+    return subnet
+
+
+TWENTY_SUBNETS_IN_DIFFERENT_AZS = [subnet_in_az(i) for i in range(20)]
+
+# Secondary EC2 subnet to expose to tests as required.
+AUX_SUBNET = {
+    "AvailabilityZone": "us-west-2a",
+    "AvailableIpAddressCount": 251,
+    "CidrBlock": "192.168.1.0/24",
+    "DefaultForAz": False,
+    "MapPublicIpOnLaunch": True,
+    "State": "available",
+    "SubnetId": "subnet-fffffff",
+    "VpcId": "vpc-fffffff",
+}
+
+# Default cluster name to expose to tests.
+DEFAULT_CLUSTER_NAME = "test-cluster-name"
+
+# Default security group settings immediately after creation
+# (prior to inbound rule configuration).
+DEFAULT_SG = {
+    "Description": "Auto-created security group for Ray workers",
+    "GroupName": ray.autoscaler._private.aws.config.RAY + "-" + DEFAULT_CLUSTER_NAME,
+    "OwnerId": "test-owner",
+    "GroupId": "sg-1234abcd",
+    "VpcId": DEFAULT_SUBNET["VpcId"],
+    "IpPermissions": [],
+    "IpPermissionsEgress": [
+        {
+            "FromPort": -1,
+            "ToPort": -1,
+            "IpProtocol": "-1",
+            "IpRanges": [{"CidrIp": "0.0.0.0/0"}],
+        }
+    ],
+    "Tags": [],
+}
+
+# Secondary security group settings after creation
+# (prior to inbound rule configuration).
+AUX_SG = copy.deepcopy(DEFAULT_SG)
+AUX_SG["GroupName"] += "-aux"
+AUX_SG["GroupId"] = "sg-dcba4321"
+
+# Default security group settings immediately after creation on aux subnet
+# (prior to inbound rule configuration).
+DEFAULT_SG_AUX_SUBNET = copy.deepcopy(DEFAULT_SG)
+DEFAULT_SG_AUX_SUBNET["VpcId"] = AUX_SUBNET["VpcId"]
+DEFAULT_SG_AUX_SUBNET["GroupId"] = AUX_SG["GroupId"]
+
+DEFAULT_IN_BOUND_RULES = [
+    {
+        "FromPort": -1,
+        "ToPort": -1,
+        "IpProtocol": "-1",
+        "UserIdGroupPairs": [{"GroupId": DEFAULT_SG["GroupId"]}],
+    },
+    {
+        "FromPort": 22,
+        "ToPort": 22,
+        "IpProtocol": "tcp",
+        "IpRanges": [{"CidrIp": "0.0.0.0/0"}],
+    },
+]
+# Default security group settings once default inbound rules are applied
+# (if used by both head and worker nodes)
+DEFAULT_SG_WITH_RULES = copy.deepcopy(DEFAULT_SG)
+DEFAULT_SG_WITH_RULES["IpPermissions"] = DEFAULT_IN_BOUND_RULES
+
+# Default security group once default inbound rules are applied
+# (if using separate security groups for head and worker nodes).
+DEFAULT_SG_DUAL_GROUP_RULES = copy.deepcopy(DEFAULT_SG_WITH_RULES)
+DEFAULT_SG_DUAL_GROUP_RULES["IpPermissions"][0]["UserIdGroupPairs"].append(
+    {"GroupId": AUX_SG["GroupId"]}
+)
+
+# Default security group on aux subnet once default inbound rules are applied.
+DEFAULT_SG_WITH_RULES_AUX_SUBNET = copy.deepcopy(DEFAULT_SG_DUAL_GROUP_RULES)
+DEFAULT_SG_WITH_RULES_AUX_SUBNET["VpcId"] = AUX_SUBNET["VpcId"]
+DEFAULT_SG_WITH_RULES_AUX_SUBNET["GroupId"] = AUX_SG["GroupId"]
+
+# Default security group with custom name
+DEFAULT_SG_WITH_NAME = copy.deepcopy(DEFAULT_SG)
+DEFAULT_SG_WITH_NAME["GroupName"] = "test_security_group_name"
+
+CUSTOM_IN_BOUND_RULES = [
+    {
+        "FromPort": 443,
+        "ToPort": 443,
+        "IpProtocol": "TCP",
+        "IpRanges": [{"CidrIp": "0.0.0.0/0"}],
+    },
+    {
+        "FromPort": 8265,
+        "ToPort": 8265,
+        "IpProtocol": "TCP",
+        "IpRanges": [{"CidrIp": "0.0.0.0/0"}],
+    },
+]
+
+# Default security group with custom name once...
+# default and custom in bound rules are applied
+DEFAULT_SG_WITH_NAME_AND_RULES = copy.deepcopy(DEFAULT_SG_WITH_NAME)
+DEFAULT_SG_WITH_NAME_AND_RULES["IpPermissions"] = (
+    DEFAULT_IN_BOUND_RULES + CUSTOM_IN_BOUND_RULES
+)
+
+# Default launch template to expose to tests.
+DEFAULT_LT = {
+    "LaunchTemplateId": "lt-00000000000000000",
+    "LaunchTemplateName": "ExampleLaunchTemplate",
+    "VersionNumber": 2,
+    "CreateTime": datetime(2020, 8, 17, 23, 30, 3),
+    "CreatedBy": DEFAULT_INSTANCE_PROFILE["Roles"][0]["Arn"],
+    "DefaultVersion": True,
+    "LaunchTemplateData": {
+        "EbsOptimized": False,
+        "IamInstanceProfile": {"Arn": DEFAULT_INSTANCE_PROFILE["Arn"]},
+        "NetworkInterfaces": [
+            {
+                "DeviceIndex": 0,
+                "Groups": [DEFAULT_SG["GroupId"]],
+                "SubnetId": DEFAULT_SUBNET["SubnetId"],
+            }
+        ],
+        "ImageId": "ami-00000000000000000",
+        "InstanceType": "m5.large",
+        "TagSpecifications": [
+            {
+                "ResourceType": "instance",
+                "Tags": [{"Key": "test-key", "Value": "test-value"}],
+            },
+            {
+                "ResourceType": "volume",
+                "Tags": [{"Key": "test-key", "Value": "test-value"}],
+            },
+        ],
+    },
+}
+
+# Default node provider tags to expose to tests.
+DEFAULT_NODE_PROVIDER_INSTANCE_TAGS = {
+    TAG_RAY_NODE_KIND: NODE_KIND_HEAD,
+    TAG_RAY_LAUNCH_CONFIG: "test-ray-launch-config",
+    TAG_RAY_USER_NODE_TYPE: "ray.head.default",
+}

python/ray/tests/aws/utils/helpers.py

@@ -0,0 +1,150 @@
+import copy
+import os
+from typing import Any, Dict
+
+import yaml
+
+import ray
+from ray.autoscaler._private.aws.cloudwatch.cloudwatch_helper import CloudwatchHelper
+from ray.autoscaler._private.aws.node_provider import AWSNodeProvider
+from ray.autoscaler._private.commands import prepare_config, validate_config
+from ray.autoscaler.tags import (
+    NODE_KIND_HEAD,
+    NODE_KIND_WORKER,
+    TAG_RAY_CLUSTER_NAME,
+    TAG_RAY_NODE_KIND,
+    TAG_RAY_USER_NODE_TYPE,
+)
+from ray.tests.aws.utils.constants import (
+    DEFAULT_CLUSTER_NAME,
+    DEFAULT_NODE_PROVIDER_INSTANCE_TAGS,
+)
+
+
+def get_aws_example_config_file_path(file_name):
+    import ray.autoscaler.aws
+
+    return os.path.join(os.path.dirname(ray.autoscaler.aws.__file__), file_name)
+
+
+def load_aws_example_config_file(file_name):
+    config_file_path = get_aws_example_config_file_path(file_name)
+    return yaml.safe_load(open(config_file_path).read())
+
+
+def fake_fillout_available_node_types_resources(config: Dict[str, Any]) -> None:
+    """A cheap way to fill out the resources field (the same way a node
+    provider would autodetect them) as far as schema validation is concerned."""
+    available_node_types = config.get("available_node_types", {})
+    for label, value in available_node_types.items():
+        value["resources"] = value.get("resources", {"filler": 1})
+
+
+def bootstrap_aws_config(config):
+    config = prepare_config(config)
+    fake_fillout_available_node_types_resources(config)
+    validate_config(config)
+    config["cluster_name"] = DEFAULT_CLUSTER_NAME
+    return ray.autoscaler._private.aws.config.bootstrap_aws(config)
+
+
+def bootstrap_aws_example_config_file(file_name):
+    config = load_aws_example_config_file(file_name)
+    return bootstrap_aws_config(config)
+
+
+def node_provider_tags(config: dict, type_name: str) -> dict:
+    """
+    Returns a copy of DEFAULT_NODE_PROVIDER_INSTANCE_TAGS with the Ray node
+    kind and Ray user node type filled in from the input config and node type
+    name.
+
+    Args:
+        config: autoscaler config
+        type_name: node type name
+    Returns:
+        tags: node provider tags
+    """
+    tags = copy.copy(DEFAULT_NODE_PROVIDER_INSTANCE_TAGS)
+    head_name = config["head_node_type"]
+    node_kind = NODE_KIND_HEAD if type_name is head_name else NODE_KIND_WORKER
+    tags[TAG_RAY_NODE_KIND] = node_kind
+    tags[TAG_RAY_USER_NODE_TYPE] = type_name
+    return tags
+
+
+def apply_node_provider_config_updates(
+    config: dict, node_cfg: dict, node_type_name: str, max_count: int
+) -> None:
+    """
+    Applies default updates made by AWSNodeProvider to node_cfg during node
+    creation. This should only be used for testing purposes.
+
+    Args:
+        config: autoscaler config
+        node_cfg: node config
+        node_type_name: node type name
+        max_count: max nodes of the given type to launch
+    """
+    tags = node_provider_tags(config, node_type_name)
+    tags[TAG_RAY_CLUSTER_NAME] = DEFAULT_CLUSTER_NAME
+    user_tag_specs = node_cfg.get("TagSpecifications", [])
+    tag_specs = [
+        {
+            "ResourceType": "instance",
+            "Tags": [{"Key": k, "Value": v} for k, v in sorted(tags.items())],
+        }
+    ]
+    node_provider_cfg_updates = {
+        "MinCount": 1,
+        "MaxCount": max_count,
+        "TagSpecifications": tag_specs,
+    }
+    tags.pop(TAG_RAY_CLUSTER_NAME)
+    node_cfg.update(node_provider_cfg_updates)
+    # merge node provider tag specs with user overrides
+    AWSNodeProvider._merge_tag_specs(tag_specs, user_tag_specs)
+
+
+def get_cloudwatch_agent_config_file_path():
+    return get_aws_example_config_file_path(
+        "cloudwatch/example-cloudwatch-agent-config.json"
+    )
+
+
+def get_cloudwatch_dashboard_config_file_path():
+    return get_aws_example_config_file_path(
+        "cloudwatch/example-cloudwatch-dashboard-config.json"
+    )
+
+
+def get_cloudwatch_alarm_config_file_path():
+    return get_aws_example_config_file_path(
+        "cloudwatch/example-cloudwatch-alarm-config.json"
+    )
+
+
+def load_cloudwatch_example_config_file():
+    config = load_aws_example_config_file("example-cloudwatch.yaml")
+    cw_cfg = config["provider"]["cloudwatch"]
+    cw_cfg["agent"]["config"] = get_cloudwatch_agent_config_file_path()
+    cw_cfg["dashboard"]["config"] = get_cloudwatch_dashboard_config_file_path()
+    cw_cfg["alarm"]["config"] = get_cloudwatch_alarm_config_file_path()
+    return config
+
+
+def get_cloudwatch_helper(node_ids):
+    config = load_cloudwatch_example_config_file()
+    config["cluster_name"] = DEFAULT_CLUSTER_NAME
+    return CloudwatchHelper(
+        config["provider"],
+        node_ids,
+        config["cluster_name"],
+    )
+
+
+def get_ssm_param_name(cluster_name, config_type):
+    ssm_config_param_name = "AmazonCloudWatch-" + "ray_{}_config_{}".format(
+        config_type, cluster_name
+    )
+    return ssm_config_param_name

python/ray/tests/aws/utils/mocks.py

@@ -0,0 +1,7 @@
+from ray.autoscaler._private.aws.config import key_pair
+from ray.tests.aws.utils.constants import DEFAULT_KEY_PAIR
+
+
+def mock_path_exists_key_pair(path):
+    key_name, key_path = key_pair(0, "us-west-2", DEFAULT_KEY_PAIR["KeyName"])
+    return path == key_path

python/ray/tests/aws/utils/stubs.py

@@ -0,0 +1,584 @@
+import copy
+import json
+from typing import Dict, List
+from unittest import mock
+from uuid import uuid4
+
+from botocore.stub import ANY
+
+import ray
+from ray.autoscaler._private.aws.cloudwatch.cloudwatch_helper import (
+    CLOUDWATCH_AGENT_INSTALLED_TAG,
+    CLOUDWATCH_CONFIG_HASH_TAG_BASE,
+)
+from ray.autoscaler._private.aws.config import key_pair
+from ray.autoscaler.tags import NODE_KIND_HEAD, TAG_RAY_NODE_KIND
+from ray.tests.aws.utils import helpers
+from ray.tests.aws.utils.constants import (
+    A_THOUSAND_SUBNETS_IN_DIFFERENT_VPCS,
+    DEFAULT_CLUSTER_NAME,
+    DEFAULT_INSTANCE_PROFILE,
+    DEFAULT_KEY_PAIR,
+    DEFAULT_LT,
+    DEFAULT_SUBNET,
+    TWENTY_SUBNETS_IN_DIFFERENT_AZS,
+)
+from ray.tests.aws.utils.helpers import (
+    get_cloudwatch_alarm_config_file_path,
+    get_cloudwatch_dashboard_config_file_path,
+)
+
+
+def configure_iam_role_default(iam_client_stub):
+    iam_client_stub.add_response(
+        "get_instance_profile",
+        expected_params={
+            "InstanceProfileName": ray.autoscaler._private.aws.config.DEFAULT_RAY_INSTANCE_PROFILE  # noqa: E501
+        },
+        service_response={"InstanceProfile": DEFAULT_INSTANCE_PROFILE},
+    )
+
+
+def configure_key_pair_default(
+    ec2_client_stub, region="us-west-2", expected_key_pair=DEFAULT_KEY_PAIR
+):
+    patcher = mock.patch("os.path.exists")
+
+    def mock_path_exists_key_pair(path):
+        _, key_path = key_pair(0, region, expected_key_pair["KeyName"])
+        return path == key_path
+
+    os_path_exists_mock = patcher.start()
+    os_path_exists_mock.side_effect = mock_path_exists_key_pair
+
+    ec2_client_stub.add_response(
+        "describe_key_pairs",
+        expected_params={
+            "Filters": [{"Name": "key-name", "Values": [expected_key_pair["KeyName"]]}]
+        },
+        service_response={"KeyPairs": [expected_key_pair]},
+    )
+
+
+def configure_subnet_default(ec2_client_stub):
+    ec2_client_stub.add_response(
+        "describe_subnets",
+        expected_params={},
+        service_response={"Subnets": [DEFAULT_SUBNET]},
+    )
+
+
+def describe_a_thousand_subnets_in_different_vpcs(ec2_client_stub):
+    ec2_client_stub.add_response(
+        "describe_subnets",
+        expected_params={},
+        service_response={"Subnets": A_THOUSAND_SUBNETS_IN_DIFFERENT_VPCS},
+    )
+
+
+def describe_twenty_subnets_in_different_azs(ec2_client_stub):
+    ec2_client_stub.add_response(
+        "describe_subnets",
+        expected_params={},
+        service_response={"Subnets": TWENTY_SUBNETS_IN_DIFFERENT_AZS},
+    )
+
+
+def skip_to_configure_sg(ec2_client_stub, iam_client_stub):
+    configure_iam_role_default(iam_client_stub)
+    configure_key_pair_default(ec2_client_stub)
+    configure_subnet_default(ec2_client_stub)
+
+
+def describe_subnets_echo(ec2_client_stub, subnets: List[Dict[str, str]]):
+    ec2_client_stub.add_response(
+        "describe_subnets",
+        expected_params={
+            "Filters": [
+                {"Name": "subnet-id", "Values": [s["SubnetId"] for s in subnets]}
+            ]
+        },
+        service_response={"Subnets": subnets},
+    )
+
+
+def describe_no_security_groups(ec2_client_stub):
+    ec2_client_stub.add_response(
+        "describe_security_groups",
+        expected_params={"Filters": ANY},
+        service_response={},
+    )
+
+
+def describe_a_security_group(ec2_client_stub, security_group):
+    ec2_client_stub.add_response(
+        "describe_security_groups",
+        expected_params={
+            "Filters": [{"Name": "group-id", "Values": [security_group["GroupId"]]}]
+        },
+        service_response={"SecurityGroups": [security_group]},
+    )
+
+
+def describe_an_sg_2(ec2_client_stub, security_group):
+    """Same as last function, different input param format.
+
+    A call with this input parameter format is made when sg.ip_permissions is
+    accessed in aws/config.py.
+    """
+    ec2_client_stub.add_response(
+        "describe_security_groups",
+        expected_params={"GroupIds": [security_group["GroupId"]]},
+        service_response={"SecurityGroups": [security_group]},
+    )
+
+
+def create_sg_echo(ec2_client_stub, security_group):
+    ec2_client_stub.add_response(
+        "create_security_group",
+        expected_params={
+            "Description": security_group["Description"],
+            "GroupName": security_group["GroupName"],
+            "VpcId": security_group["VpcId"],
+            "TagSpecifications": [
+                {
+                    "ResourceType": "security-group",
+                    "Tags": [
+                        {
+                            "Key": ray.autoscaler._private.aws.config.RAY,
+                            "Value": "true",
+                        },
+                        {"Key": "ray-cluster-name", "Value": DEFAULT_CLUSTER_NAME},
+                    ],
+                },
+            ],
+        },
+        service_response={"GroupId": security_group["GroupId"]},
+    )
+
+
+def describe_sgs_by_id(ec2_client_stub, security_group_ids, security_groups):
+    ec2_client_stub.add_response(
+        "describe_security_groups",
+        expected_params={
+            "Filters": [{"Name": "group-id", "Values": security_group_ids}]
+        },
+        service_response={"SecurityGroups": security_groups},
+    )
+
+
+def describe_sgs_on_vpc(ec2_client_stub, vpc_ids, security_groups):
+    ec2_client_stub.add_response(
+        "describe_security_groups",
+        expected_params={"Filters": [{"Name": "vpc-id", "Values": vpc_ids}]},
+        service_response={"SecurityGroups": security_groups},
+    )
+
+
+def authorize_sg_ingress(ec2_client_stub, security_group):
+    ec2_client_stub.add_response(
+        "authorize_security_group_ingress",
+        expected_params={
+            "GroupId": security_group["GroupId"],
+            "IpPermissions": security_group["IpPermissions"],
+        },
+        service_response={},
+    )
+
+
+def describe_sg_echo(ec2_client_stub, security_group):
+    ec2_client_stub.add_response(
+        "describe_security_groups",
+        expected_params={"GroupIds": [security_group["GroupId"]]},
+        service_response={"SecurityGroups": [security_group]},
+    )
+
+
+def run_instances_with_network_interfaces_consumer(ec2_client_stub, network_interfaces):
+    ec2_client_stub.add_response(
+        "run_instances",
+        expected_params={
+            "NetworkInterfaces": network_interfaces,
+            "ImageId": ANY,
+            "InstanceType": ANY,
+            "KeyName": ANY,
+            "MinCount": ANY,
+            "MaxCount": ANY,
+            "TagSpecifications": ANY,
+        },
+        service_response={},
+    )
+
+
+def run_instances_with_launch_template_consumer(
+    ec2_client_stub, config, node_cfg, node_type_name, lt_data, max_count
+):
+    # create a copy of both node config and launch template data to modify
+    lt_data_cp = copy.deepcopy(lt_data)
+    node_cfg_cp = copy.deepcopy(node_cfg)
+    # override launch template parameters with explicit node config parameters
+    lt_data_cp.update(node_cfg_cp)
+    # copy all launch template parameters back to node config
+    node_cfg_cp.update(lt_data_cp)
+    # copy all default node provider config updates to node config
+    helpers.apply_node_provider_config_updates(
+        config, node_cfg_cp, node_type_name, max_count
+    )
+    # remove any security group and subnet IDs copied from network interfaces
+    node_cfg_cp.pop("SecurityGroupIds", [])
+    node_cfg_cp.pop("SubnetIds", [])
+    ec2_client_stub.add_response(
+        "run_instances", expected_params=node_cfg_cp, service_response={}
+    )
+
+
+def describe_instances_with_any_filter_consumer(ec2_client_stub):
+    ec2_client_stub.add_response(
+        "describe_instances", expected_params={"Filters": ANY}, service_response={}
+    )
+
+
+def describe_launch_template_versions_by_id_default(ec2_client_stub, versions):
+    ec2_client_stub.add_response(
+        "describe_launch_template_versions",
+        expected_params={
+            "LaunchTemplateId": DEFAULT_LT["LaunchTemplateId"],
+            "Versions": versions,
+        },
+        service_response={"LaunchTemplateVersions": [DEFAULT_LT]},
+    )
+
+
+def describe_launch_template_versions_by_name_default(ec2_client_stub, versions):
+    ec2_client_stub.add_response(
+        "describe_launch_template_versions",
+        expected_params={
+            "LaunchTemplateName": DEFAULT_LT["LaunchTemplateName"],
+            "Versions": versions,
+        },
+        service_response={"LaunchTemplateVersions": [DEFAULT_LT]},
+    )
+
+
+def get_ec2_cwa_installed_tag_true(ec2_client_stub, node_id):
+    ec2_client_stub.add_response(
+        "describe_instances",
+        expected_params={"InstanceIds": [node_id]},
+        service_response={
+            "Reservations": [
+                {
+                    "Instances": [
+                        {
+                            "InstanceId": node_id,
+                            "Tags": [
+                                {
+                                    "Key": CLOUDWATCH_AGENT_INSTALLED_TAG,
+                                    "Value": "True",
+                                },
+                            ],
+                        }
+                    ]
+                }
+            ]
+        },
+    )
+
+
+def update_hash_tag_success(ec2_client_stub, node_id, config_type, cloudwatch_helper):
+    hash_key_value = "-".join([CLOUDWATCH_CONFIG_HASH_TAG_BASE, config_type])
+    cur_hash_value = get_sha256_hash_of_cloudwatch_config_file(
+        config_type, cloudwatch_helper
+    )
+    ec2_client_stub.add_response(
+        "create_tags",
+        expected_params={
+            "Resources": [node_id],
+            "Tags": [{"Key": hash_key_value, "Value": cur_hash_value}],
+        },
+        service_response={"ResponseMetadata": {"HTTPStatusCode": 200}},
+    )
+
+
+def add_cwa_installed_tag_response(ec2_client_stub, node_id):
+    ec2_client_stub.add_response(
+        "create_tags",
+        expected_params={
+            "Resources": node_id,
+            "Tags": [{"Key": CLOUDWATCH_AGENT_INSTALLED_TAG, "Value": "True"}],
+        },
+        service_response={"ResponseMetadata": {"HTTPStatusCode": 200}},
+    )
+
+
+def get_head_node_config_hash_different(ec2_client_stub, config_type, cwh, node_id):
+    hash_key_value = "-".join([CLOUDWATCH_CONFIG_HASH_TAG_BASE, config_type])
+    cur_hash_value = get_sha256_hash_of_cloudwatch_config_file(config_type, cwh)
+    filters = cwh._get_current_cluster_session_nodes(cwh.cluster_name)
+    filters.append(
+        {
+            "Name": "tag:{}".format(TAG_RAY_NODE_KIND),
+            "Values": [NODE_KIND_HEAD],
+        }
+    )
+    ec2_client_stub.add_response(
+        "describe_instances",
+        expected_params={"Filters": filters},
+        service_response={
+            "Reservations": [
+                {
+                    "Instances": [
+                        {
+                            "InstanceId": node_id,
+                            "Tags": [
+                                {"Key": hash_key_value, "Value": cur_hash_value},
+                            ],
+                        }
+                    ]
+                }
+            ]
+        },
+    )
+
+
+def get_cur_node_config_hash_different(ec2_client_stub, config_type, node_id):
+    hash_key_value = "-".join([CLOUDWATCH_CONFIG_HASH_TAG_BASE, config_type])
+    ec2_client_stub.add_response(
+        "describe_instances",
+        expected_params={"InstanceIds": [node_id]},
+        service_response={
+            "Reservations": [
+                {
+                    "Instances": [
+                        {
+                            "InstanceId": node_id,
+                            "Tags": [
+                                {"Key": hash_key_value, "Value": str(uuid4())},
+                            ],
+                        }
+                    ]
+                }
+            ]
+        },
+    )
+
+
+def send_command_cwa_install(ssm_client_stub, node_id):
+    command_id = str(uuid4())
+    ssm_client_stub.add_response(
+        "send_command",
+        expected_params={
+            "DocumentName": "AWS-ConfigureAWSPackage",
+            "InstanceIds": node_id,
+            "MaxConcurrency": "1",
+            "MaxErrors": "0",
+            "Parameters": {
+                "action": ["Install"],
+                "name": ["AmazonCloudWatchAgent"],
+                "version": ["latest"],
+            },
+        },
+        service_response={
+            "Command": {
+                "CommandId": command_id,
+                "DocumentName": "AWS-ConfigureAWSPackage",
+            }
+        },
+    )
+    return command_id
+
+
+def list_command_invocations_status(ssm_client_stub, node_id, cmd_id, status):
+    ssm_client_stub.add_response(
+        "list_command_invocations",
+        expected_params={"CommandId": cmd_id, "InstanceId": node_id},
+        service_response={"CommandInvocations": [{"Status": status}]},
+    )
+
+
+def list_command_invocations_failed(ssm_client_stub, node_id, cmd_id):
+    status = "Failed"
+    list_command_invocations_status(ssm_client_stub, node_id, cmd_id, status)
+
+
+def list_command_invocations_success(ssm_client_stub, node_id, cmd_id):
+    status = "Success"
+    list_command_invocations_status(ssm_client_stub, node_id, cmd_id, status)
+
+
+def put_parameter_cloudwatch_config(ssm_client_stub, cluster_name, section_name):
+    ssm_config_param_name = helpers.get_ssm_param_name(cluster_name, section_name)
+    ssm_client_stub.add_response(
+        "put_parameter",
+        expected_params={
+            "Name": ssm_config_param_name,
+            "Type": "String",
+            "Value": ANY,
+            "Overwrite": True,
+            "Tier": ANY,
+        },
+        service_response={},
+    )
+
+
+def send_command_cwa_collectd_init(ssm_client_stub, node_id):
+    command_id = str(uuid4())
+    ssm_client_stub.add_response(
+        "send_command",
+        expected_params={
+            "DocumentName": "AWS-RunShellScript",
+            "InstanceIds": [node_id],
+            "MaxConcurrency": "1",
+            "MaxErrors": "0",
+            "Parameters": {
+                "commands": [
+                    "mkdir -p /usr/share/collectd/",
+                    "touch /usr/share/collectd/types.db",
+                ],
+            },
+        },
+        service_response={"Command": {"CommandId": command_id}},
+    )
+    return command_id
+
+
+def send_command_start_cwa(ssm_client_stub, node_id, parameter_name):
+    command_id = str(uuid4())
+    ssm_client_stub.add_response(
+        "send_command",
+        expected_params={
+            "DocumentName": "AmazonCloudWatch-ManageAgent",
+            "InstanceIds": [node_id],
+            "MaxConcurrency": "1",
+            "MaxErrors": "0",
+            "Parameters": {
+                "action": ["configure"],
+                "mode": ["ec2"],
+                "optionalConfigurationSource": ["ssm"],
+                "optionalConfigurationLocation": [parameter_name],
+                "optionalRestart": ["yes"],
+            },
+        },
+        service_response={"Command": {"CommandId": command_id}},
+    )
+    return command_id
+
+
+def send_command_stop_cwa(ssm_client_stub, node_id):
+    command_id = str(uuid4())
+    ssm_client_stub.add_response(
+        "send_command",
+        expected_params={
+            "DocumentName": "AmazonCloudWatch-ManageAgent",
+            "InstanceIds": [node_id],
+            "MaxConcurrency": "1",
+            "MaxErrors": "0",
+            "Parameters": {
+                "action": ["stop"],
+                "mode": ["ec2"],
+            },
+        },
+        service_response={"Command": {"CommandId": command_id}},
+    )
+    return command_id
+
+
+def get_param_ssm_same(ssm_client_stub, ssm_param_name, cloudwatch_helper, config_type):
+    command_id = str(uuid4())
+    cw_value_json = (
+        cloudwatch_helper.CLOUDWATCH_CONFIG_TYPE_TO_CONFIG_VARIABLE_REPLACE_FUNC.get(
+            config_type
+        )(config_type)
+    )
+    ssm_client_stub.add_response(
+        "get_parameter",
+        expected_params={"Name": ssm_param_name},
+        service_response={"Parameter": {"Value": json.dumps(cw_value_json)}},
+    )
+    return command_id
+
+
+def get_sha256_hash_of_cloudwatch_config_file(config_type, cloudwatch_helper):
+    cw_value_file = cloudwatch_helper._sha256_hash_file(config_type)
+    return cw_value_file
+
+
+def get_param_ssm_different(ssm_client_stub, ssm_param_name):
+    command_id = str(uuid4())
+    ssm_client_stub.add_response(
+        "get_parameter",
+        expected_params={"Name": ssm_param_name},
+        service_response={"Parameter": {"Value": "value"}},
+    )
+    return command_id
+
+
+def get_param_ssm_exception(ssm_client_stub, ssm_param_name):
+    command_id = str(uuid4())
+    ssm_client_stub.add_client_error(
+        "get_parameter",
+        "ParameterNotFound",
+        expected_params={"Name": ssm_param_name},
+        response_meta={"Error": {"Code": "ParameterNotFound"}},
+    )
+    return command_id
+
+
+def put_cluster_dashboard_success(cloudwatch_client_stub, cloudwatch_helper):
+    widgets = []
+    json_config_path = get_cloudwatch_dashboard_config_file_path()
+    with open(json_config_path) as f:
+        dashboard_config = json.load(f)
+
+    for item in dashboard_config:
+        item_out = cloudwatch_helper._replace_all_config_variables(
+            item,
+            cloudwatch_helper.node_id,
+            cloudwatch_helper.cluster_name,
+            cloudwatch_helper.provider_config["region"],
+        )
+        widgets.append(item_out)
+
+    dashboard_name = cloudwatch_helper.cluster_name + "-" + "example-dashboard-name"
+    cloudwatch_client_stub.add_response(
+        "put_dashboard",
+        expected_params={
+            "DashboardName": dashboard_name,
+            "DashboardBody": json.dumps({"widgets": widgets}),
+        },
+        service_response={"ResponseMetadata": {"HTTPStatusCode": 200}},
+    )
+
+
+def put_cluster_alarms_success(cloudwatch_client_stub, cloudwatch_helper):
+    json_config_path = get_cloudwatch_alarm_config_file_path()
+    with open(json_config_path) as f:
+        data = json.load(f)
+    for item in data:
+        item_out = copy.deepcopy(item)
+        cloudwatch_helper._replace_all_config_variables(
+            item_out,
+            cloudwatch_helper.node_id,
+            cloudwatch_helper.cluster_name,
+            cloudwatch_helper.provider_config["region"],
+        )
+        cloudwatch_client_stub.add_response(
+            "put_metric_alarm",
+            expected_params=item_out,
+            service_response={"ResponseMetadata": {"HTTPStatusCode": 200}},
+        )
+
+
+def get_metric_alarm(cloudwatch_client_stub):
+    cloudwatch_client_stub.add_response(
+        "describe_alarms",
+        expected_params={},
+        service_response={"MetricAlarms": [{"AlarmName": "myalarm"}]},
+    )
+
+
+def delete_metric_alarms(cloudwatch_client_stub):
+    cloudwatch_client_stub.add_response(
+        "delete_alarms",
+        expected_params={"AlarmNames": ["myalarm"]},
+        service_response={"ResponseMetadata": {"HTTPStatusCode": 200}},
+    )

python/ray/tests/chaos/chaos_network_bandwidth.yaml

@@ -0,0 +1,17 @@
+# Injects a bandwidth limit to 1mbps to all traffic to the Ray nodes.
+apiVersion: chaos-mesh.org/v1alpha1
+kind: NetworkChaos
+metadata:
+  name: bandwidth
+spec:
+  action: bandwidth
+  mode: all
+  selector:
+    namespaces:
+      - default
+    labelSelectors:
+      'ray.io/cluster': 'raycluster-kuberay' # inject to all pods
+  bandwidth:
+    rate: '1mbps'
+    limit: 20971520
+    buffer: 10000

python/ray/tests/chaos/chaos_network_delay.yaml

@@ -0,0 +1,16 @@
+# Injects a 200ms delay to all traffic to the Ray nodes.
+apiVersion: chaos-mesh.org/v1alpha1
+kind: NetworkChaos
+metadata:
+  name: network-delay
+spec:
+  action: delay
+  mode: all # inject to all pods
+  selector:
+    namespaces:
+      - default
+    labelSelectors:
+      'ray.io/cluster': 'raycluster-kuberay' # inject to all pods
+  delay:
+    latency: '200ms'
+  duration: '12h'

python/ray/tests/chaos/potato_passer.py

@@ -0,0 +1,68 @@
+import argparse
+import asyncio
+
+import ray
+
+ray.init()
+
+"""
+Potato passer is a test script that lets multiple actors call each other's methods.
+Actors are wired in a round-trip fashion: actor 0 calls actor 1, which calls actor 2.
+The last actor calls actor 0. In each call, the actor sleeps for a time, occationally
+prints, and calls next actor.
+
+Note the number of tasks on-the-fly can go up to `pass-times` because the next call is
+made before exiting current call.
+"""
+
+
+@ray.remote
+class PotatoPasser:
+    def __init__(self, name, next_name, sleep_secs):
+        self.count = 0
+        self.name = name
+        self.next_name = next_name
+        self.sleep_secs = sleep_secs
+        self.print_every = 100
+
+    async def pass_potato(self, potato: int, target: int):
+        self.count += 1
+        if potato % self.print_every == 0:
+            print(
+                f"running, name {self.name}, count {self.count}, "
+                f"potato {potato}, target {target}"
+            )
+        if potato >= target:
+            print(f"target reached! name = {self.name}, count = {self.count}")
+            return target
+        next_actor = ray.get_actor(self.next_name)
+        await asyncio.sleep(self.sleep_secs)
+        return await next_actor.pass_potato.remote(potato + 1, target)
+
+
+async def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num-actors", type=int, help="Make this many actors")
+    parser.add_argument("--pass-times", type=int, help="Pass this many messages")
+    parser.add_argument(
+        "--sleep-secs",
+        type=float,
+        help="Sleep seconds before sending message to next actor",
+    )
+    args = parser.parse_args()
+
+    actors = []
+    for i in range(args.num_actors):
+        this_actor = "actor" + str(i)
+        next_actor = "actor" + str((i + 1) % args.num_actors)
+        actor = PotatoPasser.options(
+            name=this_actor, scheduling_strategy="SPREAD"
+        ).remote(this_actor, next_actor, args.sleep_secs)
+        actors.append(actor)
+
+    ret = await actors[0].pass_potato.remote(0, args.pass_times)
+    print(f"passed potato {ret} times! expected {args.pass_times} times.")
+    assert ret == args.pass_times
+
+
+asyncio.run(main())

python/ray/tests/chaos/prepare_env.sh

@@ -0,0 +1,52 @@
+#!/usr/bin/env bash
+#
+# Sets up environment for the Kubernetes chaos testing.
+# The environment consists of:
+# - a KubeRay cluster, port-forwarded to localhost:8265.
+# - a chaos-mesh operator ready to inject faults.
+
+set -euo pipefail
+
+echo "--- Preparing k8s environment."
+bash ci/k8s/prep-k8s-environment.sh
+
+kind load docker-image ray-ci:kuberay-test
+
+# Helm install KubeRay
+echo "--- Installing KubeRay operator from official Helm repo."
+helm repo add kuberay https://ray-project.github.io/kuberay-helm/
+helm install kuberay-operator kuberay/kuberay-operator
+kubectl wait pod  -l app.kubernetes.io/name=kuberay-operator \
+    --for=condition=Ready=True  --timeout=2m
+
+echo "--- Installing KubeRay cluster and port forward."
+
+helm install raycluster kuberay/ray-cluster \
+    --set image.repository=ray-ci \
+    --set image.tag=kuberay-test \
+    --set worker.replicas=2 \
+    --set worker.resources.limits.cpu=500m \
+    --set worker.resources.requests.cpu=500m \
+    --set head.resources.limits.cpu=500m \
+    --set head.resources.requests.cpu=500m \
+    --set worker.resources.limits.memory=4Gi \
+    --set worker.resources.requests.memory=4Gi \
+    --set head.resources.limits.memory=4Gi \
+    --set head.resources.requests.memory=4Gi
+
+kubectl wait pod -l ray.io/cluster=raycluster-kuberay \
+    --for=condition=Ready=True --timeout=5m
+kubectl port-forward service/raycluster-kuberay-head-svc 8265:8265 &
+
+# Helm install chaos-mesh
+echo "--- Installing chaos-mesh operator and CR."
+helm repo add chaos-mesh https://charts.chaos-mesh.org
+kubectl create ns chaos-mesh
+helm install chaos-mesh chaos-mesh/chaos-mesh -n=chaos-mesh \
+    --set chaosDaemon.runtime=containerd \
+    --set chaosDaemon.socketPath=/run/containerd/containerd.sock \
+    --version 2.6.1
+
+echo "--- Waiting for chaos-mesh to be ready."
+kubectl wait pod --namespace chaos-mesh --timeout=300s \
+    -l app.kubernetes.io/instance=chaos-mesh --for=condition=Ready=True

python/ray/tests/chaos/runtime_env.yaml

@@ -0,0 +1,2 @@
+env_vars:
+  RAY_DEDUP_LOGS: "0"

python/ray/tests/chaos/streaming_llm.py

@@ -0,0 +1,98 @@
+import argparse
+import asyncio
+import logging
+
+import requests
+from fastapi import FastAPI
+from starlette.responses import StreamingResponse
+
+import ray
+from ray import serve
+
+logger = logging.getLogger("ray.serve")
+
+fastapi_app = FastAPI()
+
+
+# Input: a prompt of words
+# Output: each word reversed and produced N times.
+@serve.deployment(
+    num_replicas=6, ray_actor_options={"num_cpus": 0.01, "memory": 10 * 1024 * 1024}
+)
+class ReverseAndDupEachWord:
+    def __init__(self, dup_times: int):
+        self.dup_times = dup_times
+
+    async def __call__(self, prompt: str):
+        for word in prompt.split():
+            rev = word[::-1]
+            for _ in range(self.dup_times):
+                await asyncio.sleep(0.001)
+                # Ideally we want to do " ".join(words), but for the sake of
+                # simplicity we also have an extra trailing space.
+                yield rev + " "
+
+
+@serve.deployment(
+    num_replicas=6, ray_actor_options={"num_cpus": 0.01, "memory": 10 * 1024 * 1024}
+)
+@serve.ingress(fastapi_app)
+class Textbot:
+    def __init__(self, llm):
+        self.llm = llm.options(stream=True)
+
+    @fastapi_app.post("/")
+    async def handle_request(self, prompt: str) -> StreamingResponse:
+        logger.info(f'Got prompt with size "{len(prompt)}"')
+        return StreamingResponse(self.llm.remote(prompt), media_type="text/plain")
+
+
+@ray.remote(num_cpus=0.1, memory=10 * 1024 * 1024)
+def make_http_query(num_words, num_queries):
+    for _ in range(num_queries):
+        words = "Lorem ipsum dolor sit amet".split()
+        prompt_words = [words[i % len(words)] for i in range(num_words)]
+        prompt = " ".join(prompt_words)
+        expected_words = [word[::-1] for word in prompt_words for _ in range(2)]
+
+        response = requests.post(f"http://localhost:8000/?prompt={prompt}", stream=True)
+        response.raise_for_status()
+        content = response.content.decode()
+        assert content == " ".join(expected_words) + " ", content
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generates HTTP workloads with Ray.")
+
+    parser.add_argument("--num_tasks", type=int, required=True, help="Number of tasks.")
+    parser.add_argument(
+        "--num_queries_per_task",
+        type=int,
+        required=True,
+        help="Number of queries per task.",
+    )
+    parser.add_argument(
+        "--num_words_per_query",
+        type=int,
+        required=True,
+        help="Number of words per query",
+    )
+
+    args = parser.parse_args()
+
+    # Run the serve, run the client, then showdown serve.
+    llm = ReverseAndDupEachWord.bind(2)
+    app = Textbot.bind(llm)
+
+    serve.run(app)
+
+    objs = [
+        make_http_query.remote(args.num_words_per_query, args.num_queries_per_task)
+        for _ in range(args.num_tasks)
+    ]
+    ray.get(objs)
+
+    serve.shutdown()
+
+
+main()

python/ray/tests/chaos/streaming_llm.yaml

@@ -0,0 +1,2 @@
+env_vars:
+  RAY_DEDUP_LOGS: "0"

python/ray/tests/chaos/test_many_job_submissions.sh

@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+#
+# Sets up environment for the Kubernetes chaos testing.
+# The environment consists of:
+# - a KubeRay cluster, port-forwarded to localhost:8265.
+# - a chaos-mesh operator ready to inject faults.
+
+set -xe
+
+for i in {1..50}; do
+    echo "submitting round ${i}"
+    ray job submit --address http://localhost:8265 --runtime-env python/ray/tests/chaos/runtime_env.yaml --working-dir python/ray/tests/chaos -- python potato_passer.py --num-actors=3 --pass-times=3 --sleep-secs=0.01
+done

python/ray/tests/chaos/test_potato_passer.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+#
+# Sets up environment for the Kubernetes chaos testing.
+# The environment consists of:
+# - a KubeRay cluster, port-forwarded to localhost:8265.
+# - a chaos-mesh operator ready to inject faults.
+
+set -xe
+
+ray job submit --address http://localhost:8265 --runtime-env python/ray/tests/chaos/runtime_env.yaml --working-dir python/ray/tests/chaos -- python potato_passer.py --num-actors=3 --pass-times=1000 --sleep-secs=0.01

python/ray/tests/chaos/test_streaming_llm.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+#
+# Sets up environment for the Kubernetes chaos testing.
+# The environment consists of:
+# - a KubeRay cluster, port-forwarded to localhost:8265.
+# - a chaos-mesh operator ready to inject faults.
+
+set -xe
+
+ray job submit --address http://localhost:8265 --runtime-env python/ray/tests/chaos/streaming_llm.yaml --working-dir python/ray/tests/chaos -- python streaming_llm.py --num_queries_per_task=100 --num_tasks=2 --num_words_per_query=100

python/ray/tests/client_test_utils.py

@@ -0,0 +1,38 @@
+import asyncio
+
+
+def create_remote_signal_actor(ray):
+    # TODO(barakmich): num_cpus=0
+    @ray.remote
+    class SignalActor:
+        def __init__(self):
+            self.ready_event = asyncio.Event()
+
+        def send(self, clear=False):
+            self.ready_event.set()
+            if clear:
+                self.ready_event.clear()
+
+        async def wait(self, should_wait=True):
+            if should_wait:
+                await self.ready_event.wait()
+
+    return SignalActor
+
+
+# See test_client::test_wrapped_actor_creation for details on usage of
+# run_wrapped_actor_creation and SomeClass.
+def run_wrapped_actor_creation():
+    import ray
+
+    RemoteClass = ray.remote(SomeClass)
+    handle = RemoteClass.remote()
+    return ray.get(handle.ready.remote())
+
+
+class SomeClass:
+    def __init__(self):
+        pass
+
+    def ready(self):
+        return 1

python/ray/tests/conftest_docker.py

@@ -0,0 +1,297 @@
+import subprocess
+import time
+from typing import List
+
+import pytest
+from pytest_docker_tools import container, fetch, network, volume, wrappers
+
+import docker
+
+from ray._common.network_utils import build_address
+
+# If you need to debug tests using fixtures in this file,
+# comment in the volume
+# mounts in the head node and worker node containers below and use
+# the repro-ci.py script to spin up an instance. The test
+# setup is a little intricate, as it uses docker-in-docker.
+# You need to ssh into the host machine, find the
+# docker-in-docker container with
+#
+# docker ps
+#
+# Log into the container with
+#
+# docker exec -it <dind-daemon container id> sh
+#
+# And run
+#
+# mkdir -p /tmp/ray
+# chmod 777 /tmp/ray
+#
+# Now you can re-run the test and the logs will show
+# up in /tmp/ray in the docker-in-docker container.
+# Good luck!
+
+
+class Container(wrappers.Container):
+    def ready(self):
+        self._container.reload()
+        if self.status == "exited":
+            from pytest_docker_tools.exceptions import ContainerFailed
+
+            raise ContainerFailed(
+                self,
+                f"Container {self.name} has already exited before "
+                "we noticed it was ready",
+            )
+
+        if self.status != "running":
+            return False
+
+        networks = self._container.attrs["NetworkSettings"]["Networks"]
+        for (_, n) in networks.items():
+            if not n["IPAddress"]:
+                return False
+
+        if "Ray runtime started" in super().logs():
+            return True
+        return False
+
+    def client(self):
+        from http.client import HTTPConnection
+
+        port = self.ports["8000/tcp"][0]
+        return HTTPConnection(f"localhost:{port}")
+
+    def print_logs(self):
+        for (name, content) in self.get_files("/tmp"):
+            print(f"===== log start:  {name} ====")
+            print(content.decode())
+
+
+# This allows us to assign static ips to docker containers
+ipam_config = docker.types.IPAMConfig(
+    pool_configs=[
+        docker.types.IPAMPool(subnet="192.168.52.0/24", gateway="192.168.52.254")
+    ]
+)
+gcs_network = network(driver="bridge", ipam=ipam_config)
+
+redis_image = fetch(repository="redis:latest")
+
+redis = container(
+    image="{redis_image.id}",
+    network="{gcs_network.name}",
+    command=("redis-server --save 60 1 --loglevel warning"),
+)
+
+head_node_vol = volume()
+worker_node_vol = volume()
+head_node_container_name = "gcs" + str(int(time.time()))
+
+
+def gen_head_node(envs):
+    return container(
+        image="rayproject/ray:ha_integration",
+        name=head_node_container_name,
+        network="{gcs_network.name}",
+        command=[
+            "ray",
+            "start",
+            "--head",
+            "--block",
+            "--num-cpus",
+            "0",
+            # Fix the port of raylet to make sure raylet restarts at the same
+            # ip:port is treated as a different raylet.
+            "--node-manager-port",
+            "9379",
+            "--dashboard-host",
+            "0.0.0.0",
+        ],
+        volumes={"{head_node_vol.name}": {"bind": "/tmp", "mode": "rw"}},
+        environment=envs,
+        wrapper_class=Container,
+        ports={
+            "8000/tcp": None,
+        },
+        # volumes={
+        #     "/tmp/ray/": {"bind": "/tmp/ray/", "mode": "rw"}
+        # },
+    )
+
+
+def gen_worker_node(envs, num_cpus):
+    return container(
+        image="rayproject/ray:ha_integration",
+        network="{gcs_network.name}",
+        command=[
+            "ray",
+            "start",
+            "--address",
+            build_address(head_node_container_name, 6379),
+            "--block",
+            # Fix the port of raylet to make sure raylet restarts at the same
+            # ip:port is treated as a different raylet.
+            "--node-manager-port",
+            "9379",
+            "--num-cpus",
+            f"{num_cpus}",
+        ],
+        volumes={"{worker_node_vol.name}": {"bind": "/tmp", "mode": "rw"}},
+        environment=envs,
+        wrapper_class=Container,
+        ports={
+            "8000/tcp": None,
+        },
+        # volumes={
+        #     "/tmp/ray/": {"bind": "/tmp/ray/", "mode": "rw"}
+        # },
+    )
+
+
+head_node = gen_head_node(
+    {
+        "RAY_REDIS_ADDRESS": "{redis.ips.primary}:6379",
+        "RAY_raylet_client_num_connect_attempts": "10",
+        "RAY_raylet_client_connect_timeout_milliseconds": "100",
+    }
+)
+
+worker_node = gen_worker_node(
+    envs={
+        "RAY_REDIS_ADDRESS": "{redis.ips.primary}:6379",
+        "RAY_raylet_client_num_connect_attempts": "10",
+        "RAY_raylet_client_connect_timeout_milliseconds": "100",
+    },
+    num_cpus=8,
+)
+
+
+@pytest.fixture
+def docker_cluster(head_node, worker_node):
+    yield (head_node, worker_node)
+
+
+def run_in_container(cmds: List[List[str]], container_id: str):
+    """Run a list of commands in the specified container.
+
+    Checks that each docker command executed without error.
+    Returns the output from each command as a list.
+    """
+
+    outputs = []
+    for cmd in cmds:
+        docker_cmd = ["docker", "exec", container_id] + cmd
+        print(f"Executing command: {docker_cmd}", time.time())
+        try:
+            resp = subprocess.check_output(docker_cmd, stderr=subprocess.STDOUT)
+            output = resp.decode("utf-8").strip()
+            print(f"Output: {output}")
+            outputs.append(output)
+        except subprocess.CalledProcessError as e:
+            error_output = e.output.decode("utf-8") if e.output else "No output"
+            print(f"Command failed with return code {e.returncode}")
+            print(f"Full error output:\n{error_output}")
+            raise
+
+    return outputs
+
+
+IMAGE_NAME = "rayproject/ray:runtime_env_container"
+# After `docker save` / `podman load`, Podman typically tags the image as below (not the
+# Docker daemon name). Use that ref for `podman create` so resolution stays local.
+PODMAN_BASE_IMAGE = "localhost/runtime_env_container:latest"
+NESTED_IMAGE_NAME = "localhost/runtime_env_container_nested:latest"
+
+
+@pytest.fixture(scope="session")
+def podman_docker_cluster():
+    start_container_command = [
+        "docker",
+        "run",
+        "-d",
+        "--privileged",
+        "-v",
+        "/var/run/docker.sock:/var/run/docker.sock",
+        "-v",
+        "/var/lib/containers:/var/lib/containers",
+        # For testing environment variables
+        "--env",
+        "RAY_TEST_ABC=1",
+        "--env",
+        "TEST_ABC=1",
+        IMAGE_NAME,
+        "tail",
+        "-f",
+        "/dev/null",
+    ]
+    try:
+        container_id = subprocess.check_output(
+            start_container_command, stderr=subprocess.STDOUT
+        ).decode("utf-8")
+    except subprocess.CalledProcessError as e:
+        error_output = e.output.decode("utf-8") if e.output else "No output"
+        print(f"Command failed with return code {e.returncode}")
+        print(f"Full error output:\n{error_output}")
+        raise
+
+    container_id = container_id.strip()
+
+    # Get group id that owns the docker socket file. Add user `ray` to
+    # group to get necessary permissions for pulling an image from
+    # docker's local storage into podman
+    docker_group_id = run_in_container(
+        [["stat", "-c", "%g", "/var/run/docker.sock"]], container_id
+    )[0]
+    run_in_container(
+        [
+            ["id"],
+            ["sudo", "groupadd", "-g", docker_group_id, "docker"],
+            ["sudo", "usermod", "-aG", "docker", "ray"],
+            [
+                "bash",
+                "-c",
+                f"docker save {IMAGE_NAME} | podman load",
+            ],
+        ],
+        container_id,
+    )
+
+    # Add custom file to new image tagged `runtime_env_container_nested`,
+    # which can be read by Ray actors / Serve deployments to verify the
+    # container runtime env plugin. Also add serve application that will
+    # be imported by the telemetry test.
+    serve_app = """
+from ray import serve
+@serve.deployment
+class Model:
+    def __call__(self):
+        with open("file.txt") as f:
+            return f.read().strip()
+app = Model.bind()
+"""
+
+    run_in_container(
+        [
+            ["bash", "-c", "echo helloworldalice >> /tmp/file.txt"],
+            ["bash", "-c", f"echo '{serve_app}' >> /tmp/serve_application.py"],
+            ["podman", "create", "--name", "tmp_container", PODMAN_BASE_IMAGE],
+            ["podman", "cp", "/tmp/file.txt", "tmp_container:/home/ray/file.txt"],
+            [
+                "podman",
+                "cp",
+                "/tmp/serve_application.py",
+                "tmp_container:/home/ray/serve_application.py",
+            ],
+            ["podman", "commit", "tmp_container", NESTED_IMAGE_NAME],
+        ],
+        container_id,
+    )
+
+    # For debugging
+    run_in_container([["podman", "image", "ls"]], container_id)
+
+    yield container_id
+
+    subprocess.check_call(["docker", "kill", container_id])

python/ray/tests/gcp/test_gcp_node_provider.py

@@ -0,0 +1,444 @@
+import logging
+import sys
+from threading import RLock
+from typing import Dict
+from unittest.mock import MagicMock, call, patch
+
+import pytest
+
+from ray.autoscaler._private.command_runner import DockerCommandRunner, SSHCommandRunner
+from ray.autoscaler._private.gcp.config import (
+    _get_num_tpu_chips,
+    _has_tpus_in_node_configs,
+    _is_single_host_tpu,
+    get_node_type,
+    tpu_accelerator_config_to_type,
+)
+from ray.autoscaler._private.gcp.node import (
+    GCPCompute,
+    GCPNode,
+    GCPNodeType,
+    GCPResource,
+)
+from ray.autoscaler._private.gcp.node_provider import GCPNodeProvider
+from ray.autoscaler._private.gcp.tpu_command_runner import (
+    TPUCommandRunner,
+    TPUVMDockerCommandRunner,
+    TPUVMSSHCommandRunner,
+)
+from ray.tests.test_autoscaler import MockProcessRunner
+
+_PROJECT_NAME = "project-one"
+_AZ = "us-west1-b"
+
+auth_config = {
+    "ssh_user": "ray",
+    "ssh_private_key": "8265.pem",
+}
+
+
+def test_create_node_returns_dict():
+    mock_node_config = {"machineType": "n2-standard-8"}
+    mock_results = [({"dict": 1}, "instance_id1"), ({"dict": 2}, "instance_id2")]
+    mock_resource = MagicMock()
+    mock_resource.create_instances.return_value = mock_results
+    expected_return_value = {"instance_id1": {"dict": 1}, "instance_id2": {"dict": 2}}
+
+    def __init__(self, provider_config: dict, cluster_name: str):
+        self.lock = RLock()
+        self.cached_nodes: Dict[str, GCPNode] = {}
+        self.resources: Dict[GCPNodeType, GCPResource] = {}
+        self.cache_stopped_nodes = False
+        self.resources[GCPNodeType.COMPUTE] = mock_resource
+
+    with patch.object(GCPNodeProvider, "__init__", __init__):
+        node_provider = GCPNodeProvider({}, "")
+        create_node_return_value = node_provider.create_node(mock_node_config, {}, 1)
+    assert create_node_return_value == expected_return_value
+
+
+def test_terminate_nodes():
+    mock_node_config = {"machineType": "n2-standard-8"}
+    node_type = GCPNodeType.COMPUTE.value
+    id1, id2 = f"instance-id1-{node_type}", f"instance-id2-{node_type}"
+    terminate_node_ids = [id1, id2]
+    mock_resource = MagicMock()
+    mock_resource.create_instances.return_value = [
+        ({"dict": 1}, id1),
+        ({"dict": 2}, id2),
+    ]
+    mock_resource.delete_instance.return_value = "test"
+
+    def __init__(self, provider_config: dict, cluster_name: str):
+        self.lock = RLock()
+        self.cached_nodes: Dict[str, GCPNode] = {}
+        self.resources: Dict[GCPNodeType, GCPResource] = {}
+        self.cache_stopped_nodes = False
+        self.resources[GCPNodeType.COMPUTE] = mock_resource
+
+    with patch.object(GCPNodeProvider, "__init__", __init__):
+        node_provider = GCPNodeProvider({}, "")
+        node_provider.create_node(mock_node_config, {}, 1)
+        node_provider.terminate_nodes(terminate_node_ids)
+
+    mock_resource.delete_instance.assert_has_calls(
+        [call(node_id=id1), call(node_id=id2)], any_order=True
+    )
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        ("n1-standard-4", f"zones/{_AZ}/machineTypes/n1-standard-4"),
+        (
+            f"zones/{_AZ}/machineTypes/n1-standard-4",
+            f"zones/{_AZ}/machineTypes/n1-standard-4",
+        ),
+    ],
+)
+def test_convert_resources_to_urls_machine(test_case):
+    gcp_compute = GCPCompute(None, _PROJECT_NAME, _AZ, "cluster_name")
+    base_machine, result_machine = test_case
+    modified_config = gcp_compute._convert_resources_to_urls(
+        {"machineType": base_machine}
+    )
+    assert modified_config["machineType"] == result_machine
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        (
+            "nvidia-tesla-k80",
+            f"projects/{_PROJECT_NAME}/zones/{_AZ}/acceleratorTypes/nvidia-tesla-k80",
+        ),
+        (
+            f"projects/{_PROJECT_NAME}/zones/{_AZ}/acceleratorTypes/nvidia-tesla-k80",
+            f"projects/{_PROJECT_NAME}/zones/{_AZ}/acceleratorTypes/nvidia-tesla-k80",
+        ),
+    ],
+)
+def test_convert_resources_to_urls_accelerators(test_case):
+    gcp_compute = GCPCompute(None, _PROJECT_NAME, _AZ, "cluster_name")
+    base_accel, result_accel = test_case
+
+    base_config = {
+        "machineType": "n1-standard-4",
+        "guestAccelerators": [{"acceleratorCount": 1, "acceleratorType": base_accel}],
+    }
+    modified_config = gcp_compute._convert_resources_to_urls(base_config)
+
+    assert modified_config["guestAccelerators"][0]["acceleratorType"] == result_accel
+
+
+def test_compute_node_list_instances_excludes_tpu():
+    mock_execute = MagicMock(return_value={"test": "abc"})
+    mock_list = MagicMock(return_value=MagicMock(execute=mock_execute))
+    mock_instances = MagicMock(return_value=MagicMock(list=mock_list))
+    mock_resource = MagicMock(instances=mock_instances)
+
+    GCPCompute(mock_resource, _PROJECT_NAME, _AZ, "cluster_name").list_instances()
+    filter_arg = mock_list.call_args.kwargs["filter"]
+
+    # Checks that the tpu negation filter is included.
+    assert "tpu_cores" in filter_arg
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        (
+            {},
+            SSHCommandRunner,
+        ),
+        (
+            {"docker_config": {"container_name": "container"}},
+            DockerCommandRunner,
+        ),
+    ],
+)
+def test_cpu_resource_returns_standard_command_runner(test_case):
+    mock_resource = MagicMock()
+
+    def __init__(self, provider_config: dict, cluster_name: str):
+        self.lock = RLock()
+        self.cached_nodes: Dict[str, GCPNode] = {}
+        self.resources: Dict[GCPNodeType, GCPResource] = {}
+        self.resources[GCPNodeType.COMPUTE] = mock_resource
+
+    with patch.object(GCPNodeProvider, "__init__", __init__):
+        node_provider = GCPNodeProvider({}, "")
+
+    optional_docker_config, expected_runner = test_case
+
+    args = {
+        "log_prefix": "test",
+        "node_id": "test-instance-compute",
+        "auth_config": auth_config,
+        "cluster_name": "test",
+        "process_runner": MockProcessRunner(),
+        "use_internal_ip": True,
+    }
+    args.update(optional_docker_config)
+    command_runner = node_provider.get_command_runner(**args)
+    assert isinstance(command_runner, expected_runner)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        (
+            {},
+            TPUVMSSHCommandRunner,
+        ),
+        (
+            {"docker_config": {"container_name": "container"}},
+            TPUVMDockerCommandRunner,
+        ),
+    ],
+)
+def test_tpu_resource_returns_tpu_command_runner(test_case):
+    mock_resource = MagicMock()
+
+    def __init__(self, provider_config: dict, cluster_name: str):
+        self.lock = RLock()
+        self.cached_nodes: Dict[str, GCPNode] = {}
+        self.resources: Dict[GCPNodeType, GCPResource] = {}
+        self.resources[GCPNodeType.COMPUTE] = mock_resource
+        self.resources[GCPNodeType.TPU] = mock_resource
+
+    with patch.object(GCPNodeProvider, "__init__", __init__):
+        node_provider = GCPNodeProvider({}, "")
+
+    optional_docker_config, expected_runner = test_case
+
+    args = {
+        "log_prefix": "test",
+        "node_id": "test-instance-tpu",
+        "auth_config": auth_config,
+        "cluster_name": "test",
+        "process_runner": MockProcessRunner(),
+        "use_internal_ip": True,
+    }
+    args.update(optional_docker_config)
+    command_runner = node_provider.get_command_runner(**args)
+    assert isinstance(command_runner, TPUCommandRunner)
+    assert isinstance(command_runner._command_runners[0], expected_runner)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        ({"acceleratorType": "v4-16"}, "TPU-v4-16-head"),
+        ({"acceleratorType": "v4-32"}, "TPU-v4-32-head"),
+        ({"acceleratorType": "v3-8"}, "TPU-v3-8-head"),
+        ({"acceleratorConfig": {"type": "V4", "topology": "2x2x2"}}, "TPU-v4-16-head"),
+        ({"acceleratorConfig": {"type": "V4", "topology": "4x4x4"}}, "TPU-v4-128-head"),
+        (
+            {"acceleratorConfig": {"type": "V5LITE_POD", "topology": "2x4"}},
+            "TPU-v5litepod-8-head",
+        ),
+        (
+            {"acceleratorConfig": {"type": "V6E", "topology": "2x4"}},
+            "TPU-v6e-8-head",
+        ),
+    ],
+)
+def test_tpu_node_fillout(test_case):
+    accelerator_config, expected_resource_str = test_case
+
+    cluster_config = {
+        "available_node_types": {
+            "ray_tpu": {
+                "resources": {"TPU": 4},
+                "node_config": {
+                    "runtimeVersion": "tpu-vm-v4-base",
+                },
+            },
+        },
+    }
+
+    cluster_config["available_node_types"]["ray_tpu"]["node_config"].update(
+        accelerator_config
+    )
+
+    new_config = GCPNodeProvider.fillout_available_node_types_resources(
+        cluster_config=cluster_config
+    )
+    resource_config = new_config["available_node_types"]["ray_tpu"]["resources"]
+    assert expected_resource_str in resource_config
+    assert resource_config[expected_resource_str] == 1
+
+
+def test_tpu_config_cannot_have_accelerator_type_and_config():
+    node = {
+        "acceleratorType": "abc",
+        "acceleratorConfig": {"abc": "def"},
+    }
+    with pytest.raises(ValueError):
+        get_node_type(node)
+
+
+@pytest.mark.parametrize(
+    "node",
+    [
+        {"acceleratorConfig": {"type": "V3", "topology": "2x2"}},
+        {"acceleratorConfig": {"type": "V2", "topology": "2x2"}},
+    ],
+)
+def test_get_node_rejects_v2_v3_accelerator_config(node):
+    with pytest.raises(ValueError):
+        get_node_type(node)
+
+
+@pytest.mark.parametrize(
+    "node_config",
+    [
+        {"acceleratorType": "vabc-12345"},
+        {"acceleratorType": "v3-abc"},
+        {"acceleratorType": "v3-8a"},
+        {"acceleratorType": "this should fail"},
+        {"acceleratorConfig": {"type": "asdf", "topology": "2x2x1"}},
+        {"acceleratorConfig": {"type": "V4", "topology": "asdf"}},
+    ],
+)
+def test_invalid_accelerator_configs(node_config):
+    with pytest.raises(ValueError):
+        get_node_type(node_config)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        ({"acceleratorType": "v2-8"}, 4, True),
+        ({"acceleratorType": "v3-8"}, 4, True),
+        ({"acceleratorType": "v4-8"}, 4, True),
+        # Note: Topology only supported in v4
+        ({"acceleratorConfig": {"type": "V4", "topology": "2x2x1"}}, 4, True),
+        ({"acceleratorType": "v2-32"}, 16, False),
+        ({"acceleratorType": "v3-128"}, 64, False),
+        ({"acceleratorType": "v4-4096"}, 2048, False),
+        ({"acceleratorConfig": {"type": "V4", "topology": "2x2x8"}}, 32, False),
+        ({"acceleratorConfig": {"type": "V4", "topology": "4x4x4"}}, 64, False),
+        ({"acceleratorConfig": {"type": "V5LITE_POD", "topology": "2x4"}}, 8, True),
+        ({"acceleratorConfig": {"type": "V6E", "topology": "2x4"}}, 8, True),
+    ],
+)
+def test_tpu_chip_calculation_single_host_logic(test_case):
+    node, expected_chips, expected_singlehost = test_case
+    assert _get_num_tpu_chips(node) == expected_chips
+    assert _is_single_host_tpu(node) == expected_singlehost
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        ({"machineType": "n2-standard-4"}, GCPNodeType.COMPUTE, False),
+        (
+            {
+                "machineType": "n2-standard-4",
+                "acceleratorType": {
+                    "guestAccelerators": {
+                        "acceleratorType": "V100",
+                        "acceleratorCount": 1,
+                    }
+                },
+            },
+            GCPNodeType.COMPUTE,
+            False,
+        ),
+        ({"acceleratorType": "v2-8"}, GCPNodeType.TPU, True),
+        ({"acceleratorType": "v3-8"}, GCPNodeType.TPU, True),
+        ({"acceleratorType": "v4-8"}, GCPNodeType.TPU, True),
+        (
+            {"acceleratorConfig": {"type": "V4", "topology": "2x2x1"}},
+            GCPNodeType.TPU,
+            True,
+        ),
+        ({"acceleratorType": "v2-32"}, GCPNodeType.TPU, True),
+        ({"acceleratorType": "v3-128"}, GCPNodeType.TPU, True),
+        ({"acceleratorType": "v4-4096"}, GCPNodeType.TPU, True),
+        (
+            {"acceleratorConfig": {"type": "V4", "topology": "2x2x8"}},
+            GCPNodeType.TPU,
+            True,
+        ),
+        (
+            {"acceleratorConfig": {"type": "V4", "topology": "4x4x4"}},
+            GCPNodeType.TPU,
+            True,
+        ),
+        (
+            {"acceleratorConfig": {"type": "V5LITE_POD", "topology": "2x4"}},
+            GCPNodeType.TPU,
+            True,
+        ),
+        (
+            {"acceleratorConfig": {"type": "V6E", "topology": "2x4"}},
+            GCPNodeType.TPU,
+            True,
+        ),
+    ],
+)
+def test_get_node_type_and_has_tpu(test_case):
+    node, expected_compute_type, expected_is_tpu = test_case
+    assert get_node_type(node) == expected_compute_type
+    config = {
+        "available_node_types": {
+            "node_type_1": {"node_config": node},
+        },
+    }
+    assert _has_tpus_in_node_configs(config) == expected_is_tpu
+
+
+@pytest.mark.parametrize(
+    "accelerator_pod_tuple",
+    [
+        ({"acceleratorType": "v2-32"}, True),
+        ({"acceleratorType": "v3-32"}, True),
+        ({"acceleratorType": "v4-32"}, True),
+        ({"acceleratorConfig": {"type": "V4", "topology": "2x2x2"}}, True),
+        ({"acceleratorType": "v2-8"}, False),
+        ({"acceleratorType": "v3-8"}, False),
+        ({"acceleratorType": "v4-8"}, False),
+        ({"acceleratorConfig": {"type": "V4", "topology": "2x2x1"}}, False),
+    ],
+)
+def test_tpu_pod_emits_warning(propagate_logs, caplog, accelerator_pod_tuple):
+    accelerator, should_emit = accelerator_pod_tuple
+
+    with caplog.at_level(
+        logging.WARNING, logger="ray.autoscaler._private.gcp.config.get_node_type"
+    ):
+        get_node_type(accelerator)
+        if should_emit:
+            assert "TPU pod detected" in caplog.text
+        else:
+            assert "TPU pod detected" not in caplog.text
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        ("v4-8", "V4", "2x2x1"),
+        ("v4-16", "V4", "2x2x2"),
+        ("v4-128", "V4", "4x4x4"),
+        ("v4-256", "V4", "4x4x8"),
+        ("v5litepod-8", "V5LITE_POD", "2x4"),
+        ("v6e-8", "V6E", "2x4"),
+    ],
+)
+def test_tpu_accelerator_config_to_type(test_case):
+    expected, accel_type, topology = test_case
+    accelerator_config = {
+        "type": accel_type,
+        "topology": topology,
+    }
+    accelerator_type = tpu_accelerator_config_to_type(
+        accelerator_config=accelerator_config
+    )
+    assert accelerator_type == expected
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tests/gcp/test_gcp_tpu_command_runner.py

@@ -0,0 +1,290 @@
+import hashlib
+import os
+import sys
+from getpass import getuser
+from unittest.mock import patch
+
+import pytest
+
+from ray._private import ray_constants
+from ray.autoscaler._private.command_runner import SSHCommandRunner
+from ray.autoscaler._private.gcp.tpu_command_runner import TPUCommandRunner
+from ray.tests.test_autoscaler import MockProcessRunner, MockProvider
+
+_MOCK_TPU_NAME = "my-tpu"
+_MOCK_ACCELERATOR_TYPE = "v4-16"
+
+auth_config = {
+    "ssh_user": "ray",
+    "ssh_private_key": "8265.pem",
+}
+
+
+class MockTpuInstance:
+    def __init__(self, num_workers: int = 1):
+        self.num_workers = num_workers
+
+    def get_internal_ip(self, worker_index: int) -> str:
+        return "0.0.0.0"
+
+    def get_external_ip(self, worker_index: int) -> str:
+        return "1.2.3.4"
+
+    def get(self, key) -> str:
+        if key == "name":
+            return _MOCK_TPU_NAME
+        elif key == "acceleratorType":
+            return _MOCK_ACCELERATOR_TYPE
+        return ""
+
+
+def test_tpu_ssh_command_runner():
+    num_workers = 2
+    process_runner = MockProcessRunner()
+    provider = MockProvider()
+    instance = MockTpuInstance(num_workers=num_workers)
+    provider.create_node({}, {}, 1)
+    cluster_name = "cluster"
+    ssh_control_hash = hashlib.sha256(cluster_name.encode()).hexdigest()
+    ssh_user_hash = hashlib.sha256(getuser().encode()).hexdigest()
+    ssh_control_path = "/tmp/ray_ssh_{}/{}".format(
+        ssh_user_hash[:10], ssh_control_hash[:10]
+    )
+    args = {
+        "instance": instance,
+        "log_prefix": "prefix",
+        "node_id": "abc",
+        "provider": provider,
+        "auth_config": auth_config,
+        "cluster_name": cluster_name,
+        "process_runner": process_runner,
+        "use_internal_ip": False,
+    }
+    env_vars = {"var1": 'quote between this " and this', "var2": "123"}
+    cmd_runner = TPUCommandRunner(**args)
+    cmd_runner.run(
+        "echo helloo", port_forward=[(8265, 8265)], environment_variables=env_vars
+    )
+
+    expected = [
+        "ssh",
+        "-tt",
+        "-L",
+        "8265:localhost:8265",
+        "-i",
+        "8265.pem",
+        "-o",
+        "StrictHostKeyChecking=no",
+        "-o",
+        "UserKnownHostsFile=/dev/null",
+        "-o",
+        "IdentitiesOnly=yes",
+        "-o",
+        "ExitOnForwardFailure=yes",
+        "-o",
+        "ServerAliveInterval=5",
+        "-o",
+        "ServerAliveCountMax=3",
+        "-o",
+        "ControlMaster=auto",
+        "-o",
+        "ControlPath={}/%C".format(ssh_control_path),
+        "-o",
+        "ControlPersist=10s",
+        "-o",
+        "ConnectTimeout=120s",
+        "ray@1.2.3.4",
+        "bash",
+        "--login",
+        "-c",
+        "-i",
+        """'source ~/.bashrc; export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && (export var1='"'"'"quote between this \\" and this"'"'"';export var2='"'"'"123"'"'"';echo helloo)'""",  # noqa: E501
+    ]
+
+    calls = process_runner.calls
+
+    # Asserts that we do make the call once per worker in the TPU pod.
+    assert len(process_runner.calls) == num_workers
+
+    # Much easier to debug this loop than the function call.
+    for i in range(num_workers):
+        for x, y in zip(calls[i], expected):
+            assert x == y
+
+
+def test_tpu_docker_command_runner():
+    num_workers = 4
+    process_runner = MockProcessRunner()
+    provider = MockProvider()
+    instance = MockTpuInstance(num_workers=num_workers)
+    provider.create_node({}, {}, 1)
+    cluster_name = "cluster"
+    ssh_control_hash = hashlib.sha256(cluster_name.encode()).hexdigest()
+    ssh_user_hash = hashlib.sha256(getuser().encode()).hexdigest()
+    ssh_control_path = "/tmp/ray_ssh_{}/{}".format(
+        ssh_user_hash[:10], ssh_control_hash[:10]
+    )
+    docker_config = {"container_name": "container"}
+    args = {
+        "instance": instance,
+        "log_prefix": "prefix",
+        "node_id": "0",
+        "provider": provider,
+        "auth_config": auth_config,
+        "cluster_name": cluster_name,
+        "process_runner": process_runner,
+        "use_internal_ip": False,
+        "docker_config": docker_config,
+    }
+    cmd_runner = TPUCommandRunner(**args)
+    env_vars = {"var1": 'quote between this " and this', "var2": "123"}
+    cmd_runner.run("echo hello", environment_variables=env_vars)
+
+    # This string is insane because there are an absurd number of embedded
+    # quotes. While this is a ridiculous string, the escape behavior is
+    # important and somewhat difficult to get right for environment variables.
+    cmd = """'source ~/.bashrc; export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && (docker exec -it  container /bin/bash -c '"'"'bash --login -c -i '"'"'"'"'"'"'"'"'source ~/.bashrc; export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && (export var1='"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"quote between this \\" and this"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"';export var2='"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"123"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"';echo hello)'"'"'"'"'"'"'"'"''"'"' )'"""  # noqa: E501
+
+    expected = [
+        "ssh",
+        "-tt",
+        "-i",
+        "8265.pem",
+        "-o",
+        "StrictHostKeyChecking=no",
+        "-o",
+        "UserKnownHostsFile=/dev/null",
+        "-o",
+        "IdentitiesOnly=yes",
+        "-o",
+        "ExitOnForwardFailure=yes",
+        "-o",
+        "ServerAliveInterval=5",
+        "-o",
+        "ServerAliveCountMax=3",
+        "-o",
+        "ControlMaster=auto",
+        "-o",
+        "ControlPath={}/%C".format(ssh_control_path),
+        "-o",
+        "ControlPersist=10s",
+        "-o",
+        "ConnectTimeout=120s",
+        "ray@1.2.3.4",
+        "bash",
+        "--login",
+        "-c",
+        "-i",
+        cmd,
+    ]
+
+    calls = process_runner.calls
+
+    # Asserts that we do make the call once per worker in the TPU pod.
+    assert len(process_runner.calls) == num_workers
+
+    # Much easier to debug this loop than the function call.
+    for i in range(num_workers):
+        for x, y in zip(calls[i], expected):
+            assert x == y
+
+
+def test_tpu_docker_run_init():
+    num_workers = 1
+    process_runner = MockProcessRunner()
+    provider = MockProvider()
+    instance = MockTpuInstance(num_workers=num_workers)
+    provider.create_node({}, {}, 1)
+    cluster_name = "cluster"
+    docker_config = {
+        "container_name": "container",
+        "image": "rayproject/ray:latest",
+    }
+    args = {
+        "instance": instance,
+        "log_prefix": "prefix",
+        "node_id": "0",
+        "provider": provider,
+        "auth_config": auth_config,
+        "cluster_name": cluster_name,
+        "process_runner": process_runner,
+        "use_internal_ip": False,
+        "docker_config": docker_config,
+    }
+    cmd_runner = TPUCommandRunner(**args)
+
+    # Taken from tests/test_command_runner.py
+    # This mocks the response of 'docker inspect' command to return an empty JSON array.
+    # This simulates the scenario where the Docker image has no set environment
+    # variables, allowing us to test the subsequent code for handling this case.
+    process_runner.respond_to_call("json .Config.Env", 2 * ["[]"])
+    cmd_runner.run_init(as_head=True, file_mounts={}, sync_run_yet=True)
+    process_runner.assert_has_call("1.2.3.4", pattern="docker")
+
+
+def test_max_active_connections_env_var():
+    num_workers = 2
+    process_runner = MockProcessRunner()
+    provider = MockProvider()
+    instance = MockTpuInstance(num_workers=num_workers)
+    provider.create_node({}, {}, 1)
+    cluster_name = "cluster"
+    docker_config = {"container_name": "container"}
+    args = {
+        "instance": instance,
+        "log_prefix": "prefix",
+        "node_id": "0",
+        "provider": provider,
+        "auth_config": auth_config,
+        "cluster_name": cluster_name,
+        "process_runner": process_runner,
+        "use_internal_ip": False,
+        "docker_config": docker_config,
+    }
+    cmd_runner = TPUCommandRunner(**args)
+    os.environ[ray_constants.RAY_TPU_MAX_CONCURRENT_CONNECTIONS_ENV_VAR] = "1"
+    num_connections = cmd_runner.num_connections
+    assert type(num_connections) is int
+    assert num_connections == 1
+
+
+def test_tpu_pod_resources():
+    num_workers = 2
+    process_runner = MockProcessRunner()
+    provider = MockProvider()
+    instance = MockTpuInstance(num_workers=num_workers)
+    provider.create_node({}, {}, 1)
+    cluster_name = "cluster"
+    args = {
+        "instance": instance,
+        "log_prefix": "prefix",
+        "node_id": "abc",
+        "provider": provider,
+        "auth_config": auth_config,
+        "cluster_name": cluster_name,
+        "process_runner": process_runner,
+        "use_internal_ip": False,
+    }
+    env_vars = {
+        ray_constants.RESOURCES_ENVIRONMENT_VARIABLE: {
+            "TPU": 4,
+            f"TPU-{_MOCK_ACCELERATOR_TYPE}-head": 1,
+        },
+    }
+
+    def test_command_run(self, environment_variables, **kwargs):
+        resources = environment_variables[ray_constants.RESOURCES_ENVIRONMENT_VARIABLE]
+        if self._worker_id == 0:
+            assert f"TPU-{_MOCK_ACCELERATOR_TYPE}-head" in resources
+        else:
+            assert f"TPU-{_MOCK_ACCELERATOR_TYPE}-head" not in resources
+
+    with patch.object(SSHCommandRunner, "run", new=test_command_run):
+        cmd_runner = TPUCommandRunner(**args)
+        cmd_runner.run(
+            "echo helloo", port_forward=[(8265, 8265)], environment_variables=env_vars
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/horovod/BUILD.bazel

@@ -0,0 +1,14 @@
+load("@rules_python//python:defs.bzl", "py_test")
+
+py_test(
+    name = "test_horovod",
+    size = "medium",
+    srcs = ["test_horovod.py"],
+    tags = [
+        "compat",
+        "exclusive",
+        "manual",
+        "team:ml",
+    ],
+    deps = ["//:ray_lib"],
+)

python/ray/tests/horovod/horovod_example.py

@@ -0,0 +1,241 @@
+# This file is duplicated in release/ml_user_tests/horovod
+import argparse
+import os
+
+import horovod.torch as hvd
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torch.utils.data.distributed
+from filelock import FileLock
+from horovod.ray import RayExecutor
+from torchvision import datasets, transforms
+
+
+def metric_average(val, name):
+    tensor = torch.tensor(val)
+    avg_tensor = hvd.allreduce(tensor, name=name)
+    return avg_tensor.item()
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
+        self.conv2_drop = nn.Dropout2d()
+        self.fc1 = nn.Linear(320, 50)
+        self.fc2 = nn.Linear(50, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x)
+
+
+def train_fn(
+    data_dir=None,
+    seed=42,
+    use_cuda=False,
+    batch_size=64,
+    use_adasum=False,
+    lr=0.01,
+    momentum=0.5,
+    num_epochs=10,
+    log_interval=10,
+):
+    # Horovod: initialize library.
+    hvd.init()
+    torch.manual_seed(seed)
+
+    if use_cuda:
+        # Horovod: pin GPU to local rank.
+        torch.cuda.set_device(hvd.local_rank())
+        torch.cuda.manual_seed(seed)
+
+    # Horovod: limit # of CPU threads to be used per worker.
+    torch.set_num_threads(1)
+
+    kwargs = {"num_workers": 1, "pin_memory": True} if use_cuda else {}
+    data_dir = data_dir or "./data"
+    with FileLock(os.path.expanduser("~/.horovod_lock")):
+        train_dataset = datasets.MNIST(
+            data_dir,
+            train=True,
+            download=True,
+            transform=transforms.Compose(
+                [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+            ),
+        )
+    # Horovod: use DistributedSampler to partition the training data.
+    train_sampler = torch.utils.data.distributed.DistributedSampler(
+        train_dataset, num_replicas=hvd.size(), rank=hvd.rank()
+    )
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=batch_size, sampler=train_sampler, **kwargs
+    )
+
+    model = Net()
+
+    # By default, Adasum doesn't need scaling up learning rate.
+    lr_scaler = hvd.size() if not use_adasum else 1
+
+    if use_cuda:
+        # Move model to GPU.
+        model.cuda()
+        # If using GPU Adasum allreduce, scale learning rate by local_size.
+        if use_adasum and hvd.nccl_built():
+            lr_scaler = hvd.local_size()
+
+    # Horovod: scale learning rate by lr_scaler.
+    optimizer = optim.SGD(model.parameters(), lr=lr * lr_scaler, momentum=momentum)
+
+    # Horovod: wrap optimizer with DistributedOptimizer.
+    optimizer = hvd.DistributedOptimizer(
+        optimizer,
+        named_parameters=model.named_parameters(),
+        op=hvd.Adasum if use_adasum else hvd.Average,
+    )
+
+    for epoch in range(1, num_epochs + 1):
+        model.train()
+        # Horovod: set epoch to sampler for shuffling.
+        train_sampler.set_epoch(epoch)
+        for batch_idx, (data, target) in enumerate(train_loader):
+            if use_cuda:
+                data, target = data.cuda(), target.cuda()
+            optimizer.zero_grad()
+            output = model(data)
+            loss = F.nll_loss(output, target)
+            loss.backward()
+            optimizer.step()
+            if batch_idx % log_interval == 0:
+                # Horovod: use train_sampler to determine the number of
+                # examples in this worker's partition.
+                print(
+                    "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
+                        epoch,
+                        batch_idx * len(data),
+                        len(train_sampler),
+                        100.0 * batch_idx / len(train_loader),
+                        loss.item(),
+                    )
+                )
+
+
+def main(
+    num_workers, use_gpu, timeout_s=30, placement_group_timeout_s=100, kwargs=None
+):
+    kwargs = kwargs or {}
+    if use_gpu:
+        kwargs["use_cuda"] = True
+    settings = RayExecutor.create_settings(
+        timeout_s=timeout_s, placement_group_timeout_s=placement_group_timeout_s
+    )
+    executor = RayExecutor(settings, use_gpu=use_gpu, num_workers=num_workers)
+    executor.start()
+    executor.run(train_fn, kwargs=kwargs)
+
+
+if __name__ == "__main__":
+    # Training settings
+    parser = argparse.ArgumentParser(
+        description="PyTorch MNIST Example",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=64,
+        metavar="N",
+        help="input batch size for training (default: 64)",
+    )
+    parser.add_argument(
+        "--num-epochs",
+        type=int,
+        default=5,
+        metavar="N",
+        help="number of epochs to train (default: 10)",
+    )
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=0.01,
+        metavar="LR",
+        help="learning rate (default: 0.01)",
+    )
+    parser.add_argument(
+        "--momentum",
+        type=float,
+        default=0.5,
+        metavar="M",
+        help="SGD momentum (default: 0.5)",
+    )
+    parser.add_argument(
+        "--use-cuda", action="store_true", default=False, help="enables CUDA training"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=42, metavar="S", help="random seed (default: 42)"
+    )
+    parser.add_argument(
+        "--log-interval",
+        type=int,
+        default=10,
+        metavar="N",
+        help="how many batches to wait before logging training status",
+    )
+    parser.add_argument(
+        "--use-adasum",
+        action="store_true",
+        default=False,
+        help="use adasum algorithm to do reduction",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=4,
+        help="Number of Ray workers to use for training.",
+    )
+    parser.add_argument(
+        "--data-dir",
+        help="location of the training dataset in the local filesystem ("
+        "will be downloaded if needed)",
+    )
+    parser.add_argument(
+        "--address",
+        required=False,
+        type=str,
+        default=None,
+        help="Address of Ray cluster.",
+    )
+
+    args = parser.parse_args()
+
+    import ray
+
+    if args.address:
+        ray.init(args.address)
+    else:
+        ray.init()
+
+    kwargs = {
+        "data_dir": args.data_dir,
+        "seed": args.seed,
+        "use_cuda": args.use_cuda if args.use_cuda else False,
+        "batch_size": args.batch_size,
+        "use_adasum": args.use_adasum if args.use_adasum else False,
+        "lr": args.lr,
+        "momentum": args.momentum,
+        "num_epochs": args.num_epochs,
+        "log_interval": args.log_interval,
+    }
+
+    main(
+        num_workers=args.num_workers,
+        use_gpu=args.use_cuda if args.use_cuda else False,
+        kwargs=kwargs,
+    )

python/ray/tests/horovod/test_horovod.py

@@ -0,0 +1,98 @@
+import sys
+
+import pytest
+import torch
+
+import ray
+from ray.util.client.ray_client_helpers import ray_start_client_server
+
+pytest.importorskip("horovod")
+
+try:
+    from horovod.common.util import gloo_built
+    from horovod.ray.runner import RayExecutor
+except ImportError:
+    pass  # This shouldn't be reached - the test should be skipped.
+
+
+# For each test, run it once with ray.init() and again with ray client.
+@pytest.fixture(params=[False, True])
+def ray_start_4_cpus(request):
+    if request.param:
+        assert not ray.util.client.ray.is_connected()
+        with ray_start_client_server(ray_init_kwargs={"num_cpus": 3}):
+            assert ray.util.client.ray.is_connected()
+            yield
+    else:
+        ray.init(num_cpus=4)
+        yield
+        ray.shutdown()
+
+
+def _train(batch_size=32, batch_per_iter=10):
+    import timeit
+
+    import horovod.torch as hvd
+    import torch.nn.functional as F
+    import torch.optim as optim
+    import torch.utils.data.distributed
+
+    hvd.init()
+
+    # Set up fixed fake data
+    data = torch.randn(batch_size, 2)
+    target = torch.LongTensor(batch_size).random_() % 2
+
+    model = torch.nn.Sequential(torch.nn.Linear(2, 2))
+    optimizer = optim.SGD(model.parameters(), lr=0.01)
+
+    # Horovod: wrap optimizer with DistributedOptimizer.
+    optimizer = hvd.DistributedOptimizer(
+        optimizer, named_parameters=model.named_parameters()
+    )
+
+    # Horovod: broadcast parameters & optimizer state.
+    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
+    hvd.broadcast_optimizer_state(optimizer, root_rank=0)
+
+    def benchmark_step():
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.cross_entropy(output, target)
+        loss.backward()
+        optimizer.step()
+
+    timeit.timeit(benchmark_step, number=batch_per_iter)
+    return hvd.local_rank()
+
+
+@pytest.mark.skipif(not gloo_built(), reason="Gloo is required for Ray integration")
+def test_train(ray_start_4_cpus):
+    def simple_fn(worker):
+        local_rank = _train()
+        return local_rank
+
+    setting = RayExecutor.create_settings(timeout_s=30)
+    hjob = RayExecutor(setting, num_workers=3, use_gpu=torch.cuda.is_available())
+    hjob.start()
+    result = hjob.execute(simple_fn)
+    assert set(result) == {0, 1, 2}
+    result = ray.get(hjob.run_remote(simple_fn, args=[None]))
+    assert set(result) == {0, 1, 2}
+    hjob.shutdown()
+
+
+@pytest.mark.skipif(not gloo_built(), reason="Gloo is required for Ray integration")
+def test_horovod_example(ray_start_4_cpus):
+    from ray.tests.horovod.horovod_example import main
+
+    kwargs = {
+        "data_dir": "./data",
+        "num_epochs": 1,
+    }
+
+    main(num_workers=1, use_gpu=False, kwargs=kwargs)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__] + sys.argv[1:]))

python/ray/tests/kuberay/Dockerfile

@@ -0,0 +1,8 @@
+# Use the latest Ray master as base.
+FROM rayproject/ray:nightly-py310
+# Invalidate the cache so that fresh code is pulled in the next step.
+ARG BUILD_DATE
+# Retrieve your development code.
+ADD . ray
+# Install symlinks to your modified Python code.
+RUN python ray/python/ray/setup-dev.py -y

python/ray/tests/kuberay/README.md

@@ -0,0 +1,42 @@
+# How to run the KubeRay autoscaling test
+
+This page provides suggestions on running the test `test_autoscaling_e2e` locally.
+You might want to do this if your PR is breaking this test in CI and you want to debug why.
+
+Running the test must happen in stages:
+
+1. Tear down any running `kind` cluster
+2. Remove the existing ray docker image that will be deployed to the cluster
+3. Build a new docker image containing the local ray repository
+4. Create a new `kind` cluster
+5. Load the docker image into the cluster
+6. Set up kuberay
+7. Run the test
+
+To help with this, there is a `Dockerfile` and a `rune2e.sh` bash script which
+together run these things for you.
+
+## Test requirements
+
+1. Ensure `kind` and `kustomize` are both installed
+2. Run `ray/autoscaler/kuberay/init-config.sh` to clone `ray-project/kuberay`,
+   which contains config files needed to set up kuberay.
+3. Finally, make sure that the `Dockerfile` is using the same python version as
+   what you're using to run the test. By default, this dockerfile is built using
+   the `rayproject/ray:nightly-py310` build.
+4. Modify `EXAMPLE_CLUSTER_PATH` in `test_autoscaling_e2e.py`.
+
+Now you're ready to run the test.
+
+## Running the test
+
+Run `./rune2e.sh` to run the test.
+
+The test itself does not tear down resources on failure; you can
+- examine a Ray cluster from a failed test (`kubectl get pods`, `kubectl get pod`, `kubectl get raycluster`)
+- view all logs (`kubectl logs <head pod name>`) or just logs associated with the autoscaler (`kubectl logs <head pod name> -c autoscaler`)
+- delete the Ray cluster (`kubectl delete raycluster -A`)
+- rerun the test without tearing the operator down (`RAY_IMAGE=<registry>/<repo>:<tag> python test_autoscaling_e2e.py`)
+- tear down the operator when you're done `python setup/teardown_kuberay.py`
+- copy files from a pod to your filesystem (`kubectl cp <pod>:/path/to/file /target/path/in/local/filesystem`)
+- access a bash prompt inside the pod (`kubectl exec -it <pod> bash`)

python/ray/tests/kuberay/rune2e.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+set -x
+RAY_IMAGE=rayproject/autoscaling_e2e_test_image
+kind delete cluster
+docker image rm $RAY_IMAGE
+
+pushd ../../../..
+docker build --progress=plain --build-arg BUILD_DATE="$(date +%Y-%m-%d:%H:%M:%S)" -t $RAY_IMAGE -f ./python/ray/tests/kuberay/Dockerfile . || exit
+popd || exit
+
+kind create cluster || exit
+kind load docker-image $RAY_IMAGE || exit
+python setup/setup_kuberay.py
+RAY_IMAGE=$RAY_IMAGE python test_autoscaling_e2e.py

python/ray/tests/kuberay/scripts/check_cpu_and_memory.py

@@ -0,0 +1,17 @@
+import ray
+
+
+def main():
+    """This script runs in a container with 1 CPU limit and 1G memory limit.
+    Validate that Ray reads the correct limits.
+    """
+    cpu_limit = ray._private.utils.get_num_cpus()
+    mem_limit_gb = round(ray._common.utils.get_system_memory() / 10**9, 2)
+    assert cpu_limit == 1, cpu_limit
+    assert mem_limit_gb == 2.00, mem_limit_gb
+    print(f"Confirmed cpu limit {cpu_limit}.")
+    print(f"Confirmed memory limit {mem_limit_gb} gigabyte.")
+
+
+if __name__ == "__main__":
+    main()

python/ray/tests/kuberay/scripts/gpu_actor_placement.py

@@ -0,0 +1,18 @@
+import ray
+
+
+def main():
+    """Requests placement of a GPU actor."""
+
+    @ray.remote(num_gpus=1, num_cpus=1)
+    class GPUActor:
+        def where_am_i(self):
+            assert len(ray.get_gpu_ids()) == 1
+            return "on-a-gpu-node"
+
+    GPUActor.options(name="gpu_actor", lifetime="detached").remote()
+
+
+if __name__ == "__main__":
+    ray.init("auto", namespace="gpu-test")
+    main()

python/ray/tests/kuberay/scripts/gpu_actor_validation.py

@@ -0,0 +1,14 @@
+import ray
+
+
+def main():
+    """Confirms placement of a GPU actor."""
+    gpu_actor = ray.get_actor("gpu_actor")
+    actor_response = ray.get(gpu_actor.where_am_i.remote())
+    return actor_response
+
+
+if __name__ == "__main__":
+    ray.init("auto", namespace="gpu-test")
+    out = main()
+    print(out)

python/ray/tests/kuberay/scripts/non_terminated_nodes_count.py

@@ -0,0 +1,26 @@
+import ray
+from ray.autoscaler._private.kuberay.autoscaling_config import _generate_provider_config
+from ray.autoscaler._private.providers import _get_node_provider
+
+
+@ray.remote
+def count_non_terminated_nodes() -> int:
+    """Get the count of non terminated nodes for the Ray cluster raycluster-autoscaler
+    in namespace default.
+    """
+    provider_config = _generate_provider_config(ray_cluster_namespace="default")
+    kuberay_node_provider = _get_node_provider(
+        provider_config=provider_config, cluster_name="raycluster-autoscaler"
+    )
+    nodes = kuberay_node_provider.non_terminated_nodes({})
+    return len(nodes)
+
+
+def main() -> int:
+    return ray.get(count_non_terminated_nodes.remote())
+
+
+if __name__ == "__main__":
+    ray.init("auto")
+    out = main()
+    print(out)

python/ray/tests/kuberay/scripts/scale_down.py

@@ -0,0 +1,37 @@
+import ray
+from ray._common import test_utils
+
+
+def main():
+    """Removes CPU request, removes GPU actor.
+    Waits for autoscaler scale-down events to get emitted to stdout.
+
+    The worker idle timeout is set to 10 seconds and the autoscaler's update interval is
+    5 seconds, so it should be enough to wait 15 seconds.
+    """
+
+    # Before scale-down.
+    cluster_resources = ray.cluster_resources()
+    assert cluster_resources.get("CPU", 0) > 0, cluster_resources
+    assert cluster_resources.get("GPU", 0) > 0, cluster_resources
+
+    # Remove resource demands
+    ray.autoscaler.sdk.request_resources(num_cpus=0)
+    gpu_actor = ray.get_actor("gpu_actor")
+    ray.kill(gpu_actor)
+
+    # Wait for scale-down to happen.
+    def verify():
+        cluster_resources = ray.cluster_resources()
+        # From head node
+        assert cluster_resources.get("CPU", 0) == 1, cluster_resources
+        assert cluster_resources.get("GPU", 0) == 0, cluster_resources
+
+        return True
+
+    test_utils.wait_for_condition(verify, timeout=60, retry_interval_ms=2000)
+
+
+if __name__ == "__main__":
+    ray.init("auto", namespace="gpu-test")
+    main()

python/ray/tests/kuberay/scripts/scale_up.py

@@ -0,0 +1,27 @@
+import ray
+from ray._common.test_utils import wait_for_condition
+
+
+def main():
+    """Submits CPU request"""
+    ray.autoscaler.sdk.request_resources(num_cpus=2)
+    from ray.autoscaler.v2.sdk import get_cluster_status
+    from ray.autoscaler.v2.utils import ClusterStatusFormatter
+
+    gcs_address = ray.get_runtime_context().gcs_address
+
+    def verify():
+        cluster_resources = ray.cluster_resources()
+
+        cluster_status = get_cluster_status(gcs_address)
+        print(ClusterStatusFormatter.format(cluster_status, verbose=True))
+        assert cluster_resources.get("CPU", 0) == 2, cluster_resources
+
+        return True
+
+    wait_for_condition(verify, timeout=60, retry_interval_ms=2000)
+
+
+if __name__ == "__main__":
+    ray.init("auto")
+    main()

python/ray/tests/kuberay/scripts/scale_up_custom.py

@@ -0,0 +1,32 @@
+import time
+
+import ray
+
+
+def main():
+    """Submits custom resource request.
+
+    Also, validates runtime env data submitted with the Ray Job that executes
+    this script.
+    """
+
+    # Workers and head are annotated as having 5 "Custom2" capacity each,
+    # so this should trigger upscaling of two workers.
+    # (One of the bundles will be "placed" on the head.)
+    ray.autoscaler.sdk.request_resources(
+        bundles=[{"Custom2": 3}, {"Custom2": 3}, {"Custom2": 3}]
+    )
+
+    while (
+        ray.cluster_resources().get("Custom2", 0) < 3
+        and ray.cluster_resources().get("Custom2", 0) < 6
+    ):
+        time.sleep(0.1)
+
+    # Output something to validate the job logs.
+    print("Submitted custom scale request!")
+
+
+if __name__ == "__main__":
+    ray.init("auto")
+    main()

python/ray/tests/kuberay/setup/raycluster_test.yaml

@@ -0,0 +1,12 @@
+---
+apiVersion: ray.io/v1
+kind: RayCluster
+metadata:
+  name: raycluster-test
+spec:
+  headGroupSpec:
+    serviceType: ClusterIP
+    template:
+      spec:
+        containers:
+          - name: ray-test

python/ray/tests/kuberay/setup/setup_kuberay.py

@@ -0,0 +1,8 @@
+from ray.tests.kuberay.utils import (
+    setup_kuberay_operator,
+    wait_for_raycluster_crd,
+)
+
+if __name__ == "__main__":
+    setup_kuberay_operator()
+    wait_for_raycluster_crd()

python/ray/tests/kuberay/setup/teardown_kuberay.py

@@ -0,0 +1,4 @@
+from ray.tests.kuberay.utils import teardown_kuberay_operator
+
+if __name__ == "__main__":
+    teardown_kuberay_operator()

python/ray/tests/kuberay/test_autoscaling_config.py

@@ -0,0 +1,959 @@
+import copy
+import platform
+import sys
+from pathlib import Path
+from typing import Any, Dict, Optional, Type
+from unittest import mock
+
+import pytest
+import requests
+import yaml
+
+from ray.autoscaler._private.kuberay.autoscaling_config import (
+    GKE_TPU_ACCELERATOR_LABEL,
+    GKE_TPU_TOPOLOGY_LABEL,
+    AutoscalingConfigProducer,
+    _derive_autoscaling_config_from_ray_cr,
+    _get_custom_resources,
+    _get_num_tpus,
+    _get_ray_resources_from_group_spec,
+    _round_up_k8s_quantity,
+)
+from ray.autoscaler._private.kuberay.utils import tpu_node_selectors_to_type
+
+AUTOSCALING_CONFIG_MODULE_PATH = "ray.autoscaler._private.kuberay.autoscaling_config"
+
+
+def get_basic_ray_cr() -> dict:
+    """Returns the example Ray CR included in the Ray documentation,
+    modified to include a GPU worker group and a TPU worker group.
+    """
+    cr_path = str(
+        Path(__file__).resolve().parents[2]
+        / "autoscaler"
+        / "kuberay"
+        / "ray-cluster.complete.yaml"
+    )
+    config = yaml.safe_load(open(cr_path).read())
+    gpu_group = copy.deepcopy(config["spec"]["workerGroupSpecs"][0])
+    gpu_group["groupName"] = "gpu-group"
+    gpu_group["template"]["spec"]["containers"][0]["resources"]["limits"].setdefault(
+        "nvidia.com/gpu", 3
+    )
+    gpu_group["maxReplicas"] = 200
+    config["spec"]["workerGroupSpecs"].append(gpu_group)
+    tpu_group = copy.deepcopy(config["spec"]["workerGroupSpecs"][0])
+    tpu_group["groupName"] = "tpu-group"
+    tpu_group["template"]["spec"]["containers"][0]["resources"]["limits"].setdefault(
+        "google.com/tpu", 4
+    )
+    tpu_group["template"]["spec"]["nodeSelector"] = {}
+    tpu_group["template"]["spec"]["nodeSelector"][
+        "cloud.google.com/gke-tpu-topology"
+    ] = "2x2x2"
+    tpu_group["template"]["spec"]["nodeSelector"][
+        "cloud.google.com/gke-tpu-accelerator"
+    ] = "tpu-v4-podslice"
+    tpu_group["maxReplicas"] = 4
+    tpu_group["numOfHosts"] = 2
+    config["spec"]["workerGroupSpecs"].append(tpu_group)
+    return config
+
+
+def _get_basic_autoscaling_config() -> dict:
+    """The expected autoscaling derived from the example Ray CR."""
+    return {
+        "cluster_name": "raycluster-complete",
+        "provider": {
+            "disable_node_updaters": True,
+            "disable_launch_config_check": True,
+            "foreground_node_launch": True,
+            "worker_liveness_check": False,
+            "namespace": "default",
+            "type": "kuberay",
+        },
+        "available_node_types": {
+            "headgroup": {
+                "labels": {},
+                "max_workers": 0,
+                "min_workers": 0,
+                "node_config": {},
+                "resources": {
+                    "CPU": 1,
+                    "memory": 1000000000,
+                    "Custom1": 1,
+                    "Custom2": 5,
+                },
+            },
+            "small-group": {
+                "labels": {},
+                "max_workers": 300,
+                "min_workers": 0,
+                "node_config": {},
+                "resources": {
+                    "CPU": 1,
+                    "memory": 536870912,
+                    "Custom2": 5,
+                    "Custom3": 1,
+                },
+            },
+            # Same as "small-group" with a GPU resource entry added
+            # and modified max_workers.
+            "gpu-group": {
+                "labels": {},
+                "max_workers": 200,
+                "min_workers": 0,
+                "node_config": {},
+                "resources": {
+                    "CPU": 1,
+                    "memory": 536870912,
+                    "Custom2": 5,
+                    "Custom3": 1,
+                    "GPU": 3,
+                },
+            },
+            # Same as "small-group" with a TPU resource entry added
+            # and modified max_workers and node_config.
+            "tpu-group": {
+                "labels": {},
+                "max_workers": 8,
+                "min_workers": 0,
+                "node_config": {},
+                "resources": {
+                    "CPU": 1,
+                    "memory": 536870912,
+                    "Custom2": 5,
+                    "Custom3": 1,
+                    "TPU": 4,
+                    "TPU-v4-16-head": 1,
+                },
+            },
+        },
+        "auth": {},
+        "cluster_synced_files": [],
+        "file_mounts": {},
+        "file_mounts_sync_continuously": False,
+        "head_node_type": "headgroup",
+        "head_setup_commands": [],
+        "head_start_ray_commands": [],
+        "idle_timeout_minutes": 1.0,
+        "initialization_commands": [],
+        "max_workers": 508,
+        "setup_commands": [],
+        "upscaling_speed": 1000,
+        "worker_setup_commands": [],
+        "worker_start_ray_commands": [],
+    }
+
+
+def _get_ray_cr_no_cpu_error() -> dict:
+    """Incorrectly formatted Ray CR without num-cpus rayStartParam and without resource
+    limits. Autoscaler should raise an error when reading this.
+    """
+    cr = get_basic_ray_cr()
+    # Verify that the num-cpus rayStartParam is not present for the worker type.
+    assert "num-cpus" not in cr["spec"]["workerGroupSpecs"][0]["rayStartParams"]
+    del cr["spec"]["workerGroupSpecs"][0]["template"]["spec"]["containers"][0][
+        "resources"
+    ]["limits"]["cpu"]
+    del cr["spec"]["workerGroupSpecs"][0]["template"]["spec"]["containers"][0][
+        "resources"
+    ]["requests"]["cpu"]
+    return cr
+
+
+def _get_no_cpu_error() -> str:
+    return (
+        "Autoscaler failed to detect `CPU` resources for group small-group."
+        "\nSet the `--num-cpus` rayStartParam and/or "
+        "the CPU resource limit for the Ray container."
+    )
+
+
+def _get_ray_cr_with_overrides() -> dict:
+    """CR with memory, cpu, and gpu overrides from rayStartParams."""
+    cr = get_basic_ray_cr()
+    cr["spec"]["workerGroupSpecs"][0]["rayStartParams"]["memory"] = "300000000"
+    # num-gpus rayStartParam with no gpus in container limits
+    cr["spec"]["workerGroupSpecs"][0]["rayStartParams"]["num-gpus"] = "100"
+    # num-gpus rayStartParam overriding gpus in container limits
+    cr["spec"]["workerGroupSpecs"][1]["rayStartParams"]["num-gpus"] = "100"
+    cr["spec"]["workerGroupSpecs"][0]["rayStartParams"]["num-cpus"] = "100"
+    return cr
+
+
+def _get_autoscaling_config_with_overrides() -> dict:
+    """Autoscaling config with memory and gpu annotations."""
+    config = _get_basic_autoscaling_config()
+    config["available_node_types"]["small-group"]["resources"]["memory"] = 300000000
+    config["available_node_types"]["small-group"]["resources"]["GPU"] = 100
+    config["available_node_types"]["small-group"]["resources"]["CPU"] = 100
+    config["available_node_types"]["gpu-group"]["resources"]["GPU"] = 100
+    return config
+
+
+def _get_ray_cr_with_autoscaler_options() -> dict:
+    cr = get_basic_ray_cr()
+    cr["spec"]["autoscalerOptions"] = {
+        "upscalingMode": "Conservative",
+        "idleTimeoutSeconds": 300,
+    }
+    return cr
+
+
+def _get_ray_cr_with_tpu_custom_resource() -> dict:
+    cr = get_basic_ray_cr()
+    cr["spec"]["workerGroupSpecs"][2]["rayStartParams"][
+        "resources"
+    ] = '"{"TPU": 4, "Custom2": 5, "Custom3": 1}"'
+    # remove google.com/tpu k8s resource Pod limit
+    del cr["spec"]["workerGroupSpecs"][2]["template"]["spec"]["containers"][0][
+        "resources"
+    ]["limits"]["google.com/tpu"]
+
+    return cr
+
+
+def _get_ray_cr_with_tpu_k8s_resource_limit_and_custom_resource() -> dict:
+    cr = get_basic_ray_cr()
+    cr["spec"]["workerGroupSpecs"][2]["rayStartParams"][
+        "resources"
+    ] = '"{"TPU": 4, "Custom2": 5, "Custom3": 1}"'
+
+    return cr
+
+
+def _get_ray_cr_with_top_level_labels() -> dict:
+    """CR with a top-level `labels` field."""
+    cr = get_basic_ray_cr()
+    # This top-level structured labels take priority.
+    cr["spec"]["workerGroupSpecs"][0]["labels"] = {"instance-type": "mx5"}
+
+    # rayStartParams labels field should be ignored.
+    cr["spec"]["workerGroupSpecs"][0]["rayStartParams"]["labels"] = "instance-type=n2"
+    return cr
+
+
+def _get_autoscaling_config_with_top_level_labels() -> dict:
+    config = _get_basic_autoscaling_config()
+    config["available_node_types"]["small-group"]["labels"] = {"instance-type": "mx5"}
+    return config
+
+
+def _get_ray_cr_with_invalid_top_level_labels() -> dict:
+    """CR with a syntactically invalid top-level `labels` field."""
+    cr = get_basic_ray_cr()
+    cr["spec"]["workerGroupSpecs"][0]["labels"] = {"!!invalid-key!!": "some-value"}
+    return cr
+
+
+def _get_ray_cr_with_top_level_resources() -> dict:
+    """CR with a top-level `resources` field to test priority."""
+    cr = get_basic_ray_cr()
+
+    # The top-level resources field should take priority.
+    cr["spec"]["workerGroupSpecs"][1]["resources"] = {
+        "CPU": "16",
+        "GPU": "8",
+        "memory": "2Gi",
+        "CustomResource": "99",
+    }
+    # These rayStartParams should be ignored.
+    cr["spec"]["workerGroupSpecs"][1]["rayStartParams"]["num-cpus"] = "1"
+    cr["spec"]["workerGroupSpecs"][1]["rayStartParams"]["memory"] = "100000"
+    cr["spec"]["workerGroupSpecs"][1]["rayStartParams"]["num-gpus"] = "2"
+    cr["spec"]["workerGroupSpecs"][1]["rayStartParams"][
+        "resources"
+    ] = '"{"Custom2": 1}"'
+    return cr
+
+
+def _get_autoscaling_config_with_top_level_resources() -> dict:
+    config = _get_basic_autoscaling_config()
+
+    config["available_node_types"]["gpu-group"]["resources"] = {
+        "CPU": 16,
+        "GPU": 8,
+        "memory": 2147483648,
+        "CustomResource": 99,
+    }
+    return config
+
+
+def _get_ray_cr_with_top_level_tpu_resource() -> dict:
+    """CR with a top-level `resources` field for the TPU custom resource."""
+    cr = _get_ray_cr_with_tpu_k8s_resource_limit_and_custom_resource()
+
+    # The top-level field should take priority.
+    cr["spec"]["workerGroupSpecs"][2]["resources"] = {"TPU": "8"}
+    return cr
+
+
+def _get_ray_cr_with_no_tpus() -> dict:
+    cr = get_basic_ray_cr()
+    # remove TPU worker group
+    cr["spec"]["workerGroupSpecs"].pop(2)
+
+    return cr
+
+
+def _get_ray_cr_with_only_requests() -> dict:
+    """CR contains only resource requests"""
+    cr = get_basic_ray_cr()
+
+    for group in [cr["spec"]["headGroupSpec"]] + cr["spec"]["workerGroupSpecs"]:
+        for container in group["template"]["spec"]["containers"]:
+            container["resources"]["requests"] = container["resources"]["limits"]
+            del container["resources"]["limits"]
+    return cr
+
+
+def _get_ray_cr_with_labels() -> dict:
+    """CR with labels in rayStartParams of head and worker groups."""
+    cr = get_basic_ray_cr()
+
+    # Pass invalid labels to the head group to test error handling.
+    cr["spec"]["headGroupSpec"]["rayStartParams"]["labels"] = "!!ray.io/node-group=,"
+    # Pass valid labels to each of the worker groups.
+    cr["spec"]["workerGroupSpecs"][0]["rayStartParams"][
+        "labels"
+    ] = "ray.io/availability-region=us-central2, ray.io/market-type=spot"
+    cr["spec"]["workerGroupSpecs"][1]["rayStartParams"][
+        "labels"
+    ] = "ray.io/accelerator-type=A100"
+    cr["spec"]["workerGroupSpecs"][2]["rayStartParams"][
+        "labels"
+    ] = "ray.io/accelerator-type=TPU-V4"
+    return cr
+
+
+def _get_autoscaling_config_with_labels() -> dict:
+    """Autoscaling config with parsed labels for each group."""
+    config = _get_basic_autoscaling_config()
+
+    # Since we passed invalid labels to the head group `rayStartParams`,
+    # we expect an empty dictionary in the autoscaling config.
+    config["available_node_types"]["headgroup"]["labels"] = {}
+    config["available_node_types"]["small-group"]["labels"] = {
+        "ray.io/availability-region": "us-central2",
+        "ray.io/market-type": "spot",
+    }
+    config["available_node_types"]["gpu-group"]["labels"] = {
+        "ray.io/accelerator-type": "A100"
+    }
+    config["available_node_types"]["tpu-group"]["labels"] = {
+        "ray.io/accelerator-type": "TPU-V4"
+    }
+    return config
+
+
+def _get_autoscaling_config_with_options() -> dict:
+    config = _get_basic_autoscaling_config()
+    config["upscaling_speed"] = 1
+    config["idle_timeout_minutes"] = 5.0
+    return config
+
+
+def _get_tpu_group_with_no_node_selectors() -> dict[str, Any]:
+    cr = get_basic_ray_cr()
+    tpu_group = cr["spec"]["workerGroupSpecs"][2]
+    tpu_group["template"]["spec"].pop("nodeSelector", None)
+    return tpu_group
+
+
+def _get_tpu_group_without_accelerator_node_selector() -> dict[str, Any]:
+    cr = get_basic_ray_cr()
+    tpu_group = cr["spec"]["workerGroupSpecs"][2]
+    tpu_group["template"]["spec"]["nodeSelector"].pop(GKE_TPU_ACCELERATOR_LABEL, None)
+    return tpu_group
+
+
+def _get_tpu_group_without_topology_node_selector() -> dict[str, Any]:
+    cr = get_basic_ray_cr()
+    tpu_group = cr["spec"]["workerGroupSpecs"][2]
+    tpu_group["template"]["spec"]["nodeSelector"].pop(GKE_TPU_TOPOLOGY_LABEL, None)
+    return tpu_group
+
+
+def _get_tpu_group_with_v7x_node_selectors() -> dict[str, Any]:
+    cr = get_basic_ray_cr()
+    tpu_group = cr["spec"]["workerGroupSpecs"][2]
+    tpu_group["template"]["spec"]["nodeSelector"][GKE_TPU_TOPOLOGY_LABEL] = "2x2x2"
+    tpu_group["template"]["spec"]["nodeSelector"][GKE_TPU_ACCELERATOR_LABEL] = "tpu7x"
+    return tpu_group
+
+
+def _get_ray_cr_with_tpu_v7x() -> dict[str, Any]:
+    cr = get_basic_ray_cr()
+    cr["spec"]["workerGroupSpecs"][2] = _get_tpu_group_with_v7x_node_selectors()
+    return cr
+
+
+def _get_autoscaling_config_with_v7x() -> dict[str, Any]:
+    config = _get_basic_autoscaling_config()
+    config["available_node_types"]["tpu-group"]["resources"]["TPU-v7x-16-head"] = 1
+    config["available_node_types"]["tpu-group"]["resources"].pop("TPU-v4-16-head", None)
+    return config
+
+
+def _get_tpu_group_with_v5litepod_node_selectors() -> dict[str, Any]:
+    cr = get_basic_ray_cr()
+    tpu_group = cr["spec"]["workerGroupSpecs"][2]
+    tpu_group["template"]["spec"]["nodeSelector"][GKE_TPU_TOPOLOGY_LABEL] = "2x4"
+    tpu_group["template"]["spec"]["nodeSelector"][
+        GKE_TPU_ACCELERATOR_LABEL
+    ] = "tpu-v5-lite-podslice"
+    return tpu_group
+
+
+def _get_ray_cr_with_tpu_v5litepod() -> dict[str, Any]:
+    cr = get_basic_ray_cr()
+    cr["spec"]["workerGroupSpecs"][2] = _get_tpu_group_with_v5litepod_node_selectors()
+    return cr
+
+
+def _get_autoscaling_config_with_v5litepod() -> dict[str, Any]:
+    config = _get_basic_autoscaling_config()
+    config["available_node_types"]["tpu-group"]["resources"]["TPU-v5litepod-8-head"] = 1
+    config["available_node_types"]["tpu-group"]["resources"].pop("TPU-v4-16-head", None)
+    return config
+
+
+@pytest.mark.parametrize(
+    "input,output",
+    [
+        # There's no particular discipline to these test cases.
+        ("100m", 1),
+        ("15001m", 16),
+        ("2", 2),
+        ("100Mi", 104857600),
+        ("1G", 1000000000),
+    ],
+)
+def test_resource_quantity(input: str, output: int):
+    assert _round_up_k8s_quantity(input) == output, output
+
+
+PARAM_ARGS = ",".join(
+    [
+        "ray_cr_in",
+        "expected_config_out",
+        "expected_error",
+        "expected_error_message",
+        "expected_log_warning",
+    ]
+)
+
+TEST_DATA = (
+    []
+    if platform.system() == "Windows"
+    else [
+        pytest.param(
+            get_basic_ray_cr(),
+            _get_basic_autoscaling_config(),
+            None,
+            None,
+            None,
+            id="basic",
+        ),
+        pytest.param(
+            _get_ray_cr_with_only_requests(),
+            _get_basic_autoscaling_config(),
+            None,
+            None,
+            None,
+            id="only-requests",
+        ),
+        pytest.param(
+            _get_ray_cr_no_cpu_error(),
+            None,
+            ValueError,
+            _get_no_cpu_error(),
+            None,
+            id="no-cpu-error",
+        ),
+        pytest.param(
+            _get_ray_cr_with_overrides(),
+            _get_autoscaling_config_with_overrides(),
+            None,
+            None,
+            None,
+            id="overrides",
+        ),
+        pytest.param(
+            _get_ray_cr_with_autoscaler_options(),
+            _get_autoscaling_config_with_options(),
+            None,
+            None,
+            None,
+            id="autoscaler-options",
+        ),
+        pytest.param(
+            _get_ray_cr_with_tpu_custom_resource(),
+            _get_basic_autoscaling_config(),
+            None,
+            None,
+            None,
+            id="tpu-custom-resource",
+        ),
+        pytest.param(
+            get_basic_ray_cr(),
+            _get_basic_autoscaling_config(),
+            None,
+            None,
+            None,
+            id="tpu-k8s-resource-limit",
+        ),
+        pytest.param(
+            _get_ray_cr_with_tpu_k8s_resource_limit_and_custom_resource(),
+            _get_basic_autoscaling_config(),
+            None,
+            None,
+            None,
+            id="tpu-k8s-resource-limit-and-custom-resource",
+        ),
+        pytest.param(
+            _get_ray_cr_with_labels(),
+            _get_basic_autoscaling_config(),
+            None,
+            None,
+            "Ignoring labels: ray.io/accelerator-type=TPU-V4 set in rayStartParams for group 'tpu-group'. Group labels are supported in the top-level Labels field starting in KubeRay v1.5",
+            id="groups-with-raystartparam-labels",
+        ),
+        pytest.param(
+            _get_ray_cr_with_top_level_labels(),
+            _get_autoscaling_config_with_top_level_labels(),
+            None,
+            None,
+            "Ignoring labels: instance-type=n2 set in rayStartParams for group 'small-group'. Group labels are supported in the top-level Labels field starting in KubeRay v1.5",
+            id="groups-with-top-level-labels",
+        ),
+        pytest.param(
+            _get_ray_cr_with_invalid_top_level_labels(),
+            _get_basic_autoscaling_config(),
+            ValueError,
+            None,
+            None,
+            id="invalid-top-level-labels",
+        ),
+        pytest.param(
+            _get_ray_cr_with_tpu_v7x(),
+            _get_autoscaling_config_with_v7x(),
+            None,
+            None,
+            None,
+            id="tpu-v7x",
+        ),
+        pytest.param(
+            _get_ray_cr_with_tpu_v5litepod(),
+            _get_autoscaling_config_with_v5litepod(),
+            None,
+            None,
+            None,
+            id="tpu-v5litepod",
+        ),
+    ]
+)
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Not relevant.")
+@pytest.mark.parametrize(PARAM_ARGS, TEST_DATA)
+def test_autoscaling_config(
+    ray_cr_in: Dict[str, Any],
+    expected_config_out: Optional[Dict[str, Any]],
+    expected_error: Optional[Type[Exception]],
+    expected_error_message: Optional[str],
+    expected_log_warning: Optional[str],
+):
+    ray_cr_in["metadata"]["namespace"] = "default"
+    # Reset log_once state to ensure each test case is independent.
+    from ray.util.debug import _logged
+
+    _logged.clear()
+    with mock.patch(f"{AUTOSCALING_CONFIG_MODULE_PATH}.logger") as mock_logger:
+        if expected_error:
+            with pytest.raises(expected_error, match=expected_error_message):
+                _derive_autoscaling_config_from_ray_cr(ray_cr_in)
+        else:
+            assert (
+                _derive_autoscaling_config_from_ray_cr(ray_cr_in) == expected_config_out
+            )
+            if expected_log_warning:
+                mock_logger.warning.assert_called_with(expected_log_warning)
+            else:
+                mock_logger.warning.assert_not_called()
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Not relevant.")
+def test_cr_image_consistency():
+    """Verify that the example config uses the same Ray image for all Ray pods."""
+    cr = get_basic_ray_cr()
+
+    group_specs = [cr["spec"]["headGroupSpec"]] + cr["spec"]["workerGroupSpecs"]
+    # Head, CPU group, GPU group, TPU group.
+    assert len(group_specs) == 4
+
+    ray_containers = [
+        group_spec["template"]["spec"]["containers"][0] for group_spec in group_specs
+    ]
+
+    # All Ray containers in the example config have "ray-" in their name.
+    assert all("ray-" in ray_container["name"] for ray_container in ray_containers)
+
+    # All Ray images are from the Ray repo.
+    assert all(
+        "rayproject/ray" in ray_container["image"] for ray_container in ray_containers
+    )
+
+    # All Ray images are the same.
+    assert len({ray_container["image"] for ray_container in ray_containers}) == 1
+
+
+@pytest.mark.parametrize("exception", [Exception, requests.HTTPError])
+@pytest.mark.parametrize("num_exceptions", range(6))
+def test_autoscaling_config_fetch_retries(exception, num_exceptions):
+    """Validates retry logic in
+    AutoscalingConfigProducer._fetch_ray_cr_from_k8s_with_retries.
+    """
+
+    class MockKubernetesHttpApiClient:
+        def __init__(self):
+            self.exception_counter = 0
+
+        def get(self, *args, **kwargs):
+            if self.exception_counter < num_exceptions:
+                self.exception_counter += 1
+                raise exception
+            else:
+                return {"ok-key": "ok-value"}
+
+    class MockAutoscalingConfigProducer(AutoscalingConfigProducer):
+        def __init__(self, *args, **kwargs):
+            self.kubernetes_api_client = MockKubernetesHttpApiClient()
+            self._ray_cr_path = "rayclusters/mock"
+
+    config_producer = MockAutoscalingConfigProducer()
+    # Patch retry backoff period.
+    with mock.patch(
+        "ray.autoscaler._private.kuberay.autoscaling_config.RAYCLUSTER_FETCH_RETRY_S",
+        0,
+    ):
+        # If you hit an exception and it's not HTTPError, expect to raise.
+        # If you hit >= 5 exceptions, expect to raise.
+        # Otherwise, don't expect to raise.
+        if (
+            num_exceptions > 0 and exception != requests.HTTPError
+        ) or num_exceptions >= 5:
+            with pytest.raises(exception):
+                config_producer._fetch_ray_cr_from_k8s_with_retries()
+        else:
+            out = config_producer._fetch_ray_cr_from_k8s_with_retries()
+            assert out == {"ok-key": "ok-value"}
+
+
+TPU_TYPES_ARGS = ",".join(
+    [
+        "accelerator",
+        "topology",
+        "expected_tpu_type",
+    ]
+)
+TPU_TYPES_DATA = (
+    []
+    if platform.system() == "Windows"
+    else [
+        pytest.param(
+            "tpu-v4-podslice",
+            None,
+            None,
+            id="tpu-none-topology",
+        ),
+        pytest.param(
+            None,
+            "2x2x2",
+            None,
+            id="tpu-none-accelerator",
+        ),
+        pytest.param(
+            "tpu-v4-podslice",
+            "2x2x2",
+            "v4-16",
+            id="tpu-v4-test",
+        ),
+        pytest.param(
+            "tpu-v5-lite-device",
+            "2x2",
+            "v5litepod-4",
+            id="tpu-v5e-device-test",
+        ),
+        pytest.param(
+            "tpu-v5-lite-podslice",
+            "2x4",
+            "v5litepod-8",
+            id="tpu-v5e-podslice-test",
+        ),
+        pytest.param(
+            "tpu-v5p-slice",
+            "2x2x4",
+            "v5p-32",
+            id="tpu-v5p-test",
+        ),
+        pytest.param(
+            "tpu-v6e-slice",
+            "16x16",
+            "v6e-256",
+            id="tpu-v6e-test",
+        ),
+        pytest.param(
+            "tpu7x",
+            "2x2x2",
+            "v7x-16",
+            id="tpu-v7x-test",
+        ),
+    ]
+)
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Not relevant.")
+@pytest.mark.parametrize(TPU_TYPES_ARGS, TPU_TYPES_DATA)
+def test_tpu_node_selectors_to_type(
+    accelerator: str, topology: str, expected_tpu_type: str
+):
+    """Verify that tpu_node_selectors_to_type correctly returns TPU type from
+    TPU nodeSelectors.
+    """
+    tpu_type = tpu_node_selectors_to_type(topology, accelerator)
+    assert expected_tpu_type == tpu_type
+
+
+TPU_PARAM_ARGS = ",".join(
+    [
+        "ray_cr_in",
+        "expected_num_tpus",
+    ]
+)
+TPU_TEST_DATA = (
+    []
+    if platform.system() == "Windows"
+    else [
+        pytest.param(
+            get_basic_ray_cr(),
+            4,
+            id="tpu-k8s-resource-limits",
+        ),
+        pytest.param(
+            _get_ray_cr_with_tpu_custom_resource(),
+            4,
+            id="tpu-custom-resource",
+        ),
+        pytest.param(
+            _get_ray_cr_with_tpu_k8s_resource_limit_and_custom_resource(),
+            4,
+            id="tpu--k8s-resource-limits-and-custom-resource",
+        ),
+        pytest.param(
+            _get_ray_cr_with_no_tpus(),
+            0,
+            id="no-tpus-requested",
+        ),
+        pytest.param(
+            _get_ray_cr_with_top_level_tpu_resource(),
+            8,
+            id="tpu-top-level-resource",
+        ),
+    ]
+)
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Not relevant.")
+@pytest.mark.parametrize(TPU_PARAM_ARGS, TPU_TEST_DATA)
+def test_get_num_tpus(ray_cr_in: Dict[str, Any], expected_num_tpus: int):
+    """Verify that _get_num_tpus correctly returns the number of requested TPUs."""
+    for worker_group in ray_cr_in["spec"]["workerGroupSpecs"]:
+        group_resources = worker_group.get("resources", {})
+        ray_start_params = worker_group["rayStartParams"]
+        custom_resources = _get_custom_resources(
+            group_resources, ray_start_params, worker_group["groupName"]
+        )
+        k8s_resources = worker_group["template"]["spec"]["containers"][0]["resources"]
+
+        num_tpus = _get_num_tpus(group_resources, custom_resources, k8s_resources)
+
+        if worker_group["groupName"] == "tpu-group":
+            assert num_tpus == expected_num_tpus
+        else:
+            assert num_tpus is None
+
+
+RAY_RESOURCES_PARAM_ARGS = ",".join(
+    [
+        "group_spec",
+        "is_head",
+        "expected_resources",
+    ]
+)
+RAY_RESOURCES_TEST_DATA = (
+    []
+    if platform.system() == "Windows"
+    else [
+        pytest.param(
+            get_basic_ray_cr()["spec"]["headGroupSpec"],
+            True,
+            {
+                "CPU": 1,
+                "memory": 1000000000,
+                "Custom1": 1,
+                "Custom2": 5,
+            },
+            id="head-group",
+        ),
+        pytest.param(
+            get_basic_ray_cr()["spec"]["workerGroupSpecs"][0],
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+            },
+            id="cpu-group",
+        ),
+        pytest.param(
+            get_basic_ray_cr()["spec"]["workerGroupSpecs"][1],
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+                "GPU": 3,
+            },
+            id="gpu-group",
+        ),
+        pytest.param(
+            get_basic_ray_cr()["spec"]["workerGroupSpecs"][2],
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+                "TPU": 4,
+                "TPU-v4-16-head": 1,
+            },
+            id="tpu-group",
+        ),
+        pytest.param(
+            _get_tpu_group_with_no_node_selectors(),
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+                "TPU": 4,
+            },
+            id="tpu-group-no-node-selectors",
+        ),
+        pytest.param(
+            _get_tpu_group_without_accelerator_node_selector(),
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+                "TPU": 4,
+            },
+            id="tpu-group-no-accelerator-node-selector",
+        ),
+        pytest.param(
+            _get_tpu_group_without_topology_node_selector(),
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+                "TPU": 4,
+            },
+            id="tpu-group-no-topology-node-selector",
+        ),
+        pytest.param(
+            _get_tpu_group_with_v7x_node_selectors(),
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+                "TPU": 4,
+                "TPU-v7x-16-head": 1,
+            },
+            id="tpu-group-v7x",
+        ),
+        pytest.param(
+            _get_tpu_group_with_v5litepod_node_selectors(),
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+                "TPU": 4,
+                "TPU-v5litepod-8-head": 1,
+            },
+            id="tpu-group-v5litepod",
+        ),
+    ]
+)
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Not relevant.")
+@pytest.mark.parametrize(RAY_RESOURCES_PARAM_ARGS, RAY_RESOURCES_TEST_DATA)
+def test_get_ray_resources_from_group_spec(
+    group_spec: Dict[str, Any],
+    is_head: bool,
+    expected_resources: Dict[str, Any],
+):
+    assert _get_ray_resources_from_group_spec(group_spec, is_head) == expected_resources
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Not relevant.")
+def test_top_level_resources_override_warnings():
+    """
+    Verify all override warnings are logged when a top-level `resources` field is used in
+    addition to specifying those resources in the rayStartParams.
+    """
+    ray_cr_in = _get_ray_cr_with_top_level_resources()
+    ray_cr_in["metadata"]["namespace"] = "default"
+
+    with mock.patch(f"{AUTOSCALING_CONFIG_MODULE_PATH}.logger") as mock_logger:
+        _derive_autoscaling_config_from_ray_cr(ray_cr_in)
+
+        expected_calls = [
+            mock.call(
+                "'CPU' specified in both the top-level 'resources' field and in 'rayStartParams'. "
+                "Using the value from 'resources': 16."
+            ),
+            mock.call(
+                "'GPU' specified in both the top-level 'resources' field and in 'rayStartParams'. "
+                "Using the value from 'resources': 8."
+            ),
+            mock.call(
+                "'memory' specified in both the top-level 'resources' field and in 'rayStartParams'. "
+                "Using the value from 'resources': 2Gi."
+            ),
+            mock.call(
+                "custom resources specified in both the top-level 'resources' field and in 'rayStartParams'. "
+                "Using the values from 'resources': {'CPU': '16', 'GPU': '8', 'memory': '2Gi', 'CustomResource': '99'}."
+            ),
+        ]
+
+        # Assert that all expected calls were made, in any order.
+        mock_logger.warning.assert_has_calls(expected_calls, any_order=True)
+        assert mock_logger.warning.call_count == 4
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tests/kuberay/test_autoscaling_e2e.py

@@ -0,0 +1,384 @@
+import base64
+import copy
+import logging
+import os
+import subprocess
+import sys
+import tempfile
+import unittest
+from typing import Any, Dict
+
+import pytest
+import yaml
+
+from ray.tests.kuberay.utils import (
+    get_pod,
+    get_pod_names,
+    get_raycluster,
+    kubectl_delete,
+    kubectl_exec_python_script,
+    kubectl_logs,
+    switch_to_ray_parent_dir,
+    wait_for_pod_to_start,
+    wait_for_pods,
+    wait_for_ray_health,
+)
+
+logger = logging.getLogger(__name__)
+
+# This image will be used for both the Ray nodes and the autoscaler.
+# The CI should pass an image built from the test branch.
+RAY_IMAGE = os.environ.get("RAY_IMAGE", "rayproject/ray:nightly-py38")
+# By default, use the same image for the autoscaler and Ray containers.
+AUTOSCALER_IMAGE = os.environ.get("AUTOSCALER_IMAGE", RAY_IMAGE)
+# Set to IfNotPresent in kind CI.
+PULL_POLICY = os.environ.get("PULL_POLICY", "IfNotPresent")
+# Set to enable autoscaler v2
+AUTOSCALER_V2 = os.environ.get("AUTOSCALER_V2", "False")
+logger.info(f"Using image `{RAY_IMAGE}` for Ray containers.")
+logger.info(f"Using image `{AUTOSCALER_IMAGE}` for Autoscaler containers.")
+logger.info(f"Using pull policy `{PULL_POLICY}` for all images.")
+logger.info(f"Using autoscaler v2: {AUTOSCALER_V2}")
+
+# Path to example config inside the rayci container.
+EXAMPLE_CLUSTER_PATH = (
+    "rayci/python/ray/tests/kuberay/test_files/ray-cluster.autoscaler-template.yaml"
+)
+EXAMPLE_CLUSTER_PATH_V2 = (
+    "rayci/python/ray/tests/kuberay/test_files/ray-cluster.autoscaler-v2-template.yaml"
+)
+
+HEAD_SERVICE = "raycluster-autoscaler-head-svc"
+HEAD_POD_PREFIX = "raycluster-autoscaler-head"
+CPU_WORKER_PREFIX = "raycluster-autoscaler-worker-small-group"
+RAY_CLUSTER_NAME = "raycluster-autoscaler"
+RAY_CLUSTER_NAMESPACE = "default"
+
+# Test runs longer than the default timeout.
+pytestmark = pytest.mark.timeout(300)
+
+
+class KubeRayAutoscalingTest(unittest.TestCase):
+    """e2e verification of autoscaling following the steps in the Ray documentation.
+    kubectl is used throughout, as that reflects the instructions in the docs.
+    """
+
+    def _get_ray_cr_config(
+        self, min_replicas=0, cpu_replicas=0, gpu_replicas=0
+    ) -> Dict[str, Any]:
+        """Get Ray CR config yaml.
+
+        - Use configurable replica fields for a CPU workerGroup.
+
+        - Add a GPU-annotated group for testing GPU upscaling.
+
+        - Fill in Ray image, autoscaler image, and image pull policies from env
+          variables.
+        """
+        if AUTOSCALER_V2 == "True":
+            with open(EXAMPLE_CLUSTER_PATH_V2) as ray_cr_config_file:
+                ray_cr_config_str = ray_cr_config_file.read()
+        else:
+            with open(EXAMPLE_CLUSTER_PATH) as ray_cr_config_file:
+                ray_cr_config_str = ray_cr_config_file.read()
+
+        for k8s_object in yaml.safe_load_all(ray_cr_config_str):
+            if k8s_object["kind"] in ["RayCluster", "RayJob", "RayService"]:
+                config = k8s_object
+                break
+        head_group = config["spec"]["headGroupSpec"]
+        if "rayStartParams" not in head_group:
+            head_group["rayStartParams"] = {}
+        head_group["rayStartParams"][
+            "resources"
+        ] = '"{\\"Custom1\\": 1, \\"Custom2\\": 5}"'
+
+        cpu_group = config["spec"]["workerGroupSpecs"][0]
+        cpu_group["replicas"] = cpu_replicas
+        cpu_group["minReplicas"] = min_replicas
+        # Keep maxReplicas big throughout the test.
+        cpu_group["maxReplicas"] = 300
+        if "rayStartParams" not in cpu_group:
+            cpu_group["rayStartParams"] = {}
+        cpu_group["rayStartParams"][
+            "resources"
+        ] = '"{\\"Custom1\\": 1, \\"Custom2\\": 5}"'
+
+        # Add a GPU-annotated group.
+        # (We're not using real GPUs, just adding a GPU annotation for the autoscaler
+        # and Ray scheduler.)
+        gpu_group = copy.deepcopy(cpu_group)
+        if "rayStartParams" not in gpu_group:
+            gpu_group["rayStartParams"] = {}
+        gpu_group["rayStartParams"]["num-gpus"] = "1"
+        gpu_group["replicas"] = gpu_replicas
+        gpu_group["minReplicas"] = 0
+        gpu_group["maxReplicas"] = 1
+        gpu_group["groupName"] = "fake-gpu-group"
+        config["spec"]["workerGroupSpecs"].append(gpu_group)
+
+        # Substitute images.
+        for group_spec in config["spec"]["workerGroupSpecs"] + [
+            config["spec"]["headGroupSpec"]
+        ]:
+            containers = group_spec["template"]["spec"]["containers"]
+
+            ray_container = containers[0]
+            # Confirm the first container in the example config is the Ray container.
+            assert ray_container["name"] in ["ray-head", "ray-worker"]
+            # ("machine-learning" is the name of the worker Ray container)
+
+            ray_container["image"] = RAY_IMAGE
+
+            for container in containers:
+                container["imagePullPolicy"] = PULL_POLICY
+
+        autoscaler_options = {
+            "image": AUTOSCALER_IMAGE,
+            "imagePullPolicy": PULL_POLICY,
+            # Allow quick scale-down for test purposes.
+            "idleTimeoutSeconds": 10,
+        }
+        config["spec"]["autoscalerOptions"] = autoscaler_options
+
+        return config
+
+    def _apply_ray_cr(
+        self,
+        min_replicas=0,
+        cpu_replicas=0,
+        gpu_replicas=0,
+        validate_replicas: bool = False,
+    ) -> None:
+        """Apply Ray CR config yaml, with configurable replica fields for the cpu
+        workerGroup.
+
+        If the CR does not yet exist, `replicas` can be set as desired.
+        If the CR does already exist, the recommended usage is this:
+            (1) Set `cpu_replicas` and `gpu_replicas` to what we currently expect them
+                to be.
+            (2) Set `validate_replicas` to True. We will then check that the replicas
+            set on the CR coincides with `replicas`.
+        """
+        if validate_replicas:
+            raycluster = get_raycluster(
+                RAY_CLUSTER_NAME, namespace=RAY_CLUSTER_NAMESPACE
+            )
+            assert raycluster["spec"]["workerGroupSpecs"][0]["replicas"] == cpu_replicas
+            assert raycluster["spec"]["workerGroupSpecs"][1]["replicas"] == gpu_replicas
+            logger.info(
+                f"Validated that cpu and gpu worker replicas for "
+                f"{RAY_CLUSTER_NAME} are currently {cpu_replicas} and"
+                f" {gpu_replicas}, respectively."
+            )
+        cr_config = self._get_ray_cr_config(
+            min_replicas=min_replicas,
+            cpu_replicas=cpu_replicas,
+            gpu_replicas=gpu_replicas,
+        )
+
+        with tempfile.NamedTemporaryFile("w") as config_file:
+            yaml.dump(cr_config, config_file)
+            config_file.flush()
+
+            subprocess.check_call(
+                ["kubectl", "apply", "-f", config_file.name],
+                stdout=sys.stdout,
+                stderr=sys.stderr,
+            )
+
+    def testAutoscaling(self):
+        """Test the following behaviors:
+
+        1. Spinning up a Ray cluster
+        2. Scaling up Ray workers via autoscaler.sdk.request_resources()
+        3. Scaling up by updating the CRD's minReplicas
+        4. Scaling down by removing the resource request and reducing maxReplicas
+        5. Autoscaler recognizes GPU annotations and Ray custom resources.
+        6. Autoscaler and operator ignore pods marked for deletion.
+        7. Autoscaler logs work. Autoscaler events are piped to the driver.
+        8. Ray utils show correct resource limits in the head container.
+
+        TODO (Dmitri): Split up the test logic.
+        Too much is stuffed into this one test case.
+
+        Resources requested by this test are safely within the bounds of an m5.xlarge
+        instance.
+
+        The resource REQUESTS are:
+        - One Ray head pod
+            - Autoscaler: .25 CPU, .5 Gi memory
+            - Ray node: .5 CPU, .5 Gi memeory
+        - Three Worker pods
+            - Ray node: .5 CPU, .5 Gi memory
+        Total: 2.25 CPU, 2.5 Gi memory.
+
+        Including operator and system pods, the total CPU requested is around 3.
+
+        The cpu LIMIT of each Ray container is 1.
+        The `num-cpus` arg to Ray start is 1 for each Ray container; thus Ray accounts
+        1 CPU for each Ray node in the test.
+        """
+        switch_to_ray_parent_dir()
+
+        # Cluster creation
+        logger.info("Creating a RayCluster with no worker pods.")
+        self._apply_ray_cr(min_replicas=0, cpu_replicas=0, gpu_replicas=0)
+
+        logger.info("Confirming presence of head.")
+        wait_for_pods(goal_num_pods=1, namespace=RAY_CLUSTER_NAMESPACE)
+
+        logger.info("Waiting for head pod to start Running.")
+        wait_for_pod_to_start(
+            pod_name_filter=HEAD_POD_PREFIX, namespace=RAY_CLUSTER_NAMESPACE
+        )
+        logger.info("Confirming Ray is up on the head pod.")
+        wait_for_ray_health(
+            pod_name_filter=HEAD_POD_PREFIX, namespace=RAY_CLUSTER_NAMESPACE
+        )
+
+        head_pod = get_pod(
+            pod_name_filter=HEAD_POD_PREFIX, namespace=RAY_CLUSTER_NAMESPACE
+        )
+        assert head_pod, "Could not find the Ray head pod."
+
+        # Confirm head pod resource allocation.
+        # (This is a misplaced test of Ray's resource detection in containers.
+        # See the TODO in the docstring.)
+        logger.info("Confirming head pod resource allocation.")
+        out = kubectl_exec_python_script(  # Interaction mode #1: `kubectl exec`
+            script_name="check_cpu_and_memory.py",
+            pod=head_pod,
+            container="ray-head",
+            namespace="default",
+        )
+
+        # Scale-up
+        logger.info("Scaling up to one worker via Ray resource request.")
+        # The request for 2 cpus should give us a 1-cpu head (already present) and a
+        # 1-cpu worker (will await scale-up).
+        kubectl_exec_python_script(  # Interaction mode #1: `kubectl exec`
+            script_name="scale_up.py",
+            pod=head_pod,
+            container="ray-head",
+            namespace="default",
+        )
+        # Check that stdout autoscaler logging is working.
+        logs = kubectl_logs(head_pod, namespace="default", container="autoscaler")
+        assert "Adding 1 node(s) of type small-group." in logs
+        logger.info("Confirming number of workers.")
+        wait_for_pods(goal_num_pods=2, namespace=RAY_CLUSTER_NAMESPACE)
+
+        # Ray CR updates.
+        logger.info("Scaling up to two workers by editing minReplicas.")
+        # replicas=1 reflects the current number of workers
+        # (which is what we expect to be already present in the Ray CR)
+        self._apply_ray_cr(
+            min_replicas=2,
+            cpu_replicas=1,
+            gpu_replicas=0,
+            # Confirm CPU, GPU replicas set on the Ray CR by the autoscaler are 1, 0:
+            validate_replicas=True,
+        )
+        logger.info("Confirming number of workers.")
+        wait_for_pods(goal_num_pods=3, namespace=RAY_CLUSTER_NAMESPACE)
+
+        # GPU upscaling.
+        # 1. Check we haven't spuriously already started a fake GPU node.
+        assert not any(
+            "gpu" in pod_name
+            for pod_name in get_pod_names(namespace=RAY_CLUSTER_NAMESPACE)
+        )
+        # 2. Trigger GPU upscaling by requesting placement of a GPU actor.
+        logger.info("Scheduling an Actor with GPU demands.")
+        kubectl_exec_python_script(
+            script_name="gpu_actor_placement.py",
+            pod=head_pod,
+            container="ray-head",
+            namespace="default",
+        )
+
+        # 3. Confirm new pod number and presence of fake GPU worker.
+        logger.info("Confirming fake GPU worker up-scaling.")
+        wait_for_pods(goal_num_pods=4, namespace=RAY_CLUSTER_NAMESPACE)
+
+        gpu_workers = [
+            pod_name
+            for pod_name in get_pod_names(namespace=RAY_CLUSTER_NAMESPACE)
+            if "gpu" in pod_name
+        ]
+        assert len(gpu_workers) == 1
+        # 4. Confirm that the GPU actor is up and that Ray believes
+        # the node the actor is on has a GPU.
+        logger.info("Confirming GPU actor placement.")
+        out = kubectl_exec_python_script(
+            script_name="gpu_actor_validation.py",
+            pod=head_pod,
+            container="ray-head",
+            namespace="default",
+        )
+
+        # Confirms the actor was placed on a GPU-annotated node.
+        # (See gpu_actor_validation.py for details.)
+        assert "on-a-gpu-node" in out
+
+        # Scale-down
+        logger.info("Reducing min workers to 0.")
+        # Max workers remains 300.
+        self._apply_ray_cr(
+            min_replicas=0,
+            cpu_replicas=2,
+            gpu_replicas=1,
+            # Confirm CPU, GPU replicas set on the Ray CR by the autoscaler are 2, 1:
+            validate_replicas=True,
+        )
+        logger.info("Removing resource demands.")
+        kubectl_exec_python_script(
+            script_name="scale_down.py",
+            pod=head_pod,
+            container="ray-head",
+            namespace="default",
+        )
+        # Autoscaler should trigger scale-down after resource demands are removed.
+        logger.info("Confirming workers are gone.")
+        # Check that stdout autoscaler logging is working.
+        logs = kubectl_logs(head_pod, namespace="default", container="autoscaler")
+        assert "Removing 1 nodes of type fake-gpu-group (idle)." in logs
+        wait_for_pods(goal_num_pods=1, namespace=RAY_CLUSTER_NAMESPACE, tries=120)
+
+        # Check custom resource upscaling.
+
+        # Submit two {"Custom2": 3} bundles to upscale two workers with 5
+        # Custom2 capacity each.
+        logger.info("Scaling up workers with request for custom resources.")
+        out = kubectl_exec_python_script(
+            script_name="scale_up_custom.py",
+            pod=head_pod,
+            container="ray-head",
+            namespace="default",
+        )
+        assert "Submitted custom scale request!" in out, out
+
+        logger.info("Confirming two workers have scaled up.")
+        wait_for_pods(goal_num_pods=3, namespace=RAY_CLUSTER_NAMESPACE)
+
+        # Cluster deletion
+        logger.info("Deleting Ray cluster.")
+        kubectl_delete(
+            kind="raycluster", name=RAY_CLUSTER_NAME, namespace=RAY_CLUSTER_NAMESPACE
+        )
+        logger.info("Confirming Ray pods are gone.")
+        wait_for_pods(goal_num_pods=0, namespace=RAY_CLUSTER_NAMESPACE)
+
+
+if __name__ == "__main__":
+    kubeconfig_base64 = os.environ.get("KUBECONFIG_BASE64")
+    if kubeconfig_base64:
+        kubeconfig_file = os.environ.get("KUBECONFIG")
+        if not kubeconfig_file:
+            raise ValueError("When KUBECONFIG_BASE64 is set, KUBECONFIG must be set.")
+
+        with open(kubeconfig_file, "wb") as f:
+            f.write(base64.b64decode(kubeconfig_base64))
+
+    sys.exit(pytest.main(["-vv", __file__]))

python/ray/tests/kuberay/test_files/podlist1.yaml

@@ -0,0 +1,607 @@
+apiVersion: v1
+items:
+- apiVersion: v1
+  kind: Pod
+  metadata:
+    annotations:
+      ray.io/ft-enabled: "false"
+      ray.io/health-state: ""
+    creationTimestamp: "2022-11-14T23:10:15Z"
+    generateName: raycluster-autoscaler-head-
+    labels:
+      app.kubernetes.io/created-by: kuberay-operator
+      app.kubernetes.io/name: kuberay
+      ray.io/cluster: raycluster-autoscaler
+      ray.io/cluster-dashboard: raycluster-autoscaler-dashboard
+      ray.io/group: headgroup
+      ray.io/identifier: raycluster-autoscaler-head
+      ray.io/is-ray-node: "yes"
+      ray.io/node-type: head
+    name: raycluster-autoscaler-head-8zsc8
+    namespace: default
+    ownerReferences:
+    - apiVersion: ray.io/v1alpha1
+      blockOwnerDeletion: true
+      controller: true
+      kind: RayCluster
+      name: raycluster-autoscaler
+      uid: ec79effb-0295-4f40-b08b-8633aa7f786a
+    resourceVersion: "4519"
+    uid: 539ea57c-8d51-4503-a395-08779efb3bf0
+  spec:
+    containers:
+    - args:
+      - 'ulimit -n 65536; ray start --head  --resources="{\"Custom1\": 1, \"Custom2\":
+        5}"  --block  --dashboard-host=0.0.0.0  --metrics-export-port=8080  --no-monitor  --num-cpus=1  --memory=1000000000 '
+      command:
+      - /bin/bash
+      - -c
+      - --
+      env:
+      - name: RAY_IP
+        value: 127.0.0.1
+      - name: RAY_PORT
+        value: "6379"
+      - name: RAY_ADDRESS
+        value: 127.0.0.1:6379
+      - name: REDIS_PASSWORD
+      image: gekho/ray
+      imagePullPolicy: Always
+      lifecycle:
+        preStop:
+          exec:
+            command:
+            - /bin/sh
+            - -c
+            - ray stop
+      name: ray-head
+      ports:
+      - containerPort: 6379
+        name: gcs
+        protocol: TCP
+      - containerPort: 8265
+        name: dashboard
+        protocol: TCP
+      - containerPort: 10001
+        name: client
+        protocol: TCP
+      - containerPort: 8080
+        name: metrics
+        protocol: TCP
+      resources:
+        limits:
+          cpu: "1"
+          memory: 1G
+        requests:
+          cpu: 500m
+          memory: 512Mi
+      terminationMessagePath: /dev/termination-log
+      terminationMessagePolicy: File
+      volumeMounts:
+      - mountPath: /dev/shm
+        name: shared-mem
+      - mountPath: /tmp/ray
+        name: ray-logs
+      - mountPath: /var/run/secrets/kubernetes.io/serviceaccount
+        name: kube-api-access-tmxvr
+        readOnly: true
+    - args:
+      - kuberay-autoscaler
+      - --cluster-name
+      - $(RAY_CLUSTER_NAME)
+      - --cluster-namespace
+      - $(RAY_CLUSTER_NAMESPACE)
+      command:
+      - ray
+      env:
+      - name: RAY_CLUSTER_NAME
+        valueFrom:
+          fieldRef:
+            apiVersion: v1
+            fieldPath: metadata.labels['ray.io/cluster']
+      - name: RAY_CLUSTER_NAMESPACE
+        valueFrom:
+          fieldRef:
+            apiVersion: v1
+            fieldPath: metadata.namespace
+      image: gekho/ray
+      imagePullPolicy: Always
+      name: autoscaler
+      resources:
+        limits:
+          cpu: 500m
+          memory: 512Mi
+        requests:
+          cpu: 500m
+          memory: 512Mi
+      terminationMessagePath: /dev/termination-log
+      terminationMessagePolicy: File
+      volumeMounts:
+      - mountPath: /tmp/ray
+        name: ray-logs
+      - mountPath: /var/run/secrets/kubernetes.io/serviceaccount
+        name: kube-api-access-tmxvr
+        readOnly: true
+    dnsPolicy: ClusterFirst
+    enableServiceLinks: true
+    nodeName: gke-cluster-1-default-pool-a5503908-181p
+    preemptionPolicy: PreemptLowerPriority
+    priority: 0
+    restartPolicy: Always
+    schedulerName: default-scheduler
+    securityContext: {}
+    serviceAccount: raycluster-autoscaler
+    serviceAccountName: raycluster-autoscaler
+    terminationGracePeriodSeconds: 30
+    tolerations:
+    - effect: NoExecute
+      key: node.kubernetes.io/not-ready
+      operator: Exists
+      tolerationSeconds: 300
+    - effect: NoExecute
+      key: node.kubernetes.io/unreachable
+      operator: Exists
+      tolerationSeconds: 300
+    volumes:
+    - emptyDir:
+        medium: Memory
+        sizeLimit: 512Mi
+      name: shared-mem
+    - emptyDir: {}
+      name: ray-logs
+    - name: kube-api-access-tmxvr
+      projected:
+        defaultMode: 420
+        sources:
+        - serviceAccountToken:
+            expirationSeconds: 3607
+            path: token
+        - configMap:
+            items:
+            - key: ca.crt
+              path: ca.crt
+            name: kube-root-ca.crt
+        - downwardAPI:
+            items:
+            - fieldRef:
+                apiVersion: v1
+                fieldPath: metadata.namespace
+              path: namespace
+  status:
+    conditions:
+    - lastProbeTime: null
+      lastTransitionTime: "2022-11-14T23:10:15Z"
+      status: "True"
+      type: Initialized
+    - lastProbeTime: null
+      lastTransitionTime: "2022-11-14T23:11:23Z"
+      status: "True"
+      type: Ready
+    - lastProbeTime: null
+      lastTransitionTime: "2022-11-14T23:11:23Z"
+      status: "True"
+      type: ContainersReady
+    - lastProbeTime: null
+      lastTransitionTime: "2022-11-14T23:10:15Z"
+      status: "True"
+      type: PodScheduled
+    containerStatuses:
+    - containerID: containerd://0b008432be839bec8dd97437d3f2be9ac8d7f017b91067a46ec45a487f141ebf
+      image: docker.io/gekho/ray:latest
+      imageID: docker.io/gekho/ray@sha256:7859a78d1a089bb88691864d5c4a2aad529f5353d7d9c82cc0274842fbda242b
+      lastState: {}
+      name: autoscaler
+      ready: true
+      restartCount: 0
+      started: true
+      state:
+        running:
+          startedAt: "2022-11-14T23:11:23Z"
+    - containerID: containerd://b2aae80ed028cc41bad1e350bb70a0a4e8ea722df098b38781efabe54adbc5ec
+      image: docker.io/gekho/ray:latest
+      imageID: docker.io/gekho/ray@sha256:7859a78d1a089bb88691864d5c4a2aad529f5353d7d9c82cc0274842fbda242b
+      lastState: {}
+      name: ray-head
+      ready: true
+      restartCount: 0
+      started: true
+      state:
+        running:
+          startedAt: "2022-11-14T23:11:22Z"
+    hostIP: 10.128.0.45
+    phase: Running
+    podIP: 10.4.2.6
+    podIPs:
+    - ip: 10.4.2.6
+    qosClass: Burstable
+    startTime: "2022-11-14T23:10:15Z"
+- apiVersion: v1
+  kind: Pod
+  metadata:
+    annotations:
+      key: value
+      ray.io/ft-enabled: "false"
+      ray.io/health-state: ""
+    creationTimestamp: "2022-11-14T23:11:45Z"
+    deletionGracePeriodSeconds: 30
+    deletionTimestamp: "2022-11-14T23:12:20Z"
+    generateName: raycluster-autoscaler-worker-small-group-
+    labels:
+      app.kubernetes.io/created-by: kuberay-operator
+      app.kubernetes.io/name: kuberay
+      key: value
+      ray.io/cluster: raycluster-autoscaler
+      ray.io/cluster-dashboard: raycluster-autoscaler-dashboard
+      ray.io/group: small-group
+      ray.io/identifier: raycluster-autoscaler-worker
+      ray.io/is-ray-node: "yes"
+      ray.io/node-type: worker
+    name: raycluster-autoscaler-worker-small-group-4wxfm
+    namespace: default
+    ownerReferences:
+    - apiVersion: ray.io/v1alpha1
+      blockOwnerDeletion: true
+      controller: true
+      kind: RayCluster
+      name: raycluster-autoscaler
+      uid: ec79effb-0295-4f40-b08b-8633aa7f786a
+    resourceVersion: "4845"
+    uid: 3698ed9b-7e06-4d47-b9f6-09e4bd08365a
+  spec:
+    containers:
+    - args:
+      - 'ulimit -n 65536; ray start  --resources="{\"Custom1\": 1, \"Custom2\": 5}"  --address=raycluster-autoscaler-head-svc:6379  --metrics-export-port=8080  --num-cpus=1  --memory=536870912  --block '
+      command:
+      - /bin/bash
+      - -c
+      - --
+      env:
+      - name: RAY_IP
+        value: raycluster-autoscaler-head-svc
+      - name: RAY_PORT
+        value: "6379"
+      - name: RAY_ADDRESS
+        value: raycluster-autoscaler-head-svc:6379
+      - name: REDIS_PASSWORD
+      image: gekho/ray
+      imagePullPolicy: Always
+      lifecycle:
+        preStop:
+          exec:
+            command:
+            - /bin/sh
+            - -c
+            - ray stop
+      name: machine-learning
+      ports:
+      - containerPort: 8080
+        name: metrics
+        protocol: TCP
+      resources:
+        limits:
+          cpu: "1"
+          memory: 512Mi
+        requests:
+          cpu: 500m
+          memory: 256Mi
+      terminationMessagePath: /dev/termination-log
+      terminationMessagePolicy: File
+      volumeMounts:
+      - mountPath: /dev/shm
+        name: shared-mem
+      - mountPath: /var/run/secrets/kubernetes.io/serviceaccount
+        name: kube-api-access-wthw9
+        readOnly: true
+    dnsPolicy: ClusterFirst
+    enableServiceLinks: true
+    initContainers:
+    - command:
+      - sh
+      - -c
+      - until nslookup $RAY_IP.$(cat /var/run/secrets/kubernetes.io/serviceaccount/namespace).svc.cluster.local;
+        do echo waiting for myservice; sleep 2; done
+      env:
+      - name: RAY_IP
+        value: raycluster-autoscaler-head-svc
+      image: busybox:1.28
+      imagePullPolicy: IfNotPresent
+      name: init-myservice
+      resources: {}
+      terminationMessagePath: /dev/termination-log
+      terminationMessagePolicy: File
+      volumeMounts:
+      - mountPath: /var/run/secrets/kubernetes.io/serviceaccount
+        name: kube-api-access-wthw9
+        readOnly: true
+    nodeName: gke-cluster-1-default-pool-a5503908-dpst
+    preemptionPolicy: PreemptLowerPriority
+    priority: 0
+    restartPolicy: Always
+    schedulerName: default-scheduler
+    securityContext: {}
+    serviceAccount: default
+    serviceAccountName: default
+    terminationGracePeriodSeconds: 30
+    tolerations:
+    - effect: NoExecute
+      key: node.kubernetes.io/not-ready
+      operator: Exists
+      tolerationSeconds: 300
+    - effect: NoExecute
+      key: node.kubernetes.io/unreachable
+      operator: Exists
+      tolerationSeconds: 300
+    volumes:
+    - emptyDir:
+        medium: Memory
+        sizeLimit: 256Mi
+      name: shared-mem
+    - name: kube-api-access-wthw9
+      projected:
+        defaultMode: 420
+        sources:
+        - serviceAccountToken:
+            expirationSeconds: 3607
+            path: token
+        - configMap:
+            items:
+            - key: ca.crt
+              path: ca.crt
+            name: kube-root-ca.crt
+        - downwardAPI:
+            items:
+            - fieldRef:
+                apiVersion: v1
+                fieldPath: metadata.namespace
+              path: namespace
+  status:
+    conditions:
+    - lastProbeTime: null
+      lastTransitionTime: "2022-11-14T23:11:47Z"
+      status: "True"
+      type: Initialized
+    - lastProbeTime: null
+      lastTransitionTime: "2022-11-14T23:11:45Z"
+      message: 'containers with unready status: [machine-learning]'
+      reason: ContainersNotReady
+      status: "False"
+      type: Ready
+    - lastProbeTime: null
+      lastTransitionTime: "2022-11-14T23:11:45Z"
+      message: 'containers with unready status: [machine-learning]'
+      reason: ContainersNotReady
+      status: "False"
+      type: ContainersReady
+    - lastProbeTime: null
+      lastTransitionTime: "2022-11-14T23:11:45Z"
+      status: "True"
+      type: PodScheduled
+    containerStatuses:
+    - image: gekho/ray
+      imageID: ""
+      lastState: {}
+      name: machine-learning
+      ready: false
+      restartCount: 0
+      started: false
+      state:
+        waiting:
+          reason: PodInitializing
+    hostIP: 10.128.0.31
+    initContainerStatuses:
+    - containerID: containerd://c7f5a0c3f63957213213ed1ebb6446cd205bd60346d010a879c5fa24e37f5145
+      image: docker.io/library/busybox:1.28
+      imageID: docker.io/library/busybox@sha256:141c253bc4c3fd0a201d32dc1f493bcf3fff003b6df416dea4f41046e0f37d47
+      lastState: {}
+      name: init-myservice
+      ready: true
+      restartCount: 0
+      state:
+        terminated:
+          containerID: containerd://c7f5a0c3f63957213213ed1ebb6446cd205bd60346d010a879c5fa24e37f5145
+          exitCode: 0
+          finishedAt: "2022-11-14T23:11:47Z"
+          reason: Completed
+          startedAt: "2022-11-14T23:11:47Z"
+    phase: Pending
+    podIP: 10.4.0.4
+    podIPs:
+    - ip: 10.4.0.4
+    qosClass: Burstable
+    startTime: "2022-11-14T23:11:45Z"
+- apiVersion: v1
+  kind: Pod
+  metadata:
+    annotations:
+      key: value
+      ray.io/ft-enabled: "false"
+      ray.io/health-state: ""
+    creationTimestamp: "2022-11-14T23:11:50Z"
+    generateName: raycluster-autoscaler-worker-small-group-
+    labels:
+      app.kubernetes.io/created-by: kuberay-operator
+      app.kubernetes.io/name: kuberay
+      key: value
+      ray.io/cluster: raycluster-autoscaler
+      ray.io/cluster-dashboard: raycluster-autoscaler-dashboard
+      ray.io/group: small-group
+      ray.io/identifier: raycluster-autoscaler-worker
+      ray.io/is-ray-node: "yes"
+      ray.io/node-type: worker
+    name: raycluster-autoscaler-worker-small-group-dkz2r
+    namespace: default
+    ownerReferences:
+    - apiVersion: ray.io/v1alpha1
+      blockOwnerDeletion: true
+      controller: true
+      kind: RayCluster
+      name: raycluster-autoscaler
+      uid: ec79effb-0295-4f40-b08b-8633aa7f786a
+    resourceVersion: "4776"
+    uid: b4fb3233-6024-48a8-9f4f-a18f5e490629
+  spec:
+    containers:
+    - args:
+      - 'ulimit -n 65536; ray start  --block  --resources="{\"Custom1\": 1, \"Custom2\":
+        5}"  --address=raycluster-autoscaler-head-svc:6379  --metrics-export-port=8080  --num-cpus=1  --memory=536870912 '
+      command:
+      - /bin/bash
+      - -c
+      - --
+      env:
+      - name: RAY_IP
+        value: raycluster-autoscaler-head-svc
+      - name: RAY_PORT
+        value: "6379"
+      - name: RAY_ADDRESS
+        value: raycluster-autoscaler-head-svc:6379
+      - name: REDIS_PASSWORD
+      image: gekho/ray
+      imagePullPolicy: Always
+      lifecycle:
+        preStop:
+          exec:
+            command:
+            - /bin/sh
+            - -c
+            - ray stop
+      name: machine-learning
+      ports:
+      - containerPort: 8080
+        name: metrics
+        protocol: TCP
+      resources:
+        limits:
+          cpu: "1"
+          memory: 512Mi
+        requests:
+          cpu: 500m
+          memory: 256Mi
+      terminationMessagePath: /dev/termination-log
+      terminationMessagePolicy: File
+      volumeMounts:
+      - mountPath: /dev/shm
+        name: shared-mem
+      - mountPath: /var/run/secrets/kubernetes.io/serviceaccount
+        name: kube-api-access-djtd9
+        readOnly: true
+    dnsPolicy: ClusterFirst
+    enableServiceLinks: true
+    initContainers:
+    - command:
+      - sh
+      - -c
+      - until nslookup $RAY_IP.$(cat /var/run/secrets/kubernetes.io/serviceaccount/namespace).svc.cluster.local;
+        do echo waiting for myservice; sleep 2; done
+      env:
+      - name: RAY_IP
+        value: raycluster-autoscaler-head-svc
+      image: busybox:1.28
+      imagePullPolicy: IfNotPresent
+      name: init-myservice
+      resources: {}
+      terminationMessagePath: /dev/termination-log
+      terminationMessagePolicy: File
+      volumeMounts:
+      - mountPath: /var/run/secrets/kubernetes.io/serviceaccount
+        name: kube-api-access-djtd9
+        readOnly: true
+    nodeName: gke-cluster-1-default-pool-a5503908-j51d
+    preemptionPolicy: PreemptLowerPriority
+    priority: 0
+    restartPolicy: Always
+    schedulerName: default-scheduler
+    securityContext: {}
+    serviceAccount: default
+    serviceAccountName: default
+    terminationGracePeriodSeconds: 30
+    tolerations:
+    - effect: NoExecute
+      key: node.kubernetes.io/not-ready
+      operator: Exists
+      tolerationSeconds: 300
+    - effect: NoExecute
+      key: node.kubernetes.io/unreachable
+      operator: Exists
+      tolerationSeconds: 300
+    volumes:
+    - emptyDir:
+        medium: Memory
+        sizeLimit: 256Mi
+      name: shared-mem
+    - name: kube-api-access-djtd9
+      projected:
+        defaultMode: 420
+        sources:
+        - serviceAccountToken:
+            expirationSeconds: 3607
+            path: token
+        - configMap:
+            items:
+            - key: ca.crt
+              path: ca.crt
+            name: kube-root-ca.crt
+        - downwardAPI:
+            items:
+            - fieldRef:
+                apiVersion: v1
+                fieldPath: metadata.namespace
+              path: namespace
+  status:
+    conditions:
+    - lastProbeTime: null
+      lastTransitionTime: "2022-11-14T23:11:51Z"
+      status: "True"
+      type: Initialized
+    - lastProbeTime: null
+      lastTransitionTime: "2022-11-14T23:11:50Z"
+      message: 'containers with unready status: [machine-learning]'
+      reason: ContainersNotReady
+      status: "False"
+      type: Ready
+    - lastProbeTime: null
+      lastTransitionTime: "2022-11-14T23:11:50Z"
+      message: 'containers with unready status: [machine-learning]'
+      reason: ContainersNotReady
+      status: "False"
+      type: ContainersReady
+    - lastProbeTime: null
+      lastTransitionTime: "2022-11-14T23:11:50Z"
+      status: "True"
+      type: PodScheduled
+    containerStatuses:
+    - image: gekho/ray
+      imageID: ""
+      lastState: {}
+      name: machine-learning
+      ready: false
+      restartCount: 0
+      started: false
+      state:
+        waiting:
+          reason: PodInitializing
+    hostIP: 10.128.0.43
+    initContainerStatuses:
+    - containerID: containerd://672d9a5836e27a17f57a4e15e1d86431cfee6f2edef1210d60e864e3c510aac0
+      image: docker.io/library/busybox:1.28
+      imageID: docker.io/library/busybox@sha256:141c253bc4c3fd0a201d32dc1f493bcf3fff003b6df416dea4f41046e0f37d47
+      lastState: {}
+      name: init-myservice
+      ready: true
+      restartCount: 0
+      state:
+        terminated:
+          containerID: containerd://672d9a5836e27a17f57a4e15e1d86431cfee6f2edef1210d60e864e3c510aac0
+          exitCode: 0
+          finishedAt: "2022-11-14T23:11:51Z"
+          reason: Completed
+          startedAt: "2022-11-14T23:11:51Z"
+    phase: Pending
+    podIP: 10.4.1.8
+    podIPs:
+    - ip: 10.4.1.8
+    qosClass: Burstable
+    startTime: "2022-11-14T23:11:50Z"
+kind: List
+metadata:
+  resourceVersion: ""

python/ray/tests/kuberay/test_files/ray-cluster.autoscaler-template.yaml

@@ -0,0 +1,73 @@
+# `ray-cluster.autoscaler-template.yaml` is a template for the RayCluster CR and
+# is used by the function `_get_ray_cr_config` in `test_autoscaling_e2e.py`.
+# [Note]
+# (1) The VM test runner only has 4 CPUs, so we lower the CPU requests.
+# (2) `test_autoscaling_e2e.py` assumes that each Ray Pod has 1 logical CPU.
+apiVersion: ray.io/v1
+kind: RayCluster
+metadata:
+  name: raycluster-autoscaler
+spec:
+  # The version of Ray you are using. Make sure all Ray containers are running this version of Ray.
+  rayVersion: '2.46.0'
+  # If `enableInTreeAutoscaling` is true, the Autoscaler sidecar will be added to the Ray head pod.
+  # Ray Autoscaler integration is Beta with KubeRay >= 0.3.0 and Ray >= 2.0.0.
+  enableInTreeAutoscaling: true
+  autoscalerOptions:
+    upscalingMode: Default
+    idleTimeoutSeconds: 60
+    imagePullPolicy: IfNotPresent
+    resources:
+      limits:
+        cpu: "500m"
+        memory: "512Mi"
+      requests:
+        cpu: "500m"
+        memory: "512Mi"
+  headGroupSpec:
+    template:
+      spec:
+        containers:
+        - name: ray-head
+          image: rayproject/ray:2.46.0
+          ports:
+          - containerPort: 6379
+            name: gcs
+          - containerPort: 8265
+            name: dashboard
+          - containerPort: 10001
+            name: client
+          imagePullPolicy: IfNotPresent
+          lifecycle:
+            preStop:
+              exec:
+                command: ["/bin/sh","-c","ray stop"]
+          resources:
+            limits:
+              cpu: "1"
+              memory: "2G"
+            requests:
+              cpu: "500m"
+              memory: "2G"
+  workerGroupSpecs:
+  - replicas: 1
+    minReplicas: 1
+    maxReplicas: 10
+    groupName: small-group
+    template:
+      spec:
+        containers:
+        - name: ray-worker
+          image: rayproject/ray:2.46.0
+          lifecycle:
+            preStop:
+              exec:
+                command: ["/bin/sh","-c","ray stop"]
+          imagePullPolicy: IfNotPresent
+          resources:
+            limits:
+              cpu: "1"
+              memory: "1G"
+            requests:
+              cpu: "500m"
+              memory: "1G"

python/ray/tests/kuberay/test_files/ray-cluster.autoscaler-v2-template.yaml

@@ -0,0 +1,76 @@
+# This is a copy from the `ray-cluster.autoscaler-template.yaml` with modifications needed
+# to make kuberary work with Ray's Autoscaler V2.
+# See more at `ray-cluster.autoscaler-template.yaml` for the non autoscaler-v2 definition.
+apiVersion: ray.io/v1
+kind: RayCluster
+metadata:
+  name: raycluster-autoscaler
+spec:
+  # The version of Ray you are using. Make sure all Ray containers are running this version of Ray.
+  rayVersion: '2.46.0'
+  # If `enableInTreeAutoscaling` is true, the Autoscaler sidecar will be added to the Ray head pod.
+  # Ray Autoscaler integration is Beta with KubeRay >= 0.3.0 and Ray >= 2.0.0.
+  enableInTreeAutoscaling: true
+  autoscalerOptions:
+    # Use version: v2 instead of env var RAY_enable_autoscaler_v2 and restartPolicy: Never below if
+    # you're using KubeRay >= 1.4.0.
+    version: v2
+    upscalingMode: Default
+    idleTimeoutSeconds: 60
+    imagePullPolicy: IfNotPresent
+    resources:
+      limits:
+        cpu: "500m"
+        memory: "512Mi"
+      requests:
+        cpu: "500m"
+        memory: "512Mi"
+  headGroupSpec:
+    template:
+      spec:
+        containers:
+        - name: ray-head
+          image: rayproject/ray:2.46.0
+          ports:
+          - containerPort: 6379
+            name: gcs
+          - containerPort: 8265
+            name: dashboard
+          - containerPort: 10001
+            name: client
+          imagePullPolicy: IfNotPresent
+          lifecycle:
+            preStop:
+              exec:
+                command: ["/bin/sh","-c","ray stop"]
+          resources:
+            limits:
+              cpu: "1"
+              memory: "2G"
+            requests:
+              cpu: "500m"
+              memory: "2G"
+        restartPolicy: Never # No restart to avoid reuse of pod for different ray nodes.
+  workerGroupSpecs:
+  - replicas: 1
+    minReplicas: 1
+    maxReplicas: 10
+    groupName: small-group
+    template:
+      spec:
+        containers:
+        - name: ray-worker
+          image: rayproject/ray:2.46.0
+          lifecycle:
+            preStop:
+              exec:
+                command: ["/bin/sh","-c","ray stop"]
+          imagePullPolicy: IfNotPresent
+          resources:
+            limits:
+              cpu: "1"
+              memory: "1G"
+            requests:
+              cpu: "500m"
+              memory: "1G"
+        restartPolicy: Never # Never restart a pod to avoid pod reuse

python/ray/tests/kuberay/test_kuberay_node_provider.py

@@ -0,0 +1,371 @@
+import copy
+import sys
+from collections import defaultdict
+from pathlib import Path
+from typing import List, Set
+from unittest import mock
+
+import jsonpatch
+import pytest
+import yaml
+
+from ray.autoscaler._private.kuberay.node_provider import (
+    KubeRayNodeProvider,
+    ScaleRequest,
+    _worker_group_index,
+    _worker_group_max_replicas,
+    _worker_group_replicas,
+)
+from ray.autoscaler._private.util import NodeID
+from ray.autoscaler.batching_node_provider import NodeData
+from ray.tests.kuberay.test_autoscaling_config import get_basic_ray_cr
+
+
+def _get_basic_ray_cr_workers_to_delete(
+    cpu_workers_to_delete: List[NodeID],
+    gpu_workers_to_delete: List[NodeID],
+    tpu_workers_to_delete: List[NodeID],
+):
+    """Generate a Ray cluster with non-empty workersToDelete field."""
+    raycluster = get_basic_ray_cr()
+    raycluster["spec"]["workerGroupSpecs"][0]["scaleStrategy"] = {
+        "workersToDelete": cpu_workers_to_delete
+    }
+    raycluster["spec"]["workerGroupSpecs"][1]["scaleStrategy"] = {
+        "workersToDelete": gpu_workers_to_delete
+    }
+    raycluster["spec"]["workerGroupSpecs"][2]["scaleStrategy"] = {
+        "workersToDelete": tpu_workers_to_delete
+    }
+    return raycluster
+
+
+def _get_test_yaml(file_name):
+    file_path = str(Path(__file__).resolve().parent / "test_files" / file_name)
+    return yaml.safe_load(open(file_path).read())
+
+
+@pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.")
+@pytest.mark.parametrize(
+    "group_name,expected_index", [("small-group", 0), ("gpu-group", 1)]
+)
+def test_worker_group_index(group_name, expected_index):
+    """Basic unit test for _worker_group_index.
+
+    Uses a RayCluster CR with worker groups "small-group" and "gpu-group",
+    listed in that order.
+    """
+    raycluster_cr = get_basic_ray_cr()
+    assert _worker_group_index(raycluster_cr, group_name) == expected_index
+
+
+@pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.")
+@pytest.mark.parametrize(
+    "group_index,expected_max_replicas,expected_replicas",
+    [(0, 300, 1), (1, 200, 1), (2, 4, 1), (3, None, 0)],
+)
+def test_worker_group_replicas(group_index, expected_max_replicas, expected_replicas):
+    """Basic unit test for _worker_group_max_replicas and _worker_group_replicas
+
+    Uses a RayCluster CR with worker groups with 300 maxReplicas, 200 maxReplicas,
+    and unspecified maxReplicas, in that order.
+    """
+    raycluster = get_basic_ray_cr()
+
+    # Add a worker group without maxReplicas to confirm behavior
+    # when maxReplicas is not specified.
+    no_max_replicas_group = copy.deepcopy(raycluster["spec"]["workerGroupSpecs"][0])
+    no_max_replicas_group["groupName"] = "no-max-replicas"
+    del no_max_replicas_group["maxReplicas"]
+    # Also, replicas field, just for the sake of testing.
+    no_max_replicas_group["replicas"] = 0
+    raycluster["spec"]["workerGroupSpecs"].append(no_max_replicas_group)
+
+    assert _worker_group_max_replicas(raycluster, group_index) == expected_max_replicas
+    assert _worker_group_replicas(raycluster, group_index) == expected_replicas
+
+
+@pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.")
+@pytest.mark.parametrize(
+    "attempted_target_replica_count,expected_target_replica_count",
+    [(200, 200), (250, 250), (300, 300), (400, 300), (1000, 300)],
+)
+def test_create_node_cap_at_max(
+    attempted_target_replica_count: int, expected_target_replica_count: int
+):
+    """Validates that KubeRayNodeProvider does not attempt to create more nodes than
+    allowed by maxReplicas. For the config in this test, maxReplicas is fixed at 300.
+
+    Args:
+        attempted_target_replica_count: The mocked desired replica count for a given
+            worker group.
+        expected_target_replica_count: The actual requested replicaCount. Should be
+            capped at maxReplicas (300, for the config in this test.)
+    """
+    raycluster = get_basic_ray_cr()
+    with mock.patch.object(KubeRayNodeProvider, "__init__", return_value=None):
+        kr_node_provider = KubeRayNodeProvider(provider_config={}, cluster_name="fake")
+        scale_request = ScaleRequest(
+            workers_to_delete=set(),
+            desired_num_workers={"small-group": attempted_target_replica_count},
+        )
+        patch = kr_node_provider._scale_request_to_patch_payload(
+            scale_request=scale_request, raycluster=raycluster
+        )
+        assert patch[0]["value"] == expected_target_replica_count
+
+
+@pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.")
+@pytest.mark.parametrize(
+    "podlist_file,expected_node_data",
+    [
+        (
+            # Pod list obtained by running kubectl get pod -o yaml at runtime.
+            "podlist1.yaml",
+            {
+                "raycluster-autoscaler-head-8zsc8": NodeData(
+                    kind="head",
+                    type="headgroup",
+                    replica_index=None,
+                    ip="10.4.2.6",
+                    status="up-to-date",
+                ),  # up-to-date status because the Ray container is in running status
+                "raycluster-autoscaler-worker-small-group-dkz2r": NodeData(
+                    kind="worker",
+                    type="small-group",
+                    replica_index=None,
+                    ip="10.4.1.8",
+                    status="waiting",
+                ),  # waiting status, because Ray container's state is "waiting".
+                # The pod list includes a worker with non-null deletion timestamp.
+                # It is excluded from the node data because it is considered
+                # "terminated".
+            },
+        ),
+        (
+            # Pod list obtained by running kubectl get pod -o yaml at runtime.
+            "podlist2.yaml",
+            {
+                "raycluster-autoscaler-head-8zsc8": NodeData(
+                    kind="head",
+                    type="headgroup",
+                    replica_index=None,
+                    ip="10.4.2.6",
+                    status="up-to-date",
+                ),
+                "raycluster-autoscaler-worker-fake-gpu-group-2qnhv": NodeData(
+                    kind="worker",
+                    type="fake-gpu-group",
+                    replica_index=None,
+                    ip="10.4.0.6",
+                    status="up-to-date",
+                ),
+                "raycluster-autoscaler-worker-small-group-dkz2r": NodeData(
+                    kind="worker",
+                    type="small-group",
+                    replica_index=None,
+                    ip="10.4.1.8",
+                    status="up-to-date",
+                ),
+                "raycluster-autoscaler-worker-small-group-lbfm4": NodeData(
+                    kind="worker",
+                    type="small-group",
+                    replica_index=None,
+                    ip="10.4.0.5",
+                    status="up-to-date",
+                ),
+                "raycluster-autoscaler-tpu-group-worker-s8jhq": NodeData(
+                    kind="worker",
+                    type="tpu-group",
+                    replica_index="tpu-group-0",
+                    ip="10.24.9.4",
+                    status="up-to-date",
+                ),
+                "raycluster-autoscaler-tpu-group-worker-jd69f": NodeData(
+                    kind="worker",
+                    type="tpu-group",
+                    replica_index="tpu-group-0",
+                    ip="10.24.8.4",
+                    status="up-to-date",
+                ),
+            },
+        ),
+    ],
+)
+def test_get_node_data(podlist_file: str, expected_node_data):
+    """Test translation of a K8s pod list into autoscaler node data."""
+    pod_list = _get_test_yaml(podlist_file)
+
+    def mock_get(node_provider, path):
+        if "pods" in path:
+            return pod_list
+        elif "raycluster" in path:
+            return get_basic_ray_cr()
+        else:
+            raise ValueError("Invalid path.")
+
+    with mock.patch.object(
+        KubeRayNodeProvider, "__init__", return_value=None
+    ), mock.patch.object(KubeRayNodeProvider, "_get", mock_get):
+        kr_node_provider = KubeRayNodeProvider(provider_config={}, cluster_name="fake")
+        kr_node_provider.cluster_name = "fake"
+        kr_node_provider.replica_index_to_nodes = defaultdict(list[str])
+        nodes = kr_node_provider.non_terminated_nodes({})
+        assert kr_node_provider.node_data_dict == expected_node_data
+        assert set(nodes) == set(expected_node_data.keys())
+
+
+@pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.")
+@pytest.mark.parametrize(
+    "node_data_dict,scale_request,expected_patch_payload",
+    [
+        (
+            {
+                "raycluster-autoscaler-head-8zsc8": NodeData(
+                    kind="head",
+                    type="headgroup",
+                    replica_index=None,
+                    ip="10.4.2.6",
+                    status="up-to-date",
+                ),
+                "raycluster-autoscaler-worker-fake-gpu-group-2qnhv": NodeData(
+                    kind="worker",
+                    type="fake-gpu-group",
+                    replica_index=None,
+                    ip="10.4.0.6",
+                    status="up-to-date",
+                ),
+                "raycluster-autoscaler-worker-small-group-dkz2r": NodeData(
+                    kind="worker",
+                    type="small-group",
+                    replica_index=None,
+                    ip="10.4.1.8",
+                    status="up-to-date",
+                ),
+                "raycluster-autoscaler-worker-small-group-lbfm4": NodeData(
+                    kind="worker",
+                    type="small-group",
+                    replica_index=None,
+                    ip="10.4.0.5",
+                    status="up-to-date",
+                ),
+            },
+            ScaleRequest(
+                desired_num_workers={
+                    "small-group": 1,  # Delete 1
+                    "gpu-group": 1,  # Don't touch
+                    "blah-group": 5,  # Create 5
+                },
+                workers_to_delete={
+                    "raycluster-autoscaler-worker-small-group-dkz2r",
+                },
+            ),
+            [
+                {
+                    "op": "replace",
+                    "path": "/spec/workerGroupSpecs/3/replicas",
+                    "value": 5,
+                },
+                {
+                    "op": "replace",
+                    "path": "/spec/workerGroupSpecs/0/scaleStrategy",
+                    "value": {
+                        "workersToDelete": [
+                            "raycluster-autoscaler-worker-small-group-dkz2r"
+                        ]
+                    },
+                },
+            ],
+        ),
+    ],
+)
+def test_submit_scale_request(node_data_dict, scale_request, expected_patch_payload):
+    """Test the KubeRayNodeProvider's RayCluster patch payload given a dict
+    of current node counts and a scale request.
+    """
+    raycluster = get_basic_ray_cr()
+    # Add another worker group for the sake of this test.
+    blah_group = copy.deepcopy(raycluster["spec"]["workerGroupSpecs"][1])
+    blah_group["groupName"] = "blah-group"
+    raycluster["spec"]["workerGroupSpecs"].append(blah_group)
+    with mock.patch.object(KubeRayNodeProvider, "__init__", return_value=None):
+        kr_node_provider = KubeRayNodeProvider(provider_config={}, cluster_name="fake")
+        kr_node_provider.node_data_dict = node_data_dict
+        patch_payload = kr_node_provider._scale_request_to_patch_payload(
+            scale_request=scale_request, raycluster=raycluster
+        )
+        assert patch_payload == expected_patch_payload
+
+
+@pytest.mark.parametrize("node_set", [{"A", "B", "C", "D", "E"}])
+@pytest.mark.parametrize("cpu_workers_to_delete", ["A", "Z"])
+@pytest.mark.parametrize("gpu_workers_to_delete", ["B", "Y"])
+@pytest.mark.parametrize("tpu_workers_to_delete", ["C", "X"])
+@pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.")
+def test_safe_to_scale(
+    node_set: Set[NodeID],
+    cpu_workers_to_delete: List[NodeID],
+    gpu_workers_to_delete: List[NodeID],
+    tpu_workers_to_delete: List[NodeID],
+):
+    # NodeData values unimportant for this test.
+    mock_node_data = NodeData("-", "-", "-", "-", "-")
+    node_data_dict = {node_id: mock_node_data for node_id in node_set}
+
+    raycluster = _get_basic_ray_cr_workers_to_delete(
+        cpu_workers_to_delete, gpu_workers_to_delete, tpu_workers_to_delete
+    )
+
+    def mock_patch(kuberay_provider, path, patch_payload):
+        patch = jsonpatch.JsonPatch(patch_payload)
+        kuberay_provider._patched_raycluster = patch.apply(kuberay_provider._raycluster)
+
+    with mock.patch.object(
+        KubeRayNodeProvider, "__init__", return_value=None
+    ), mock.patch.object(KubeRayNodeProvider, "_patch", mock_patch):
+        kr_node_provider = KubeRayNodeProvider(provider_config={}, cluster_name="fake")
+        kr_node_provider.cluster_name = "fake"
+        kr_node_provider._patched_raycluster = raycluster
+        kr_node_provider._raycluster = raycluster
+        kr_node_provider.node_data_dict = node_data_dict
+        actual_safe = kr_node_provider.safe_to_scale()
+
+    expected_safe = (
+        not any(
+            cpu_worker_to_delete in node_set
+            for cpu_worker_to_delete in cpu_workers_to_delete
+        )
+        and not any(
+            gpu_worker_to_delete in node_set
+            for gpu_worker_to_delete in gpu_workers_to_delete
+        )
+        and not any(
+            tpu_worker_to_delete in node_set
+            for tpu_worker_to_delete in tpu_workers_to_delete
+        )
+    )
+    assert expected_safe is actual_safe
+    patched_cpu_workers_to_delete = kr_node_provider._patched_raycluster["spec"][
+        "workerGroupSpecs"
+    ][0]["scaleStrategy"]["workersToDelete"]
+    patched_gpu_workers_to_delete = kr_node_provider._patched_raycluster["spec"][
+        "workerGroupSpecs"
+    ][1]["scaleStrategy"]["workersToDelete"]
+    patched_tpu_workers_to_delete = kr_node_provider._patched_raycluster["spec"][
+        "workerGroupSpecs"
+    ][2]["scaleStrategy"]["workersToDelete"]
+
+    if expected_safe:
+        # Cleaned up workers to delete
+        assert patched_cpu_workers_to_delete == []
+        assert patched_gpu_workers_to_delete == []
+        assert patched_tpu_workers_to_delete == []
+    else:
+        # Did not clean up workers to delete
+        assert patched_cpu_workers_to_delete == cpu_workers_to_delete
+        assert patched_gpu_workers_to_delete == gpu_workers_to_delete
+        assert patched_tpu_workers_to_delete == tpu_workers_to_delete
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tests/kuberay/utils.py

@@ -0,0 +1,484 @@
+"""Utilities for e2e tests of KubeRay/Ray integration.
+For consistency, all K8s interactions use kubectl through subprocess calls.
+"""
+import atexit
+import contextlib
+import logging
+import os
+import pathlib
+import subprocess
+import tempfile
+import time
+from typing import Any, Dict, Generator, List, Optional
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+SCRIPTS_DIR = pathlib.Path(__file__).resolve().parent / "scripts"
+TEST_CR_PATH = (
+    pathlib.Path(__file__).resolve().parent / "setup" / "raycluster_test.yaml"
+)
+TEST_CLUSTER_NAME = "raycluster-test"
+
+# Parent directory of Ray repository
+RAY_PARENT = str(pathlib.Path(__file__).resolve().parents[5])
+
+RAYCLUSTERS_QUALIFIED = "rayclusters.ray.io"
+
+LOG_FORMAT = "[%(levelname)s %(asctime)s] %(filename)s: %(lineno)d  %(message)s"
+
+
+def switch_to_ray_parent_dir():
+    # Switch to parent of Ray repo, because that's what the doc examples do.
+    logger.info("Switching to parent of Ray directory.")
+    os.chdir(RAY_PARENT)
+
+
+def setup_kuberay_operator():
+    """Set up KubeRay operator and Ray autoscaler RBAC."""
+    switch_to_ray_parent_dir()
+    logger.info("Cloning KubeRay and setting up KubeRay configuration.")
+    # For faster run-time when triggering the test locally, don't run the init
+    # script if it has already been run.
+    subprocess.check_call(
+        [
+            "bash",
+            "-c",
+            (
+                "ls ray/python/ray/autoscaler/kuberay/config ||"
+                " ./ray/python/ray/autoscaler/kuberay/init-config.sh"
+            ),
+        ]
+    )
+    logger.info("Creating KubeRay operator.")
+    subprocess.check_call(
+        [
+            "kubectl",
+            "create",
+            "-k",
+            "ray/python/ray/autoscaler/kuberay/config/default",
+        ]
+    )
+
+
+def teardown_kuberay_operator():
+    logger.info("Switching to parent of Ray directory.")
+    os.chdir(RAY_PARENT)
+
+    logger.info("Deleting operator.")
+    subprocess.check_call(
+        [
+            "kubectl",
+            "delete",
+            "--ignore-not-found",
+            "-k",
+            "ray/python/ray/autoscaler/kuberay/config/default",
+        ]
+    )
+
+    logger.info("Double-checking no pods left over in namespace ray-system.")
+    wait_for_pods(goal_num_pods=0, namespace="ray-system")
+
+
+def wait_for_raycluster_crd(tries=60, backoff_s=5):
+    """CRD creation can take a bit of time after the client request.
+    This function waits until the crd with the provided name is registered.
+    """
+    switch_to_ray_parent_dir()
+    logger.info("Making sure RayCluster CRD has been registered.")
+    for i in range(tries):
+        get_crd_output = subprocess.check_output(["kubectl", "get", "crd"]).decode()
+        if RAYCLUSTERS_QUALIFIED in get_crd_output:
+            logger.info("Confirmed existence of RayCluster CRD.")
+            break
+        elif i < tries - 1:
+            logger.info("Still waiting to register RayCluster CRD.")
+            time.sleep(backoff_s)
+        else:
+            raise Exception("Failed to register RayCluster CRD.")
+
+    # Create a test RayCluster CR to make sure that the CRD is fully registered.
+    for i in range(tries):
+        try:
+            subprocess.check_call(["kubectl", "apply", "-f", TEST_CR_PATH])
+            break
+        except subprocess.CalledProcessError as e:
+            logger.info("Can't create RayCluster CR.")
+            if i < tries - 1:
+                logger.info("Retrying.")
+                time.sleep(backoff_s)
+            else:
+                logger.info("Giving up.")
+                raise e from None
+
+    # Confirm the test RayCluster exists.
+    out = subprocess.check_output(["kubectl", "get", RAYCLUSTERS_QUALIFIED]).decode()
+    assert TEST_CLUSTER_NAME in out, out
+
+    # Delete the test RayCluster.
+    subprocess.check_call(["kubectl", "delete", "-f", TEST_CR_PATH])
+    # Make sure the associated resources are gone before proceeding.
+    wait_for_pods(goal_num_pods=0, namespace="default")
+
+
+def wait_for_pods(goal_num_pods: int, namespace: str, tries=60, backoff_s=5) -> None:
+    """Wait for the number of pods in the `namespace` to be exactly `num_pods`.
+
+    Raise an exception after exceeding `tries` attempts with `backoff_s` second waits.
+    """
+    for i in range(tries):
+        cur_num_pods = _get_num_pods(namespace)
+        if cur_num_pods == goal_num_pods:
+            logger.info(f"Confirmed {goal_num_pods} pod(s) in namespace {namespace}.")
+            return
+        elif i < tries - 1:
+            logger.info(
+                f"The number of pods in namespace {namespace} is {cur_num_pods}."
+                f" Waiting until the number of pods is {goal_num_pods}."
+                f"{get_pod_names(namespace)}"
+            )
+            time.sleep(backoff_s)
+        else:
+            raise Exception(
+                f"Failed to scale to {goal_num_pods} pod(s) in namespace {namespace}."
+            )
+
+
+def _get_num_pods(namespace: str) -> int:
+    return len(get_pod_names(namespace))
+
+
+def get_pod_names(namespace: str) -> List[str]:
+    """Get the list of pod names in the namespace."""
+    get_pods_output = (
+        subprocess.check_output(
+            [
+                "kubectl",
+                "-n",
+                namespace,
+                "get",
+                "pods",
+                "-o",
+                "custom-columns=POD:metadata.name",
+                "--no-headers",
+            ]
+        )
+        .decode()
+        .strip()
+    )
+
+    # If there aren't any pods, the output is any empty string.
+    if not get_pods_output:
+        return []
+    else:
+        return get_pods_output.split("\n")
+
+
+def wait_for_pod_to_start(
+    pod_name_filter: str, namespace: str, tries=60, backoff_s=5
+) -> None:
+    """Waits for a pod to have Running status.phase.
+
+    More precisely, waits until there is a pod with name containing `pod_name_filter`
+    and the pod has Running status.phase."""
+    for i in range(tries):
+        pod = get_pod(pod_name_filter=pod_name_filter, namespace=namespace)
+        if not pod:
+            # We didn't get a matching pod.
+            continue
+        pod_status = (
+            subprocess.check_output(
+                [
+                    "kubectl",
+                    "-n",
+                    namespace,
+                    "get",
+                    "pod",
+                    pod,
+                    "-o",
+                    "custom-columns=POD:status.phase",
+                    "--no-headers",
+                ]
+            )
+            .decode()
+            .strip()
+        )
+        # "not found" is part of the kubectl output if the pod's not there.
+        if "not found" in pod_status:
+            raise Exception(f"Pod {pod} not found.")
+        elif pod_status == "Running":
+            logger.info(f"Confirmed pod {pod} is Running.")
+            return
+        elif i < tries - 1:
+            logger.info(
+                f"Pod {pod} has status {pod_status}. Waiting for the pod to enter "
+                "Running status."
+            )
+            time.sleep(backoff_s)
+        else:
+            raise Exception(f"Timed out waiting for pod {pod} to enter Running status.")
+
+
+def wait_for_ray_health(
+    pod_name_filter: str,
+    namespace: str,
+    tries=60,
+    backoff_s=5,
+    ray_container="ray-head",
+) -> None:
+    """Waits until a Ray pod passes `ray health-check`.
+
+    More precisely, waits until a Ray pod whose name includes the string
+    `pod_name_filter` passes `ray health-check`.
+    (Ensures Ray has completely started in the pod.)
+
+    Use case: Wait until there is a Ray head pod with Ray running on it.
+    """
+    for i in range(tries):
+        try:
+            pod = get_pod(pod_name_filter=pod_name_filter, namespace="default")
+            assert pod, f"Couldn't find a pod matching {pod_name_filter}."
+            # `ray health-check` yields 0 exit status iff it succeeds
+            kubectl_exec(
+                ["ray", "health-check"], pod, namespace, container=ray_container
+            )
+            logger.info(f"ray health check passes for pod {pod}")
+            return
+        except subprocess.CalledProcessError as e:
+            logger.info(f"Failed ray health check for pod {pod}.")
+            if i < tries - 1:
+                logger.info("Trying again.")
+                time.sleep(backoff_s)
+            else:
+                logger.info("Giving up.")
+                raise e from None
+
+
+def get_pod(pod_name_filter: str, namespace: str) -> Optional[str]:
+    """Gets pods in the `namespace`.
+
+    Returns the first pod that has `pod_name_filter` as a
+    substring of its name. Returns None if there are no matches.
+    """
+    pod_names = get_pod_names(namespace)
+    matches = [pod_name for pod_name in pod_names if pod_name_filter in pod_name]
+    if not matches:
+        logger.warning(f"No match for `{pod_name_filter}` in namespace `{namespace}`.")
+        return None
+    return matches[0]
+
+
+def kubectl_exec(
+    command: List[str],
+    pod: str,
+    namespace: str,
+    container: Optional[str] = None,
+) -> str:
+    """kubectl exec the `command` in the given `pod` in the given `namespace`.
+    If a `container` is specified, will specify that container for kubectl.
+
+    Prints and return kubectl's output as a string.
+    """
+    container_option = ["-c", container] if container else []
+    kubectl_exec_command = (
+        ["kubectl", "exec", "-it", pod] + container_option + ["--"] + command
+    )
+    # Print for debugging convenience.
+    try:
+        out = subprocess.check_output(kubectl_exec_command).decode().strip()
+    except subprocess.CalledProcessError as e:
+        logger.error(f"Error running command {kubectl_exec_command}.")
+        logger.error(f"Output: {e.output.decode()}")
+        raise e from None
+    print(out)
+    return out
+
+
+def kubectl_logs(
+    pod: str,
+    namespace: str,
+    container: Optional[str] = None,
+) -> str:
+    """Wrapper for kubectl logs.
+
+    Returns the logs as a string.
+    """
+    container_option = ["-c", container] if container else []
+    kubectl_logs_command = ["kubectl", "logs", pod] + container_option
+    out = subprocess.check_output(kubectl_logs_command).decode().strip()
+    return out
+
+
+def kubectl_exec_python_script(
+    script_name: str,
+    pod: str,
+    namespace: str,
+    container: Optional[str] = None,
+) -> str:
+    """
+    Runs a python script in a container via `kubectl exec`.
+    Scripts live in `tests/kuberay/scripts`.
+
+    Prints and return kubectl's output as a string.
+    """
+    script_path = SCRIPTS_DIR / script_name
+    with open(script_path) as script_file:
+        script_string = script_file.read()
+    return kubectl_exec(["python", "-c", script_string], pod, namespace, container)
+
+
+def get_raycluster(raycluster: str, namespace: str) -> Dict[str, Any]:
+    """Gets the Ray CR with name `raycluster` in namespace `namespace`.
+
+    Returns the CR as a nested Dict.
+    """
+    get_raycluster_output = (
+        subprocess.check_output(
+            ["kubectl", "-n", namespace, "get", "raycluster", raycluster, "-o", "yaml"]
+        )
+        .decode()
+        .strip()
+    )
+    return yaml.safe_load(get_raycluster_output)
+
+
+def _get_service_port(service: str, namespace: str, target_port: int) -> int:
+    """Given a K8s service and a port targetted by the service, returns the
+    corresponding port exposed by the service.
+
+    Args:
+        service: Name of a K8s service.
+        namespace: Namespace to which the service belongs.
+        target_port: Port targeted by the service.
+
+    Returns:
+        service_port: The port exposed by the service.
+    """
+    service_str = (
+        subprocess.check_output(
+            ["kubectl", "-n", namespace, "get", "service", service, "-o", "yaml"]
+        )
+        .decode()
+        .strip()
+    )
+    service_dict = yaml.safe_load(service_str)
+    service_ports: List = service_dict["spec"]["ports"]
+    matching_ports = [
+        port for port in service_ports if port["targetPort"] == target_port
+    ]
+    assert matching_ports
+    service_port = matching_ports[0]["port"]
+    return service_port
+
+
+@contextlib.contextmanager
+def _kubectl_port_forward(
+    service: str, namespace: str, target_port: int, local_port: Optional[int] = None
+) -> Generator[int, None, None]:
+    """Context manager which creates a kubectl port-forward process targeting a
+    K8s service.
+
+    Terminates the port-forwarding process upon exit.
+
+    Args:
+        service: Name of a K8s service.
+        namespace: Namespace to which the service belongs.
+        target_port: The port targeted by the service.
+        local_port: Forward from this port. Optional. By default, uses the port exposed
+            by the service.
+
+    Yields:
+        int: The local port. The service can then be accessed at
+            127.0.0.1:<local_port>.
+    """
+    # First, figure out which port the service exposes for the given target port.
+    service_port = _get_service_port(service, namespace, target_port)
+    if not local_port:
+        local_port = service_port
+
+    process = subprocess.Popen(
+        [
+            "kubectl",
+            "-n",
+            namespace,
+            "port-forward",
+            f"service/{service}",
+            f"{local_port}:{service_port}",
+        ]
+    )
+
+    def terminate_process():
+        process.terminate()
+        # Wait 10 seconds for the process to terminate.
+        # This cleans up the zombie entry from the process table.
+        # 10 seconds is a deliberately excessive amount of time to wait.
+        process.wait(timeout=10)
+
+    # Ensure clean-up in case of interrupt.
+    atexit.register(terminate_process)
+    # terminate_process is ok to execute multiple times.
+
+    try:
+        yield local_port
+    finally:
+        terminate_process()
+
+
+def kubectl_patch(
+    kind: str,
+    name: str,
+    namespace: str,
+    patch: Dict[str, Any],
+    patch_type: str = "strategic",
+):
+    """Wrapper for kubectl patch.
+
+    Args:
+        kind: Kind of the K8s resource (e.g. pod)
+        name: Name of the K8s resource.
+        namespace: Namespace of the K8s resource.
+        patch: The patch to apply, as a dict.
+        patch_type: json, merge, or strategic
+    """
+    with tempfile.NamedTemporaryFile("w") as patch_file:
+        yaml.dump(patch, patch_file)
+        patch_file.flush()
+        subprocess.check_call(
+            [
+                "kubectl",
+                "-n",
+                f"{namespace}",
+                "patch",
+                f"{kind}",
+                f"{name}",
+                "--patch-file",
+                f"{patch_file.name}",
+                "--type",
+                f"{patch_type}",
+            ]
+        )
+
+
+def kubectl_delete(kind: str, name: str, namespace: str, wait: bool = True):
+    """Wrapper for kubectl delete.
+
+    Args:
+        kind: Kind of the K8s resource (e.g. pod)
+        name: Name of the K8s resource.
+        namespace: Namespace of the K8s resource.
+        wait: Whether to pass ``--wait=true`` so ``kubectl`` blocks until the
+            resource is fully removed.
+    """
+    wait_str = "true" if wait else "false"
+    subprocess.check_output(
+        [
+            "kubectl",
+            "-n",
+            f"{namespace}",
+            "delete",
+            f"{kind}",
+            f"{name}",
+            f"--wait={wait_str}",
+        ]
+    )

python/ray/tests/ludwig/BUILD.bazel

@@ -0,0 +1,12 @@
+load("@rules_python//python:defs.bzl", "py_test")
+
+py_test(
+    name = "test_ludwig",
+    size = "medium",
+    srcs = ["test_ludwig.py"],
+    tags = [
+        "exclusive",
+        "team:ml",
+    ],
+    deps = ["//:ray_lib"],
+)

python/ray/tests/ludwig/ludwig_test_utils.py

@@ -0,0 +1,643 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2019 Uber Technologies, Inc.
+#
+# 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.
+# ==============================================================================
+
+# This file is copied and adapted from
+# https://github.com/ludwig-ai/ludwig/blob/master/tests/integration_tests/utils.py
+
+import multiprocessing
+import os
+import random
+import shutil
+import sys
+import traceback
+import unittest
+import uuid
+from distutils.util import strtobool
+from typing import Any, Dict, List, Optional
+
+import cloudpickle
+import numpy as np
+import pandas as pd
+from ludwig.api import LudwigModel
+from ludwig.backend import LocalBackend
+from ludwig.constants import COLUMN, NAME, PROC_COLUMN, VECTOR
+from ludwig.data.dataset_synthesizer import DATETIME_FORMATS, build_synthetic_dataset
+from ludwig.experiment import experiment_cli
+from ludwig.features.feature_utils import compute_feature_hash
+from ludwig.utils.data_utils import read_csv, replace_file_extension
+
+ENCODERS = [
+    "embed",
+    "rnn",
+    "parallel_cnn",
+    "cnnrnn",
+    "stacked_parallel_cnn",
+    "stacked_cnn",
+    "transformer",
+]
+
+HF_ENCODERS_SHORT = ["distilbert"]
+
+HF_ENCODERS = [
+    "bert",
+    "gpt",
+    "gpt2",
+    # "transformer_xl",
+    "xlnet",
+    "xlm",
+    "roberta",
+    "distilbert",
+    "ctrl",
+    "camembert",
+    "albert",
+    "t5",
+    "xlmroberta",
+    "longformer",
+    "flaubert",
+    "electra",
+    "mt5",
+]
+
+
+class LocalTestBackend(LocalBackend):
+    @property
+    def supports_multiprocessing(self):
+        return False
+
+
+def parse_flag_from_env(key, default=False):
+    try:
+        value = os.environ[key]
+    except KeyError:
+        # KEY isn't set, default to `default`.
+        _value = default
+    else:
+        # KEY is set, convert it to True or False.
+        try:
+            _value = strtobool(value)
+        except ValueError:
+            # More values are supported, but let's keep the message simple.
+            raise ValueError("If set, {} must be yes or no.".format(key))
+    return _value
+
+
+_run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False)
+
+
+def slow(test_case):
+    """
+    Decorator marking a test as slow.
+
+    Slow tests are skipped by default. Set the RUN_SLOW environment variable
+    to a truth value to run them.
+
+    """
+    if not _run_slow_tests:
+        test_case = unittest.skip("Skipping: this test is too slow")(test_case)
+    return test_case
+
+
+def generate_data(
+    input_features: List[Dict[str, Any]],
+    output_features: List[Dict[str, Any]],
+    filename: str = "test_csv.csv",
+    num_examples: int = 25,
+) -> str:
+    """Generate synthetic data based on input/output feature specs.
+
+    Args:
+        input_features: Input feature schema.
+        output_features: Output feature schema.
+        filename: Path to the file where data is stored.
+        num_examples: Number of examples to generate.
+
+    Returns:
+        The path to the file where the generated data was written.
+    """
+    features = input_features + output_features
+    df = build_synthetic_dataset(num_examples, features)
+    data = [next(df) for _ in range(num_examples)]
+
+    dataframe = pd.DataFrame(data[1:], columns=data[0])
+    dataframe.to_csv(filename, index=False)
+
+    return filename
+
+
+def random_string(length=5):
+    return uuid.uuid4().hex[:length].upper()
+
+
+def numerical_feature(normalization=None, **kwargs):
+    feature = {
+        "name": "num_" + random_string(),
+        "type": "number",
+        "preprocessing": {"normalization": normalization},
+    }
+    feature.update(kwargs)
+    feature[COLUMN] = feature[NAME]
+    feature[PROC_COLUMN] = compute_feature_hash(feature)
+    return feature
+
+
+def category_feature(**kwargs):
+    feature = {
+        "type": "category",
+        "name": "category_" + random_string(),
+        "vocab_size": 10,
+        "embedding_size": 5,
+    }
+    feature.update(kwargs)
+    feature[COLUMN] = feature[NAME]
+    feature[PROC_COLUMN] = compute_feature_hash(feature)
+    return feature
+
+
+def text_feature(**kwargs):
+    feature = {
+        "name": "text_" + random_string(),
+        "type": "text",
+        "reduce_input": None,
+        "vocab_size": 5,
+        "min_len": 7,
+        "max_len": 7,
+        "embedding_size": 8,
+        "state_size": 8,
+    }
+    feature.update(kwargs)
+    feature[COLUMN] = feature[NAME]
+    feature[PROC_COLUMN] = compute_feature_hash(feature)
+    return feature
+
+
+def set_feature(**kwargs):
+    feature = {
+        "type": "set",
+        "name": "set_" + random_string(),
+        "vocab_size": 10,
+        "max_len": 5,
+        "embedding_size": 5,
+    }
+    feature.update(kwargs)
+    feature[COLUMN] = feature[NAME]
+    feature[PROC_COLUMN] = compute_feature_hash(feature)
+    return feature
+
+
+def sequence_feature(**kwargs):
+    feature = {
+        "type": "sequence",
+        "name": "sequence_" + random_string(),
+        "vocab_size": 10,
+        "max_len": 7,
+        "encoder": "embed",
+        "embedding_size": 8,
+        "fc_size": 8,
+        "state_size": 8,
+        "num_filters": 8,
+        "hidden_size": 8,
+    }
+    feature.update(kwargs)
+    feature[COLUMN] = feature[NAME]
+    feature[PROC_COLUMN] = compute_feature_hash(feature)
+    return feature
+
+
+def image_feature(folder, **kwargs):
+    feature = {
+        "type": "image",
+        "name": "image_" + random_string(),
+        "encoder": "resnet",
+        "preprocessing": {
+            "in_memory": True,
+            "height": 12,
+            "width": 12,
+            "num_channels": 3,
+        },
+        "resnet_size": 8,
+        "destination_folder": folder,
+        "fc_size": 8,
+        "num_filters": 8,
+    }
+    feature.update(kwargs)
+    feature[COLUMN] = feature[NAME]
+    feature[PROC_COLUMN] = compute_feature_hash(feature)
+    return feature
+
+
+def audio_feature(folder, **kwargs):
+    feature = {
+        "name": "audio_" + random_string(),
+        "type": "audio",
+        "preprocessing": {
+            "audio_feature": {
+                "type": "fbank",
+                "window_length_in_s": 0.04,
+                "window_shift_in_s": 0.02,
+                "num_filter_bands": 80,
+            },
+            "audio_file_length_limit_in_s": 3.0,
+        },
+        "encoder": "stacked_cnn",
+        "should_embed": False,
+        "conv_layers": [
+            {
+                "filter_size": 400,
+                "pool_size": 16,
+                "num_filters": 32,
+                "regularize": "false",
+            },
+            {
+                "filter_size": 40,
+                "pool_size": 10,
+                "num_filters": 64,
+                "regularize": "false",
+            },
+        ],
+        "fc_size": 256,
+        "destination_folder": folder,
+    }
+    feature.update(kwargs)
+    feature[COLUMN] = feature[NAME]
+    feature[PROC_COLUMN] = compute_feature_hash(feature)
+    return feature
+
+
+def timeseries_feature(**kwargs):
+    feature = {
+        "name": "timeseries_" + random_string(),
+        "type": "timeseries",
+        "max_len": 7,
+    }
+    feature.update(kwargs)
+    feature[COLUMN] = feature[NAME]
+    feature[PROC_COLUMN] = compute_feature_hash(feature)
+    return feature
+
+
+def binary_feature(**kwargs):
+    feature = {"name": "binary_" + random_string(), "type": "binary"}
+    feature.update(kwargs)
+    feature[COLUMN] = feature[NAME]
+    feature[PROC_COLUMN] = compute_feature_hash(feature)
+    return feature
+
+
+def bag_feature(**kwargs):
+    feature = {
+        "name": "bag_" + random_string(),
+        "type": "bag",
+        "max_len": 5,
+        "vocab_size": 10,
+        "embedding_size": 5,
+    }
+    feature.update(kwargs)
+    feature[COLUMN] = feature[NAME]
+    feature[PROC_COLUMN] = compute_feature_hash(feature)
+    return feature
+
+
+def date_feature(**kwargs):
+    feature = {
+        "name": "date_" + random_string(),
+        "type": "date",
+        "preprocessing": {
+            "datetime_format": random.choice(list(DATETIME_FORMATS.keys()))
+        },
+    }
+    feature.update(kwargs)
+    feature[COLUMN] = feature[NAME]
+    feature[PROC_COLUMN] = compute_feature_hash(feature)
+    return feature
+
+
+def h3_feature(**kwargs):
+    feature = {"name": "h3_" + random_string(), "type": "h3"}
+    feature.update(kwargs)
+    feature[COLUMN] = feature[NAME]
+    feature[PROC_COLUMN] = compute_feature_hash(feature)
+    return feature
+
+
+def vector_feature(**kwargs):
+    feature = {"type": VECTOR, "vector_size": 5, "name": "vector_" + random_string()}
+    feature.update(kwargs)
+    feature[COLUMN] = feature[NAME]
+    feature[PROC_COLUMN] = compute_feature_hash(feature)
+    return feature
+
+
+def run_experiment(
+    input_features: Optional[List[Dict[str, Any]]],
+    output_features: Optional[List[Dict[str, Any]]],
+    skip_save_processed_input: bool = True,
+    config: Optional[Dict[str, Any]] = None,
+    backend: Optional[LocalBackend] = None,
+    **kwargs,
+) -> None:
+    """Run an experiment and clean up artifacts saved to disk.
+
+    Args:
+        input_features: List of input feature dictionaries.
+        output_features: List of output feature dictionaries.
+        skip_save_processed_input: Whether to skip persisting processed input
+            to disk.
+        config: Optional Ludwig configuration dictionary. If unset, a default
+            config is constructed from ``input_features`` and
+            ``output_features``.
+        backend: Optional Ludwig backend to use. Defaults to
+            ``LocalTestBackend()``.
+        **kwargs: Extra keyword arguments forwarded to the underlying
+            ``experiment_cli`` call.
+    """
+    if input_features is not None and output_features is not None:
+        # This if is necessary so that the caller can call with
+        # config_file (and not config)
+        config = {
+            "input_features": input_features,
+            "output_features": output_features,
+            "combiner": {"type": "concat", "fc_size": 14},
+            "training": {"epochs": 2},
+        }
+
+    args = {
+        "config": config,
+        "backend": backend or LocalTestBackend(),
+        "skip_save_training_description": True,
+        "skip_save_training_statistics": True,
+        "skip_save_processed_input": skip_save_processed_input,
+        "skip_save_progress": True,
+        "skip_save_unprocessed_output": True,
+        "skip_save_model": True,
+        "skip_save_predictions": True,
+        "skip_save_eval_stats": True,
+        "skip_collect_predictions": True,
+        "skip_collect_overall_stats": True,
+        "skip_save_log": True,
+    }
+    args.update(kwargs)
+
+    _, _, _, _, exp_dir_name = experiment_cli(**args)
+    shutil.rmtree(exp_dir_name, ignore_errors=True)
+
+
+def generate_output_features_with_dependencies(main_feature, dependencies):
+    # helper function to generate multiple output features specifications
+    # with dependencies, support for 'test_experiment_multiple_seq_seq` unit
+    # test
+    # Parameters:
+    # main_feature: feature identifier, valid values 'feat1', 'feat2', 'feat3'
+    # dependencies: list of dependencies for 'main_feature', do not li
+    # Example:
+    #  generate_output_features_with_dependencies('feat2', ['feat1', 'feat3'])
+
+    output_features = [
+        category_feature(vocab_size=2, reduce_input="sum"),
+        sequence_feature(vocab_size=10, max_len=5),
+        numerical_feature(),
+    ]
+
+    # value portion of dictionary is a tuple: (position, feature_name)
+    #   position: location of output feature in the above output_features list
+    #   feature_name: Ludwig generated feature name
+    feature_names = {
+        "feat1": (0, output_features[0]["name"]),
+        "feat2": (1, output_features[1]["name"]),
+        "feat3": (2, output_features[2]["name"]),
+    }
+
+    # generate list of dependencies with real feature names
+    generated_dependencies = [feature_names[feat_name][1] for feat_name in dependencies]
+
+    # specify dependencies for the main_feature
+    output_features[feature_names[main_feature][0]][
+        "dependencies"
+    ] = generated_dependencies
+
+    return output_features
+
+
+def _subproc_wrapper(fn, queue, *args, **kwargs):
+    fn = cloudpickle.loads(fn)
+    try:
+        results = fn(*args, **kwargs)
+    except Exception as e:
+        traceback.print_exc(file=sys.stderr)
+        results = e
+    queue.put(results)
+
+
+def spawn(fn):
+    def wrapped_fn(*args, **kwargs):
+        ctx = multiprocessing.get_context("spawn")
+        queue = ctx.Queue()
+
+        p = ctx.Process(
+            target=_subproc_wrapper,
+            args=(cloudpickle.dumps(fn), queue, *args),
+            kwargs=kwargs,
+        )
+
+        p.start()
+        p.join()
+        results = queue.get()
+        if isinstance(results, Exception):
+            raise RuntimeError(
+                f"Spawned subprocess raised {type(results).__name__}, "
+                f"check log output above for stack trace."
+            )
+        return results
+
+    return wrapped_fn
+
+
+def run_api_experiment(
+    input_features: List[Dict[str, Any]],
+    output_features: List[Dict[str, Any]],
+    data_csv: str,
+) -> None:
+    """Run an experiment through Ludwig's Python API.
+
+    Args:
+        input_features: Input schema.
+        output_features: Output schema.
+        data_csv: Path to data.
+    """
+    config = {
+        "input_features": input_features,
+        "output_features": output_features,
+        "combiner": {"type": "concat", "fc_size": 14},
+        "training": {"epochs": 2},
+    }
+
+    model = LudwigModel(config)
+    output_dir = None
+
+    try:
+        # Training with csv
+        _, _, output_dir = model.train(
+            dataset=data_csv,
+            skip_save_processed_input=True,
+            skip_save_progress=True,
+            skip_save_unprocessed_output=True,
+        )
+        model.predict(dataset=data_csv)
+
+        model_dir = os.path.join(output_dir, "model")
+        loaded_model = LudwigModel.load(model_dir)
+
+        # Necessary before call to get_weights() to materialize the weights
+        loaded_model.predict(dataset=data_csv)
+
+        model_weights = model.model.get_weights()
+        loaded_weights = loaded_model.model.get_weights()
+        for model_weight, loaded_weight in zip(model_weights, loaded_weights):
+            assert np.allclose(model_weight, loaded_weight)
+    finally:
+        # Remove results/intermediate data saved to disk
+        shutil.rmtree(output_dir, ignore_errors=True)
+
+    try:
+        # Training with dataframe
+        data_df = read_csv(data_csv)
+        _, _, output_dir = model.train(
+            dataset=data_df,
+            skip_save_processed_input=True,
+            skip_save_progress=True,
+            skip_save_unprocessed_output=True,
+        )
+        model.predict(dataset=data_df)
+    finally:
+        shutil.rmtree(output_dir, ignore_errors=True)
+
+
+def create_data_set_to_use(data_format, raw_data):
+    # helper function for generating training and test data with specified
+    # format handles all data formats except for hdf5
+    # assumes raw_data is a csv dataset generated by
+    # tests.integration_tests.utils.generate_data() function
+
+    # support for writing to a fwf dataset based on this stackoverflow posting:
+    # https://stackoverflow.com/questions/16490261/python-pandas-write-dataframe-to-fixed-width-file-to-fwf
+    from ray._private.thirdparty.tabulate.tabulate import tabulate
+
+    def to_fwf(df, fname):
+        content = tabulate(df.values.tolist(), list(df.columns), tablefmt="plain")
+        open(fname, "w").write(content)
+
+    pd.DataFrame.to_fwf = to_fwf
+
+    dataset_to_use = None
+
+    if data_format == "csv":
+        dataset_to_use = raw_data
+
+    elif data_format in {"df", "dict"}:
+        dataset_to_use = pd.read_csv(raw_data)
+        if data_format == "dict":
+            dataset_to_use = dataset_to_use.to_dict(orient="list")
+
+    elif data_format == "excel":
+        dataset_to_use = replace_file_extension(raw_data, "xlsx")
+        pd.read_csv(raw_data).to_excel(dataset_to_use, index=False)
+
+    elif data_format == "excel_xls":
+        dataset_to_use = replace_file_extension(raw_data, "xls")
+        pd.read_csv(raw_data).to_excel(dataset_to_use, index=False)
+
+    elif data_format == "feather":
+        dataset_to_use = replace_file_extension(raw_data, "feather")
+        pd.read_csv(raw_data).to_feather(dataset_to_use)
+
+    elif data_format == "fwf":
+        dataset_to_use = replace_file_extension(raw_data, "fwf")
+        pd.read_csv(raw_data).to_fwf(dataset_to_use)
+
+    elif data_format == "html":
+        dataset_to_use = replace_file_extension(raw_data, "html")
+        pd.read_csv(raw_data).to_html(dataset_to_use, index=False)
+
+    elif data_format == "json":
+        dataset_to_use = replace_file_extension(raw_data, "json")
+        pd.read_csv(raw_data).to_json(dataset_to_use, orient="records")
+
+    elif data_format == "jsonl":
+        dataset_to_use = replace_file_extension(raw_data, "jsonl")
+        pd.read_csv(raw_data).to_json(dataset_to_use, orient="records", lines=True)
+
+    elif data_format == "parquet":
+        dataset_to_use = replace_file_extension(raw_data, "parquet")
+        pd.read_csv(raw_data).to_parquet(dataset_to_use, index=False)
+
+    elif data_format == "pickle":
+        dataset_to_use = replace_file_extension(raw_data, "pickle")
+        pd.read_csv(raw_data).to_pickle(dataset_to_use)
+
+    elif data_format == "stata":
+        dataset_to_use = replace_file_extension(raw_data, "stata")
+        pd.read_csv(raw_data).to_stata(dataset_to_use)
+
+    elif data_format == "tsv":
+        dataset_to_use = replace_file_extension(raw_data, "tsv")
+        pd.read_csv(raw_data).to_csv(dataset_to_use, sep="\t", index=False)
+
+    else:
+        ValueError("'{}' is an unrecognized data format".format(data_format))
+
+    return dataset_to_use
+
+
+def train_with_backend(
+    backend,
+    config,
+    dataset=None,
+    training_set=None,
+    validation_set=None,
+    test_set=None,
+    predict=True,
+    evaluate=True,
+):
+    model = LudwigModel(config, backend=backend)
+    output_dir = None
+
+    ret = False
+    try:
+        _, _, output_dir = model.train(
+            dataset=dataset,
+            training_set=training_set,
+            validation_set=validation_set,
+            test_set=test_set,
+            skip_save_processed_input=True,
+            skip_save_progress=True,
+            skip_save_unprocessed_output=True,
+        )
+
+        if dataset is None:
+            dataset = training_set
+
+        if predict:
+            preds, _ = model.predict(dataset=dataset)
+            assert backend.df_engine.compute(preds) is not None
+
+        if evaluate:
+            _, eval_preds, _ = model.evaluate(dataset=dataset)
+            assert backend.df_engine.compute(eval_preds) is not None
+
+        ret = True
+    finally:
+        # Remove results/intermediate data saved to disk
+        shutil.rmtree(output_dir, ignore_errors=True)
+    return ret

python/ray/tests/ludwig/test_ludwig.py

@@ -0,0 +1,168 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2019 Uber Technologies, Inc.
+#
+# 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.
+# ==============================================================================
+
+# This file is copied and adapted from
+# https://github.com/ludwig-ai/ludwig/blob/master/tests/integration_tests/test_ray.py
+
+import contextlib
+import os
+import sys
+import tempfile
+
+import pytest
+
+import ray
+
+ludwig_installed = True
+tf_installed = True
+
+try:
+    import ludwig  # noqa: F401
+except (ImportError, ModuleNotFoundError):
+    ludwig_installed = False
+
+try:
+    import tensorflow as tf  # noqa: F401
+except (ImportError, ModuleNotFoundError):
+    tf_installed = False
+
+skip = not ludwig_installed or not tf_installed
+
+# These tests are written for versions of Modin that require python 3.7+
+pytestmark = pytest.mark.skipif(skip, reason="Missing Ludwig dependency")
+
+if not skip:
+    from ludwig.backend.ray import RayBackend, get_horovod_kwargs
+
+    from ray.tests.ludwig.ludwig_test_utils import (
+        bag_feature,
+        binary_feature,
+        category_feature,
+        create_data_set_to_use,
+        date_feature,
+        generate_data,
+        h3_feature,
+        numerical_feature,
+        sequence_feature,
+        set_feature,
+        spawn,
+        train_with_backend,
+        vector_feature,
+    )
+
+else:
+
+    def spawn(func):
+        return func
+
+
+@contextlib.contextmanager
+def ray_start_2_cpus():
+    is_ray_initialized = ray.is_initialized()
+    with tempfile.TemporaryDirectory() as tmpdir:
+        if not is_ray_initialized:
+            res = ray.init(
+                num_cpus=2,
+                include_dashboard=False,
+                object_store_memory=150 * 1024 * 1024,
+                _temp_dir=tmpdir,
+            )
+        else:
+            res = None
+        try:
+            yield res
+        finally:
+            if not is_ray_initialized:
+                ray.shutdown()
+
+
+def run_api_experiment(config, data_parquet):
+    # Sanity check that we get 4 slots over 1 host
+    kwargs = get_horovod_kwargs()
+    assert kwargs.get("num_workers") == 2
+
+    # Train on Parquet
+    dask_backend = RayBackend()
+    assert train_with_backend(
+        dask_backend, config, dataset=data_parquet, evaluate=False
+    )
+
+
+@spawn
+def run_test_parquet(
+    input_features,
+    output_features,
+    num_examples=100,
+    run_fn=run_api_experiment,
+    expect_error=False,
+):
+    tf.config.experimental_run_functions_eagerly(True)
+    with ray_start_2_cpus():
+        config = {
+            "input_features": input_features,
+            "output_features": output_features,
+            "combiner": {"type": "concat", "fc_size": 14},
+            "training": {"epochs": 2, "batch_size": 8},
+        }
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            csv_filename = os.path.join(tmpdir, "dataset.csv")
+            dataset_csv = generate_data(
+                input_features, output_features, csv_filename, num_examples=num_examples
+            )
+            dataset_parquet = create_data_set_to_use("parquet", dataset_csv)
+
+            if expect_error:
+                with pytest.raises(ValueError):
+                    run_fn(config, data_parquet=dataset_parquet)
+            else:
+                run_fn(config, data_parquet=dataset_parquet)
+
+
+def test_ray_tabular():
+    input_features = [
+        sequence_feature(reduce_output="sum"),
+        numerical_feature(normalization="zscore"),
+        set_feature(),
+        binary_feature(),
+        bag_feature(),
+        vector_feature(),
+        h3_feature(),
+        date_feature(),
+    ]
+    output_features = [
+        category_feature(vocab_size=2, reduce_input="sum"),
+        binary_feature(),
+        set_feature(max_len=3, vocab_size=5),
+        numerical_feature(normalization="zscore"),
+        vector_feature(),
+    ]
+    run_test_parquet(input_features, output_features)
+
+
+def test_ray_tabular_client():
+    from ray.util.client.ray_client_helpers import ray_start_client_server
+
+    with ray_start_2_cpus():
+        assert not ray.util.client.ray.is_connected()
+        with ray_start_client_server():
+            assert ray.util.client.ray.is_connected()
+
+            test_ray_tabular()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/tests/mock_s3_server.py

@@ -0,0 +1,93 @@
+# extracted from aioboto3
+#    https://github.com/terrycain/aioboto3/blob/16a1a1085191ebe6d40ee45d9588b2173738af0c/tests/mock_server.py
+import shutil
+import signal
+import subprocess as sp
+import time
+
+import pytest
+import requests
+
+from ray._common.network_utils import build_address
+
+_proxy_bypass = {
+    "http": None,
+    "https": None,
+}
+
+
+def start_service(service_name, host, port):
+    moto_svr_path = shutil.which("moto_server")
+    # moto 5.x no longer accepts a service name argument - all services
+    # are served on a single endpoint
+    args = [moto_svr_path, "-H", host, "-p", str(port)]
+    process = sp.Popen(
+        args, stdin=sp.PIPE, stdout=sp.DEVNULL, stderr=sp.DEVNULL
+    )  # shell=True
+    url = f"http://{build_address(host, port)}"
+
+    for i in range(0, 30):
+        output = process.poll()
+        if output is not None:
+            print("moto_server exited status {0}".format(output))
+            stdout, stderr = process.communicate()
+            print("moto_server stdout: {0}".format(stdout))
+            print("moto_server stderr: {0}".format(stderr))
+            pytest.fail("Can not start service: {}".format(service_name))
+
+        try:
+            # we need to bypass the proxies due to monkeypatches
+            requests.get(url, timeout=5, proxies=_proxy_bypass)
+            break
+        except requests.exceptions.ConnectionError:
+            time.sleep(0.5)
+    else:
+        stop_process(process)  # pytest.fail doesn't call stop_process
+        pytest.fail("Can not start service: {}".format(service_name))
+
+    return process
+
+
+def stop_process(process):
+    try:
+        process.send_signal(signal.SIGTERM)
+        process.communicate(timeout=20)
+    except sp.TimeoutExpired:
+        process.kill()
+        outs, errors = process.communicate(timeout=20)
+        exit_code = process.returncode
+        msg = "Child process finished {} not in clean way: {} {}".format(
+            exit_code, outs, errors
+        )
+        raise RuntimeError(msg)
+
+
+@pytest.fixture(scope="session")
+def dynamodb2_server():
+    host = "localhost"
+    port = 5001
+    url = f"http://{build_address(host, port)}"
+    process = start_service("dynamodb2", host, port)
+    yield url
+    stop_process(process)
+
+
+@pytest.fixture(scope="session")
+def s3_server():
+    host = "localhost"
+    port = 5002
+    url = f"http://{build_address(host, port)}"
+    process = start_service("s3", host, port)
+    yield url
+    stop_process(process)
+
+
+@pytest.fixture(scope="session")
+def kms_server():
+    host = "localhost"
+    port = 5003
+    url = f"http://{build_address(host, port)}"
+    process = start_service("kms", host, port)
+
+    yield url
+    stop_process(process)

python/ray/tests/mock_setup_worker.py

@@ -0,0 +1,46 @@
+import argparse
+import os
+import sys
+
+parser = argparse.ArgumentParser(
+    description=("Set up the environment for a Ray worker and launch the worker.")
+)
+
+parser.add_argument(
+    "--worker-setup-hook",
+    type=str,
+    help="the module path to a Python function to run to set up the "
+    "environment for a worker and launch the worker.",
+)
+parser.add_argument(
+    "--serialized-runtime-env", type=str, help="the serialized parsed runtime env dict"
+)
+
+parser.add_argument(
+    "--serialized-runtime-env-context",
+    type=str,
+    help="the serialized runtime env context",
+)
+
+parser.add_argument(
+    "--allocated-instances-serialized-json",
+    type=str,
+    help="the worker allocated resource",
+)
+
+# The worker is not set up yet, so we can't get session_dir from the worker.
+parser.add_argument(
+    "--session-dir", type=str, help="the directory for the current session"
+)
+
+parser.add_argument("--language", type=str, help="the language type of the worker")
+
+args, remaining_args = parser.parse_known_args()
+
+py_executable: str = sys.executable
+command_str = " ".join([f"exec {py_executable}"] + remaining_args)
+child_pid = os.fork()
+if child_pid == 0:
+    # child process
+    os.execvp("bash", ["bash", "-c", command_str])
+os.waitpid(child_pid, 0)

python/ray/tests/modin/BUILD.bazel

@@ -0,0 +1,15 @@
+load("@rules_python//python:defs.bzl", "py_test")
+
+py_test(
+    name = "test_modin",
+    size = "small",
+    srcs = ["test_modin.py"],
+    tags = [
+        "exclusive",
+        "team:core",
+    ],
+    deps = [
+        "//:ray_lib",
+        "//python/ray/tests:conftest",
+    ],
+)

python/ray/tests/modin/modin_test_utils.py

@@ -0,0 +1,151 @@
+# Licensed to Modin Development Team under one or more contributor license
+# agreements. See the NOTICE file distributed with this work for additional
+# information regarding copyright ownership.  The Modin Development Team
+# licenses this file to you 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.
+#
+# This file is copied and adapted from
+# http://github.com/modin-project/modin/master/modin/pandas/test/utils.py
+
+from typing import Any
+
+import modin.pandas as pd
+import numpy as np
+import pandas
+
+# to_pandas moved from modin.utils to modin.pandas.io in modin 0.26.0,
+from modin.pandas.io import to_pandas
+from pandas.testing import (
+    assert_extension_array_equal,
+    assert_frame_equal,
+    assert_index_equal,
+    assert_series_equal,
+)
+
+
+def categories_equals(left, right):
+    assert (left.ordered and right.ordered) or (not left.ordered and not right.ordered)
+    assert_extension_array_equal(left, right)
+
+
+def df_categories_equals(df1, df2):
+    if not hasattr(df1, "select_dtypes"):
+        if isinstance(df1, pandas.CategoricalDtype):
+            return categories_equals(df1, df2)
+        elif isinstance(df1.dtype, pandas.CategoricalDtype) and isinstance(
+            df1.dtype, pandas.CategoricalDtype
+        ):
+            return categories_equals(df1.dtype, df2.dtype)
+        else:
+            return True
+
+    categories_columns = df1.select_dtypes(include="category").columns
+    for column in categories_columns:
+        assert_extension_array_equal(
+            df1[column].values,
+            df2[column].values,
+            check_dtype=False,
+        )
+
+
+def df_equals(df1: Any, df2: Any) -> bool:
+    """Tests if df1 and df2 are equal.
+
+    Args:
+        df1: (pandas or modin DataFrame or series) dataframe to test if equal.
+        df2: (pandas or modin DataFrame or series) dataframe to test if equal.
+
+    Returns:
+        True if df1 is equal to df2.
+    """
+    # Gets AttributError if modin's groupby object is not import like this
+    from modin.pandas.groupby import DataFrameGroupBy
+
+    groupby_types = (pandas.core.groupby.DataFrameGroupBy, DataFrameGroupBy)
+
+    # The typing behavior of how pandas treats its index is not consistent when
+    # the length of the DataFrame or Series is 0, so we just verify that the
+    # contents are the same.
+    if (
+        hasattr(df1, "index")
+        and hasattr(df2, "index")
+        and len(df1) == 0
+        and len(df2) == 0
+    ):
+        if type(df1).__name__ == type(df2).__name__:
+            if hasattr(df1, "name") and hasattr(df2, "name") and df1.name == df2.name:
+                return
+            if (
+                hasattr(df1, "columns")
+                and hasattr(df2, "columns")
+                and df1.columns.equals(df2.columns)
+            ):
+                return
+        assert False
+
+    if isinstance(df1, (list, tuple)) and all(
+        isinstance(d, (pd.DataFrame, pd.Series, pandas.DataFrame, pandas.Series))
+        for d in df1
+    ):
+        assert isinstance(df2, type(df1)), "Different type of collection"
+        assert len(df1) == len(df2), "Different length result"
+        return (df_equals(d1, d2) for d1, d2 in zip(df1, df2))
+
+    # Convert to pandas
+    if isinstance(df1, (pd.DataFrame, pd.Series)):
+        df1 = to_pandas(df1)
+    if isinstance(df2, (pd.DataFrame, pd.Series)):
+        df2 = to_pandas(df2)
+
+    if isinstance(df1, pandas.DataFrame) and isinstance(df2, pandas.DataFrame):
+        if (df1.empty and not df2.empty) or (df2.empty and not df1.empty):
+            assert False, "One of the passed frames is empty, when other isn't"
+        elif df1.empty and df2.empty and type(df1) is not type(df2):
+            assert (
+                False
+            ), f"Empty frames have different types: {type(df1)} != {type(df2)}"
+
+    if isinstance(df1, pandas.DataFrame) and isinstance(df2, pandas.DataFrame):
+        assert_frame_equal(
+            df1,
+            df2,
+            check_dtype=False,
+            check_datetimelike_compat=True,
+            check_index_type=False,
+            check_column_type=False,
+            check_categorical=False,
+        )
+        df_categories_equals(df1, df2)
+    elif isinstance(df1, pandas.Index) and isinstance(df2, pandas.Index):
+        assert_index_equal(df1, df2)
+    elif isinstance(df1, pandas.Series) and isinstance(df2, pandas.Series):
+        assert_series_equal(df1, df2, check_dtype=False, check_series_type=False)
+    elif isinstance(df1, groupby_types) and isinstance(df2, groupby_types):
+        for g1, g2 in zip(df1, df2):
+            assert g1[0] == g2[0]
+            df_equals(g1[1], g2[1])
+    elif (
+        isinstance(df1, pandas.Series)
+        and isinstance(df2, pandas.Series)
+        and df1.empty
+        and df2.empty
+    ):
+        assert all(df1.index == df2.index)
+        assert df1.dtypes == df2.dtypes
+    elif isinstance(df1, pandas.core.arrays.numpy_.PandasArray):
+        assert isinstance(df2, pandas.core.arrays.numpy_.PandasArray)
+        assert df1 == df2
+    elif isinstance(df1, np.recarray) and isinstance(df2, np.recarray):
+        np.testing.assert_array_equal(df1, df2)
+    else:
+        if df1 != df2:
+            np.testing.assert_almost_equal(df1, df2)

python/ray/tests/modin/test_modin.py

@@ -0,0 +1,418 @@
+# Licensed to Modin Development Team under one or more contributor license
+# agreements. See the NOTICE file distributed with this work for additional
+# information regarding copyright ownership.  The Modin Development Team
+# licenses this file to you 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.
+#
+# This file is copied and adapted from:
+# http://github.com/modin-project/modin/master/modin/pandas/test/test_general.py
+
+import sys
+
+import numpy as np
+import pandas
+import pytest
+from numpy.testing import assert_array_equal
+
+from ray.tests.conftest import ray_start_regular_shared  # noqa F401
+
+modin_installed = True
+
+try:
+    import modin  # noqa: F401
+except ModuleNotFoundError:
+    modin_installed = False
+
+skip = not modin_installed
+
+# These tests are written for versions of Modin that require python 3.8+
+pytestmark = pytest.mark.skipif(skip, reason="Outdated or missing Modin dependency")
+
+if not skip:
+    import modin.pandas as pd
+
+    from ray.tests.modin.modin_test_utils import df_equals
+
+
+@pytest.fixture(autouse=True)
+def connect_to_ray_cluster(ray_start_regular_shared):  # noqa F811
+    yield
+
+
+random_state = np.random.RandomState(seed=42)
+
+# Size of test dataframes
+NCOLS, NROWS = (2**6, 2**8)
+
+# Range for values for test data
+RAND_LOW = 0
+RAND_HIGH = 100
+
+# Input data and functions for the tests
+# The test data that we will test our code against
+test_data = {
+    "int_data": {
+        "col{}".format(int((i - NCOLS / 2) % NCOLS + 1)): random_state.randint(
+            RAND_LOW, RAND_HIGH, size=(NROWS)
+        )
+        for i in range(NCOLS)
+    },
+    "float_nan_data": {
+        "col{}".format(int((i - NCOLS / 2) % NCOLS + 1)): [
+            x
+            if (j % 4 == 0 and i > NCOLS // 2) or (j != i and i <= NCOLS // 2)
+            else np.nan
+            for j, x in enumerate(
+                random_state.uniform(RAND_LOW, RAND_HIGH, size=(NROWS))
+            )
+        ]
+        for i in range(NCOLS)
+    },
+}
+
+test_data["int_data"]["index"] = test_data["int_data"].pop(
+    "col{}".format(int(NCOLS / 2))
+)
+
+for col in test_data["float_nan_data"]:
+    for row in range(NROWS // 2):
+        if row % 16 == 0:
+            test_data["float_nan_data"][col][row] = np.nan
+
+test_data_values = list(test_data.values())
+test_data_keys = list(test_data.keys())
+
+
+@pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
+def test_isna(data):
+    pandas_df = pandas.DataFrame(data)
+    modin_df = pd.DataFrame(data)
+
+    pandas_result = pandas.isna(pandas_df)
+    modin_result = pd.isna(modin_df)
+    df_equals(modin_result, pandas_result)
+
+    modin_result = pd.isna(pd.Series([1, np.nan, 2]))
+    pandas_result = pandas.isna(pandas.Series([1, np.nan, 2]))
+    df_equals(modin_result, pandas_result)
+
+    assert pd.isna(np.nan) == pandas.isna(np.nan)
+
+
+@pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
+def test_isnull(data):
+    pandas_df = pandas.DataFrame(data)
+    modin_df = pd.DataFrame(data)
+
+    pandas_result = pandas.isnull(pandas_df)
+    modin_result = pd.isnull(modin_df)
+    df_equals(modin_result, pandas_result)
+
+    modin_result = pd.isnull(pd.Series([1, np.nan, 2]))
+    pandas_result = pandas.isnull(pandas.Series([1, np.nan, 2]))
+    df_equals(modin_result, pandas_result)
+
+    assert pd.isna(np.nan) == pandas.isna(np.nan)
+
+
+@pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
+def test_notna(data):
+    pandas_df = pandas.DataFrame(data)
+    modin_df = pd.DataFrame(data)
+
+    pandas_result = pandas.notna(pandas_df)
+    modin_result = pd.notna(modin_df)
+    df_equals(modin_result, pandas_result)
+
+    modin_result = pd.notna(pd.Series([1, np.nan, 2]))
+    pandas_result = pandas.notna(pandas.Series([1, np.nan, 2]))
+    df_equals(modin_result, pandas_result)
+
+    assert pd.isna(np.nan) == pandas.isna(np.nan)
+
+
+@pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
+def test_notnull(data):
+    pandas_df = pandas.DataFrame(data)
+    modin_df = pd.DataFrame(data)
+
+    pandas_result = pandas.notnull(pandas_df)
+    modin_result = pd.notnull(modin_df)
+    df_equals(modin_result, pandas_result)
+
+    modin_result = pd.notnull(pd.Series([1, np.nan, 2]))
+    pandas_result = pandas.notnull(pandas.Series([1, np.nan, 2]))
+    df_equals(modin_result, pandas_result)
+
+    assert pd.isna(np.nan) == pandas.isna(np.nan)
+
+
+def test_merge():
+    frame_data = {
+        "col1": [0, 1, 2, 3],
+        "col2": [4, 5, 6, 7],
+        "col3": [8, 9, 0, 1],
+        "col4": [2, 4, 5, 6],
+    }
+
+    modin_df = pd.DataFrame(frame_data)
+    pandas_df = pandas.DataFrame(frame_data)
+
+    frame_data2 = {"col1": [0, 1, 2], "col2": [1, 5, 6]}
+    modin_df2 = pd.DataFrame(frame_data2)
+    pandas_df2 = pandas.DataFrame(frame_data2)
+
+    join_types = ["outer", "inner"]
+    for how in join_types:
+        # Defaults
+        modin_result = pd.merge(modin_df, modin_df2, how=how)
+        pandas_result = pandas.merge(pandas_df, pandas_df2, how=how)
+        df_equals(modin_result, pandas_result)
+
+        # left_on and right_index
+        modin_result = pd.merge(
+            modin_df, modin_df2, how=how, left_on="col1", right_index=True
+        )
+        pandas_result = pandas.merge(
+            pandas_df, pandas_df2, how=how, left_on="col1", right_index=True
+        )
+        df_equals(modin_result, pandas_result)
+
+        # left_index and right_on
+        modin_result = pd.merge(
+            modin_df, modin_df2, how=how, left_index=True, right_on="col1"
+        )
+        pandas_result = pandas.merge(
+            pandas_df, pandas_df2, how=how, left_index=True, right_on="col1"
+        )
+        df_equals(modin_result, pandas_result)
+
+        # left_on and right_on col1
+        modin_result = pd.merge(
+            modin_df, modin_df2, how=how, left_on="col1", right_on="col1"
+        )
+        pandas_result = pandas.merge(
+            pandas_df, pandas_df2, how=how, left_on="col1", right_on="col1"
+        )
+        df_equals(modin_result, pandas_result)
+
+        # left_on and right_on col2
+        modin_result = pd.merge(
+            modin_df, modin_df2, how=how, left_on="col2", right_on="col2"
+        )
+        pandas_result = pandas.merge(
+            pandas_df, pandas_df2, how=how, left_on="col2", right_on="col2"
+        )
+        df_equals(modin_result, pandas_result)
+
+        # left_index and right_index
+        modin_result = pd.merge(
+            modin_df, modin_df2, how=how, left_index=True, right_index=True
+        )
+        pandas_result = pandas.merge(
+            pandas_df, pandas_df2, how=how, left_index=True, right_index=True
+        )
+        df_equals(modin_result, pandas_result)
+
+    s = pd.Series(frame_data.get("col1"))
+    with pytest.raises(ValueError):
+        pd.merge(s, modin_df2)
+
+    with pytest.raises(TypeError):
+        pd.merge("Non-valid type", modin_df2)
+
+
+def test_pivot():
+    test_df = pd.DataFrame(
+        {
+            "foo": ["one", "one", "one", "two", "two", "two"],
+            "bar": ["A", "B", "C", "A", "B", "C"],
+            "baz": [1, 2, 3, 4, 5, 6],
+            "zoo": ["x", "y", "z", "q", "w", "t"],
+        }
+    )
+
+    df = pd.pivot(test_df, index="foo", columns="bar", values="baz")
+    assert isinstance(df, pd.DataFrame)
+
+    with pytest.raises(ValueError):
+        pd.pivot(test_df["bar"], index="foo", columns="bar", values="baz")
+
+
+def test_pivot_table():
+    test_df = pd.DataFrame(
+        {
+            "A": ["foo", "foo", "foo", "foo", "foo", "bar", "bar", "bar", "bar"],
+            "B": ["one", "one", "one", "two", "two", "one", "one", "two", "two"],
+            "C": [
+                "small",
+                "large",
+                "large",
+                "small",
+                "small",
+                "large",
+                "small",
+                "small",
+                "large",
+            ],
+            "D": [1, 2, 2, 3, 3, 4, 5, 6, 7],
+            "E": [2, 4, 5, 5, 6, 6, 8, 9, 9],
+        }
+    )
+
+    df = pd.pivot_table(
+        test_df, values="D", index=["A", "B"], columns=["C"], aggfunc=np.sum
+    )
+    assert isinstance(df, pd.DataFrame)
+
+    with pytest.raises(ValueError):
+        pd.pivot_table(
+            test_df["C"], values="D", index=["A", "B"], columns=["C"], aggfunc=np.sum
+        )
+
+
+def test_unique():
+    modin_result = pd.unique([2, 1, 3, 3])
+    pandas_result = pandas.unique([2, 1, 3, 3])
+    assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
+
+    modin_result = pd.unique(pd.Series([2] + [1] * 5))
+    pandas_result = pandas.unique(pandas.Series([2] + [1] * 5))
+    assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
+
+    modin_result = pd.unique(
+        pd.Series([pd.Timestamp("20160101"), pd.Timestamp("20160101")])
+    )
+    pandas_result = pandas.unique(
+        pandas.Series([pandas.Timestamp("20160101"), pandas.Timestamp("20160101")])
+    )
+    assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
+
+    modin_result = pd.unique(
+        pd.Series(
+            [
+                pd.Timestamp("20160101", tz="US/Eastern"),
+                pd.Timestamp("20160101", tz="US/Eastern"),
+            ]
+        )
+    )
+    pandas_result = pandas.unique(
+        pandas.Series(
+            [
+                pandas.Timestamp("20160101", tz="US/Eastern"),
+                pandas.Timestamp("20160101", tz="US/Eastern"),
+            ]
+        )
+    )
+    assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
+
+    modin_result = pd.unique(
+        pd.Index(
+            [
+                pd.Timestamp("20160101", tz="US/Eastern"),
+                pd.Timestamp("20160101", tz="US/Eastern"),
+            ]
+        )
+    )
+    pandas_result = pandas.unique(
+        pandas.Index(
+            [
+                pandas.Timestamp("20160101", tz="US/Eastern"),
+                pandas.Timestamp("20160101", tz="US/Eastern"),
+            ]
+        )
+    )
+    assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
+
+    modin_result = pd.unique(pd.Series(pd.Categorical(list("baabc"))))
+    pandas_result = pandas.unique(pandas.Series(pandas.Categorical(list("baabc"))))
+    assert_array_equal(modin_result, pandas_result)
+    assert modin_result.shape == pandas_result.shape
+
+
+def test_to_datetime():
+    # DataFrame input for to_datetime
+    modin_df = pd.DataFrame({"year": [2015, 2016], "month": [2, 3], "day": [4, 5]})
+    pandas_df = pandas.DataFrame({"year": [2015, 2016], "month": [2, 3], "day": [4, 5]})
+    df_equals(pd.to_datetime(modin_df), pandas.to_datetime(pandas_df))
+
+    # Series input for to_datetime
+    modin_s = pd.Series(["3/11/2000", "3/12/2000", "3/13/2000"] * 1000)
+    pandas_s = pandas.Series(["3/11/2000", "3/12/2000", "3/13/2000"] * 1000)
+    df_equals(pd.to_datetime(modin_s), pandas.to_datetime(pandas_s))
+
+    # Other inputs for to_datetime
+    value = 1490195805
+    assert pd.to_datetime(value, unit="s") == pandas.to_datetime(value, unit="s")
+    value = 1490195805433502912
+    assert pd.to_datetime(value, unit="ns") == pandas.to_datetime(value, unit="ns")
+    value = [1, 2, 3]
+    assert pd.to_datetime(value, unit="D", origin=pd.Timestamp("2000-01-01")).equals(
+        pandas.to_datetime(value, unit="D", origin=pandas.Timestamp("2000-01-01"))
+    )
+
+
+@pytest.mark.parametrize(
+    "data, errors, downcast",
+    [
+        (["1.0", "2", -3], "raise", None),
+        (["1.0", "2", -3], "raise", "float"),
+        (["1.0", "2", -3], "raise", "signed"),
+        (["apple", "1.0", "2", -3], "ignore", None),
+        (["apple", "1.0", "2", -3], "coerce", None),
+    ],
+)
+def test_to_numeric(data, errors, downcast):
+    modin_series = pd.Series(data)
+    pandas_series = pandas.Series(data)
+    modin_result = pd.to_numeric(modin_series, errors=errors, downcast=downcast)
+    pandas_result = pandas.to_numeric(pandas_series, errors=errors, downcast=downcast)
+    df_equals(modin_result, pandas_result)
+
+
+def test_to_pandas_indices():
+    data = test_data_values[0]
+
+    md_df = pd.DataFrame(data)
+    index = pandas.MultiIndex.from_tuples(
+        [(i, i * 2) for i in np.arange(len(md_df) + 1)], names=["A", "B"]
+    ).drop(0)
+    columns = pandas.MultiIndex.from_tuples(
+        [(i, i * 2) for i in np.arange(len(md_df.columns) + 1)], names=["A", "B"]
+    ).drop(0)
+
+    md_df.index = index
+    md_df.columns = columns
+
+    pd_df = md_df._to_pandas()
+
+    for axis in [0, 1]:
+        assert md_df.axes[axis].equals(
+            pd_df.axes[axis]
+        ), f"Indices at axis {axis} are different!"
+        assert md_df.axes[axis].equal_levels(
+            pd_df.axes[axis]
+        ), f"Levels of indices at axis {axis} are different!"
+
+
+def test_empty_dataframe():
+    df = pd.DataFrame(columns=["a", "b"])
+    df[(df.a == 1) & (df.b == 2)]
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/tests/nsight_fake/nsys_fake.py

@@ -0,0 +1,30 @@
+import argparse
+import random
+import subprocess
+
+
+def run():
+    parser = argparse.ArgumentParser(description="Your Script Description")
+    parser.add_argument("command", help="The command to profile")
+    parser.add_argument("-t", default="", help="Cuda APIs to be profiled")
+    parser.add_argument("--stop-on-exit", default="true", help="profile finish on exit")
+    parser.add_argument("--cudabacktrace", default="none", help="profile cuda apis")
+    parser.add_argument("-o", default="report%p", help="output filename")
+    parser.add_argument("-d", default="0", help="profiling duration")
+    parser.add_argument("script", help="Python script to profile")
+    parser.add_argument(
+        "script_args", nargs=argparse.REMAINDER, help="Arguments for the Python script"
+    )
+
+    args = parser.parse_args()
+    random_id = random.randint(1, 9999)
+
+    if args.command != "profile":
+        print("nothing")
+    else:
+        output_filepath = args.o.replace("%p", str(random_id)) + ".nsys-rep"
+        with open(output_filepath, "w") as file:
+            file.write("Mock file.\n")
+
+    command = [args.script] + args.script_args
+    subprocess.run(command)

python/ray/tests/nsight_fake/setup.py

@@ -0,0 +1,7 @@
+from setuptools import setup
+
+setup(
+    name="nsys",
+    version="0.0.1",
+    entry_points={"console_scripts": ["nsys=nsys_fake:run"]},
+)

python/ray/tests/perf_integration_tests/test_perf_integration.py

@@ -0,0 +1,104 @@
+import numpy as np
+import pytest
+
+import ray
+from ray.tests.conftest import _ray_start_cluster
+
+num_tasks_submitted = [10**n for n in range(0, 6)]
+num_tasks_ids = ["{}_tasks".format(i) for i in num_tasks_submitted]
+
+
+@ray.remote
+def dummy_task(val):
+    return val
+
+
+def benchmark_task_submission(num_tasks):
+    total_tasks = 100000
+    for _ in range(total_tasks // num_tasks):
+        ray.get([dummy_task.remote(i) for i in range(num_tasks)])
+
+
+def warmup():
+    x = np.zeros(10**6, dtype=np.uint8)
+    for _ in range(5):
+        for _ in range(5):
+            ray.put(x)
+        for _ in range(5):
+            ray.get([dummy_task.remote(0) for _ in range(1000)])
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize("num_tasks", num_tasks_submitted, ids=num_tasks_ids)
+def test_task_submission(benchmark, num_tasks):
+    num_cpus = 16
+    ray.init(
+        num_cpus=num_cpus,
+        object_store_memory=150 * 1024 * 1024,
+        ignore_reinit_error=True,
+    )
+    # warm up the plasma store
+    warmup()
+    benchmark(benchmark_task_submission, num_tasks)
+    ray.shutdown()
+
+
+def benchmark_task_forward(f, num_tasks):
+    ray.get([f.remote() for _ in range(num_tasks)])
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize(
+    "num_tasks",
+    [10**3, 10**4],
+    ids=[str(num) + "_tasks" for num in [10**3, 10**4]],
+)
+def test_task_forward(benchmark, num_tasks):
+    with _ray_start_cluster(
+        do_init=True,
+        num_nodes=1,
+        num_cpus=16,
+        object_store_memory=150 * 1024 * 1024,
+    ) as cluster:
+        cluster.add_node(
+            num_cpus=16,
+            object_store_memory=150 * 1024 * 1024,
+            resources={"my_resource": 100},
+        )
+
+        @ray.remote(resources={"my_resource": 0.001})
+        def f():
+            return 1
+
+        # Warm up
+        ray.get([f.remote() for _ in range(100)])
+        benchmark(benchmark_task_forward, f, num_tasks)
+
+
+def benchmark_transfer_object(actor, object_refs):
+    ray.get(actor.f.remote(object_refs))
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize(
+    "object_number, data_size", [(10000, 500), (10000, 5000), (1000, 500), (1000, 5000)]
+)
+def test_transfer_performance(
+    benchmark, ray_start_cluster_head, object_number, data_size
+):
+    cluster = ray_start_cluster_head
+    cluster.add_node(resources={"my_resource": 1}, object_store_memory=10**9)
+
+    @ray.remote(resources={"my_resource": 1})
+    class ObjectActor:
+        def f(self, object_refs):
+            ray.get(object_refs)
+
+    # setup remote actor
+    actor = ObjectActor.remote()
+    actor.f.remote([])
+
+    data = bytes(1) * data_size
+    object_refs = [ray.put(data) for _ in range(object_number)]
+
+    benchmark(benchmark_transfer_object, actor, object_refs)

python/ray/tests/project_files/docker_project/ray-project/project.yaml

@@ -0,0 +1,8 @@
+name: testproject1
+description: "Test project for docker environment"
+
+environment:
+  dockerfile: "Dockerfile"
+
+cluster:
+  config: cluster.yaml

python/ray/tests/project_files/no_project1/ray-project/project.yaml

@@ -0,0 +1,4 @@
+name: testmissingcluster
+
+environment:
+  shell: "one command"

python/ray/tests/project_files/no_project2/.rayproject/project.yaml

@@ -0,0 +1,4 @@
+name: testmissingyaml
+
+cluster:
+  config: cluster.yaml

python/ray/tests/project_files/no_project3/.rayproject/project.yaml

@@ -0,0 +1,9 @@
+name: testproject3
+
+environment:
+  dockerfile: "Dockerfile"
+
+  dockerimage: "some docker image"
+
+cluster:
+  config: cluster.yaml

python/ray/tests/project_files/project1/ray-project/project.yaml

@@ -0,0 +1,11 @@
+name: "project1"
+
+cluster:
+  config: ray-project/cluster.yaml
+
+environment:
+  requirements: requirements.txt
+
+commands:
+  - name: default
+    command: ls