/python/ray/serve/tests/unit/test_deployment_scheduler.py

import random
import sys
from collections import defaultdict
from typing import List
from unittest import mock
from unittest.mock import Mock

import pytest

import ray
from ray._raylet import NodeID
from ray.serve._private import default_impl
from ray.serve._private.common import (
    GANG_PG_NAME_PREFIX,
    CreatePlacementGroupRequest,
    DeploymentID,
    GangPlacementGroupRequest,
    ReplicaID,
)
from ray.serve._private.config import ReplicaConfig
from ray.serve._private.constants import (
    RAY_SERVE_USE_PACK_SCHEDULING_STRATEGY,
)
from ray.serve._private.deployment_scheduler import (
    AvailableNodeResources,
    DeploymentDownscaleRequest,
    DeploymentSchedulingInfo,
    ReplicaSchedulingRequest,
    ReplicaSchedulingRequestStatus,
    RequestedResources,
    Resources,
    SpreadDeploymentSchedulingPolicy,
)
from ray.serve._private.deployment_state import DeploymentStateManager
from ray.serve._private.test_utils import (
    MockActorClass,
    MockClusterNodeInfoCache,
    MockPlacementGroup,
)
from ray.tests.conftest import *  # noqa
from ray.util.scheduling_strategies import (
    In,
    NodeAffinitySchedulingStrategy,
    NodeLabelSchedulingStrategy,
    PlacementGroupSchedulingStrategy,
)


def dummy():
    pass


def rconfig(**config_opts):
    return ReplicaConfig.create(dummy, **config_opts)


def get_random_resources(n: int) -> List[dict]:
    """Gets n random resources."""

    resources = {
        "CPU": lambda: random.randint(0, 10),
        "GPU": lambda: random.randint(0, 10),
        "memory": lambda: random.randint(0, 10),
        "custom_A": lambda: random.randint(0, 10),
    }

    res = list()
    for _ in range(n):
        resource_dict = dict()
        for resource, callable in resources.items():
            if random.randint(0, 1) == 0:
                resource_dict[resource] = callable()

        res.append(resource_dict)

    return res


class TestResources:
    def test_base_resources_cannot_be_instantiated(self):
        with pytest.raises(TypeError, match="cannot be instantiated directly"):
            Resources()

    @pytest.mark.parametrize("resource_type", ["CPU", "GPU", "memory"])
    def test_basic(self, resource_type: str):
        # basic resources
        a = AvailableNodeResources({resource_type: 1})
        b = RequestedResources({resource_type: 0})
        assert a.can_fit(b)

        b = AvailableNodeResources({resource_type: 0})
        a = RequestedResources({resource_type: 1})
        assert not b.can_fit(a)

    def test_neither_bigger(self):
        a = AvailableNodeResources({"CPU": 1, "GPU": 0})
        b = RequestedResources({"CPU": 0, "GPU": 1})
        assert not a == b
        assert not a.can_fit(b)

    combos = [tuple(get_random_resources(20)[i : i + 2]) for i in range(0, 20, 2)]

    @pytest.mark.parametrize("resource_A,resource_B", combos)
    @pytest.mark.parametrize(
        "resource_class", [AvailableNodeResources, RequestedResources]
    )
    def test_soft_resources_consistent_comparison(
        self, resource_A, resource_B, resource_class
    ):
        """Resources should have consistent comparison. Either A==B, A<B, or A>B."""

        assert (
            resource_class(resource_A) == resource_class(resource_B)
            or resource_class(resource_A) > resource_class(resource_B)
            or resource_class(resource_A) < resource_class(resource_B)
        )

    @pytest.mark.parametrize(
        "resource_class", [AvailableNodeResources, RequestedResources]
    )
    def test_compare_resources(self, resource_class):
        # Prioritize GPU
        a = resource_class({"GPU": 1, "CPU": 10, "memory": 10, "custom": 10})
        b = resource_class({"GPU": 2, "CPU": 0, "memory": 0, "custom": 0})
        assert b > a

        # Then CPU
        a = resource_class({"GPU": 1, "CPU": 1, "memory": 10, "custom": 10})
        b = resource_class({"GPU": 1, "CPU": 2, "memory": 0, "custom": 0})
        assert b > a

        # Then memory
        a = resource_class({"GPU": 1, "CPU": 1, "memory": 1, "custom": 10})
        b = resource_class({"GPU": 1, "CPU": 1, "memory": 2, "custom": 0})
        assert b > a

        # Then custom resources
        a = resource_class({"GPU": 1, "CPU": 1, "memory": 1, "custom": 1})
        b = resource_class({"GPU": 1, "CPU": 1, "memory": 1, "custom": 2})
        assert b > a

    @pytest.mark.parametrize(
        "resource_class", [AvailableNodeResources, RequestedResources]
    )
    def test_sort_resources(self, resource_class):
        """Prioritize GPUs, CPUs, memory, then custom resources when sorting."""

        a = resource_class({"GPU": 0, "CPU": 4, "memory": 99, "A": 10})
        b = resource_class({"GPU": 0, "CPU": 2, "memory": 100})
        c = resource_class({"GPU": 1, "CPU": 1, "memory": 50})
        d = resource_class({"GPU": 2, "CPU": 0, "memory": 0})
        e = resource_class({"GPU": 3, "CPU": 8, "memory": 10000, "A": 6})
        f = resource_class({"GPU": 3, "CPU": 8, "memory": 10000, "A": 2})

        for _ in range(10):
            resources = [a, b, c, d, e, f]
            random.shuffle(resources)
            resources.sort(reverse=True)
            assert resources == [e, f, d, c, a, b]

    def test_custom_resources(self):
        a = AvailableNodeResources({"alice": 3})
        b = AvailableNodeResources({"alice": 2})
        assert b < a
        assert a.can_fit(b)
        assert a + b == AvailableNodeResources(**{"alice": 5})

        a = AvailableNodeResources({"bob": 2})
        b = AvailableNodeResources({"CPU": 4})
        assert a + b == AvailableNodeResources(**{"CPU": 4, "bob": 2})

    def test_implicit_resources(self):
        r = AvailableNodeResources()
        # Implicit resources
        assert r.get(f"{ray._raylet.IMPLICIT_RESOURCE_PREFIX}random") == 1
        # Everything else
        assert r.get("CPU") == 0
        assert r.get("GPU") == 0
        assert r.get("memory") == 0
        assert r.get("random_custom") == 0

        # Arithmetric with implicit resources
        implicit_resource_1 = f"{ray._raylet.IMPLICIT_RESOURCE_PREFIX}whatever"
        implicit_resource_2 = f"{ray._raylet.IMPLICIT_RESOURCE_PREFIX}whatever2"
        a = AvailableNodeResources()
        b = RequestedResources({implicit_resource_1: 0.5})
        c = RequestedResources({implicit_resource_2: 0.25})
        assert a.get(implicit_resource_1) == 1
        assert a.can_fit(b)

        a -= b
        assert a.get(implicit_resource_1) == 0.5
        assert a.can_fit(b)

        a -= b
        assert a.get(implicit_resource_1) == 0
        assert not a.can_fit(b)

        for i in range(4):
            assert a.can_fit(c)

            a -= c
            assert a.get(implicit_resource_1) == 0
            assert a.get(implicit_resource_2) == 1 - 0.25 * (i + 1)

        # Implicit resources exhausted
        assert not a.can_fit(c)


def test_deployment_scheduling_info():
    info = DeploymentSchedulingInfo(
        deployment_id=DeploymentID("a", "b"),
        scheduling_policy=SpreadDeploymentSchedulingPolicy,
        actor_resources=RequestedResources({"CPU": 2, "GPU": 1}),
    )
    assert info.required_resources == RequestedResources({"CPU": 2, "GPU": 1})
    assert not info.is_non_strict_pack_pg()

    info = DeploymentSchedulingInfo(
        deployment_id=DeploymentID("a", "b"),
        scheduling_policy=SpreadDeploymentSchedulingPolicy,
        actor_resources=RequestedResources({"CPU": 2, "GPU": 1}),
        placement_group_bundles=[
            RequestedResources({"CPU": 100}),
            RequestedResources({"GPU": 100}),
        ],
        placement_group_strategy="STRICT_PACK",
    )
    assert info.required_resources == RequestedResources({"CPU": 100, "GPU": 100})
    assert not info.is_non_strict_pack_pg()

    info = DeploymentSchedulingInfo(
        deployment_id=DeploymentID("a", "b"),
        scheduling_policy=SpreadDeploymentSchedulingPolicy,
        actor_resources=RequestedResources({"CPU": 1}),
        placement_group_bundles=[
            # Bundle 0 hosts the actor's CPU plus the CPU+GPU for a child
            # task/actor captured into the PG.
            RequestedResources({"CPU": 2, "GPU": 1}),
            RequestedResources({"CPU": 1, "GPU": 1}),
        ],
        placement_group_strategy="PACK",
    )
    # Actor is pinned as a subset of bundle 0, so required_resources is
    # bundle 0's full reservation, not just actor_resources.
    assert info.required_resources == RequestedResources({"CPU": 2, "GPU": 1})
    assert info.is_non_strict_pack_pg()


def test_deployment_scheduling_info_required_resources_no_mutation():
    dep_id = DeploymentID("app", "name")
    actor = RequestedResources({"CPU": 1})
    info = DeploymentSchedulingInfo(
        deployment_id=dep_id,
        scheduling_policy=SpreadDeploymentSchedulingPolicy,
        actor_resources=actor,
        max_replicas_per_node=2,
    )
    implicit = (
        f"{ray._raylet.IMPLICIT_RESOURCE_PREFIX}" f"{dep_id.app_name}:{dep_id.name}"
    )
    assert info.required_resources == RequestedResources({"CPU": 1, implicit: 0.5})
    assert actor == RequestedResources({"CPU": 1})
    assert implicit not in actor


