disallow logging option for local procs (meta-pytorch#825)

Summary:

there is no way to tail local procs; simply disallow it; I don't quite
like the new interface introduced. RFC is needed.

Differential Revision: D80063700

Author: James Sun

hyperactor_mesh/src/logging.rs

@@ -34,7 +34,6 @@ use hyperactor::Unbind;
 use hyperactor::channel;
 use hyperactor::channel::ChannelAddr;
 use hyperactor::channel::ChannelRx;
-use hyperactor::channel::ChannelTransport;
 use hyperactor::channel::ChannelTx;
 use hyperactor::channel::Rx;
 use hyperactor::channel::Tx;
@@ -539,6 +538,10 @@ impl<T: LogSender + Unpin + 'static, S: io::AsyncWrite + Send + Unpin + 'static>
                 // Since LogSender::send takes &self, we don't need to clone it
                 if let Err(e) = this.log_sender.send(output_target, data_to_send) {
                     tracing::error!("error sending log: {}", e);
+                    return Poll::Ready(Err(io::Error::other(format!(
+                        "error sending write message: {}",
+                        e
+                    ))));
                 }
                 // Return success with the full buffer size
                 Poll::Ready(Ok(buf.len()))
@@ -610,37 +613,15 @@ impl Actor for LogForwardActor {
     async fn new(logging_client_ref: Self::Params) -> Result<Self> {
         let log_channel: ChannelAddr = match std::env::var(BOOTSTRAP_LOG_CHANNEL) {
             Ok(channel) => channel.parse()?,
-            Err(err) => {
-                tracing::debug!(
-                    "log forwarder actor failed to read env var {}: {}",
-                    BOOTSTRAP_LOG_CHANNEL,
-                    err
-                );
-                // TODO: an empty channel to serve
-                ChannelAddr::any(ChannelTransport::Unix)
-            }
+            Err(err) => return Err(err.into()),
         };
         tracing::info!(
             "log forwarder {} serve at {}",
             std::process::id(),
             log_channel
         );
 
-        let rx = match channel::serve(log_channel.clone()).await {
-            Ok((_, rx)) => rx,
-            Err(err) => {
-                // This can happen if we are not spanwed on a separate process like local.
-                // For local mesh, log streaming anyway is not needed.
-                tracing::error!(
-                    "log forwarder actor failed to bootstrap on given channel {}: {}",
-                    log_channel,
-                    err
-                );
-                channel::serve(ChannelAddr::any(ChannelTransport::Unix))
-                    .await?
-                    .1
-            }
-        };
+        let (_, rx) = channel::serve(log_channel.clone()).await?;
         Ok(Self {
             rx,
             logging_client_ref,
@@ -743,15 +724,7 @@ impl Actor for LogFlushActor {
     async fn new(_: ()) -> Result<Self, anyhow::Error> {
         let log_channel: ChannelAddr = match std::env::var(BOOTSTRAP_LOG_CHANNEL) {
             Ok(channel) => channel.parse()?,
-            Err(err) => {
-                tracing::debug!(
-                    "log forwarder actor failed to read env var {}: {}",
-                    BOOTSTRAP_LOG_CHANNEL,
-                    err
-                );
-                // TODO: this should error out; it can only happen with local proc; we need to fix it.
-                ChannelAddr::any(ChannelTransport::Unix)
-            }
+            Err(err) => return Err(err.into()),
         };
         let tx = channel::dial::<LogMessage>(log_channel)?;
 
@@ -1025,6 +998,7 @@ mod tests {
 
     use hyperactor::channel;
     use hyperactor::channel::ChannelAddr;
+    use hyperactor::channel::ChannelTransport;
     use hyperactor::channel::ChannelTx;
     use hyperactor::channel::Tx;
     use hyperactor::id;

python/monarch/_src/actor/allocator.py

@@ -33,6 +33,7 @@
 class AllocHandle(DeprecatedNotAFuture):
     _hy_alloc: "Shared[Alloc]"
     _extent: Dict[str, int]
+    _fork_processes: bool
 
     @property
     def initialized(self) -> Future[Literal[True]]:
@@ -48,6 +49,10 @@ async def task() -> Literal[True]:
 
         return Future(coro=task())
 
+    @property
+    def fork_processes(self) -> bool:
+        return self._fork_processes
+
 
 class AllocateMixin(abc.ABC):
     @abc.abstractmethod
@@ -63,7 +68,14 @@ def allocate(self, spec: AllocSpec) -> "AllocHandle":
         Returns:
         - A future that will be fulfilled when the requested allocation is fulfilled.
         """
-        return AllocHandle(self.allocate_nonblocking(spec).spawn(), spec.extent)
+        return AllocHandle(
+            self.allocate_nonblocking(spec).spawn(),
+            spec.extent,
+            self._fork_processes(),
+        )
+
+    @abc.abstractmethod
+    def _fork_processes(self) -> bool: ...
 
 
 @final
@@ -72,20 +84,29 @@ class ProcessAllocator(ProcessAllocatorBase, AllocateMixin):
     An allocator that allocates by spawning local processes.
     """
 
+    def _fork_processes(self) -> bool:
+        return True
+
 
 @final
 class LocalAllocator(LocalAllocatorBase, AllocateMixin):
     """
     An allocator that allocates by spawning actors into the current process.
     """
 
+    def _fork_processes(self) -> bool:
+        return False
+
 
 @final
 class SimAllocator(SimAllocatorBase, AllocateMixin):
     """
     An allocator that allocates by spawning actors into the current process using simulated channels for transport
     """
 
+    def _fork_processes(self) -> bool:
+        return False
+
 
 class RemoteAllocInitializer(abc.ABC):
     """Subclass-able Python interface for `hyperactor_mesh::alloc::remoteprocess:RemoteProcessAllocInitializer`.
@@ -219,3 +240,6 @@ class RemoteAllocator(RemoteAllocatorBase, AllocateMixin):
     An allocator that allocates by spawning actors on a remote host.
     The remote host must be running hyperactor's remote-process-allocator.
     """
+
+    def _fork_processes(self) -> bool:
+        return True

python/monarch/_src/actor/proc_mesh.py

@@ -142,6 +142,7 @@ def __init__(
         self,
         hy_proc_mesh: "Shared[HyProcMesh]",
         shape: Shape,
+        _fork_processes: bool,
         _device_mesh: Optional["DeviceMesh"] = None,
     ) -> None:
         self._proc_mesh = hy_proc_mesh
@@ -155,6 +156,7 @@ def __init__(
         self._code_sync_client: Optional[CodeSyncMeshClient] = None
         self._logging_mesh_client: Optional[LoggingMeshClient] = None
         self._maybe_device_mesh: Optional["DeviceMesh"] = _device_mesh
+        self._fork_processes = _fork_processes
         self._stopped = False
 
     @property
@@ -172,27 +174,34 @@ async def task() -> Literal[True]:
 
         return Future(coro=task())
 
-    def _init_manager_actors(self, setup: Callable[[], None] | None = None) -> None:
+    def _init_manager_actors(
+        self, setup: Callable[[], None] | None = None, _fork_processes: bool = True
+    ) -> None:
         self._proc_mesh = PythonTask.from_coroutine(
-            self._init_manager_actors_coro(self._proc_mesh, setup)
+            self._init_manager_actors_coro(self._proc_mesh, setup, _fork_processes)
         ).spawn()
 
     async def _init_manager_actors_coro(
         self,
         proc_mesh_: "Shared[HyProcMesh]",
         setup: Callable[[], None] | None = None,
+        _fork_processes: bool = True,
     ) -> "HyProcMesh":
         # WARNING: it is unsafe to await self._proc_mesh here
         # because self._proc_mesh is the result of this function itself!
 
         proc_mesh = await proc_mesh_
 
-        self._logging_mesh_client = await LoggingMeshClient.spawn(proc_mesh=proc_mesh)
-        self._logging_mesh_client.set_mode(
-            stream_to_client=True,
-            aggregate_window_sec=3,
-            level=logging.INFO,
-        )
+        if _fork_processes:
+            # logging mesh is only makes sense with forked (remote or local) processes
+            self._logging_mesh_client = await LoggingMeshClient.spawn(
+                proc_mesh=proc_mesh
+            )
+            self._logging_mesh_client.set_mode(
+                stream_to_client=True,
+                aggregate_window_sec=3,
+                level=logging.INFO,
+            )
 
         _rdma_manager = (
             # type: ignore[16]
@@ -229,7 +238,12 @@ def _new_with_shape(self, shape: Shape) -> "ProcMesh":
             if self._maybe_device_mesh is None
             else self._device_mesh._new_with_shape(shape)
         )
-        pm = ProcMesh(self._proc_mesh, shape, _device_mesh=device_mesh)
+        pm = ProcMesh(
+            self._proc_mesh,
+            shape,
+            _device_mesh=device_mesh,
+            _fork_processes=self._fork_processes,
+        )
         pm._slice = True
         return pm
 
@@ -301,10 +315,12 @@ async def task() -> HyProcMesh:
             list(alloc._extent.keys()),
             Slice.new_row_major(list(alloc._extent.values())),
         )
-        pm = ProcMesh(PythonTask.from_coroutine(task()).spawn(), shape)
+        pm = ProcMesh(
+            PythonTask.from_coroutine(task()).spawn(), shape, alloc.fork_processes
+        )
 
         if _init_manager_actors:
-            pm._init_manager_actors(setup)
+            pm._init_manager_actors(setup, alloc.fork_processes)
             # we do this here rather than in _init_manager_actors
             # because serializing debug_client() requires waiting on
             # its actor to spawn which would block the tokio event loop inside
@@ -450,6 +466,11 @@ async def logging_option(
         Returns:
             None
         """
+        if not self._fork_processes:
+            raise RuntimeError(
+                "Logging option is only available for allocators that fork processes. Allocators like LocalAllocator are not supported."
+            )
+
         if level < 0 or level > 255:
             raise ValueError("Invalid logging level: {}".format(level))
         await self.initialized

python/tests/test_python_actors.py

@@ -1092,3 +1092,12 @@ def test_mesh_len():
     proc_mesh = local_proc_mesh(gpus=12).get()
     s = proc_mesh.spawn("sync_actor", SyncActor).get()
     assert 12 == len(s)
+
+
+async def test_logging_option_on_local_procs() -> None:
+    proc_mesh = local_proc_mesh(gpus=1)
+    with pytest.raises(
+        RuntimeError,
+        match="Logging option is only available for allocators that fork processes",
+    ):
+        await proc_mesh.logging_option(stream_to_client=True)