local_proc_mesh with sim channels (#475)

Summary:
Pull Request resolved: #475

We want to instantiate local_proc_meshes that use ChannelTransport::SIm instead of ChannelTranport::Local so that the simulator can intercept and control the delivery of messages

To preserve to `allocate()` interface so that we can reuse existing test generation macros we will create a wrapper class for this around `LocalAlloc`

Reviewed By: colin2328

Differential Revision: D77941640

fbshipit-source-id: 0be3a11b408d6a2b838de0508f1d4501ccb5afbc

Author: thomasywang

hyperactor_mesh/src/actor_mesh.rs

@@ -992,4 +992,10 @@ mod tests {
             buck_resources::get("monarch/hyperactor_mesh/bootstrap").unwrap()
         )));
     }
+
+    mod sim {
+        use crate::alloc::sim::SimAllocator;
+
+        actor_mesh_test_suite!(SimAllocator::new_and_start_simnet());
+    }
 }

hyperactor_mesh/src/alloc.rs

@@ -13,6 +13,7 @@ pub mod local;
 pub(crate) mod logtailer;
 pub mod process;
 pub mod remoteprocess;
+pub mod sim;
 
 use std::collections::HashMap;
 use std::fmt;

hyperactor_mesh/src/alloc/local.rs

@@ -76,10 +76,15 @@ pub struct LocalAlloc {
     todo_rx: mpsc::UnboundedReceiver<Action>,
     stopped: bool,
     failed: bool,
+    transport: ChannelTransport,
 }
 
 impl LocalAlloc {
     fn new(spec: AllocSpec) -> Self {
+        Self::new_with_transport(spec, ChannelTransport::Local)
+    }
+
+    pub(crate) fn new_with_transport(spec: AllocSpec, transport: ChannelTransport) -> Self {
         let name = ShortUuid::generate();
         let (todo_tx, todo_rx) = mpsc::unbounded_channel();
         for rank in 0..spec.shape.slice().len() {
@@ -95,6 +100,7 @@ impl LocalAlloc {
             todo_rx,
             stopped: false,
             failed: false,
+            transport,
         }
     }
 
@@ -124,7 +130,7 @@ impl LocalAlloc {
         &self.name
     }
 
-    fn size(&self) -> usize {
+    pub(crate) fn size(&self) -> usize {
         self.spec.shape.slice().len()
     }
 }
@@ -250,7 +256,7 @@ impl Alloc for LocalAlloc {
     }
 
     fn transport(&self) -> ChannelTransport {
-        ChannelTransport::Local
+        self.transport.clone()
     }
 
     async fn log_source(&self) -> Result<LogSource, AllocatorError> {

hyperactor_mesh/src/alloc/sim.rs

@@ -0,0 +1,127 @@
+/*
+ * 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.
+ */
+
+//! Support for allocating procs in the local process with simulated channels.
+
+#![allow(dead_code)] // until it is used outside of testing
+
+use async_trait::async_trait;
+use hyperactor::WorldId;
+use hyperactor::channel::ChannelAddr;
+use hyperactor::channel::ChannelTransport;
+use hyperactor::mailbox::MailboxServerHandle;
+use hyperactor::proc::Proc;
+use ndslice::Shape;
+
+use super::ProcStopReason;
+use crate::alloc::Alloc;
+use crate::alloc::AllocSpec;
+use crate::alloc::Allocator;
+use crate::alloc::AllocatorError;
+use crate::alloc::LocalAlloc;
+use crate::alloc::ProcState;
+use crate::log_source::LogSource;
+use crate::shortuuid::ShortUuid;
+
+/// An allocator that runs procs in the local process with network traffic going through simulated channels.
+/// Other than transport, the underlying implementation is an inner LocalAlloc.
+pub struct SimAllocator;
+
+#[async_trait]
+impl Allocator for SimAllocator {
+    type Alloc = SimAlloc;
+
+    async fn allocate(&mut self, spec: AllocSpec) -> Result<Self::Alloc, AllocatorError> {
+        Ok(SimAlloc::new(spec))
+    }
+}
+
+impl SimAllocator {
+    #[cfg(test)]
+    pub(crate) fn new_and_start_simnet() -> Self {
+        hyperactor::simnet::start();
+        Self
+    }
+}
+
+struct SimProc {
+    proc: Proc,
+    addr: ChannelAddr,
+    handle: MailboxServerHandle,
+}
+
+/// A simulated allocation. It is a collection of procs that are running in the local process.
+pub struct SimAlloc {
+    inner: LocalAlloc,
+}
+
+impl SimAlloc {
+    fn new(spec: AllocSpec) -> Self {
+        Self {
+            inner: LocalAlloc::new_with_transport(
+                spec,
+                ChannelTransport::Sim(Box::new(ChannelTransport::Unix)),
+            ),
+        }
+    }
+    /// A chaos monkey that can be used to stop procs at random.
+    pub(crate) fn chaos_monkey(&self) -> impl Fn(usize, ProcStopReason) + 'static {
+        self.inner.chaos_monkey()
+    }
+
+    /// A function to shut down the alloc for testing purposes.
+    pub(crate) fn stopper(&self) -> impl Fn() + 'static {
+        self.inner.stopper()
+    }
+
+    pub(crate) fn name(&self) -> &ShortUuid {
+        self.inner.name()
+    }
+
+    fn size(&self) -> usize {
+        self.inner.size()
+    }
+}
+
+#[async_trait]
+impl Alloc for SimAlloc {
+    async fn next(&mut self) -> Option<ProcState> {
+        self.inner.next().await
+    }
+
+    fn shape(&self) -> &Shape {
+        self.inner.shape()
+    }
+
+    fn world_id(&self) -> &WorldId {
+        self.inner.world_id()
+    }
+
+    fn transport(&self) -> ChannelTransport {
+        self.inner.transport()
+    }
+
+    async fn stop(&mut self) -> Result<(), AllocatorError> {
+        self.inner.stop().await
+    }
+
+    async fn log_source(&self) -> Result<LogSource, AllocatorError> {
+        self.inner.log_source().await
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_allocator_basic() {
+        hyperactor::simnet::start();
+        crate::alloc::testing::test_allocator_basic(SimAllocator).await;
+    }
+}

monarch_hyperactor/src/alloc.rs

@@ -28,6 +28,7 @@ use hyperactor_mesh::alloc::ProcessAllocator;
 use hyperactor_mesh::alloc::remoteprocess::RemoteProcessAlloc;
 use hyperactor_mesh::alloc::remoteprocess::RemoteProcessAllocHost;
 use hyperactor_mesh::alloc::remoteprocess::RemoteProcessAllocInitializer;
+use hyperactor_mesh::alloc::sim::SimAllocator;
 use hyperactor_mesh::log_source::LogSource;
 use hyperactor_mesh::shape::Shape;
 use ndslice::Slice;
@@ -227,6 +228,53 @@ impl PyLocalAllocator {
     }
 }
 
+#[pyclass(
+    name = "SimAllocatorBase",
+    module = "monarch._rust_bindings.monarch_hyperactor.alloc",
+    subclass
+)]
+pub struct PySimAllocator;
+
+#[pymethods]
+impl PySimAllocator {
+    #[new]
+    fn new() -> Self {
+        PySimAllocator {}
+    }
+
+    fn allocate_nonblocking<'py>(
+        &self,
+        py: Python<'py>,
+        spec: &PyAllocSpec,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        // We could use Bound here, and acquire the GIL inside of `future_into_py`, but
+        // it is rather awkward with the current APIs, and we can anyway support Arc/Mutex
+        // pretty easily.
+        let spec = spec.inner.clone();
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            SimAllocator
+                .allocate(spec)
+                .await
+                .map(|inner| PyAlloc::new(Box::new(inner)))
+                .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))
+        })
+    }
+
+    fn allocate_blocking<'py>(&self, py: Python<'py>, spec: &PyAllocSpec) -> PyResult<PyAlloc> {
+        // We could use Bound here, and acquire the GIL inside of
+        // `signal_safe_block_on`, but it is rather awkward with the current
+        // APIs, and we can anyway support Arc/Mutex pretty easily.
+        let spec = spec.inner.clone();
+        signal_safe_block_on(py, async move {
+            SimAllocator
+                .allocate(spec)
+                .await
+                .map(|inner| PyAlloc::new(Box::new(inner)))
+                .map_err(|e| PyRuntimeError::new_err(format!("{:?}", e)))
+        })?
+    }
+}
+
 #[pyclass(
     name = "ProcessAllocatorBase",
     module = "monarch._rust_bindings.monarch_hyperactor.alloc",