def test_max_replicas_per_node_zero_skips_implicit_resource():
    """Falsy max_replicas_per_node (e.g. 0) must not trigger 1.0 / 0."""
    dep_id = DeploymentID("app", "name")
    implicit = (
        f"{ray._raylet.IMPLICIT_RESOURCE_PREFIX}" f"{dep_id.app_name}:{dep_id.name}"
    )

    info = DeploymentSchedulingInfo(
        deployment_id=dep_id,
        scheduling_policy=SpreadDeploymentSchedulingPolicy,
        actor_resources=RequestedResources({"CPU": 1}),
        max_replicas_per_node=0,
    )
    assert info.required_resources == RequestedResources({"CPU": 1})
    assert implicit not in info.required_resources

    req = ReplicaSchedulingRequest(
        replica_id=ReplicaID("r0", dep_id),
        actor_def=MockActorClass(),
        actor_resources={"CPU": 1},
        actor_options={"name": "r0"},
        actor_init_args=(),
        on_scheduled=lambda *args, **kwargs: None,
        max_replicas_per_node=0,
    )
    assert req.requested_resources == RequestedResources({"CPU": 1})
    assert implicit not in req.requested_resources


def test_get_available_resources_per_node():
    d_id = DeploymentID("a", "b")

    cluster_node_info_cache = MockClusterNodeInfoCache()
    cluster_node_info_cache.add_node(
        "node1", {"GPU": 10, "CPU": 32, "memory": 1024, "customx": 1}
    )

    scheduler = default_impl.create_deployment_scheduler(
        cluster_node_info_cache,
        head_node_id_override="fake-head-node-id",
        create_placement_group_fn_override=None,
    )
    scheduler.on_deployment_created(d_id, SpreadDeploymentSchedulingPolicy())
    scheduler.on_deployment_deployed(
        d_id,
        ReplicaConfig.create(
            dummy,
            ray_actor_options={
                "num_gpus": 1,
                "num_cpus": 3,
                "resources": {"customx": 0.1},
            },
            max_replicas_per_node=4,
        ),
    )

    # Without updating cluster node info cache, when a replica is marked
    # as launching, the resources it uses should decrease the scheduler's
    # view of current available resources per node in the cluster
    scheduler._on_replica_launching(
        ReplicaID(unique_id="replica0", deployment_id=d_id), target_node_id="node1"
    )
    assert scheduler._get_available_resources_per_node().get(
        "node1"
    ) == AvailableNodeResources(
        **{
            "GPU": 9,
            "CPU": 29,
            "memory": 1024,
            "customx": 0.9,
            f"{ray._raylet.IMPLICIT_RESOURCE_PREFIX}b:a": 0.75,
        }
    )

    # Similarly when a replica is marked as running, the resources it
    # uses should decrease current available resources per node
    scheduler.on_replica_running(
        ReplicaID(unique_id="replica1", deployment_id=d_id), node_id="node1"
    )
    assert scheduler._get_available_resources_per_node().get(
        "node1"
    ) == AvailableNodeResources(
        **{
            "GPU": 8,
            "CPU": 26,
            "memory": 1024,
            "customx": 0.8,
            f"{ray._raylet.IMPLICIT_RESOURCE_PREFIX}b:a": 0.5,
        }
    )

    # Get updated info from GCS that available MEMORY has dropped,
    # the decreased memory should reflect in current available resources
    # per node, while also keeping track of the CPU, GPU, custom resources
    # used by launching and running replicas
    cluster_node_info_cache.set_available_resources_per_node(
        "node1", {"GPU": 10, "CPU": 32, "memory": 256, "customx": 1}
    )
    assert scheduler._get_available_resources_per_node().get(
        "node1"
    ) == AvailableNodeResources(
        **{
            "GPU": 8,
            "CPU": 26,
            "memory": 256,
            "customx": 0.8,
            f"{ray._raylet.IMPLICIT_RESOURCE_PREFIX}b:a": 0.5,
        }
    )


def test_get_node_to_running_replicas():
    """Test DeploymentScheduler._get_node_to_running_replicas()."""

    d_id = DeploymentID("a", "b")
    scheduler = default_impl.create_deployment_scheduler(
        MockClusterNodeInfoCache(),
        head_node_id_override="fake-head-node-id",
        create_placement_group_fn_override=None,
    )
    scheduler.on_deployment_created(d_id, SpreadDeploymentSchedulingPolicy())
    scheduler.on_deployment_deployed(d_id, rconfig())

    # Test simple fixed case
    scheduler.on_replica_running(ReplicaID("r1", d_id), "node1")
    scheduler.on_replica_running(ReplicaID("r2", d_id), "node1")
    scheduler.on_replica_running(ReplicaID("r3", d_id), "node2")
    assert scheduler._get_node_to_running_replicas() == {
        "node1": {ReplicaID("r1", d_id), ReplicaID("r2", d_id)},
        "node2": {ReplicaID("r3", d_id)},
    }
    scheduler.on_replica_stopping(ReplicaID("r1", d_id))
    scheduler.on_replica_stopping(ReplicaID("r2", d_id))
    scheduler.on_replica_stopping(ReplicaID("r3", d_id))

    # Test random case
    node_to_running_replicas = defaultdict(set)
    for i in range(40):
        node_id = f"node{random.randint(0,5)}"
        r_id = ReplicaID(f"r{i}", d_id)
        node_to_running_replicas[node_id].add(r_id)
        scheduler.on_replica_running(r_id, node_id)
    assert scheduler._get_node_to_running_replicas() == node_to_running_replicas


def test_get_available_resources_per_node_pg():
    """Test DeploymentScheduler._get_available_resources_per_node()."""

    d_id = DeploymentID("a", "b")

    cluster_node_info_cache = MockClusterNodeInfoCache()
    cluster_node_info_cache.add_node(
        "node1", {"GPU": 10, "CPU": 32, "memory": 1024, "customx": 1}
    )

    scheduler = default_impl.create_deployment_scheduler(
        cluster_node_info_cache,
        head_node_id_override="fake-head-node-id",
        create_placement_group_fn_override=None,
    )
    scheduler.on_deployment_created(d_id, SpreadDeploymentSchedulingPolicy())
    scheduler.on_deployment_deployed(
        d_id,
        ReplicaConfig.create(
            dummy,
            ray_actor_options={"num_cpus": 0},
            placement_group_bundles=[{"GPU": 1}, {"CPU": 3}, {"customx": 0.1}],
            placement_group_strategy="STRICT_PACK",
        ),
    )

    # Without updating cluster node info cache, when a replica is marked
    # as launching, the resources it uses should decrease the scheduler's
    # view of current available resources per node in the cluster
    scheduler._on_replica_launching(
        ReplicaID(unique_id="replica0", deployment_id=d_id), target_node_id="node1"
    )
    assert scheduler._get_available_resources_per_node().get(
        "node1"
    ) == AvailableNodeResources(
        **{
            "GPU": 9,
            "CPU": 29,
            "memory": 1024,
            "customx": 0.9,
        }
    )

    # Similarly when a replica is marked as running, the resources it
    # uses should decrease current available resources per node
    scheduler.on_replica_running(
        ReplicaID(unique_id="replica1", deployment_id=d_id), node_id="node1"
    )
    assert scheduler._get_available_resources_per_node().get(
        "node1"
    ) == AvailableNodeResources(
        **{
            "GPU": 8,
            "CPU": 26,
            "memory": 1024,
            "customx": 0.8,
        }
    )

    # Get updated info from GCS that available MEMORY has dropped,
    # the decreased memory should reflect in current available resources
    # per node, while also keeping track of the CPU, GPU, custom resources
    # used by launching and running replicas
    cluster_node_info_cache.set_available_resources_per_node(
        "node1", {"GPU": 10, "CPU": 32, "memory": 256, "customx": 1}
    )
    assert scheduler._get_available_resources_per_node().get(
        "node1"
    ) == AvailableNodeResources(
        **{
            "GPU": 8,
            "CPU": 26,
            "memory": 256,
            "customx": 0.8,
        }
    )


def test_best_fit_node():
    """Test DeploymentScheduler._best_fit_node()."""

    scheduler = default_impl.create_deployment_scheduler(
        MockClusterNodeInfoCache(),
        head_node_id_override="fake-head-node-id",
        create_placement_group_fn_override=None,
    )

    # None of the nodes can schedule the replica
    assert (
        scheduler._best_fit_node(
            required_resources=RequestedResources(GPU=1, CPU=1, customx=0.1),
            available_resources={
                "node1": AvailableNodeResources(GPU=3, CPU=3),
                "node2": AvailableNodeResources(CPU=3, customx=1),
            },
        )
        is None
    )

    # Only node2 can fit the replica
    assert "node2" == scheduler._best_fit_node(
        required_resources=RequestedResources(GPU=1, CPU=1, customx=0.1),
        available_resources={
            "node1": AvailableNodeResources(CPU=3),
            "node2": AvailableNodeResources(GPU=1, CPU=3, customx=1),
            "node3": AvailableNodeResources(CPU=3, customx=1),
        },
    )

    # We should prioritize minimizing fragementation of GPUs over CPUs
    assert "node1" == scheduler._best_fit_node(
        required_resources=RequestedResources(GPU=1, CPU=1, customx=0.1),
        available_resources={
            "node1": AvailableNodeResources(GPU=2, CPU=10, customx=1),
            "node2": AvailableNodeResources(GPU=10, CPU=2, customx=1),
        },
    )

    # When GPU is the same, should prioritize minimizing fragmentation
    # of CPUs over customer resources
    assert "node2" == scheduler._best_fit_node(
        required_resources=RequestedResources(GPU=1, CPU=1, customx=0.1),
        available_resources={
            "node1": AvailableNodeResources(GPU=10, CPU=5, customx=0.1),
            "node2": AvailableNodeResources(GPU=10, CPU=2, customx=10),
        },
    )

    # Custom resource prioritization: customx is more important than customy
    with mock.patch(
        "ray.serve._private.deployment_scheduler.RAY_SERVE_HIGH_PRIORITY_CUSTOM_RESOURCES",
        "customx,customy",
    ):
        original = Resources.CUSTOM_PRIORITY
        Resources.CUSTOM_PRIORITY = ["customx", "customy"]

        assert "node2" == scheduler._best_fit_node(
            required_resources=RequestedResources(customx=1, customy=1),
            available_resources={
                "node1": AvailableNodeResources(customx=2, customy=5),
                "node2": AvailableNodeResources(customx=2, customy=1),
            },
        )

        # If customx and customy are equal, GPU should determine best fit
        assert "node2" == scheduler._best_fit_node(
            required_resources=RequestedResources(customx=1, customy=1, GPU=1),
            available_resources={
                "node1": AvailableNodeResources(customx=2, customy=2, GPU=10),
                "node2": AvailableNodeResources(customx=2, customy=2, GPU=2),
            },
        )

        # restore
        Resources.CUSTOM_PRIORITY = original


