1#![allow(clippy::result_large_err)]
47
48use hyperactor::ActorRef;
49use hyperactor::Endpoint as _;
50use hyperactor::Gateway;
51use hyperactor::Handler;
52use hyperactor::accum::IdleFlushReducerOpts;
53use hyperactor::channel::ChannelTransport;
54use hyperactor::id::Label;
55use hyperactor::id::Uid;
56use hyperactor_cast::cast_actor::CastActor;
57use hyperactor_config::CONFIG;
58use hyperactor_config::ConfigAttr;
59use hyperactor_config::NonZeroUsize;
60use hyperactor_config::attrs::declare_attrs;
61use ndslice::view::CollectMeshExt;
62
63use crate::mesh_admin::MeshAdminAgent;
64use crate::supervision::MeshFailure;
65
66pub mod host_agent;
67
68use std::collections::HashSet;
69use std::ops::Deref;
70use std::ops::DerefMut;
71use std::str::FromStr;
72use std::sync::Arc;
73use std::time::Duration;
74
75use hyperactor::ActorAddr;
76use hyperactor::ActorId;
77use hyperactor::ProcAddr;
78use hyperactor::ProcId;
79use hyperactor::channel::ChannelAddr;
80use hyperactor::context;
81use hyperactor_cast::cast_actor::CAST_ACTOR_NAME;
82use ndslice::Extent;
83use ndslice::Region;
84use ndslice::ViewExt;
85use ndslice::extent;
86use ndslice::view;
87use ndslice::view::Ranked;
88use ndslice::view::RegionParseError;
89use serde::Deserialize;
90use serde::Serialize;
91use tracing::Instrument;
92use typeuri::Named;
93
94use crate::ActorMeshRef;
95use crate::Bootstrap;
96use crate::ProcMesh;
97use crate::ValueMesh;
98use crate::bootstrap::BootstrapCommand;
99use crate::bootstrap::BootstrapProcManager;
100use crate::bootstrap::HostBootstrapReady;
101use crate::bootstrap::ProcBind;
102use crate::host::Host;
103use crate::host::LocalProcManager;
104use crate::host::SERVICE_PROC_NAME;
105use crate::host::legacy_service_proc_id;
106pub use crate::host_mesh::host_agent::HostAgent;
107use crate::host_mesh::host_agent::ProcManagerSpawnFn;
108use crate::host_mesh::host_agent::ProcState;
109use crate::mesh_controller::ProcMeshController;
110use crate::mesh_id::ActorMeshId;
111use crate::mesh_id::HostMeshId;
112use crate::mesh_id::ProcMeshId;
113use crate::mesh_id::ResourceId;
114use crate::proc_agent::ProcAgent;
115use crate::proc_mesh::ProcMeshRef;
116use crate::resource;
117use crate::resource::GetRankStatus;
118use crate::resource::RankedValues;
119use crate::resource::Status;
120use crate::resource::WaitRankStatusClient;
121use crate::transport::DEFAULT_TRANSPORT;
122
123pub const PROC_MESH_CONTROLLER_NAME: &str = "proc_mesh_controller";
125
126declare_attrs! {
127 @meta(CONFIG = ConfigAttr::new(
130 Some("HYPERACTOR_MESH_PROC_SPAWN_MAX_IDLE".to_string()),
131 Some("mesh_proc_spawn_max_idle".to_string()),
132 ))
133 pub attr PROC_SPAWN_MAX_IDLE: Duration = Duration::from_secs(30);
134
135 @meta(CONFIG = ConfigAttr::new(
138 Some("HYPERACTOR_MESH_PROC_STOP_MAX_IDLE".to_string()),
139 Some("proc_stop_max_idle".to_string()),
140 ))
141 pub attr PROC_STOP_MAX_IDLE: Duration = Duration::from_secs(30);
142
143 @meta(CONFIG = ConfigAttr::new(
146 Some("HYPERACTOR_MESH_GET_PROC_STATE_MAX_IDLE".to_string()),
147 Some("get_proc_state_max_idle".to_string()),
148 ))
149 pub attr GET_PROC_STATE_MAX_IDLE: Duration = Duration::from_mins(1);
150}
151
152fn legacy_service_proc_addr(host_addr: ChannelAddr) -> ProcAddr {
153 ProcAddr::new(legacy_service_proc_id(), host_addr.into_dial_addr().into())
154}
155
156pub(crate) fn host_agent_ref(service_proc: ProcAddr) -> ActorRef<HostAgent> {
157 ActorRef::attest(service_proc.actor_addr(host_agent::HOST_MESH_AGENT_ACTOR_NAME))
158}
159
160pub(crate) fn legacy_host_agent_ref(host_addr: ChannelAddr) -> ActorRef<HostAgent> {
161 host_agent_ref(legacy_service_proc_addr(host_addr))
162}
163
164fn named_proc_on_host(agent: &ActorRef<HostAgent>, id: &ResourceId) -> ProcAddr {
165 let location = agent
166 .actor_addr()
167 .proc_addr()
168 .location()
169 .clone()
170 .append_via(id.uid().clone());
171 id.proc_addr(location)
172}
173
174#[derive(Debug, thiserror::Error)]
184pub enum ConfigPushFailure {
185 #[error("send failed: {0}")]
188 SendFailed(#[source] Box<hyperactor::mailbox::MailboxSenderError>),
189
190 #[error("cast failed: {0}")]
193 CastFailed(String),
194
195 #[error("reply timed out after MESH_ATTACH_CONFIG_TIMEOUT")]
202 ReplyTimedOut,
203
204 #[error("reply channel closed before reply")]
207 ReplyChannelClosed,
208}
209
210#[derive(Debug)]
218pub struct ConfigPushError {
219 pub failures: Vec<(ChannelAddr, ConfigPushFailure)>,
221}
222
223impl std::fmt::Display for ConfigPushError {
224 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
225 write!(
226 f,
227 "config push failed during attach on {} host(s):",
228 self.failures.len()
229 )?;
230 for (host, failure) in &self.failures {
231 write!(f, "\n - {}: {}", host, failure)?;
232 }
233 Ok(())
234 }
235}
236
237impl std::error::Error for ConfigPushError {
238 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
239 None
243 }
244}
245
246pub struct HostMesh {
257 id: HostMeshId,
258 extent: Extent,
259 current_ref: HostMeshRef,
260 shutdown_on_drop: bool,
263}
264
265impl HostMesh {
266 fn notify_created(&self) {
268 let name_str = self.id.to_string();
269 let mesh_id_hash = hyperactor_telemetry::hash_to_u64(&self.id);
270
271 hyperactor_telemetry::notify_mesh_created(hyperactor_telemetry::MeshEvent {
272 id: mesh_id_hash,
273 timestamp: std::time::SystemTime::now(),
274 class: "Host".to_string(),
275 given_name: self
276 .id
277 .display_label()
278 .map(|l| l.as_str())
279 .unwrap_or("unnamed")
280 .to_string(),
281 full_name: name_str,
282 shape_json: serde_json::to_string(&self.extent).unwrap_or_default(),
283 parent_mesh_id: None,
284 parent_view_json: None,
285 });
286
287 let now = std::time::SystemTime::now();
290 for (rank, actor) in self.current_ref.host_agent_mesh.values().enumerate() {
291 hyperactor_telemetry::notify_actor_created(hyperactor_telemetry::ActorEvent {
292 id: hyperactor_telemetry::hash_to_u64(actor.actor_addr().id()),
293 timestamp: now,
294 mesh_id: mesh_id_hash,
295 rank: rank as u64,
296 full_name: actor.actor_addr().to_string(),
297 display_name: None,
298 });
299 }
300 }
301
302 pub async fn local() -> crate::Result<HostMesh> {
325 Self::local_with_bootstrap(BootstrapCommand::current()?).await
326 }
327
328 pub async fn local_with_bootstrap(bootstrap_cmd: BootstrapCommand) -> crate::Result<HostMesh> {
336 if let Ok(Some(boot)) = Bootstrap::get_from_env() {
337 let result = boot.bootstrap().await;
338 if let Err(err) = result {
339 tracing::error!("failed to bootstrap local host mesh process: {}", err);
340 }
341 std::process::exit(1);
342 }
343
344 let addr = hyperactor_config::global::get_cloned(DEFAULT_TRANSPORT).binding_addr();
345
346 let manager = BootstrapProcManager::new(bootstrap_cmd)?;
347 let host = Host::new_with_gateway(
352 manager,
353 ProcAddr::new(legacy_service_proc_id(), addr.into()),
354 None,
355 Gateway::new(),
356 None,
357 )
358 .await?;
359 let addr = host.addr().clone();
360 let system_proc = host.system_proc().clone();
361 let host_mesh_agent = system_proc
362 .spawn_bound_with_uid(
363 Uid::singleton(Label::new(host_agent::HOST_MESH_AGENT_ACTOR_NAME).unwrap()),
364 HostAgent::new_process(host, None),
365 )
366 .map_err(crate::Error::SingletonActorSpawnError)?;
367 HostAgent::wait_initialized(&host_mesh_agent).await?;
368 system_proc
369 .spawn_bound_with_uid(
370 Uid::singleton(Label::strip(CAST_ACTOR_NAME)),
371 CastActor::default(),
372 )
373 .map_err(crate::Error::SingletonActorSpawnError)?;
374
375 let host_mesh_ref = HostMeshRef::new(
376 HostMeshId::instance(Label::new("local").unwrap()),
377 extent!(hosts = 1).into(),
378 vec![addr],
379 )?;
380 Ok(HostMesh::take(host_mesh_ref))
381 }
382
383 pub async fn local_in_process() -> crate::Result<HostMesh> {
393 let addr = hyperactor_config::global::get_cloned(DEFAULT_TRANSPORT).binding_addr();
394 Ok(HostMesh::take(Self::local_n_in_process(vec![addr]).await?))
395 }
396
397 pub(crate) async fn local_n_in_process(
406 host_addrs: Vec<ChannelAddr>,
407 ) -> crate::Result<HostMeshRef> {
408 let n = host_addrs.len();
409 let mut in_process_host_addrs = Vec::with_capacity(n);
410 for host_addr in host_addrs {
411 in_process_host_addrs.push(Self::create_in_process_host(host_addr).await?);
412 }
413 HostMeshRef::new(
414 HostMeshId::instance(Label::new("local").unwrap()),
415 extent!(hosts = n).into(),
416 in_process_host_addrs,
417 )
418 }
419
420 async fn create_in_process_host(addr: ChannelAddr) -> crate::Result<ChannelAddr> {
423 let spawn: ProcManagerSpawnFn =
424 Box::new(|proc| Box::pin(std::future::ready(ProcAgent::boot_v1(proc, None))));
425 let manager = LocalProcManager::new(spawn);
426 let host = Host::new_with_gateway(
431 manager,
432 ProcAddr::new(legacy_service_proc_id(), addr.into()),
433 None,
434 Gateway::new(),
435 None,
436 )
437 .await?;
438 let addr = host.addr().clone();
439 let system_proc = host.system_proc().clone();
440 let host_mesh_agent = system_proc
441 .spawn_bound_with_uid(
442 Uid::singleton(Label::new(host_agent::HOST_MESH_AGENT_ACTOR_NAME).unwrap()),
443 HostAgent::new_local(host),
444 )
445 .map_err(crate::Error::SingletonActorSpawnError)?;
446 HostAgent::wait_initialized(&host_mesh_agent).await?;
447
448 system_proc
449 .spawn_bound_with_uid(
450 Uid::singleton(Label::strip(CAST_ACTOR_NAME)),
451 CastActor::default(),
452 )
453 .map_err(crate::Error::SingletonActorSpawnError)?;
454
455 Ok(addr)
456 }
457
458 pub async fn process(extent: Extent, command: BootstrapCommand) -> crate::Result<HostMesh> {
469 if let Ok(Some(boot)) = Bootstrap::get_from_env() {
470 let result = boot.bootstrap().await;
471 if let Err(err) = result {
472 tracing::error!("failed to bootstrap process host mesh process: {}", err);
473 }
474 std::process::exit(1);
475 }
476
477 let bind_spec = hyperactor_config::global::get_cloned(DEFAULT_TRANSPORT);
478 let mut host_addrs = Vec::with_capacity(extent.num_ranks());
479 let mut ready = Vec::with_capacity(extent.num_ranks());
480 for _ in 0..extent.num_ranks() {
481 let addr = bind_spec.binding_addr();
482 let callback = HostBootstrapReady::new(addr.clone())?;
483 let bootstrap = Bootstrap::Host {
484 addr: addr.clone(),
485 callback_addr: callback.callback_addr(),
486 command: Some(command.clone()),
487 config: Some(hyperactor_config::global::attrs()),
488 exit_on_shutdown: false,
489 };
490
491 let mut cmd = command.new();
492 bootstrap.to_env(&mut cmd);
493 cmd.spawn()?;
494 host_addrs.push(addr);
495 ready.push(callback);
496 }
497 for callback in ready {
498 callback.wait().await?;
499 }
500
501 let host_mesh_ref = HostMeshRef::new(
502 HostMeshId::instance(Label::new("process").unwrap()),
503 extent.into(),
504 host_addrs,
505 )?;
506 Ok(HostMesh::take(host_mesh_ref))
507 }
508 pub fn take(mesh: HostMeshRef) -> Self {
513 let id = mesh.id.clone();
514 let extent = mesh.region().extent().clone();
515
516 let result = Self {
517 current_ref: mesh,
518 id,
519 extent,
520 shutdown_on_drop: true,
521 };
522 result.notify_created();
523 result
524 }
525
526 pub async fn attach(
540 cx: &impl context::Actor,
541 id: HostMeshId,
542 addresses: Vec<ChannelAddr>,
543 ) -> crate::Result<Self> {
544 let mesh_ref = HostMeshRef::from_hosts(id, addresses);
545 let config = hyperactor_config::global::propagatable_attrs();
546 mesh_ref.push_config(cx, config).await?;
547 Ok(Self::take(mesh_ref))
548 }
549
550 pub async fn attach_with_service_procs(
552 cx: &impl context::Actor,
553 id: HostMeshId,
554 service_procs: Vec<ProcAddr>,
555 ) -> crate::Result<Self> {
556 let agents = service_procs.into_iter().map(host_agent_ref).collect();
557 let mesh_ref = HostMeshRef::from_host_agents(id, agents)?;
558 let config = hyperactor_config::global::propagatable_attrs();
559 mesh_ref.push_config(cx, config).await?;
560 Ok(Self::take(mesh_ref))
561 }
562
563 #[hyperactor::instrument(fields(host_mesh=self.id.to_string()))]
572 pub async fn shutdown(&mut self, cx: &impl hyperactor::context::Actor) -> anyhow::Result<()> {
573 let t0 = std::time::Instant::now();
574 tracing::info!(name = "HostMeshStatus", status = "Shutdown::Attempt");
575
576 if let Err(e) = self.current_ref.cast_drain(cx, None).await {
579 tracing::warn!(
580 name = "HostMeshStatus",
581 status = "Shutdown::Drain::Failed",
582 drain_ms = t0.elapsed().as_millis(),
583 error = %e,
584 "failed to cast DrainHost barrier"
585 );
586 }
587 let drain_ms = t0.elapsed().as_millis();
588
589 let t1 = std::time::Instant::now();
591 let shutdown_result = self.current_ref.cast_shutdown(cx).await;
592 let shutdown_ack_ms = t1.elapsed().as_millis();
593 let total_ms = t0.elapsed().as_millis();
594 if let Err(e) = shutdown_result {
595 tracing::warn!(
596 name = "HostMeshStatus",
597 status = "Shutdown::Ack::Failed",
598 drain_ms,
599 shutdown_ack_ms,
600 total_ms,
601 error = %e,
602 "failed waiting for ShutdownHost acknowledgment barrier"
603 );
604 } else {
605 tracing::info!(
606 name = "HostMeshStatus",
607 status = "Shutdown::Success",
608 drain_ms,
609 shutdown_ack_ms,
610 total_ms
611 );
612 }
613
614 self.shutdown_on_drop = false;
615 Ok(())
616 }
617
618 pub fn shutdown_guard(self) -> HostMeshShutdownGuard {
621 HostMeshShutdownGuard(self)
622 }
623
624 #[hyperactor::instrument(fields(host_mesh=self.id.to_string()))]
631 pub async fn stop(&mut self, cx: &impl hyperactor::context::Actor) -> anyhow::Result<()> {
632 let t0 = std::time::Instant::now();
633 tracing::info!(name = "HostMeshStatus", status = "Stop::Attempt");
634
635 let result = self.current_ref.cast_drain(cx, Some(self.id.clone())).await;
636 let total_ms = t0.elapsed().as_millis();
637 match result {
638 Ok(()) => {
639 tracing::info!(name = "HostMeshStatus", status = "Stop::Success", total_ms,);
640 }
641 Err(e) => tracing::warn!(
642 name = "HostMeshStatus",
643 status = "Stop::Drain::Failed",
644 total_ms,
645 error = %e,
646 "failed waiting for DrainHost acknowledgment barrier"
647 ),
648 }
649
650 self.shutdown_on_drop = false;
653
654 Ok(())
655 }
656}
657
658impl HostMesh {
659 pub fn set_bootstrap(&mut self, cmd: BootstrapCommand) {
664 self.current_ref = self.current_ref.clone().with_bootstrap(cmd);
665 }
666}
667
668impl Deref for HostMesh {
669 type Target = HostMeshRef;
670
671 fn deref(&self) -> &Self::Target {
672 &self.current_ref
673 }
674}
675
676impl AsRef<HostMeshRef> for HostMesh {
677 fn as_ref(&self) -> &HostMeshRef {
678 self
679 }
680}
681
682impl AsRef<HostMeshRef> for HostMeshRef {
683 fn as_ref(&self) -> &HostMeshRef {
684 self
685 }
686}
687
688pub struct HostMeshShutdownGuard(pub HostMesh);
690
691impl Deref for HostMeshShutdownGuard {
692 type Target = HostMesh;
693
694 fn deref(&self) -> &HostMesh {
695 &self.0
696 }
697}
698
699impl DerefMut for HostMeshShutdownGuard {
700 fn deref_mut(&mut self) -> &mut HostMesh {
701 &mut self.0
702 }
703}
704
705impl Drop for HostMeshShutdownGuard {
706 fn drop(&mut self) {
725 if !self.0.shutdown_on_drop {
726 tracing::debug!(
727 name = "HostMeshStatus",
728 host_mesh = %self.0.id,
729 status = "DropCleanup::Skipped",
730 "hostmesh drop-cleanup skipped after explicit stop/shutdown"
731 );
732 return;
733 }
734
735 tracing::info!(
736 name = "HostMeshStatus",
737 host_mesh = %self.0.id,
738 status = "Dropping",
739 );
740 let current_ref = self.0.current_ref.clone();
741 let host_count = current_ref.region().num_ranks();
742
743 if host_count == 0 {
745 tracing::debug!(
746 host_mesh = %self.0.id,
747 "HostMesh drop cleanup skipped because no owned hosts remain"
748 );
749 } else if let Ok(handle) = tokio::runtime::Handle::try_current() {
750 let mesh_id = self.0.id.clone();
751 let span = tracing::info_span!(
752 "hostmesh_drop_cleanup",
753 host_mesh = %mesh_id,
754 hosts = host_count,
755 );
756
757 handle.spawn(
758 async move {
759 match hyperactor::Proc::direct(
762 ChannelTransport::Unix.any(),
763 "hostmesh-drop".to_string(),
764 ) {
765 Err(e) => {
766 tracing::warn!(
767 error = %e,
768 "failed to construct ephemeral Proc for drop-cleanup; \
769 relying on PDEATHSIG/manager Drop"
770 );
771 }
772 Ok(proc) => {
773 let client = proc.client("drop");
774 if let Err(e) = current_ref.cast_shutdown(&client).await {
775 tracing::warn!(
776 error = %e,
777 "drop-cleanup: failed to cast ShutdownHost"
778 );
779 } else {
780 tracing::info!(
781 hosts = host_count,
782 "hostmesh drop-cleanup shutdown barrier complete"
783 );
784 }
785 }
786 }
787 }
788 .instrument(span),
789 );
790 } else {
791 tracing::warn!(
794 host_mesh = %self.0.id,
795 hosts = host_count,
796 "HostMesh dropped without a Tokio runtime; skipping \
797 best-effort shutdown. This indicates that .shutdown() \
798 on this mesh has not been called before program exit \
799 (perhaps due to a missing call to \
800 'monarch.actor.shutdown_context()'?) This in turn can \
801 lead to backtrace output due to folly SIGTERM \
802 handlers."
803 );
804 }
805
806 tracing::info!(
807 name = "HostMeshStatus",
808 host_mesh = %self.0.id,
809 status = "Dropped",
810 );
811 }
812}
813
814pub(crate) fn mesh_to_rankedvalues_with_default<T, F>(
823 mesh: &ValueMesh<T>,
824 default: T,
825 is_sentinel: F,
826 len: usize,
827) -> RankedValues<T>
828where
829 T: Eq + Clone + 'static,
830 F: Fn(&T) -> bool,
831{
832 let mut out = RankedValues::from((0..len, default));
833 for (i, s) in mesh.values().enumerate() {
834 if !is_sentinel(&s) {
835 out.merge_from(RankedValues::from((i..i + 1, s)));
836 }
837 }
838 out
839}
840
841#[derive(Debug, Clone, Named, Serialize, Deserialize, PartialEq, Eq, Hash)]
860pub struct HostMeshRef {
861 id: HostMeshId,
862 host_agent_mesh: ActorMeshRef<HostAgent>,
863 #[serde(default)]
868 pub bootstrap_command: Option<BootstrapCommand>,
869}
870
871pub type PerRankBootstrapFn = dyn Fn(view::Point) -> anyhow::Result<BootstrapCommand> + Send + Sync;
877
878wirevalue::register_type!(HostMeshRef);
879
880impl HostMeshRef {
881 #[allow(clippy::result_large_err)]
884 fn new(id: HostMeshId, region: Region, host_addrs: Vec<ChannelAddr>) -> crate::Result<Self> {
885 if region.num_ranks() != host_addrs.len() {
886 return Err(crate::Error::InvalidRankCardinality {
887 expected: region.num_ranks(),
888 actual: host_addrs.len(),
889 });
890 }
891 let host_agent_mesh = Self::host_agent_mesh_ref_from_addrs(®ion, host_addrs)?;
892 Ok(Self {
893 id,
894 host_agent_mesh,
895 bootstrap_command: None,
896 })
897 }
898
899 pub fn from_hosts(id: HostMeshId, host_addrs: Vec<ChannelAddr>) -> Self {
902 let region = extent!(hosts = host_addrs.len()).into();
903 let host_agent_mesh = Self::host_agent_mesh_ref_from_addrs(®ion, host_addrs)
904 .expect("host rank cardinality must match generated region");
905 Self {
906 id,
907 host_agent_mesh,
908 bootstrap_command: None,
909 }
910 }
911
912 pub fn from_host_agents(
914 id: HostMeshId,
915 agents: Vec<ActorRef<HostAgent>>,
916 ) -> crate::Result<Self> {
917 let region = extent!(hosts = agents.len()).into();
918 let host_agent_mesh = Self::host_agent_mesh_ref_from_agents(®ion, agents)?;
919 Ok(Self {
920 id,
921 host_agent_mesh,
922 bootstrap_command: None,
923 })
924 }
925
926 pub fn from_host_agent(id: HostMeshId, agent: ActorRef<HostAgent>) -> crate::Result<Self> {
928 Self::from_host_agents(id, vec![agent])
929 }
930
931 pub fn with_bootstrap(self, cmd: BootstrapCommand) -> Self {
934 Self {
935 bootstrap_command: Some(cmd),
936 ..self
937 }
938 }
939
940 fn host_agent_mesh_ref_from_addrs(
941 region: &Region,
942 host_addrs: Vec<ChannelAddr>,
943 ) -> crate::Result<ActorMeshRef<HostAgent>> {
944 let agents = host_addrs.into_iter().map(legacy_host_agent_ref).collect();
945 Self::host_agent_mesh_ref_from_agents(region, agents)
946 }
947
948 fn host_agent_mesh_ref_from_agents(
949 region: &Region,
950 agents: Vec<ActorRef<HostAgent>>,
951 ) -> crate::Result<ActorMeshRef<HostAgent>> {
952 let members = Arc::new(
953 agents
954 .into_iter()
955 .map(|agent| agent.actor_addr().clone())
956 .collect_mesh::<ValueMesh<_>>(region.clone())
957 .map_err(|error| crate::Error::ConfigurationError(error.into()))?,
958 );
959
960 Ok(ActorMeshRef::new_managed(
961 ActorMeshId::singleton(Label::strip(host_agent::HOST_MESH_AGENT_ACTOR_NAME)),
962 None,
964 region.clone(),
965 None,
966 members,
967 ))
968 }
969
970 async fn cast_drain(
971 &self,
972 cx: &impl context::Actor,
973 host_mesh_id: Option<HostMeshId>,
974 ) -> anyhow::Result<()> {
975 let region = self.region().clone();
976 let num_hosts = region.num_ranks();
977 if num_hosts == 0 {
978 return Ok(());
979 }
980
981 let (reply, rx) = cx.mailbox().open_idle_flush_accum_port(
985 crate::StatusMesh::from_single(region.clone(), Status::NotExist),
986 IdleFlushReducerOpts {
987 idle_timeout: Duration::from_millis(50),
988 abandon_timeout: Duration::from_secs(30),
989 expected_updates_per_destination: NonZeroUsize::MIN,
990 },
991 );
992 let mut reply = reply.bind();
993 reply.return_undeliverable(false);
994
995 let terminate_timeout =
996 hyperactor_config::global::get(crate::bootstrap::MESH_TERMINATE_TIMEOUT);
997
998 self.host_agent_mesh.cast(
999 cx,
1000 host_agent::DrainHost {
1001 timeout: terminate_timeout,
1002 max_in_flight: hyperactor_config::global::get(
1003 crate::bootstrap::MESH_TERMINATE_CONCURRENCY,
1004 )
1005 .clamp(1, 256),
1006 host_mesh_id,
1007 rank: Default::default(),
1008 reply,
1009 },
1010 )?;
1011
1012 let barrier_timeout = terminate_timeout.saturating_add(std::time::Duration::from_secs(30));
1015
1016 match GetRankStatus::wait(rx, num_hosts, barrier_timeout, region).await {
1017 Ok(_) => Ok(()),
1018 Err(partial) => {
1019 let missing: Vec<usize> = partial
1020 .values()
1021 .enumerate()
1022 .filter(|(_, status)| status.is_not_exist())
1023 .map(|(rank, _)| rank)
1024 .collect();
1025 anyhow::bail!(
1026 "DrainHost barrier timed out after {:?}; {} of {} hosts did not acknowledge (host ranks {:?})",
1027 barrier_timeout,
1028 missing.len(),
1029 num_hosts,
1030 missing,
1031 )
1032 }
1033 }
1034 }
1035
1036 async fn cast_shutdown(&self, cx: &impl context::Actor) -> anyhow::Result<()> {
1037 let num_hosts = self.region().num_ranks();
1038 if num_hosts == 0 {
1039 return Ok(());
1040 }
1041
1042 let (ack, mut rx) = cx.mailbox().open_port::<usize>();
1048 let mut ack = ack.bind().unsplit();
1053 ack.return_undeliverable(false);
1054
1055 let terminate_timeout =
1056 hyperactor_config::global::get(crate::bootstrap::MESH_TERMINATE_TIMEOUT);
1057
1058 self.host_agent_mesh.cast(
1059 cx,
1060 host_agent::ShutdownHost {
1061 timeout: terminate_timeout,
1062 max_in_flight: hyperactor_config::global::get(
1063 crate::bootstrap::MESH_TERMINATE_CONCURRENCY,
1064 )
1065 .clamp(1, 256),
1066 rank: Default::default(),
1067 ack,
1068 },
1069 )?;
1070
1071 let barrier_timeout = terminate_timeout.saturating_add(std::time::Duration::from_secs(30));
1074
1075 let mut acked = std::collections::HashSet::new();
1076
1077 while acked.len() < num_hosts {
1078 match tokio::time::timeout(barrier_timeout, rx.recv()).await {
1079 Ok(Ok(rank)) => {
1080 acked.insert(rank);
1081 }
1082 Ok(Err(err)) => return Err(anyhow::Error::from(err)),
1083 Err(_) => {
1084 let missing: Vec<usize> =
1085 (0..num_hosts).filter(|r| !acked.contains(r)).collect();
1086
1087 anyhow::bail!(
1088 "ShutdownHost barrier timed out after {:?}; {} of {} hosts did not acknowledge shutdown (host ranks {:?})",
1089 barrier_timeout,
1090 missing.len(),
1091 num_hosts,
1092 missing,
1093 );
1094 }
1095 }
1096 }
1097 Ok(())
1098 }
1099
1100 pub(crate) fn cast_stream_state(
1104 &self,
1105 cx: &impl context::Actor,
1106 id: ResourceId,
1107 subscriber: hyperactor::PortRef<resource::RankedState<ProcState>>,
1108 ) -> anyhow::Result<()> {
1109 Ok(self.host_agent_mesh.cast(
1110 cx,
1111 resource::StreamState::<ProcState> {
1112 id,
1113 subscriber_rank: resource::Rank::default(),
1114 subscriber,
1115 },
1116 )?)
1117 }
1118
1119 pub(crate) fn forward_wait_rank_status(
1120 &self,
1121 cx: &impl context::Actor,
1122 procs: impl IntoIterator<Item = ProcAddr>,
1123 region: Region,
1124 message: resource::WaitRankStatus,
1125 ) -> anyhow::Result<()> {
1126 for (view_rank, (create_rank, proc_id)) in region.slice().iter().zip(procs).enumerate() {
1127 let mut message = message.clone();
1128 message.id = ResourceId::new(proc_id.uid().clone(), proc_id.label().cloned());
1129 message.rank = resource::Rank::new(view_rank);
1130 self.host_agent_for_proc_rank(®ion, create_rank)
1131 .post(cx, message);
1132 }
1133 Ok(())
1134 }
1135
1136 pub(crate) fn host_entries(&self) -> Vec<(String, ActorRef<HostAgent>)> {
1140 self.host_agent_mesh
1141 .values()
1142 .map(|agent| (agent.actor_addr().addr().to_string(), agent.clone()))
1143 .collect()
1144 }
1145
1146 pub(crate) async fn push_config(
1158 &self,
1159 cx: &impl context::Actor,
1160 attrs: hyperactor_config::attrs::Attrs,
1161 ) -> Result<(), ConfigPushError> {
1162 let timeout = hyperactor_config::global::get(crate::config::MESH_ATTACH_CONFIG_TIMEOUT);
1163 let host_addrs = self.host_addrs();
1164 let num_hosts = host_addrs.len();
1165
1166 if num_hosts == 0 {
1167 tracing::info!(success = 0, "push_config complete");
1168 return Ok(());
1169 }
1170
1171 fn failures_for_host_addrs(
1172 host_addrs: &[ChannelAddr],
1173 mut make_failure: impl FnMut() -> ConfigPushFailure,
1174 ) -> ConfigPushError {
1175 ConfigPushError {
1176 failures: host_addrs
1177 .iter()
1178 .cloned()
1179 .map(|host_addr| (host_addr, make_failure()))
1180 .collect(),
1181 }
1182 }
1183
1184 let region = self.region().clone();
1185
1186 let (reply, rx) = cx.mailbox().open_idle_flush_accum_port(
1190 crate::StatusMesh::from_single(region.clone(), Status::NotExist),
1191 IdleFlushReducerOpts {
1192 idle_timeout: Duration::from_millis(50),
1193 abandon_timeout: Duration::from_secs(30),
1194 expected_updates_per_destination: NonZeroUsize::MIN,
1195 },
1196 );
1197 let mut reply = reply.bind();
1198 reply.return_undeliverable(false);
1199
1200 if let Err(err) = self.host_agent_mesh.cast(
1207 cx,
1208 host_agent::SetClientConfig {
1209 attrs,
1210 rank: Default::default(),
1211 reply,
1212 },
1213 ) {
1214 let error = err.to_string();
1215
1216 tracing::warn!(error = %error, "config push cast failed");
1217
1218 return Err(failures_for_host_addrs(&host_addrs, || {
1222 ConfigPushFailure::CastFailed(error.clone())
1223 }));
1224 }
1225
1226 match GetRankStatus::wait(rx, num_hosts, timeout, region).await {
1227 Ok(_) => {
1228 tracing::info!(success = num_hosts, "push_config complete");
1229 Ok(())
1230 }
1231 Err(partial) => {
1232 let failures: Vec<(ChannelAddr, ConfigPushFailure)> = partial
1236 .values()
1237 .enumerate()
1238 .filter(|(_, status)| status.is_not_exist())
1239 .map(|(rank, _)| (host_addrs[rank].clone(), ConfigPushFailure::ReplyTimedOut))
1240 .collect();
1241
1242 tracing::info!(
1243 success = num_hosts - failures.len(),
1244 failed = failures.len(),
1245 "push_config complete with failures"
1246 );
1247
1248 Err(ConfigPushError { failures })
1249 }
1250 }
1251 }
1252
1253 #[allow(clippy::result_large_err)]
1272 pub async fn spawn<C: context::Actor>(
1273 &self,
1274 cx: &C,
1275 name: &str,
1276 per_host: Extent,
1277 proc_bind: Option<Vec<ProcBind>>,
1278 per_rank_bootstrap: Option<Box<PerRankBootstrapFn>>,
1279 ) -> crate::Result<ProcMesh>
1280 where
1281 C::A: Handler<MeshFailure>,
1282 {
1283 self.spawn_inner(
1284 cx,
1285 ProcMeshId::instance(Label::strip(name)),
1286 per_host,
1287 proc_bind,
1288 per_rank_bootstrap,
1289 )
1290 .await
1291 }
1292
1293 #[hyperactor::instrument(fields(host_mesh=self.id.to_string(), proc_mesh=proc_mesh_id.to_string()))]
1294 async fn spawn_inner<C: context::Actor>(
1295 &self,
1296 cx: &C,
1297 proc_mesh_id: ProcMeshId,
1298 per_host: Extent,
1299 proc_bind: Option<Vec<ProcBind>>,
1300 per_rank_bootstrap: Option<Box<PerRankBootstrapFn>>,
1301 ) -> crate::Result<ProcMesh>
1302 where
1303 C::A: Handler<MeshFailure>,
1304 {
1305 tracing::info!(name = "HostMeshStatus", status = "ProcMesh::Spawn::Attempt");
1306 tracing::info!(name = "ProcMeshStatus", status = "Spawn::Attempt",);
1307 let result = self
1308 .spawn_inner_inner(cx, proc_mesh_id, per_host, proc_bind, per_rank_bootstrap)
1309 .await;
1310 match &result {
1311 Ok(_) => {
1312 tracing::info!(name = "HostMeshStatus", status = "ProcMesh::Spawn::Success");
1313 tracing::info!(name = "ProcMeshStatus", status = "Spawn::Success");
1314 }
1315 Err(error) => {
1316 tracing::error!(name = "HostMeshStatus", status = "ProcMesh::Spawn::Failed", %error);
1317 tracing::error!(name = "ProcMeshStatus", status = "Spawn::Failed", %error);
1318 }
1319 }
1320 result
1321 }
1322
1323 async fn spawn_inner_inner<C: context::Actor>(
1324 &self,
1325 cx: &C,
1326 proc_mesh_id: ProcMeshId,
1327 per_host: Extent,
1328 proc_bind: Option<Vec<ProcBind>>,
1329 per_rank_bootstrap: Option<Box<PerRankBootstrapFn>>,
1330 ) -> crate::Result<ProcMesh>
1331 where
1332 C::A: Handler<MeshFailure>,
1333 {
1334 let per_host_labels = per_host.labels().iter().collect::<HashSet<_>>();
1335 let host_labels = self.region().labels().iter().collect::<HashSet<_>>();
1336 if !per_host_labels
1337 .intersection(&host_labels)
1338 .collect::<Vec<_>>()
1339 .is_empty()
1340 {
1341 return Err(crate::Error::ConfigurationError(anyhow::anyhow!(
1342 "per_host dims overlap with existing dims when spawning proc mesh"
1343 )));
1344 }
1345 if let Some(proc_bind) = proc_bind.as_ref()
1346 && proc_bind.len() != per_host.num_ranks()
1347 {
1348 return Err(crate::Error::ConfigurationError(anyhow::anyhow!(
1349 "proc_bind length does not match per_host extent"
1350 )));
1351 }
1352
1353 let extent = self
1354 .region()
1355 .extent()
1356 .concat(&per_host)
1357 .map_err(|err| crate::Error::ConfigurationError(err.into()))?;
1358
1359 let region: Region = extent.clone().into();
1360
1361 tracing::info!(
1362 name = "ProcMeshStatus",
1363 status = "Spawn::Attempt",
1364 %region,
1365 "spawning proc mesh"
1366 );
1367
1368 let mut procs = Vec::new();
1369 let num_ranks = region.num_ranks();
1370 let expected_updates_per_destination = NonZeroUsize::new(per_host.num_ranks())
1371 .expect("a per-host region contains at least one rank");
1372 let (port, rx) = cx.mailbox().open_idle_flush_accum_port(
1375 crate::StatusMesh::from_single(region.clone(), Status::NotExist),
1376 IdleFlushReducerOpts {
1377 idle_timeout: Duration::from_millis(50),
1378 abandon_timeout: Duration::from_secs(30),
1379 expected_updates_per_destination,
1380 },
1381 );
1382
1383 let mut proc_names = Vec::new();
1386 let client_config_override = hyperactor_config::global::propagatable_attrs();
1387 for (host_rank, agent) in self.host_agent_mesh.values().enumerate() {
1388 for (_, create_rank, proc_name) in
1389 host_agent::proc_slots(&proc_mesh_id, host_rank, per_host.num_ranks())
1390 {
1391 proc_names.push(proc_name.clone());
1392 let proc_id = named_proc_on_host(&agent, &proc_name);
1393 let proc_agent =
1394 ActorRef::attest(proc_id.actor_addr(crate::proc_agent::PROC_AGENT_ACTOR_NAME));
1395 tracing::info!(
1396 name = "ProcMeshStatus",
1397 status = "Spawn::CreatingProc",
1398 %proc_id,
1399 rank = create_rank,
1400 );
1401 procs.push(crate::proc_mesh::ProcRef::new(
1402 proc_id,
1403 create_rank,
1404 proc_agent,
1406 ));
1407 }
1408 }
1409
1410 let mut reply_port = port.bind();
1411
1412 reply_port.return_undeliverable(false);
1413
1414 let total_procs = self.region().num_ranks() * per_host.num_ranks();
1415
1416 let bootstrap_commands = match per_rank_bootstrap.as_ref() {
1417 Some(per_rank_bootstrap) => Some(
1418 (0..total_procs)
1419 .map(|create_rank| {
1420 per_rank_bootstrap(
1421 extent
1422 .point_of_rank(create_rank)
1423 .expect("rank in combined extent"),
1424 )
1425 .map(Some)
1426 .map_err(crate::Error::ConfigurationError)
1427 })
1428 .collect::<crate::Result<Vec<_>>>()?,
1429 ),
1430 None => None,
1431 };
1432
1433 self.host_agent_mesh.cast(
1436 cx,
1437 host_agent::SpawnProcs {
1438 rank: resource::Rank::default(),
1439 proc_mesh_id: proc_mesh_id.clone(),
1440 num_per_host: per_host.num_ranks(),
1441 client_config_override,
1442 host_mesh_id: Some(self.id.clone()),
1443 default_bootstrap_command: self.bootstrap_command.clone(),
1444 proc_bind,
1445 bootstrap_commands,
1446 },
1447 )?;
1448 self.host_agent_mesh.cast(
1449 cx,
1450 host_agent::WaitProcs {
1451 rank: resource::Rank::default(),
1452 proc_mesh_id: proc_mesh_id.clone(),
1453 num_per_host: per_host.num_ranks(),
1454 status_reply: reply_port,
1455 },
1456 )?;
1457
1458 let start_time = tokio::time::Instant::now();
1459
1460 match GetRankStatus::wait(
1463 rx,
1464 num_ranks,
1465 hyperactor_config::global::get(PROC_SPAWN_MAX_IDLE),
1466 region.clone(), )
1468 .await
1469 {
1470 Ok(statuses) => {
1471 if let Some((rank, status)) = statuses
1474 .values()
1475 .enumerate()
1476 .find(|(_, s)| s.is_terminating())
1477 {
1478 let proc_name = &proc_names[rank];
1479 let host_rank = rank / per_host.num_ranks();
1480 let mesh_agent = self
1481 .host_agent_mesh
1482 .get(host_rank)
1483 .expect("host rank must be in host agent mesh")
1484 .clone();
1485 let (reply_tx, mut reply_rx) = cx.mailbox().open_port();
1486 let mut reply_tx = reply_tx.bind();
1487 reply_tx.return_undeliverable(false);
1490 mesh_agent.post(
1491 cx,
1492 resource::GetState {
1493 id: proc_name.clone(),
1494 reply: reply_tx,
1495 },
1496 );
1497 let state = match tokio::time::timeout(
1498 hyperactor_config::global::get(PROC_SPAWN_MAX_IDLE),
1499 reply_rx.recv(),
1500 )
1501 .await
1502 {
1503 Ok(Ok(state)) => state,
1504 _ => resource::State {
1505 id: proc_name.clone(),
1506 status,
1507 state: None,
1508 generation: 0,
1509 timestamp: std::time::SystemTime::now(),
1510 },
1511 };
1512
1513 tracing::error!(
1514 name = "ProcMeshStatus",
1515 status = "Spawn::GetRankStatus",
1516 rank = host_rank,
1517 "rank {} is terminating with state: {}",
1518 host_rank,
1519 state
1520 );
1521
1522 return Err(crate::Error::ProcCreationError {
1523 state: Box::new(state),
1524 host_rank,
1525 mesh_agent,
1526 });
1527 }
1528 }
1529 Err(complete) => {
1530 tracing::error!(
1531 name = "ProcMeshStatus",
1532 status = "Spawn::GetRankStatus",
1533 "timeout after {:?} when waiting for procs being created",
1534 hyperactor_config::global::get(PROC_SPAWN_MAX_IDLE),
1535 );
1536 let legacy = mesh_to_rankedvalues_with_default(
1539 &complete,
1540 Status::Timeout(start_time.elapsed()),
1541 Status::is_not_exist,
1542 num_ranks,
1543 );
1544 return Err(crate::Error::ProcSpawnError { statuses: legacy });
1545 }
1546 }
1547
1548 let mut mesh = ProcMesh::create(proc_mesh_id, extent, self.clone(), procs);
1549 if let Ok(ref mut mesh) = mesh {
1550 let mesh_ref: ProcMeshRef = (**mesh).clone();
1554 let region = ndslice::view::Ranked::region(&mesh_ref).clone();
1555 let initial_statuses: crate::ValueMesh<resource::Status> =
1556 std::iter::repeat_n(resource::Status::Running, region.num_ranks())
1557 .collect_mesh::<crate::ValueMesh<_>>(region)?;
1558 let controller = ProcMeshController::new(mesh_ref, None, None, initial_statuses);
1559 let controller_name = format!("{}_{}", PROC_MESH_CONTROLLER_NAME, mesh.id());
1562 let controller_handle = cx.spawn_with_label(&controller_name, controller);
1563 let controller_ref: ActorRef<ProcMeshController> = controller_handle.bind();
1568 mesh.set_controller(Some(controller_ref));
1569 }
1570 mesh
1571 }
1572
1573 pub fn id(&self) -> &HostMeshId {
1575 &self.id
1576 }
1577
1578 pub(crate) fn agent_mesh(&self) -> &ActorMeshRef<HostAgent> {
1583 &self.host_agent_mesh
1584 }
1585
1586 pub fn host_addrs(&self) -> Vec<ChannelAddr> {
1588 self.host_agent_mesh
1589 .values()
1590 .map(|agent| agent.actor_addr().addr().clone())
1591 .collect()
1592 }
1593
1594 fn host_agent_for_proc_rank(
1595 &self,
1596 proc_region: &Region,
1597 create_rank: usize,
1598 ) -> &ActorRef<HostAgent> {
1599 let host_dims = self.region().labels().len();
1600 assert!(
1601 proc_region.labels().starts_with(self.region().labels()),
1602 "proc mesh host dimensions must match its originating host mesh"
1603 );
1604
1605 let host_rank = if host_dims == 0 {
1606 0
1607 } else {
1608 let procs_per_host = proc_region.slice().strides()[host_dims - 1];
1612 create_rank / procs_per_host
1613 };
1614 self.get(host_rank)
1615 .expect("proc rank must map to a host in its originating host mesh")
1616 }
1617
1618 #[hyperactor::instrument(fields(host_mesh=self.id.to_string(), proc_mesh=proc_mesh_id.to_string()))]
1631 pub(crate) async fn stop_proc_mesh(
1632 &self,
1633 cx: &impl hyperactor::context::Actor,
1634 proc_mesh_id: &ProcMeshId,
1635 procs: impl IntoIterator<Item = ProcAddr>,
1636 region: Region,
1637 reason: String,
1638 ) -> crate::Result<crate::StatusMesh> {
1639 let procs = procs.into_iter().collect::<Vec<_>>();
1640 assert_eq!(
1641 procs.len(),
1642 region.num_ranks(),
1643 "proc addresses must match the proc mesh region"
1644 );
1645 let proc_names = procs
1646 .iter()
1647 .map(|proc_id| ResourceId::new(proc_id.uid().clone(), proc_id.label().cloned()))
1648 .collect::<Vec<_>>();
1649 let num_ranks = region.num_ranks();
1650 let (port, rx) = cx.mailbox().open_idle_flush_accum_port(
1651 crate::StatusMesh::from_single(region.clone(), Status::NotExist),
1652 IdleFlushReducerOpts {
1653 idle_timeout: Duration::from_millis(50),
1654 abandon_timeout: Duration::from_secs(30),
1655 expected_updates_per_destination: NonZeroUsize::MIN,
1656 },
1657 );
1658 let mut reply = port.bind();
1659 reply.return_undeliverable(false);
1660 for (view_rank, (create_rank, proc_id)) in region.slice().iter().zip(procs).enumerate() {
1661 let proc_resource_id = ResourceId::new(proc_id.uid().clone(), proc_id.label().cloned());
1662 let host_agent = self.host_agent_for_proc_rank(®ion, create_rank);
1663 host_agent.post(
1664 cx,
1665 resource::Stop {
1666 id: proc_resource_id.clone(),
1667 reason: reason.clone(),
1668 },
1669 );
1670 host_agent
1671 .wait_rank_status(
1672 cx,
1673 proc_resource_id,
1674 resource::Rank::new(view_rank),
1675 Status::Stopped,
1676 reply.clone(),
1677 )
1678 .await
1679 .map_err(|error| crate::Error::CallError(host_agent.actor_addr().clone(), error))?;
1680 }
1681 tracing::info!(
1682 name = "HostMeshStatus",
1683 status = "ProcMesh::Stop::Sent",
1684 "sending Stop to proc mesh for {} procs: {}",
1685 proc_names.len(),
1686 proc_names
1687 .iter()
1688 .map(|n| n.to_string())
1689 .collect::<Vec<_>>()
1690 .join(", ")
1691 );
1692
1693 let start_time = tokio::time::Instant::now();
1694
1695 match GetRankStatus::wait(
1696 rx,
1697 num_ranks,
1698 hyperactor_config::global::get(PROC_STOP_MAX_IDLE),
1699 region.clone(), )
1701 .await
1702 {
1703 Ok(statuses) => {
1704 let all_stopped = statuses.values().all(|s| s.is_terminated());
1705 if !all_stopped {
1706 let legacy = mesh_to_rankedvalues_with_default(
1707 &statuses,
1708 Status::NotExist,
1709 Status::is_not_exist,
1710 num_ranks,
1711 );
1712 tracing::error!(
1713 name = "ProcMeshStatus",
1714 status = "FailedToStop",
1715 "failed to terminate proc mesh: {:?}",
1716 statuses,
1717 );
1718 return Err(crate::Error::ProcMeshStopError { statuses: legacy });
1719 }
1720 tracing::info!(name = "ProcMeshStatus", status = "Stopped");
1721 Ok(statuses)
1722 }
1723 Err(complete) => {
1724 let legacy = mesh_to_rankedvalues_with_default(
1727 &complete,
1728 Status::Timeout(start_time.elapsed()),
1729 Status::is_not_exist,
1730 num_ranks,
1731 );
1732 tracing::error!(
1733 name = "ProcMeshStatus",
1734 status = "StoppingTimeout",
1735 "failed to terminate proc mesh {} before timeout: {:?}",
1736 proc_mesh_id,
1737 legacy,
1738 );
1739 Err(crate::Error::ProcMeshStopError { statuses: legacy })
1740 }
1741 }
1742 }
1743}
1744
1745#[derive(Clone, Debug, PartialEq, Eq, Hash)]
1750pub(crate) enum HostAgentKey {
1751 Id(ActorId),
1753 Legacy(ActorAddr),
1755}
1756
1757impl HostAgentKey {
1758 pub(crate) fn new(addr: &ActorAddr) -> Self {
1760 if addr.proc_id() == &legacy_service_proc_id() {
1761 Self::Legacy(addr.clone())
1762 } else {
1763 Self::Id(addr.id().clone())
1764 }
1765 }
1766}
1767
1768struct HostSet {
1776 seen: HashSet<HostAgentKey>,
1777 entries: Vec<(String, ActorRef<HostAgent>)>,
1778}
1779
1780impl HostSet {
1781 fn new() -> Self {
1782 Self {
1783 seen: HashSet::new(),
1784 entries: Vec::new(),
1785 }
1786 }
1787
1788 fn insert(&mut self, addr: String, agent_ref: ActorRef<HostAgent>) {
1791 if self.seen.insert(HostAgentKey::new(agent_ref.actor_addr())) {
1792 self.entries.push((addr, agent_ref));
1793 }
1794 }
1795
1796 fn extend_from_mesh(&mut self, mesh: &HostMeshRef) {
1798 for (addr, agent) in mesh.host_entries() {
1799 self.insert(addr, agent);
1800 }
1801 }
1802
1803 fn into_vec(self) -> Vec<(String, ActorRef<HostAgent>)> {
1804 self.entries
1805 }
1806}
1807
1808fn aggregate_hosts(
1814 meshes: &[impl AsRef<HostMeshRef>],
1815 client_host_entries: Option<Vec<(String, ActorRef<HostAgent>)>>,
1816) -> Vec<(String, ActorRef<HostAgent>)> {
1817 let mut set = HostSet::new();
1818
1819 for mesh in meshes {
1821 set.extend_from_mesh(mesh.as_ref());
1822 }
1823
1824 if let Some(entries) = client_host_entries {
1826 for (addr, agent_ref) in entries {
1827 set.insert(addr, agent_ref);
1828 }
1829 }
1830
1831 set.into_vec()
1832}
1833
1834pub async fn spawn_admin(
1848 meshes: impl IntoIterator<Item = impl AsRef<HostMeshRef>>,
1849 cx: &impl hyperactor::context::Actor,
1850 admin_addr: Option<std::net::SocketAddr>,
1851 telemetry_url: Option<String>,
1852) -> anyhow::Result<ActorRef<MeshAdminAgent>> {
1853 let meshes: Vec<_> = meshes.into_iter().collect();
1854 anyhow::ensure!(!meshes.is_empty(), "at least one mesh is required (SA-1)");
1855 for (i, mesh) in meshes.iter().enumerate() {
1856 anyhow::ensure!(
1857 mesh.as_ref().region().num_ranks() != 0,
1858 "mesh at index {} has no hosts (SA-2)",
1859 i,
1860 );
1861 }
1862
1863 let client_entries =
1864 crate::global_context::try_this_host().map(|client_host| client_host.host_entries());
1865 let hosts = aggregate_hosts(&meshes, client_entries);
1866
1867 let root_client_id = cx.mailbox().actor_addr().clone();
1868
1869 let local_proc = cx.instance().proc();
1872 let agent_handle = local_proc.spawn_bound_with_uid(
1873 Uid::singleton(Label::new(crate::mesh_admin::MESH_ADMIN_ACTOR_NAME).unwrap()),
1874 crate::mesh_admin::MeshAdminAgent::new(
1875 hosts,
1876 Some(root_client_id),
1877 admin_addr,
1878 telemetry_url,
1879 ),
1880 )?;
1881 let admin_ref = agent_handle.bind();
1882 Ok(admin_ref)
1883}
1884
1885impl view::Ranked for HostMeshRef {
1886 type Item = ActorRef<HostAgent>;
1887
1888 fn region(&self) -> &Region {
1889 self.host_agent_mesh.region()
1890 }
1891
1892 fn get(&self, rank: usize) -> Option<&Self::Item> {
1893 self.host_agent_mesh.get(rank)
1894 }
1895}
1896
1897impl view::RankedSliceable for HostMeshRef {
1898 fn sliced(&self, region: Region) -> Self {
1899 Self {
1900 id: self.id.clone(),
1901 host_agent_mesh: self.host_agent_mesh.sliced(region),
1902 bootstrap_command: self.bootstrap_command.clone(),
1903 }
1904 }
1905}
1906
1907impl std::fmt::Display for HostMeshRef {
1908 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1909 let legacy_service_proc_id = ResourceId::from_name(SERVICE_PROC_NAME).proc_id();
1910 write!(f, "{}:", self.id)?;
1911 for (rank, agent) in self.host_agent_mesh.values().enumerate() {
1912 if rank > 0 {
1913 write!(f, ",")?;
1914 }
1915 let proc_addr = agent.actor_addr().proc_addr();
1916 if proc_addr.id() == &legacy_service_proc_id {
1917 write!(f, "{}", proc_addr.addr())?;
1918 } else {
1919 write!(f, "{}", proc_addr)?;
1920 }
1921 }
1922 write!(f, "@{}", self.region())
1923 }
1924}
1925
1926#[derive(thiserror::Error, Debug)]
1928pub enum HostMeshRefParseError {
1929 #[error(transparent)]
1930 RegionParseError(#[from] RegionParseError),
1931
1932 #[error("invalid host mesh ref: missing region")]
1933 MissingRegion,
1934
1935 #[error("invalid host mesh ref: missing id")]
1936 MissingId,
1937
1938 #[error(transparent)]
1939 InvalidId(#[from] crate::mesh_id::ResourceIdParseError),
1940
1941 #[error(transparent)]
1942 InvalidHostMeshRef(#[from] Box<crate::Error>),
1943
1944 #[error(transparent)]
1945 Other(#[from] anyhow::Error),
1946}
1947
1948impl From<crate::Error> for HostMeshRefParseError {
1949 fn from(err: crate::Error) -> Self {
1950 Self::InvalidHostMeshRef(Box::new(err))
1951 }
1952}
1953
1954fn parse_host_agent_ref(host_ref: &str) -> anyhow::Result<ActorRef<HostAgent>> {
1955 if let Some((proc_id, host_addr)) = host_ref.split_once('@')
1956 && let Ok(proc_id) = ProcId::from_str(proc_id)
1957 {
1958 let location = hyperactor::Location::from_str(host_addr)?;
1959 return Ok(host_agent_ref(ProcAddr::new(proc_id, location)));
1960 }
1961
1962 let host_addr = ChannelAddr::from_str(host_ref)?;
1963 Ok(legacy_host_agent_ref(host_addr))
1964}
1965
1966impl FromStr for HostMeshRef {
1967 type Err = HostMeshRefParseError;
1968
1969 fn from_str(s: &str) -> Result<Self, Self::Err> {
1970 let (id_str, rest) = s.split_once(':').ok_or(HostMeshRefParseError::MissingId)?;
1971
1972 let id = HostMeshId::from_str(id_str)?;
1973
1974 let (host_refs, region) = rest
1975 .rsplit_once('@')
1976 .ok_or(HostMeshRefParseError::MissingRegion)?;
1977 let host_agents = if host_refs.trim().is_empty() {
1978 Vec::new()
1979 } else {
1980 host_refs
1981 .split(',')
1982 .map(str::trim)
1983 .map(parse_host_agent_ref)
1984 .collect::<anyhow::Result<Vec<_>>>()?
1985 };
1986 let region = region.parse()?;
1987 let host_agent_mesh = Self::host_agent_mesh_ref_from_agents(®ion, host_agents)?;
1988 Ok(Self {
1989 id,
1990 host_agent_mesh,
1991 bootstrap_command: None,
1992 })
1993 }
1994}
1995
1996#[cfg(test)]
1997mod tests {
1998 #[cfg(fbcode_build)]
1999 use std::assert_matches;
2000
2001 #[cfg(fbcode_build)]
2002 use hyperactor_config::attrs::Attrs;
2003 use ndslice::ViewExt;
2004 use ndslice::extent;
2005 #[cfg(fbcode_build)]
2006 use timed_test::assert_no_process_leak;
2007 #[cfg(fbcode_build)]
2008 use tokio::process::Command;
2009 #[cfg(fbcode_build)]
2010 use tracing_test::traced_test;
2011
2012 use super::*;
2013 #[cfg(fbcode_build)]
2014 use crate::ActorMesh;
2015 #[cfg(fbcode_build)]
2016 use crate::Bootstrap;
2017 #[cfg(fbcode_build)]
2018 use crate::bootstrap::MESH_TAIL_LOG_LINES;
2019 #[cfg(fbcode_build)]
2020 use crate::resource::Status;
2021 #[cfg(fbcode_build)]
2022 use crate::testactor;
2023 #[cfg(fbcode_build)]
2024 use crate::testactor::GetConfigAttrs;
2025 #[cfg(fbcode_build)]
2026 use crate::testactor::SetConfigAttrs;
2027 use crate::testing;
2028
2029 #[test]
2030 fn test_host_mesh_subset() {
2031 let hosts: HostMeshRef = "test:local:1,local:2,local:3,local:4@replica=2/2,host=2/1"
2032 .parse()
2033 .unwrap();
2034 assert_eq!(
2035 hosts.range("replica", 1).unwrap().to_string(),
2036 "test:local:3,local:4@2+replica=1/2,host=2/1"
2037 );
2038 }
2039
2040 #[test]
2041 fn test_host_mesh_ref_parse_roundtrip() {
2042 let host_mesh_ref = HostMeshRef::new(
2043 HostMeshId::singleton(Label::new("test").unwrap()),
2044 extent!(replica = 2, host = 2).into(),
2045 vec![
2046 "tcp:127.0.0.1:123".parse().unwrap(),
2047 "tcp:127.0.0.1:123".parse().unwrap(),
2048 "tcp:127.0.0.1:123".parse().unwrap(),
2049 "tcp:127.0.0.1:123".parse().unwrap(),
2050 ],
2051 )
2052 .unwrap();
2053
2054 let parsed: HostMeshRef = host_mesh_ref.to_string().parse().unwrap();
2055 assert_eq!(parsed.id().to_string(), host_mesh_ref.id().to_string());
2056 assert_eq!(parsed.region(), host_mesh_ref.region());
2057 assert_eq!(parsed.host_addrs(), host_mesh_ref.host_addrs());
2058 assert_eq!(parsed.bootstrap_command, host_mesh_ref.bootstrap_command);
2059 }
2060
2061 fn free_localhost_addr() -> ChannelAddr {
2067 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
2068 ChannelAddr::Tcp(listener.local_addr().unwrap())
2069 }
2070
2071 #[cfg(fbcode_build)]
2072 async fn execute_extrinsic_allocation(config: &hyperactor_config::global::ConfigLock) {
2073 let _guard = config.override_key(crate::bootstrap::MESH_BOOTSTRAP_ENABLE_PDEATHSIG, false);
2074
2075 let program = crate::testresource::get("monarch/hyperactor_mesh/bootstrap");
2076
2077 let hosts = vec![free_localhost_addr(), free_localhost_addr()];
2078
2079 let mut children = Vec::new();
2080 let mut ready = Vec::new();
2081 for host in hosts.iter() {
2082 let callback = HostBootstrapReady::new(host.clone()).unwrap();
2083 let mut cmd = Command::new(program.clone());
2084 let boot = Bootstrap::Host {
2085 addr: host.clone(),
2086 callback_addr: callback.callback_addr(),
2087 command: None, config: None,
2089 exit_on_shutdown: false,
2090 };
2091 boot.to_env(&mut cmd);
2092 cmd.kill_on_drop(true);
2093 children.push(cmd.spawn().unwrap());
2094 ready.push(callback);
2095 }
2096 for callback in ready {
2097 callback.wait().await.unwrap();
2098 }
2099
2100 let instance = testing::instance();
2101 let host_mesh =
2102 HostMeshRef::from_hosts(HostMeshId::singleton(Label::new("test").unwrap()), hosts);
2103
2104 let proc_mesh = host_mesh
2105 .spawn(&testing::instance(), "test", extent!(gpus = 4), None, None)
2106 .await
2107 .unwrap();
2108
2109 let actor_mesh: ActorMesh<testactor::TestActor> = proc_mesh
2110 .spawn(&testing::instance(), "test", &())
2111 .await
2112 .unwrap();
2113
2114 testactor::assert_mesh_shape(actor_mesh).await;
2115
2116 HostMesh::take(host_mesh)
2117 .shutdown(&instance)
2118 .await
2119 .expect("hosts shutdown");
2120 }
2121
2122 #[tokio::test]
2123 #[cfg(fbcode_build)]
2124 async fn test_extrinsic_allocation() {
2125 let config = hyperactor_config::global::lock();
2126 execute_extrinsic_allocation(&config).await;
2127 }
2128
2129 #[expect(
2134 clippy::await_holding_invalid_type,
2135 reason = "tracing-test's #[traced_test] enters a span whose Entered guard is held across awaits; this is inherent to the macro and harmless in a test"
2136 )]
2137 #[traced_test]
2138 #[tokio::test]
2139 #[cfg(fbcode_build)]
2140 async fn test_shutdown_succeeds() {
2141 let config = hyperactor_config::global::lock();
2142 execute_extrinsic_allocation(&config).await;
2143
2144 assert!(
2145 logs_contain("Shutdown::Success"),
2146 "Shutdown::Success status log not found after shutting down host mesh"
2147 );
2148 }
2149
2150 #[tokio::test]
2151 #[cfg(fbcode_build)]
2152 async fn test_failing_proc_allocation() {
2153 let lock = hyperactor_config::global::lock();
2154 let _guard = lock.override_key(MESH_TAIL_LOG_LINES, 100);
2155
2156 let program = crate::testresource::get("monarch/hyperactor_mesh/bootstrap");
2157
2158 let hosts = vec![free_localhost_addr(), free_localhost_addr()];
2159
2160 let mut children = Vec::new();
2161 let mut ready = Vec::new();
2162 for host in hosts.iter() {
2163 let callback = HostBootstrapReady::new(host.clone()).unwrap();
2164 let mut cmd = Command::new(program.clone());
2165 let boot = Bootstrap::Host {
2166 addr: host.clone(),
2167 callback_addr: callback.callback_addr(),
2168 config: None,
2169 command: Some(BootstrapCommand::from("false")),
2171 exit_on_shutdown: false,
2172 };
2173 boot.to_env(&mut cmd);
2174 cmd.kill_on_drop(true);
2175 children.push(cmd.spawn().unwrap());
2176 ready.push(callback);
2177 }
2178 for callback in ready {
2179 callback.wait().await.unwrap();
2180 }
2181 let host_mesh =
2182 HostMeshRef::from_hosts(HostMeshId::singleton(Label::new("test").unwrap()), hosts);
2183
2184 let instance = testing::instance();
2185
2186 let err = host_mesh
2187 .spawn(&instance, "test", Extent::unity(), None, None)
2188 .await
2189 .unwrap_err();
2190 assert_matches!(
2191 err,
2192 crate::Error::ProcCreationError { state, .. }
2193 if matches!(state.status, resource::Status::Failed(ref msg) if msg.contains("failed to configure process: Ready(Terminal(Stopped { exit_code: 1"))
2194 );
2195 }
2196
2197 #[cfg(fbcode_build)]
2198 #[assert_no_process_leak]
2199 #[tokio::test]
2200 async fn test_halting_proc_allocation() {
2201 let config = hyperactor_config::global::lock();
2202 let _guard1 = config.override_key(PROC_SPAWN_MAX_IDLE, Duration::from_secs(20));
2203
2204 let program = crate::testresource::get("monarch/hyperactor_mesh/bootstrap");
2205
2206 let hosts = vec![free_localhost_addr(), free_localhost_addr()];
2207
2208 let mut children = Vec::new();
2209 let mut ready = Vec::new();
2210
2211 for (index, host) in hosts.iter().enumerate() {
2212 let callback = HostBootstrapReady::new(host.clone()).unwrap();
2213 let mut cmd = Command::new(program.clone());
2214 let command = if index == 0 {
2215 let mut command = BootstrapCommand::from("sleep");
2216 command.args.push("60".to_string());
2217 Some(command)
2218 } else {
2219 None
2220 };
2221 let boot = Bootstrap::Host {
2222 addr: host.clone(),
2223 callback_addr: callback.callback_addr(),
2224 config: None,
2225 command,
2226 exit_on_shutdown: false,
2227 };
2228 boot.to_env(&mut cmd);
2229 cmd.kill_on_drop(true);
2230 children.push(cmd.spawn().unwrap());
2231 ready.push(callback);
2232 }
2233 for callback in ready {
2234 callback.wait().await.unwrap();
2235 }
2236 let host_mesh =
2237 HostMeshRef::from_hosts(HostMeshId::singleton(Label::new("test").unwrap()), hosts);
2238
2239 let instance = testing::instance();
2240
2241 let err = host_mesh
2242 .spawn(&instance, "test", Extent::unity(), None, None)
2243 .await
2244 .unwrap_err();
2245 let statuses = err.into_proc_spawn_error().unwrap();
2246 assert_matches!(
2247 &statuses.materialized_iter(2).cloned().collect::<Vec<_>>()[..],
2248 &[Status::Timeout(_), Status::Running]
2249 );
2250 }
2251
2252 #[cfg(fbcode_build)]
2253 #[assert_no_process_leak]
2254 #[tokio::test]
2255 async fn test_proc_mesh_stop_routes_to_originating_host_agents() {
2256 let instance = testing::instance();
2257 let mut host_mesh = testing::host_mesh(2).await;
2258 let mut proc_mesh = host_mesh
2259 .spawn(instance, "stop-routing", extent!(gpus = 2), None, None)
2260 .await
2261 .expect("spawn proc mesh");
2262
2263 proc_mesh
2264 .stop(instance, "test complete".to_string())
2265 .await
2266 .expect("stop proc mesh");
2267 host_mesh
2268 .shutdown(instance)
2269 .await
2270 .expect("shutdown host mesh");
2271 }
2272
2273 #[tokio::test]
2274 #[cfg(fbcode_build)]
2275 async fn test_client_config_override() {
2276 let config = hyperactor_config::global::lock();
2277 let _guard1 = config.override_key(crate::bootstrap::MESH_BOOTSTRAP_ENABLE_PDEATHSIG, false);
2278 let _guard2 = config.override_key(
2279 hyperactor::config::HOST_SPAWN_READY_TIMEOUT,
2280 Duration::from_mins(2),
2281 );
2282 let _guard3 = config.override_key(
2283 hyperactor::config::MESSAGE_DELIVERY_TIMEOUT,
2284 Duration::from_mins(1),
2285 );
2286 let _guard4 = config.override_key(PROC_SPAWN_MAX_IDLE, Duration::from_mins(2));
2287
2288 unsafe {
2293 std::env::remove_var("HYPERACTOR_HOST_SPAWN_READY_TIMEOUT");
2294 std::env::remove_var("HYPERACTOR_MESSAGE_DELIVERY_TIMEOUT");
2295 }
2296
2297 let instance = testing::instance();
2298
2299 let mut hm = testing::host_mesh(2).await;
2300 let proc_mesh = hm
2301 .spawn(instance, "test", Extent::unity(), None, None)
2302 .await
2303 .unwrap();
2304 let proc_ids = proc_mesh
2305 .proc_ids()
2306 .map(|proc_addr| proc_addr.id().clone())
2307 .collect::<Vec<_>>();
2308 let unique_proc_ids = proc_ids.iter().collect::<std::collections::HashSet<_>>();
2309
2310 assert_eq!(proc_ids.len(), 2);
2311 assert_eq!(unique_proc_ids.len(), proc_ids.len());
2312
2313 let actor_mesh: ActorMesh<testactor::TestActor> =
2314 proc_mesh.spawn(instance, "test", &()).await.unwrap();
2315
2316 let mut attrs_override = Attrs::new();
2317 attrs_override.set(
2318 hyperactor::config::HOST_SPAWN_READY_TIMEOUT,
2319 Duration::from_mins(3),
2320 );
2321 actor_mesh
2322 .cast(
2323 instance,
2324 SetConfigAttrs(
2325 bincode::serde::encode_to_vec(&attrs_override, bincode::config::legacy())
2326 .unwrap(),
2327 ),
2328 )
2329 .unwrap();
2330
2331 let (tx, mut rx) = instance.open_port();
2332 actor_mesh
2333 .cast(instance, GetConfigAttrs(tx.bind()))
2334 .unwrap();
2335 let actual_attrs = rx.recv().await.unwrap();
2336 let actual_attrs =
2337 bincode::serde::decode_from_slice::<Attrs, _>(&actual_attrs, bincode::config::legacy())
2338 .map(|(v, _)| v)
2339 .unwrap();
2340
2341 assert_eq!(
2342 *actual_attrs
2343 .get(hyperactor::config::HOST_SPAWN_READY_TIMEOUT)
2344 .unwrap(),
2345 Duration::from_mins(3)
2346 );
2347 assert_eq!(
2348 *actual_attrs
2349 .get(hyperactor::config::MESSAGE_DELIVERY_TIMEOUT)
2350 .unwrap(),
2351 Duration::from_mins(1)
2352 );
2353
2354 let _ = hm.shutdown(instance).await;
2355 }
2356
2357 #[tokio::test]
2375 async fn test_attach_fails_closed_on_unreachable_host() {
2376 let config = hyperactor_config::global::lock();
2377 let _guard = config.override_key(
2380 crate::config::MESH_ATTACH_CONFIG_TIMEOUT,
2381 Duration::from_millis(500),
2382 );
2383
2384 let instance = testing::instance();
2385
2386 let unreachable = free_localhost_addr();
2391
2392 let id = HostMeshId::instance(Label::new("hm_test").unwrap());
2393 let result = HostMesh::attach(instance, id, vec![unreachable.clone()]).await;
2394
2395 let err = match result {
2397 Ok(_) => panic!("HM-2: attach must fail when a host is unreachable"),
2398 Err(e) => e,
2399 };
2400
2401 let push_err = match err {
2403 crate::Error::ConfigPushFailed(e) => e,
2404 other => panic!("expected ConfigPushFailed, got: {other:?}"),
2405 };
2406 assert_eq!(push_err.failures.len(), 1);
2407 let (failed_host, _failure) = &push_err.failures[0];
2408 assert_eq!(
2409 failed_host, &unreachable,
2410 "HM-4: failure entry must identify the unreachable host"
2411 );
2412 }
2424
2425 #[test]
2426 fn test_host_mesh_ref_canonicalizes_alias_to_dial_addr() {
2427 let dial_to = ChannelAddr::from_zmq_url("tcp://127.0.0.1:26600").unwrap();
2428 let alias = ChannelAddr::from_zmq_url("tcp://127.0.0.1:26600@tcp://0.0.0.0:26600").unwrap();
2429
2430 let mesh = HostMeshRef::from_hosts(
2431 HostMeshId::singleton(Label::new("alias").unwrap()),
2432 vec![alias],
2433 );
2434
2435 assert_eq!(mesh.host_addrs(), vec![dial_to.clone()]);
2436 assert_eq!(
2437 ndslice::view::Ranked::get(&mesh, 0)
2438 .expect("host rank should exist")
2439 .actor_addr()
2440 .proc_addr()
2441 .addr(),
2442 &dial_to
2443 );
2444 }
2445
2446 #[test]
2447 fn test_host_mesh_ref_serialization_preserves_service_proc_identity() {
2448 let dial_addr = ChannelAddr::from_zmq_url("tcp://127.0.0.1:26600").unwrap();
2449 let service_proc_id: ProcId = "service<2>".parse().unwrap();
2450 let agent = host_agent_ref(ProcAddr::new(
2451 service_proc_id.clone(),
2452 dial_addr.clone().into(),
2453 ));
2454 let mesh = HostMeshRef::from_host_agents(
2455 HostMeshId::singleton(Label::new("explicit-service").unwrap()),
2456 vec![agent],
2457 )
2458 .unwrap();
2459
2460 let serialized = serde_json::to_string(&mesh).unwrap();
2461 let round_trip: HostMeshRef = serde_json::from_str(&serialized).unwrap();
2462 let round_trip_agent = ndslice::view::Ranked::get(&round_trip, 0).unwrap();
2463
2464 assert_eq!(
2465 round_trip_agent.actor_addr().proc_addr().id(),
2466 &service_proc_id
2467 );
2468 assert_eq!(round_trip_agent.actor_addr().proc_addr().addr(), &dial_addr);
2469 }
2470
2471 #[test]
2472 fn test_host_mesh_ref_from_host_agent_preserves_identity_and_location() {
2473 let dial_addr = ChannelAddr::from_zmq_url("tcp://127.0.0.1:26600").unwrap();
2474 let service_proc_id: ProcId = "service<2>".parse().unwrap();
2475 let via_uid = ProcId::instance(Label::strip("client-gateway"))
2476 .uid()
2477 .clone();
2478 let service_location = hyperactor::Location::from(dial_addr.clone()).with_via(via_uid);
2479 let agent = host_agent_ref(ProcAddr::new(
2480 service_proc_id.clone(),
2481 service_location.clone(),
2482 ));
2483
2484 let mesh = HostMeshRef::from_host_agent(
2485 HostMeshId::singleton(Label::new("explicit-service").unwrap()),
2486 agent,
2487 )
2488 .unwrap();
2489 let round_trip_agent = ndslice::view::Ranked::get(&mesh, 0).unwrap();
2490
2491 assert_eq!(
2492 round_trip_agent.actor_addr().proc_addr().id(),
2493 &service_proc_id
2494 );
2495 assert_eq!(
2496 round_trip_agent.actor_addr().proc_addr().location(),
2497 &service_location
2498 );
2499 }
2500
2501 #[test]
2502 fn test_proc_rank_maps_to_originating_host_agent() {
2503 let dial_addr = ChannelAddr::from_zmq_url("tcp://127.0.0.1:26600").unwrap();
2504 let service_proc_ids = [
2505 "service<2>".parse::<ProcId>().unwrap(),
2506 "service<3>".parse::<ProcId>().unwrap(),
2507 ];
2508 let agents = service_proc_ids
2509 .iter()
2510 .cloned()
2511 .map(|proc_id| host_agent_ref(ProcAddr::new(proc_id, dial_addr.clone().into())))
2512 .collect();
2513 let hosts = HostMeshRef::from_host_agents(
2514 HostMeshId::singleton(Label::new("hosts").unwrap()),
2515 agents,
2516 )
2517 .unwrap();
2518 let proc_region: Region = extent!(hosts = 2, gpus = 2).into();
2519 let sliced = proc_region.range("gpus", 1..2).unwrap();
2520
2521 let owners = sliced
2522 .slice()
2523 .iter()
2524 .map(|create_rank| {
2525 hosts
2526 .host_agent_for_proc_rank(&sliced, create_rank)
2527 .actor_addr()
2528 .proc_addr()
2529 .id()
2530 .clone()
2531 })
2532 .collect::<Vec<_>>();
2533
2534 assert_eq!(owners, service_proc_ids);
2535 }
2536
2537 #[test]
2538 fn test_host_mesh_ref_string_preserves_service_proc_identity() {
2539 let dial_addr = ChannelAddr::from_zmq_url("tcp://127.0.0.1:26600").unwrap();
2540 let service_proc_ids = vec![
2541 "service<2>".parse::<ProcId>().unwrap(),
2542 "service<3>".parse::<ProcId>().unwrap(),
2543 ];
2544 let agents = service_proc_ids
2545 .iter()
2546 .cloned()
2547 .map(|proc_id| host_agent_ref(ProcAddr::new(proc_id, dial_addr.clone().into())))
2548 .collect();
2549 let mesh = HostMeshRef::from_host_agents(
2550 HostMeshId::singleton(Label::new("explicit-services").unwrap()),
2551 agents,
2552 )
2553 .unwrap();
2554
2555 let round_trip: HostMeshRef = mesh.to_string().parse().unwrap();
2556 let round_trip_proc_ids = round_trip
2557 .host_agent_mesh
2558 .values()
2559 .map(|agent| agent.actor_addr().proc_addr().id().clone())
2560 .collect::<Vec<_>>();
2561
2562 assert_eq!(round_trip_proc_ids, service_proc_ids);
2563 assert_eq!(round_trip.host_addrs(), vec![dial_addr.clone(), dial_addr]);
2564 }
2565
2566 #[test]
2567 fn test_host_mesh_ref_string_parses_unix_abstract_addresses() {
2568 let dial_addr: ChannelAddr = "unix:@host-agent".parse().unwrap();
2569 let service_proc_id: ProcId = "service<2>".parse().unwrap();
2570 let agents = vec![
2571 legacy_host_agent_ref(dial_addr.clone()),
2572 host_agent_ref(ProcAddr::new(
2573 service_proc_id.clone(),
2574 dial_addr.clone().into(),
2575 )),
2576 ];
2577 let mesh = HostMeshRef::from_host_agents(
2578 HostMeshId::singleton(Label::new("unix-services").unwrap()),
2579 agents,
2580 )
2581 .unwrap();
2582
2583 let round_trip: HostMeshRef = mesh.to_string().parse().unwrap();
2584 let round_trip_proc_ids = round_trip
2585 .host_agent_mesh
2586 .values()
2587 .map(|agent| agent.actor_addr().proc_addr().id().clone())
2588 .collect::<Vec<_>>();
2589
2590 assert_eq!(
2591 round_trip_proc_ids,
2592 vec![
2593 ResourceId::from_name(SERVICE_PROC_NAME).proc_id(),
2594 service_proc_id,
2595 ]
2596 );
2597 assert_eq!(round_trip.host_addrs(), vec![dial_addr.clone(), dial_addr]);
2598 }
2599
2600 #[tokio::test]
2601 async fn test_sa1_empty_mesh_set_rejected() {
2602 let instance = testing::instance();
2603 let result = spawn_admin(std::iter::empty::<&HostMeshRef>(), instance, None, None).await;
2604 let err = result.unwrap_err().to_string();
2605 assert!(err.contains("SA-1"), "expected SA-1 error, got: {err}");
2606 }
2607
2608 #[tokio::test]
2609 async fn test_sa2_empty_hosts_rejected() {
2610 let instance = testing::instance();
2611 let mesh =
2612 HostMeshRef::from_hosts(HostMeshId::singleton(Label::new("empty").unwrap()), vec![]);
2613 let result = spawn_admin([&mesh], instance, None, None).await;
2614 let err = result.unwrap_err().to_string();
2615 assert!(err.contains("SA-2"), "expected SA-2 error, got: {err}");
2616 }
2617
2618 #[test]
2623 fn test_sa3_host_set_insert_idempotent() {
2624 let addr_a: ChannelAddr = "tcp:127.0.0.1:2001".parse().unwrap();
2625 let addr_b: ChannelAddr = "tcp:127.0.0.1:2002".parse().unwrap();
2626
2627 let ref_a = legacy_host_agent_ref(addr_a.clone());
2628 let ref_b = legacy_host_agent_ref(addr_b.clone());
2629
2630 let mut set = HostSet::new();
2631 set.insert(addr_a.to_string(), ref_a.clone());
2632 set.insert(addr_b.to_string(), ref_b.clone());
2633 set.insert("duplicate_addr".to_string(), ref_a.clone());
2635
2636 let result = set.into_vec();
2637 assert_eq!(
2638 result.len(),
2639 2,
2640 "SA-3: duplicate host identity must not add entry"
2641 );
2642 assert_eq!(
2643 result[0].0,
2644 addr_a.to_string(),
2645 "SA-3: first-seen order preserved"
2646 );
2647 assert_eq!(
2648 result[1].0,
2649 addr_b.to_string(),
2650 "SA-3: first-seen order preserved"
2651 );
2652 }
2653
2654 #[test]
2655 fn test_sa3_aggregate_hosts_dedup() {
2656 let addr_a: ChannelAddr = "tcp:127.0.0.1:1001".parse().unwrap();
2657 let addr_b: ChannelAddr = "tcp:127.0.0.1:1002".parse().unwrap();
2658 let addr_c: ChannelAddr = "tcp:127.0.0.1:1003".parse().unwrap();
2659
2660 let mesh_a = HostMeshRef::from_hosts(
2662 HostMeshId::singleton(Label::new("mesh-a").unwrap()),
2663 vec![addr_a.clone(), addr_b.clone()],
2664 );
2665 let mesh_b = HostMeshRef::from_hosts(
2667 HostMeshId::singleton(Label::new("mesh-b").unwrap()),
2668 vec![addr_b.clone(), addr_c.clone()],
2669 );
2670
2671 let result = aggregate_hosts(&[&mesh_a, &mesh_b], None);
2672
2673 assert_eq!(result.len(), 3, "expected 3 hosts, got {:?}", result);
2675
2676 let addrs: Vec<String> = result.iter().map(|(a, _)| a.clone()).collect();
2678 assert_eq!(addrs[0], addr_a.to_string());
2679 assert_eq!(addrs[1], addr_b.to_string());
2680 assert_eq!(addrs[2], addr_c.to_string());
2681 }
2682
2683 #[test]
2686 fn test_sa6_ch1_client_host_dedup() {
2687 let addr_a: ChannelAddr = "tcp:127.0.0.1:1001".parse().unwrap();
2688 let addr_b: ChannelAddr = "tcp:127.0.0.1:1002".parse().unwrap();
2689
2690 let mesh = HostMeshRef::from_hosts(
2691 HostMeshId::singleton(Label::new("mesh").unwrap()),
2692 vec![addr_a.clone(), addr_b.clone()],
2693 );
2694
2695 let client_ref = legacy_host_agent_ref(addr_a.clone());
2697 let client_entries = vec![("client_addr".to_string(), client_ref)];
2698
2699 let result = aggregate_hosts(&[&mesh], Some(client_entries));
2700
2701 assert_eq!(result.len(), 2, "expected 2 hosts, got {:?}", result);
2703 let addrs: Vec<String> = result.iter().map(|(a, _)| a.clone()).collect();
2704 assert_eq!(addrs[0], addr_a.to_string());
2705 assert_eq!(addrs[1], addr_b.to_string());
2706 }
2707}