[Data] Add autoscaler_scales_up_when_compute_bound release test

release/nightly_tests/dataset/autoscaling/autoscaler_scales_up_when_compute_bound.py

@@ -0,0 +1,115 @@
+import argparse
+import math
+import time
+from typing import Any, Dict, Iterator
+
+import numpy as np
+import pyarrow as pa
+
+import ray
+
+from benchmark import Benchmark
+from cluster_resource_monitor import ClusterResourceMonitor
+
+# With 1000 inputs this takes ~30 minutes, long enough that node provisioning
+# speed doesn't make the test flaky.
+NUM_INPUTS = 1000
+BLOCKS_PER_INPUT = 4
+PRODUCE_SLEEP_S = 5
+CONSUME_SLEEP_S = 1
+BLOCK_SHAPE = (128, 1024, 1024)
+ROWS_PER_BLOCK = BLOCK_SHAPE[0]
+CONSUME_BATCH_SIZE = 2 * ROWS_PER_BLOCK
+
+# From the compute config.
+MAX_GPU_NODES = 10
+CPUS_PER_NODE = 8
+
+EXPECTED_GPU_NODES = MAX_GPU_NODES
+
+# Calculation for MIN_CPU_NODES.
+_PRODUCE_BLOCKS_PER_S = 1 / PRODUCE_SLEEP_S
+_BLOCKS_PER_CONSUME_BATCH = CONSUME_BATCH_SIZE / ROWS_PER_BLOCK
+_CONSUME_BLOCKS_PER_S = _BLOCKS_PER_CONSUME_BATCH / CONSUME_SLEEP_S
+_PRODUCE_WORKERS_PER_CONSUME_WORKER = _CONSUME_BLOCKS_PER_S / _PRODUCE_BLOCKS_PER_S
+
+_CPUS_NEEDED = EXPECTED_GPU_NODES * _PRODUCE_WORKERS_PER_CONSUME_WORKER
+_CPUS_FROM_GPU_NODES = EXPECTED_GPU_NODES * CPUS_PER_NODE
+MIN_CPU_NODES = math.ceil((_CPUS_NEEDED - _CPUS_FROM_GPU_NODES) / CPUS_PER_NODE)
+assert MIN_CPU_NODES > 0, (
+    "The GPU nodes' own CPUs already satisfy the pipeline, so this test no "
+    "longer exercises CPU scale-up. Adjust the sleeps or the compute config."
+)
+
+
+def produce(_: Dict[str, np.ndarray]) -> Iterator[Dict[str, np.ndarray]]:
+    for _ in range(BLOCKS_PER_INPUT):
+        time.sleep(PRODUCE_SLEEP_S)
+        yield {"data": np.zeros(BLOCK_SHAPE, dtype=np.uint8)}
+
+
+def consume(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+    time.sleep(CONSUME_SLEEP_S)
+    return batch
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--assert-max-cpu-nodes",
+        type=int,
+        default=None,
+        help="Assert that the autoscaler provisions at most this many CPU nodes.",
+    )
+    return parser.parse_args()
+
+
+def main(args: argparse.Namespace) -> Dict[str, Any]:
+    """This test checks if the cluster scales up enough to balance the pipeline."""
+    if not ray.is_initialized():
+        ray.init()
+
+    input_blocks = [
+        pa.Table.from_pydict({"input_id": [input_id]}) for input_id in range(NUM_INPUTS)
+    ]
+
+    with ClusterResourceMonitor() as monitor:
+        ds = (
+            ray.data.from_blocks(input_blocks)
+            .map_batches(produce)
+            .map_batches(consume, num_gpus=1, batch_size=CONSUME_BATCH_SIZE)
+        )
+        # Don't materialize, so blocks are freed as they're consumed.
+        for _ in ds.iter_internal_ref_bundles():
+            pass
+
+    peak_nodes = monitor.get_peak_node_counts()
+    print(f"Peak worker nodes: {peak_nodes.cpu} CPU, {peak_nodes.gpu} GPU")
+
+    assert peak_nodes.gpu == EXPECTED_GPU_NODES, (
+        f"Expected the autoscaler to provision {EXPECTED_GPU_NODES} GPU nodes, "
+        f"but it provisioned {peak_nodes.gpu}"
+    )
+    assert peak_nodes.cpu >= MIN_CPU_NODES, (
+        f"Expected the autoscaler to provision at least {MIN_CPU_NODES} CPU nodes "
+        f"to balance the pipeline, but it provisioned {peak_nodes.cpu}"
+    )
+    if args.assert_max_cpu_nodes is not None:
+        assert peak_nodes.cpu <= args.assert_max_cpu_nodes, (
+            f"Expected the autoscaler to provision at most "
+            f"{args.assert_max_cpu_nodes} CPU nodes, but it provisioned "
+            f"{peak_nodes.cpu}"
+        )
+
+    return {
+        "peak_cpu_nodes": peak_nodes.cpu,
+        "peak_gpu_nodes": peak_nodes.gpu,
+    }
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    benchmark = Benchmark()
+    benchmark.run_fn("main", main, args)
+    benchmark.write_result()

release/nightly_tests/dataset/autoscaling/autoscaler_scales_up_when_compute_bound_cluster_compute.yaml

