Skip to main content

hyperactor_mesh/
resource.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
9//! This modules defines a set of common message types used for managing resources
10//! in hyperactor meshes.
11
12pub mod mesh;
13
14use core::slice::GetDisjointMutIndex as _;
15use std::collections::HashMap;
16use std::fmt;
17use std::fmt::Debug;
18use std::hash::Hash;
19use std::mem::replace;
20use std::mem::take;
21use std::ops::Deref;
22use std::ops::DerefMut;
23use std::ops::Range;
24use std::time::Duration;
25
26use enum_as_inner::EnumAsInner;
27use hyperactor::HandleClient;
28use hyperactor::Handler;
29use hyperactor::PortRef;
30use hyperactor::RefClient;
31use hyperactor::RemoteMessage;
32use hyperactor::mailbox::PortReceiver;
33use hyperactor_config::attrs::Attrs;
34use ndslice::Region;
35use ndslice::ViewExt;
36use serde::Deserialize;
37use serde::Serialize;
38use typeuri::Named;
39
40use crate::StatusOverlay;
41use crate::ValueMesh;
42use crate::bootstrap;
43use crate::bootstrap::BootstrapCommand;
44use crate::bootstrap::ProcBind;
45use crate::host_mesh::host_agent::ProcState;
46use crate::mesh_id::ResourceId;
47use crate::proc_agent::ActorSpec;
48use crate::proc_agent::ActorState;
49
50/// The current lifecycle status of a resource.
51#[derive(
52    Clone,
53    Debug,
54    Serialize,
55    Deserialize,
56    Named,
57    PartialOrd,
58    Ord,
59    PartialEq,
60    Eq,
61    Hash,
62    EnumAsInner,
63    strum::Display
64)]
65pub enum Status {
66    /// The resource does not exist.
67    NotExist,
68    /// The resource is being created.
69    Initializing,
70    /// The resource is running.
71    Running,
72    /// The resource is being stopped.
73    Stopping,
74    /// The resource is stopped.
75    Stopped,
76    /// The resource has failed, with an error message.
77    #[strum(to_string = "Failed({0})")]
78    Failed(String),
79    /// The resource has been declared failed after a timeout.
80    #[strum(to_string = "Timeout({0:?})")]
81    Timeout(Duration),
82    /// The resource exists but its status is not known.
83    Unknown,
84}
85
86impl Status {
87    /// Returns whether the status is a terminating status (includes `Stopping`).
88    pub fn is_terminating(&self) -> bool {
89        matches!(
90            self,
91            Status::Stopping | Status::Stopped | Status::Failed(_) | Status::Timeout(_)
92        )
93    }
94
95    /// Tells whether the status represents a failure. A failure is both terminating
96    /// (the resource is not running), but also means abnormal exit (the resource
97    /// did not stop cleanly).
98    pub fn is_failure(&self) -> bool {
99        matches!(self, Self::Failed(_) | Self::Timeout(_))
100    }
101
102    /// Returns whether the status is fully terminal (the resource has
103    /// stopped, failed, or timed out — but NOT merely `Stopping`).
104    pub fn is_terminated(&self) -> bool {
105        matches!(
106            self,
107            Status::Stopped | Status::Failed(_) | Status::Timeout(_)
108        )
109    }
110
111    pub fn is_healthy(&self) -> bool {
112        matches!(self, Status::Initializing | Status::Running)
113    }
114
115    /// Ensure this status is at least as terminal as `floor`.
116    ///
117    /// If `floor` is a terminating status (Stopping, Stopped, Failed,
118    /// Timeout) and `self` is not, returns `floor`. Otherwise returns
119    /// `self` unchanged. This is used to prevent a child resource
120    /// from appearing healthier than its parent.
121    pub fn clamp_min(self, floor: Status) -> Status {
122        if floor.is_terminating() && !self.is_terminating() {
123            floor
124        } else {
125            self
126        }
127    }
128}
129
130impl From<bootstrap::ProcStatus> for Status {
131    fn from(status: bootstrap::ProcStatus) -> Self {
132        use bootstrap::ProcStatus;
133        match status {
134            ProcStatus::Starting => Status::Initializing,
135            ProcStatus::Running { .. } | ProcStatus::Ready { .. } => Status::Running,
136            ProcStatus::Stopping { .. } => Status::Stopping,
137            ProcStatus::Stopped { .. } => Status::Stopped,
138            ProcStatus::Failed { reason } => Status::Failed(reason),
139            ProcStatus::Killed { .. } => Status::Failed(format!("{}", status)),
140        }
141    }
142}
143
144impl From<crate::host::LocalProcStatus> for Status {
145    fn from(status: crate::host::LocalProcStatus) -> Self {
146        match status {
147            crate::host::LocalProcStatus::Stopping => Status::Stopping,
148            crate::host::LocalProcStatus::Stopped => Status::Stopped,
149        }
150    }
151}
152
153/// Data type used to communicate ranks.
154/// Serialized as a typed multipart part so comm actors can replace instances
155/// with the delivered rank while the message is in transit.
156#[derive(Clone, Debug, Named, PartialEq, Eq, Default)]
157pub struct Rank(pub Option<usize>);
158wirevalue::register_type!(Rank);
159
160/// Serialized representation for [`Rank`] multipart parts.
161#[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq, Eq)]
162pub struct RankRepr(pub Option<usize>);
163
164impl TryFrom<&Rank> for RankRepr {
165    type Error = serde_multipart::Error;
166
167    fn try_from(rank: &Rank) -> serde_multipart::Result<Self> {
168        Ok(Self(rank.0))
169    }
170}
171
172impl TryFrom<RankRepr> for Rank {
173    type Error = serde_multipart::Error;
174
175    fn try_from(repr: RankRepr) -> serde_multipart::Result<Self> {
176        Ok(Self(repr.0))
177    }
178}
179
180serde_multipart::part_codec! {
181    impl Rank
182    {
183        type Repr = RankRepr;
184    }
185}
186
187impl Rank {
188    /// Create a new rank with the provided value.
189    pub fn new(rank: usize) -> Self {
190        Self(Some(rank))
191    }
192
193    /// Unwrap the rank; panics if not set.
194    pub fn unwrap(&self) -> usize {
195        self.0.unwrap()
196    }
197}
198
199/// Get the status of a resource across the mesh.
200///
201/// This message is cast to all ranks; each rank replies with a sparse
202/// status **overlay**. The comm reducer merges overlays (right-wins)
203/// and the accumulator applies them to produce **full StatusMesh
204/// snapshots** on the receiver side.
205#[derive(
206    Clone,
207    Debug,
208    Serialize,
209    Deserialize,
210    Named,
211    Handler,
212    HandleClient,
213    RefClient
214)]
215pub struct GetRankStatus {
216    /// The resource identifier.
217    pub id: ResourceId,
218    /// Sparse status updates (overlays) from a rank.
219    pub reply: PortRef<StatusOverlay>,
220}
221
222/// Like [`GetRankStatus`], but the handler defers its reply until the
223/// resource's status is >= `min_status`. This avoids the race where
224/// the caller sees `Stopping` before the process has actually exited.
225#[derive(
226    Clone,
227    Debug,
228    Serialize,
229    Deserialize,
230    Named,
231    Handler,
232    HandleClient,
233    RefClient
234)]
235pub struct WaitRankStatus {
236    /// The resource identifier.
237    pub id: ResourceId,
238    /// The minimum status the caller wants to observe.
239    /// The handler will not reply until the resource's status
240    /// is >= this threshold.
241    pub min_status: Status,
242    /// Sparse status updates (overlays) from a rank.
243    pub reply: PortRef<StatusOverlay>,
244}
245
246/// Collect an accumulated [`ValueMesh<T>`] from `rx` until every rank has been
247/// reported or `max_idle_time` elapses with no update.
248///
249/// `reported` decides whether a rank's current value counts as "reported" (e.g.
250/// moved off a `NotExist`/`Timeout` placeholder). The accumulator emits the full
251/// mesh over the target region, so completion is "every cell reported". Returns
252/// `Ok(mesh)` once all ranks are reported, or `Err(mesh)` carrying the latest
253/// snapshot (or `fallback` if nothing arrived) on idle timeout / channel close.
254pub async fn wait_mesh<T>(
255    mut rx: PortReceiver<ValueMesh<T>>,
256    max_idle_time: Duration,
257    fallback: ValueMesh<T>,
258    reported: impl Fn(&T) -> bool,
259) -> Result<ValueMesh<T>, ValueMesh<T>>
260where
261    T: Send + Sync + Clone + 'static,
262{
263    let mut alarm = hyperactor::time::Alarm::new();
264    alarm.arm(max_idle_time);
265
266    // Latest-wins snapshot; `fallback` stands in until the first update arrives.
267    let mut snapshot = fallback;
268
269    loop {
270        let mut sleeper = alarm.sleeper();
271        tokio::select! {
272            _ = sleeper.sleep() => return Err(snapshot),
273            next = rx.recv() => {
274                match next {
275                    Ok(mesh) => { snapshot = mesh; }   // latest-wins snapshot
276                    Err(_)   => return Err(snapshot),
277                }
278            }
279        }
280
281        alarm.arm(max_idle_time);
282
283        // Short-circuit: done as soon as no rank is still unreported.
284        if snapshot.values().all(|v| reported(&v)) {
285            break Ok(snapshot);
286        }
287    }
288}
289
290impl GetRankStatus {
291    pub async fn wait(
292        rx: PortReceiver<crate::StatusMesh>,
293        num_ranks: usize,
294        max_idle_time: Duration,
295        region: Region, // used only for fallback
296    ) -> Result<crate::StatusMesh, crate::StatusMesh> {
297        debug_assert_eq!(region.num_ranks(), num_ranks, "region/num_ranks mismatch");
298
299        let fallback = crate::StatusMesh::from_single(region, crate::resource::Status::NotExist);
300        wait_mesh(rx, max_idle_time, fallback, |s| {
301            !matches!(s, crate::resource::Status::NotExist)
302        })
303        .await
304    }
305}
306
307/// The state of a resource.
308#[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq, Eq, Handler)]
309pub struct State<S> {
310    /// The resource identifier.
311    pub id: ResourceId,
312    /// Its status.
313    pub status: Status,
314    /// Optionally, a resource-defined state.
315    pub state: Option<S>,
316    /// Monotonic generation counter for last-writer-wins ordering.
317    pub generation: u64,
318    /// Wall-clock timestamp for debugging context.
319    pub timestamp: std::time::SystemTime,
320}
321wirevalue::register_type!(State<ActorState>);
322wirevalue::register_type!(State<ProcState>);
323
324impl<S: Serialize> fmt::Display for State<S> {
325    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
326        // Use serde_json to serialize the struct to a compact JSON string
327        match serde_json::to_string(self) {
328            Ok(json) => write!(f, "{}", json),
329            Err(e) => write!(f, "<state: serde_json error: {}>", e),
330        }
331    }
332}
333
334/// Create or update a resource according to a spec.
335#[derive(
336    Debug,
337    Clone,
338    Serialize,
339    Deserialize,
340    Named,
341    Handler,
342    HandleClient,
343    RefClient
344)]
345pub struct CreateOrUpdate<S> {
346    /// The resource identifier.
347    pub id: ResourceId,
348    /// The rank of the resource, when available.
349    pub rank: Rank,
350    /// The specification of the resource.
351    pub spec: S,
352}
353wirevalue::register_type!(CreateOrUpdate<ProcSpec>);
354wirevalue::register_type!(CreateOrUpdate<ActorSpec>);
355
356/// Stop a resource according to a spec.
357#[derive(
358    Debug,
359    Clone,
360    Serialize,
361    Deserialize,
362    Named,
363    Handler,
364    HandleClient,
365    RefClient
366)]
367pub struct Stop {
368    /// The resource identifier.
369    pub id: ResourceId,
370    /// The reason for stopping the resource.
371    pub reason: String,
372}
373wirevalue::register_type!(Stop);
374
375/// Stop all resources owned by the receiver of this message.
376/// No reply, this just issues the stop command.
377/// Use GetRankStatus to determine if it has successfully stopped.
378#[derive(
379    Debug,
380    Clone,
381    Serialize,
382    Deserialize,
383    Named,
384    Handler,
385    HandleClient,
386    RefClient
387)]
388pub struct StopAll {
389    /// The reason for stopping.
390    pub reason: String,
391}
392wirevalue::register_type!(StopAll);
393
394/// Retrieve the current state of the resource.
395#[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
396#[serde(bound(serialize = "S: Named", deserialize = "S: Named"))]
397pub struct GetState<S> {
398    /// The resource identifier.
399    pub id: ResourceId,
400    /// A reply containing the state.
401    #[reply]
402    pub reply: PortRef<State<S>>,
403}
404wirevalue::register_type!(GetState<ProcState>);
405wirevalue::register_type!(GetState<ActorState>);
406
407impl<S> Clone for GetState<S>
408where
409    S: RemoteMessage,
410{
411    fn clone(&self) -> Self {
412        Self {
413            id: self.id.clone(),
414            reply: self.reply.clone(),
415        }
416    }
417}
418
419/// Same as GetState, but additionally tells the receiver that the owner is still alive.
420/// If the receiver does not receive this message for a while, it might assume the owner is dead.
421#[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
422#[serde(bound(serialize = "S: Named", deserialize = "S: Named"))]
423pub struct KeepaliveGetState<S> {
424    /// The time at which the actor should be considered expired if no further
425    /// keepalive is received.
426    pub expires_after: std::time::SystemTime,
427    pub get_state: GetState<S>,
428}
429wirevalue::register_type!(KeepaliveGetState<ProcState>);
430wirevalue::register_type!(KeepaliveGetState<ActorState>);
431
432impl<S> Clone for KeepaliveGetState<S>
433where
434    S: RemoteMessage,
435{
436    fn clone(&self) -> Self {
437        Self {
438            expires_after: self.expires_after,
439            get_state: self.get_state.clone(),
440        }
441    }
442}
443
444/// Subscribe to streaming state updates for a named resource.
445/// The subscriber port will receive `State<S>` whenever the resource's
446/// state changes. The current state is sent immediately upon subscription.
447#[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
448#[serde(bound(serialize = "S: Named", deserialize = "S: Named"))]
449pub struct StreamState<S> {
450    /// The resource identifier.
451    pub id: ResourceId,
452    /// A streaming port that will receive state updates.
453    pub subscriber: PortRef<State<S>>,
454}
455wirevalue::register_type!(StreamState<ActorState>);
456wirevalue::register_type!(StreamState<ProcState>);
457
458impl<S> Clone for StreamState<S>
459where
460    S: RemoteMessage,
461{
462    fn clone(&self) -> Self {
463        Self {
464            id: self.id.clone(),
465            subscriber: self.subscriber.clone(),
466        }
467    }
468}
469
470/// List the set of resources managed by the controller.
471#[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
472pub struct List {
473    /// List of resource names managed by this controller.
474    #[reply]
475    pub reply: PortRef<Vec<ResourceId>>,
476}
477wirevalue::register_type!(List);
478
479/// A trait that bundles a set of types that together define a resource.
480pub trait Resource {
481    /// The spec specification for this resource.
482    type Spec: typeuri::Named
483        + Serialize
484        + for<'de> Deserialize<'de>
485        + Send
486        + Sync
487        + std::fmt::Debug;
488
489    /// The state for this resource.
490    type State: typeuri::Named
491        + Serialize
492        + for<'de> Deserialize<'de>
493        + Send
494        + Sync
495        + std::fmt::Debug;
496}
497
498// A behavior defining the interface for a mesh controller.
499hyperactor::behavior!(
500    Controller<R: Resource>,
501    CreateOrUpdate<R::Spec>,
502    GetState<R::State>,
503    Stop,
504);
505
506/// RankedValues compactly represents rank-indexed values of type T.
507/// It stores contiguous values in a set of intervals; thus it is
508/// efficient and compact when the cardinality of T-typed values is
509/// low.
510#[derive(Debug, Clone, Named, Serialize, Deserialize)]
511pub struct RankedValues<T> {
512    intervals: Vec<(Range<usize>, T)>,
513}
514
515impl<T: PartialEq> PartialEq for RankedValues<T> {
516    fn eq(&self, other: &Self) -> bool {
517        self.intervals == other.intervals
518    }
519}
520
521impl<T: Eq> Eq for RankedValues<T> {}
522
523impl<T> Default for RankedValues<T> {
524    fn default() -> Self {
525        Self {
526            intervals: Vec::new(),
527        }
528    }
529}
530
531impl<T> RankedValues<T> {
532    /// Iterate over contiguous rank intervals of values.
533    pub fn iter(&self) -> impl Iterator<Item = &(Range<usize>, T)> + '_ {
534        self.intervals.iter()
535    }
536
537    /// The (set) rank of the RankedValues is the number of values stored with
538    /// rank less than `value`.
539    pub fn rank(&self, value: usize) -> usize {
540        self.iter()
541            .take_while(|(ranks, _)| ranks.start <= value)
542            .map(|(ranks, _)| ranks.end.min(value) - ranks.start)
543            .sum()
544    }
545}
546
547impl<T: Clone> RankedValues<T> {
548    pub fn materialized_iter(&self, until: usize) -> impl Iterator<Item = &T> + '_ {
549        assert_eq!(self.rank(until), until, "insufficient rank");
550        self.iter()
551            .flat_map(|(range, value)| std::iter::repeat_n(value, range.end - range.start))
552    }
553}
554
555impl<T: Hash + Eq + Clone> RankedValues<T> {
556    /// Invert this ranked values into a [`ValuesByRank<T>`].
557    pub fn invert(&self) -> ValuesByRank<T> {
558        let mut inverted: HashMap<T, Vec<Range<usize>>> = HashMap::new();
559        for (range, value) in self.iter() {
560            inverted
561                .entry(value.clone())
562                .or_default()
563                .push(range.clone());
564        }
565        ValuesByRank { values: inverted }
566    }
567}
568
569impl<T: Eq + Clone> RankedValues<T> {
570    /// Merge `other` into this set of ranked values. Values in `other` that overlap
571    /// with `self` take prededence.
572    ///
573    /// This currently uses a simple algorithm that merges the full set of RankedValues.
574    /// This remains efficient when the cardinality of T-typed values is low. However,
575    /// it does not efficiently merge high cardinality value sets. Consider using interval
576    /// trees or bitmap techniques like Roaring Bitmaps in these cases.
577    pub fn merge_from(&mut self, other: Self) {
578        let mut left_iter = take(&mut self.intervals).into_iter();
579        let mut right_iter = other.intervals.into_iter();
580
581        let mut left = left_iter.next();
582        let mut right = right_iter.next();
583
584        while left.is_some() && right.is_some() {
585            let (left_ranks, left_value) = left.as_mut().unwrap();
586            let (right_ranks, right_value) = right.as_mut().unwrap();
587
588            if left_ranks.is_overlapping(right_ranks) {
589                if left_value == right_value {
590                    let ranks = left_ranks.start.min(right_ranks.start)..right_ranks.end;
591                    let (_, value) = replace(&mut right, right_iter.next()).unwrap();
592                    left_ranks.start = ranks.end;
593                    if left_ranks.is_empty() {
594                        left = left_iter.next();
595                    }
596                    self.append(ranks, value);
597                } else if left_ranks.start < right_ranks.start {
598                    let ranks = left_ranks.start..right_ranks.start;
599                    left_ranks.start = ranks.end;
600                    // TODO: get rid of clone
601                    self.append(ranks, left_value.clone());
602                } else {
603                    let (ranks, value) = replace(&mut right, right_iter.next()).unwrap();
604                    left_ranks.start = ranks.end;
605                    if left_ranks.is_empty() {
606                        left = left_iter.next();
607                    }
608                    self.append(ranks, value);
609                }
610            } else if left_ranks.start < right_ranks.start {
611                let (ranks, value) = replace(&mut left, left_iter.next()).unwrap();
612                self.append(ranks, value);
613            } else {
614                let (ranks, value) = replace(&mut right, right_iter.next()).unwrap();
615                self.append(ranks, value);
616            }
617        }
618
619        while let Some((left_ranks, left_value)) = left {
620            self.append(left_ranks, left_value);
621            left = left_iter.next();
622        }
623        while let Some((right_ranks, right_value)) = right {
624            self.append(right_ranks, right_value);
625            right = right_iter.next();
626        }
627    }
628
629    /// Merge the contents of this RankedValues into another RankedValues.
630    pub fn merge_into(self, other: &mut Self) {
631        other.merge_from(self);
632    }
633
634    fn append(&mut self, range: Range<usize>, value: T) {
635        if let Some(last) = self.intervals.last_mut()
636            && last.0.end == range.start
637            && last.1 == value
638        {
639            last.0.end = range.end;
640        } else {
641            self.intervals.push((range, value));
642        }
643    }
644}
645
646impl RankedValues<Status> {
647    pub fn first_terminating(&self) -> Option<(usize, Status)> {
648        self.intervals
649            .iter()
650            .find(|(_, status)| status.is_terminating())
651            .map(|(range, status)| (range.start, status.clone()))
652    }
653
654    pub fn first_failed(&self) -> Option<(usize, Status)> {
655        self.intervals
656            .iter()
657            .find(|(_, status)| matches!(status, Status::Failed(_) | Status::Timeout(_)))
658            .map(|(range, status)| (range.start, status.clone()))
659    }
660}
661
662impl<T> From<(usize, T)> for RankedValues<T> {
663    fn from((rank, value): (usize, T)) -> Self {
664        Self {
665            intervals: vec![(rank..rank + 1, value)],
666        }
667    }
668}
669
670impl<T> From<(Range<usize>, T)> for RankedValues<T> {
671    fn from((range, value): (Range<usize>, T)) -> Self {
672        Self {
673            intervals: vec![(range, value)],
674        }
675    }
676}
677
678/// An inverted index of RankedValues, providing all ranks for
679/// which each unique T-typed value appears.
680#[derive(Clone, Debug)]
681pub struct ValuesByRank<T> {
682    values: HashMap<T, Vec<Range<usize>>>,
683}
684
685impl<T: Eq + Hash> PartialEq for ValuesByRank<T> {
686    fn eq(&self, other: &Self) -> bool {
687        self.values == other.values
688    }
689}
690
691impl<T: Eq + Hash> Eq for ValuesByRank<T> {}
692
693impl<T: fmt::Display> fmt::Display for ValuesByRank<T> {
694    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
695        let mut first_value = true;
696        for (value, ranges) in self.iter() {
697            if first_value {
698                first_value = false;
699            } else {
700                write!(f, ";")?;
701            }
702            write!(f, "{}=", value)?;
703            let mut first_range = true;
704            for range in ranges.iter() {
705                if first_range {
706                    first_range = false;
707                } else {
708                    write!(f, ",")?;
709                }
710                write!(f, "{}..{}", range.start, range.end)?;
711            }
712        }
713        Ok(())
714    }
715}
716
717impl<T> Deref for ValuesByRank<T> {
718    type Target = HashMap<T, Vec<Range<usize>>>;
719
720    fn deref(&self) -> &Self::Target {
721        &self.values
722    }
723}
724
725impl<T> DerefMut for ValuesByRank<T> {
726    fn deref_mut(&mut self) -> &mut Self::Target {
727        &mut self.values
728    }
729}
730
731/// Enabled for test only because we have to guarantee that the input
732/// iterator is well-formed.
733#[cfg(test)]
734impl<T> FromIterator<(Range<usize>, T)> for RankedValues<T> {
735    fn from_iter<I: IntoIterator<Item = (Range<usize>, T)>>(iter: I) -> Self {
736        Self {
737            intervals: iter.into_iter().collect(),
738        }
739    }
740}
741
742/// Spec for a host mesh agent to use when spawning a new proc.
743#[derive(Clone, Debug, Serialize, Deserialize, Named, Default)]
744pub(crate) struct ProcSpec {
745    /// Config values to set on the spawned proc's global config,
746    /// at the `ClientOverride` layer.
747    pub(crate) client_config_override: Attrs,
748    /// Optional per-process CPU/NUMA binding configuration.
749    pub(crate) proc_bind: Option<ProcBind>,
750    /// Optional bootstrap command override. When set, this command is used
751    /// to spawn the proc instead of the host agent's default bootstrap command.
752    pub(crate) bootstrap_command: Option<BootstrapCommand>,
753    /// The id of the HostMesh that owns this proc. Used by
754    /// `DrainHost` to selectively drain only procs belonging to a
755    /// specific mesh.
756    pub(crate) host_mesh_id: Option<crate::mesh_id::HostMeshId>,
757    /// The id of the ProcMesh that owns this proc. A host agent can hold procs
758    /// from several proc meshes, so this lets per-mesh queries (e.g.
759    /// `StreamState`) scope to a single mesh rather than every proc on the host.
760    pub(crate) proc_mesh_id: Option<crate::mesh_id::ProcMeshId>,
761}
762wirevalue::register_type!(ProcSpec);
763
764#[cfg(test)]
765mod tests {
766    use hyperactor::port::Port;
767
768    use super::*;
769
770    #[test]
771    fn handler_ports_are_distinct_for_resource_messages() {
772        assert_ne!(
773            Port::handler::<CreateOrUpdate<ProcSpec>>(),
774            Port::handler::<Stop>(),
775        );
776    }
777
778    #[test]
779    fn test_ranked_values_merge() {
780        #[derive(PartialEq, Debug, Eq, Clone)]
781        enum Side {
782            Left,
783            Right,
784            Both,
785        }
786        use Side::Both;
787        use Side::Left;
788        use Side::Right;
789
790        let mut left: RankedValues<Side> = [
791            (0..10, Left),
792            (15..20, Left),
793            (30..50, Both),
794            (60..70, Both),
795        ]
796        .into_iter()
797        .collect();
798
799        let right: RankedValues<Side> = [
800            (9..12, Right),
801            (25..30, Right),
802            (30..40, Both),
803            (40..50, Right),
804            (50..60, Both),
805        ]
806        .into_iter()
807        .collect();
808
809        left.merge_from(right);
810        assert_eq!(
811            left.iter().cloned().collect::<Vec<_>>(),
812            vec![
813                (0..9, Left),
814                (9..12, Right),
815                (15..20, Left),
816                (25..30, Right),
817                (30..40, Both),
818                (40..50, Right),
819                // Merge consecutive:
820                (50..70, Both)
821            ]
822        );
823
824        assert_eq!(left.rank(5), 5);
825        assert_eq!(left.rank(10), 10);
826        assert_eq!(left.rank(16), 13);
827        assert_eq!(left.rank(70), 62);
828        assert_eq!(left.rank(100), 62);
829    }
830
831    #[test]
832    fn test_equality() {
833        assert_eq!(
834            RankedValues::from((0..10, 123)),
835            RankedValues::from((0..10, 123))
836        );
837        assert_eq!(
838            RankedValues::from((0..10, Status::Failed("foo".to_string()))),
839            RankedValues::from((0..10, Status::Failed("foo".to_string()))),
840        );
841    }
842
843    #[test]
844    fn test_default_through_merging() {
845        let values: RankedValues<usize> =
846            [(0..10, 1), (15..20, 1), (30..50, 1)].into_iter().collect();
847
848        let mut default = RankedValues::from((0..50, 0));
849        default.merge_from(values);
850
851        assert_eq!(
852            default.iter().cloned().collect::<Vec<_>>(),
853            vec![
854                (0..10, 1),
855                (10..15, 0),
856                (15..20, 1),
857                (20..30, 0),
858                (30..50, 1)
859            ]
860        );
861    }
862}