def test_schedule_replica():
    """Test DeploymentScheduler._schedule_replica()"""

    d_id = DeploymentID("deployment1", "app1")
    cluster_node_info_cache = MockClusterNodeInfoCache()
    scheduler = default_impl.create_deployment_scheduler(
        cluster_node_info_cache,
        head_node_id_override="fake-head-node-id",
        create_placement_group_fn_override=lambda request: MockPlacementGroup(request),
    )

    scheduler.on_deployment_created(d_id, SpreadDeploymentSchedulingPolicy())
    scheduler.on_deployment_deployed(d_id, rconfig(ray_actor_options={"num_cpus": 1}))

    scheduling_strategy = None

    def set_scheduling_strategy(actor_handle, placement_group):
        nonlocal scheduling_strategy
        scheduling_strategy = actor_handle._options["scheduling_strategy"]

    # Placement group without target node id
    r0_id = ReplicaID(unique_id="r0", deployment_id=d_id)
    scheduling_request = ReplicaSchedulingRequest(
        replica_id=r0_id,
        actor_def=MockActorClass(),
        actor_resources={"CPU": 1},
        placement_group_bundles=[{"CPU": 1}, {"CPU": 1}],
        placement_group_strategy="STRICT_PACK",
        actor_options={"name": "r0"},
        actor_init_args=(),
        on_scheduled=set_scheduling_strategy,
    )
    scheduler._pending_replicas[d_id][r0_id] = scheduling_request
    scheduler._schedule_replica(
        scheduling_request=scheduling_request,
        default_scheduling_strategy="some_default",
        target_node_id=None,
        target_labels={"abc": In("xyz")},
    )
    assert isinstance(scheduling_strategy, PlacementGroupSchedulingStrategy)
    assert len(scheduler._launching_replicas[d_id]) == 1
    assert not scheduler._launching_replicas[d_id][r0_id].target_labels
    assert not scheduler._launching_replicas[d_id][r0_id].target_node_id

    # Placement group with target node id
    r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)
    scheduling_request = ReplicaSchedulingRequest(
        replica_id=r1_id,
        actor_def=MockActorClass(),
        actor_resources={"CPU": 1},
        placement_group_bundles=[{"CPU": 1}, {"CPU": 1}],
        placement_group_strategy="STRICT_PACK",
        actor_options={"name": "r1"},
        actor_init_args=(),
        on_scheduled=set_scheduling_strategy,
    )
    scheduler._pending_replicas[d_id][r1_id] = scheduling_request
    node_id_1 = NodeID.from_random().hex()
    scheduler._schedule_replica(
        scheduling_request=scheduling_request,
        default_scheduling_strategy="some_default",
        target_node_id=node_id_1,
        target_labels={"abc": In("xyz")},  # this should get ignored
    )
    assert isinstance(scheduling_strategy, PlacementGroupSchedulingStrategy)
    assert len(scheduler._launching_replicas[d_id]) == 2
    assert not scheduler._launching_replicas[d_id][r1_id].target_labels
    assert scheduler._launching_replicas[d_id][r1_id].target_node_id == node_id_1

    # Target node id without placement group
    r2_id = ReplicaID(unique_id="r2", deployment_id=d_id)
    scheduling_request = ReplicaSchedulingRequest(
        replica_id=r2_id,
        actor_def=MockActorClass(),
        actor_resources={"CPU": 1},
        actor_options={"name": "r2"},
        actor_init_args=(),
        on_scheduled=set_scheduling_strategy,
    )
    scheduler._pending_replicas[d_id][r2_id] = scheduling_request
    scheduler._schedule_replica(
        scheduling_request=scheduling_request,
        default_scheduling_strategy="some_default",
        target_node_id=node_id_1,
        target_labels={"abc": In("xyz")},  # this should get ignored
    )
    assert isinstance(scheduling_strategy, NodeAffinitySchedulingStrategy)
    assert scheduling_strategy.node_id == node_id_1
    assert len(scheduler._launching_replicas[d_id]) == 3
    assert not scheduler._launching_replicas[d_id][r2_id].target_labels
    assert scheduler._launching_replicas[d_id][r2_id].target_node_id == node_id_1

    # Target labels
    r3_id = ReplicaID(unique_id="r3", deployment_id=d_id)
    scheduling_request = ReplicaSchedulingRequest(
        replica_id=r3_id,
        actor_def=MockActorClass(),
        actor_resources={"CPU": 1},
        actor_options={"name": "r3"},
        actor_init_args=(),
        on_scheduled=set_scheduling_strategy,
    )
    scheduler._pending_replicas[d_id][r3_id] = scheduling_request
    scheduler._schedule_replica(
        scheduling_request=scheduling_request,
        default_scheduling_strategy="some_default",
        target_node_id=None,
        target_labels={"abc": In("xyz")},
    )
    assert isinstance(scheduling_strategy, NodeLabelSchedulingStrategy)
    assert scheduling_strategy.soft
    assert len(scheduler._launching_replicas[d_id]) == 4
    assert not scheduler._launching_replicas[d_id][r3_id].target_node_id
    assert len(scheduler._launching_replicas[d_id][r3_id].target_labels.keys()) == 1
    operator = scheduler._launching_replicas[d_id][r3_id].target_labels["abc"]
    assert isinstance(operator, In) and operator.values == ["xyz"]

    # internal implicit resource with max_replicas_per_node
    r4_id = ReplicaID(unique_id="r4", deployment_id=d_id)
    scheduling_request = ReplicaSchedulingRequest(
        replica_id=r4_id,
        actor_def=MockActorClass(),
        actor_resources={"my_rs": 1, "CPU": 1},
        placement_group_bundles=None,
        placement_group_strategy=None,
        actor_options={"name": "r4", "num_cpus": 1, "resources": {"my_rs": 1}},
        actor_init_args=(),
        on_scheduled=set_scheduling_strategy,
        max_replicas_per_node=10,
    )
    scheduler._pending_replicas[d_id][r4_id] = scheduling_request
    scheduler._schedule_replica(
        scheduling_request=scheduling_request,
        default_scheduling_strategy="some_default",
        target_node_id=None,
        target_labels=None,
    )
    assert scheduling_strategy == "some_default"
    assert len(scheduler._launching_replicas[d_id]) == 5
    assert scheduling_request.actor_options == {
        "name": "r4",
        "num_cpus": 1,
        "resources": {"my_rs": 1},
    }


def test_downscale_multiple_deployments():
    """Test to make sure downscale prefers replicas without node id
    and then replicas on a node with fewest replicas of all deployments.
    """

    cluster_node_info_cache = MockClusterNodeInfoCache()
    scheduler = default_impl.create_deployment_scheduler(
        cluster_node_info_cache,
        head_node_id_override="fake-head-node-id",
        create_placement_group_fn_override=None,
    )

    d1_id = DeploymentID(name="deployment1")
    d2_id = DeploymentID(name="deployment2")
    d1_r1_id = ReplicaID(
        unique_id="replica1",
        deployment_id=d1_id,
    )
    d1_r2_id = ReplicaID(
        unique_id="replica2",
        deployment_id=d1_id,
    )
    d1_r3_id = ReplicaID(
        unique_id="replica3",
        deployment_id=d1_id,
    )
    d2_r1_id = ReplicaID(
        unique_id="replica1",
        deployment_id=d2_id,
    )
    d2_r2_id = ReplicaID(
        unique_id="replica2",
        deployment_id=d2_id,
    )
    d2_r3_id = ReplicaID(
        unique_id="replica3",
        deployment_id=d2_id,
    )
    d2_r4_id = ReplicaID(
        unique_id="replica4",
        deployment_id=d2_id,
    )
    scheduler.on_deployment_created(d1_id, SpreadDeploymentSchedulingPolicy())
    scheduler.on_deployment_created(d2_id, SpreadDeploymentSchedulingPolicy())
    scheduler.on_replica_running(d1_r1_id, "node1")
    scheduler.on_replica_running(d1_r2_id, "node2")
    scheduler.on_replica_running(d1_r3_id, "node2")
    scheduler.on_replica_running(d2_r1_id, "node1")
    scheduler.on_replica_running(d2_r2_id, "node2")
    scheduler.on_replica_running(d2_r3_id, "node1")
    scheduler.on_replica_running(d2_r4_id, "node1")
    deployment_to_replicas_to_stop = scheduler.schedule(
        upscales={},
        downscales={
            d1_id: DeploymentDownscaleRequest(deployment_id=d1_id, num_to_stop=1)
        },
    )
    assert len(deployment_to_replicas_to_stop) == 1
    # Even though node1 has fewest replicas of deployment1
    # but it has more replicas of all deployments so
    # we should stop replicas from node2.
    assert len(deployment_to_replicas_to_stop[d1_id]) == 1
    assert deployment_to_replicas_to_stop[d1_id].issubset({d1_r2_id, d1_r3_id})

    scheduler.on_replica_stopping(d1_r3_id)
    scheduler.on_replica_stopping(d2_r3_id)
    scheduler.on_replica_stopping(d2_r4_id)

    deployment_to_replicas_to_stop = scheduler.schedule(
        upscales={},
        downscales={
            d1_id: DeploymentDownscaleRequest(deployment_id=d1_id, num_to_stop=1),
            d2_id: DeploymentDownscaleRequest(deployment_id=d2_id, num_to_stop=1),
        },
    )
    assert len(deployment_to_replicas_to_stop) == 2
    # We should stop replicas from the same node.
    assert len(deployment_to_replicas_to_stop[d1_id]) == 1
    assert {r.unique_id for r in deployment_to_replicas_to_stop[d1_id]} == {
        r.unique_id for r in deployment_to_replicas_to_stop[d2_id]
    }

    scheduler.on_replica_stopping(d1_r1_id)
    scheduler.on_replica_stopping(d1_r2_id)
    scheduler.on_replica_stopping(d2_r1_id)
    scheduler.on_replica_stopping(d2_r2_id)
    scheduler.on_deployment_deleted(d1_id)
    scheduler.on_deployment_deleted(d2_id)


