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__]))