chore: import upstream snapshot with attribution

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

@@ -0,0 +1,478 @@
+import asyncio
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+from freezegun import freeze_time
+
+from ray.data._internal.cluster_autoscaler.default_autoscaling_coordinator import (
+    ResourceRequestPriority,
+)
+from ray.train.v2._internal.execution.callback import ControllerCallback
+from ray.train.v2._internal.execution.scaling_policy import (
+    AUTOSCALING_REQUESTS_EXPIRE_TIME_S,
+    AUTOSCALING_REQUESTS_INTERVAL_S,
+    NoopDecision,
+    ResizeDecision,
+)
+from ray.train.v2._internal.execution.scaling_policy.elastic import (
+    ElasticScalingPolicy,
+)
+from ray.train.v2._internal.execution.worker_group import (
+    WorkerGroupPollStatus,
+    WorkerGroupState,
+    WorkerStatus,
+)
+from ray.train.v2._internal.util import time_monotonic
+from ray.train.v2.api.config import ScalingConfig
+
+
+@pytest.fixture(autouse=True)
+def mock_autoscaling_coordinator(monkeypatch):
+    mock_coordinator = MagicMock()
+    mock_coordinator._allocated_resources = None
+    mock_coordinator.get_allocated_resources.remote = MagicMock(
+        side_effect=lambda _: mock_coordinator._allocated_resources
+    )
+
+    monkeypatch.setattr(
+        ElasticScalingPolicy, "_autoscaling_coordinator", mock_coordinator
+    )
+
+
+@pytest.fixture(autouse=True)
+def patch_ray_get():
+    with patch(
+        "ray.get",
+        side_effect=lambda x, **_: x,
+    ):
+        yield
+
+
+def _get_mock_worker_group_status(num_workers: int) -> WorkerGroupPollStatus:
+    return WorkerGroupPollStatus(
+        worker_statuses={
+            i: WorkerStatus(running=True, error=None) for i in range(num_workers)
+        },
+    )
+
+
+def _get_mock_worker_group_state(
+    num_workers: int, start_time: float
+) -> WorkerGroupState:
+    return WorkerGroupState(
+        start_time=start_time,
+        placement_group_handle=MagicMock(),
+        workers=[MagicMock() for _ in range(num_workers)],
+        sync_actor=MagicMock(),
+    )
+
+
+@patch.object(ElasticScalingPolicy, "GET_ALLOCATED_RESOURCES_INTERVAL_S", 0.0)
+def test_non_running_worker_group_decision():
+    """Test decisions being made when the worker group is initializing/restarting.
+    Ensures that the policy will resize the worker group as soon as resources are available.
+    """
+    min_workers, max_workers = 4, 64
+    resources_per_worker = {"CPU": 8, "GPU": 1}
+
+    scaling_config = ScalingConfig(
+        num_workers=(min_workers, max_workers),
+        resources_per_worker=resources_per_worker,
+        use_gpu=True,
+    )
+    policy = ElasticScalingPolicy(scaling_config)
+    mock_coordinator = policy._autoscaling_coordinator
+
+    # No resources are available at the start
+    decision = policy.make_decision_for_non_running_worker_group()
+    assert isinstance(decision, NoopDecision)
+
+    # Resources for < min workers are available
+    mock_coordinator._allocated_resources = [resources_per_worker] * (min_workers - 1)
+    decision = policy.make_decision_for_non_running_worker_group()
+    assert isinstance(decision, NoopDecision)
+
+    # Resources for >= min workers are available
+    mock_coordinator._allocated_resources = [resources_per_worker] * min_workers
+    decision = policy.make_decision_for_non_running_worker_group()
+    assert isinstance(decision, ResizeDecision)
+    assert decision.num_workers == min_workers
+
+    # Resources for >= max workers are available
+    mock_coordinator._allocated_resources = [resources_per_worker] * max_workers
+    decision = policy.make_decision_for_non_running_worker_group()
+    assert isinstance(decision, ResizeDecision)
+    assert decision.num_workers == max_workers
+
+
+def test_before_controller_abort():
+    """Test that before_controller_abort sends a cancel request to the AutoscalingCoordinator."""
+    resources_per_worker = {"CPU": 4, "GPU": 1}
+    scaling_config = ScalingConfig(
+        num_workers=(2, 4),
+        resources_per_worker=resources_per_worker,
+        use_gpu=True,
+    )
+    policy = ElasticScalingPolicy(scaling_config)
+    mock_coordinator = policy._autoscaling_coordinator
+
+    # Call before_controller_abort and check that cancel_request is called with the requester_id
+    policy.before_controller_abort()
+    mock_coordinator.cancel_request.remote.assert_called_once_with(
+        requester_id=policy._requester_id
+    )
+
+
+def test_get_allocated_resources_interval():
+    """Tests that remote calls to the AutoscalingCoordinator are spaced out by a minimum time interval."""
+    min_workers, max_workers = 4, 64
+    resources_per_worker = {"CPU": 8, "GPU": 1}
+    get_allocated_resources_interval_s = (
+        ElasticScalingPolicy.GET_ALLOCATED_RESOURCES_INTERVAL_S
+    )
+
+    scaling_config = ScalingConfig(
+        num_workers=(min_workers, max_workers),
+        resources_per_worker=resources_per_worker,
+        use_gpu=True,
+    )
+    policy = ElasticScalingPolicy(scaling_config)
+    mock_coordinator = policy._autoscaling_coordinator
+
+    with freeze_time() as frozen_time:
+        # No resources are available at the start
+        allocated_resources = policy._get_allocated_resources()
+        assert allocated_resources is None
+
+        # Resources for < min workers are available
+        frozen_time.tick(get_allocated_resources_interval_s)
+        mock_coordinator._allocated_resources = [resources_per_worker] * (
+            min_workers - 1
+        )
+        allocated_resources = policy._get_allocated_resources()
+        assert allocated_resources == [resources_per_worker] * (min_workers - 1)
+
+        # Resources for >= min workers are available, but get_allocated_resources interval
+        # has not yet passed.
+        mock_coordinator._allocated_resources = [resources_per_worker] * min_workers
+        allocated_resources = policy._get_allocated_resources()
+        assert allocated_resources == [resources_per_worker] * (min_workers - 1)
+
+        # Resources for >= min workers are available and the get_allocated_resources
+        # interval has passed.
+        frozen_time.tick(get_allocated_resources_interval_s)
+        mock_coordinator._allocated_resources = [resources_per_worker] * min_workers
+        allocated_resources = policy._get_allocated_resources()
+        assert allocated_resources == [resources_per_worker] * min_workers
+
+        # Resources for >= max workers are available but the get_allocated_resources
+        # interval has not yet passed.
+        mock_coordinator._allocated_resources = [resources_per_worker] * max_workers
+        allocated_resources = policy._get_allocated_resources()
+        assert allocated_resources == [resources_per_worker] * min_workers
+
+        # Resources for >= max workers are available and the get_allocated_resources
+        # interval has passed.
+        frozen_time.tick(get_allocated_resources_interval_s)
+        allocated_resources = policy._get_allocated_resources()
+        assert allocated_resources == [resources_per_worker] * max_workers
+
+
+@patch.object(ElasticScalingPolicy, "GET_ALLOCATED_RESOURCES_INTERVAL_S", 0.0)
+def test_running_worker_group_decision():
+    """Test decisions being made when the worker group is running.
+    Ensures that the policy will resize the worker group when there is a change
+    in available resources.
+    """
+    min_workers, max_workers = 4, 64
+    resources_per_worker = {"CPU": 8, "GPU": 1}
+
+    scaling_config = ScalingConfig(
+        num_workers=(min_workers, max_workers),
+        resources_per_worker=resources_per_worker,
+        use_gpu=True,
+        # NOTE: This test just asserts the policy decisions, not the monitor interval.
+        elastic_resize_monitor_interval_s=0.0,
+    )
+    policy = ElasticScalingPolicy(scaling_config)
+    mock_coordinator = policy._autoscaling_coordinator
+
+    # The worker group just started
+    worker_group_state = _get_mock_worker_group_state(min_workers, time_monotonic())
+    worker_group_status = _get_mock_worker_group_status(min_workers)
+
+    # No change in resources
+    mock_coordinator._allocated_resources = [resources_per_worker] * min_workers
+    decision = policy.make_decision_for_running_worker_group(
+        worker_group_state=worker_group_state,
+        worker_group_status=worker_group_status,
+    )
+    assert isinstance(decision, NoopDecision)
+
+    # Resources for < min workers are available
+    mock_coordinator._allocated_resources = [resources_per_worker] * (min_workers - 1)
+    decision = policy.make_decision_for_running_worker_group(
+        worker_group_state=worker_group_state,
+        worker_group_status=worker_group_status,
+    )
+    assert isinstance(decision, NoopDecision)
+
+    # More resources are available.
+    mock_coordinator._allocated_resources = [resources_per_worker] * max_workers
+    decision = policy.make_decision_for_running_worker_group(
+        worker_group_state=worker_group_state,
+        worker_group_status=worker_group_status,
+    )
+    assert isinstance(decision, ResizeDecision)
+    assert decision.num_workers == max_workers
+
+
+def test_monitor_recently_started_worker_group():
+    """Test monitor decisions being made when the worker group is running.
+    Ensures that resizing decisions are not made too soon after the worker group starts.
+    """
+    min_workers, max_workers = 4, 64
+    monitor_interval_s = 60
+    resources_per_worker = {"CPU": 8, "GPU": 1}
+
+    scaling_config = ScalingConfig(
+        num_workers=(min_workers, max_workers),
+        resources_per_worker=resources_per_worker,
+        use_gpu=True,
+        elastic_resize_monitor_interval_s=monitor_interval_s,
+    )
+    policy = ElasticScalingPolicy(scaling_config)
+    mock_coordinator = policy._autoscaling_coordinator
+
+    with freeze_time() as frozen_time:
+        # The worker group just started
+        worker_group_state = _get_mock_worker_group_state(min_workers, time_monotonic())
+        worker_group_status = _get_mock_worker_group_status(min_workers)
+
+        # Advance time partway through the monitor interval
+        frozen_time.tick(delta=monitor_interval_s / 2)
+
+        # Even though there are new resources available, we should not resize yet
+        # because the monitor interval has not passed since
+        mock_coordinator._allocated_resources = [resources_per_worker] * (
+            max_workers - 1
+        )
+
+        assert isinstance(
+            policy.make_decision_for_running_worker_group(
+                worker_group_state=worker_group_state,
+                worker_group_status=worker_group_status,
+            ),
+            NoopDecision,
+        )
+
+        frozen_time.tick(delta=monitor_interval_s / 2)
+
+        # The monitor interval has passed, should detect resources and resize
+        decision = policy.make_decision_for_running_worker_group(
+            worker_group_state=worker_group_state,
+            worker_group_status=worker_group_status,
+        )
+        assert isinstance(decision, ResizeDecision)
+        assert decision.num_workers == max_workers - 1
+
+
+def test_monitor_long_running_worker_group():
+    """Test monitor decisions being made when the worker group is running.
+    Ensures that the resizing considerations are not made too frequently.
+    """
+    min_workers, max_workers = 4, 64
+    monitor_interval_s = 60
+    resources_per_worker = {"CPU": 8, "GPU": 1}
+
+    scaling_config = ScalingConfig(
+        num_workers=(min_workers, max_workers),
+        resources_per_worker=resources_per_worker,
+        use_gpu=True,
+        elastic_resize_monitor_interval_s=monitor_interval_s,
+    )
+    policy = ElasticScalingPolicy(scaling_config)
+    mock_coordinator = policy._autoscaling_coordinator
+
+    with freeze_time() as frozen_time:
+        worker_group_state = _get_mock_worker_group_state(min_workers, time_monotonic())
+        worker_group_status = _get_mock_worker_group_status(min_workers)
+        mock_coordinator._allocated_resources = [resources_per_worker] * min_workers
+
+        # The worker group has been running for a while at the same size
+        frozen_time.tick(monitor_interval_s * 60)
+
+        # Consider resizing.
+        decision = policy.make_decision_for_running_worker_group(
+            worker_group_state=worker_group_state,
+            worker_group_status=worker_group_status,
+        )
+        assert isinstance(decision, NoopDecision)
+
+        # We recently considered resizing, so we should wait until the next interval
+        # to consider again --> no-op even if new resources are available
+        mock_coordinator._allocated_resources = [resources_per_worker] * max_workers
+        frozen_time.tick(monitor_interval_s / 2)
+        decision = policy.make_decision_for_running_worker_group(
+            worker_group_state=worker_group_state,
+            worker_group_status=worker_group_status,
+        )
+        assert isinstance(decision, NoopDecision)
+
+        frozen_time.tick(monitor_interval_s / 2)
+        decision = policy.make_decision_for_running_worker_group(
+            worker_group_state=worker_group_state,
+            worker_group_status=worker_group_status,
+        )
+        assert isinstance(decision, ResizeDecision)
+        assert decision.num_workers == max_workers
+
+
+def test_count_possible_workers():
+    """Test counting the number of workers that can be started with
+    available node resources."""
+    resources_per_worker = {"CPU": 8, "GPU": 1}
+    scaling_config = ScalingConfig(
+        num_workers=(1, 8),
+        use_gpu=True,
+        resources_per_worker=resources_per_worker,
+    )
+    policy = ElasticScalingPolicy(scaling_config)
+
+    # No resources
+    assert policy._count_possible_workers([]) == 0
+
+    # Single node
+    assert policy._count_possible_workers([{"CPU": 8, "GPU": 1}]) == 1
+    assert policy._count_possible_workers([{"CPU": 16, "GPU": 2}]) == 2
+    assert policy._count_possible_workers([{"CPU": 16, "GPU": 1}]) == 1
+
+    # Multinode
+    assert policy._count_possible_workers([{"CPU": 7, "GPU": 1}] * 2) == 0
+    assert policy._count_possible_workers([{"CPU": 9, "GPU": 2}] * 8) == 8
+    assert policy._count_possible_workers([{"CPU": 16, "GPU": 2}] * 2) == 4
+    assert policy._count_possible_workers([{"CPU": 8, "GPU": 1}] * 4) == 4
+
+    # If there are excess resources, the number of workers is still capped at max_workers
+    assert policy._count_possible_workers([{"CPU": 16, "GPU": 2}] * 10) == 8
+
+
+def test_count_possible_workers_with_zero_resources():
+    max_workers = 4
+    scaling_config = ScalingConfig(
+        num_workers=(1, max_workers),
+        resources_per_worker={"CPU": 0, "GPU": 0, "memory": 0},
+    )
+    policy = ElasticScalingPolicy(scaling_config)
+
+    assert (
+        policy._count_possible_workers([{"CPU": 1, "GPU": 1, "memory": 1}])
+        == max_workers
+    )
+
+
+def test_request_and_clear():
+    """Tests that the policy makes resource requests and clears the requests."""
+    resources_per_worker = {"CPU": 8, "GPU": 1}
+    policy = ElasticScalingPolicy(
+        scaling_config=ScalingConfig(
+            use_gpu=True, resources_per_worker=resources_per_worker, num_workers=(2, 4)
+        )
+    )
+    assert isinstance(policy, ControllerCallback)
+    mock_coordinator = policy._autoscaling_coordinator
+
+    def assert_resource_request_called_with():
+        nonlocal mock_coordinator
+
+        mock_coordinator.request_resources.remote.assert_called_with(
+            requester_id=policy._requester_id,
+            resources=[resources_per_worker] * 4,
+            label_selectors=None,
+            expire_after_s=AUTOSCALING_REQUESTS_EXPIRE_TIME_S,
+            priority=ResourceRequestPriority.HIGH,
+        )
+
+    with freeze_time() as frozen_time:
+        worker_group_state = _get_mock_worker_group_state(2, time_monotonic())
+        worker_group_status = _get_mock_worker_group_status(2)
+
+        # Test request_resources is called when the controller starts.
+        policy.after_controller_start(train_run_context=MagicMock())
+        assert mock_coordinator.request_resources.remote.call_count == 1
+        assert_resource_request_called_with()
+
+        # Test request_resources is only called in
+        # `make_decision_for_running_worker_group`,
+        # if `AUTOSCALING_REQUESTS_INTERVAL_S` has passed.
+        frozen_time.tick(AUTOSCALING_REQUESTS_INTERVAL_S / 2)
+        policy.make_decision_for_running_worker_group(
+            worker_group_state=worker_group_state,
+            worker_group_status=worker_group_status,
+        )
+        assert mock_coordinator.request_resources.remote.call_count == 1
+
+        frozen_time.tick(AUTOSCALING_REQUESTS_INTERVAL_S / 2)
+        policy.make_decision_for_running_worker_group(
+            worker_group_state=worker_group_state,
+            worker_group_status=worker_group_status,
+        )
+        assert mock_coordinator.request_resources.remote.call_count == 2
+        assert_resource_request_called_with()
+
+    # Test cancel_request is called when the controller is shutting down.
+    asyncio.run(policy.before_controller_shutdown())
+    mock_coordinator.cancel_request.remote.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    "num_autoscaler_nodes, mock_gcs_intact_slices, expected_workers",
+    [
+        # No resources -> 0 workers
+        (0, 0, 0),
+        # Autoscaler sees 3 nodes but slice requires 4 -> rounds to 0 workers
+        (3, 0, 0),
+        # Autoscaler sees 4 nodes and 1 intact slice -> 4 workers
+        (4, 1, 4),
+        # Autoscaler sees 8 nodes (2 slices worth) but only 1 slice is intact
+        # (e.g. 1 host down in the other slice) -> capped at 4 workers
+        (8, 1, 4),
+        # Autoscaler sees 8 nodes and both slices are intact -> 8 workers
+        (8, 2, 8),
+        # Autoscaler sees 12 nodes (3 slices), all 3 intact -> 12 workers (max)
+        (12, 3, 12),
+    ],
+)
+@patch("ray.util.tpu.get_num_tpu_slices")
+def test_count_possible_workers_tpu_slice_rounding(
+    mock_get_intact_slices,
+    num_autoscaler_nodes,
+    mock_gcs_intact_slices,
+    expected_workers,
+):
+    """
+    Test that TPU scaling correctly floors to the nearest complete, physically
+    intact TPU slice. Intact slices are counted regardless of whether they are
+    currently idle or in use.
+    """
+    mock_get_intact_slices.return_value = mock_gcs_intact_slices
+
+    # Scaling config for TPU v6e 4x4 between 1 and 3 slices.
+    scaling_config = ScalingConfig(
+        use_tpu=True,
+        accelerator_type="TPU-V6E",
+        topology="4x4",
+        num_workers=(4, 12),
+        resources_per_worker={"TPU": 4, "CPU": 1},
+    )
+    policy = ElasticScalingPolicy(scaling_config)
+
+    tpu_node = {"TPU": 4, "CPU": 1, "accelerator_type:TPU-V6E": 1}
+    allocated_resources = [tpu_node] * num_autoscaler_nodes
+
+    assert policy._count_possible_workers(allocated_resources) == expected_workers
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))