def test_downscale_head_node():
    """Test to make sure downscale deprioritizes replicas on the head node."""

    head_node_id = "fake-head-node-id"
    dep_id = DeploymentID(name="deployment1")
    cluster_node_info_cache = MockClusterNodeInfoCache()
    scheduler = default_impl.create_deployment_scheduler(
        cluster_node_info_cache,
        head_node_id_override=head_node_id,
        create_placement_group_fn_override=None,
    )

    r1_id = ReplicaID(
        unique_id="replica1",
        deployment_id=dep_id,
    )
    r2_id = ReplicaID(
        unique_id="replica2",
        deployment_id=dep_id,
    )
    r3_id = ReplicaID(
        unique_id="replica3",
        deployment_id=dep_id,
    )
    scheduler.on_deployment_created(dep_id, SpreadDeploymentSchedulingPolicy())
    scheduler.on_replica_running(r1_id, head_node_id)
    scheduler.on_replica_running(r2_id, "node2")
    scheduler.on_replica_running(r3_id, "node2")
    deployment_to_replicas_to_stop = scheduler.schedule(
        upscales={},
        downscales={
            dep_id: DeploymentDownscaleRequest(deployment_id=dep_id, num_to_stop=1)
        },
    )
    assert len(deployment_to_replicas_to_stop) == 1
    assert deployment_to_replicas_to_stop[dep_id].issubset({r2_id, r3_id})
    scheduler.on_replica_stopping(deployment_to_replicas_to_stop[dep_id].pop())

    deployment_to_replicas_to_stop = scheduler.schedule(
        upscales={},
        downscales={
            dep_id: DeploymentDownscaleRequest(deployment_id=dep_id, num_to_stop=1)
        },
    )
    assert len(deployment_to_replicas_to_stop) == 1
    assert deployment_to_replicas_to_stop[dep_id] < {r2_id, r3_id}
    scheduler.on_replica_stopping(deployment_to_replicas_to_stop[dep_id].pop())

    deployment_to_replicas_to_stop = scheduler.schedule(
        upscales={},
        downscales={
            dep_id: DeploymentDownscaleRequest(deployment_id=dep_id, num_to_stop=1)
        },
    )
    assert len(deployment_to_replicas_to_stop) == 1
    assert deployment_to_replicas_to_stop[dep_id] == {r1_id}
    scheduler.on_replica_stopping(r1_id)
    scheduler.on_deployment_deleted(dep_id)


def test_downscale_single_deployment():
    """Test to make sure downscale prefers replicas without node id
    and then replicas on a node with fewest replicas of all deployments.
    """

    dep_id = DeploymentID(name="deployment1")
    cluster_node_info_cache = MockClusterNodeInfoCache()
    cluster_node_info_cache.add_node("node1")
    cluster_node_info_cache.add_node("node2")
    scheduler = default_impl.create_deployment_scheduler(
        cluster_node_info_cache,
        head_node_id_override="fake-head-node-id",
        create_placement_group_fn_override=None,
    )

    scheduler.on_deployment_created(dep_id, SpreadDeploymentSchedulingPolicy())
    scheduler.on_deployment_deployed(
        dep_id, ReplicaConfig.create(lambda x: x, ray_actor_options={"num_cpus": 0})
    )
    r1_id = ReplicaID(
        unique_id="replica1",
        deployment_id=dep_id,
    )
    r2_id = ReplicaID(
        unique_id="replica2",
        deployment_id=dep_id,
    )
    r3_id = ReplicaID(
        unique_id="replica3",
        deployment_id=dep_id,
    )
    r4_id = ReplicaID(
        unique_id="replica4",
        deployment_id=dep_id,
    )
    scheduler.on_replica_running(r1_id, "node1")
    scheduler.on_replica_running(r2_id, "node1")
    scheduler.on_replica_running(r3_id, "node2")
    scheduler.on_replica_recovering(r4_id)
    deployment_to_replicas_to_stop = scheduler.schedule(
        upscales={},
        downscales={
            dep_id: DeploymentDownscaleRequest(deployment_id=dep_id, num_to_stop=1)
        },
    )
    assert len(deployment_to_replicas_to_stop) == 1
    # Prefer replica without node id
    assert deployment_to_replicas_to_stop[dep_id] == {r4_id}
    scheduler.on_replica_stopping(r4_id)

    r5_id = ReplicaID(
        unique_id="replica5",
        deployment_id=dep_id,
    )
    deployment_to_replicas_to_stop = scheduler.schedule(
        upscales={
            dep_id: [
                ReplicaSchedulingRequest(
                    replica_id=r5_id,
                    actor_def=Mock(),
                    actor_resources={"CPU": 1},
                    actor_options={},
                    actor_init_args=(),
                    on_scheduled=lambda actor_handle, placement_group: actor_handle,
                ),
            ]
        },
        downscales={},
    )
    assert not deployment_to_replicas_to_stop
    deployment_to_replicas_to_stop = scheduler.schedule(
        upscales={},
        downscales={
            dep_id: DeploymentDownscaleRequest(deployment_id=dep_id, num_to_stop=1)
        },
    )
    assert len(deployment_to_replicas_to_stop) == 1
    # Prefer replica without node id
    assert deployment_to_replicas_to_stop[dep_id] == {r5_id}
    scheduler.on_replica_stopping(r5_id)

    deployment_to_replicas_to_stop = scheduler.schedule(
        upscales={},
        downscales={
            dep_id: DeploymentDownscaleRequest(deployment_id=dep_id, num_to_stop=1)
        },
    )
    assert len(deployment_to_replicas_to_stop) == 1
    # Prefer replica on a node with fewest replicas of all deployments.
    assert deployment_to_replicas_to_stop[dep_id] == {r3_id}
    scheduler.on_replica_stopping(r3_id)

    deployment_to_replicas_to_stop = scheduler.schedule(
        upscales={},
        downscales={
            dep_id: DeploymentDownscaleRequest(deployment_id=dep_id, num_to_stop=2)
        },
    )
    assert len(deployment_to_replicas_to_stop) == 1
    assert deployment_to_replicas_to_stop[dep_id] <= {r1_id, r2_id}
    scheduler.on_replica_stopping(r1_id)
    scheduler.on_replica_stopping(r2_id)
    scheduler.on_deployment_deleted(dep_id)


def test_schedule_passes_placement_group_options():
    """Test that bundle_label_selector is passed to CreatePlacementGroupRequest."""
    cluster_node_info_cache = MockClusterNodeInfoCache()
    captured_requests = []

    def mock_create_pg(request):
        captured_requests.append(request)

        class MockPG:
            def wait(self, *args):
                return True

        return MockPG()

    scheduler = default_impl.create_deployment_scheduler(
        cluster_node_info_cache,
        head_node_id_override="fake-head-node-id",
        create_placement_group_fn_override=mock_create_pg,
    )

    dep_id = DeploymentID(name="pg_options_test")
    # Use Spread policy here, but the logic is shared across policies.
    scheduler.on_deployment_created(dep_id, SpreadDeploymentSchedulingPolicy())

    test_labels = [{"region": "us-west"}]
    # Create a request with the new options
    req = ReplicaSchedulingRequest(
        replica_id=ReplicaID("r1", dep_id),
        actor_def=MockActorClass(),
        actor_resources={"CPU": 1},
        actor_options={"name": "r1"},
        actor_init_args=(),
        on_scheduled=lambda *args, **kwargs: None,
        placement_group_bundles=[{"CPU": 1}],
        placement_group_bundle_label_selector=test_labels,
        placement_group_strategy="STRICT_PACK",
    )

    scheduler.schedule(upscales={dep_id: [req]}, downscales={})

    # Verify the PlacementGroupSchedulingRequest is created.
    assert len(captured_requests) == 1
    pg_request = captured_requests[0]

    # bundle_label_selector should be passed to request.
    assert pg_request.bundle_label_selector == test_labels


def test_schedule_pins_actor_to_bundle_0():
    """Replicas with a placement group are scheduled with placement_group_bundle_index=0."""
    cluster_node_info_cache = MockClusterNodeInfoCache()
    scheduler = default_impl.create_deployment_scheduler(
        cluster_node_info_cache,
        head_node_id_override="fake-head-node-id",
        create_placement_group_fn_override=lambda request: MockPlacementGroup(request),
    )

    dep_id = DeploymentID(name="pin_test")
    scheduler.on_deployment_created(dep_id, SpreadDeploymentSchedulingPolicy())

    captured_handles = []
    req = ReplicaSchedulingRequest(
        replica_id=ReplicaID("r1", dep_id),
        actor_def=MockActorClass(),
        actor_resources={"CPU": 1},
        actor_options={"name": "r1"},
        actor_init_args=(),
        on_scheduled=lambda handle, **kwargs: captured_handles.append(handle),
        placement_group_bundles=[{"CPU": 1, "GPU": 1}, {"CPU": 1, "GPU": 1}],
        placement_group_strategy="PACK",
    )
    scheduler.schedule(upscales={dep_id: [req]}, downscales={})

    assert len(captured_handles) == 1
    strategy = captured_handles[0]._options["scheduling_strategy"]
    assert isinstance(strategy, PlacementGroupSchedulingStrategy)
    assert strategy.placement_group_bundle_index == 0


def test_filter_nodes_by_label_selector():
    """Test _filter_nodes_by_label_selector logic used by _find_best_fit_node_for_pack
    when bin-packing, such that label constraints are enforced for the preferred node."""

    class MockScheduler(default_impl.DefaultDeploymentScheduler):
        def __init__(self):
            pass

    scheduler = MockScheduler()

    nodes = {
        "n1": AvailableNodeResources(),
        "n2": AvailableNodeResources(),
        "n3": AvailableNodeResources(),
    }
    node_labels = {
        "n1": {"region": "us-west", "gpu": "T4", "env": "prod"},
        "n2": {"region": "us-east", "gpu": "A100", "env": "dev"},
        "n3": {"region": "me-central", "env": "staging"},  # No GPU label
    }

    # equals operator
    filtered = scheduler._filter_nodes_by_label_selector(
        nodes, {"region": "us-west"}, node_labels
    )
    assert set(filtered.keys()) == {"n1"}

    # not equals operator
    filtered = scheduler._filter_nodes_by_label_selector(
        nodes, {"region": "!us-west"}, node_labels
    )
    assert set(filtered.keys()) == {"n2", "n3"}

    # in operator
    filtered = scheduler._filter_nodes_by_label_selector(
        nodes, {"region": "in(us-west,us-east)"}, node_labels
    )
    assert set(filtered.keys()) == {"n1", "n2"}

    # !in operator
    filtered = scheduler._filter_nodes_by_label_selector(
        nodes, {"env": "!in(dev,staging)"}, node_labels
    )
    assert set(filtered.keys()) == {"n1"}

    # Missing labels treated as not a match for equality.
    filtered = scheduler._filter_nodes_by_label_selector(
        nodes, {"gpu": "A100"}, node_labels
    )
    assert set(filtered.keys()) == {"n2"}

    # Not equal should match node with missing labels.
    filtered = scheduler._filter_nodes_by_label_selector(
        nodes, {"gpu": "!T4"}, node_labels
    )
    assert set(filtered.keys()) == {"n2", "n3"}