@@ -479,6 +527,7 @@ pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResul
     hyperactor_mod.add_class::<PyAllocSpec>()?;
     hyperactor_mod.add_class::<PyProcessAllocator>()?;
     hyperactor_mod.add_class::<PyLocalAllocator>()?;
+    hyperactor_mod.add_class::<PySimAllocator>()?;
     hyperactor_mod.add_class::<PyRemoteAllocator>()?;
 
     Ok(())

python/monarch/__init__.py

@@ -116,6 +116,8 @@
     "timer": ("monarch.timer", "timer"),
     "ProcessAllocator": ("monarch._src.actor.allocator", "ProcessAllocator"),
     "LocalAllocator": ("monarch._src.actor.allocator", "LocalAllocator"),
+    "SimAllocator": ("monarch._src_actor.allocator", "SimAllocator"),
+    "ActorFuture": ("monarch.future", "ActorFuture"),
     "builtins": ("monarch.builtins", "builtins"),
 }
 
@@ -184,6 +186,8 @@ def __getattr__(name):
     "timer",
     "ProcessAllocator",
     "LocalAllocator",
+    "SimAllocator",
+    "ActorFuture",
     "builtins",
 ]
 assert sorted(__all__) == sorted(_public_api)

python/monarch/_rust_bindings/monarch_hyperactor/alloc.pyi

