wehub-resource-sync
1983 lines
73 KiB
Python
1983 lines
73 KiB
Python
import os
|
||
import sys
|
||
import threading
|
||
import time
|
||
|
||
import pytest
|
||
|
||
import ray
|
||
from ray import serve
|
||
from ray._common.test_utils import SignalActor, wait_for_condition
|
||
from ray.serve._private.common import GANG_PG_NAME_PREFIX, DeploymentID, ReplicaState
|
||
from ray.serve._private.constants import SERVE_DEFAULT_APP_NAME
|
||
from ray.serve._private.test_utils import (
|
||
Accumulator,
|
||
FailedGangReplicaStore,
|
||
check_apps_running,
|
||
check_num_replicas_eq,
|
||
)
|
||
from ray.serve._private.utils import get_all_live_placement_group_names
|
||
from ray.serve.config import GangPlacementStrategy, GangSchedulingConfig
|
||
from ray.serve.context import _get_global_client
|
||
from ray.tests.conftest import * # noqa
|
||
from ray.util.placement_group import get_current_placement_group, placement_group_table
|
||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||
|
||
|
||
def _get_running_replicas(deployment_id: DeploymentID):
|
||
"""Return RUNNING replicas for a deployment from controller state."""
|
||
controller = _get_global_client()._controller
|
||
replicas = ray.get(
|
||
controller._dump_replica_states_for_testing.remote(deployment_id)
|
||
)
|
||
return replicas.get([ReplicaState.RUNNING])
|
||
|
||
|
||
def _get_gang_ids_from_running(running) -> set:
|
||
return {r.gang_context.gang_id for r in running if r.gang_context is not None}
|
||
|
||
|
||
def _get_node_ids_from_running(running) -> set:
|
||
return {r.actor_node_id for r in running if r.actor_node_id}
|
||
|
||
|
||
class TestGangScheduling:
|
||
"""Tests for gang scheduling with placement groups."""
|
||
|
||
def test_sufficient_resources(self, ray_cluster):
|
||
"""Verifies that gang scheduling succeeds when cluster has sufficient resources."""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment(
|
||
num_replicas=8,
|
||
ray_actor_options={"num_cpus": 0.25},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=4),
|
||
)
|
||
class GangDeployment:
|
||
def __call__(self):
|
||
return ray.get_runtime_context().get_node_id()
|
||
|
||
handle = serve.run(GangDeployment.bind(), name="gang_app_success")
|
||
wait_for_condition(
|
||
check_apps_running,
|
||
apps=["gang_app_success"],
|
||
)
|
||
|
||
# Verify all replicas are running and responding
|
||
refs = [handle.remote() for _ in range(8)]
|
||
results = [ref.result() for ref in refs]
|
||
assert len(results) == 8
|
||
|
||
serve.delete("gang_app_success")
|
||
serve.shutdown()
|
||
|
||
def test_sufficient_resources_with_options(self, ray_cluster):
|
||
"""Verifies gang scheduling via .options() succeeds and responds to requests."""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment(num_replicas=1, ray_actor_options={"num_cpus": 0})
|
||
class GangDeployment:
|
||
def __call__(self):
|
||
return ray.get_runtime_context().get_node_id()
|
||
|
||
app = GangDeployment.options(
|
||
num_replicas=8,
|
||
ray_actor_options={"num_cpus": 0.25},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=4),
|
||
).bind()
|
||
|
||
handle = serve.run(app, name="gang_app_options")
|
||
wait_for_condition(
|
||
check_apps_running,
|
||
apps=["gang_app_options"],
|
||
)
|
||
|
||
# Verify all replicas are running and responding
|
||
refs = [handle.remote() for _ in range(8)]
|
||
results = [ref.result() for ref in refs]
|
||
assert len(results) == 8
|
||
|
||
serve.delete("gang_app_options")
|
||
serve.shutdown()
|
||
|
||
def test_incomplete_deployment(self, ray_cluster):
|
||
"""
|
||
Verifies that schedulable gangs serve traffic while unschedulable gangs wait for resources.
|
||
"""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment
|
||
class IncompleteGangDeployment:
|
||
def __call__(self):
|
||
return ray.get_runtime_context().get_node_id()
|
||
|
||
app = IncompleteGangDeployment.options(
|
||
num_replicas=12,
|
||
ray_actor_options={"num_cpus": 0.25},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=4),
|
||
).bind()
|
||
|
||
handle = serve._run(app, name="gang_partial_app", _blocking=False)
|
||
|
||
# The deployment should NOT fail. 2 of 3 gangs should be scheduled,
|
||
# and those 8 replicas should serve traffic. The deployment stays
|
||
# DEPLOYING because it hasn't reached 12 replicas.
|
||
def check_replicas_running(expected_count: int):
|
||
try:
|
||
app_status = serve.status().applications["gang_partial_app"]
|
||
# Should be DEPLOYING
|
||
if app_status.status == "DEPLOY_FAILED":
|
||
raise AssertionError(
|
||
"Deployment should not fail with partial gang scheduling"
|
||
)
|
||
# Check that some replicas are running
|
||
dep_status = list(app_status.deployments.values())[0]
|
||
running = dep_status.replica_states.get("RUNNING", 0)
|
||
assert running == expected_count
|
||
return True
|
||
except KeyError:
|
||
return False
|
||
|
||
wait_for_condition(check_replicas_running, expected_count=8, timeout=60)
|
||
|
||
# Verify the running replicas can serve traffic.
|
||
results = set()
|
||
for _ in range(40):
|
||
results.add(handle.remote().result())
|
||
assert len(results) > 0
|
||
|
||
# Verify deployment is still DEPLOYING
|
||
app_status = serve.status().applications["gang_partial_app"]
|
||
assert app_status.status == "DEPLOYING"
|
||
|
||
# Now add a 3rd node so the remaining gang can be scheduled.
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
|
||
# The deployment should become RUNNING with all 12 replicas.
|
||
wait_for_condition(
|
||
check_apps_running,
|
||
apps=["gang_partial_app"],
|
||
timeout=60,
|
||
)
|
||
|
||
# Verify all 12 replicas are running across 3 nodes (controller state,
|
||
# not handle routing, which may only hit local replicas).
|
||
dep_id = DeploymentID(
|
||
name="IncompleteGangDeployment", app_name="gang_partial_app"
|
||
)
|
||
running = _get_running_replicas(dep_id)
|
||
assert len(running) == 12
|
||
assert len(_get_node_ids_from_running(running)) == 3
|
||
|
||
serve.delete("gang_partial_app")
|
||
serve.shutdown()
|
||
|
||
def test_no_partial_gang(self, ray_cluster):
|
||
"""Verifies atomic gang scheduling: no partial gangs are created."""
|
||
cluster = ray_cluster
|
||
# 2 CPUs total: enough for 2 full gangs (1.6 CPUs) but not 3 (2.4 CPUs).
|
||
# The leftover 0.4 CPUs must NOT produce a partial gang.
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment
|
||
class AtomicGangDeployment:
|
||
def __call__(self):
|
||
return ray.get_runtime_context().get_node_id()
|
||
|
||
app = AtomicGangDeployment.options(
|
||
num_replicas=12,
|
||
ray_actor_options={"num_cpus": 0.2},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=4),
|
||
).bind()
|
||
|
||
handle = serve._run(app, name="atomic_gang_app", _blocking=False)
|
||
|
||
# Wait until exactly 8 replicas (2 gangs) are running.
|
||
def check_replicas_running(expected_count: int):
|
||
try:
|
||
app_status = serve.status().applications["atomic_gang_app"]
|
||
if app_status.status == "DEPLOY_FAILED":
|
||
raise AssertionError(
|
||
"Deployment should not fail — partial gangs should "
|
||
"serve traffic while waiting for resources."
|
||
)
|
||
dep_status = list(app_status.deployments.values())[0]
|
||
running = dep_status.replica_states.get("RUNNING", 0)
|
||
assert running == expected_count
|
||
return True
|
||
except KeyError:
|
||
return False
|
||
|
||
wait_for_condition(check_replicas_running, expected_count=8, timeout=60)
|
||
|
||
# Deployment should still be DEPLOYING (not RUNNING, not DEPLOY_FAILED).
|
||
app_status = serve.status().applications["atomic_gang_app"]
|
||
assert app_status.status == "DEPLOYING"
|
||
|
||
# Verify the 8 running replicas can serve traffic.
|
||
results = set()
|
||
for _ in range(80):
|
||
results.add(handle.remote().result())
|
||
assert len(results) > 0
|
||
|
||
# Add 1 more CPU so the 3rd gang (0.8 CPUs) can be scheduled.
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
|
||
# The deployment should become RUNNING with all 12 replicas.
|
||
wait_for_condition(check_apps_running, apps=["atomic_gang_app"], timeout=60)
|
||
|
||
# All 12 replicas should now serve traffic.
|
||
app_status = serve.status().applications["atomic_gang_app"]
|
||
dep_status = list(app_status.deployments.values())[0]
|
||
running = dep_status.replica_states.get("RUNNING", 0)
|
||
assert running == 12
|
||
|
||
serve.delete("atomic_gang_app")
|
||
serve.shutdown()
|
||
|
||
def test_pack_strategy(self, ray_cluster):
|
||
"""Verifies that PACK strategy places gang replicas on the same node."""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment
|
||
def PackDeployment():
|
||
return os.environ.get(
|
||
"RAY_NODE_ID", ray.get_runtime_context().get_node_id()
|
||
)
|
||
|
||
# 1 gang with PACK strategy - all replicas should be on same node
|
||
app = PackDeployment.options(
|
||
num_replicas=4,
|
||
ray_actor_options={"num_cpus": 0.25},
|
||
gang_scheduling_config=GangSchedulingConfig(
|
||
gang_size=4,
|
||
gang_placement_strategy=GangPlacementStrategy.PACK,
|
||
),
|
||
).bind()
|
||
|
||
handle = serve.run(app, name="gang_pack_app")
|
||
wait_for_condition(check_apps_running, apps=["gang_pack_app"])
|
||
|
||
# Query multiple times to hit all replicas and collect node IDs.
|
||
# Intentionally handle-based: the assertion is that all replicas share
|
||
# a single node, and handle routing can only ever surface a subset of
|
||
# the nodes actually used. Locality-aware routing can therefore never
|
||
# inflate this count, so it cannot cause a false failure here.
|
||
node_ids = set()
|
||
for _ in range(40):
|
||
result = handle.remote().result()
|
||
node_ids.add(result)
|
||
|
||
# With PACK strategy, all 4 replicas should be on the same node
|
||
assert len(node_ids) == 1
|
||
|
||
serve.delete("gang_pack_app")
|
||
serve.shutdown()
|
||
|
||
def test_gang_scheduling_spread_strategy(self, ray_cluster):
|
||
"""Verifies that SPREAD strategy places gang replicas on different nodes."""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment
|
||
def SpreadDeployment():
|
||
return os.environ.get(
|
||
"RAY_NODE_ID", ray.get_runtime_context().get_node_id()
|
||
)
|
||
|
||
# 1 gang with SPREAD strategy - replicas should be on different nodes
|
||
app = SpreadDeployment.options(
|
||
num_replicas=2,
|
||
ray_actor_options={"num_cpus": 0.25},
|
||
gang_scheduling_config=GangSchedulingConfig(
|
||
gang_size=2,
|
||
gang_placement_strategy=GangPlacementStrategy.SPREAD,
|
||
),
|
||
).bind()
|
||
|
||
serve.run(app, name="gang_spread_app")
|
||
wait_for_condition(check_apps_running, apps=["gang_spread_app"])
|
||
|
||
# With SPREAD strategy, 2 replicas should be on 2 different nodes.
|
||
dep_id = DeploymentID(name="SpreadDeployment", app_name="gang_spread_app")
|
||
running = _get_running_replicas(dep_id)
|
||
assert len(running) == 2
|
||
assert len(_get_node_ids_from_running(running)) == 2
|
||
|
||
serve.delete("gang_spread_app")
|
||
serve.shutdown()
|
||
|
||
def test_gang_context(self, ray_cluster):
|
||
"""Verifies GangContext is correctly populated in ReplicaContext."""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment
|
||
class GangContextDeployment:
|
||
def __call__(self):
|
||
return ray.get_runtime_context().get_node_id()
|
||
|
||
app = GangContextDeployment.options(
|
||
num_replicas=4,
|
||
ray_actor_options={"num_cpus": 0.25},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
).bind()
|
||
|
||
serve.run(app, name="gang_context_app")
|
||
wait_for_condition(check_apps_running, apps=["gang_context_app"])
|
||
|
||
# Read gang context from controller replica state instead of handle
|
||
# routing (which may only hit local replicas under locality-aware
|
||
# routing). The controller stores the exact GangContext each replica
|
||
# reports from its own ReplicaContext, so this verifies the same values.
|
||
dep_id = DeploymentID(name="GangContextDeployment", app_name="gang_context_app")
|
||
running = _get_running_replicas(dep_id)
|
||
assert len(running) == 4
|
||
assert all(r.gang_context is not None for r in running)
|
||
|
||
# Group replicas by gang_id.
|
||
gangs = {}
|
||
for r in running:
|
||
gangs.setdefault(r.gang_context.gang_id, []).append(r)
|
||
|
||
assert len(gangs) == 2
|
||
|
||
for gang_id, members in gangs.items():
|
||
assert len(members) == 2
|
||
assert all(m.gang_context.world_size == 2 for m in members)
|
||
assert (
|
||
members[0].gang_context.member_replica_ids
|
||
== members[1].gang_context.member_replica_ids
|
||
)
|
||
|
||
expected_ids = sorted([m.replica_id.unique_id for m in members])
|
||
actual_ids = sorted(members[0].gang_context.member_replica_ids)
|
||
assert actual_ids == expected_ids
|
||
|
||
ranks = sorted([m.gang_context.rank for m in members])
|
||
assert ranks == [0, 1]
|
||
|
||
# Across gangs: gang_ids should be different.
|
||
gang_ids = list(gangs.keys())
|
||
assert gang_ids[0] != gang_ids[1]
|
||
|
||
# Across gangs: member_replica_ids should be different
|
||
gang_members_list = list(gangs.values())
|
||
assert sorted(
|
||
gang_members_list[0][0].gang_context.member_replica_ids
|
||
) != sorted(gang_members_list[1][0].gang_context.member_replica_ids)
|
||
|
||
serve.delete("gang_context_app")
|
||
serve.shutdown()
|
||
|
||
def test_gang_placement_groups_cleanup_on_deletion(self, ray_cluster):
|
||
"""Verifies serve.delete() removes reserved gang placement groups."""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment(
|
||
num_replicas=4,
|
||
ray_actor_options={"num_cpus": 0.25},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class GangDeleteCleanupDeployment:
|
||
def __call__(self):
|
||
return "ok"
|
||
|
||
app_name = "gang_delete_cleanup_app"
|
||
deployment_name = "GangDeleteCleanupDeployment"
|
||
pg_name_prefix = f"{GANG_PG_NAME_PREFIX}{app_name}_{deployment_name}_"
|
||
|
||
serve.run(GangDeleteCleanupDeployment.bind(), name=app_name)
|
||
wait_for_condition(check_apps_running, apps=[app_name])
|
||
|
||
wait_for_condition(
|
||
lambda: any(
|
||
name.startswith(pg_name_prefix)
|
||
for name in get_all_live_placement_group_names()
|
||
),
|
||
timeout=60,
|
||
)
|
||
|
||
serve.delete(app_name)
|
||
wait_for_condition(
|
||
lambda: not any(
|
||
name.startswith(pg_name_prefix)
|
||
for name in get_all_live_placement_group_names()
|
||
),
|
||
timeout=60,
|
||
)
|
||
serve.shutdown()
|
||
|
||
def test_multiple_gang_deployments_in_one_app(self, ray_cluster):
|
||
"""Verifies two gang deployments run together under one Serve app."""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment(
|
||
num_replicas=4,
|
||
ray_actor_options={"num_cpus": 0.25},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class GangA:
|
||
def __init__(self, gang_b):
|
||
self._gang_b = gang_b
|
||
|
||
def __call__(self):
|
||
return "a"
|
||
|
||
@serve.deployment(
|
||
num_replicas=4,
|
||
ray_actor_options={"num_cpus": 0.25},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class GangB:
|
||
def __call__(self):
|
||
return "b"
|
||
|
||
app_name = "multi_gang_app"
|
||
serve.run(GangA.bind(GangB.bind()), name=app_name)
|
||
wait_for_condition(check_apps_running, apps=[app_name])
|
||
|
||
app_status = serve.status().applications[app_name]
|
||
assert app_status.deployments["GangA"].replica_states.get("RUNNING", 0) == 4
|
||
assert app_status.deployments["GangB"].replica_states.get("RUNNING", 0) == 4
|
||
|
||
serve.delete(app_name)
|
||
serve.shutdown()
|
||
|
||
|
||
class TestGangResourceReservation:
|
||
@pytest.mark.parametrize(
|
||
"ray_actor_options, placement_group_bundles, gang_placement_strategy, "
|
||
"expected_bundles, expected_strategy, expect_same_node",
|
||
[
|
||
# Case 1: Only ray_actor_options — one flat bundle per replica, PACK
|
||
(
|
||
{"num_cpus": 0.25},
|
||
None,
|
||
"PACK",
|
||
[{"CPU": 0.25}, {"CPU": 0.25}],
|
||
"PACK",
|
||
True,
|
||
),
|
||
# Case 2: placement_group_bundles — flattened into the gang PG, PACK
|
||
(
|
||
{"num_cpus": 0},
|
||
[{"CPU": 0.25}] * 2,
|
||
"PACK",
|
||
[{"CPU": 0.25}] * 4,
|
||
"PACK",
|
||
True,
|
||
),
|
||
# Case 3: placement_group_bundles + SPREAD strategy
|
||
(
|
||
{"num_cpus": 0},
|
||
[{"CPU": 0.25}] * 2,
|
||
"SPREAD",
|
||
[{"CPU": 0.25}] * 4,
|
||
"SPREAD",
|
||
False,
|
||
),
|
||
],
|
||
)
|
||
def test_gang_resource_reservation(
|
||
self,
|
||
ray_cluster,
|
||
ray_actor_options,
|
||
placement_group_bundles,
|
||
gang_placement_strategy,
|
||
expected_bundles,
|
||
expected_strategy,
|
||
expect_same_node,
|
||
):
|
||
"""Verifies the gang PG has the correct bundles, strategy, and
|
||
that per-replica bundles are placed according to the strategy."""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
deployment_kwargs = {
|
||
"num_replicas": 2,
|
||
"ray_actor_options": ray_actor_options,
|
||
"gang_scheduling_config": GangSchedulingConfig(
|
||
gang_size=2,
|
||
gang_placement_strategy=gang_placement_strategy,
|
||
),
|
||
}
|
||
if placement_group_bundles is not None:
|
||
deployment_kwargs["placement_group_bundles"] = placement_group_bundles
|
||
|
||
@serve.deployment(**deployment_kwargs)
|
||
class GangDeployment:
|
||
def get_pg_info(self):
|
||
pg = get_current_placement_group()
|
||
if pg is None:
|
||
return None
|
||
pg_table = placement_group_table(pg)
|
||
return {
|
||
"bundle_specs": pg.bundle_specs,
|
||
"strategy": pg_table["strategy"],
|
||
"bundles_to_node_id": pg_table["bundles_to_node_id"],
|
||
}
|
||
|
||
def __call__(self):
|
||
return "ok"
|
||
|
||
app = GangDeployment.bind()
|
||
handle = serve.run(app, name="gang_reservation_app")
|
||
wait_for_condition(
|
||
check_apps_running,
|
||
apps=["gang_reservation_app"],
|
||
)
|
||
|
||
# Intentionally handle-based: each response is a self-contained
|
||
# per-replica invariant (bundle specs, strategy, per-replica bundle
|
||
# placement), so validating any sampled subset is sufficient. This
|
||
# never needs to enumerate all replicas, so locality-aware routing
|
||
# cannot cause a false failure.
|
||
for _ in range(20):
|
||
pg_info = handle.get_pg_info.remote().result()
|
||
assert pg_info is not None
|
||
assert pg_info["bundle_specs"] == expected_bundles
|
||
assert pg_info["strategy"] == expected_strategy
|
||
|
||
bundles_per_replica = (
|
||
len(placement_group_bundles) if placement_group_bundles else 1
|
||
)
|
||
gang_size = 2
|
||
|
||
for replica_idx in range(gang_size):
|
||
start = replica_idx * bundles_per_replica
|
||
replica_nodes = {
|
||
pg_info["bundles_to_node_id"][i]
|
||
for i in range(start, start + bundles_per_replica)
|
||
}
|
||
if expect_same_node:
|
||
assert len(replica_nodes) == 1
|
||
else:
|
||
assert len(replica_nodes) == bundles_per_replica
|
||
|
||
serve.delete("gang_reservation_app")
|
||
serve.shutdown()
|
||
|
||
def test_gang_label_selector(self, ray_cluster):
|
||
"""
|
||
Verifies that placement_group_bundle_label_selector steers gang bundles
|
||
onto the labeled node.
|
||
"""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.add_node(num_cpus=1, labels={"accelerator": "tpu"})
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment(
|
||
num_replicas=2,
|
||
ray_actor_options={"num_cpus": 0},
|
||
placement_group_bundles=[{"CPU": 0.25}],
|
||
placement_group_bundle_label_selector=[{"accelerator": "tpu"}],
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class LabeledGangDeployment:
|
||
def get_pg_info(self):
|
||
pg = get_current_placement_group()
|
||
if pg is None:
|
||
return None
|
||
pg_table = placement_group_table(pg)
|
||
return {
|
||
"bundle_specs": pg.bundle_specs,
|
||
"bundles_to_node_id": pg_table["bundles_to_node_id"],
|
||
"node_labels": ray.get_runtime_context().get_node_labels(),
|
||
}
|
||
|
||
def __call__(self):
|
||
return "ok"
|
||
|
||
app = LabeledGangDeployment.bind()
|
||
handle = serve.run(app, name="label_selector_app")
|
||
wait_for_condition(
|
||
check_apps_running,
|
||
apps=["label_selector_app"],
|
||
)
|
||
|
||
labeled_node_id = None
|
||
for node in ray.nodes():
|
||
if node["Labels"].get("accelerator") == "tpu":
|
||
labeled_node_id = node["NodeID"]
|
||
break
|
||
assert labeled_node_id is not None
|
||
|
||
# Intentionally handle-based: each response is a self-contained
|
||
# per-replica invariant (all bundles on the labeled node), so
|
||
# validating any sampled subset is sufficient and locality-aware
|
||
# routing cannot cause a false failure.
|
||
for _ in range(20):
|
||
pg_info = handle.get_pg_info.remote().result()
|
||
assert pg_info is not None
|
||
assert pg_info["bundle_specs"] == [{"CPU": 0.25}, {"CPU": 0.25}]
|
||
# Replica actor itself should be on the labeled node
|
||
assert pg_info["node_labels"].get("accelerator") == "tpu"
|
||
# All bundles in the gang PG should be on the labeled node
|
||
for node_id in pg_info["bundles_to_node_id"].values():
|
||
assert node_id == labeled_node_id
|
||
|
||
serve.delete("label_selector_app")
|
||
serve.shutdown()
|
||
|
||
|
||
class TestGangConstructorFailure:
|
||
"""Tests for gang scheduling with constructor failures."""
|
||
|
||
def test_consistent_constructor_failure(self, ray_shutdown):
|
||
"""Validates gang deployment where all replicas consistently fail their constructor."""
|
||
ray.init(num_cpus=1)
|
||
serve.start()
|
||
|
||
@serve.deployment(
|
||
num_replicas=4,
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class GangConstructorFailure:
|
||
def __init__(self):
|
||
raise RuntimeError("Intentionally failing gang replica constructor")
|
||
|
||
async def __call__(self, request):
|
||
return "hi"
|
||
|
||
with pytest.raises(RuntimeError):
|
||
serve.run(GangConstructorFailure.bind())
|
||
|
||
client = serve.context._get_global_client()
|
||
deployment_dict = ray.get(client._controller._all_running_replicas.remote())
|
||
deployment_id = DeploymentID(name="GangConstructorFailure")
|
||
assert len(deployment_dict[deployment_id]) == 0
|
||
app_status = serve.status().applications[SERVE_DEFAULT_APP_NAME]
|
||
assert app_status.status == "DEPLOY_FAILED"
|
||
assert (
|
||
app_status.deployments["GangConstructorFailure"].status == "DEPLOY_FAILED"
|
||
)
|
||
|
||
def test_partial_constructor_failure(self, ray_shutdown):
|
||
"""Validates gang deployment where one replica consistently fails."""
|
||
ray.init(num_cpus=1)
|
||
serve.start()
|
||
|
||
failed_replica_store = FailedGangReplicaStore.remote()
|
||
|
||
@serve.deployment(
|
||
num_replicas=4,
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class GangPartialConstructorFailure:
|
||
def __init__(self, store):
|
||
gang_id = serve.get_replica_context().gang_context.gang_id
|
||
is_first_fail = ray.get(store.mark_first_failing_gang.remote(gang_id))
|
||
if is_first_fail:
|
||
raise RuntimeError("Consistently throwing on same replica.")
|
||
should_fail = ray.get(store.mark_retry_failing_gang.remote(gang_id))
|
||
if should_fail:
|
||
raise RuntimeError("Keep failing the replica")
|
||
|
||
async def __call__(self, request):
|
||
return "hi"
|
||
|
||
serve._run(
|
||
GangPartialConstructorFailure.bind(failed_replica_store),
|
||
_blocking=False,
|
||
)
|
||
|
||
deployment_name = "GangPartialConstructorFailure"
|
||
|
||
def _one_gang_running_and_updating() -> bool:
|
||
app_status = serve.status().applications[SERVE_DEFAULT_APP_NAME]
|
||
dep = app_status.deployments[deployment_name]
|
||
return (
|
||
dep.replica_states.get("RUNNING", 0) == 2 and dep.status == "UPDATING"
|
||
)
|
||
|
||
wait_for_condition(_one_gang_running_and_updating, timeout=30)
|
||
|
||
# Wait well past the failed-to-start threshold
|
||
# (max(num_replicas * 3, 6) = 12 for 4 replicas) to prove the
|
||
# deployment stays stuck in UPDATING.
|
||
def _enough_retries_and_still_stable() -> bool:
|
||
failed_gangs = ray.get(failed_replica_store.get_failed_gang_count.remote())
|
||
return failed_gangs >= 15 and _one_gang_running_and_updating()
|
||
|
||
wait_for_condition(_enough_retries_and_still_stable, timeout=90)
|
||
assert serve.status().applications[SERVE_DEFAULT_APP_NAME].status == "DEPLOYING"
|
||
|
||
def test_transient_constructor_failure(self, ray_shutdown):
|
||
"""Validates gang deployment where the first constructor call fails then succeeds."""
|
||
ray.init(num_cpus=1)
|
||
serve.start()
|
||
|
||
failed_replica_store = FailedGangReplicaStore.remote()
|
||
|
||
@serve.deployment(
|
||
num_replicas=4,
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class GangTransientConstructorFailure:
|
||
def __init__(self, store):
|
||
gang_id = serve.get_replica_context().gang_context.gang_id
|
||
is_first_fail = ray.get(store.mark_first_failing_gang.remote(gang_id))
|
||
if is_first_fail:
|
||
raise RuntimeError("Intentionally throw on first try.")
|
||
|
||
async def __call__(self, request):
|
||
return "hi"
|
||
|
||
serve.run(GangTransientConstructorFailure.bind(failed_replica_store))
|
||
|
||
client = serve.context._get_global_client()
|
||
deployment_id = DeploymentID(name="GangTransientConstructorFailure")
|
||
deployment_dict = ray.get(client._controller._all_running_replicas.remote())
|
||
assert len(deployment_dict[deployment_id]) == 4
|
||
app_status = serve.status().applications[SERVE_DEFAULT_APP_NAME]
|
||
assert app_status.status == "RUNNING"
|
||
assert (
|
||
app_status.deployments["GangTransientConstructorFailure"].status
|
||
== "HEALTHY"
|
||
)
|
||
|
||
|
||
class TestGangFailureRecovery:
|
||
def test_startup_failure_stops_entire_gang(self, ray_shutdown):
|
||
"""Startup failure stops both replicas in the affected gang."""
|
||
ray.init(num_cpus=1)
|
||
serve.start()
|
||
failed_replica_store = FailedGangReplicaStore.remote()
|
||
recovery_signal = SignalActor.remote()
|
||
|
||
@serve.deployment(
|
||
num_replicas=4,
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class StartupFailureDeployment:
|
||
def __init__(self, failed_replica_store, recovery_signal):
|
||
gang_id = serve.get_replica_context().gang_context.gang_id
|
||
is_first_failure = ray.get(
|
||
failed_replica_store.mark_first_failing_gang.remote(gang_id)
|
||
)
|
||
if is_first_failure:
|
||
raise RuntimeError("Fail one startup to trigger gang cleanup.")
|
||
should_hold = ray.get(
|
||
failed_replica_store.mark_retry_failing_gang.remote(gang_id)
|
||
)
|
||
if should_hold:
|
||
# Hold failed replica retry until the intermediate state is asserted.
|
||
ray.get(recovery_signal.wait.remote())
|
||
|
||
def __call__(self):
|
||
ctx = serve.get_replica_context()
|
||
gc = ctx.gang_context
|
||
return {
|
||
"replica_id": ctx.replica_id.unique_id,
|
||
"gang_id": gc.gang_id,
|
||
}
|
||
|
||
app_name = "gang_startup_cleanup_app"
|
||
deployment_name = "StartupFailureDeployment"
|
||
handle = serve._run(
|
||
StartupFailureDeployment.bind(failed_replica_store, recovery_signal),
|
||
name=app_name,
|
||
_blocking=False,
|
||
)
|
||
|
||
# The unaffected gang should reach 2 RUNNING while the failed
|
||
# gang is being cleaned up and retried.
|
||
wait_for_condition(
|
||
lambda: (
|
||
serve.status()
|
||
.applications[app_name]
|
||
.deployments[deployment_name]
|
||
.replica_states.get("RUNNING", 0)
|
||
== 2
|
||
),
|
||
timeout=60,
|
||
)
|
||
|
||
# The 2 running replicas must belong to the SAME gang,
|
||
# proving no partial gang survived. Intentionally handle-based: this is
|
||
# a single-node cluster (ray.init(num_cpus=1)), so every replica is
|
||
# local to the caller and locality-aware routing still reaches all of
|
||
# them.
|
||
contexts = {}
|
||
for _ in range(50):
|
||
result = handle.remote().result()
|
||
contexts.setdefault(result["replica_id"], result)
|
||
if len(contexts) == 2:
|
||
break
|
||
assert len(contexts) == 2
|
||
assert len({ctx["gang_id"] for ctx in contexts.values()}) == 1
|
||
|
||
# Release constructor retry gate so the failed gang can recover.
|
||
ray.get(recovery_signal.send.remote())
|
||
|
||
# After retry, all 4 replicas should be RUNNING.
|
||
wait_for_condition(check_apps_running, apps=[app_name], timeout=60)
|
||
app_status = serve.status().applications[app_name]
|
||
dep_status = app_status.deployments[deployment_name]
|
||
assert dep_status.replica_states.get("RUNNING", 0) == 4
|
||
|
||
serve.delete(app_name)
|
||
serve.shutdown()
|
||
|
||
def test_health_failure_restarts_gang(self, ray_shutdown):
|
||
"""Single health check failure tears down and restarts the entire gang."""
|
||
ray.init(num_cpus=1)
|
||
serve.start()
|
||
target_replica_collector = Accumulator.remote()
|
||
|
||
@serve.deployment(
|
||
num_replicas=4,
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
health_check_period_s=1,
|
||
health_check_timeout_s=1,
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class HealthFailureDeployment:
|
||
def __call__(self):
|
||
ctx = serve.get_replica_context()
|
||
gc = ctx.gang_context
|
||
return {
|
||
"replica_id": ctx.replica_id.unique_id,
|
||
"gang_id": gc.gang_id,
|
||
}
|
||
|
||
def check_health(self):
|
||
targets = ray.get(target_replica_collector.get.remote())
|
||
if not targets:
|
||
return
|
||
target_id = targets[-1]
|
||
# Only 1 replica fails; its sibling stays healthy.
|
||
# The gang-aware cleanup must stop the sibling too.
|
||
ctx = serve.get_replica_context()
|
||
if ctx.replica_id.unique_id == target_id:
|
||
raise RuntimeError("Intentional health check failure.")
|
||
|
||
app_name = "gang_health_failure_app"
|
||
deployment_name = "HealthFailureDeployment"
|
||
handle = serve.run(HealthFailureDeployment.bind(), name=app_name)
|
||
wait_for_condition(check_apps_running, apps=[app_name], timeout=60)
|
||
|
||
# Discover all 4 replica contexts. Intentionally handle-based: this is
|
||
# a single-node cluster (ray.init(num_cpus=1)), so every replica is
|
||
# local to the caller and locality-aware routing still reaches all of
|
||
# them (unlike the multi-node placement checks that read controller
|
||
# state).
|
||
contexts_by_replica = {}
|
||
for _ in range(120):
|
||
result = handle.remote().result()
|
||
contexts_by_replica.setdefault(result["replica_id"], result)
|
||
if len(contexts_by_replica) == 4:
|
||
break
|
||
assert len(contexts_by_replica) == 4
|
||
|
||
# Pick 1 replica to fail health checks.
|
||
target_ctx = next(iter(contexts_by_replica.values()))
|
||
target_gang_id = target_ctx["gang_id"]
|
||
|
||
target_gang_replica_ids = {
|
||
ctx["replica_id"]
|
||
for ctx in contexts_by_replica.values()
|
||
if ctx["gang_id"] == target_gang_id
|
||
}
|
||
unaffected_replica_ids = (
|
||
set(contexts_by_replica.keys()) - target_gang_replica_ids
|
||
)
|
||
assert len(target_gang_replica_ids) == 2
|
||
assert len(unaffected_replica_ids) == 2
|
||
|
||
# Trigger failure for only 1 replica in the target gang.
|
||
ray.get(target_replica_collector.add.remote(target_ctx["replica_id"]))
|
||
|
||
client = serve.context._get_global_client()
|
||
deployment_id = DeploymentID(name=deployment_name, app_name=app_name)
|
||
|
||
def check_target_gang_restarted():
|
||
replicas = ray.get(
|
||
client._controller._dump_replica_states_for_testing.remote(
|
||
deployment_id
|
||
)
|
||
)
|
||
running_replicas = replicas.get([ReplicaState.RUNNING])
|
||
running_ids = {r.replica_id.unique_id for r in running_replicas}
|
||
# Both old gang members must be gone (not just the one that
|
||
# failed), and the unaffected gang must be untouched.
|
||
return (
|
||
len(running_ids) == 4
|
||
and len(running_ids & target_gang_replica_ids) == 0
|
||
and len(running_ids & unaffected_replica_ids) == 2
|
||
)
|
||
|
||
wait_for_condition(check_target_gang_restarted, timeout=90)
|
||
wait_for_condition(check_apps_running, apps=[app_name], timeout=60)
|
||
serve.delete(app_name)
|
||
serve.shutdown()
|
||
|
||
|
||
class TestGangChildSpawnPlacementGroup:
|
||
@ray.remote(num_cpus=0.1)
|
||
class ChildActor:
|
||
def get_pg(self):
|
||
return get_current_placement_group()
|
||
|
||
@ray.remote(num_cpus=0)
|
||
def child_task_get_pg():
|
||
return get_current_placement_group()
|
||
|
||
@pytest.mark.parametrize("child_type", ["actor", "task"])
|
||
def test_child_in_gang_pg(self, ray_cluster, child_type):
|
||
"""Spawn a child actor/task inside a gang replica and verify it shares the gang placement group."""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=2)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
ChildActor = TestGangChildSpawnPlacementGroup.ChildActor
|
||
child_task_get_pg = TestGangChildSpawnPlacementGroup.child_task_get_pg
|
||
|
||
@serve.deployment(
|
||
num_replicas=2,
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
# Extra bundle per replica so the child actor has resources
|
||
# inside the gang PG (the first bundle is consumed by the replica).
|
||
placement_group_bundles=[{"CPU": 0.1}, {"CPU": 0.1}],
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class GangWithChild:
|
||
def test_child_in_pg(self):
|
||
parent_pg = get_current_placement_group()
|
||
if child_type == "actor":
|
||
child = ChildActor.remote()
|
||
child_pg = ray.get(child.get_pg.remote())
|
||
else:
|
||
child_pg = ray.get(child_task_get_pg.remote())
|
||
return {
|
||
"parent_pg_id": parent_pg.id.hex() if parent_pg else None,
|
||
"child_pg_id": child_pg.id.hex() if child_pg else None,
|
||
}
|
||
|
||
def __call__(self):
|
||
return "ok"
|
||
|
||
app_name = "gang_child_app"
|
||
handle = serve.run(GangWithChild.bind(), name=app_name)
|
||
wait_for_condition(check_apps_running, apps=[app_name])
|
||
|
||
for _ in range(20):
|
||
result = handle.test_child_in_pg.remote().result()
|
||
assert result["parent_pg_id"] is not None
|
||
assert result["child_pg_id"] is not None
|
||
assert result["child_pg_id"] == result["parent_pg_id"]
|
||
|
||
serve.delete(app_name)
|
||
serve.shutdown()
|
||
|
||
def test_child_actor_gang_pg_bundles_bounded(self, ray_cluster):
|
||
"""Gang replicas with placement_group_bundles: verify child actors are resource-bounded by the gang PG."""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=2)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
ChildActor = TestGangChildSpawnPlacementGroup.ChildActor
|
||
|
||
@serve.deployment(
|
||
num_replicas=1,
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
# Replica consumes the first bundle (0.1 CPU). Worker bundle (0.1
|
||
# CPU) fits exactly one ChildActor, so a second child is blocked.
|
||
placement_group_bundles=[{"CPU": 0.1}, {"CPU": 0.1}],
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=1),
|
||
)
|
||
class GangWithBundlesAndChild:
|
||
def test_second_worker_blocked(self):
|
||
"""The second child actor shouldn't fit in this replica's bundle slice."""
|
||
w1 = ChildActor.remote()
|
||
w2 = ChildActor.remote()
|
||
ready, _ = ray.wait([w2.get_pg.remote()], timeout=1)
|
||
ray.kill(w1)
|
||
ray.kill(w2)
|
||
return len(ready) == 0
|
||
|
||
def __call__(self):
|
||
return "ok"
|
||
|
||
app_name = "gang_bundles_child_app"
|
||
handle = serve.run(GangWithBundlesAndChild.bind(), name=app_name)
|
||
wait_for_condition(check_apps_running, apps=[app_name])
|
||
|
||
# Verify resource limits are enforced within the gang PG bundle slice.
|
||
for _ in range(4):
|
||
assert handle.test_second_worker_blocked.remote().result() is True
|
||
|
||
serve.delete(app_name)
|
||
serve.shutdown()
|
||
|
||
def test_child_actor_opt_out_gang_pg(self, ray_cluster):
|
||
"""Verify a child actor can opt out of the gang PG by passing placement_group=None."""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=2)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
ChildActor = TestGangChildSpawnPlacementGroup.ChildActor
|
||
|
||
@serve.deployment(
|
||
num_replicas=2,
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class GangWithEscapedChild:
|
||
def get_child_outside_pg(self):
|
||
parent_pg = get_current_placement_group()
|
||
child = ChildActor.options(
|
||
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
||
placement_group=None, # Explicitly schedule outside the placement group
|
||
)
|
||
).remote()
|
||
child_pg = ray.get(child.get_pg.remote())
|
||
return {
|
||
"parent_pg_id": parent_pg.id.hex() if parent_pg else None,
|
||
"child_pg_id": child_pg.id.hex() if child_pg else None,
|
||
}
|
||
|
||
def __call__(self):
|
||
return "ok"
|
||
|
||
app_name = "gang_escaped_child_app"
|
||
handle = serve.run(GangWithEscapedChild.bind(), name=app_name)
|
||
wait_for_condition(check_apps_running, apps=[app_name])
|
||
|
||
for _ in range(20):
|
||
result = handle.get_child_outside_pg.remote().result()
|
||
assert result["parent_pg_id"] is not None
|
||
assert result["child_pg_id"] is None
|
||
|
||
serve.delete(app_name)
|
||
serve.shutdown()
|
||
|
||
|
||
class TestGangControllerRecovery:
|
||
def test_gang_context_recovery(self, ray_cluster):
|
||
"""Verifies that the controller recovers all app and deployment states
|
||
after a crash, including gang_context for gang deployments and normal
|
||
replicas for non-gang deployments.
|
||
"""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment(
|
||
num_replicas=4,
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class Gang1:
|
||
def __call__(self):
|
||
return "ok"
|
||
|
||
@serve.deployment(
|
||
num_replicas=2,
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class Gang2:
|
||
def __call__(self):
|
||
return "ok"
|
||
|
||
@serve.deployment(
|
||
num_replicas=2,
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
)
|
||
class NoGang:
|
||
def __call__(self):
|
||
return "ok"
|
||
|
||
app_names = ["gang_app1", "gang_app2", "no_gang_app"]
|
||
serve.run(Gang1.bind(), name="gang_app1", route_prefix="/gang1")
|
||
serve.run(Gang2.bind(), name="gang_app2", route_prefix="/gang2")
|
||
serve.run(NoGang.bind(), name="no_gang_app", route_prefix="/no_gang")
|
||
wait_for_condition(check_apps_running, apps=app_names)
|
||
|
||
gang_deployment_ids = [
|
||
DeploymentID(name="Gang1", app_name="gang_app1"),
|
||
DeploymentID(name="Gang2", app_name="gang_app2"),
|
||
]
|
||
no_gang_deployment_id = DeploymentID(name="NoGang", app_name="no_gang_app")
|
||
controller = serve.context._get_global_client()._controller
|
||
|
||
# Record controller-side gang_context before crash
|
||
gang_ctx_before = {}
|
||
for dep_id in gang_deployment_ids:
|
||
replicas = ray.get(
|
||
controller._dump_replica_states_for_testing.remote(dep_id)
|
||
)
|
||
running = replicas.get([ReplicaState.RUNNING])
|
||
for r in running:
|
||
assert r.gang_context is not None
|
||
gang_ctx_before[r.replica_id.unique_id] = r.gang_context
|
||
|
||
# Record non-gang replica count
|
||
no_gang_replicas = ray.get(
|
||
controller._dump_replica_states_for_testing.remote(no_gang_deployment_id)
|
||
)
|
||
no_gang_count_before = len(no_gang_replicas.get([ReplicaState.RUNNING]))
|
||
assert no_gang_count_before == 2
|
||
|
||
# Kill the controller and wait for recovery of all apps
|
||
ray.kill(controller, no_restart=False)
|
||
wait_for_condition(check_apps_running, apps=app_names, timeout=60)
|
||
|
||
new_controller = serve.context._get_global_client()._controller
|
||
|
||
def all_states_recovered():
|
||
# Verify gang_context recovered for all gang deployments
|
||
for dep_id in gang_deployment_ids:
|
||
replicas = ray.get(
|
||
new_controller._dump_replica_states_for_testing.remote(dep_id)
|
||
)
|
||
running = replicas.get([ReplicaState.RUNNING])
|
||
for r in running:
|
||
before = gang_ctx_before.get(r.replica_id.unique_id)
|
||
if r.gang_context is None or r.gang_context != before:
|
||
return False
|
||
|
||
# Verify non-gang deployment recovered
|
||
replicas = ray.get(
|
||
new_controller._dump_replica_states_for_testing.remote(
|
||
no_gang_deployment_id
|
||
)
|
||
)
|
||
if len(replicas.get([ReplicaState.RUNNING])) != no_gang_count_before:
|
||
return False
|
||
|
||
return True
|
||
|
||
wait_for_condition(all_states_recovered, timeout=60)
|
||
|
||
# Verify application and deployment statuses after recovery
|
||
status = serve.status()
|
||
for app_name in app_names:
|
||
app_status = status.applications[app_name]
|
||
assert app_status.status == "RUNNING"
|
||
for dep_name, dep_status in app_status.deployments.items():
|
||
assert dep_status.status == "HEALTHY"
|
||
|
||
for app_name in app_names:
|
||
serve.delete(app_name)
|
||
serve.shutdown()
|
||
|
||
@pytest.mark.parametrize("same_gang", [True, False])
|
||
def test_gang_replica_crash_during_controller_downtime(
|
||
self, ray_cluster, same_gang
|
||
):
|
||
"""When gang replicas crash while the controller is down, the controller
|
||
recovers and reschedules the affected gangs.
|
||
"""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=2)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment(
|
||
num_replicas=4,
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
health_check_period_s=1,
|
||
)
|
||
class GangApp:
|
||
def __call__(self):
|
||
return os.getpid()
|
||
|
||
app_name = "gang_crash_app"
|
||
dep_id = DeploymentID(name="GangApp", app_name=app_name)
|
||
serve.run(GangApp.bind(), name=app_name)
|
||
wait_for_condition(check_apps_running, apps=[app_name])
|
||
|
||
controller = serve.context._get_global_client()._controller
|
||
|
||
# Record initial replicas and group by gang.
|
||
replicas = ray.get(controller._dump_replica_states_for_testing.remote(dep_id))
|
||
running = replicas.get([ReplicaState.RUNNING])
|
||
assert len(running) == 4
|
||
|
||
gangs = {}
|
||
for r in running:
|
||
gangs.setdefault(r.gang_context.gang_id, []).append(r)
|
||
gang_ids = list(gangs.keys())
|
||
assert len(gang_ids) == 2
|
||
|
||
# Pick 2 victims
|
||
if same_gang:
|
||
victims = gangs[gang_ids[0]]
|
||
else:
|
||
victims = [gangs[gang_ids[0]][0], gangs[gang_ids[1]][0]]
|
||
|
||
victim_ids = {v.replica_id.unique_id for v in victims}
|
||
|
||
# Record the controller pid so we can confirm it actually restarts.
|
||
# ray.kill(..., no_restart=False) is asynchronous, so without this wait
|
||
# the checks below can read stale pre-crash state from the still-alive
|
||
# old controller (which still lists the victims as RUNNING).
|
||
original_controller_pid = ray.get(controller.get_pid.remote())
|
||
|
||
# Kill the controller, then kill the victims while it is down.
|
||
ray.kill(controller, no_restart=False)
|
||
for v in victims:
|
||
handle = ray.get_actor(v.replica_id.to_full_id_str(), namespace="serve")
|
||
ray.kill(handle, no_restart=True)
|
||
|
||
# Wait for the controller process to actually restart before checking
|
||
# recovery, otherwise we may observe the old controller's stale state.
|
||
def controller_restarted():
|
||
try:
|
||
pid = ray.get(controller.get_pid.remote(), timeout=5)
|
||
return pid != original_controller_pid
|
||
except Exception:
|
||
return False
|
||
|
||
wait_for_condition(controller_restarted, timeout=60)
|
||
wait_for_condition(check_apps_running, apps=[app_name], timeout=60)
|
||
|
||
new_controller = serve.context._get_global_client()._controller
|
||
|
||
# The affected gangs must be fully rescheduled: 4 RUNNING replicas, all
|
||
# with gang_context, and none of them the killed victims. Folding the
|
||
# victim check into the wait avoids racing the controller's reconcile.
|
||
def recovered_without_victims():
|
||
replicas = ray.get(
|
||
new_controller._dump_replica_states_for_testing.remote(dep_id)
|
||
)
|
||
running = replicas.get([ReplicaState.RUNNING])
|
||
if len(running) != 4:
|
||
return False
|
||
running_ids = {r.replica_id.unique_id for r in running}
|
||
return victim_ids.isdisjoint(running_ids) and all(
|
||
r.gang_context is not None for r in running
|
||
)
|
||
|
||
wait_for_condition(recovered_without_victims, timeout=60)
|
||
|
||
serve.delete(app_name)
|
||
serve.shutdown()
|
||
|
||
|
||
class TestGangNodeFailure:
|
||
@pytest.mark.parametrize(
|
||
"strategy", [GangPlacementStrategy.SPREAD, GangPlacementStrategy.PACK]
|
||
)
|
||
def test_worker_node_failure_restarts_gang(self, ray_cluster, strategy):
|
||
"""Killing a node restarts the affected gang with no request
|
||
downtime (surviving gang keeps serving) and no leaked PGs.
|
||
"""
|
||
cluster = ray_cluster
|
||
# Head and workers each get 1 CPU. Each PG (2 × 0.5 CPU = 1.0 CPU)
|
||
# fills exactly one node for PACK, and SPREAD distributes one bundle
|
||
# per node. At most 1 PG fits on the head, so at least 1 PG must
|
||
# land on a worker, giving the test a killable target. The extra
|
||
# nodes provide recovery capacity after one worker is removed.
|
||
cluster.add_node(num_cpus=1)
|
||
workers = [cluster.add_node(num_cpus=1) for _ in range(3)]
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
num_replicas = 4
|
||
gang_size = 2
|
||
|
||
@serve.deployment(
|
||
num_replicas=num_replicas,
|
||
ray_actor_options={"num_cpus": 0.5},
|
||
gang_scheduling_config=GangSchedulingConfig(
|
||
gang_size=gang_size,
|
||
gang_placement_strategy=strategy,
|
||
),
|
||
health_check_period_s=1,
|
||
)
|
||
class GangApp:
|
||
def __call__(self):
|
||
return ray.get_runtime_context().get_node_id()
|
||
|
||
app_name = "node_kill_app"
|
||
dep_id = DeploymentID(name="GangApp", app_name=app_name)
|
||
handle = serve.run(GangApp.bind(), name=app_name)
|
||
wait_for_condition(check_apps_running, apps=[app_name])
|
||
|
||
controller = serve.context._get_global_client()._controller
|
||
replicas = ray.get(controller._dump_replica_states_for_testing.remote(dep_id))
|
||
running = replicas.get([ReplicaState.RUNNING])
|
||
assert len(running) == num_replicas
|
||
|
||
# Group replicas by gang
|
||
gangs = {}
|
||
for r in running:
|
||
gangs.setdefault(r.gang_context.gang_id, []).append(r)
|
||
|
||
# Pick a worker to kill that (a) hosts at least one gang member and
|
||
# (b) leaves at least one gang fully intact on the surviving nodes.
|
||
node_to_kill = None
|
||
affected_gangs = []
|
||
for worker in workers:
|
||
affected_gangs = [
|
||
gid
|
||
for gid, members in gangs.items()
|
||
if any(r.actor_node_id == worker.node_id for r in members)
|
||
]
|
||
surviving = [
|
||
gid
|
||
for gid, members in gangs.items()
|
||
if all(r.actor_node_id != worker.node_id for r in members)
|
||
]
|
||
if affected_gangs and surviving:
|
||
node_to_kill = worker
|
||
break
|
||
|
||
assert node_to_kill is not None
|
||
assert len(affected_gangs) > 0
|
||
|
||
# Continuously send requests in a background thread.
|
||
stop = threading.Event()
|
||
recovered = threading.Event()
|
||
errors_before_recovery = []
|
||
errors_after_recovery = []
|
||
successes = []
|
||
|
||
def send_requests():
|
||
while not stop.is_set():
|
||
try:
|
||
result = handle.remote().result()
|
||
successes.append(result)
|
||
except Exception as e:
|
||
if recovered.is_set():
|
||
errors_after_recovery.append(e)
|
||
else:
|
||
errors_before_recovery.append(e)
|
||
time.sleep(0.1)
|
||
|
||
sender = threading.Thread(target=send_requests, daemon=True)
|
||
sender.start()
|
||
time.sleep(1)
|
||
|
||
cluster.remove_node(node_to_kill)
|
||
|
||
expected_num_pgs = num_replicas // gang_size
|
||
|
||
def fully_recovered():
|
||
replicas = ray.get(
|
||
controller._dump_replica_states_for_testing.remote(dep_id)
|
||
)
|
||
running = replicas.get([ReplicaState.RUNNING])
|
||
if len(running) != num_replicas:
|
||
return False
|
||
for r in running:
|
||
if r.gang_context is None:
|
||
return False
|
||
# Verify PG count has converged: the old affected PG should be removed
|
||
# and the replacement PG should be created.
|
||
gang_pg_names = [
|
||
n
|
||
for n in get_all_live_placement_group_names()
|
||
if n.startswith(GANG_PG_NAME_PREFIX)
|
||
]
|
||
if len(gang_pg_names) != expected_num_pgs:
|
||
return False
|
||
return True
|
||
|
||
wait_for_condition(fully_recovered, timeout=60)
|
||
recovered.set()
|
||
|
||
# Wait for at least one post-recovery success
|
||
successes_at_recovery = len(successes)
|
||
wait_for_condition(
|
||
lambda: len(successes) > successes_at_recovery,
|
||
timeout=10,
|
||
)
|
||
stop.set()
|
||
sender.join(timeout=5)
|
||
|
||
# Requests may fail during the brief disruption window: the node
|
||
# is dead but the handle may still route to the dead replica actors
|
||
# until the controller detects the failure and restarts them.
|
||
# After full recovery, no errors should occur.
|
||
assert len(errors_after_recovery) == 0
|
||
assert len(successes) > 0
|
||
|
||
wait_for_condition(check_apps_running, apps=[app_name])
|
||
|
||
serve.delete(app_name)
|
||
serve.shutdown()
|
||
|
||
|
||
class TestGangScaling:
|
||
@pytest.mark.parametrize(
|
||
"initial_num_replicas, final_num_replicas",
|
||
[
|
||
(4, 2), # Manual downscale: serve deploy num_replicas = 4 -> 2
|
||
(8, 4), # Downscaling
|
||
(4, 8), # Upscaling
|
||
],
|
||
)
|
||
def test_scale_gang_boundary(
|
||
self, ray_cluster, initial_num_replicas, final_num_replicas
|
||
):
|
||
"""Validates that scaling preserves complete gangs."""
|
||
GANG_SIZE = 2
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment(
|
||
name="D",
|
||
num_replicas=initial_num_replicas,
|
||
ray_actor_options={"num_cpus": 0.25},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=GANG_SIZE),
|
||
)
|
||
class D:
|
||
def __call__(self):
|
||
ctx = ray.serve.context._get_internal_replica_context()
|
||
gc = ctx.gang_context
|
||
return {"pid": os.getpid(), "gang_id": gc.gang_id if gc else None}
|
||
|
||
D = D.options(_internal=True, version="v1")
|
||
handle = serve.run(D.bind(), name="app")
|
||
wait_for_condition(check_apps_running, apps=["app"])
|
||
|
||
initial_num_gangs = initial_num_replicas // GANG_SIZE
|
||
deployment_id = DeploymentID(name="D", app_name="app")
|
||
|
||
initial_running = _get_running_replicas(deployment_id)
|
||
assert len(initial_running) == initial_num_replicas
|
||
initial_gang_ids = _get_gang_ids_from_running(initial_running)
|
||
assert len(initial_gang_ids) == initial_num_gangs
|
||
|
||
# Monitor requests during scaling to ensure zero downtime
|
||
errors, successes = [], []
|
||
stop_event = threading.Event()
|
||
|
||
def send_requests():
|
||
while not stop_event.is_set():
|
||
try:
|
||
handle.remote().result(timeout_s=10)
|
||
successes.append(True)
|
||
except Exception as e:
|
||
errors.append(str(e))
|
||
time.sleep(0.1)
|
||
|
||
t = threading.Thread(target=send_requests, daemon=True)
|
||
t.start()
|
||
|
||
# Scale to the final replica count.
|
||
handle = serve.run(
|
||
D.options(num_replicas=final_num_replicas).bind(), name="app"
|
||
)
|
||
wait_for_condition(check_apps_running, apps=["app"])
|
||
|
||
deployment = list(serve.status().applications["app"].deployments.values())[0]
|
||
assert deployment.replica_states.get("RUNNING", 0) == final_num_replicas
|
||
|
||
stop_event.set()
|
||
t.join(timeout=5)
|
||
|
||
# Scaling should be zero-downtime: no requests should fail.
|
||
assert len(errors) == 0
|
||
assert len(successes) > 0
|
||
|
||
final_num_gangs = final_num_replicas // GANG_SIZE
|
||
|
||
final_running = _get_running_replicas(deployment_id)
|
||
assert len(final_running) == final_num_replicas
|
||
final_gang_ids = _get_gang_ids_from_running(final_running)
|
||
assert len(final_gang_ids) == final_num_gangs
|
||
|
||
smaller, larger = sorted([initial_gang_ids, final_gang_ids], key=len)
|
||
assert smaller.issubset(larger)
|
||
|
||
serve.delete("app")
|
||
serve.shutdown()
|
||
|
||
|
||
class TestGangRollingUpdate:
|
||
def test_rolling_update(self, ray_cluster):
|
||
"""Verifies that rolling update replaces complete gangs atomically.
|
||
|
||
During the update, RUNNING replicas must always form complete gangs.
|
||
"""
|
||
GANG_SIZE = 2
|
||
NUM_REPLICAS = 4
|
||
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment(
|
||
name="D",
|
||
num_replicas=NUM_REPLICAS,
|
||
ray_actor_options={"num_cpus": 0.25},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=GANG_SIZE),
|
||
)
|
||
class V1:
|
||
def __call__(self):
|
||
return "v1"
|
||
|
||
handle = serve.run(V1.bind(), name="app")
|
||
wait_for_condition(check_apps_running, apps=["app"])
|
||
assert handle.remote().result() == "v1"
|
||
|
||
client = _get_global_client()
|
||
controller = client._controller
|
||
deployment_id = DeploymentID(name="D", app_name="app")
|
||
|
||
# Collect initial gang_ids.
|
||
replicas = ray.get(
|
||
controller._dump_replica_states_for_testing.remote(deployment_id)
|
||
)
|
||
running = replicas.get([ReplicaState.RUNNING])
|
||
assert len(running) == NUM_REPLICAS
|
||
initial_gang_ids = {r.gang_context.gang_id for r in running}
|
||
assert len(initial_gang_ids) == NUM_REPLICAS // GANG_SIZE
|
||
|
||
# Gate V2 startup behind a signal so we can deterministically
|
||
# observe mixed old/new gang state during the rolling update.
|
||
signal = SignalActor.remote()
|
||
|
||
# New code version triggers requires_actor_restart -> rolling update
|
||
@serve.deployment(
|
||
name="D",
|
||
num_replicas=NUM_REPLICAS,
|
||
ray_actor_options={"num_cpus": 0.25},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=GANG_SIZE),
|
||
)
|
||
class V2:
|
||
def __init__(self):
|
||
ray.get(signal.wait.remote())
|
||
|
||
def __call__(self):
|
||
return "v2"
|
||
|
||
# Issue the update without blocking so we can poll
|
||
# intermediate controller state from the main thread.
|
||
serve._run(V2.bind(), name="app", _blocking=False)
|
||
|
||
# Wait until we observe mixed state: at least one old gang still
|
||
# RUNNING and at least one new gang in STARTING (blocked on signal).
|
||
def mixed_state_observed():
|
||
replicas = ray.get(
|
||
controller._dump_replica_states_for_testing.remote(deployment_id)
|
||
)
|
||
running = replicas.get([ReplicaState.RUNNING])
|
||
starting = replicas.get([ReplicaState.STARTING])
|
||
|
||
running_gang_ids = {
|
||
r.gang_context.gang_id for r in running if r.gang_context is not None
|
||
}
|
||
starting_gang_ids = {
|
||
r.gang_context.gang_id for r in starting if r.gang_context is not None
|
||
}
|
||
|
||
has_old_running = bool(running_gang_ids & initial_gang_ids)
|
||
has_new_starting = bool(starting_gang_ids - initial_gang_ids)
|
||
|
||
# While old gangs are still running, they must be complete.
|
||
gang_counts: dict = {}
|
||
for r in running:
|
||
if r.gang_context is not None:
|
||
gid = r.gang_context.gang_id
|
||
gang_counts[gid] = gang_counts.get(gid, 0) + 1
|
||
for gid, count in gang_counts.items():
|
||
if gid in initial_gang_ids:
|
||
assert count == GANG_SIZE
|
||
|
||
return has_old_running and has_new_starting
|
||
|
||
wait_for_condition(mixed_state_observed, timeout=30)
|
||
|
||
# Unblock V2 constructors and wait for the update to finish.
|
||
signal.send.remote()
|
||
|
||
def update_complete():
|
||
replicas = ray.get(
|
||
controller._dump_replica_states_for_testing.remote(deployment_id)
|
||
)
|
||
running = replicas.get([ReplicaState.RUNNING])
|
||
if len(running) != NUM_REPLICAS:
|
||
return False
|
||
current_gang_ids = {
|
||
r.gang_context.gang_id for r in running if r.gang_context is not None
|
||
}
|
||
return current_gang_ids and not (current_gang_ids & initial_gang_ids)
|
||
|
||
wait_for_condition(update_complete, timeout=60)
|
||
|
||
# Confirm all replicas serve the new version.
|
||
for _ in range(20):
|
||
assert handle.remote().result() == "v2"
|
||
|
||
serve.delete("app")
|
||
serve.shutdown()
|
||
|
||
|
||
class TestGangAutoscaling:
|
||
def test_gang_autoscaling(self, ray_cluster):
|
||
"""Verifies that autoscaling with gang scheduling scales in complete gangs."""
|
||
GANG_SIZE = 2
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=2)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
signal = SignalActor.remote()
|
||
|
||
@serve.deployment(
|
||
num_replicas="auto",
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=GANG_SIZE),
|
||
autoscaling_config={
|
||
"min_replicas": 2,
|
||
"max_replicas": 8,
|
||
# Lower delays/windows so the test observes scaling within seconds
|
||
"upscale_delay_s": 0.1,
|
||
"downscale_delay_s": 0.1,
|
||
"look_back_period_s": 0.1,
|
||
"target_ongoing_requests": 1,
|
||
},
|
||
max_ongoing_requests=20,
|
||
)
|
||
class GangAutoscale:
|
||
async def __call__(self):
|
||
await signal.wait.remote()
|
||
return os.getpid()
|
||
|
||
handle = serve.run(GangAutoscale.bind(), name="gang_autoscale_app")
|
||
wait_for_condition(check_apps_running, apps=["gang_autoscale_app"])
|
||
|
||
wait_for_condition(
|
||
check_num_replicas_eq,
|
||
name="GangAutoscale",
|
||
target=2,
|
||
app_name="gang_autoscale_app",
|
||
use_controller=True,
|
||
)
|
||
|
||
# Send enough requests to trigger upscaling
|
||
results = [handle.remote() for _ in range(20)]
|
||
|
||
# Wait for scale-up to 8 replicas (4 complete gangs).
|
||
wait_for_condition(
|
||
check_num_replicas_eq,
|
||
name="GangAutoscale",
|
||
target=8,
|
||
app_name="gang_autoscale_app",
|
||
timeout=60,
|
||
use_controller=True,
|
||
)
|
||
|
||
# Replica count should always be a multiple of gang_size
|
||
deployment = (
|
||
serve.status()
|
||
.applications["gang_autoscale_app"]
|
||
.deployments["GangAutoscale"]
|
||
)
|
||
running = deployment.replica_states.get("RUNNING")
|
||
assert running % GANG_SIZE == 0
|
||
|
||
# Release all requests to allow traffic to drain
|
||
signal.send.remote()
|
||
for res in results:
|
||
res.result()
|
||
|
||
# As the queue is drained, we should scale back down
|
||
wait_for_condition(
|
||
check_num_replicas_eq,
|
||
name="GangAutoscale",
|
||
target=2,
|
||
app_name="gang_autoscale_app",
|
||
timeout=60,
|
||
use_controller=True,
|
||
)
|
||
|
||
deployment = (
|
||
serve.status()
|
||
.applications["gang_autoscale_app"]
|
||
.deployments["GangAutoscale"]
|
||
)
|
||
running = deployment.replica_states.get("RUNNING")
|
||
assert running % GANG_SIZE == 0
|
||
|
||
serve.delete("gang_autoscale_app")
|
||
serve.shutdown()
|
||
|
||
def test_gang_autoscaling_unaligned_upscale(self, ray_cluster):
|
||
GANG_SIZE = 3
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=2)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
signal = SignalActor.remote()
|
||
|
||
@serve.deployment(
|
||
num_replicas="auto",
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=GANG_SIZE),
|
||
autoscaling_config={
|
||
"min_replicas": 3,
|
||
"max_replicas": 9,
|
||
"metrics_interval_s": 0.5,
|
||
"upscale_delay_s": 0.1,
|
||
"downscale_delay_s": 0.1,
|
||
"look_back_period_s": 1,
|
||
"target_ongoing_requests": 2,
|
||
},
|
||
max_ongoing_requests=50,
|
||
)
|
||
class UnalignedUpscale:
|
||
async def __call__(self):
|
||
await signal.wait.remote()
|
||
return os.getpid()
|
||
|
||
handle = serve.run(UnalignedUpscale.bind(), name="unaligned_upscale_app")
|
||
wait_for_condition(check_apps_running, apps=["unaligned_upscale_app"])
|
||
wait_for_condition(
|
||
check_num_replicas_eq,
|
||
name="UnalignedUpscale",
|
||
target=3,
|
||
app_name="unaligned_upscale_app",
|
||
use_controller=True,
|
||
)
|
||
|
||
# Send 9 blocking requests. With target_ongoing_requests=2:
|
||
# desired = ceil(9/2) = 5 (unaligned).
|
||
# Gang policy rounds up: ceil(5/3)*3 = 6.
|
||
results = [handle.remote() for _ in range(9)]
|
||
|
||
def upscaled_and_aligned():
|
||
deployment = (
|
||
serve.status()
|
||
.applications["unaligned_upscale_app"]
|
||
.deployments["UnalignedUpscale"]
|
||
)
|
||
running = deployment.replica_states.get("RUNNING", 0)
|
||
assert running == 6
|
||
return True
|
||
|
||
wait_for_condition(upscaled_and_aligned, timeout=60)
|
||
|
||
# Release all requests so the queue drains.
|
||
signal.send.remote()
|
||
for res in results:
|
||
res.result()
|
||
|
||
# The autoscaler should scale back down to min_replicas=3.
|
||
wait_for_condition(
|
||
check_num_replicas_eq,
|
||
name="UnalignedUpscale",
|
||
target=3,
|
||
app_name="unaligned_upscale_app",
|
||
timeout=60,
|
||
use_controller=True,
|
||
)
|
||
|
||
serve.delete("unaligned_upscale_app")
|
||
serve.shutdown()
|
||
|
||
def test_gang_autoscaling_unaligned_downscale(self, ray_cluster):
|
||
GANG_SIZE = 3
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=2)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
signal = SignalActor.remote()
|
||
|
||
@serve.deployment(
|
||
num_replicas="auto",
|
||
ray_actor_options={"num_cpus": 0.1},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=GANG_SIZE),
|
||
autoscaling_config={
|
||
"min_replicas": 6,
|
||
"max_replicas": 9,
|
||
"initial_replicas": 9,
|
||
"metrics_interval_s": 0.5,
|
||
"upscale_delay_s": 5,
|
||
# Must be long enough for all 9 gang-scheduled replicas to
|
||
# start before the autoscaler can trigger a downscale.
|
||
"downscale_delay_s": 20,
|
||
"look_back_period_s": 1,
|
||
"target_ongoing_requests": 2,
|
||
},
|
||
max_ongoing_requests=50,
|
||
)
|
||
class UnalignedDownscale:
|
||
async def __call__(self):
|
||
await signal.wait.remote()
|
||
return os.getpid()
|
||
|
||
handle = serve.run(UnalignedDownscale.bind(), name="unaligned_downscale_app")
|
||
wait_for_condition(check_apps_running, apps=["unaligned_downscale_app"])
|
||
wait_for_condition(
|
||
check_num_replicas_eq,
|
||
name="UnalignedDownscale",
|
||
target=9,
|
||
app_name="unaligned_downscale_app",
|
||
use_controller=True,
|
||
timeout=60,
|
||
)
|
||
|
||
# Send 10 blocking requests. With target_ongoing_requests=2:
|
||
# desired = ceil(10/2) = 5 (unaligned).
|
||
# Gang policy rounds up: ceil(5/3)*3 = 6.
|
||
results = [handle.remote() for _ in range(10)]
|
||
|
||
# Wait for downscale from 9 to 6.
|
||
def downscaled_and_aligned():
|
||
deployment = (
|
||
serve.status()
|
||
.applications["unaligned_downscale_app"]
|
||
.deployments["UnalignedDownscale"]
|
||
)
|
||
running = deployment.replica_states.get("RUNNING", 0)
|
||
assert running == 6
|
||
return True
|
||
|
||
wait_for_condition(downscaled_and_aligned, timeout=60)
|
||
|
||
# Release all requests so the queue drains.
|
||
signal.send.remote()
|
||
for res in results:
|
||
res.result()
|
||
|
||
serve.delete("unaligned_downscale_app")
|
||
serve.shutdown()
|
||
|
||
|
||
class TestGangMigration:
|
||
def test_gang_migration(self, ray_cluster):
|
||
"""Verifies that when a node drains, entire gangs migrate together."""
|
||
cluster = ray_cluster
|
||
cluster.add_node(num_cpus=1)
|
||
node_to_drain = cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
ray.init(address=cluster.address)
|
||
serve.start()
|
||
|
||
@serve.deployment(
|
||
name="D",
|
||
num_replicas=4,
|
||
ray_actor_options={"num_cpus": 0.25},
|
||
gang_scheduling_config=GangSchedulingConfig(gang_size=2),
|
||
)
|
||
class D:
|
||
def __call__(self):
|
||
ctx = ray.serve.context._get_internal_replica_context()
|
||
gc = ctx.gang_context
|
||
return {
|
||
"pid": os.getpid(),
|
||
"gang_id": gc.gang_id if gc else None,
|
||
"node_id": ray.get_runtime_context().get_node_id(),
|
||
}
|
||
|
||
D = D.options(_internal=True, version="v1")
|
||
serve.run(D.bind(), name="app")
|
||
wait_for_condition(check_apps_running, apps=["app"])
|
||
|
||
deployment_id = DeploymentID(name="D", app_name="app")
|
||
running = _get_running_replicas(deployment_id)
|
||
assert len(running) == 4
|
||
assert len(_get_gang_ids_from_running(running)) == 2
|
||
|
||
# Add another node for replicas to migrate to, then drain a node
|
||
cluster.add_node(num_cpus=1)
|
||
cluster.wait_for_nodes()
|
||
cluster.remove_node(node_to_drain)
|
||
wait_for_condition(check_apps_running, apps=["app"], timeout=120)
|
||
|
||
deployment = list(serve.status().applications["app"].deployments.values())[0]
|
||
assert deployment.replica_states.get("RUNNING", 0) == 4
|
||
|
||
def check_complete_gangs():
|
||
running = _get_running_replicas(deployment_id)
|
||
assert len(running) == 4
|
||
gang_ids = {
|
||
r.gang_context.gang_id for r in running if r.gang_context is not None
|
||
}
|
||
assert len(gang_ids) == 2
|
||
return True
|
||
|
||
wait_for_condition(check_complete_gangs, timeout=60)
|
||
|
||
serve.delete("app")
|
||
serve.shutdown()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(pytest.main(["-v", "-s", __file__]))
|