def test_build_pack_placement_candidates():
    """Test strategy generation logic in DefaultDeploymentScheduler._build_pack_placement_candidates,
    verifying that the scheduler correctly generates a list of (resources, labels) tuples to
    attempt for scheduling."""

    # Setup scheduler with mocks
    cluster_node_info_cache = MockClusterNodeInfoCache()
    scheduler = default_impl.create_deployment_scheduler(
        cluster_node_info_cache,
        head_node_id_override="head_node",
        create_placement_group_fn_override=None,
    )

    # Basic Ray Actor
    req_basic = ReplicaSchedulingRequest(
        replica_id=ReplicaID("r1", DeploymentID(name="d1")),
        actor_def=MockActorClass(),
        actor_resources={"CPU": 1},
        actor_options={},
        actor_init_args=(),
        on_scheduled=Mock(),
    )
    strategies = scheduler._build_pack_placement_candidates(req_basic)
    assert len(strategies) == 1
    assert strategies[0][0] == {"CPU": 1}
    assert strategies[0][1] == []

    # Actor with label_selector and fallback_strategy
    req_fallback = ReplicaSchedulingRequest(
        replica_id=ReplicaID("r2", DeploymentID(name="d1")),
        actor_def=MockActorClass(),
        actor_resources={"CPU": 1},
        actor_options={
            "label_selector": {"region": "us-west"},
            "fallback_strategy": [{"label_selector": {"region": "us-east"}}],
        },
        actor_init_args=(),
        on_scheduled=Mock(),
    )
    strategies = scheduler._build_pack_placement_candidates(req_fallback)
    assert len(strategies) == 2

    assert strategies[0][0] == {"CPU": 1}
    assert strategies[0][1] == [{"region": "us-west"}]
    assert strategies[1][0] == {"CPU": 1}
    assert strategies[1][1] == [{"region": "us-east"}]

    # Scheduling replica with placement group PACK strategy and bundle_label_selector
    req_pack = ReplicaSchedulingRequest(
        replica_id=ReplicaID("r4", DeploymentID(name="d1")),
        actor_def=MockActorClass(),
        actor_resources={"CPU": 0.1},
        actor_options={},
        actor_init_args=(),
        on_scheduled=Mock(),
        placement_group_bundles=[{"CPU": 5}],
        placement_group_strategy="PACK",
        placement_group_bundle_label_selector=[
            {"accelerator-type": "H100"},
            {"accelerator-type": "H100"},
        ],
    )

    with pytest.raises(NotImplementedError):
        scheduler._build_pack_placement_candidates(req_pack)

    # Scheduling replica with placement group STRICT_PACK strategy and bundle_label_selector
    req_pg = ReplicaSchedulingRequest(
        replica_id=ReplicaID("r3", DeploymentID(name="d1")),
        actor_def=MockActorClass(),
        actor_resources={},
        actor_options={},
        actor_init_args=(),
        on_scheduled=Mock(),
        placement_group_bundles=[{"CPU": 2}],
        placement_group_strategy="STRICT_PACK",
        placement_group_bundle_label_selector=[{"accelerator-type": "A100"}],
    )
    strategies = scheduler._build_pack_placement_candidates(req_pg)
    assert len(strategies) == 1

    assert strategies[0][0] == {"CPU": 2}
    assert strategies[0][1] == [{"accelerator-type": "A100"}]


def test_build_pack_placement_candidates_pg_fallback_error():
    """
    Test that providing placement_group_fallback_strategy raises NotImplementedError.
    """
    cluster_node_info_cache = MockClusterNodeInfoCache()
    scheduler = default_impl.create_deployment_scheduler(
        cluster_node_info_cache,
        head_node_id_override="head_node",
        create_placement_group_fn_override=None,
    )

    # Create a request with placement_group_fallback_strategy defined.
    req = ReplicaSchedulingRequest(
        replica_id=ReplicaID("r1", DeploymentID(name="d1")),
        actor_def=MockActorClass(),
        actor_resources={},
        actor_options={},
        actor_init_args=(),
        on_scheduled=Mock(),
        placement_group_bundles=[{"CPU": 1}],
        placement_group_strategy="STRICT_PACK",
        # Raises NotImplementedError since not added to placement group options yet.
        placement_group_fallback_strategy=[{"label_selector": {"zone": "us-east-1a"}}],
    )

    # Verify the scheduler raises the expected error
    with pytest.raises(NotImplementedError, match="not yet supported"):
        scheduler._build_pack_placement_candidates(req)


