lazy init log client (#525)

Summary:
Pull Request resolved: #525

like all other actors (_rdma_manager, _debug_manager, _rsync_mesh_client, etc
in the ProcMesh object, let's also make the log client optioanl during init. Otherwise,
it will init multiple times as ProcMesh is only a reference. Only init the client once
we need it.)

Reviewed By: dcci

Differential Revision: D78229069

fbshipit-source-id: 7ba10d0a4bc0cfbe2de03c7a2cd972b8cf7d2db1

Author: James Sun

monarch_extension/src/logging.rs

@@ -28,7 +28,6 @@ use pyo3::types::PyModule;
     module = "monarch._rust_bindings.monarch_extension.logging"
 )]
 pub struct LoggingMeshClient {
-    // TODO: add more interfaces so that users can opt in/out streaming anytime they want
     actor_mesh: SharedCell<RootActorMesh<'static, LogForwardActor>>,
 }
 
@@ -41,10 +40,10 @@ impl LoggingMeshClient {
         signal_safe_block_on(py, async move {
             let client_actor: ActorRef<LogClientActor> = proc_mesh
                 .client_proc()
-                .spawn("logging_client", ())
+                .spawn("log_client", ())
                 .await?
                 .bind();
-            let actor_mesh = proc_mesh.spawn("logging_source", &client_actor).await?;
+            let actor_mesh = proc_mesh.spawn("log_forwarder", &client_actor).await?;
             Ok(Self { actor_mesh })
         })?
     }

python/monarch/_src/actor/proc_mesh.py

@@ -100,9 +100,7 @@ def __init__(
         self._mailbox: Mailbox = self._proc_mesh.client
         self._rsync_mesh_client: Optional[RsyncMeshClient] = None
         self._auto_reload_actor: Optional[AutoReloadActor] = None
-        self._logging_mesh_client: LoggingMeshClient = LoggingMeshClient.spawn_blocking(
-            proc_mesh=self._proc_mesh,
-        )
+        self._logging_mesh_client: Optional[LoggingMeshClient] = None
         self._maybe_device_mesh: Optional["DeviceMesh"] = _device_mesh
         self._stopped = False
         if _mock_shape is None and HAS_TENSOR_ENGINE:
@@ -281,6 +279,10 @@ def logging_option(self, stream_to_client: bool = False) -> None:
         Returns:
             None
         """
+        if self._logging_mesh_client is None:
+            self._logging_mesh_client = LoggingMeshClient.spawn_blocking(
+                proc_mesh=self._proc_mesh,
+            )
         self._logging_mesh_client.set_mode(stream_to_client)
 
     async def stop(self) -> None: