Repository: ray-project/ray
File: python/ray/train/v2/tests/test_data_resource_cleanup.py (Python)

Code before:
    # shutdown_executor called on SplitCoordinator
    remote_mock.assert_called_once()


def test_after_worker_group_shutdown():
    callback = DatasetsCallback(create_dummy_run_context())

    # Mock SplitCoordinator shutdown_executor method
    coord_mock = create_autospec(SplitCoordinator)
    remote_mock = MagicMock()
    coord_mock.shutdown_executor.remote = remote_mock
    callback._coordinator_actors = [coord_mock]

    dummy_wg_context = WorkerGroupContext(
        run_attempt_id="test",
        train_fn_ref=DummyObjectRefWrapper(lambda: None),
        num_workers=4,
        resources_per_worker={"CPU": 1},
    )
    callback.after_worker_group_shutdown(dummy_wg_context)

    # shutdown_executor called on SplitCoordinator
    remote_mock.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_resources_when_updated(requester, prev_requests=None, timeout=3.0):
        """Retrieve resource requests within the specified timeout. Returns after a new request is made or when the time expires."""
        prev_requests = prev_requests or {}
        deadline = time.time() + timeout
        requests = {}
        while time.time() < deadline:
            requests = ray.get(
                requester.__ray_call__.remote(lambda r: r._resource_requests)
            )
            if requests != prev_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))

    # Explicity trigger autoscaling

Review comment (refactor):
Possibly out-of-scope for this PR since it's an existing issue, but this test breaks Ray Core and Ray Data abstraction barriers. I think we should refactor it to test against a higher level of abstraction (e.g., the `ClusterAutoscaler` interface).

Code after:
    # shutdown_executor called on SplitCoordinator
    remote_mock.assert_called_once()


def test_after_worker_group_shutdown():
    callback = DatasetsCallback(create_dummy_run_context())

    # Mock SplitCoordinator shutdown_executor method
    coord_mock = create_autospec(SplitCoordinator)
    remote_mock = MagicMock()
    coord_mock.shutdown_executor.remote = remote_mock
    callback._coordinator_actors = [coord_mock]

    dummy_wg_context = WorkerGroupContext(
        run_attempt_id="test",
        train_fn_ref=DummyObjectRefWrapper(lambda: None),
        num_workers=4,
        resources_per_worker={"CPU": 1},
    )
    callback.after_worker_group_shutdown(dummy_wg_context)

    # shutdown_executor called on SplitCoordinator
    remote_mock.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(