chore: import upstream snapshot with attribution

python/ray/autoscaler/v2/tests/test_priority_selection.py

@@ -0,0 +1,328 @@
+import os
+import sys
+import time
+
+import pytest
+
+from ray.autoscaler.v2.instance_manager.config import NodeTypeConfig
+from ray.autoscaler.v2.instance_manager.subscribers.cloud_resource_monitor import (
+    CloudResourceMonitor,
+)
+from ray.autoscaler.v2.scheduler import (
+    ResourceDemandScheduler,
+    ResourceRequestSource,
+    SchedulingNode,
+    SchedulingNodeStatus,
+    SchedulingRequest,
+)
+from ray.core.generated.autoscaler_pb2 import ResourceRequest as PBResourceRequest
+from ray.core.generated.instance_manager_pb2 import (
+    Instance,
+    InstanceUpdateEvent,
+    NodeKind,
+)
+
+
+def test_recovery_scoring():
+    monitor = CloudResourceMonitor()
+    node_type = "gpu-node"
+
+    # Mock failure
+    event = InstanceUpdateEvent(
+        instance_type=node_type, new_instance_status=Instance.ALLOCATION_TIMEOUT
+    )
+    monitor.notify([event])
+
+    # Immediately after failure, score should be 0.0
+    scores = monitor.get_recoverable_resource_availabilities()
+    assert scores[node_type] == 0.0
+
+    # After safety floor (e.g., 11s, default safety floor is 10s)
+    monitor._last_unavailable_timestamp[node_type] = time.time() - 11
+    scores = monitor.get_recoverable_resource_availabilities()
+    assert 0.0 < scores[node_type] < 0.1
+
+    # Halfway through recovery window (600s / 2 = 300s)
+    monitor._last_unavailable_timestamp[node_type] = time.time() - 300
+    scores = monitor.get_recoverable_resource_availabilities()
+    # 300 / 600 = 0.5
+    assert pytest.approx(scores[node_type], 0.01) == 0.5
+
+    # After recovery window
+    monitor._last_unavailable_timestamp[node_type] = time.time() - 601
+    scores = monitor.get_recoverable_resource_availabilities()
+    assert scores[node_type] == 1.0
+
+
+def test_scheduler_priority_tie_breaking():
+    # Two node types with identical resources
+    resources = {"CPU": 4}
+    node_type_1 = "high-priority"
+    node_type_2 = "low-priority"
+
+    config_1 = NodeTypeConfig(
+        name=node_type_1,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=10,
+    )
+    config_2 = NodeTypeConfig(
+        name=node_type_2,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=0,
+    )
+
+    node_1 = SchedulingNode.from_node_config(
+        config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+    node_2 = SchedulingNode.from_node_config(
+        config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+
+    request = PBResourceRequest(resources_bundle={"CPU": 1})
+
+    # Utilization and Availability are equal (all perfect)
+    # Priority should break the tie
+    best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
+        [request], [node_2, node_1], ResourceRequestSource.PENDING_DEMAND, {}, {}
+    )
+
+    assert best_node.node_type == node_type_1
+
+
+def test_schedule_context_propagation():
+    resources = {"CPU": 4}
+    node_type = "gpu-node"
+
+    config = NodeTypeConfig(
+        name=node_type,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=7,
+    )
+
+    # Mock cloud availabilities
+    cloud_availabilities = {node_type: 0.1}  # failure recency
+    recoverable_availabilities = {node_type: 0.5}
+
+    req = SchedulingRequest(
+        node_type_configs={node_type: config},
+        cloud_resource_availabilities=cloud_availabilities,
+        recoverable_resource_availabilities=recoverable_availabilities,
+        disable_launch_config_check=True,
+    )
+
+    ctx = ResourceDemandScheduler.ScheduleContext.from_schedule_request(req)
+
+    # Check if a new node created from this context has the correct priority
+    node_pools = [
+        SchedulingNode.from_node_config(
+            ctx.get_node_type_configs()[nt],
+            status=SchedulingNodeStatus.TO_LAUNCH,
+            node_kind=NodeKind.WORKER,
+        )
+        for nt, num_available in ctx.get_node_type_available().items()
+    ]
+
+    assert len(node_pools) == 1
+    node = node_pools[0]
+    assert node.priority == 7
+    # Dynamic scores are in the context, not the node
+    assert ctx.get_recoverable_resource_availabilities()[node_type] == 0.5
+    assert ctx.get_cloud_resource_availabilities()[node_type] == 0.1
+
+
+def test_scheduler_availability_over_priority():
+    # High priority node is recovering (score 0.5)
+    # Low priority node is available (score 1.0)
+    resources = {"CPU": 4}
+    node_type_1 = "high-priority"
+    node_type_2 = "low-priority"
+
+    config_1 = NodeTypeConfig(
+        name=node_type_1,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=10,
+    )
+    config_2 = NodeTypeConfig(
+        name=node_type_2,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=0,
+    )
+
+    node_1 = SchedulingNode.from_node_config(
+        config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+    node_2 = SchedulingNode.from_node_config(
+        config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+
+    request = PBResourceRequest(resources_bundle={"CPU": 1})
+
+    # Recoverable Availability is higher for node_2, so it should be chosen despite lower priority
+    best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
+        [request],
+        [node_1, node_2],
+        ResourceRequestSource.PENDING_DEMAND,
+        cloud_resource_availabilities={},
+        recoverable_resource_availabilities={node_type_1: 0.5, node_type_2: 1.0},
+    )
+
+    assert best_node.node_type == node_type_2
+
+
+def test_scheduler_failure_recency_tie_breaking():
+    # Same priority, same recoverable availability (1.0)
+    # One has an older failure than the other.
+    resources = {"CPU": 4}
+    node_type_1 = "older-failure"
+    node_type_2 = "newer-failure"
+
+    config_1 = NodeTypeConfig(
+        name=node_type_1,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=5,
+    )
+    config_2 = NodeTypeConfig(
+        name=node_type_2,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=5,
+    )
+
+    node_1 = SchedulingNode.from_node_config(
+        config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+    node_2 = SchedulingNode.from_node_config(
+        config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+
+    request = PBResourceRequest(resources_bundle={"CPU": 1})
+
+    best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
+        [request],
+        [node_1, node_2],
+        ResourceRequestSource.PENDING_DEMAND,
+        cloud_resource_availabilities={node_type_1: 0.9, node_type_2: 0.1},
+        recoverable_resource_availabilities={node_type_1: 1.0, node_type_2: 1.0},
+    )
+
+    assert best_node.node_type == node_type_1
+
+
+def test_recovery_integration():
+    monitor = CloudResourceMonitor()
+    node_type_1 = "high-priority"
+    node_type_2 = "low-priority"
+
+    # Mock failure for node_1 (high priority)
+    event = InstanceUpdateEvent(
+        instance_type=node_type_1, new_instance_status=Instance.ALLOCATION_TIMEOUT
+    )
+    monitor.notify([event])
+
+    # Score should be 0.0 for node_1 immediately
+    scores = monitor.get_recoverable_resource_availabilities()
+    assert scores[node_type_1] == 0.0
+
+    # Setup scheduler structures
+    resources = {"CPU": 4}
+    config_1 = NodeTypeConfig(
+        name=node_type_1,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=10,
+    )
+    config_2 = NodeTypeConfig(
+        name=node_type_2,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=0,
+    )
+
+    node_1 = SchedulingNode.from_node_config(
+        config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+    node_2 = SchedulingNode.from_node_config(
+        config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+
+    request = PBResourceRequest(resources_bundle={"CPU": 1})
+
+    # Pass scores from monitor to scheduler
+    best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
+        [request],
+        [node_1, node_2],
+        ResourceRequestSource.PENDING_DEMAND,
+        cloud_resource_availabilities={},
+        recoverable_resource_availabilities=scores,
+    )
+
+    # Node 2 should be chosen because Node 1 is recovering, despite Node 1 having higher priority
+    assert best_node.node_type == node_type_2
+
+
+def test_scheduler_utilization_over_priority():
+    # Node 1: 4 CPUs, priority 10
+    # Node 2: 2 CPUs, priority 0
+    # Request: 2 CPUs
+    # Node 2 should be selected because it fits perfectly (utilization score is higher),
+    # even though Node 1 has higher priority.
+    resources_1 = {"CPU": 4}
+    resources_2 = {"CPU": 2}
+    node_type_1 = "large-node"
+    node_type_2 = "small-node"
+
+    config_1 = NodeTypeConfig(
+        name=node_type_1,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources_1,
+        priority=10,
+    )
+    config_2 = NodeTypeConfig(
+        name=node_type_2,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources_2,
+        priority=0,
+    )
+
+    node_1 = SchedulingNode.from_node_config(
+        config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+    node_2 = SchedulingNode.from_node_config(
+        config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+
+    request = PBResourceRequest(resources_bundle={"CPU": 2})
+
+    best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
+        [request],
+        [node_1, node_2],
+        ResourceRequestSource.PENDING_DEMAND,
+        cloud_resource_availabilities={},
+        recoverable_resource_availabilities={},
+    )
+
+    assert best_node.node_type == node_type_2
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))