chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_data_resource_cleanup.py

@@ -0,0 +1,110 @@
+import sys
+import time
+from unittest.mock import create_autospec
+
+import pytest
+
+import ray
+from ray.data._internal.cluster_autoscaler.default_autoscaling_coordinator import (
+    get_or_create_autoscaling_coordinator,
+)
+from ray.data._internal.iterator.stream_split_iterator import (
+    SplitCoordinator,
+)
+from ray.train.v2._internal.callbacks.datasets import (
+    DatasetsCallback,
+    RayDatasetShardProvider,
+)
+from ray.train.v2._internal.execution.worker_group import WorkerGroupContext
+from ray.train.v2.tests.util import DummyObjectRefWrapper, create_dummy_run_context
+
+pytestmark = pytest.mark.usefixtures("mock_runtime_context")
+
+
+def _dummy_worker_group_context() -> WorkerGroupContext:
+    return WorkerGroupContext(
+        run_attempt_id="test",
+        train_fn_ref=DummyObjectRefWrapper(lambda: None),
+        num_workers=4,
+        resources_per_worker={"CPU": 1},
+    )
+
+
+def test_after_worker_group_shutdown():
+    """The callback delegates shutdown to the dataset shard provider."""
+    callback = DatasetsCallback(
+        train_run_context=create_dummy_run_context(), datasets={}
+    )
+    shard_provider = create_autospec(RayDatasetShardProvider)
+    callback._dataset_shard_provider = shard_provider
+
+    callback.after_worker_group_shutdown(
+        worker_group_context=_dummy_worker_group_context()
+    )
+    shard_provider.shutdown_data_executors.assert_called_once()
+
+
+def test_after_worker_group_abort():
+    """The callback delegates abort cleanup to the dataset shard provider."""
+    callback = DatasetsCallback(
+        train_run_context=create_dummy_run_context(), datasets={}
+    )
+    shard_provider = create_autospec(RayDatasetShardProvider)
+    callback._dataset_shard_provider = shard_provider
+
+    callback.after_worker_group_abort(
+        worker_group_context=_dummy_worker_group_context()
+    )
+    shard_provider.shutdown_data_executors.assert_called_once()
+
+
+def test_split_coordinator_shutdown_executor(ray_start_4_cpus):
+    """Tests that the SplitCoordinator properly requests resources for the data executor and cleans up after it is shutdown"""
+
+    def get_ongoing_requests(coordinator, timeout=3.0):
+        """Retrieve ongoing requests from the AutoscalingCoordinator."""
+        deadline = time.time() + timeout
+        requests = {}
+        while time.time() < deadline:
+            requests = ray.get(
+                coordinator.__ray_call__.remote(lambda c: dict(c._ongoing_reqs))
+            )
+            if requests:
+                break
+            time.sleep(0.05)
+        return requests
+
+    # Start coordinator and executor
+    NUM_SPLITS = 1
+    dataset = ray.data.range(100)
+    coord = SplitCoordinator.options(name="test_split_coordinator").remote(
+        dataset, NUM_SPLITS, None
+    )
+    ray.get(coord.start_epoch.remote(0))
+
+    # Explicitly trigger autoscaling
+    ray.get(
+        coord.__ray_call__.remote(
+            lambda coord: coord._current_executor._cluster_autoscaler.try_trigger_scaling()
+        )
+    )
+
+    # Collect requests from the AutoscalingCoordinator
+    coordinator = get_or_create_autoscaling_coordinator()
+    requests = get_ongoing_requests(coordinator)
+
+    # One request made (V2 registers with the coordinator)
+    assert len(requests) == 1
+    requester_id = list(requests.keys())[0]
+    assert requester_id.startswith("data-")
+
+    # Shutdown data executor
+    ray.get(coord.shutdown_executor.remote())
+
+    # Verify that the request is cancelled (removed from ongoing requests)
+    requests = ray.get(coordinator.__ray_call__.remote(lambda c: dict(c._ongoing_reqs)))
+    assert len(requests) == 0, "Resource request was not cancelled"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-x", __file__]))