: v1/logging: avoid spinning up LogForwardActor mesh when redundant

Differential Revision: D85919326

Author: shayne-fletcher

hyperactor_mesh/Cargo.toml

@@ -63,6 +63,7 @@ hyperactor_mesh_macros = { version = "0.0.0", path = "../hyperactor_mesh_macros"
 hyperactor_telemetry = { version = "0.0.0", path = "../hyperactor_telemetry" }
 libc = "0.2.139"
 mockall = "0.13.1"
+monarch_hyperactor = { version = "0.0.0", path = "../monarch_hyperactor" }
 ndslice = { version = "0.0.0", path = "../ndslice" }
 nix = { version = "0.30.1", features = ["dir", "event", "hostname", "inotify", "ioctl", "mman", "mount", "net", "poll", "ptrace", "reboot", "resource", "sched", "signal", "term", "time", "user", "zerocopy"] }
 opentelemetry = "0.29"

hyperactor_mesh/test/bootstrap.rs

@@ -6,6 +6,19 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+// Make sure this binary pulls in monarch_hyperactor's attr registry
+// (inventory::submit! side effects). With LTO / linker GC, just
+// adding the crate as a dep is not enough: if nothing in this binary
+// references it, the linker will drop it and its AttrKeyInfo
+// submissions will never register. Touching a symbol makes the linker
+// keep the crate. Issue
+// https://github.com/dtolnay/inventory/issues/7 says this shouldn't
+// be necessary (see also D71096199).
+#[allow(dead_code)]
+fn _link_monarch_hyperactor_attrs() {
+    let _ = monarch_hyperactor::pytokio::UNAWAITED_PYTOKIO_TRACEBACK.name();
+}
+
 /// This is an "empty shell" bootstrap process,
 /// simply invoking [`hyperactor_mesh::bootstrap_or_die`].
 #[tokio::main]

monarch_hyperactor/Cargo.toml

@@ -52,6 +52,7 @@ tokio-util = { version = "0.7.15", features = ["full"] }
 tracing = { version = "0.1.41", features = ["attributes", "valuable"] }
 
 [dev-dependencies]
+buck-resources = "1"
 dir-diff = "0.3"
 
 [features]

monarch_hyperactor/src/v1/logging.rs

@@ -37,10 +37,20 @@ use crate::v1::proc_mesh::PyProcMesh;
     module = "monarch._rust_bindings.monarch_hyperactor.v1.logging"
 )]
 pub struct LoggingMeshClient {
-    // handles remote process log forwarding; no python runtime
-    forwarder_mesh: ActorMesh<LogForwardActor>,
-    // handles python logger; has python runtime
+    // Per-proc LogForwardActor mesh (optional). When enabled, each
+    // remote proc forwards its stdout/stderr back to the client. This
+    // actor does not interact with the embedded Python runtime.
+    forwarder_mesh: Option<ActorMesh<LogForwardActor>>,
+
+    // Per-proc LoggerRuntimeActor mesh. Runs on every proc in the
+    // mesh and drives that proc's Python logging configuration (log
+    // level, handlers, etc.). If the proc isn't running embedded
+    // Python, this is effectively a no-op.
     logger_mesh: ActorMesh<LoggerRuntimeActor>,
+
+    // Client-side LogClientActor. Lives in the client process;
+    // receives forwarded output, aggregates and buffers it, and
+    // coordinates sync flush barriers.
     client_actor: ActorHandle<LogClientActor>,
 }
 