@@ -100,6 +100,26 @@ class LocalAllocatorBase:
         """
         ...
 
+class SimAllocatorBase:
+    async def allocate_nonblocking(self, spec: AllocSpec) -> Alloc:
+        """
+        Allocate a process according to the provided spec.
+
+        Arguments:
+        - `spec`: The spec to allocate according to.
+        """
+        ...
+
+    def allocate_blocking(self, spec: AllocSpec) -> Alloc:
+        """
+        Allocate a process according to the provided spec, blocking until an
+        alloc is returned.
+
+        Arguments:
+        - `spec`: The spec to allocate according to.
+        """
+        ...
+
 class RemoteAllocatorBase:
     def __new__(
         cls,

python/monarch/_src/actor/allocator.py

@@ -16,6 +16,7 @@
     LocalAllocatorBase,
     ProcessAllocatorBase,
     RemoteAllocatorBase,
+    SimAllocatorBase,
 )
 
 from monarch._src.actor.future import Future
@@ -69,6 +70,28 @@ def allocate(self, spec: AllocSpec) -> Future[Alloc]:
         )
 
 
+@final
+class SimAllocator(SimAllocatorBase):
+    """
+    An allocator that allocates by spawning actors into the current process using simulated channels for transport
+    """
+
+    def allocate(self, spec: AllocSpec) -> Future[Alloc]:
+        """
+        Allocate a process according to the provided spec.
+
+        Arguments:
+        - `spec`: The spec to allocate according to.
+
+        Returns:
+        - A future that will be fulfilled when the requested allocation is fulfilled.
+        """
+        return Future(
+            lambda: self.allocate_nonblocking(spec),
+            lambda: self.allocate_blocking(spec),
+        )
+
+
 class RemoteAllocInitializer(abc.ABC):
     """Subclass-able Python interface for `hyperactor_mesh::alloc::remoteprocess:RemoteProcessAllocInitializer`.
 

python/monarch/_src/actor/proc_mesh.py

@@ -35,7 +35,7 @@
 )
 from monarch._rust_bindings.monarch_hyperactor.shape import Shape, Slice
 from monarch._src.actor.actor_mesh import _Actor, _ActorMeshRefImpl, Actor, ActorMeshRef
-from monarch._src.actor.allocator import LocalAllocator, ProcessAllocator
+from monarch._src.actor.allocator import LocalAllocator, ProcessAllocator, SimAllocator
 from monarch._src.actor.code_sync import RsyncMeshClient, WorkspaceLocation
 from monarch._src.actor.code_sync.auto_reload import AutoReloadActor
 from monarch._src.actor.debugger import (
@@ -329,6 +329,33 @@ def local_proc_mesh(*, gpus: Optional[int] = None, hosts: int = 1) -> Future[Pro
     )
 
 
+async def sim_proc_mesh_nonblocking(
+    *, gpus: Optional[int] = None, hosts: int = 1
+) -> ProcMesh:
+    if gpus is None:
+        gpus = _local_device_count()
+    spec = AllocSpec(AllocConstraints(), gpus=gpus, hosts=hosts)
+    allocator = SimAllocator()
+    alloc = await allocator.allocate(spec)
+    return await ProcMesh.from_alloc(alloc)
+
+
+def sim_proc_mesh_blocking(*, gpus: Optional[int] = None, hosts: int = 1) -> ProcMesh:
+    if gpus is None:
+        gpus = _local_device_count()
+    spec = AllocSpec(AllocConstraints(), gpus=gpus, hosts=hosts)
+    allocator = SimAllocator()
+    alloc = allocator.allocate(spec).get()
+    return ProcMesh.from_alloc(alloc).get()
+
+
+def sim_proc_mesh(*, gpus: Optional[int] = None, hosts: int = 1) -> Future[ProcMesh]:
+    return Future(
+        lambda: sim_proc_mesh_nonblocking(gpus=gpus, hosts=hosts),
+        lambda: sim_proc_mesh_blocking(gpus=gpus, hosts=hosts),
+    )
+
+
 _BOOTSTRAP_MAIN = "monarch._src.actor.bootstrap_main"