chore: import upstream snapshot with attribution

python/ray/tests/test_batch_node_provider_integration.py

@@ -0,0 +1,178 @@
+"""Integration/e2e test for BatchingNodeProvider.
+Adapts FakeMultiNodeProvider tests.
+"""
+import logging
+import sys
+from copy import deepcopy
+
+import pytest
+
+import ray
+from ray._common.test_utils import wait_for_condition
+from ray.autoscaler._private.constants import FOREGROUND_NODE_LAUNCH_KEY
+from ray.autoscaler._private.fake_multi_node.node_provider import FakeMultiNodeProvider
+from ray.autoscaler.batching_node_provider import (
+    BatchingNodeProvider,
+    NodeData,
+    ScaleRequest,
+)
+from ray.autoscaler.tags import (
+    NODE_KIND_WORKER,
+    STATUS_UP_TO_DATE,
+    TAG_RAY_NODE_KIND,
+    TAG_RAY_NODE_STATUS,
+    TAG_RAY_USER_NODE_TYPE,
+)
+from ray.cluster_utils import AutoscalingCluster
+
+logger = logging.getLogger(__name__)
+
+
+class FakeBatchingNodeProvider(BatchingNodeProvider):
+    """Class for e2e local testing of BatchingNodeProvider.
+    Uses FakeMultiNodeProvider as a proxy for managing the nodes.
+
+    This node provider requires the "available_node_types" section of the
+    autoscaling config to be copied into the "provider" section.
+    That's needed so that node resources can be accessed as
+    provider_config["available_node_types"][node_type]["resources"].
+
+    See the create_node_with_resources call in submit_scale_request.
+    See class BatchingAutoscaler below.
+    """
+
+    def __init__(self, provider_config, cluster_name):
+        BatchingNodeProvider.__init__(self, provider_config, cluster_name)
+        self.fake_multi_node_provider = FakeMultiNodeProvider(
+            provider_config, cluster_name
+        )
+
+    # Manually "inherit" internal utility functions.
+    # I prefer this over attempting multiple inheritance.
+    def _next_hex_node_id(self):
+        return self.fake_multi_node_provider._next_hex_node_id()
+
+    def _terminate_node(self, node):
+        return self.fake_multi_node_provider._terminate_node(node)
+
+    def get_node_data(self):
+        node_data_dict = {}
+        for node_id in self.fake_multi_node_provider._nodes:
+            tags = self.fake_multi_node_provider._nodes[node_id]["tags"]
+            node_data_dict[node_id] = NodeData(
+                kind=tags[TAG_RAY_NODE_KIND],
+                type=tags[TAG_RAY_USER_NODE_TYPE],
+                status=tags[TAG_RAY_NODE_STATUS],
+                ip=node_id,
+            )
+        return node_data_dict
+
+    def submit_scale_request(self, scale_request: ScaleRequest):
+        worker_counts = self.cur_num_workers()
+        for worker_to_delete in scale_request.workers_to_delete:
+            node_type = self.node_tags(worker_to_delete)[TAG_RAY_USER_NODE_TYPE]
+            FakeMultiNodeProvider.terminate_node(self, worker_to_delete)
+            worker_counts[node_type] -= 1
+        for node_type in scale_request.desired_num_workers:
+            diff = (
+                scale_request.desired_num_workers[node_type] - worker_counts[node_type]
+            )
+            # It is non-standard for "available_node_types" to be included in the
+            # provider config, but it is necessary for this node provider.
+            resources = self.provider_config["available_node_types"][node_type][
+                "resources"
+            ]
+            labels = self.provider_config["available_node_types"][node_type].get(
+                "labels", {}
+            )
+            tags = {
+                TAG_RAY_NODE_KIND: NODE_KIND_WORKER,
+                TAG_RAY_USER_NODE_TYPE: node_type,
+                TAG_RAY_NODE_STATUS: STATUS_UP_TO_DATE,
+            }
+            self.fake_multi_node_provider.create_node_with_resources_and_labels(
+                node_config={},
+                tags=tags,
+                count=diff,
+                resources=resources,
+                labels=labels,
+            )
+
+
+class BatchingAutoscalingCluster(AutoscalingCluster):
+    """Class used for e2e testing of BatchingNodeProvider.
+    Like AutoscalingCluster but uses a BatchingNodePorvider.
+    """
+
+    def _generate_config(self, head_resources, worker_node_types, **config_kwargs):
+        config = AutoscalingCluster._generate_config(
+            self, head_resources, worker_node_types
+        )
+        # Need this for resource data.
+        config["provider"]["available_node_types"] = deepcopy(
+            config["available_node_types"]
+        )
+        # Load the node provider class above.
+        config["provider"]["type"] = "external"
+        config["provider"][
+            "module"
+        ] = "ray.tests.test_batch_node_provider_integration.FakeBatchingNodeProvider"
+        # Need to run in single threaded mode to use BatchingNodeProvider.
+        config["provider"][FOREGROUND_NODE_LAUNCH_KEY] = True
+        return config
+
+
+@pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.")
+def test_fake_batching_autoscaler_e2e(shutdown_only):
+    cluster = BatchingAutoscalingCluster(
+        head_resources={"CPU": 2},
+        worker_node_types={
+            "cpu_node": {
+                "resources": {
+                    "CPU": 4,
+                    "object_store_memory": 1024 * 1024 * 1024,
+                },
+                "node_config": {},
+                "min_workers": 0,
+                "max_workers": 2,
+            },
+            "gpu_node": {
+                "resources": {
+                    "CPU": 2,
+                    "GPU": 1,
+                    "object_store_memory": 1024 * 1024 * 1024,
+                },
+                "node_config": {},
+                "min_workers": 0,
+                "max_workers": 2,
+            },
+        },
+    )
+
+    try:
+        cluster.start()
+        ray.init("auto")
+
+        # Triggers the addition of a GPU node.
+        @ray.remote(num_gpus=1)
+        def f():
+            print("gpu ok")
+
+        # Triggers the addition of a CPU node.
+        @ray.remote(num_cpus=3)
+        def g():
+            print("cpu ok")
+
+        ray.get(f.remote())
+        ray.get(g.remote())
+        # Wait for scale-down.
+        wait_for_condition(
+            lambda: ray.cluster_resources().get("CPU", 0) == 2, timeout=30
+        )
+        ray.shutdown()
+    finally:
+        cluster.shutdown()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))