[RFC]: Coordinator-Based Control Plane for Mooncake PG

## 1. Summary

Mooncake PG currently conflated communication outcomes with group membership: if a collective operation failed, the underlying TE connection was torn down, and the rank was immediately evicted from `active_ranks`. This decentralized, imperative approach made transient network failures indistinguishable from permanent node deaths, forcing upper-layer systems into heavy fault-recovery workflows and causing split-brain scenarios during network partitions.

This RFC proposes a control-plane design with three core changes:
1. **Decouple Execution State from Configuration:** Operations will return a rank-local `failed_ranks_hint` tensor, leaving the logical `active_ranks` configuration unchanged.
2. **Introduce Explicit Physical States:** Introduce `OFFLINE`, `SYNCED`, and `HEALTHY` to accurately reflect a rank's physical readiness.
3. **Centralized Coordinator:** Shift from decentralized mutation to a "single-point proposal, Coordinator broadcasts" model to guarantee strict global consistency.

Together, these changes make Mooncake PG robust, safe by default, and capable of supporting advanced resilience features, such as seamless in-place rejoin.

## 2. Motivation

As we attempted to implement in-place rejoin to recover from transient network failures gracefully, we hit several fundamental roadblocks in the current architecture:

### 2.1 The Conflation of Configuration and State

Currently, `active_ranks` represents both *who should participate* (Configuration) and *who actually succeeded* (Execution State). By implicitly removing ranks from `active_ranks` on any RDMA failure, PG strips the upper layer of the opportunity to perform lightweight retries.

### 2.2 The Lack of Explicit Intermediate States

We realized that `peerConnected = 1` is insufficient to safely recover a rank. A rank might have a physical RDMA connection, but its metadata might be stale or incorrect. To safely execute an in-place rejoin, PG natively needs to understand:
* **Is the rank completely dead?** (`OFFLINE` - e.g., process crashed, cannot connect to Rank 0).
* **Is its metadata synced, but TE is not ready?** (`SYNCED` - e.g., RPC is connected, but RDMA paths are still warming up or partitioned).
* **Is it fully connected and ready for collectives?** (`HEALTHY`).

Without these intermediate states, recovering a rank safely becomes incredibly complex and error-prone.
### 2.3 The Split-Brain Danger of Decentralized Membership

If membership changes are decided independently by multiple ranks, what happens during a network partition?
* If Rank 1-4 want to deactivate Rank 7, but Rank 5 wants to deactivate Rank 6 and 7, who wins?
* If we require all ranks to agree, the system deadlocks.
* If we take the most conservative approach, a localized network drop might evict the entire cluster.

We need a central authority inside PG to validate these requests against the physical reality of the network and enforce a single source of truth.

## 3. Goals & Non-Goals

### Goals
* **Authoritative Source of Truth:** Guarantee consistent configuration. The Coordinator deterministically finds the maximum connected sub-graph and evicts the rest during partitions.
* **Unified Recovery Workflow:** Ensure that both full process replacement and in-place rejoin (recovering from a transient drop) share the exact same internal state machine logic.
* **Safe Single-Point Proposals:** Allow a single rank (or upper-layer instance) to propose an `activate`/`deactivate` operation. The Coordinator acts as a safety net, validating and broadcasting it to everyone.
* **Declarative Control Plane:** The Coordinator broadcasts states (e.g., "Rank X is OFFLINE"), rather than issuing imperative commands (e.g., "Tear down connection to X").

### Non-Goals

* **High Availability Coordinator:** For this first iteration, we implement a `CentralizedCoordinator` running on Rank 0. HA coordinators such as Raft-based variants are future work, but this RFC lays the exact architectural groundwork for it.

## 4. Core Concepts: State vs. Configuration

A major conceptual shift in this RFC is strictly separating *Physical State* (managed by the Coordinator) from *Logical Configuration* (managed by the User/Upper Layer).

### 4.1 Physical State: The "Healthy Set"

Maintained authoritatively by the Coordinator and mirrored by local Agents:
1. **`OFFLINE`:** The rank's control plane cannot communicate with Rank 0's Coordinator (RPC down or process dead).
2. **`SYNCED`:** The rank can communicate with Rank 0. Its metadata is fully synchronized, but the underlying TE (RDMA/NVLink/...) is not yet verified as globally functional.
3. **`HEALTHY`:** The rank is `SYNCED` *and* the Coordinator has verified it belongs to the globally connected "Healthy Set" (a fully connected sub-graph of TE links).

### 4.2 Logical Configuration: `active_ranks`

`active_ranks` is now strictly a membership configuration indicating who *should* participate in the collective.
* If a rank is `active` but temporarily becomes `unhealthy` (e.g., a transient link failure), PG leaves `active_ranks` alone (unless `auto_deactivate_on_failure=True`). The operation will return `failed_ranks_hint`, allowing upper layers to retry.
* An `inactive` rank can **only** be activated if its physical state is `HEALTHY`.

## 5. System Architecture

To realize this safely, the system is divided into three layers:

1. **Coordinator (Rank 0):** The central brain. It receives heartbeats and observations from agents, calculates the maximum connected TE sub-graph, and broadcasts `ViewUpdate`, `RankStateUpdate`, ...
2. **Agent (Every Rank):** A local control-plane client. It strictly follows the Coordinator's authoritative state.
3. **Data Plane (`MooncakeWorker` / `P2PProxy`):** Executes transfers. **Crucially, the data plane no longer modifies membership or tears down connections on failure.** It simply records `failed_ranks_hint` and passes them to the Agent.

### 5.1 Pure State Machine & Execution Host

