disallow logging option for local procs (#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;
@@ -550,6 +549,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()))
@@ -621,37 +624,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,
@@ -754,15 +735,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)?;
 
@@ -1036,6 +1009,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

@@ -179,6 +179,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
@@ -193,6 +194,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
         self._controller_controller: Optional["_ControllerController"] = None
 
@@ -225,7 +227,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,
+            _fork_processes=self._fork_processes,
+            _device_mesh=device_mesh,
+        )
         pm._slice = True
         return pm
 
@@ -300,7 +307,12 @@ async def task() -> HyProcMesh:
 
         hy_proc_mesh = PythonTask.from_coroutine(task()).spawn()
 
-        pm = ProcMesh(hy_proc_mesh, shape)
+        fork_processes: bool = alloc.fork_processes
+        pm = ProcMesh(
+            hy_proc_mesh,
+            shape,
+            _fork_processes=fork_processes,
+        )
 
         async def task(
             pm: "ProcMesh",
@@ -309,14 +321,15 @@ async def task(
         ) -> HyProcMesh:
             hy_proc_mesh = await hy_proc_mesh_task
 
-            pm._logging_mesh_client = await LoggingMeshClient.spawn(
-                proc_mesh=hy_proc_mesh
-            )
-            pm._logging_mesh_client.set_mode(
-                stream_to_client=True,
-                aggregate_window_sec=3,
-                level=logging.INFO,
-            )
+            if fork_processes:
+                pm._logging_mesh_client = await LoggingMeshClient.spawn(
+                    proc_mesh=hy_proc_mesh
+                )
+                pm._logging_mesh_client.set_mode(
+                    stream_to_client=True,
+                    aggregate_window_sec=3,
+                    level=logging.INFO,
+                )
 
             if setup_actor is not None:
                 await setup_actor.setup.call()
@@ -483,6 +496,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

@@ -1101,3 +1101,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)