Client trace id for connecting events on backend and frontend (#897)

Summary:
Pull Request resolved: #897

Introducing a client context, which currently only contains a trace id, which can be used to connect events on the client and worker side. The trace id will be logged to scuba, and is saved as an env var on the process.

get_trace_id_or_create_new is first called when we initialize the rust bindings and the trace_id itself is first assigned to be the same as the execution id.

Then, as part of the alloc we also pass this trace id along. For RemoteProcessAllocator, the Allocate message will take the client context, and on receiving it, set it in the AllocSpec before allocating the new processes.

The ProcessAlloc will set the env var after reading the trace_id from the match_labels of the AllocSpec. It is only set when the client has set it and passed it with the RemoteAlloc. It will default to "".

Reviewed By: eliothedeman, vidhyav

Differential Revision: D79537187

fbshipit-source-id: fbea0889f65616f2faecd959cae7000d0dbcbfe8

Author: jayasi

hyperactor_mesh/src/alloc/process.rs

@@ -33,6 +33,8 @@ use hyperactor::sync::monitor;
 use ndslice::view::Extent;
 use nix::sys::signal;
 use nix::unistd::Pid;
+use serde::Deserialize;
+use serde::Serialize;
 use tokio::io;
 use tokio::process::Command;
 use tokio::sync::Mutex;
@@ -56,6 +58,8 @@ use crate::shortuuid::ShortUuid;
 /// The maximum number of log lines to tail keep for managed processes.
 const MAX_TAIL_LOG_LINES: usize = 100;
 
+pub const CLIENT_TRACE_ID_LABEL: &str = "CLIENT_TRACE_ID";
+
 /// An allocator that allocates procs by executing managed (local)
 /// processes. ProcessAllocator is configured with a [`Command`] (template)
 /// to spawn external processes. These processes must invoke [`hyperactor_mesh::bootstrap`] or
@@ -102,10 +106,26 @@ impl Allocator for ProcessAllocator {
             children: JoinSet::new(),
             running: true,
             failed: false,
+            client_context: ClientContext {
+                trace_id: spec
+                    .constraints
+                    .match_labels
+                    .get(CLIENT_TRACE_ID_LABEL)
+                    .cloned()
+                    .unwrap_or_else(|| "".to_string()),
+            },
         })
     }
 }
 