@@ -0,0 +1,24 @@
+# Two autoscaling worker groups that both start empty. The GPU group uses CPU
+# instances each with a single fake GPU, so this test can exercise heterogeneous
+# autoscaling without using real GPUs.
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.2xlarge
+
+worker_nodes:
+    # Runs the produce tasks.
+    - name: cpu-node
+      instance_type: m5.2xlarge
+      min_nodes: 0
+      max_nodes: 10
+
+    # Runs the consume tasks. Only `GPU` is overridden, so these nodes still
+    # expose the 8 vCPUs of an m5.2xlarge and can run produce tasks too. The
+    # test's CPU-node arithmetic depends on that.
+    - name: gpu-node
+      instance_type: m5.2xlarge
+      min_nodes: 0
+      max_nodes: 10
+      resources:
+        GPU: 1

release/nightly_tests/dataset/cluster_resource_monitor.py

@@ -1,16 +1,45 @@
+import logging
 import time
 import threading
-from typing import Tuple, Optional
+from typing import NamedTuple, Tuple, Optional
 
 import ray
+from ray._common.constants import HEAD_NODE_RESOURCE_NAME
 from ray.data._internal.execution.interfaces import ExecutionResources
 
+logger = logging.getLogger(__name__)
+
+
+class NodeCounts(NamedTuple):
+    """The number of alive worker nodes, by whether they have a GPU."""
+
+    cpu: int
+    gpu: int
+
+
+def _count_worker_nodes() -> NodeCounts:
+    """Count the alive worker nodes, excluding the head node.
+
+    A node counts as a GPU node if it has any GPU resource.
+    """
+    cpu_nodes = 0
+    gpu_nodes = 0
+    for node in ray.nodes():
+        if not node.get("Alive", False):
+            continue
+        resources = node.get("Resources", {})
+        if HEAD_NODE_RESOURCE_NAME in resources:
+            continue
+        if resources.get("GPU", 0) > 0:
+            gpu_nodes += 1
+        else:
+            cpu_nodes += 1
+    return NodeCounts(cpu=cpu_nodes, gpu=gpu_nodes)
+
 
 class ClusterResourceMonitor:
     """Monitor and validate cluster resources during benchmark execution.
 
-    This class tracks the peak number of cluster resources during execution.
-
     This can be used to validate that the autoscaler behaves well.
     """
 
@@ -23,6 +52,7 @@ def __init__(self):
 
         self._peak_cpu_count: float = 0
         self._peak_gpu_count: float = 0
+        self._peak_node_counts = NodeCounts(cpu=0, gpu=0)
 
     def __repr__(self):
         return "ClusterResourceMonitor()"
@@ -37,20 +67,36 @@ def __enter__(self):
     def get_peak_cluster_resources(self) -> ExecutionResources:
         return ExecutionResources(cpu=self._peak_cpu_count, gpu=self._peak_gpu_count)
 
+    def get_peak_node_counts(self) -> NodeCounts:
+        """Get the peak number of alive worker nodes, excluding the head node."""
+        return self._peak_node_counts
+
     def _start_background_thread(
         self, interval_s: float = 5.0
     ) -> Tuple[threading.Thread, threading.Event]:
         stop_event = threading.Event()
 
         def monitor_cluster_resources():
             while not stop_event.is_set():
-                resources = ray.cluster_resources()
-                self._peak_cpu_count = max(
-                    self._peak_cpu_count, resources.get("CPU", 0)
-                )
-                self._peak_gpu_count = max(
-                    self._peak_gpu_count, resources.get("GPU", 0)
-                )
+                # These query the GCS, so a transient failure shouldn't kill the
+                # thread and leave the peaks frozen for the rest of the run.
+                try:
+                    resources = ray.cluster_resources()
+                    self._peak_cpu_count = max(
+                        self._peak_cpu_count, resources.get("CPU", 0)
+                    )
+                    self._peak_gpu_count = max(
+                        self._peak_gpu_count, resources.get("GPU", 0)
+                    )
+
+                    node_counts = _count_worker_nodes()
+                    self._peak_node_counts = NodeCounts(
+                        cpu=max(self._peak_node_counts.cpu, node_counts.cpu),
+                        gpu=max(self._peak_node_counts.gpu, node_counts.gpu),
+                    )
+                except Exception:
+                    logger.warning("Failed to sample cluster state.", exc_info=True)
+
                 time.sleep(interval_s)
 
         thread = threading.Thread(target=monitor_cluster_resources, daemon=True)

release/release_data_tests.yaml

@@ -1317,3 +1317,19 @@
       env: aws
       cluster:
         cluster_compute: dataset/cross_az_250_350_compute_aws.yaml
+
+###################
+# Autoscaling tests
+###################
+
+- name: autoscaler_scales_up_when_compute_bound
+  python: "3.10"
+  group: data-autoscaling
+
+  cluster:
+    anyscale_sdk_2026: true
+    cluster_compute: autoscaling/autoscaler_scales_up_when_compute_bound_cluster_compute.yaml
+
+  run:
+    timeout: 3600
+    script: python autoscaling/autoscaler_scales_up_when_compute_bound.py