Managing state consistency via conventional multithreading introduces high risks of race conditions and deadlocks. To mitigate this complexity, both the Coordinator and the Agent are architected as **Pure State Machines paired with a decoupled Execution Host**.
* The state machines contain zero threads, zero locks, and zero I/O. They take an event (e.g., a heartbeat) and return a list of `Effects` (e.g., Broadcast View).
* An "Execution Host" running on a single `SerializedExecutor` thread pushes events into the state machine, catches the resulting Effects, and executes the actual RPCs.
* *Benefit:* Beyond making the system highly testable, this functional design eliminates race conditions natively and provides strict determinism and replayability.

### 5.2 Barrier-style view updates

Membership changes are applied through a coordinated, two-phase proposal flow:

1. A rank or upper-layer component submits an `activate` or `deactivate` proposal.
2. The Coordinator validates the proposal against the current physical state.
3. The Coordinator broadcasts the new `GroupView`.
4. Relevant participants ACK the update.
5. If a participant times out during the ACK phase, the Coordinator marks it OFFLINE and commits the view for the remaining healthy ranks (reporting the timeout to the upper layer).

This provides a strict, deterministic barrier rather than best-effort propagation, ensuring that if an `activate` or `deactivate` call succeeds, PG is immediately and globally ready to execute collectives.

## 6. API Impact

### 6.1 Exposing `failed_ranks_hint`

Every `Work` object returned by a collective or P2P operation now carries a `failed_ranks_hint` tensor.

```python
from mooncake import pg
import torch.distributed as dist

work = dist.all_reduce(tensor, op=dist.ReduceOp.SUM, async_op=True)
work.wait()

# Returns a 1-D int32 tensor, shape = (max_group_size,)
# failed[i] == 1 means rank `i` failed in THIS specific operation.
failed = pg.get_failed_ranks_hint(work)
```

### 6.2 Explicit Membership & Control APIs

We introduce explicit APIs for the upper layer to manage membership and transport timeouts, fully decoupled from implicit failure handling.

* `pg.deactivate_ranks(group, ranks)`: Logically removes ranks from the active membership.
* `pg.activate_ranks(group, ranks)`: Adds ranks back to the active membership (alias for `recover_ranks`). The Coordinator guarantees this only succeeds if the target ranks are physically `HEALTHY`.
* Timeout Controls: `pg.set_collective_timeout_us(us)` and `pg.set_p2p_timeout_us(us)` to control data-plane liveness probes.

Legacy APIs like `pg.extend_group_size_to` are now deprecated. They will yield a warning log but have no side effects, as the new control plane inherently manages the underlying state synchronization.

## 7. Backward Compatibility

### 7.1 `auto_deactivate_on_failure`

To preserve backward compatibility with existing use cases, PG introduces an `auto_deactivate_on_failure` option (default: `True`).

* **Legacy Mode (`True`):** On failure, **Coordinator** automatically kicks the unhealthy rank out of `active_ranks` and broadcasts the new view. This exactly matches the previous behavior, ensuring legacy applications work unchanged, but now with global consistency guarantees.
* **Fine-Grained Mode (`False`):** PG only populates the `failed_ranks_hint` tensor. `active_ranks` is never mutated by communication failures. The upper layer takes full control over the retry or deactivate logic.

### 7.2 Single-Point Proposals & Deduplication

Under the old decentralized model, upper-layer systems were forced to execute `recover_ranks` calls on *every single rank* simultaneously to maintain a consistent view.

With the new control plane, a proposal from *any* single rank is sufficient. However, to maintain full backward compatibility, the Coordinator is designed to seamlessly handle redundant requests. If an existing application broadcasts an `activate_ranks` call across all 8 ranks, the Coordinator receives 8 identical proposals, safely **deduplicates** them, and issues a single `ViewUpdate` broadcast. This guarantees that existing recovery code works out-of-the-box while instantly benefiting from the new strict-consistency safety net.

## 8. Rationale

**Why not push ALL consistency responsibilities to the upper layer?**

For fault tolerance and inference elasticity, upper layers (like application frameworks) usually maintain their own control plane to ensure cross-rank consistency. If the upper layer already has a control plane, why build another one into PG?

An alternative architectural approach is to design PG purely as a stateless thin TE wrapper, delegating all consistency guarantees to the upper layer. However, this creates a fundamental mismatch: application frameworks are agnostic to physical TE connection states, yet robust fault tolerance and recovery inherently require this low-level visibility. Forcing upper layers to manage these low-level states would tightly couple the application framework to PG's internal implementations. This tight coupling degrades portability and raises the barrier to integration -- ultimately limiting PG's capability to serve as a clean, framework-agnostic communication primitives in the broader ecosystem.

We chose to build a dedicated control plane inside PG because:
1. **Robustness & Safety:** Relying on the upper layer to perfectly synchronize `failed_ranks_hint` and execute symmetric membership updates makes PG fragile. The Coordinator acts as a strict safety net, rejecting invalid, conflicting, or asymmetric requests from the upper layer.
2. **Simplified Recovery & Ease of Use:** In-place rejoin becomes straightforward. The PG control plane handles metadata synchronization and TE warmup automatically, drastically reducing the implementation burden on upper-layer recovery paths.
3. **Future-Proofing:** By establishing a hard boundary between the Control Plane and Data Plane now, we lay the necessary groundwork for implementing true, multi-node High Availability in the future.

## 9. Related PR

- #2338
- #3011
- #2455

#2338 first introduced the separation between per-operation communication failures and logical membership. It was subsequently reverted by #3011 because its per-operation failed-ranks allocation was incompatible with CUDA Graph capture and replay.

#2455 introduces the initial implementation of the control plane. It also resolves the CUDA Graph compatibility issue that led to #2338 being reverted by #3011.