wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
ray-project--ray/python/ray/serve/tests/unit/test_deployment_scheduler.py
T
wehub-resource-sync f1825c8ceb chore: import upstream snapshot with attribution
2026-07-13 13:17:40 +08:00

2180 lines
82 KiB
Python
Raw Permalink Blame History

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__]))
Reference in New Issue View Git Blame Copy Permalink