@pytest.mark.skipif(
    not RAY_SERVE_USE_PACK_SCHEDULING_STRATEGY, reason="Needs pack strategy."
)
class TestPackScheduling:
    def test_basic(self):
        d_id1 = DeploymentID(name="deployment1")
        d_id2 = DeploymentID(name="deployment2")
        node_id_1 = NodeID.from_random().hex()
        node_id_2 = NodeID.from_random().hex()

        cluster_node_info_cache = MockClusterNodeInfoCache()
        cluster_node_info_cache.add_node(node_id_1, {"CPU": 3})
        cluster_node_info_cache.add_node(node_id_2, {"CPU": 2})
        scheduler = default_impl.create_deployment_scheduler(
            cluster_node_info_cache,
            head_node_id_override="fake-head-node-id",
            create_placement_group_fn_override=None,
        )

        scheduler.on_deployment_created(d_id1, SpreadDeploymentSchedulingPolicy())
        scheduler.on_deployment_created(d_id2, SpreadDeploymentSchedulingPolicy())

        scheduler.on_deployment_deployed(
            d_id1,
            ReplicaConfig.create(dummy, ray_actor_options={"num_cpus": 1}),
        )
        scheduler.on_deployment_deployed(
            d_id2,
            ReplicaConfig.create(dummy, ray_actor_options={"num_cpus": 3}),
        )

        on_scheduled_mock = Mock()
        on_scheduled_mock2 = Mock()
        scheduler.schedule(
            upscales={
                d_id1: [
                    ReplicaSchedulingRequest(
                        replica_id=ReplicaID(unique_id=f"r{i}", deployment_id=d_id1),
                        actor_def=MockActorClass(),
                        actor_resources={"CPU": 1},
                        actor_options={},
                        actor_init_args=(),
                        on_scheduled=on_scheduled_mock,
                    )
                    for i in range(2)
                ],
                d_id2: [
                    ReplicaSchedulingRequest(
                        replica_id=ReplicaID(unique_id="r2", deployment_id=d_id2),
                        actor_def=MockActorClass(),
                        actor_resources={"CPU": 3},
                        actor_options={},
                        actor_init_args=(),
                        on_scheduled=on_scheduled_mock2,
                    )
                ],
            },
            downscales={},
        )

        assert len(on_scheduled_mock.call_args_list) == 2
        for call in on_scheduled_mock.call_args_list:
            assert call.kwargs == {"placement_group": None}
            assert len(call.args) == 1
            scheduling_strategy = call.args[0]._options["scheduling_strategy"]
            assert isinstance(scheduling_strategy, NodeAffinitySchedulingStrategy)
            assert scheduling_strategy.node_id == node_id_2

        assert len(on_scheduled_mock2.call_args_list) == 1
        call = on_scheduled_mock2.call_args_list[0]
        assert call.kwargs == {"placement_group": None}
        assert len(call.args) == 1
        scheduling_strategy = call.args[0]._options["scheduling_strategy"]
        assert isinstance(scheduling_strategy, NodeAffinitySchedulingStrategy)
        assert scheduling_strategy.node_id == node_id_1

    def test_placement_groups(self):
        d_id1 = DeploymentID(name="deployment1")
        d_id2 = DeploymentID(name="deployment2")

        cluster_node_info_cache = MockClusterNodeInfoCache()
        cluster_node_info_cache.add_node("node1", {"CPU": 3})
        cluster_node_info_cache.add_node("node2", {"CPU": 2})
        scheduler = default_impl.create_deployment_scheduler(
            cluster_node_info_cache,
            head_node_id_override="fake-head-node-id",
            create_placement_group_fn_override=lambda *args, **kwargs: MockPlacementGroup(  # noqa
                *args, **kwargs
            ),
        )

        _ = ray.util.placement_group
        scheduler.on_deployment_created(d_id1, SpreadDeploymentSchedulingPolicy())
        scheduler.on_deployment_created(d_id2, SpreadDeploymentSchedulingPolicy())

        scheduler.on_deployment_deployed(
            d_id1,
            ReplicaConfig.create(
                dummy,
                ray_actor_options={"num_cpus": 0},
                placement_group_bundles=[{"CPU": 0.5}, {"CPU": 0.5}],
                placement_group_strategy="STRICT_PACK",
            ),
        )
        scheduler.on_deployment_deployed(
            d_id2,
            ReplicaConfig.create(
                dummy,
                ray_actor_options={"num_cpus": 0},
                placement_group_bundles=[{"CPU": 0.5}, {"CPU": 2.5}],
                placement_group_strategy="STRICT_PACK",
            ),
        )

        on_scheduled_mock = Mock()
        on_scheduled_mock2 = Mock()
        scheduler.schedule(
            upscales={
                d_id1: [
                    ReplicaSchedulingRequest(
                        replica_id=ReplicaID(unique_id=f"r{i}", deployment_id=d_id1),
                        actor_def=MockActorClass(),
                        actor_resources={"CPU": 0},
                        placement_group_bundles=[{"CPU": 0.5}, {"CPU": 0.5}],
                        placement_group_strategy="STRICT_PACK",
                        actor_options={"name": "random_replica"},
                        actor_init_args=(),
                        on_scheduled=on_scheduled_mock,
                    )
                    for i in range(2)
                ],
                d_id2: [
                    ReplicaSchedulingRequest(
                        replica_id=ReplicaID(unique_id="r2", deployment_id=d_id2),
                        actor_def=MockActorClass(),
                        actor_resources={"CPU": 0},
                        placement_group_bundles=[{"CPU": 0.5}, {"CPU": 2.5}],
                        placement_group_strategy="STRICT_PACK",
                        actor_options={"name": "some_replica"},
                        actor_init_args=(),
                        on_scheduled=on_scheduled_mock2,
                    )
                ],
            },
            downscales={},
        )

        assert len(on_scheduled_mock.call_args_list) == 2
        for call in on_scheduled_mock.call_args_list:
            assert len(call.args) == 1
            scheduling_strategy = call.args[0]._options["scheduling_strategy"]
            assert isinstance(scheduling_strategy, PlacementGroupSchedulingStrategy)
            assert call.kwargs.get("placement_group")._soft_target_node_id == "node2"

        assert len(on_scheduled_mock2.call_args_list) == 1
        call = on_scheduled_mock2.call_args_list[0]
        assert len(call.args) == 1
        scheduling_strategy = call.args[0]._options["scheduling_strategy"]
        assert isinstance(scheduling_strategy, PlacementGroupSchedulingStrategy)
        assert call.kwargs.get("placement_group")._soft_target_node_id == "node1"

    def test_heterogeneous_resources(self):
        d_id1 = DeploymentID(name="deployment1")
        d_id2 = DeploymentID(name="deployment2")
        node_id_1 = NodeID.from_random().hex()
        node_id_2 = NodeID.from_random().hex()

        cluster_node_info_cache = MockClusterNodeInfoCache()
        cluster_node_info_cache.add_node(node_id_1, {"GPU": 4, "CPU": 6})
        cluster_node_info_cache.add_node(node_id_2, {"GPU": 10, "CPU": 2})
        scheduler = default_impl.create_deployment_scheduler(
            cluster_node_info_cache,
            head_node_id_override="fake-head-node-id",
            create_placement_group_fn_override=None,
        )

        scheduler.on_deployment_created(d_id1, SpreadDeploymentSchedulingPolicy())
        scheduler.on_deployment_created(d_id2, SpreadDeploymentSchedulingPolicy())
        scheduler.on_deployment_deployed(
            d_id1,
            ReplicaConfig.create(
                dummy, ray_actor_options={"num_gpus": 2, "num_cpus": 2}
            ),
        )
        scheduler.on_deployment_deployed(
            d_id2,
            ReplicaConfig.create(
                dummy, ray_actor_options={"num_gpus": 1, "num_cpus": 1}
            ),
        )

        on_scheduled_mock = Mock()
        scheduler.schedule(
            upscales={
                d_id1: [
                    ReplicaSchedulingRequest(
                        replica_id=ReplicaID(unique_id="r0", deployment_id=d_id1),
                        actor_def=MockActorClass(),
                        actor_resources={"GPU": 2, "CPU": 2},
                        actor_options={},
                        actor_init_args=(),
                        on_scheduled=on_scheduled_mock,
                    )
                ],
                d_id2: [
                    ReplicaSchedulingRequest(
                        replica_id=ReplicaID(unique_id=f"r{i+1}", deployment_id=d_id2),
                        actor_def=MockActorClass(),
                        actor_resources={"GPU": 1, "CPU": 1},
                        actor_options={},
                        actor_init_args=(),
                        on_scheduled=on_scheduled_mock,
                    )
                    for i in range(2)
                ],
            },
            downscales={},
        )

        # Even though scheduling on node 2 would minimize fragmentation
        # of CPU resources, we should prioritize minimizing fragmentation
        # of GPU resources first, so all 3 replicas should be scheduled
        # to node 1
        assert len(on_scheduled_mock.call_args_list) == 3
        for call in on_scheduled_mock.call_args_list:
            assert len(call.args) == 1
            scheduling_strategy = call.args[0]._options["scheduling_strategy"]
            assert isinstance(scheduling_strategy, NodeAffinitySchedulingStrategy)
            assert scheduling_strategy.node_id == node_id_1
            assert call.kwargs == {"placement_group": None}

    def test_max_replicas_per_node(self):
        """Test that at most `max_replicas_per_node` number of replicas
        are scheduled onto a node even if that node has more resources.
        """

        d_id1 = DeploymentID(name="deployment1")
        node_id_1 = NodeID.from_random().hex()
        node_id_2 = NodeID.from_random().hex()
        cluster_node_info_cache = MockClusterNodeInfoCache()
        # Should try to schedule on node1 to minimize fragmentation
        cluster_node_info_cache.add_node(node_id_1, {"CPU": 20})
        cluster_node_info_cache.add_node(node_id_2, {"CPU": 21})

        scheduler = default_impl.create_deployment_scheduler(
            cluster_node_info_cache,
            head_node_id_override="fake-head-node-id",
            create_placement_group_fn_override=lambda *args, **kwargs: MockPlacementGroup(  # noqa
                *args, **kwargs
            ),
        )
        scheduler.on_deployment_created(d_id1, SpreadDeploymentSchedulingPolicy())
        scheduler.on_deployment_deployed(
            d_id1,
            ReplicaConfig.create(
                dummy, max_replicas_per_node=4, ray_actor_options={"num_cpus": 2}
            ),
        )

        state = defaultdict(int)

        def on_scheduled(actor_handle, placement_group):
            scheduling_strategy = actor_handle._options["scheduling_strategy"]
            if isinstance(scheduling_strategy, NodeAffinitySchedulingStrategy):
                state[scheduling_strategy.node_id] += 1
            elif isinstance(scheduling_strategy, PlacementGroupSchedulingStrategy):
                state[placement_group._soft_target_node_id] += 1

        scheduler.schedule(
            upscales={
                d_id1: [
                    ReplicaSchedulingRequest(
                        replica_id=ReplicaID(
                            unique_id=f"replica{i}", deployment_id=d_id1
                        ),
                        actor_def=MockActorClass(),
                        actor_resources={"CPU": 2},
                        max_replicas_per_node=4,
                        actor_options={"name": "random"},
                        actor_init_args=(),
                        on_scheduled=on_scheduled,
                    )
                    for i in range(5)
                ]
            },
            downscales={},
        )
        assert state[node_id_1] == 4
        assert state[node_id_2] == 1

    def test_heterogeneous_resources_with_max_replicas_per_node(self):
        d_id1 = DeploymentID(name="deployment1")
        d_id2 = DeploymentID(name="deployment2")
        max_replicas_per_node = {d_id1: 2, d_id2: 3}

        cluster_node_info_cache = MockClusterNodeInfoCache()
        node1 = NodeID.from_random().hex()
        node2 = NodeID.from_random().hex()

        cluster_node_info_cache.add_node(node1, {"GPU": 8, "CPU": 32})
        cluster_node_info_cache.add_node(node2, {"GPU": 10, "CPU": 32})
        scheduler = default_impl.create_deployment_scheduler(
            cluster_node_info_cache,
            head_node_id_override="fake-head-node-id",
            create_placement_group_fn_override=None,
        )
        scheduler.on_deployment_created(d_id1, SpreadDeploymentSchedulingPolicy())
        scheduler.on_deployment_created(d_id2, SpreadDeploymentSchedulingPolicy())
        scheduler.on_deployment_deployed(
            d_id1,
            ReplicaConfig.create(
                dummy,
                max_replicas_per_node=max_replicas_per_node[d_id1],
                ray_actor_options={"num_gpus": 2, "num_cpus": 1},
            ),
        )
        scheduler.on_deployment_deployed(
            d_id2,
            ReplicaConfig.create(
                dummy,
                max_replicas_per_node=max_replicas_per_node[d_id2],
                ray_actor_options={"num_gpus": 2, "num_cpus": 1},
            ),
        )

        state = defaultdict(list)

        def on_scheduled(actor_handle, *args, **kwargs):
            scheduling_strategy = actor_handle._options["scheduling_strategy"]
            assert isinstance(scheduling_strategy, NodeAffinitySchedulingStrategy)
            state[scheduling_strategy.node_id].append(
                actor_handle._options["deployment"]
            )

        # Schedule one d1 and one d2
        scheduler.schedule(
            upscales={
                d_id1: [
                    ReplicaSchedulingRequest(
                        replica_id=ReplicaID(unique_id="replica0", deployment_id=d_id1),
                        actor_def=MockActorClass(),
                        actor_resources={"GPU": 2},
                        max_replicas_per_node=max_replicas_per_node[d_id1],
                        actor_options={"name": "random", "deployment": d_id1},
                        actor_init_args=(),
                        on_scheduled=on_scheduled,
                    ),
                ],
                d_id2: [
                    ReplicaSchedulingRequest(
                        replica_id=ReplicaID(unique_id="replica1", deployment_id=d_id2),
                        actor_def=MockActorClass(),
                        actor_resources={"GPU": 2},
                        max_replicas_per_node=max_replicas_per_node[d_id2],
                        actor_options={"name": "random", "deployment": d_id2},
                        actor_init_args=(),
                        on_scheduled=on_scheduled,
                    )
                ],
            },
            downscales={},
        )
        assert state[node1].count(d_id1) == 1
        assert state[node1].count(d_id2) == 1
        assert len(state[node2]) == 0

        # Schedule two more d1
        scheduler.schedule(
            upscales={
                d_id1: [
                    ReplicaSchedulingRequest(
                        replica_id=ReplicaID(
                            unique_id=f"replica{i+2}", deployment_id=d_id1
                        ),
                        actor_def=MockActorClass(),
                        actor_resources={"GPU": 2},
                        max_replicas_per_node=max_replicas_per_node[d_id1],
                        actor_options={"name": "random", "deployment": d_id1},
                        actor_init_args=(),
                        on_scheduled=on_scheduled,
                    )
                    for i in range(2)
                ],
            },
            downscales={},
        )
        # 2 d1 + 1 d2 on node1
        assert state[node1].count(d_id1) == 2
        assert state[node1].count(d_id2) == 1

        # 1 d1 on node2 because of max_replicas_per_node=2 (otherwise node1 could have fit both new d1 replicas)
        assert state[node2].count(d_id1) == 1
        assert state[node2].count(d_id2) == 0

    def test_custom_resources(self):
        d_id = DeploymentID(name="deployment1")
        node_id_1 = NodeID.from_random().hex()
        node_id_2 = NodeID.from_random().hex()
        cluster_node_info_cache = MockClusterNodeInfoCache()
        cluster_node_info_cache.add_node(node_id_1, {"CPU": 3})
        cluster_node_info_cache.add_node(node_id_2, {"CPU": 100, "customA": 1})

        scheduler = default_impl.create_deployment_scheduler(
            cluster_node_info_cache,
            head_node_id_override="fake-head-node-id",
            create_placement_group_fn_override=lambda *args, **kwargs: MockPlacementGroup(  # noqa
                *args, **kwargs
            ),
        )
        scheduler.on_deployment_created(d_id, SpreadDeploymentSchedulingPolicy())
        scheduler.on_deployment_deployed(
            d_id,
            ReplicaConfig.create(
                dummy, ray_actor_options={"num_cpus": 2, "resources": {"customA": 0.1}}
            ),
        )

        # Despite trying to schedule on node that minimizes fragmentation,
        # should respect custom resources and schedule onto node2
        def on_scheduled(actor_handle, placement_group):
            scheduling_strategy = actor_handle._options["scheduling_strategy"]
            assert isinstance(scheduling_strategy, NodeAffinitySchedulingStrategy)
            assert scheduling_strategy.node_id == node_id_2

        scheduler.schedule(
            upscales={
                d_id: [
                    ReplicaSchedulingRequest(
                        replica_id=ReplicaID(unique_id="r0", deployment_id=d_id),
                        actor_def=MockActorClass(),
                        actor_resources={"CPU": 2, "customA": 0.1},
                        actor_options={"name": "random"},
                        actor_init_args=(),
                        on_scheduled=on_scheduled,
                    )
                ]
            },
            downscales={},
        )

    def test_actor_creation_failure_does_not_decrement_resources(self):
        """When actor creation fails for a replica, available resources
        should not be decremented so subsequent replicas in the same
        scheduling batch can still use that node.
        """

        d_id = DeploymentID(name="deployment1")
        node_id = NodeID.from_random().hex()

        cluster_node_info_cache = MockClusterNodeInfoCache()
        # Node has exactly 1 CPU — enough for one 1-CPU replica.
        cluster_node_info_cache.add_node(node_id, {"CPU": 1})

        scheduler = default_impl.create_deployment_scheduler(
            cluster_node_info_cache,
            head_node_id_override="fake-head-node-id",
            create_placement_group_fn_override=None,
        )
        scheduler.on_deployment_created(d_id, SpreadDeploymentSchedulingPolicy())
        scheduler.on_deployment_deployed(
            d_id,
            ReplicaConfig.create(dummy, ray_actor_options={"num_cpus": 1}),
        )

        # Create a mock actor class whose .options().remote() raises on the
        # first call (simulating actor creation failure) but succeeds after.
        call_count = 0

        class FailOnceMockActorClass(MockActorClass):
            def remote(self, *args):
                nonlocal call_count
                call_count += 1
                if call_count == 1:
                    raise RuntimeError("Simulated actor creation failure")
                return super().remote(*args)

        on_scheduled_mock = Mock()
        r0_id = ReplicaID(unique_id="r0", deployment_id=d_id)
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)

        req0 = ReplicaSchedulingRequest(
            replica_id=r0_id,
            actor_def=FailOnceMockActorClass(),
            actor_resources={"CPU": 1},
            actor_options={},
            actor_init_args=(),
            on_scheduled=on_scheduled_mock,
        )
        req1 = ReplicaSchedulingRequest(
            replica_id=r1_id,
            actor_def=MockActorClass(),
            actor_resources={"CPU": 1},
            actor_options={},
            actor_init_args=(),
            on_scheduled=on_scheduled_mock,
        )

        scheduler.schedule(
            upscales={d_id: [req0, req1]},
            downscales={},
        )

        # The first replica should have failed.
        assert req0.status == ReplicaSchedulingRequestStatus.ACTOR_CREATION_FAILED

        # The second replica should have succeeded and been scheduled to the
        # node.
        assert req1.status == ReplicaSchedulingRequestStatus.SUCCEEDED
        assert on_scheduled_mock.call_count == 1
        call = on_scheduled_mock.call_args_list[0]
        scheduling_strategy = call.args[0]._options["scheduling_strategy"]
        assert isinstance(scheduling_strategy, NodeAffinitySchedulingStrategy)
        assert scheduling_strategy.node_id == node_id

    def test_pg_creation_failure_does_not_decrement_resources(self):
        """When placement group creation fails for a replica, available
        resources should not be decremented so subsequent replicas in the
        same scheduling batch can still use that node.
        """

        d_id = DeploymentID(name="deployment1")
        node_id = NodeID.from_random().hex()

        cluster_node_info_cache = MockClusterNodeInfoCache()
        # Node has exactly 1 CPU — enough for one replica with 1-CPU PG.
        cluster_node_info_cache.add_node(node_id, {"CPU": 1})

        call_count = 0

        def fail_once_create_pg(request):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                raise RuntimeError("Simulated PG creation failure")
            return MockPlacementGroup(request)

        scheduler = default_impl.create_deployment_scheduler(
            cluster_node_info_cache,
            head_node_id_override="fake-head-node-id",
            create_placement_group_fn_override=fail_once_create_pg,
        )
        scheduler.on_deployment_created(d_id, SpreadDeploymentSchedulingPolicy())
        scheduler.on_deployment_deployed(
            d_id,
            ReplicaConfig.create(
                dummy,
                ray_actor_options={"num_cpus": 0},
                placement_group_bundles=[{"CPU": 1}],
                placement_group_strategy="STRICT_PACK",
            ),
        )

        on_scheduled_mock = Mock()
        r0_id = ReplicaID(unique_id="r0", deployment_id=d_id)
        r1_id = ReplicaID(unique_id="r1", deployment_id=d_id)

        req0 = ReplicaSchedulingRequest(
            replica_id=r0_id,
            actor_def=MockActorClass(),
            actor_resources={"CPU": 0},
            placement_group_bundles=[{"CPU": 1}],
            placement_group_strategy="STRICT_PACK",
            actor_options={"name": "r0"},
            actor_init_args=(),
            on_scheduled=on_scheduled_mock,
        )
        req1 = ReplicaSchedulingRequest(
            replica_id=r1_id,
            actor_def=MockActorClass(),
            actor_resources={"CPU": 0},
            placement_group_bundles=[{"CPU": 1}],
            placement_group_strategy="STRICT_PACK",
            actor_options={"name": "r1"},
            actor_init_args=(),
            on_scheduled=on_scheduled_mock,
        )

        scheduler.schedule(
            upscales={d_id: [req0, req1]},
            downscales={},
        )

        # The first replica should have failed at PG creation.
        assert (
            req0.status
            == ReplicaSchedulingRequestStatus.PLACEMENT_GROUP_CREATION_FAILED
        )

        # The second replica should still succeed.
        assert req1.status == ReplicaSchedulingRequestStatus.SUCCEEDED
        assert on_scheduled_mock.call_count == 1
        call = on_scheduled_mock.call_args_list[0]
        scheduling_strategy = call.args[0]._options["scheduling_strategy"]
        assert isinstance(scheduling_strategy, PlacementGroupSchedulingStrategy)

    def test_pack_prefers_newly_non_idle_node(self):
        """After scheduling a replica to a previously idle node, subsequent
        replicas in the same batch should prefer that node (now non-idle)
        over other idle nodes, even if the idle node is a tighter fit.

        Regression test: without updating node_to_running_replicas after
        each scheduling, the PACK scheduler would treat all initially-idle
        nodes as idle for the entire batch, falling through to pure
        best-fit and potentially spreading replicas across nodes.
        """

        d_id1 = DeploymentID(name="deployment1")
        d_id2 = DeploymentID(name="deployment2")
        node_id_1 = NodeID.from_random().hex()
        node_id_2 = NodeID.from_random().hex()

        cluster_node_info_cache = MockClusterNodeInfoCache()
        # Node 1 has GPU + CPU; node 2 has only CPU.
        # After the GPU replica is placed on node 1, node 2 would be
        # a tighter best-fit for a CPU-only replica (2 CPU remaining
        # vs 4 CPU on node 1). But PACK should prefer node 1 because
        # it is now non-idle.
        cluster_node_info_cache.add_node(node_id_1, {"GPU": 1, "CPU": 4})
        cluster_node_info_cache.add_node(node_id_2, {"CPU": 2})

        scheduler = default_impl.create_deployment_scheduler(
            cluster_node_info_cache,
            head_node_id_override="fake-head-node-id",
            create_placement_group_fn_override=None,
        )

        scheduler.on_deployment_created(d_id1, SpreadDeploymentSchedulingPolicy())
        scheduler.on_deployment_created(d_id2, SpreadDeploymentSchedulingPolicy())

        scheduler.on_deployment_deployed(
            d_id1,
            ReplicaConfig.create(
                dummy, ray_actor_options={"num_gpus": 1, "num_cpus": 0}
            ),
        )
        scheduler.on_deployment_deployed(
            d_id2,
            ReplicaConfig.create(dummy, ray_actor_options={"num_cpus": 1}),
        )

        on_scheduled_mock1 = Mock()
        on_scheduled_mock2 = Mock()
        scheduler.schedule(
            upscales={
                d_id1: [
                    ReplicaSchedulingRequest(
                        replica_id=ReplicaID(unique_id="r0", deployment_id=d_id1),
                        actor_def=MockActorClass(),
                        actor_resources={"GPU": 1},
                        actor_options={},
                        actor_init_args=(),
                        on_scheduled=on_scheduled_mock1,
                    )
                ],
                d_id2: [
                    ReplicaSchedulingRequest(
                        replica_id=ReplicaID(unique_id="r1", deployment_id=d_id2),
                        actor_def=MockActorClass(),
                        actor_resources={"CPU": 1},
                        actor_options={},
                        actor_init_args=(),
                        on_scheduled=on_scheduled_mock2,
                    )
                ],
            },
            downscales={},
        )

        # The GPU replica must go to node 1 (only node with GPU).
        assert len(on_scheduled_mock1.call_args_list) == 1
        call1 = on_scheduled_mock1.call_args_list[0]
        strategy1 = call1.args[0]._options["scheduling_strategy"]
        assert isinstance(strategy1, NodeAffinitySchedulingStrategy)
        assert strategy1.node_id == node_id_1
        assert call1.kwargs == {"placement_group": None}

        # The CPU replica should also go to node 1 (now non-idle) rather
        # than node 2 (idle but tighter fit). The PACK scheduler prefers
        # non-idle nodes to consolidate replicas onto fewer nodes.
        assert len(on_scheduled_mock2.call_args_list) == 1
        call2 = on_scheduled_mock2.call_args_list[0]
        strategy2 = call2.args[0]._options["scheduling_strategy"]
        assert isinstance(strategy2, NodeAffinitySchedulingStrategy)
        assert strategy2.node_id == node_id_1
        assert call2.kwargs == {"placement_group": None}


