chore: import upstream snapshot with attribution

python/ray/tests/test_autoscaler_drain_node_api.py

@@ -0,0 +1,118 @@
+import logging
+import platform
+import sys
+import time
+
+import pytest
+
+import ray
+import ray._private.ray_constants as ray_constants
+from ray._common.test_utils import wait_for_condition
+from ray._private.test_utils import (
+    get_error_message,
+    init_error_pubsub,
+)
+from ray.autoscaler._private.fake_multi_node.node_provider import FakeMultiNodeProvider
+from ray.cluster_utils import AutoscalingCluster
+
+logger = logging.getLogger(__name__)
+
+
+class MockFakeProvider(FakeMultiNodeProvider):
+    """FakeMultiNodeProvider, with Ray node process termination mocked out.
+
+    Used to check that a Ray node can be terminated by DrainNode API call
+    from the autoscaler.
+    """
+
+    def _kill_ray_processes(self, node):
+        logger.info("Leaving Raylet termination to autoscaler Drain API!")
+
+
+class MockAutoscalingCluster(AutoscalingCluster):
+    """AutoscalingCluster modified to used the above MockFakeProvider."""
+
+    def _generate_config(
+        self, head_resources, worker_node_types, autoscaler_v2: bool = False
+    ):
+        config = super()._generate_config(
+            head_resources, worker_node_types, autoscaler_v2=autoscaler_v2
+        )
+        config["provider"]["type"] = "external"
+        config["provider"][
+            "module"
+        ] = "ray.tests.test_autoscaler_drain_node_api.MockFakeProvider"
+        return config
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Failing on Windows.")
+@pytest.mark.parametrize("autoscaler_v2", [False, True], ids=["v1", "v2"])
+def test_drain_api(autoscaler_v2, shutdown_only):
+    """E2E test of the autoscaler's use of the DrainNode API.
+
+    Adapted from test_autoscaler_fake_multinode.py.
+
+    The strategy is to mock out Ray node process termination in
+    FakeMultiNodeProvider, leaving node termination to the DrainNode API.
+
+    Scale-down is verified by `ray.cluster_resources`. It is verified that
+    no removed_node errors are issued adter scale-down.
+
+    Validity of this test depends on the current implementation of DrainNode.
+    DrainNode currently works by asking the GCS to de-register and shut down
+    Ray nodes.
+    """
+    # Autoscaling cluster with Ray process termination mocked out in the node
+    # provider.
+    cluster = MockAutoscalingCluster(
+        head_resources={"CPU": 1},
+        worker_node_types={
+            "gpu_node": {
+                "resources": {
+                    "CPU": 1,
+                    "GPU": 1,
+                    "object_store_memory": 1024 * 1024 * 1024,
+                },
+                "node_config": {},
+                "min_workers": 0,
+                "max_workers": 2,
+            },
+        },
+        autoscaler_v2=autoscaler_v2,
+    )
+
+    try:
+        cluster.start()
+        ray.init("auto")
+
+        # Triggers the addition of a GPU node.
+        @ray.remote(num_gpus=1)
+        def f():
+            print("gpu ok")
+
+        ray.get(f.remote())
+
+        # Verify scale-up
+        wait_for_condition(lambda: ray.cluster_resources().get("GPU", 0) == 1)
+        # Sleep for double the idle timeout of 6 seconds.
+        time.sleep(12)
+
+        # Verify scale-down
+        wait_for_condition(lambda: ray.cluster_resources().get("GPU", 0) == 0)
+
+        # Check that no errors were raised while draining nodes.
+        # (Logic copied from test_failure4::test_gcs_drain.)
+        try:
+            p = init_error_pubsub()
+            errors = get_error_message(
+                p, 1, ray_constants.REMOVED_NODE_ERROR, timeout=5
+            )
+            assert len(errors) == 0
+        finally:
+            p.close()
+    finally:
+        cluster.shutdown()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))