@@ -74,11 +84,40 @@ impl LoggingMeshClient {
 
 #[pymethods]
 impl LoggingMeshClient {
+    /// Initialize logging for a `ProcMesh` and return a
+    /// `LoggingMeshClient`.
+    ///
+    /// This wires up three pieces of logging infrastructure:
+    ///
+    /// 1. A single `LogClientActor` in the *client* process. This
+    ///    actor receives forwarded stdout/stderr, buffers and
+    ///    aggregates it, and coordinates sync flush barriers.
+    ///
+    /// 2. (Optional) A `LogForwardActor` on every remote proc in the
+    ///    mesh. These forwarders read that proc's stdout/stderr and
+    ///    stream it back to the client. We only spawn this mesh if
+    ///    `MESH_ENABLE_LOG_FORWARDING` was `true` in the config. If
+    ///    forwarding is disabled at startup, we do not spawn these
+    ///    actors and `forwarder_mesh` will be `None`.
+    ///
+    /// 3. A `LoggerRuntimeActor` on every remote proc in the mesh.
+    ///    This actor controls the Python logging runtime (log level,
+    ///    handlers, etc.) in that process. This is always spawned,
+    ///    even if log forwarding is disabled.
+    ///
+    /// The returned `LoggingMeshClient` holds handles to those
+    /// actors. Later, `set_mode(...)` can adjust per-proc log level
+    /// and (if forwarding was enabled) toggle whether remote output
+    /// is actually streamed back to the client. If forwarding was
+    /// disabled by config, requests to enable streaming will fail.
     #[staticmethod]
     fn spawn(instance: &PyInstance, proc_mesh: &PyProcMesh) -> PyResult<PyPythonTask> {
         let proc_mesh = proc_mesh.mesh_ref()?;
         let instance = instance.clone();
+
         PyPythonTask::new(async move {
+            // 1. Spawn the client-side coordinator actor (lives in
+            // the caller's process).
             let client_actor: ActorHandle<LogClientActor> =
                 instance_dispatch!(instance, async move |cx_instance| {
                     cx_instance
@@ -87,12 +126,30 @@ impl LoggingMeshClient {
                         .await
                 })?;
             let client_actor_ref = client_actor.bind();
-            let forwarder_mesh = instance_dispatch!(instance, async |cx_instance| {
-                proc_mesh
-                    .spawn(cx_instance, "log_forwarder", &client_actor_ref)
-                    .await
-            })
-            .map_err(anyhow::Error::from)?;
+
+            // Read config to decide if we stand up per-proc
+            // stdout/stderr forwarding.
+            let forwarding_enabled = hyperactor::config::global::get(
+                hyperactor_mesh::bootstrap::MESH_ENABLE_LOG_FORWARDING,
+            );
+
+            // 2. Optionally spawn per-proc `LogForwardActor` mesh
+            // (stdout/stderr forwarders).
+            let forwarder_mesh = if forwarding_enabled {
+                // Spawn a `LogFwdActor` on every proc.
+                let mesh = instance_dispatch!(instance, async |cx_instance| {
+                    proc_mesh
+                        .spawn(cx_instance, "log_forwarder", &client_actor_ref)
+                        .await
+                })
+                .map_err(anyhow::Error::from)?;
+
+                Some(mesh)
+            } else {
+                None
+            };
+
+            // 3. Always spawn a `LoggerRuntimeActor` on every proc.
             let logger_mesh = instance_dispatch!(instance, async |cx_instance| {
                 proc_mesh.spawn(cx_instance, "logger", &()).await
             })
@@ -106,31 +163,87 @@ impl LoggingMeshClient {
         })
     }
 
+    /// Update logging behavior for this mesh.
+    ///
+    /// `stream_to_client` controls whether remote procs actively
+    /// stream their stdout/stderr back to the client process.
+    ///
+    /// - If log forwarding was enabled at startup, `forwarder_mesh`
+    ///   is `Some` and we propagate this flag to every per-proc
+    ///   `LogForwardActor`.
+    /// - If log forwarding was disabled at startup, `forwarder_mesh`
+    ///   is `None`.
+    ///   In that case:
+    ///     * requesting `stream_to_client = false` is a no-op
+    ///       (accepted),
+    ///     * requesting `stream_to_client = true` is rejected,
+    ///       because we did not spawn forwarders and we don't
+    ///       dynamically create them later.
+    ///
+    /// `aggregate_window_sec` configures how the client-side
+    /// `LogClientActor` batches forwarded output. It is only
+    /// meaningful when streaming is enabled. Calling this with
+    /// `Some(..)` while `stream_to_client == false` is invalid and
+    /// returns an error.
+    ///
+    /// `level` is the desired Python logging level. We always
+    /// broadcast this to the per-proc `LoggerRuntimeActor` mesh so
+    /// each remote process can update its own Python logger
+    /// configuration, regardless of whether stdout/stderr forwarding
+    /// is active.
     fn set_mode(
         &self,
         instance: &PyInstance,
         stream_to_client: bool,
         aggregate_window_sec: Option<u64>,
         level: u8,
     ) -> PyResult<()> {
+        // We can't ask for an aggregation window if we're not
+        // streaming.
         if aggregate_window_sec.is_some() && !stream_to_client {
             return Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                 "cannot set aggregate window without streaming to client".to_string(),
             ));
         }
 
-        instance_dispatch!(instance, |cx_instance| {
-            self.forwarder_mesh
-                .cast(cx_instance, LogForwardMessage::SetMode { stream_to_client })
-        })
-        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+        // Handle the forwarder side (stdout/stderr streaming back to
+        // client).
+        match (&self.forwarder_mesh, stream_to_client) {
+            // Forwarders exits (config enabled at startup). We can
+            // toggle live.
+            (Some(fwd_mesh), _) => {
+                instance_dispatch!(instance, |cx_instance| {
+                    fwd_mesh.cast(cx_instance, LogForwardMessage::SetMode { stream_to_client })
+                })
+                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+            }
+
+            // Forwarders were never spawned (global forwarding
+            // disabled) and the caller is asking NOT to stream.
+            // That's effectively a no-op so we silently accept.
+            (None, false) => {
+                // Nothing to do.
+            }
+
+            // Forwarders were never spawned, but caller is asking to
+            // stream. We can't satisfy this request without
+            // re-spawning infra, which we deliberately don't do at
+            // runtime.
+            (None, true) => {
+                return Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
+                    "log forwarding disabled by config at startup; cannot enable streaming_to_client",
+                ));
+            }
+        }
 
+        // Always update the per-proc Python logging level.
         instance_dispatch!(instance, |cx_instance| {
             self.logger_mesh
                 .cast(cx_instance, LoggerRuntimeMessage::SetLogging { level })
         })
         .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
 
+        // Always update the client actor's aggregation window.
         self.client_actor
             .send(LogClientMessage::SetAggregate {
                 aggregate_window_sec,
@@ -140,13 +253,28 @@ impl LoggingMeshClient {
         Ok(())
     }
 
-    // A sync flush mechanism for the client make sure all the stdout/stderr are streamed back and flushed.
+    /// Force a sync flush of remote stdout/stderr back to the client,
+    /// and wait for completion.
+    ///
+    /// If log forwarding was disabled at startup (so we never spawned
+    /// any `LogForwardActor`s), this becomes a no-op success: there's
+    /// nothing to flush from remote procs in that mode, and we don't
+    /// try to manufacture it dynamically.
     fn flush(&self, instance: &PyInstance) -> PyResult<PyPythonTask> {
-        let forwarder_mesh = self.forwarder_mesh.deref().clone();
+        let forwarder_mesh_opt = self
+            .forwarder_mesh
+            .as_ref()
+            .map(|mesh| mesh.deref().clone());
         let client_actor = self.client_actor.clone();
         let instance = instance.clone();
 
         PyPythonTask::new(async move {
+            // If there's no forwarer mesh (forwarding disabled by
+            // config), we just succeed immediately.
+            let Some(forwarder_mesh) = forwarder_mesh_opt else {
+                return Ok(());
+            };
+
             instance_dispatch!(instance, async move |cx_instance| {
                 Self::flush_internal(cx_instance, client_actor, forwarder_mesh).await
             })
@@ -171,3 +299,97 @@ pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
     module.add_class::<LoggingMeshClient>()?;
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use hyperactor::Instance;
+    use hyperactor::channel::ChannelTransport;
+    use hyperactor::proc::Proc;
+    use hyperactor_mesh::alloc::AllocSpec;
+    use hyperactor_mesh::alloc::Allocator;
+    use hyperactor_mesh::alloc::ProcessAllocator;
+    use hyperactor_mesh::v1::ProcMesh;
+    use hyperactor_mesh::v1::host_mesh::HostMesh;
+    use ndslice::Extent;
+    use ndslice::View; // .region(), .num_ranks() etc.
+    use ndslice::extent;
+    use tokio::process::Command;
+
+    /// Bring up a minimal "world" suitable for integration-style
+    /// tests.
+    ///
+    /// This does all the real work:
+    ///   - spawns a root `Proc` locally,
+    ///   - creates a client `Instance<()>` on that proc for
+    ///     sending/receiving messages,
+    ///   - allocates a 1-host `HostMesh` via the bootstrap binary,
+    ///   - spawns a 1-rank `ProcMesh` on that host.
+    #[cfg(fbcode_build)] // Uses buck resources.
+    pub async fn test_world() -> anyhow::Result<(
+        Proc,         // root proc
+        Instance<()>, // client instance handle for messaging
+        HostMesh,     // allocated host mesh
+        ProcMesh,     // spawned proc mesh
+    )> {
+        let proc = Proc::direct(ChannelTransport::Unix.any(), "root".to_string())
+            .await
+            .expect("failed to start root Proc");
+
+        let (instance, _handle) = proc
+            .instance("client")
+            .expect("failed to create Proc instance");
+
+        let mut allocator = ProcessAllocator::new(Command::new(
+            buck_resources::get("monarch/monarch_hyperactor/bootstrap")
+                .expect("missing bootstrap exe"),
+        ));
+
+        let alloc = allocator
+            .allocate(AllocSpec {
+                extent: extent!(replicas = 1),
+                constraints: Default::default(),
+                proc_name: None,
+                transport: ChannelTransport::Unix,
+            })
+            .await
+            .expect("allocator.allocate failed");
+
+        let host_mesh = HostMesh::allocate(&instance, Box::new(alloc), "test", None)
+            .await
+            .expect("HostMesh::allocate failed");
+
+        let proc_mesh = host_mesh
+            .spawn(&instance, "p0", Extent::unity())
+            .await
+            .expect("host_mesh.spawn failed");
+
+        Ok((proc, instance, host_mesh, proc_mesh))
+    }
+
+    #[cfg(fbcode_build)]
+    #[tokio::test]
+    async fn test_world_smoke() {
+        let (proc, instance, host_mesh, proc_mesh) =
+            test_world().await.expect("bootstrap_cannonical failed");
+
+        assert_eq!(
+            host_mesh.region().num_ranks(),
+            1,
+            "bootstrap_cannonical() should allocate exactly one host"
+        );
+
+        assert_eq!(
+            proc_mesh.region().num_ranks(),
+            1,
+            "bootstrap_cannonical() should spawn exactly one proc"
+        );
+
+        assert_eq!(
+            instance.self_id().proc_id(),
+            proc.proc_id(),
+            "returned Instance<()> should be bound to the root Proc"
+        );
+
+        host_mesh.shutdown(&instance).await.expect("host shutdown");
+    }
+}