+// Client Context is saved in ProcessAlloc, and is also passed in
+// the RemoteProcessAllocator's Allocate method
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ClientContext {
+    /// Trace ID for correlating logs across client and worker processes
+    pub trace_id: String,
+}
+
 /// An allocation produced by [`ProcessAllocator`].
 pub struct ProcessAlloc {
     name: ShortUuid,
@@ -121,6 +141,7 @@ pub struct ProcessAlloc {
     children: JoinSet<(usize, ProcStopReason)>,
     running: bool,
     failed: bool,
+    client_context: ClientContext,
 }
 
 #[derive(EnumAsInner)]
@@ -389,6 +410,10 @@ impl ProcessAlloc {
             bootstrap::BOOTSTRAP_ADDR_ENV,
             self.bootstrap_addr.to_string(),
         );
+        cmd.env(
+            bootstrap::CLIENT_TRACE_ID_ENV,
+            self.client_context.trace_id.as_str(),
+        );
         cmd.env(bootstrap::BOOTSTRAP_INDEX_ENV, index.to_string());
         cmd.env(bootstrap::BOOTSTRAP_LOG_CHANNEL, log_channel.to_string());
         cmd.stdout(Stdio::piped());

hyperactor_mesh/src/alloc/remoteprocess.rs

@@ -55,12 +55,15 @@ use tokio_stream::wrappers::WatchStream;
 use tokio_util::sync::CancellationToken;
 
 use crate::alloc::Alloc;
+use crate::alloc::AllocConstraints;
 use crate::alloc::AllocSpec;
 use crate::alloc::Allocator;
 use crate::alloc::AllocatorError;
 use crate::alloc::ProcState;
 use crate::alloc::ProcStopReason;
 use crate::alloc::ProcessAllocator;
+use crate::alloc::process::CLIENT_TRACE_ID_LABEL;
+use crate::alloc::process::ClientContext;
 
 /// Control messages sent from remote process allocator to local allocator.
 #[derive(Debug, Clone, Serialize, Deserialize, Named)]
@@ -74,6 +77,10 @@ pub enum RemoteProcessAllocatorMessage {
         /// Ordered list of hosts in this allocation. Can be used to
         /// pre-populate the any local configurations such as torch.dist.
         hosts: Vec<String>,
+        /// Client context which is passed to the ProcessAlloc
+        /// Todo: Once RemoteProcessAllocator moves to mailbox,
+        /// the client_context will go to the message header instead
+        client_context: Option<ClientContext>,
     },
     /// Stop allocation.
     Stop,
@@ -196,15 +203,25 @@ impl RemoteProcessAllocator {
                             view,
                             bootstrap_addr,
                             hosts,
+                            client_context,
                         }) => {
                             tracing::info!("received allocation request for view: {}", view);
                             ensure_previous_alloc_stopped(&mut active_allocation).await;
                             tracing::info!("allocating...");
 
                             // Create the corresponding local allocation spec.
+                            let mut constraints: AllocConstraints = Default::default();
+                            if let Some(context) = &client_context {
+                                constraints = AllocConstraints {
+                                    match_labels: HashMap::from([(
+                                    CLIENT_TRACE_ID_LABEL.to_string(),
+                                    context.trace_id.to_string(),
+                                    )]
+                                )};
+                            }
                             let spec = AllocSpec {
                                 extent: view.extent(),
-                                constraints: Default::default(),
+                                constraints,
                             };
 
                             match process_allocator.allocate(spec.clone()).await {
@@ -749,10 +766,15 @@ impl RemoteProcessAlloc {
                     "failed to dial remote {} for host {}",
                     remote_addr, host.id
                 ))?;
+
+            let trace_id = hyperactor_telemetry::trace::get_or_create_trace_id();
+            let client_context = Some(ClientContext { trace_id });
+
             tx.post(RemoteProcessAllocatorMessage::Allocate {
                 view,
                 bootstrap_addr: self.bootstrap_addr.clone(),
                 hosts: hostnames.clone(),
+                client_context,
             });
 
             self.hosts_by_offset.insert(offset, host.id.clone());
@@ -1280,6 +1302,7 @@ mod test {
             view: extent.clone().into(),
             bootstrap_addr,
             hosts: vec![],
+            client_context: None,
         })
         .await
         .unwrap();
@@ -1419,6 +1442,7 @@ mod test {
             view: extent.clone().into(),
             bootstrap_addr,
             hosts: vec![],
+            client_context: None,
         })
         .await
         .unwrap();
@@ -1519,6 +1543,7 @@ mod test {
             view: extent.clone().into(),
             bootstrap_addr: bootstrap_addr.clone(),
             hosts: vec![],
+            client_context: None,
         })
         .await
         .unwrap();
@@ -1539,6 +1564,7 @@ mod test {
             view: extent.clone().into(),
             bootstrap_addr,
             hosts: vec![],
+            client_context: None,
         })
         .await
         .unwrap();
@@ -1632,6 +1658,7 @@ mod test {
             view: extent.clone().into(),
             bootstrap_addr,
             hosts: vec![],
+            client_context: None,
         })
         .await
         .unwrap();
@@ -1721,6 +1748,7 @@ mod test {
             view: extent.clone().into(),
             bootstrap_addr,
             hosts: vec![],
+            client_context: None,
         })
         .await
         .unwrap();
@@ -1747,6 +1775,150 @@ mod test {
         remote_allocator.terminate();
         handle.await.unwrap().unwrap();
     }
