local_proc_mesh with sim channels

Summary:
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

Rollback Plan:

Differential Revision: D77941640

Author: thomasywang

hyperactor_mesh/src/alloc/local.rs

@@ -58,6 +58,12 @@ impl Allocator for LocalAllocator {
     }
 }
 
+impl LocalAllocator {
+    pub async fn allocate_sim(&mut self, spec: AllocSpec) -> Result<LocalAlloc, AllocatorError> {
+        Ok(LocalAlloc::new_sim(spec))
+    }
+}
+
 struct LocalProc {
     proc: Proc,
     addr: ChannelAddr,
@@ -75,10 +81,22 @@ 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)
+    }
+
+    fn new_sim(spec: AllocSpec) -> Self {
+        Self::new_with_transport(
+            spec,
+            ChannelTransport::Sim(Box::new(ChannelTransport::Unix)),
+        )
+    }
+
+    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() {
@@ -94,6 +112,7 @@ impl LocalAlloc {
             todo_rx,
             stopped: false,
             failed: false,
+            transport,
         }
     }
 
@@ -249,7 +268,7 @@ impl Alloc for LocalAlloc {
     }
 
     fn transport(&self) -> ChannelTransport {
-        ChannelTransport::Local
+        self.transport.clone()
     }
 
     async fn stop(&mut self) -> Result<(), AllocatorError> {

monarch_hyperactor/src/alloc.rs

@@ -81,6 +81,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 {
+            LocalAllocator
+                .allocate_sim(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 {
+            LocalAllocator
+                .allocate_sim(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",
@@ -330,6 +377,7 @@ impl PyRemoteAllocator {
 pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
     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

@@ -113,6 +113,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"),
 }
 
@@ -181,6 +183,8 @@ def __getattr__(name):
     "timer",
     "ProcessAllocator",
     "LocalAllocator",
+    "SimAllocator",
+    "ActorFuture",
     "builtins",
 ]
 assert sorted(__all__) == sorted(_public_api)

python/monarch/_rust_bindings/monarch_hyperactor/alloc.pyi

@@ -69,6 +69,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

@@ -19,6 +19,7 @@
     LocalAllocatorBase,
     ProcessAllocatorBase,
     RemoteAllocatorBase,
+    SimAllocatorBase,
 )
 
 from monarch._src.actor.future import Future
@@ -72,6 +73,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

@@ -34,7 +34,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
 
@@ -287,29 +287,35 @@ def __del__(self) -> None:
 
 
 async def local_proc_mesh_nonblocking(
-    *, gpus: Optional[int] = None, hosts: int = 1
+    *, gpus: Optional[int] = None, hosts: int = 1, simulated: bool = False
 ) -> ProcMesh:
     if gpus is None:
         gpus = _local_device_count()
     spec = AllocSpec(AllocConstraints(), gpus=gpus, hosts=hosts)
-    allocator = LocalAllocator()
+    allocator = LocalAllocator() if not simulated else SimAllocator()
     alloc = await allocator.allocate(spec)
     return await ProcMesh.from_alloc(alloc)
 
 
-def local_proc_mesh_blocking(*, gpus: Optional[int] = None, hosts: int = 1) -> ProcMesh:
+def local_proc_mesh_blocking(
+    *, gpus: Optional[int] = None, hosts: int = 1, simulated: bool = False
+) -> ProcMesh:
     if gpus is None:
         gpus = _local_device_count()
     spec = AllocSpec(AllocConstraints(), gpus=gpus, hosts=hosts)
-    allocator = LocalAllocator()
+    allocator = LocalAllocator() if not simulated else SimAllocator()
     alloc = allocator.allocate(spec).get()
     return ProcMesh.from_alloc(alloc).get()
 
 
-def local_proc_mesh(*, gpus: Optional[int] = None, hosts: int = 1) -> Future[ProcMesh]:
+def local_proc_mesh(
+    *, gpus: Optional[int] = None, hosts: int = 1, simulated: bool = False
+) -> Future[ProcMesh]:
     return Future(
-        lambda: local_proc_mesh_nonblocking(gpus=gpus, hosts=hosts),
-        lambda: local_proc_mesh_blocking(gpus=gpus, hosts=hosts),
+        lambda: local_proc_mesh_nonblocking(
+            gpus=gpus, hosts=hosts, simulated=simulated
+        ),
+        lambda: local_proc_mesh_blocking(gpus=gpus, hosts=hosts, simulated=simulated),
     )