Log ProcMeshStatus with status field (#1882)

Summary:
Pull Request resolved: #1882

The goal is that we can get a `ProcMeshStatus` by searching scuba with:

> mesh_name=<name> and name=ProcMeshStatus

Reviewed By: mariusae

Differential Revision: D86829124

fbshipit-source-id: ce961f8d5d5dc827507e4d93288eea4553b47691

Author: pzhan9

hyperactor/src/channel.rs

@@ -828,8 +828,9 @@ impl<M: RemoteMessage> Rx<M> for ChannelRx<M> {
 /// if the channel cannot be established. The underlying connection is
 /// dropped whenever the returned Tx is dropped.
 #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `ChannelError`.
+#[track_caller]
 pub fn dial<M: RemoteMessage>(addr: ChannelAddr) -> Result<ChannelTx<M>, ChannelError> {
-    tracing::debug!(name = "dial", "dialing channel {}", addr);
+    tracing::debug!(name = "dial", caller = %Location::caller(), %addr, "dialing channel {}", addr);
     let inner = match addr {
         ChannelAddr::Local(port) => ChannelTxKind::Local(local::dial(port)?),
         ChannelAddr::Tcp(addr) => ChannelTxKind::Tcp(net::tcp::dial(addr)),

hyperactor_mesh/src/alloc.rs

@@ -418,7 +418,11 @@ impl<A: ?Sized + Send + Alloc> AllocExt for A {
                 return Err(AllocatorError::Incomplete(self.extent().clone()));
             };
 
-            let name = state.arm().unwrap_or("unknown");
+            let name = tracing::Span::current()
+                .metadata()
+                .map(|m| m.name())
+                .unwrap_or("initialize");
+            let status = format!("ProcState:{}", state.arm().unwrap_or("unknown"));
 
             match state {
                 ProcState::Created {
@@ -427,17 +431,20 @@ impl<A: ?Sized + Send + Alloc> AllocExt for A {
                     let rank = point.rank();
                     if let Some(old_create_key) = created.insert(rank, create_key.clone()) {
                         tracing::warn!(
+                            name,
+                            status,
+                            rank,
                             "rank {rank} reassigned from {old_create_key} to {create_key}"
                         );
                     }
                     tracing::info!(
-                        name = name,
-                        rank = rank,
+                        name,
+                        status,
+                        rank,
                         "proc with create key {}, rank {}: created",
                         create_key,
                         rank
                     );
-                    // tracing::info!("created: {} rank {}: created", create_key, rank);
                 }
                 ProcState::Running {
                     create_key,
@@ -447,7 +454,9 @@ impl<A: ?Sized + Send + Alloc> AllocExt for A {
                 } => {
                     let Some(rank) = created.rank(&create_key) else {
                         tracing::warn!(
-                            name = name,
+                            name,
+                            %proc_id,
+                            status,
                             "proc id {proc_id} with create key {create_key} \
                             is running, but was not created"
                         );
@@ -463,14 +472,19 @@ impl<A: ?Sized + Send + Alloc> AllocExt for A {
                     if let Some(old_allocated_proc) = running.insert(*rank, allocated_proc.clone())
                     {
                         tracing::warn!(
-                            name = name,
+                            name,
+                            %proc_id,
+                            status,
+                            rank,
                             "duplicate running notifications for {rank}: \
                             old:{old_allocated_proc}; \
                             new:{allocated_proc}"
                         )
                     }
                     tracing::info!(
-                        name = name,
+                        name,
+                        %proc_id,
+                        status,
                         "proc {} rank {}: running at addr:{addr} mesh_agent:{mesh_agent}",
                         proc_id,
                         rank
@@ -481,7 +495,8 @@ impl<A: ?Sized + Send + Alloc> AllocExt for A {
                 // ProcState::Failed to fail the whole allocation.
                 ProcState::Stopped { create_key, reason } => {
                     tracing::error!(
-                        name = name,
+                        name,
+                        status,
                         "allocation failed for proc with create key {}: {}",
                         create_key,
                         reason
@@ -493,7 +508,8 @@ impl<A: ?Sized + Send + Alloc> AllocExt for A {
                     description,
                 } => {
                     tracing::error!(
-                        name = name,
+                        name,
+                        status,
                         "allocation failed for world {}: {}",
                         world_id,
                         description

hyperactor_mesh/src/v1/proc_mesh.rs

@@ -51,8 +51,6 @@ use serde::Deserialize;
 use serde::Serialize;
 use tokio::sync::Notify;
 use tracing::Instrument;
-use tracing::Level;
-use tracing::span;
 
 use crate::CommActor;
 use crate::alloc::Alloc;
@@ -295,28 +293,40 @@ impl ProcMesh {
 
     /// Allocate a new ProcMesh from the provided alloc.
     /// Allocate does not require an owning actor because references are not owned.
-    #[tracing::instrument(skip_all)]
     #[track_caller]
     pub async fn allocate(
         cx: &impl context::Actor,
-        mut alloc: Box<dyn Alloc + Send + Sync + 'static>,
+        alloc: Box<dyn Alloc + Send + Sync + 'static>,
         name: &str,
+    ) -> v1::Result<Self> {
+        let caller = Location::caller();
+        Self::allocate_inner(cx, alloc, Name::new(name), caller).await
+    }
+
+    // Use allocate_inner to set field mesh_name in span
+    #[hyperactor::instrument(fields(mesh_name=name.to_string()))]
+    async fn allocate_inner(
+        cx: &impl context::Actor,
+        mut alloc: Box<dyn Alloc + Send + Sync + 'static>,
+        name: Name,
+        caller: &'static Location<'static>,
     ) -> v1::Result<Self> {
         let alloc_id = Self::alloc_counter().fetch_add(1, Ordering::Relaxed) + 1;
         tracing::info!(
-            name = "ProcMesh::Allocate::Attempt",
+            name = "ProcMeshStatus",
+            status = "Allocate::Attempt",
+            %caller,
             alloc_id,
-            caller = %Location::caller(),
             shape = ?alloc.shape(),
             "allocating proc mesh"
         );
 
         let running = alloc
             .initialize()
-            .instrument(span!(
-                Level::INFO,
-                "ProcMesh::Allocate::Initialize",
-                alloc_id
+            .instrument(tracing::info_span!(
+                "ProcMeshStatus::Allocate::Initialize",
+                alloc_id,
+                mesh_name = name.to_string()
             ))
             .await?;
 
@@ -330,17 +340,18 @@ impl ProcMesh {
         // First make sure we can serve the proc:
         let proc_channel_addr = {
             let _guard =
-                tracing::span!(Level::INFO, "allocate_serve_proc", proc_id = %proc.proc_id())
-                    .entered();
+                tracing::info_span!("allocate_serve_proc", proc_id = %proc.proc_id()).entered();
             let (addr, rx) = channel::serve(ChannelAddr::any(alloc.transport()))?;
             proc.clone().serve(rx);
+            tracing::info!(
+                name = "ProcMeshStatus",
+                status = "Allocate::ChannelServe",
+                mesh_name = name.to_string(),
+                %addr,
+                "proc started listening on addr: {addr}"
+            );
             addr
         };
-        tracing::info!(
-            name = "ProcMesh::Allocate::ChannelServe",
-            alloc_id = alloc_id,
-            "proc started listening on addr: {proc_channel_addr}"
-        );
 
         let bind_allocated_procs = |router: &DialMailboxRouter| {
             // Route all of the allocated procs:
@@ -423,7 +434,6 @@ impl ProcMesh {
             .collect();
 
         let stop = Arc::new(Notify::new());
-        let name = Name::new(name);
         let extent = alloc.extent().clone();
 
         {
@@ -453,10 +463,10 @@ impl ProcMesh {
                         }
                     }
                 }
-            });
+            }.instrument(tracing::info_span!("alloc_monitor")));
         }
 
-        Self::create(
+        let mesh = Self::create(
             cx,
             name,
             ProcMeshAllocation::Allocated {
@@ -466,14 +476,21 @@ impl ProcMesh {
             },
             true, // alloc-based meshes support comm actors
         )
-        .await
+        .await;
+        match &mesh {
+            Ok(_) => tracing::info!(name = "ProcMeshStatus", status = "Allocate::Created"),
+            Err(error) => {
+                tracing::info!(name = "ProcMeshStatus", status = "Allocate::Failed", %error)
+            }
+        }
+        mesh
     }
 
     /// Detach the proc mesh from the lifetime of `self`, and return its reference.
-    #[allow(dead_code)]
+    #[cfg(test)]
     pub(crate) fn detach(self) -> ProcMeshRef {
         // This also keeps the ProcMeshAllocation::Allocated alloc task alive.
-        self.current_ref
+        self.current_ref.clone()
     }
 
     /// Stop this mesh gracefully.
@@ -508,6 +525,16 @@ impl Deref for ProcMesh {
     }
 }
 
+impl Drop for ProcMesh {
+    fn drop(&mut self) {
+        tracing::info!(
+            name = "ProcMeshStatus",
+            mesh_name = %self.name,
+            status = "Dropped",
+        );
+    }
+}
+
 /// Represents different ways ProcMeshes can be allocated.
 enum ProcMeshAllocation {
     /// A mesh that has been allocated from an `Alloc`.