class TestScheduleGangPlacementGroups:
    def test_schedule_gang_placement_groups(self, mock_deployment_state_manager):
        """Creates gangs successfully and verifies placement requests include expected bundles and strategy."""
        captured_requests = []
        gang_size = 2
        num_gangs = 2
        num_replicas_to_add = gang_size * num_gangs
        replica_resource_dict = {"CPU": 2.0, "GPU": 1.0}
        gang_strategy = "SPREAD"

        def create_pg_fn(request: CreatePlacementGroupRequest, *args, **kwargs):
            captured_requests.append(request)
            return Mock()

        create_dsm, _, _, _ = mock_deployment_state_manager
        dsm: DeploymentStateManager = create_dsm(
            create_placement_group_fn_override=create_pg_fn,
        )
        scheduler = dsm._deployment_scheduler
        deployment_id = DeploymentID(name="d1", app_name="app1")
        gang_request = GangPlacementGroupRequest(
            deployment_id,
            gang_size,
            gang_strategy,
            num_replicas_to_add,
            replica_resource_dict=replica_resource_dict,
        )

        result = scheduler.schedule_gang_placement_groups({deployment_id: gang_request})

        assert deployment_id in result
        reservation = result[deployment_id]
        assert reservation.success
        assert len(reservation.gang_pgs) == num_gangs
        assert len(captured_requests) == num_gangs
        for req in captured_requests:
            assert isinstance(req, CreatePlacementGroupRequest)
            assert req.bundles == [replica_resource_dict] * gang_size
            assert req.strategy == gang_strategy

        assert len(reservation.gang_ids) == num_gangs
        assert len(reservation.gang_pg_names) == num_gangs
        assert len(set(reservation.gang_ids)) == num_gangs
        for pg_name in reservation.gang_pg_names:
            assert pg_name.startswith(GANG_PG_NAME_PREFIX)

    def test_schedule_gang_placement_groups_invalid_gang_size(
        self, mock_deployment_state_manager
    ):
        """Returns failure when desired replicas cannot be evenly divided by gang size."""
        gang_size = 3
        num_replicas_to_add = 4
        create_pg_fn = Mock()
        create_dsm, _, _, _ = mock_deployment_state_manager
        dsm: DeploymentStateManager = create_dsm(
            create_placement_group_fn_override=create_pg_fn,
        )
        scheduler = dsm._deployment_scheduler
        deployment_id = DeploymentID(name="d2", app_name="app2")
        gang_request = GangPlacementGroupRequest(
            deployment_id,
            gang_size,
            "STRICT_PACK",
            num_replicas_to_add,
            {"CPU": 1.0},
        )

        result = scheduler.schedule_gang_placement_groups({deployment_id: gang_request})

        assert not result[deployment_id].success
        assert "not divisible by gang_size" in result[deployment_id].error_message
        create_pg_fn.assert_not_called()

    def test_schedule_gang_placement_groups_all_pg_creation_failures(
        self, mock_deployment_state_manager
    ):
        """Reports failure when every gang placement group creation attempt raises exceptions."""
        gang_size = 2
        num_gangs = 2
        num_replicas_to_add = gang_size * num_gangs

        def create_pg_fn(request: CreatePlacementGroupRequest, *args, **kwargs):
            raise RuntimeError("simulated placement group creation failure")

        create_dsm, _, _, _ = mock_deployment_state_manager
        dsm: DeploymentStateManager = create_dsm(
            create_placement_group_fn_override=create_pg_fn,
        )
        scheduler = dsm._deployment_scheduler
        deployment_id = DeploymentID(name="d3", app_name="app3")
        gang_request = GangPlacementGroupRequest(
            deployment_id,
            gang_size,
            "STRICT_PACK",
            num_replicas_to_add,
            {"CPU": 1.0},
        )

        result = scheduler.schedule_gang_placement_groups({deployment_id: gang_request})

        assert not result[deployment_id].success
        assert (
            "Failed to create any gang placement groups"
            in result[deployment_id].error_message
        )

    def test_schedule_gang_placement_groups_partial_pg_creation_failures(
        self, mock_deployment_state_manager
    ):
        """Keeps successful gang reservations when only a subset of placement groups fail."""
        gang_size = 2
        num_gangs = 2
        num_replicas_to_add = gang_size * num_gangs
        failed_gangs = 1
        num_calls = 0

        def create_pg_fn(request: CreatePlacementGroupRequest, *args, **kwargs):
            nonlocal num_calls
            num_calls += 1
            if num_calls == 1:
                raise RuntimeError("fail first gang only")
            return Mock()

        create_dsm, _, _, _ = mock_deployment_state_manager
        dsm: DeploymentStateManager = create_dsm(
            create_placement_group_fn_override=create_pg_fn,
        )
        scheduler = dsm._deployment_scheduler
        deployment_id = DeploymentID(name="d4", app_name="app4")
        gang_request = GangPlacementGroupRequest(
            deployment_id,
            gang_size,
            "STRICT_PACK",
            num_replicas_to_add,
            {"CPU": 1.0},
        )

        result = scheduler.schedule_gang_placement_groups({deployment_id: gang_request})

        assert result[deployment_id].success
        assert len(result[deployment_id].gang_pgs) == num_gangs - failed_gangs

    def test_schedule_gang_placement_groups_with_per_replica_bundles(
        self, mock_deployment_state_manager
    ):
        """Flattens per-replica bundles and propagates label selectors and fallback strategies correctly."""
        captured_requests = []
        gang_size = 2
        num_gangs = 4
        num_replicas_to_add = num_gangs * gang_size

        def create_pg_fn(request: CreatePlacementGroupRequest, *args, **kwargs):
            captured_requests.append(request)
            return Mock()

        create_dsm, _, _, _ = mock_deployment_state_manager
        dsm: DeploymentStateManager = create_dsm(
            create_placement_group_fn_override=create_pg_fn,
        )
        scheduler = dsm._deployment_scheduler
        deployment_id = DeploymentID(name="d5", app_name="app5")
        per_replica_bundles = [{"GPU": 1.0, "CPU": 1.0}, {"CPU": 1.0}]
        per_replica_label_selector = [{"gpu": "a100"}, {"zone": "z1"}]
        per_replica_fallback = [{"allow_soft": True}, {"allow_soft": False}]
        gang_request = GangPlacementGroupRequest(
            deployment_id,
            gang_size,
            "STRICT_PACK",
            num_replicas_to_add,
            replica_resource_dict={"CPU": 1.0},
            replica_placement_group_bundles=per_replica_bundles,
            replica_pg_bundle_label_selector=per_replica_label_selector,
            replica_pg_fallback_strategy=per_replica_fallback,
        )

        result = scheduler.schedule_gang_placement_groups({deployment_id: gang_request})

        assert result[deployment_id].success
        assert len(captured_requests) == num_gangs
        expected_bundles = per_replica_bundles * gang_size
        expected_label_selector = per_replica_label_selector * gang_size
        expected_fallback = per_replica_fallback * gang_size
        for req in captured_requests:
            assert req.bundles == expected_bundles
            assert req.bundle_label_selector == expected_label_selector
            assert req.fallback_strategy == expected_fallback

    def test_schedule_gang_placement_groups_without_per_replica_bundles_uses_resource_dict(
        self, mock_deployment_state_manager
    ):
        """Uses replica resource dict for each gang bundle without optional selectors."""
        captured_requests = []
        gang_size = 3
        num_gangs = 2
        num_replicas_to_add = gang_size * num_gangs
        replica_resource_dict = {"CPU": 2.0, "GPU": 0.5}

        def create_pg_fn(request: CreatePlacementGroupRequest, *args, **kwargs):
            captured_requests.append(request)
            return Mock()

        create_dsm, _, _, _ = mock_deployment_state_manager
        dsm: DeploymentStateManager = create_dsm(
            create_placement_group_fn_override=create_pg_fn,
        )
        scheduler = dsm._deployment_scheduler
        deployment_id = DeploymentID(name="d6", app_name="app6")
        gang_request = GangPlacementGroupRequest(
            deployment_id,
            gang_size,
            "STRICT_PACK",
            num_replicas_to_add,
            replica_resource_dict=replica_resource_dict,
        )

        result = scheduler.schedule_gang_placement_groups({deployment_id: gang_request})

        assert result[deployment_id].success
        assert len(captured_requests) == num_gangs
        for req in captured_requests:
            assert req.bundles == [replica_resource_dict] * gang_size
            assert req.bundle_label_selector is None
            assert req.fallback_strategy is None

    def test_schedule_gang_placement_groups_multiple_deployments(
        self, mock_deployment_state_manager
    ):
        """Schedules gang placement groups for multiple deployments and returns independent results."""
        create_pg_fn = Mock(return_value=Mock())
        gang_size_1 = 2
        num_gangs_1 = 2
        num_replicas_to_add_1 = gang_size_1 * num_gangs_1
        gang_size_2 = 3
        num_gangs_2 = 2
        num_replicas_to_add_2 = gang_size_2 * num_gangs_2
        create_dsm, _, _, _ = mock_deployment_state_manager
        dsm: DeploymentStateManager = create_dsm(
            create_placement_group_fn_override=create_pg_fn,
        )
        scheduler = dsm._deployment_scheduler
        deployment_id_1 = DeploymentID(name="d7", app_name="app7")
        deployment_id_2 = DeploymentID(name="d8", app_name="app8")
        gang_requests = {
            deployment_id_1: GangPlacementGroupRequest(
                deployment_id_1,
                gang_size_1,
                "STRICT_PACK",
                num_replicas_to_add_1,
                {"CPU": 1.0},
            ),
            deployment_id_2: GangPlacementGroupRequest(
                deployment_id_2,
                gang_size_2,
                "STRICT_PACK",
                num_replicas_to_add_2,
                {"CPU": 1.0},
            ),
        }

        result = scheduler.schedule_gang_placement_groups(gang_requests)

        assert set(result.keys()) == {deployment_id_1, deployment_id_2}
        assert result[deployment_id_1].success
        assert result[deployment_id_2].success
        assert len(result[deployment_id_1].gang_pgs) == num_gangs_1
        assert len(result[deployment_id_2].gang_pgs) == num_gangs_2
        assert create_pg_fn.call_count == num_gangs_1 + num_gangs_2


if __name__ == "__main__":
    sys.exit(pytest.main(["-v", "-s", __file__]))