+
+    #[timed_test::async_timed_test(timeout_secs = 15)]
+    async fn test_trace_id_propagation() {
+        let config = hyperactor::config::global::lock();
+        let _guard = config.override_key(
+            hyperactor::config::REMOTE_ALLOCATOR_HEARTBEAT_INTERVAL,
+            Duration::from_secs(60),
+        );
+        hyperactor_telemetry::initialize_logging(ClockKind::default());
+        let serve_addr = ChannelAddr::any(ChannelTransport::Unix);
+        let bootstrap_addr = ChannelAddr::any(ChannelTransport::Unix);
+        let (_, mut rx) = channel::serve(bootstrap_addr.clone()).await.unwrap();
+
+        let extent = extent!(host = 1, gpu = 1);
+        let tx = channel::dial(serve_addr.clone()).unwrap();
+        let test_world_id: WorldId = id!(test_world_id);
+        let test_trace_id = "test_trace_id_12345";
+
+        // Create a mock alloc that we can verify receives the correct trace id
+        let mut alloc = MockAlloc::new();
+        alloc.expect_world_id().return_const(test_world_id.clone());
+        alloc.expect_extent().return_const(extent.clone());
+        alloc.expect_next().return_const(None);
+
+        // Create a mock allocator that captures the AllocSpec passed to it
+        let mut allocator = MockAllocator::new();
+        allocator
+            .expect_allocate()
+            .times(1)
+            .withf(move |spec: &AllocSpec| {
+                // Verify that the trace id is correctly set in the constraints
+                spec.constraints
+                    .match_labels
+                    .get(CLIENT_TRACE_ID_LABEL)
+                    .is_some_and(|trace_id| trace_id == test_trace_id)
+            })
+            .return_once(|_| Ok(MockAllocWrapper::new(alloc)));
+
+        let remote_allocator = RemoteProcessAllocator::new();
+        let handle = tokio::spawn({
+            let remote_allocator = remote_allocator.clone();
+            async move {
+                remote_allocator
+                    .start_with_allocator(serve_addr, allocator, None)
+                    .await
+            }
+        });
+
+        // Send allocate message with client context containing trace id
+        tx.send(RemoteProcessAllocatorMessage::Allocate {
+            view: extent.clone().into(),
+            bootstrap_addr,
+            hosts: vec![],
+            client_context: Some(ClientContext {
+                trace_id: test_trace_id.to_string(),
+            }),
+        })
+        .await
+        .unwrap();
+
+        // Verify we get the allocated message
+        let m = rx.recv().await.unwrap();
+        assert_matches!(
+            m,
+            RemoteProcessProcStateMessage::Allocated { world_id, view }
+            if world_id == test_world_id && view.extent() == extent
+        );
+
+        // Verify we get the done message since the mock alloc returns None immediately
+        let m = rx.recv().await.unwrap();
+        assert_matches!(m, RemoteProcessProcStateMessage::Done(world_id) if world_id == test_world_id);
+
+        remote_allocator.terminate();
+        handle.await.unwrap().unwrap();
+    }
+
+    #[timed_test::async_timed_test(timeout_secs = 15)]
+    async fn test_trace_id_propagation_no_client_context() {
+        let config = hyperactor::config::global::lock();
+        let _guard = config.override_key(
+            hyperactor::config::REMOTE_ALLOCATOR_HEARTBEAT_INTERVAL,
+            Duration::from_secs(60),
+        );
+        hyperactor_telemetry::initialize_logging(ClockKind::default());
+        let serve_addr = ChannelAddr::any(ChannelTransport::Unix);
+        let bootstrap_addr = ChannelAddr::any(ChannelTransport::Unix);
+        let (_, mut rx) = channel::serve(bootstrap_addr.clone()).await.unwrap();
+
+        let extent = extent!(host = 1, gpu = 1);
+        let tx = channel::dial(serve_addr.clone()).unwrap();
+        let test_world_id: WorldId = id!(test_world_id);
+
+        // Create a mock alloc
+        let mut alloc = MockAlloc::new();
+        alloc.expect_world_id().return_const(test_world_id.clone());
+        alloc.expect_extent().return_const(extent.clone());
+        alloc.expect_next().return_const(None);
+
+        // Create a mock allocator that verifies no trace id is set when client_context is None
+        let mut allocator = MockAllocator::new();
+        allocator
+            .expect_allocate()
+            .times(1)
+            .withf(move |spec: &AllocSpec| {
+                // Verify that no trace id is set in the constraints when client_context is None
+                spec.constraints.match_labels.is_empty()
+            })
+            .return_once(|_| Ok(MockAllocWrapper::new(alloc)));
+
+        let remote_allocator = RemoteProcessAllocator::new();
+        let handle = tokio::spawn({
+            let remote_allocator = remote_allocator.clone();
+            async move {
+                remote_allocator
+                    .start_with_allocator(serve_addr, allocator, None)
+                    .await
+            }
+        });
+
+        // Send allocate message without client context
+        tx.send(RemoteProcessAllocatorMessage::Allocate {
+            view: extent.clone().into(),
+            bootstrap_addr,
+            hosts: vec![],
+            client_context: None,
+        })
+        .await
+        .unwrap();
+
+        // Verify we get the allocated message
+        let m = rx.recv().await.unwrap();
+        assert_matches!(
+            m,
+            RemoteProcessProcStateMessage::Allocated { world_id, view }
+            if world_id == test_world_id && view.extent() == extent
+        );
+
+        // Verify we get the done message since the mock alloc returns None immediately
+        let m = rx.recv().await.unwrap();
+        assert_matches!(m, RemoteProcessProcStateMessage::Done(world_id) if world_id == test_world_id);
+
+        remote_allocator.terminate();
+        handle.await.unwrap().unwrap();
+    }
 }
 
 #[cfg(test)]

