1pub 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#[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 NotExist,
68 Initializing,
70 Running,
72 Stopping,
74 Stopped,
76 #[strum(to_string = "Failed({0})")]
78 Failed(String),
79 #[strum(to_string = "Timeout({0:?})")]
81 Timeout(Duration),
82 Unknown,
84}
85
86impl Status {
87 pub fn is_terminating(&self) -> bool {
89 matches!(
90 self,
91 Status::Stopping | Status::Stopped | Status::Failed(_) | Status::Timeout(_)
92 )
93 }
94
95 pub fn is_failure(&self) -> bool {
99 matches!(self, Self::Failed(_) | Self::Timeout(_))
100 }
101
102 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 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#[derive(Clone, Debug, Named, PartialEq, Eq, Default)]
157pub struct Rank(pub Option<usize>);
158wirevalue::register_type!(Rank);
159
160#[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 pub fn new(rank: usize) -> Self {
190 Self(Some(rank))
191 }
192
193 pub fn unwrap(&self) -> usize {
195 self.0.unwrap()
196 }
197}
198
199#[derive(
206 Clone,
207 Debug,
208 Serialize,
209 Deserialize,
210 Named,
211 Handler,
212 HandleClient,
213 RefClient
214)]
215pub struct GetRankStatus {
216 pub id: ResourceId,
218 pub reply: PortRef<StatusOverlay>,
220}
221
222#[derive(
226 Clone,
227 Debug,
228 Serialize,
229 Deserialize,
230 Named,
231 Handler,
232 HandleClient,
233 RefClient
234)]
235pub struct WaitRankStatus {
236 pub id: ResourceId,
238 pub min_status: Status,
242 pub reply: PortRef<StatusOverlay>,
244}
245
246pub 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 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; } Err(_) => return Err(snapshot),
277 }
278 }
279 }
280
281 alarm.arm(max_idle_time);
282
283 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, ) -> 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#[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq, Eq, Handler)]
309pub struct State<S> {
310 pub id: ResourceId,
312 pub status: Status,
314 pub state: Option<S>,
316 pub generation: u64,
318 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 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#[derive(
336 Debug,
337 Clone,
338 Serialize,
339 Deserialize,
340 Named,
341 Handler,
342 HandleClient,
343 RefClient
344)]
345pub struct CreateOrUpdate<S> {
346 pub id: ResourceId,
348 pub rank: Rank,
350 pub spec: S,
352}
353wirevalue::register_type!(CreateOrUpdate<ProcSpec>);
354wirevalue::register_type!(CreateOrUpdate<ActorSpec>);
355
356#[derive(
358 Debug,
359 Clone,
360 Serialize,
361 Deserialize,
362 Named,
363 Handler,
364 HandleClient,
365 RefClient
366)]
367pub struct Stop {
368 pub id: ResourceId,
370 pub reason: String,
372}
373wirevalue::register_type!(Stop);
374
375#[derive(
379 Debug,
380 Clone,
381 Serialize,
382 Deserialize,
383 Named,
384 Handler,
385 HandleClient,
386 RefClient
387)]
388pub struct StopAll {
389 pub reason: String,
391}
392wirevalue::register_type!(StopAll);
393
394#[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
396#[serde(bound(serialize = "S: Named", deserialize = "S: Named"))]
397pub struct GetState<S> {
398 pub id: ResourceId,
400 #[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#[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
422#[serde(bound(serialize = "S: Named", deserialize = "S: Named"))]
423pub struct KeepaliveGetState<S> {
424 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#[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
448#[serde(bound(serialize = "S: Named", deserialize = "S: Named"))]
449pub struct StreamState<S> {
450 pub id: ResourceId,
452 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#[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
472pub struct List {
473 #[reply]
475 pub reply: PortRef<Vec<ResourceId>>,
476}
477wirevalue::register_type!(List);
478
479pub trait Resource {
481 type Spec: typeuri::Named
483 + Serialize
484 + for<'de> Deserialize<'de>
485 + Send
486 + Sync
487 + std::fmt::Debug;
488
489 type State: typeuri::Named
491 + Serialize
492 + for<'de> Deserialize<'de>
493 + Send
494 + Sync
495 + std::fmt::Debug;
496}
497
498hyperactor::behavior!(
500 Controller<R: Resource>,
501 CreateOrUpdate<R::Spec>,
502 GetState<R::State>,
503 Stop,
504);
505
506#[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 pub fn iter(&self) -> impl Iterator<Item = &(Range<usize>, T)> + '_ {
534 self.intervals.iter()
535 }
536
537 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 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 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 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 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#[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#[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#[derive(Clone, Debug, Serialize, Deserialize, Named, Default)]
744pub(crate) struct ProcSpec {
745 pub(crate) client_config_override: Attrs,
748 pub(crate) proc_bind: Option<ProcBind>,
750 pub(crate) bootstrap_command: Option<BootstrapCommand>,
753 pub(crate) host_mesh_id: Option<crate::mesh_id::HostMeshId>,
757 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 (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}