Skip to main content

hyperactor_mesh/
mesh_controller.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use std::collections::HashMap;
10use std::collections::HashSet;
11use std::fmt::Debug;
12use std::time::SystemTime;
13
14use async_trait::async_trait;
15use hyperactor::Actor;
16use hyperactor::ActorAddr;
17use hyperactor::Context;
18use hyperactor::Endpoint as _;
19use hyperactor::Handler;
20use hyperactor::Instance;
21use hyperactor::RemoteEndpoint as _;
22use hyperactor::actor::ActorError;
23use hyperactor::actor::ActorErrorKind;
24use hyperactor::actor::ActorStatus;
25use hyperactor::actor::Referable;
26use hyperactor::actor::handle_undeliverable_message;
27use hyperactor::context;
28use hyperactor::kv_pairs;
29use hyperactor::mailbox::MessageEnvelope;
30use hyperactor::mailbox::RemoteMessage;
31use hyperactor::mailbox::Undeliverable;
32use hyperactor::mailbox::UndeliverableReason;
33use hyperactor::supervision::ActorSupervisionEvent;
34use hyperactor_config::CONFIG;
35use hyperactor_config::ConfigAttr;
36use hyperactor_config::Flattrs;
37use hyperactor_config::attrs::declare_attrs;
38use hyperactor_telemetry::declare_static_counter;
39use ndslice::ViewExt;
40use ndslice::view::CollectMeshExt;
41use ndslice::view::Point;
42use ndslice::view::Ranked;
43use opentelemetry::metrics::Counter;
44use serde::Deserialize;
45use serde::Serialize;
46use tokio::time::Duration;
47use typeuri::Named;
48
49use crate::ValueMesh;
50use crate::actor_mesh::ActorMeshRef;
51use crate::bootstrap::ProcStatus;
52use crate::casting::CAST_ACTOR_MESH_ID;
53use crate::casting::update_undeliverable_envelope_for_casting;
54use crate::mesh_id::ResourceId;
55use crate::proc_agent::ActorState;
56use crate::proc_agent::MESH_ORPHAN_TIMEOUT;
57use crate::proc_mesh::ProcMeshRef;
58use crate::resource;
59use crate::supervision::MeshFailure;
60use crate::supervision::Unhealthy;
61
62/// Actor name for `ActorMeshController` when spawned as a named child.
63pub const ACTOR_MESH_CONTROLLER_NAME: &str = "actor_mesh_controller";
64
65declare_attrs! {
66    /// Time between checks of actor states to create supervision events for
67    /// owners. The longer this is, the longer it will take to detect a failure
68    /// and report it to all subscribers; however, shorter intervals will send
69    /// more frequent messages and heartbeats just to see everything is still running.
70    /// The default is chosen to balance these two objectives.
71    /// This also controls how frequently the healthy heartbeat is sent out to
72    /// subscribers if there are no failures encountered.
73    @meta(CONFIG = ConfigAttr::new(
74        Some("HYPERACTOR_MESH_SUPERVISION_POLL_FREQUENCY".to_string()),
75        None,
76    ))
77    pub attr SUPERVISION_POLL_FREQUENCY: Duration = Duration::from_secs(10);
78}
79
80declare_static_counter!(
81    ACTOR_MESH_CONTROLLER_SUPERVISION_STALLS,
82    "actor.actor_mesh_controller.num_stalls"
83);
84
85declare_static_counter!(
86    PROC_MESH_CONTROLLER_SUPERVISION_STALLS,
87    "actor.proc_mesh_controller.num_stalls"
88);
89
90/// Aggregated health and subscriber bookkeeping for a single
91/// `ResourceController`. Tracks the most recently observed status of every
92/// rank in the controlled mesh, the latched unhealthy event (if any), the
93/// owner port (notified on failures), and the set of streaming subscribers
94/// (notified on both stop and failure events). The generation counter on
95/// each status entry provides last-writer-wins ordering between streamed
96/// and polled updates.
97#[derive(Debug)]
98pub struct HealthState {
99    /// The status of each rank in the controlled mesh, paired with the
100    /// generation counter from the most recent update. The generation is
101    /// used for last-writer-wins ordering between streamed and polled updates.
102    statuses: HashMap<Point, (resource::Status, u64)>,
103    /// The latched unhealthy event for the mesh, if any. Once set, this is
104    /// surfaced to new subscribers on subscribe and to `GetState` callers.
105    unhealthy_event: Option<Unhealthy>,
106    /// Per-rank supervision events for ranks that have crashed. Used to build
107    /// region-scoped failure reports.
108    crashed_ranks: HashMap<usize, ActorSupervisionEvent>,
109    /// The single owner of the controlled mesh, notified on failure events
110    /// (but not on clean stops).
111    owner: Option<hyperactor::PortRef<MeshFailure>>,
112    /// Streaming subscribers, notified on both stop and failure events as
113    /// well as periodic heartbeats.
114    subscribers: HashSet<hyperactor::PortRef<Option<MeshFailure>>>,
115}
116
117impl HealthState {
118    fn new(
119        statuses: HashMap<Point, resource::Status>,
120        owner: Option<hyperactor::PortRef<MeshFailure>>,
121    ) -> Self {
122        Self {
123            statuses: statuses
124                .into_iter()
125                .map(|(point, status)| (point, (status, 0)))
126                .collect(),
127            unhealthy_event: None,
128            crashed_ranks: HashMap::new(),
129            owner,
130            subscribers: HashSet::new(),
131        }
132    }
133
134    /// Try to update the status at `point`. Returns `true` if the status
135    /// was newly inserted or changed; `false` if dominated by a higher
136    /// generation or unchanged.
137    fn maybe_update(&mut self, point: Point, status: resource::Status, generation: u64) -> bool {
138        use std::collections::hash_map::Entry;
139        match self.statuses.entry(point) {
140            Entry::Occupied(mut entry) => {
141                let (old_status, old_gen) = entry.get();
142                // Once a resource enters a terminating state (including Stopping),
143                // its status is frozen — later updates are ignored.
144                if old_status.is_terminating() || *old_gen > generation {
145                    return false;
146                }
147                let changed = *old_status != status;
148                *entry.get_mut() = (status, generation);
149                changed
150            }
151            Entry::Vacant(entry) => {
152                entry.insert((status, generation));
153                true
154            }
155        }
156    }
157
158    /// True when every tracked rank has reached a terminating status.
159    fn all_terminating(&self) -> bool {
160        self.statuses.values().all(|(s, _)| s.is_terminating())
161    }
162
163    /// True when at least one tracked rank has reached a terminating status.
164    fn any_terminating(&self) -> bool {
165        self.statuses.values().any(|(s, _)| s.is_terminating())
166    }
167
168    /// The lowest rank that has not yet consumed one terminal state slot.
169    fn first_non_terminating_rank(&self) -> Option<usize> {
170        self.statuses
171            .iter()
172            .filter(|(_, (status, _))| !status.is_terminating())
173            .map(|(point, _)| point.rank())
174            .min()
175    }
176
177    /// Mark `rank` with a terminal status. Returns `true` only if this rank
178    /// was not already terminal.
179    fn mark_rank_terminating(&mut self, rank: usize, status: resource::Status) -> bool {
180        assert!(status.is_terminating(), "rank status must be terminating");
181        let point = self
182            .statuses
183            .keys()
184            .find(|point| point.rank() == rank)
185            .cloned()
186            .unwrap_or_else(|| panic!("rank {rank} is not tracked by health state"));
187        self.maybe_update(point, status, u64::MAX)
188    }
189
190    /// Apply status updates from polled resource states and invoke `on_change`
191    /// for each rank whose status actually changed. The `Point` passed to
192    /// `on_change` is the rank in the *current* (possibly sliced) view — the
193    /// same key used to update health state — so callers position notifications
194    /// by it rather than by the payload's first-creation rank. Returns `true`
195    /// if `on_change` reported at least one notification (used to decide whether
196    /// a heartbeat is needed).
197    pub(crate) fn apply_updates_and_notify<S: Clone + 'static>(
198        &mut self,
199        states: &ValueMesh<resource::State<S>>,
200        mut on_change: impl FnMut(Point, resource::State<S>, &mut HealthState) -> bool,
201    ) -> bool {
202        let mut did_notify = false;
203        for (point, state) in states.iter() {
204            let status = state.status.clone();
205            let generation = state.generation;
206            if self.maybe_update(point.clone(), status, generation) && on_change(point, state, self)
207            {
208                did_notify = true;
209            }
210        }
211        did_notify
212    }
213}
214
215/// Outcome of the mesh-specific polling phase inside `CheckState`.
216pub enum PollResult {
217    /// An error or early condition was handled internally; just reschedule.
218    Reschedule,
219    /// A terminal polling failure was reported; stop periodic polling.
220    StopMonitoring,
221    /// States were polled and processed. `did_notify` is true if at least
222    /// one subscriber/owner notification was sent.
223    Processed { did_notify: bool },
224}
225
226/// Compute the keepalive expiry from `MESH_ORPHAN_TIMEOUT`, or `None` if
227/// the timeout is disabled.
228fn compute_keepalive() -> Option<SystemTime> {
229    hyperactor_config::global::get(MESH_ORPHAN_TIMEOUT).map(|d| SystemTime::now() + d)
230}
231
232/// Subscribe me to updates about a mesh. If a duplicate is subscribed, only a single
233/// message is sent.
234/// Will send None if there are no failures on the mesh periodically. This guarantees
235/// the listener that the controller is still alive. Make sure to filter such events
236/// out as not useful.
237#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named)]
238pub struct Subscribe(pub hyperactor::PortRef<Option<MeshFailure>>);
239wirevalue::register_type!(Subscribe);
240
241/// Unsubscribe me to future updates about a mesh. Should be the same port used in
242/// the Subscribe message.
243#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named)]
244pub struct Unsubscribe(pub hyperactor::PortRef<Option<MeshFailure>>);
245wirevalue::register_type!(Unsubscribe);
246
247/// Query the number of active supervision subscribers on this controller.
248#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named)]
249pub struct GetSubscriberCount(pub hyperactor::PortRef<usize>);
250wirevalue::register_type!(GetSubscriberCount);
251
252/// Check state of the actors in the mesh. This is used as a self message to
253/// periodically check.
254/// Stores the next time we expect to start running a check state message.
255/// Used to check for stalls in message handling.
256#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named)]
257pub struct CheckState(pub SystemTime);
258wirevalue::register_type!(CheckState);
259
260declare_attrs! {
261    /// If present in a message header, the message is from a mesh supervision
262    /// source to a subscriber and can be safely dropped if it is returned as
263    /// undeliverable.
264    pub attr ACTOR_MESH_SUBSCRIBER_MESSAGE: bool;
265}
266
267fn send_subscriber_message(
268    cx: &impl context::Actor,
269    subscriber: &hyperactor::PortRef<Option<MeshFailure>>,
270    message: MeshFailure,
271) {
272    let mut headers = Flattrs::new();
273    headers.set(ACTOR_MESH_SUBSCRIBER_MESSAGE, true);
274    subscriber.post_with_headers(cx, headers, Some(message.clone()));
275    tracing::info!(event = %message, "sent supervision failure message to subscriber {}", subscriber.port_addr());
276}
277
278/// Like send_state_change, but when there was no state change that occurred.
279/// Will send a None message to subscribers, and there is no state to change.
280/// Is not sent to the owner, because the owner is only watching for failures.
281/// Should be called once every so often so subscribers can discern the difference
282/// between "no messages because no errors" and "no messages because controller died".
283/// Without sending these hearbeats, subscribers will assume the mesh is dead.
284fn send_heartbeat(cx: &impl context::Actor, health_state: &HealthState) {
285    tracing::debug!(
286        num_subscribers = health_state.subscribers.len(),
287        "sending heartbeat to subscribers",
288    );
289
290    for subscriber in health_state.subscribers.iter() {
291        let mut headers = Flattrs::new();
292        headers.set(ACTOR_MESH_SUBSCRIBER_MESSAGE, true);
293        subscriber.post_with_headers(cx, headers, None);
294    }
295}
296
297/// Sends a MeshFailure to the owner and subscribers of this mesh,
298/// and changes the health state stored unhealthy_event.
299/// Owners are sent a message only for Failure events, not for Stopped events.
300/// Subscribers are sent both Stopped and Failure events.
301fn send_state_change(
302    cx: &impl context::Actor,
303    rank: usize,
304    event: ActorSupervisionEvent,
305    mesh_name: &ResourceId,
306    is_proc_stopped: bool,
307    health_state: &mut HealthState,
308) {
309    // This does not include the Stopped status, which is a state that occurs when the
310    // user calls stop() on a proc or actor mesh.
311    let is_failed = event.is_error();
312    if is_failed {
313        tracing::warn!(
314            name = "SupervisionEvent",
315            actor_mesh = %mesh_name,
316            %event,
317            "detected supervision error on monitored mesh: name={mesh_name}",
318        );
319    } else {
320        tracing::debug!(
321            name = "SupervisionEvent",
322            actor_mesh = %mesh_name,
323            %event,
324            "detected non-error supervision event on monitored mesh: name={mesh_name}",
325        );
326    }
327
328    let failure_message = MeshFailure {
329        actor_mesh_name: Some(mesh_name.to_string()),
330        event: event.clone(),
331        crashed_ranks: vec![rank],
332        // MFCA-1/MFCA-2: stamp this controller's id so the owner attributes the
333        // failure by controller, whatever the event subject is.
334        reporting_controller: Some(cx.instance().self_addr().id().clone()),
335    };
336    health_state.crashed_ranks.insert(rank, event.clone());
337    health_state.unhealthy_event = Some(if is_proc_stopped {
338        Unhealthy::StreamClosed(failure_message.clone())
339    } else {
340        Unhealthy::Crashed(failure_message.clone())
341    });
342    // Send a notification to the owning actor of this mesh, if there is one.
343    // Don't send a message to the owner for non-failure events such as "stopped".
344    // Those events are always initiated by the owner, who don't need to be
345    // told that they were stopped.
346    if is_failed && let Some(owner) = &health_state.owner {
347        owner.post(cx, failure_message.clone());
348        tracing::info!(actor_mesh = %mesh_name, %event, "sent supervision failure message to owner {}", owner.port_addr());
349    }
350    // Subscribers get all messages, even for non-failures like Stopped, because
351    // they need to know if the owner stopped the mesh.
352    for subscriber in health_state.subscribers.iter() {
353        send_subscriber_message(cx, subscriber, failure_message.clone());
354    }
355}
356
357fn send_poll_failure(
358    cx: &impl context::Actor,
359    event: ActorSupervisionEvent,
360    mesh_name: &ResourceId,
361    health_state: &mut HealthState,
362) -> PollResult {
363    let Some(rank) = health_state.first_non_terminating_rank() else {
364        return PollResult::StopMonitoring;
365    };
366    health_state.mark_rank_terminating(rank, resource::Status::Failed(event.to_string()));
367    send_state_change(cx, rank, event, mesh_name, false, health_state);
368    PollResult::StopMonitoring
369}
370
371fn actor_state_to_supervision_events(
372    state: resource::State<ActorState>,
373    expected_actor_id: Option<ActorAddr>,
374) -> Vec<ActorSupervisionEvent> {
375    // Prefer the payload's actor id; fall back to the controller's expected
376    // actor at this rank when the payload carries no `ActorState` (a missing or
377    // spawn-failed proc), so a failure is still attributed to the right actor.
378    let (actor_id, events) = match state.state {
379        Some(inner) => (Some(inner.actor_id), inner.supervision_events.clone()),
380        None => (expected_actor_id, vec![]),
381    };
382    if !events.is_empty() {
383        return events;
384    }
385    // No reported events: synthesize one for terminal/failed statuses. Without
386    // any actor identity there is nothing to attribute, so emit nothing.
387    let Some(actor_id) = actor_id else {
388        return vec![];
389    };
390    match state.status {
391        // Killed/stopped/missing/timeout: the actor may not report a Failed
392        // status or events, and can't tell us which rank it was.
393        resource::Status::NotExist | resource::Status::Stopped | resource::Status::Timeout(_) => {
394            vec![ActorSupervisionEvent::new(
395                actor_id,
396                None,
397                ActorStatus::Stopped(
398                    format!(
399                        "actor status is {}; actor may have been killed",
400                        state.status
401                    )
402                    .to_string(),
403                ),
404                None,
405            )]
406        }
407        // Failed with no reported event: synthesize one so the failure is not
408        // silently swallowed while health is updated.
409        resource::Status::Failed(ref reason) => {
410            vec![ActorSupervisionEvent::new(
411                actor_id,
412                None,
413                ActorStatus::generic_failure(reason.clone()),
414                None,
415            )]
416        }
417        // All other states are successful.
418        _ => vec![],
419    }
420}
421
422/// Map a process-level [`ProcStatus`] to an actor-level [`ActorStatus`].
423///
424/// When the supervision poll discovers that a process is terminating, this
425/// function decides whether to treat it as a clean stop or a failure.
426/// Notably, [`ProcStatus::Stopping`] (SIGTERM sent, process not yet exited)
427/// is mapped to [`ActorStatus::Stopped`] rather than [`ActorStatus::Failed`]
428/// so that a graceful shutdown in progress does not trigger unhandled
429/// supervision errors.
430fn proc_status_to_actor_status(proc_status: Option<ProcStatus>) -> ActorStatus {
431    match proc_status {
432        Some(ProcStatus::Stopped { exit_code: 0, .. }) => {
433            ActorStatus::Stopped("process exited cleanly".to_string())
434        }
435        Some(ProcStatus::Stopped { exit_code, .. }) => {
436            ActorStatus::Failed(ActorErrorKind::Generic(format!(
437                "the process this actor was running on exited with non-zero code {}",
438                exit_code
439            )))
440        }
441        // Stopping is a transient state during graceful shutdown. Treat it the
442        // same as a clean stop rather than a failure.
443        Some(ProcStatus::Stopping { .. }) => {
444            ActorStatus::Stopped("process is stopping".to_string())
445        }
446        // Conservatively treat lack of status as stopped
447        None => ActorStatus::Stopped("no status received from process".to_string()),
448        Some(status) => ActorStatus::Failed(ActorErrorKind::Generic(format!(
449            "the process this actor was running on failed: {}",
450            status
451        ))),
452    }
453}
454
455/// Log a warning and bump `counter` if the supervision loop is running late.
456///
457/// "Late" means the current wall-clock time exceeds `expected_time` by more
458/// than one full poll interval, i.e. 2x the expected period.
459fn check_stall(expected_time: SystemTime, actor_id: &hyperactor::ActorId, counter: &Counter<u64>) {
460    let now = SystemTime::now();
461    let poll_frequency = hyperactor_config::global::get(SUPERVISION_POLL_FREQUENCY);
462    let Ok(mut stalled_by) = now.duration_since(expected_time + poll_frequency) else {
463        return;
464    };
465    // Add back before using, as the addition is only used to avoid logging unless it's
466    // stalled by at least one cycle.
467    stalled_by += poll_frequency;
468    counter.add(
469        1,
470        // hyperactor_telemetry::Value only supports i64, not u64.
471        kv_pairs!("actor_id" => actor_id.to_string(), "stalled_by_seconds" => stalled_by.as_secs() as i64),
472    );
473    tracing::warn!(
474        %actor_id,
475        "Handler<CheckState> is stalled by {}",
476        humantime::format_duration(stalled_by),
477    );
478}
479
480/// Mesh-specific behavior required by the generic `ResourceController`.
481///
482/// Each variant of resource mesh (actor, proc) implements this trait to
483/// provide the details that cannot be shared by the generic controller:
484/// the state type carried in `resource::State<_>`, how to query or stream
485/// that state from the underlying agents, how to stop the resources, and
486/// how to notify observers when the state changes.
487#[async_trait]
488pub trait Controlled: Clone + Debug + Send + Sync + 'static {
489    /// Inner payload carried in `resource::State<Self::StateInner>`.
490    type StateInner: RemoteMessage + Clone + Debug + 'static;
491
492    /// Counter bumped when the supervision loop detects a stall.
493    fn stall_counter() -> &'static Counter<u64>;
494
495    /// The mesh's resource identifier.
496    fn id(&self) -> &ResourceId;
497
498    /// The region of ranks in this mesh.
499    fn region(&self) -> &ndslice::Region;
500
501    /// Subscribe the given port to positioned state updates from the
502    /// underlying agents.
503    fn subscribe_to_stream(
504        &self,
505        cx: &impl context::Actor,
506        subscriber: hyperactor::PortRef<resource::RankedState<Self::StateInner>>,
507    ) -> anyhow::Result<()>;
508
509    /// Forward a `WaitRankStatus` message to the underlying agents.
510    fn forward_wait_rank_status(
511        &self,
512        cx: &impl context::Actor,
513        msg: resource::WaitRankStatus,
514    ) -> anyhow::Result<()>;
515
516    /// Mesh-specific polling step for the supervision loop. Implementations
517    /// may do pre-checks (such as the actor mesh's proc-aliveness check)
518    /// before querying rank states; updates to `health_state` happen
519    /// in-place. `supervision_display_name` is used for synthesised
520    /// supervision events (e.g., when a proc dies).
521    async fn poll_states(
522        &self,
523        cx: &impl context::Actor,
524        supervision_display_name: &str,
525        health_state: &mut HealthState,
526    ) -> PollResult;
527
528    /// Process a single streamed or polled state. Updates the health state
529    /// and notifies owner/subscribers as appropriate. Returns `true` if a
530    /// notification was emitted (used to suppress heartbeats).
531    fn process_state(
532        &self,
533        cx: &impl context::Actor,
534        state: resource::RankedState<Self::StateInner>,
535        health_state: &mut HealthState,
536    ) -> bool;
537
538    /// Perform the mesh-specific stop: issue stop messages to the underlying
539    /// agents and, where appropriate, update `health_state` and notify
540    /// subscribers. The caller has already taken the monitor and logged.
541    async fn handle_stop_request(
542        &self,
543        cx: &impl context::Actor,
544        supervision_display_name: &str,
545        reason: String,
546        health_state: &mut HealthState,
547    ) -> anyhow::Result<()>;
548
549    /// Stop this mesh on controller cleanup (when `Stop` was not received
550    /// but the actor is shutting down).
551    async fn cleanup_stop(&self, cx: &impl context::Actor, reason: String) -> anyhow::Result<()>;
552}
553
554/// Generic controller for a mesh of resources. Actor meshes instantiate this
555/// as `ActorMeshController<A> = ResourceController<ActorMeshControlPlane<A>>`.
556/// All shared behavior lives here; mesh-specific behavior is delegated through
557/// the `Controlled` trait.
558///
559/// `resource::mesh::Spec<()>` and `resource::mesh::State<()>` (instead of
560/// `Spec<T::Spec>` / `State<T::StateInner>`) are used because the
561/// controller participates in the mesh `resource` protocol only at the
562/// outer layer: callers of `GetState` on the controller want the
563/// per-rank statuses and the mesh-wide status that `resource::mesh::State`
564/// already carries, not the inner `T::StateInner` payload (which is
565/// available rank-by-rank via the
566/// `resource::RankedState<T::StateInner>` stream).
567/// The unit type is the explicit "no extra payload" choice.
568#[hyperactor::export(
569    handlers=[
570        Subscribe,
571        Unsubscribe,
572        GetSubscriberCount,
573        CheckState,
574        resource::WaitRankStatus,
575        resource::CreateOrUpdate<resource::mesh::Spec<()>>,
576        resource::GetState<resource::mesh::State<()>>,
577        resource::Stop,
578        resource::RankedState<T::StateInner>,
579    ]
580)]
581pub struct ResourceController<T: Controlled> {
582    mesh: T,
583    /// Supervision display name used in telemetry and fake supervision
584    /// events. If `None`, falls back to `mesh.id()`.
585    supervision_display_name: Option<String>,
586    /// Shared health state for the monitor and responding to queries.
587    health_state: HealthState,
588    /// The monitor which continuously runs in the background to refresh
589    /// state. If None, the controller has stopped monitoring.
590    monitor: Option<()>,
591}
592
593/// Controller-side state for an actor mesh.
594///
595/// `ActorMeshRef` is moving toward being a self-contained, serializable mesh
596/// descriptor. The controller is the one place that needs both the actor mesh
597/// and its backing proc mesh: it streams actor state through proc agents and
598/// stops actors through the proc mesh control plane.
599pub struct ActorMeshControlPlane<A: Referable> {
600    actor_mesh: ActorMeshRef<A>,
601    proc_mesh: ProcMeshRef,
602}
603
604impl<A: Referable> ActorMeshControlPlane<A> {
605    pub(crate) fn new(actor_mesh: ActorMeshRef<A>, proc_mesh: ProcMeshRef) -> Self {
606        Self {
607            actor_mesh,
608            proc_mesh,
609        }
610    }
611}
612
613impl<A: Referable> Clone for ActorMeshControlPlane<A> {
614    fn clone(&self) -> Self {
615        Self {
616            actor_mesh: self.actor_mesh.clone(),
617            proc_mesh: self.proc_mesh.clone(),
618        }
619    }
620}
621
622impl<A: Referable> Debug for ActorMeshControlPlane<A> {
623    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
624        f.debug_struct("ActorMeshControlPlane")
625            .field("actor_mesh", &self.actor_mesh)
626            .field("proc_mesh", &self.proc_mesh)
627            .finish()
628    }
629}
630
631impl<A: Referable> Named for ActorMeshControlPlane<A> {
632    fn typename() -> &'static str {
633        wirevalue::intern_typename!(Self, "hyperactor_mesh::ActorMeshControlPlane<{}>", A)
634    }
635}
636
637/// Controller for an actor mesh.
638pub type ActorMeshController<A> = ResourceController<ActorMeshControlPlane<A>>;
639
640impl<T: Controlled> ResourceController<T> {
641    /// Create a new controller over the given mesh.
642    pub(crate) fn new(
643        mesh: T,
644        supervision_display_name: Option<String>,
645        owner: Option<hyperactor::PortRef<MeshFailure>>,
646        initial_statuses: ValueMesh<resource::Status>,
647    ) -> Self {
648        Self {
649            mesh,
650            supervision_display_name,
651            health_state: HealthState::new(initial_statuses.iter().collect(), owner),
652            monitor: None,
653        }
654    }
655
656    /// The display name to use for supervision events and telemetry.
657    pub(crate) fn supervision_display_name(&self) -> String {
658        self.supervision_display_name
659            .clone()
660            .unwrap_or_else(|| self.mesh.id().to_string())
661    }
662
663    /// Schedule the next `CheckState` self-message if the monitor is active.
664    ///
665    /// `send_fn` bridges the type gap: the caller passes a closure that
666    /// captures the typed `Instance`/`Context` and calls
667    /// `post_after`.
668    fn schedule_next_check(&self, send_fn: impl FnOnce(CheckState, Duration)) {
669        if self.monitor.is_some() {
670            let delay = hyperactor_config::global::get(SUPERVISION_POLL_FREQUENCY);
671            send_fn(CheckState(SystemTime::now() + delay), delay);
672        }
673    }
674
675    /// Derive the mesh-level status from health state and monitor presence.
676    fn mesh_status(&self) -> resource::Status {
677        if let Some(Unhealthy::Crashed(e)) = &self.health_state.unhealthy_event {
678            resource::Status::Failed(e.to_string())
679        } else if let Some(Unhealthy::StreamClosed(_)) = &self.health_state.unhealthy_event {
680            resource::Status::Stopped
681        } else if self.monitor.is_none() {
682            resource::Status::Stopped
683        } else {
684            resource::Status::Running
685        }
686    }
687
688    /// Build and send the `GetState<resource::mesh::State<()>>` response.
689    fn handle_get_state_msg(
690        &self,
691        cx: &impl context::Actor,
692        message: resource::GetState<resource::mesh::State<()>>,
693    ) -> anyhow::Result<()> {
694        let status = self.mesh_status();
695        let mut statuses = self
696            .health_state
697            .statuses
698            .iter()
699            .map(|(p, (s, _))| (p.clone(), s.clone()))
700            .collect::<Vec<_>>();
701        statuses.sort_by_key(|(p, _)| p.rank());
702        let statuses: ValueMesh<resource::Status> =
703            statuses
704                .into_iter()
705                .map(|(_, s)| s)
706                .collect_mesh::<ValueMesh<_>>(self.mesh.region().clone())?;
707        let state = resource::mesh::State {
708            statuses,
709            state: (),
710        };
711        message.reply.post(
712            cx,
713            resource::State {
714                id: message.id,
715                status,
716                state: Some(state),
717                generation: 0,
718                timestamp: SystemTime::now(),
719            },
720        );
721        Ok(())
722    }
723
724    /// Drop the monitor if every tracked rank has reached a terminal status.
725    fn stop_if_all_terminating(&mut self) {
726        if self.health_state.all_terminating() {
727            self.monitor.take();
728        }
729    }
730
731    async fn handle_check_state(
732        &mut self,
733        cx: &Context<'_, Self>,
734        expected_time: SystemTime,
735    ) -> anyhow::Result<()>
736    where
737        resource::State<T::StateInner>: RemoteMessage,
738    {
739        if self.monitor.is_none() {
740            return Ok(());
741        }
742        check_stall(expected_time, cx.self_addr().id(), T::stall_counter());
743
744        let display = self.supervision_display_name();
745        let result = self
746            .mesh
747            .poll_states(cx, &display, &mut self.health_state)
748            .await;
749
750        match result {
751            PollResult::Reschedule => {
752                self.schedule_next_check(|msg, delay| cx.post_after(cx, msg, delay));
753            }
754            PollResult::StopMonitoring => {
755                self.monitor.take();
756            }
757            PollResult::Processed { did_notify } => {
758                // Suppress heartbeats once any rank is terminating: the mesh is on
759                // its way down and subscribers will get a real state-change message
760                // for the terminal transition.
761                if !did_notify && !self.health_state.any_terminating() {
762                    send_heartbeat(cx, &self.health_state);
763                }
764                if !self.health_state.all_terminating() {
765                    self.schedule_next_check(|msg, delay| cx.post_after(cx, msg, delay));
766                } else {
767                    self.monitor.take();
768                }
769            }
770        }
771        Ok(())
772    }
773}
774
775impl<T: Controlled> Debug for ResourceController<T> {
776    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
777        f.debug_struct("ResourceController")
778            .field("mesh", &self.mesh)
779            .field("health_state", &self.health_state)
780            .field("monitor", &self.monitor)
781            .finish()
782    }
783}
784
785impl<T: Controlled> resource::mesh::Mesh for ResourceController<T> {
786    type Spec = ();
787    type State = ();
788}
789
790#[async_trait]
791impl<T: Controlled> Actor for ResourceController<T>
792where
793    resource::State<T::StateInner>: RemoteMessage,
794{
795    async fn init(&mut self, this: &Instance<Self>) -> Result<(), anyhow::Error> {
796        this.set_system();
797
798        // Subscribe to streaming state updates from the underlying agents so
799        // the controller receives state changes in real time, complementing
800        // the existing polling loop. Must happen before starting the monitor
801        // so that the first CheckState does not race the initial StreamState
802        // cast.
803        //
804        // Bind the handler port before handing the `PortRef` to the agent.
805        // The later external bind races against `init`, but both calls bind
806        // the same well-known handler port and are idempotent. No `.unsplit()`:
807        // the subscriber port is split by the cast tree so streamed states
808        // reduce up it (see the `StreamState` cast).
809        self.mesh.subscribe_to_stream(this, this.port().bind())?;
810
811        // Start the monitor task.
812        self.monitor = Some(());
813        self.schedule_next_check(|msg, delay| this.post_after(this, msg, delay));
814
815        let owner = if let Some(owner) = &self.health_state.owner {
816            owner.to_string()
817        } else {
818            String::from("None")
819        };
820        tracing::info!(
821            actor_id = %this.self_addr(),
822            %owner,
823            "started resource controller for {}",
824            self.mesh.id()
825        );
826        Ok(())
827    }
828
829    async fn cleanup(
830        &mut self,
831        this: &Instance<Self>,
832        _err: Option<&ActorError>,
833    ) -> Result<(), anyhow::Error> {
834        if self.monitor.take().is_some() {
835            tracing::info!(
836                actor_id = %this.self_addr(),
837                mesh = %self.mesh.id(),
838                "starting cleanup for ResourceController, stopping mesh",
839            );
840            self.mesh
841                .cleanup_stop(this, "resource controller cleanup".to_string())
842                .await?;
843        }
844        Ok(())
845    }
846
847    async fn handle_undeliverable_message(
848        &mut self,
849        cx: &Instance<Self>,
850        reason: UndeliverableReason,
851        mut envelope: Undeliverable<MessageEnvelope>,
852    ) -> Result<(), anyhow::Error> {
853        envelope = update_undeliverable_envelope_for_casting(envelope);
854        let Some(returned) = envelope.as_message() else {
855            return handle_undeliverable_message(cx, reason, envelope);
856        };
857        if let Some(true) = returned.headers().get(ACTOR_MESH_SUBSCRIBER_MESSAGE) {
858            // Remove from the subscriber list (if it existed) so we don't
859            // send to this subscriber again.
860            // NOTE: The only part of the port that is used for equality checks is
861            // the port id, so create a new one just for the comparison.
862            let dest_port_id = returned.dest().clone();
863            let port = hyperactor::PortRef::<Option<MeshFailure>>::attest(dest_port_id);
864            let did_exist = self.health_state.subscribers.remove(&port);
865            if did_exist {
866                tracing::debug!(
867                    actor_id = %cx.self_addr(),
868                    num_subscribers = self.health_state.subscribers.len(),
869                    "ResourceController: removed subscriber {} from mesh controller",
870                    port.port_addr()
871                );
872            }
873            Ok(())
874        } else if returned.headers().get(CAST_ACTOR_MESH_ID).is_some() {
875            // A cast message we sent (e.g. StreamState or KeepaliveGetState)
876            // was returned by cast routing because it could not be forwarded.
877            // This is expected when the network session is broken. Log and
878            // continue — the supervision polling loop will detect the failure.
879            tracing::warn!(
880                actor_id = %cx.self_addr(),
881                dest = %returned.dest(),
882                "ResourceController: ignoring undeliverable cast message",
883            );
884            Ok(())
885        } else {
886            handle_undeliverable_message(cx, reason, envelope)
887        }
888    }
889
890    async fn handle_invalid_reference(
891        &mut self,
892        cx: &Instance<Self>,
893        invalid: hyperactor::mailbox::InvalidReference,
894        envelope: Undeliverable<MessageEnvelope>,
895    ) -> Result<(), anyhow::Error> {
896        let envelope = update_undeliverable_envelope_for_casting(envelope);
897        let Some(returned) = envelope.as_message() else {
898            return hyperactor::actor::handle_invalid_reference(cx, invalid, envelope);
899        };
900        if let Some(true) = returned.headers().get(ACTOR_MESH_SUBSCRIBER_MESSAGE) {
901            let dest_port_id = returned.dest().clone();
902            let port = hyperactor::PortRef::<Option<MeshFailure>>::attest(dest_port_id);
903            let did_exist = self.health_state.subscribers.remove(&port);
904            if did_exist {
905                tracing::debug!(
906                    actor_id = %cx.self_addr(),
907                    num_subscribers = self.health_state.subscribers.len(),
908                    "ResourceController: removed subscriber {} from mesh controller",
909                    port.port_addr()
910                );
911            }
912            Ok(())
913        } else if returned.headers().get(CAST_ACTOR_MESH_ID).is_some() {
914            tracing::warn!(
915                actor_id = %cx.self_addr(),
916                dest = %returned.dest(),
917                "ResourceController: ignoring undeliverable cast message",
918            );
919            Ok(())
920        } else {
921            hyperactor::actor::handle_invalid_reference(cx, invalid, envelope)
922        }
923    }
924}
925
926#[async_trait]
927impl<T: Controlled> Handler<Subscribe> for ResourceController<T>
928where
929    resource::State<T::StateInner>: RemoteMessage,
930{
931    async fn handle(&mut self, cx: &Context<Self>, message: Subscribe) -> anyhow::Result<()> {
932        // If there are any crashed ranks, replay a failure event so the new
933        // subscriber learns about the current health state. We send a single
934        // message with all crashed ranks so the subscriber's filter can check
935        // overlap with its slice region. This avoids the watch-channel
936        // coalescing problem (sending per-rank messages would lose all but
937        // the last one).
938        if let Some(unhealthy) = &self.health_state.unhealthy_event {
939            let msg = match unhealthy {
940                Unhealthy::StreamClosed(msg) | Unhealthy::Crashed(msg) => msg,
941            };
942            let mut replay_msg = msg.clone();
943            replay_msg.crashed_ranks = self.health_state.crashed_ranks.keys().copied().collect();
944            send_subscriber_message(cx, &message.0, replay_msg);
945        }
946        let port_id = message.0.port_addr().clone();
947        if self.health_state.subscribers.insert(message.0) {
948            tracing::debug!(
949                actor_id = %cx.self_addr(),
950                num_subscribers = self.health_state.subscribers.len(),
951                "added subscriber {} to mesh controller",
952                port_id
953            );
954        }
955        Ok(())
956    }
957}
958
959#[async_trait]
960impl<T: Controlled> Handler<Unsubscribe> for ResourceController<T>
961where
962    resource::State<T::StateInner>: RemoteMessage,
963{
964    async fn handle(&mut self, cx: &Context<Self>, message: Unsubscribe) -> anyhow::Result<()> {
965        if self.health_state.subscribers.remove(&message.0) {
966            tracing::debug!(
967                actor_id = %cx.self_addr(),
968                num_subscribers = self.health_state.subscribers.len(),
969                "removed subscriber {} from mesh controller",
970                message.0.port_addr()
971            );
972        }
973        Ok(())
974    }
975}
976
977#[async_trait]
978impl<T: Controlled> Handler<GetSubscriberCount> for ResourceController<T>
979where
980    resource::State<T::StateInner>: RemoteMessage,
981{
982    async fn handle(
983        &mut self,
984        cx: &Context<Self>,
985        message: GetSubscriberCount,
986    ) -> anyhow::Result<()> {
987        message.0.post(cx, self.health_state.subscribers.len());
988        Ok(())
989    }
990}
991
992#[async_trait]
993impl<T: Controlled> Handler<resource::CreateOrUpdate<resource::mesh::Spec<()>>>
994    for ResourceController<T>
995where
996    resource::State<T::StateInner>: RemoteMessage,
997{
998    /// Currently a no-op as there's nothing to create or update, but allows
999    /// `ResourceController` to implement the resource mesh behavior.
1000    async fn handle(
1001        &mut self,
1002        _cx: &Context<Self>,
1003        _message: resource::CreateOrUpdate<resource::mesh::Spec<()>>,
1004    ) -> anyhow::Result<()> {
1005        Ok(())
1006    }
1007}
1008
1009#[async_trait]
1010impl<T: Controlled> Handler<resource::GetState<resource::mesh::State<()>>> for ResourceController<T>
1011where
1012    resource::State<T::StateInner>: RemoteMessage,
1013{
1014    async fn handle(
1015        &mut self,
1016        cx: &Context<Self>,
1017        message: resource::GetState<resource::mesh::State<()>>,
1018    ) -> anyhow::Result<()> {
1019        self.handle_get_state_msg(cx, message)
1020    }
1021}
1022
1023#[async_trait]
1024impl<T: Controlled> Handler<resource::Stop> for ResourceController<T>
1025where
1026    resource::State<T::StateInner>: RemoteMessage,
1027{
1028    async fn handle(&mut self, cx: &Context<Self>, message: resource::Stop) -> anyhow::Result<()> {
1029        let mesh_name = self.mesh.id().clone();
1030        tracing::info!(
1031            name = "ResourceControllerStatus",
1032            %mesh_name,
1033            reason = %message.reason,
1034            "stopping mesh"
1035        );
1036        if self.monitor.take().is_none() {
1037            tracing::debug!(
1038                actor_id = %cx.self_addr(),
1039                %mesh_name,
1040                "duplicate stop request, mesh is already stopped",
1041            );
1042            return Ok(());
1043        }
1044        let display = self.supervision_display_name();
1045        self.mesh
1046            .handle_stop_request(cx, &display, message.reason, &mut self.health_state)
1047            .await
1048    }
1049}
1050
1051#[async_trait]
1052impl<T: Controlled> Handler<resource::WaitRankStatus> for ResourceController<T>
1053where
1054    resource::State<T::StateInner>: RemoteMessage,
1055{
1056    /// Forward WaitRankStatus to the underlying agents. Each agent replies
1057    /// directly to the caller's accumulator port when its resource reaches
1058    /// the requested status.
1059    async fn handle(
1060        &mut self,
1061        cx: &Context<Self>,
1062        msg: resource::WaitRankStatus,
1063    ) -> anyhow::Result<()> {
1064        self.mesh.forward_wait_rank_status(cx, msg)
1065    }
1066}
1067
1068#[async_trait]
1069impl<T: Controlled> Handler<CheckState> for ResourceController<T>
1070where
1071    resource::State<T::StateInner>: RemoteMessage,
1072{
1073    async fn handle(
1074        &mut self,
1075        cx: &Context<Self>,
1076        CheckState(expected_time): CheckState,
1077    ) -> Result<(), anyhow::Error> {
1078        self.handle_check_state(cx, expected_time).await
1079    }
1080}
1081
1082#[async_trait]
1083impl<T: Controlled> Handler<resource::RankedState<T::StateInner>> for ResourceController<T>
1084where
1085    resource::RankedState<T::StateInner>: RemoteMessage,
1086{
1087    async fn handle(
1088        &mut self,
1089        cx: &Context<Self>,
1090        state: resource::RankedState<T::StateInner>,
1091    ) -> anyhow::Result<()> {
1092        self.mesh.process_state(cx, state, &mut self.health_state);
1093        self.stop_if_all_terminating();
1094        Ok(())
1095    }
1096}
1097
1098/// `Controlled` implementation for an actor mesh.
1099#[async_trait]
1100impl<A: Referable> Controlled for ActorMeshControlPlane<A> {
1101    type StateInner = ActorState;
1102
1103    fn stall_counter() -> &'static Counter<u64> {
1104        &ACTOR_MESH_CONTROLLER_SUPERVISION_STALLS
1105    }
1106
1107    fn id(&self) -> &ResourceId {
1108        self.actor_mesh
1109            .as_managed()
1110            .expect("actor mesh control plane requires a managed mesh")
1111            .id()
1112            .resource_id()
1113    }
1114
1115    fn region(&self) -> &ndslice::Region {
1116        ndslice::view::Ranked::region(&self.actor_mesh)
1117    }
1118
1119    fn subscribe_to_stream(
1120        &self,
1121        cx: &impl context::Actor,
1122        subscriber: hyperactor::PortRef<resource::RankedState<ActorState>>,
1123    ) -> anyhow::Result<()> {
1124        self.proc_mesh.agent_mesh().cast(
1125            cx,
1126            resource::StreamState::<ActorState> {
1127                id: self
1128                    .actor_mesh
1129                    .as_managed()
1130                    .expect("actor mesh control plane requires a managed mesh")
1131                    .id()
1132                    .resource_id()
1133                    .clone(),
1134                subscriber_rank: resource::Rank::default(),
1135                subscriber,
1136            },
1137        )?;
1138        Ok(())
1139    }
1140
1141    fn forward_wait_rank_status(
1142        &self,
1143        cx: &impl context::Actor,
1144        msg: resource::WaitRankStatus,
1145    ) -> anyhow::Result<()> {
1146        self.proc_mesh.agent_mesh().cast(cx, msg)?;
1147        Ok(())
1148    }
1149
1150    async fn poll_states(
1151        &self,
1152        cx: &impl context::Actor,
1153        supervision_display_name: &str,
1154        health_state: &mut HealthState,
1155    ) -> PollResult {
1156        let mesh_name = Controlled::id(self);
1157
1158        // Actor-specific: first check if the proc mesh is dead before
1159        // trying to query their agents.
1160        let proc_states = self.proc_mesh.states(cx, None).await;
1161        if let Err(e) = proc_states {
1162            return send_poll_failure(
1163                cx,
1164                ActorSupervisionEvent::new(
1165                    cx.instance().self_addr().clone(),
1166                    None,
1167                    ActorStatus::generic_failure(format!(
1168                        "unable to query for proc states: {:?}",
1169                        e
1170                    )),
1171                    None,
1172                ),
1173                mesh_name,
1174                health_state,
1175            );
1176        }
1177        if let Some(proc_states) = proc_states.unwrap() {
1178            // Check if the proc mesh is still alive.
1179            if let Some((point, state)) = proc_states
1180                .iter()
1181                .find(|(_rank, state)| state.status.is_terminating())
1182            {
1183                // TODO: allow "actor supervision event" to be general, and
1184                // make the proc failure the cause. It is a hack to try to determine
1185                // the correct status based on process exit status.
1186                let actor_status =
1187                    proc_status_to_actor_status(state.state.and_then(|s| s.proc_status));
1188                let stop_monitoring = actor_status.is_failed();
1189                let display = crate::actor_display_name(supervision_display_name, &point);
1190                let event = ActorSupervisionEvent::new(
1191                    // Attribute this to the monitored actor, even if the underlying
1192                    // cause is a proc_failure. We propagate the cause explicitly.
1193                    self.actor_mesh
1194                        .get(point.rank())
1195                        .unwrap()
1196                        .actor_addr()
1197                        .clone(),
1198                    Some(display),
1199                    actor_status,
1200                    None,
1201                );
1202                if stop_monitoring {
1203                    if health_state.mark_rank_terminating(
1204                        point.rank(),
1205                        resource::Status::Failed(event.to_string()),
1206                    ) {
1207                        send_state_change(cx, point.rank(), event, mesh_name, true, health_state);
1208                    }
1209                    return PollResult::StopMonitoring;
1210                } else {
1211                    send_state_change(cx, point.rank(), event, mesh_name, true, health_state);
1212                    return PollResult::Reschedule;
1213                }
1214            }
1215        }
1216
1217        // Query resource states with keepalive.
1218        let actor_states = self
1219            .proc_mesh
1220            .actor_states_with_keepalive(
1221                cx,
1222                self.actor_mesh
1223                    .as_managed()
1224                    .expect("actor mesh control plane requires a managed mesh")
1225                    .id()
1226                    .clone(),
1227                compute_keepalive(),
1228            )
1229            .await;
1230        match actor_states {
1231            Err(e) => send_poll_failure(
1232                cx,
1233                ActorSupervisionEvent::new(
1234                    cx.instance().self_addr().clone(),
1235                    Some(supervision_display_name.to_string()),
1236                    ActorStatus::generic_failure(format!(
1237                        "unable to query for actor states: {:?}",
1238                        e
1239                    )),
1240                    None,
1241                ),
1242                mesh_name,
1243                health_state,
1244            ),
1245            Ok(states) => {
1246                let did_notify =
1247                    health_state.apply_updates_and_notify(&states, |point, state, health_state| {
1248                        let expected_actor_id = self
1249                            .actor_mesh
1250                            .get(point.rank())
1251                            .map(|r| r.actor_addr().clone());
1252                        let events = actor_state_to_supervision_events(state, expected_actor_id);
1253                        if events.is_empty() {
1254                            return false;
1255                        }
1256                        // Notify at this update's current-view rank, not the
1257                        // payload's historical `create_rank`.
1258                        send_state_change(
1259                            cx,
1260                            point.rank(),
1261                            events[0].clone(),
1262                            mesh_name,
1263                            false,
1264                            health_state,
1265                        );
1266                        true
1267                    });
1268                PollResult::Processed { did_notify }
1269            }
1270        }
1271    }
1272
1273    fn process_state(
1274        &self,
1275        cx: &impl context::Actor,
1276        state: resource::RankedState<ActorState>,
1277        health_state: &mut HealthState,
1278    ) -> bool {
1279        let resource::RankedState { rank, state } = state;
1280        let rank = rank.unwrap();
1281        let Ok(point) = Controlled::region(self).extent().point_of_rank(rank) else {
1282            return false;
1283        };
1284        // The controller's expected actor at this rank, used as the fallback
1285        // identity when the payload carries no `ActorState`.
1286        let expected_actor_id = self.actor_mesh.get(rank).map(|r| r.actor_addr().clone());
1287        let events = actor_state_to_supervision_events(state.clone(), expected_actor_id);
1288
1289        let changed = health_state.maybe_update(point, state.status, state.generation);
1290
1291        if changed && !events.is_empty() {
1292            send_state_change(
1293                cx,
1294                rank,
1295                events[0].clone(),
1296                Controlled::id(self),
1297                false,
1298                health_state,
1299            );
1300            true
1301        } else {
1302            false
1303        }
1304    }
1305
1306    async fn handle_stop_request(
1307        &self,
1308        cx: &impl context::Actor,
1309        _supervision_display_name: &str,
1310        reason: String,
1311        health_state: &mut HealthState,
1312    ) -> anyhow::Result<()> {
1313        let mesh_name = Controlled::id(self);
1314        tracing::info!(
1315            actor_id = %cx.instance().self_addr(),
1316            actor_mesh = %mesh_name,
1317            "forwarding stop request from ActorMeshController to proc mesh"
1318        );
1319
1320        // Let the client know that the controller has stopped. Since the monitor
1321        // is cancelled, it will not alert the owner or the subscribers.
1322        // We use a placeholder rank to get an actor id, but really there should
1323        // be a stop event for every rank in the mesh. Since every rank has the
1324        // same owner, we assume the rank doesn't matter, and the owner can just
1325        // assume the stop happened on all actors.
1326        let rank = 0usize;
1327        let event = ActorSupervisionEvent::new(
1328            self.actor_mesh
1329                .get(rank)
1330                .expect("mesh must have at least one rank")
1331                .actor_addr()
1332                .clone(),
1333            None,
1334            ActorStatus::Stopped("ActorMeshController received explicit stop request".to_string()),
1335            None,
1336        );
1337        let failure_message = MeshFailure {
1338            actor_mesh_name: Some(mesh_name.to_string()),
1339            event,
1340            crashed_ranks: vec![],
1341            // MFCA-1: controller-generated stop report.
1342            reporting_controller: Some(cx.instance().self_addr().id().clone()),
1343        };
1344        health_state.unhealthy_event = Some(Unhealthy::StreamClosed(failure_message.clone()));
1345        // We don't send a message to the owner on stops, because only the owner
1346        // can request a stop. We just send to subscribers instead, as they did
1347        // not request the stop themselves.
1348        for subscriber in health_state.subscribers.iter() {
1349            send_subscriber_message(cx, subscriber, failure_message.clone());
1350        }
1351
1352        // max_rank and extent are only needed for the deprecated RankedValues.
1353        // TODO: add cmp::Ord to Point for a max() impl.
1354        let max_rank = health_state.statuses.keys().map(|p| p.rank()).max();
1355        let extent = health_state
1356            .statuses
1357            .keys()
1358            .next()
1359            .map(|p| p.extent().clone());
1360
1361        // Cannot use "ActorMesh::stop" as it tries to message the controller.
1362        let result = self
1363            .proc_mesh
1364            .stop_actor_by_id(
1365                cx,
1366                self.actor_mesh
1367                    .as_managed()
1368                    .expect("actor mesh control plane requires a managed mesh")
1369                    .id()
1370                    .clone(),
1371                reason,
1372            )
1373            .await;
1374
1375        match result {
1376            Ok(statuses) => {
1377                // All stops successful, set actor status on health state.
1378                for (rank, status) in statuses.iter() {
1379                    health_state
1380                        .statuses
1381                        .entry(rank)
1382                        .and_modify(move |s| *s = (status, u64::MAX));
1383                }
1384            }
1385            Err(crate::Error::ActorStopError { statuses }) => {
1386                if let Some(max_rank) = max_rank {
1387                    let extent = extent.expect("no actors in mesh");
1388                    for (rank, status) in statuses.materialized_iter(max_rank).enumerate() {
1389                        *health_state
1390                            .statuses
1391                            .get_mut(&extent.point_of_rank(rank).expect("illegal rank"))
1392                            .unwrap() = (status.clone(), u64::MAX);
1393                    }
1394                }
1395            }
1396            Err(e) => {
1397                return Err(e.into());
1398            }
1399        }
1400
1401        tracing::info!(
1402            actor_id = %cx.instance().self_addr(),
1403            actor_mesh = %mesh_name,
1404            "stopped mesh"
1405        );
1406        Ok(())
1407    }
1408
1409    async fn cleanup_stop(&self, cx: &impl context::Actor, reason: String) -> anyhow::Result<()> {
1410        self.proc_mesh
1411            .stop_actor_by_id(
1412                cx,
1413                self.actor_mesh
1414                    .as_managed()
1415                    .expect("actor mesh control plane requires a managed mesh")
1416                    .id()
1417                    .clone(),
1418                reason,
1419            )
1420            .await?;
1421        Ok(())
1422    }
1423}
1424
1425/// Controller for a proc mesh.
1426pub(crate) type ProcMeshController = ResourceController<ProcMeshRef>;
1427
1428/// `Controlled` implementation for a proc mesh.
1429#[async_trait]
1430impl Controlled for ProcMeshRef {
1431    type StateInner = crate::host_mesh::host_agent::ProcState;
1432
1433    fn stall_counter() -> &'static Counter<u64> {
1434        &PROC_MESH_CONTROLLER_SUPERVISION_STALLS
1435    }
1436
1437    fn id(&self) -> &ResourceId {
1438        ProcMeshRef::id(self).resource_id()
1439    }
1440
1441    fn region(&self) -> &ndslice::Region {
1442        ndslice::view::Ranked::region(self)
1443    }
1444
1445    fn subscribe_to_stream(
1446        &self,
1447        cx: &impl context::Actor,
1448        subscriber: hyperactor::PortRef<resource::RankedState<Self::StateInner>>,
1449    ) -> anyhow::Result<()> {
1450        // A `ProcMeshController` is only ever created for host-backed proc
1451        // meshes (host_mesh.rs spawn path), so a host mesh should always be
1452        // present. Surface a violation as an error rather than panicking — this
1453        // runs from `init`, so a panic would abort the actor. Cast a single
1454        // StreamState to the host-agent mesh so each host streams its procs'
1455        // state back through the cast tree (fanning in at cast actor 0) instead
1456        // of every host dialing the subscriber directly.
1457        let host_mesh = self.hosts().ok_or_else(|| {
1458            anyhow::anyhow!(
1459                "ProcMeshController has no host mesh; it must run on a host-backed proc mesh"
1460            )
1461        })?;
1462
1463        host_mesh.cast_stream_state(cx, ProcMeshRef::id(self).resource_id().clone(), subscriber)?;
1464        Ok(())
1465    }
1466
1467    fn forward_wait_rank_status(
1468        &self,
1469        cx: &impl context::Actor,
1470        msg: resource::WaitRankStatus,
1471    ) -> anyhow::Result<()> {
1472        let hosts = self.hosts().ok_or_else(|| {
1473            anyhow::anyhow!("ProcMeshController cannot wait without a backing host mesh")
1474        })?;
1475        hosts.forward_wait_rank_status(cx, self.proc_ids(), Ranked::region(self).clone(), msg)
1476    }
1477
1478    async fn poll_states(
1479        &self,
1480        cx: &impl context::Actor,
1481        supervision_display_name: &str,
1482        health_state: &mut HealthState,
1483    ) -> PollResult {
1484        let mesh_name = Controlled::id(self);
1485
1486        match self.states(cx, compute_keepalive()).await {
1487            Err(e) => send_poll_failure(
1488                cx,
1489                ActorSupervisionEvent::new(
1490                    cx.instance().self_addr().clone(),
1491                    Some(supervision_display_name.to_string()),
1492                    ActorStatus::generic_failure(format!(
1493                        "unable to query for proc states: {:?}",
1494                        e
1495                    )),
1496                    None,
1497                ),
1498                mesh_name,
1499                health_state,
1500            ),
1501            Ok(None) => PollResult::Processed { did_notify: false },
1502            Ok(Some(states)) => {
1503                let did_notify = health_state.apply_updates_and_notify(
1504                    &states,
1505                    |_point, state, health_state| {
1506                        self.notify_proc_state_change(
1507                            cx,
1508                            supervision_display_name,
1509                            state,
1510                            health_state,
1511                        )
1512                    },
1513                );
1514                PollResult::Processed { did_notify }
1515            }
1516        }
1517    }
1518
1519    fn process_state(
1520        &self,
1521        cx: &impl context::Actor,
1522        state: resource::RankedState<Self::StateInner>,
1523        health_state: &mut HealthState,
1524    ) -> bool {
1525        let resource::RankedState { rank, state } = state;
1526        let Ok(point) = Controlled::region(self)
1527            .extent()
1528            .point_of_rank(rank.unwrap())
1529        else {
1530            return false;
1531        };
1532        let changed = health_state.maybe_update(point, state.status.clone(), state.generation);
1533        if !changed {
1534            return false;
1535        }
1536        let display = Controlled::id(self).to_string();
1537        self.notify_proc_state_change(cx, &display, state, health_state)
1538    }
1539
1540    async fn handle_stop_request(
1541        &self,
1542        cx: &impl context::Actor,
1543        _supervision_display_name: &str,
1544        reason: String,
1545        health_state: &mut HealthState,
1546    ) -> anyhow::Result<()> {
1547        let mesh_name = Controlled::id(self);
1548        tracing::info!(
1549            actor_id = %cx.instance().self_addr(),
1550            proc_mesh = %mesh_name,
1551            "ProcMeshController stopping proc mesh"
1552        );
1553        // Marker so subscribers know the mesh is being torn down on request.
1554        let event = ActorSupervisionEvent::new(
1555            cx.instance().self_addr().clone(),
1556            None,
1557            ActorStatus::Stopped("ProcMeshController received explicit stop request".to_string()),
1558            None,
1559        );
1560        let failure_message = MeshFailure {
1561            actor_mesh_name: Some(mesh_name.to_string()),
1562            event,
1563            crashed_ranks: vec![],
1564            // MFCA-1: controller-generated stop report.
1565            reporting_controller: Some(cx.instance().self_addr().id().clone()),
1566        };
1567        health_state.unhealthy_event = Some(Unhealthy::StreamClosed(failure_message.clone()));
1568        for subscriber in health_state.subscribers.iter() {
1569            send_subscriber_message(cx, subscriber, failure_message.clone());
1570        }
1571
1572        let names = self.proc_ids().collect::<Vec<hyperactor::ProcAddr>>();
1573        let region = Ranked::region(self).clone();
1574        let Some(hosts) = self.hosts() else {
1575            return Ok(());
1576        };
1577        // stop_proc_mesh waits for every rank to reach a terminating state
1578        // before returning Ok, so we can apply its returned StatusMesh
1579        // verbatim. On error we still got per-rank statuses for whatever
1580        // ranks the host agents reported on; apply those too so health
1581        // state stays as accurate as we can make it.
1582        let max_rank = health_state.statuses.keys().map(|p| p.rank()).max();
1583        let extent = health_state
1584            .statuses
1585            .keys()
1586            .next()
1587            .map(|p| p.extent().clone());
1588        match hosts
1589            .stop_proc_mesh(cx, self.id(), names, region, reason)
1590            .await
1591        {
1592            Ok(statuses) => {
1593                for (rank, status) in statuses.iter() {
1594                    health_state
1595                        .statuses
1596                        .entry(rank)
1597                        .and_modify(move |s| *s = (status, u64::MAX));
1598                }
1599                Ok(())
1600            }
1601            Err(crate::Error::ProcMeshStopError { statuses }) => {
1602                if let (Some(max_rank), Some(extent)) = (max_rank, extent) {
1603                    for (rank, status) in statuses.materialized_iter(max_rank).enumerate() {
1604                        if let Ok(point) = extent.point_of_rank(rank) {
1605                            health_state
1606                                .statuses
1607                                .entry(point)
1608                                .and_modify(|s| *s = (status.clone(), u64::MAX));
1609                        }
1610                    }
1611                }
1612                Err(crate::Error::ProcMeshStopError { statuses }.into())
1613            }
1614            Err(e) => Err(e.into()),
1615        }
1616    }
1617
1618    async fn cleanup_stop(&self, cx: &impl context::Actor, reason: String) -> anyhow::Result<()> {
1619        let names = self.proc_ids().collect::<Vec<hyperactor::ProcAddr>>();
1620        let region = Ranked::region(self).clone();
1621        if let Some(hosts) = self.hosts() {
1622            hosts
1623                .stop_proc_mesh(cx, self.id(), names, region, reason)
1624                .await?;
1625        }
1626        Ok(())
1627    }
1628}
1629
1630impl ProcMeshRef {
1631    /// Translate a polled or streamed `State<ProcState>` into a supervision
1632    /// event on this proc-mesh controller. Returns `true` if a notification
1633    /// was sent (which suppresses the heartbeat path).
1634    fn notify_proc_state_change(
1635        &self,
1636        cx: &impl context::Actor,
1637        supervision_display_name: &str,
1638        state: resource::State<crate::host_mesh::host_agent::ProcState>,
1639        health_state: &mut HealthState,
1640    ) -> bool {
1641        let create_rank = state.state.as_ref().map(|s| s.create_rank);
1642        let actor_status = proc_status_to_actor_status(state.state.and_then(|s| s.proc_status));
1643        let event = ActorSupervisionEvent::new(
1644            cx.instance().self_addr().clone(),
1645            Some(supervision_display_name.to_string()),
1646            actor_status,
1647            None,
1648        );
1649        let rank = create_rank
1650            .and_then(|r| {
1651                ndslice::view::Ranked::region(self)
1652                    .extent()
1653                    .point_of_rank(r)
1654                    .ok()
1655            })
1656            .map(|p| p.rank())
1657            .unwrap_or(0);
1658        send_state_change(cx, rank, event, Controlled::id(self), true, health_state);
1659        true
1660    }
1661}
1662
1663#[cfg(test)]
1664mod tests {
1665    use std::ops::Deref;
1666    use std::time::Duration;
1667
1668    use hyperactor::actor::ActorErrorKind;
1669    use hyperactor::actor::ActorStatus;
1670    use hyperactor::channel::ChannelAddr;
1671    use hyperactor::id::Label;
1672    use hyperactor::supervision::ActorSupervisionEvent;
1673    use ndslice::Extent;
1674    use ndslice::ViewExt;
1675    use ndslice::view::CollectMeshExt;
1676
1677    use super::HealthState;
1678    use super::PollResult;
1679    #[cfg(fbcode_build)]
1680    use super::SUPERVISION_POLL_FREQUENCY;
1681    use super::actor_state_to_supervision_events;
1682    use super::proc_status_to_actor_status;
1683    use super::send_poll_failure;
1684    use super::send_state_change;
1685    use crate::ActorMesh;
1686    use crate::bootstrap::ProcStatus;
1687    #[cfg(fbcode_build)]
1688    use crate::host_mesh::PROC_SPAWN_MAX_IDLE;
1689    use crate::mesh_id::ActorMeshId;
1690    #[cfg(fbcode_build)]
1691    use crate::mesh_id::HostMeshId;
1692    use crate::mesh_id::ResourceId;
1693    use crate::proc_agent::ActorState;
1694    use crate::proc_agent::MESH_ORPHAN_TIMEOUT;
1695    use crate::proc_mesh::USE_DIRECT_SPAWN;
1696    use crate::resource;
1697    use crate::supervision::MeshFailure;
1698    use crate::test_utils::local_host_mesh;
1699    use crate::testactor;
1700    use crate::testing;
1701
1702    #[tokio::test]
1703    async fn poll_failure_consumes_one_terminal_rank_for_owner_notification_bound() {
1704        let instance = testing::instance();
1705        let (owner_port, mut owner_rx) = instance.open_port::<MeshFailure>();
1706        let mesh_name = ResourceId::instance(Label::new("workers").unwrap());
1707        let region: ndslice::Region = ndslice::extent!(gpus = 3).into();
1708        let statuses = (0..3)
1709            .map(|rank| {
1710                (
1711                    region.extent().point_of_rank(rank).unwrap(),
1712                    resource::Status::Running,
1713                )
1714            })
1715            .collect();
1716        let mut health_state = HealthState::new(statuses, Some(owner_port.bind()));
1717
1718        let rank0_event = failed_event(0, "rank 0 failed");
1719        assert!(
1720            health_state
1721                .mark_rank_terminating(0, resource::Status::Failed("rank 0 failed".to_string()))
1722        );
1723        send_state_change(
1724            &instance,
1725            0,
1726            rank0_event.clone(),
1727            &mesh_name,
1728            false,
1729            &mut health_state,
1730        );
1731        let rank0_failure = owner_rx.recv().await.unwrap();
1732        assert_eq!(rank0_failure.crashed_ranks, vec![0]);
1733        assert_eq!(rank0_failure.event, rank0_event);
1734
1735        let poll_event = failed_event(99, "unable to query for actor states");
1736        assert!(matches!(
1737            send_poll_failure(&instance, poll_event.clone(), &mesh_name, &mut health_state),
1738            PollResult::StopMonitoring
1739        ));
1740        let poll_failure = owner_rx.recv().await.unwrap();
1741        assert_eq!(poll_failure.crashed_ranks, vec![1]);
1742        assert_eq!(poll_failure.event, poll_event);
1743
1744        let late_rank1_event = failed_event(1, "late rank 1 failure");
1745        if health_state.mark_rank_terminating(
1746            1,
1747            resource::Status::Failed("late rank 1 failure".to_string()),
1748        ) {
1749            send_state_change(
1750                &instance,
1751                1,
1752                late_rank1_event,
1753                &mesh_name,
1754                false,
1755                &mut health_state,
1756            );
1757        }
1758        assert_eq!(owner_rx.try_recv().unwrap(), None);
1759
1760        let rank2_event = failed_event(2, "rank 2 failed");
1761        assert!(
1762            health_state
1763                .mark_rank_terminating(2, resource::Status::Failed("rank 2 failed".to_string()))
1764        );
1765        send_state_change(
1766            &instance,
1767            2,
1768            rank2_event.clone(),
1769            &mesh_name,
1770            false,
1771            &mut health_state,
1772        );
1773        let rank2_failure = owner_rx.recv().await.unwrap();
1774        assert_eq!(rank2_failure.crashed_ranks, vec![2]);
1775        assert_eq!(rank2_failure.event, rank2_event);
1776
1777        assert_eq!(health_state.first_non_terminating_rank(), None);
1778        assert_eq!(owner_rx.try_recv().unwrap(), None);
1779    }
1780
1781    // MFCA-1/MFCA-3: `send_state_change` stamps the reporting controller's id on
1782    // the owner post (and the subscriber copy) even when the event subject is a
1783    // ProcAgent rather than the controller or a mesh member.
1784    #[tokio::test]
1785    async fn owner_post_carries_reporting_controller_with_procagent_subject() {
1786        let instance = testing::instance();
1787        let (owner_port, mut owner_rx) = instance.open_port::<MeshFailure>();
1788        let (sub_port, mut sub_rx) = instance.open_port::<Option<MeshFailure>>();
1789        let mesh_name = ResourceId::instance(Label::new("workers").unwrap());
1790        let region: ndslice::Region = ndslice::extent!(gpus = 1).into();
1791        let statuses = std::iter::once((
1792            region.extent().point_of_rank(0).unwrap(),
1793            resource::Status::Running,
1794        ))
1795        .collect();
1796        let mut health_state = HealthState::new(statuses, Some(owner_port.bind()));
1797        health_state.subscribers.insert(sub_port.bind());
1798
1799        // Event subject is a ProcAgent, distinct from the controller (the
1800        // synthesized-timeout shape that defeats event-subject inference).
1801        let procagent = ResourceId::proc_addr_from_name(ChannelAddr::Local(0), "agent_proc")
1802            .actor_addr("proc_agent");
1803        let event = ActorSupervisionEvent::new(
1804            procagent.clone(),
1805            None,
1806            ActorStatus::Failed(ActorErrorKind::Generic(
1807                "unable to query for actor states".to_string(),
1808            )),
1809            None,
1810        );
1811        assert!(
1812            health_state.mark_rank_terminating(0, resource::Status::Failed("boom".to_string()))
1813        );
1814        send_state_change(&instance, 0, event, &mesh_name, false, &mut health_state);
1815
1816        let controller_id = instance.self_addr().id().clone();
1817        let owner_failure = owner_rx.recv().await.unwrap();
1818        // MFCA-1: the owner post carries the controller id.
1819        assert_eq!(
1820            owner_failure.reporting_controller,
1821            Some(controller_id.clone())
1822        );
1823        // MFCA-3: the event subject stays the ProcAgent, distinct from the reporter.
1824        assert_eq!(owner_failure.event.actor_id, procagent);
1825        assert_ne!(owner_failure.event.actor_id.id(), &controller_id);
1826
1827        // The subscriber copy carries the same reporter after the clone.
1828        let sub_failure = sub_rx
1829            .recv()
1830            .await
1831            .unwrap()
1832            .expect("subscriber failure is Some");
1833        assert_eq!(sub_failure.reporting_controller, Some(controller_id));
1834    }
1835
1836    fn failed_event(rank: usize, message: &str) -> ActorSupervisionEvent {
1837        ActorSupervisionEvent::new(
1838            ResourceId::proc_addr_from_name(ChannelAddr::Local(0), "test_proc")
1839                .actor_addr(format!("worker_{rank}")),
1840            None,
1841            ActorStatus::Failed(ActorErrorKind::Generic(message.to_string())),
1842            None,
1843        )
1844    }
1845
1846    #[test]
1847    fn polled_actor_state_uses_current_view_point() {
1848        let region: ndslice::Region = ndslice::extent!(gpus = 1).into();
1849        let point = region.extent().point_of_rank(0).unwrap();
1850        let actor_id = failed_event(0, "actor id").actor_id;
1851        let states = vec![resource::State {
1852            id: ResourceId::instance(Label::new("worker").unwrap()),
1853            status: resource::Status::Running,
1854            state: Some(ActorState {
1855                actor_id,
1856                create_rank: 9,
1857                supervision_events: vec![],
1858            }),
1859            generation: 1,
1860            timestamp: std::time::SystemTime::now(),
1861        }]
1862        .into_iter()
1863        .collect_mesh::<crate::ValueMesh<_>>(region.clone())
1864        .unwrap();
1865        let mut health_state = HealthState::new(
1866            vec![(point, resource::Status::Initializing)]
1867                .into_iter()
1868                .collect(),
1869            None,
1870        );
1871        let mut observed_rank = None;
1872
1873        health_state.apply_updates_and_notify(&states, |point, state, _| {
1874            observed_rank = Some(point.rank());
1875            assert_eq!(state.state.unwrap().create_rank, 9);
1876            true
1877        });
1878
1879        assert_eq!(observed_rank, Some(0));
1880    }
1881
1882    #[test]
1883    fn failed_actor_state_without_payload_uses_expected_actor() {
1884        let actor_id = failed_event(4, "actor id").actor_id;
1885        let state = resource::State {
1886            id: ResourceId::instance(Label::new("worker").unwrap()),
1887            status: resource::Status::Failed("spawn failed".to_string()),
1888            state: None,
1889            generation: 1,
1890            timestamp: std::time::SystemTime::now(),
1891        };
1892
1893        let events = actor_state_to_supervision_events(state, Some(actor_id.clone()));
1894
1895        assert_eq!(events.len(), 1);
1896        assert_eq!(events[0].actor_id, actor_id);
1897        assert!(matches!(events[0].actor_status, ActorStatus::Failed(_)));
1898    }
1899
1900    /// Wraps a host mesh's shutdown guard and the spawned host child
1901    /// processes so tests can simulate an unclean host crash by killing
1902    /// the children directly rather than asking an in-mesh actor to
1903    /// `process::exit`, which can also tear down the test binary.
1904    #[cfg(fbcode_build)]
1905    struct TestHostMesh {
1906        guard: crate::host_mesh::HostMeshShutdownGuard,
1907        children: Vec<tokio::process::Child>,
1908    }
1909
1910    #[cfg(fbcode_build)]
1911    impl TestHostMesh {
1912        async fn kill_hosts(&mut self) {
1913            for child in &mut self.children {
1914                let _ = child.start_kill();
1915                let _ = child.wait().await;
1916            }
1917            self.children.clear();
1918        }
1919    }
1920
1921    #[cfg(fbcode_build)]
1922    impl std::ops::Deref for TestHostMesh {
1923        type Target = crate::host_mesh::HostMeshShutdownGuard;
1924
1925        fn deref(&self) -> &Self::Target {
1926            &self.guard
1927        }
1928    }
1929
1930    #[cfg(fbcode_build)]
1931    impl std::ops::DerefMut for TestHostMesh {
1932        fn deref_mut(&mut self) -> &mut Self::Target {
1933            &mut self.guard
1934        }
1935    }
1936
1937    /// Verify that actors spawned without a controller are cleaned up
1938    /// when their keepalive expiry lapses. We:
1939    ///   1. Enable the orphan timeout on the `ProcMeshAgent`.
1940    ///   2. Spawn actors as *system actors* (no `ActorMeshController`).
1941    ///   3. Send a single keepalive with a short expiry time.
1942    ///   4. Wait for the expiry to pass and `SelfCheck` to fire.
1943    ///   5. Assert that the actors are now stopped.
1944    #[tokio::test]
1945    async fn test_orphaned_actors_are_cleaned_up() {
1946        let config = hyperactor_config::global::lock();
1947        let _mode = config.override_key(USE_DIRECT_SPAWN, false);
1948        // Short orphan timeout so SelfCheck fires frequently.
1949        let _orphan = config.override_key(MESH_ORPHAN_TIMEOUT, Some(Duration::from_secs(1)));
1950
1951        let instance = testing::instance();
1952        let host_mesh = local_host_mesh(2).await;
1953        let proc_mesh = host_mesh
1954            .spawn(instance, "test", Extent::unity(), None, None)
1955            .await
1956            .unwrap();
1957
1958        let actor_name = ActorMeshId::instance(Label::new("orphan_test").unwrap());
1959        // Spawn as a system actor so no controller is created. This lets us
1960        // control keepalive messages directly without the controller
1961        // interfering.
1962        let actor_mesh: ActorMesh<testactor::TestActor> = proc_mesh
1963            .spawn_with_name(instance, actor_name.clone(), &(), None, true)
1964            .await
1965            .unwrap();
1966        assert!(
1967            actor_mesh.deref().extent().num_ranks() > 0,
1968            "should have spawned at least one actor"
1969        );
1970
1971        // Send a keepalive with a short expiry. This is what the
1972        // ActorMeshController would normally do on each supervision poll.
1973        let states = proc_mesh
1974            .actor_states_with_keepalive(
1975                instance,
1976                actor_name.clone(),
1977                Some(std::time::SystemTime::now() + Duration::from_secs(2)),
1978            )
1979            .await
1980            .unwrap();
1981        // All actors should be running right now.
1982        for state in states.values() {
1983            assert_eq!(
1984                state.status,
1985                resource::Status::Running,
1986                "actor should be running before expiry"
1987            );
1988        }
1989
1990        // Poll until all actors are stopped, rather than sleeping a
1991        // fixed duration. The expiry is 2s and SelfCheck fires every 1s,
1992        // so this should converge quickly, but we allow a generous timeout
1993        // for slow CI environments.
1994        let deadline = tokio::time::Instant::now() + Duration::from_secs(30);
1995        loop {
1996            let states = proc_mesh
1997                .actor_states(instance, actor_name.clone())
1998                .await
1999                .unwrap();
2000            if states
2001                .values()
2002                .all(|s| s.status == resource::Status::Stopped)
2003            {
2004                break;
2005            }
2006            assert!(
2007                tokio::time::Instant::now() < deadline,
2008                "timed out waiting for actors to be stopped after keepalive expiry"
2009            );
2010            tokio::time::sleep(Duration::from_millis(200)).await;
2011        }
2012    }
2013
2014    /// Create a multi-process host mesh that propagates the current
2015    /// process's config overrides to child processes via Bootstrap.
2016    #[cfg(fbcode_build)]
2017    async fn host_mesh_with_config(n: usize) -> TestHostMesh {
2018        use hyperactor::channel::ChannelTransport;
2019        use tokio::process::Command;
2020
2021        let program = crate::testresource::get("monarch/hyperactor_mesh/bootstrap");
2022        let mut host_addrs = vec![];
2023        let mut children = Vec::new();
2024        let mut ready = Vec::new();
2025        for _ in 0..n {
2026            host_addrs.push(ChannelTransport::Unix.any());
2027        }
2028
2029        for host in host_addrs.iter() {
2030            let callback = crate::bootstrap::HostBootstrapReady::new(host.clone()).unwrap();
2031            let mut cmd = Command::new(program.clone());
2032            let boot = crate::Bootstrap::Host {
2033                addr: host.clone(),
2034                callback_addr: callback.callback_addr(),
2035                command: None,
2036                config: Some(hyperactor_config::global::attrs()),
2037                exit_on_shutdown: false,
2038            };
2039            boot.to_env(&mut cmd);
2040            cmd.kill_on_drop(false);
2041            // SAFETY: pre_exec sets PR_SET_PDEATHSIG so the child is
2042            // cleaned up if the parent (test) process dies.
2043            unsafe {
2044                cmd.pre_exec(crate::bootstrap::install_pdeathsig_kill);
2045            }
2046            children.push(cmd.spawn().unwrap());
2047            ready.push(callback);
2048        }
2049        for callback in ready {
2050            callback.wait().await.unwrap();
2051        }
2052
2053        let host_mesh = crate::HostMeshRef::from_hosts(
2054            HostMeshId::instance(Label::new("test").unwrap()),
2055            host_addrs,
2056        );
2057        TestHostMesh {
2058            guard: crate::host_mesh::HostMesh::take(host_mesh).shutdown_guard(),
2059            children,
2060        }
2061    }
2062
2063    /// Verify that actors are cleaned up via the orphan timeout when the
2064    /// `ActorMeshController`'s host process crashes. Unlike the system-actor
2065    /// test above, this spawns actors through a real controller (via
2066    /// `WrapperActor`) and then kills the controller's host process
2067    /// uncleanly. The agents on the surviving proc mesh detect the expired
2068    /// keepalive and stop the actors.
2069    #[tokio::test]
2070    #[cfg(fbcode_build)]
2071    async fn test_orphaned_actors_cleaned_up_on_controller_crash() {
2072        let config = hyperactor_config::global::lock();
2073        let _mode = config.override_key(USE_DIRECT_SPAWN, false);
2074        let _orphan = config.override_key(MESH_ORPHAN_TIMEOUT, Some(Duration::from_secs(2)));
2075        let _poll = config.override_key(SUPERVISION_POLL_FREQUENCY, Duration::from_secs(1));
2076        let _proc_spawn = config.override_key(PROC_SPAWN_MAX_IDLE, Duration::from_secs(60));
2077        let _host_spawn = config.override_key(
2078            hyperactor::config::HOST_SPAWN_READY_TIMEOUT,
2079            Duration::from_secs(60),
2080        );
2081
2082        let instance = testing::instance();
2083        let num_replicas = 1;
2084
2085        // Host mesh for the test actors (these survive the crash).
2086        // host_mesh_with_config propagates config overrides to child
2087        // processes via Bootstrap, so agents boot with
2088        // MESH_ORPHAN_TIMEOUT=2s and start the SelfCheck loop.
2089        let mut actor_hm = host_mesh_with_config(num_replicas).await;
2090        let actor_proc_mesh = actor_hm
2091            .spawn(instance, "actors", Extent::unity(), None, None)
2092            .await
2093            .unwrap();
2094
2095        // Host mesh for the wrapper + controller (will be killed).
2096        let mut controller_hm = host_mesh_with_config(1).await;
2097        let controller_proc_mesh = controller_hm
2098            .spawn(instance, "controller", Extent::unity(), None, None)
2099            .await
2100            .unwrap();
2101
2102        let child_name = ActorMeshId::instance(Label::new("orphan_child").unwrap());
2103
2104        // Supervision port required by WrapperActor params.
2105        let (supervision_port, _supervision_receiver) = instance.open_port::<MeshFailure>();
2106        let supervisor = supervision_port.bind();
2107
2108        // Spawn WrapperActor on controller_proc_mesh. Its init() spawns
2109        // ActorMesh<TestActor> on actor_proc_mesh with a real
2110        // ActorMeshController co-located on the controller's process.
2111        let _wrapper_mesh: ActorMesh<testactor::WrapperActor> = controller_proc_mesh
2112            .spawn(
2113                instance,
2114                "wrapper",
2115                &(
2116                    actor_proc_mesh.deref().clone(),
2117                    supervisor,
2118                    child_name.clone(),
2119                ),
2120            )
2121            .await
2122            .unwrap();
2123
2124        // Give the controller time to run at least one CheckState cycle
2125        // (polling every 1s). This is what sends `KeepaliveGetState` to
2126        // each agent, and is what arms the agent's `expiry_time` so the
2127        // agent's `SelfCheck` reaper can cull the actors after the
2128        // controller dies. Polling for `Running` is not enough: actors
2129        // reach `Running` at spawn time before the controller's first
2130        // poll, and if we kill the controller in that window the agents
2131        // never received a keepalive and the orphan timeout never trips.
2132        tokio::time::sleep(Duration::from_secs(3)).await;
2133        let states = actor_proc_mesh
2134            .actor_states(instance, child_name.clone())
2135            .await
2136            .unwrap();
2137        for state in states.values() {
2138            assert_eq!(
2139                state.status,
2140                resource::Status::Running,
2141                "actor should be running before controller crash"
2142            );
2143        }
2144
2145        // Kill the controller's host process uncleanly. The TestActors on
2146        // actor_proc_mesh survive. Killing the host (rather than asking the
2147        // wrapper actor to `process::exit`) is critical: the wrapper runs
2148        // in this same test binary's address space when the host mesh is
2149        // co-located, so an in-process exit would also tear down the test
2150        // runner.
2151        controller_hm.kill_hosts().await;
2152
2153        // Poll until all actors are stopped via the orphan timeout. The
2154        // configured timeout is 2s and `SelfCheck` fires every 2s, so this
2155        // converges quickly; allow generous slack for slow CI environments.
2156        let deadline = tokio::time::Instant::now() + Duration::from_secs(30);
2157        loop {
2158            let states = actor_proc_mesh
2159                .actor_states(instance, child_name.clone())
2160                .await
2161                .unwrap();
2162            if states
2163                .values()
2164                .all(|s| s.status == resource::Status::Stopped)
2165            {
2166                break;
2167            }
2168            assert!(
2169                tokio::time::Instant::now() < deadline,
2170                "timed out waiting for actors to be stopped after controller crash and orphan timeout"
2171            );
2172            tokio::time::sleep(Duration::from_millis(200)).await;
2173        }
2174
2175        let _ = actor_hm.shutdown(instance).await;
2176    }
2177
2178    #[test]
2179    fn test_proc_status_to_actor_status_stopped_cleanly() {
2180        let status = proc_status_to_actor_status(Some(ProcStatus::Stopped {
2181            exit_code: 0,
2182            stderr_tail: vec![],
2183        }));
2184        assert!(
2185            matches!(status, ActorStatus::Stopped(ref msg) if msg.contains("cleanly")),
2186            "expected Stopped, got {:?}",
2187            status
2188        );
2189    }
2190
2191    #[test]
2192    fn test_proc_status_to_actor_status_nonzero_exit() {
2193        let status = proc_status_to_actor_status(Some(ProcStatus::Stopped {
2194            exit_code: 1,
2195            stderr_tail: vec![],
2196        }));
2197        assert!(
2198            matches!(status, ActorStatus::Failed(_)),
2199            "expected Failed, got {:?}",
2200            status
2201        );
2202    }
2203
2204    #[test]
2205    fn test_proc_status_to_actor_status_stopping_is_not_a_failure() {
2206        let status = proc_status_to_actor_status(Some(ProcStatus::Stopping {
2207            started_at: std::time::SystemTime::now(),
2208        }));
2209        assert!(
2210            matches!(status, ActorStatus::Stopped(ref msg) if msg.contains("stopping")),
2211            "expected Stopped, got {:?}",
2212            status
2213        );
2214    }
2215
2216    #[test]
2217    fn test_proc_status_to_actor_status_none() {
2218        let status = proc_status_to_actor_status(None);
2219        assert!(
2220            matches!(status, ActorStatus::Stopped(_)),
2221            "expected Stopped, got {:?}",
2222            status
2223        );
2224    }
2225
2226    #[test]
2227    fn test_proc_status_to_actor_status_killed() {
2228        let status = proc_status_to_actor_status(Some(ProcStatus::Killed {
2229            signal: 9,
2230            core_dumped: false,
2231        }));
2232        assert!(
2233            matches!(status, ActorStatus::Failed(_)),
2234            "expected Failed, got {:?}",
2235            status
2236        );
2237    }
2238
2239    #[test]
2240    fn test_proc_status_to_actor_status_failed() {
2241        let status = proc_status_to_actor_status(Some(ProcStatus::Failed {
2242            reason: "oom".to_string(),
2243        }));
2244        assert!(
2245            matches!(status, ActorStatus::Failed(_)),
2246            "expected Failed, got {:?}",
2247            status
2248        );
2249    }
2250
2251    /// Force the supervision loop to look stalled by handing `check_stall` an
2252    /// `expected_time` several poll intervals in the past, then assert it both
2253    /// logs the warning and records the stall duration. This is the cheapest
2254    /// way to eyeball the warning the controller emits in production without
2255    /// standing up a real mesh.
2256    #[tracing_test::traced_test]
2257    #[test]
2258    fn test_check_stall_logs_when_late() {
2259        use std::time::SystemTime;
2260
2261        use hyperactor::id::ActorId;
2262        use hyperactor::id::ProcId;
2263
2264        let poll = hyperactor_config::global::get(super::SUPERVISION_POLL_FREQUENCY);
2265        // Pretend the CheckState handler was due five poll intervals ago.
2266        let expected_time = SystemTime::now() - poll * 5;
2267
2268        let proc = ProcId::instance(Label::new("stall_demo").unwrap());
2269        let actor_id = ActorId::singleton(Label::new("controller").unwrap(), proc);
2270
2271        super::check_stall(
2272            expected_time,
2273            &actor_id,
2274            &super::ACTOR_MESH_CONTROLLER_SUPERVISION_STALLS,
2275        );
2276
2277        // The reported lateness is `now - expected_time` (~5 poll intervals),
2278        // which humantime renders at full precision (e.g. "50s 343us 134ns").
2279        // Assert only on the stable prefix so the test stays deterministic.
2280        assert!(
2281            logs_contain("Handler<CheckState> is stalled by"),
2282            "expected a stall warning to be logged"
2283        );
2284    }
2285}