hyperactor_mesh/src/bootstrap.rs

@@ -29,6 +29,7 @@ use crate::proc_mesh::mesh_agent::MeshAgent;
 
 pub const BOOTSTRAP_ADDR_ENV: &str = "HYPERACTOR_MESH_BOOTSTRAP_ADDR";
 pub const BOOTSTRAP_INDEX_ENV: &str = "HYPERACTOR_MESH_INDEX";
+pub const CLIENT_TRACE_ID_ENV: &str = "MONARCH_CLIENT_TRACE_ID";
 /// A channel used by each process to receive its own stdout and stderr
 /// Because stdout and stderr can only be obtained by the parent process,
 /// they need to be streamed back to the process.

hyperactor_telemetry/src/lib.rs

@@ -34,6 +34,7 @@ mod pool;
 pub mod recorder;
 mod spool;
 pub mod sqlite;
+pub mod trace;
 use std::io::IsTerminal;
 use std::io::Write;
 use std::str::FromStr;

hyperactor_telemetry/src/trace.rs

@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use std::env;
+
+use crate::env::execution_id;
+
+const MONARCH_CLIENT_TRACE_ID: &str = "MONARCH_CLIENT_TRACE_ID";
+
+/// Returns the current trace ID if it exists, or None if it doesn't.
+/// This function does not create a new trace ID if one doesn't exist.
+/// Todo: Eventually use Message Headers to relay this traceid instead of
+/// env vars
+pub fn get_trace_id() -> Option<String> {
+    if let Ok(env_id) = env::var(MONARCH_CLIENT_TRACE_ID) {
+        if !env_id.is_empty() {
+            return Some(env_id);
+        }
+    }
+
+    // No trace ID found
+    None
+}
+
+/// Returns the current trace ID, if none exists, set the current execution id as the trace id.
+/// This ensures that the client trace id and execution id is the same.
+/// The trace ID remains the same as long as it is in the same process.
+/// Use this method only on the client side.
+pub fn get_or_create_trace_id() -> String {
+    if let Ok(existing_trace_id) = std::env::var(MONARCH_CLIENT_TRACE_ID) {
+        if !existing_trace_id.is_empty() {
+            return existing_trace_id;
+        }
+    }
+
+    let trace_id = execution_id().clone();
+    // Safety: Can be unsound if there are multiple threads
+    // reading and writing the environment.
+    unsafe {
+        std::env::set_var(MONARCH_CLIENT_TRACE_ID, trace_id.clone());
+    }
+
+    trace_id
+}