[monarch/ext] fold monarch_hyperactor into actor_extension

actor_extension/Cargo.toml

@@ -1,7 +1,7 @@
-# @generated by autocargo from //monarch/actor_extension:actor_extension-lib
+# @generated by autocargo from //monarch/actor_extension:[actor_extension-lib,process_allocator-oss]
 
 [package]
-name = "_extension"
+name = "actor_extension_lib"
 version = "0.0.0"
 authors = ["Meta"]
 edition = "2021"
@@ -14,14 +14,23 @@ crate-type = ["cdylib"]
 
 [dependencies]
 anyhow = "1.0.98"
+async-trait = "0.1.86"
 bincode = "1.3.3"
+clap = { version = "4.5.38", features = ["derive", "env", "string", "unicode", "wrap_help"] }
+erased-serde = "0.3.27"
+fbinit = { version = "0.2.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
 hyperactor = { version = "0.0.0", path = "../hyperactor" }
 hyperactor_mesh = { version = "0.0.0", path = "../hyperactor_mesh" }
+hyperactor_multiprocess = { version = "0.0.0", path = "../hyperactor_multiprocess" }
 hyperactor_telemetry = { version = "0.0.0", path = "../hyperactor_telemetry" }
+inventory = "0.3.8"
 libc = "0.2.139"
-monarch_hyperactor = { version = "0.0.0", path = "../monarch_hyperactor" }
+monarch_types = { version = "0.0.0", path = "../monarch_types" }
+ndslice = { version = "0.0.0", path = "../ndslice" }
 pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods"] }
 pyo3-async-runtimes = { version = "0.24", features = ["attributes", "tokio-runtime"] }
 serde = { version = "1.0.185", features = ["derive", "rc"] }
+serde_bytes = "0.11"
+thiserror = "2.0.12"
 tokio = { version = "1.45.0", features = ["full", "test-util", "tracing"] }
 tracing = { version = "0.1.41", features = ["attributes", "valuable"] }

monarch_hyperactor/src/actor.rs → actor_extension/src/actor.rs

@@ -52,10 +52,7 @@ use crate::proc::PySerialized;
 use crate::runtime::signal_safe_block_on;
 use crate::shape::PyShape;
 
-#[pyclass(
-    frozen,
-    module = "monarch._src.actor._extension.monarch_hyperactor.actor"
-)]
+#[pyclass(frozen, module = "monarch._src.actor._extension.actor")]
 #[derive(Serialize, Deserialize, Named)]
 pub struct PickledMessage {
     sender_actor_id: ActorId,
@@ -99,7 +96,7 @@ impl PickledMessage {
     }
 }
 
-#[pyclass(module = "monarch._src.actor._extension.monarch_hyperactor.actor")]
+#[pyclass(module = "monarch._src.actor._extension.actor")]
 pub struct PickledMessageClientActor {
     instance: Arc<Mutex<InstanceWrapper<PickledMessage>>>,
 }
@@ -174,10 +171,7 @@ impl PickledMessageClientActor {
     }
 }
 
-#[pyclass(
-    frozen,
-    module = "monarch._src.actor._extension.monarch_hyperactor.actor"
-)]
+#[pyclass(frozen, module = "monarch._src.actor._extension.actor")]
 #[derive(Default, Clone, Serialize, Deserialize, Named, PartialEq)]
 pub struct PythonMessage {
     pub(crate) method: String,
@@ -266,7 +260,7 @@ impl PythonMessage {
     }
 }
 
-#[pyclass(module = "monarch._src.actor._extension.monarch_hyperactor.actor")]
+#[pyclass(module = "monarch._src.actor._extension.actor")]
 pub(super) struct PythonActorHandle {
     pub(super) inner: ActorHandle<PythonActor>,
 }
@@ -391,7 +385,7 @@ fn create_task_locals() -> pyo3_async_runtimes::TaskLocals {
 // in Python, can catch the PanicException and notify the Rust awaiter manually.
 // In this way we can guarantee that the awaiter will complete even if the
 // `PyTaskCompleter` callback explodes.
-#[pyclass(module = "monarch._src.actor._extension.monarch_hyperactor.actor")]
+#[pyclass(module = "monarch._src.actor._extension.actor")]
 struct PanicFlag {
     sender: Option<tokio::sync::oneshot::Sender<PyObject>>,
 }

monarch_hyperactor/src/actor_mesh.rs → actor_extension/src/actor_mesh.rs

@@ -26,7 +26,7 @@ use crate::shape::PyShape;
 
 #[pyclass(
     name = "PythonActorMesh",
-    module = "monarch._src.actor._extension.monarch_hyperactor.actor_mesh"
+    module = "monarch._src.actor._extension.actor_mesh"
 )]
 pub struct PythonActorMesh {
     pub(super) inner: SharedCell<RootActorMesh<'static, PythonActor>>,

monarch_hyperactor/src/alloc.rs → actor_extension/src/alloc.rs

@@ -41,7 +41,7 @@ use crate::runtime::signal_safe_block_on;
 
 #[pyclass(
     name = "LocalAllocatorBase",
-    module = "monarch._src.actor._extension.monarch_hyperactor.alloc",
+    module = "monarch._src.actor._extension.alloc",
     subclass
 )]
 pub struct PyLocalAllocator;
@@ -88,7 +88,7 @@ impl PyLocalAllocator {
 
 #[pyclass(
     name = "ProcessAllocatorBase",
-    module = "monarch._src.actor._extension.monarch_hyperactor.alloc",
+    module = "monarch._src.actor._extension.alloc",
     subclass
 )]
 pub struct PyProcessAllocator {
@@ -156,7 +156,7 @@ impl PyProcessAllocator {
 /// Basically follows https://pyo3.rs/v0.25.0/trait-bounds.html.
 /// The Python subclass should implement `def initialize_alloc(self) -> list[str]`.
 pub struct PyRemoteProcessAllocInitializer {
-    // instance of a Python subclass of `monarch._src.actor._extension.monarch_hyperactor.alloc.RemoteProcessAllocInitializer`.
+    // instance of a Python subclass of `monarch._src.actor._extension.alloc.RemoteProcessAllocInitializer`.
     py_inner: Py<PyAny>,
 
     // allocation constraints passed onto the allocator's allocate call and passed along to python initializer.
@@ -234,7 +234,7 @@ impl RemoteProcessAllocInitializer for PyRemoteProcessAllocInitializer {
 
 #[pyclass(
     name = "RemoteAllocatorBase",
-    module = "monarch._src.actor._extension.monarch_hyperactor.alloc",
+    module = "monarch._src.actor._extension.alloc",
     subclass
 )]
 pub struct PyRemoteAllocator {
@@ -335,10 +335,7 @@ impl PyRemoteAllocator {
 /// A python class that wraps a Rust Alloc trait object. It represents what
 /// is shown on the python side. Internals are not exposed.
 /// It ensures that the Alloc is only used once (i.e. moved) in rust.
-#[pyclass(
-    name = "Alloc",
-    module = "monarch._src.actor._extension.monarch_hyperactor.alloc"
-)]
+#[pyclass(name = "Alloc", module = "monarch._src.actor._extension.alloc")]
 pub struct PyAlloc {
     pub inner: Arc<Mutex<Option<PyAllocWrapper>>>,
 }
@@ -399,7 +396,7 @@ impl Alloc for PyAllocWrapper {
 
 #[pyclass(
     name = "AllocConstraints",
-    module = "monarch._src.actor._extension.monarch_hyperactor.alloc"
+    module = "monarch._src.actor._extension.alloc"
 )]
 pub struct PyAllocConstraints {
     inner: AllocConstraints,
@@ -419,10 +416,7 @@ impl PyAllocConstraints {
     }
 }
 
-#[pyclass(
-    name = "AllocSpec",
-    module = "monarch._src.actor._extension.monarch_hyperactor.alloc"
-)]
+#[pyclass(name = "AllocSpec", module = "monarch._src.actor._extension.alloc")]
 pub struct PyAllocSpec {
     pub inner: AllocSpec,
 }

monarch_hyperactor/src/bootstrap.rs → actor_extension/src/bootstrap.rs

@@ -38,10 +38,7 @@ pub fn bootstrap_main(py: Python) -> PyResult<Bound<PyAny>> {
 
 pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
     let f = wrap_pyfunction!(bootstrap_main, hyperactor_mod)?;
-    f.setattr(
-        "__module__",
-        "monarch._src.actor._extension.monarch_hyperactor.bootstrap",
-    )?;
+    f.setattr("__module__", "monarch._src.actor._extension.bootstrap")?;
     hyperactor_mod.add_function(f)?;
 
     Ok(())

monarch_hyperactor/src/channel.rs → actor_extension/src/channel.rs

@@ -16,7 +16,7 @@ use pyo3::prelude::*;
 /// Python binding for [`hyperactor::channel::ChannelTransport`]
 #[pyclass(
     name = "ChannelTransport",
-    module = "monarch._src.actor._extension.monarch_hyperactor.channel",
+    module = "monarch._src.actor._extension.channel",
     eq
 )]
 #[derive(PartialEq, Clone, Copy, Debug)]
@@ -28,10 +28,7 @@ pub enum PyChannelTransport {
     // Sim(/*proxy address:*/ ChannelAddr), TODO kiuk@ add support
 }
 
-#[pyclass(
-    name = "ChannelAddr",
-    module = "monarch._src.actor._extension.monarch_hyperactor.channel"
-)]
+#[pyclass(name = "ChannelAddr", module = "monarch._src.actor._extension.channel")]
 pub struct PyChannelAddr {
     inner: ChannelAddr,
 }

actor_extension/src/code_sync.rs

@@ -16,9 +16,6 @@ use hyperactor_mesh::code_sync::WorkspaceLocation;
 use hyperactor_mesh::code_sync::rsync;
 use hyperactor_mesh::shape::Shape;
 use hyperactor_mesh::shared_cell::SharedCell;
-use monarch_hyperactor::proc_mesh::PyProcMesh;
-use monarch_hyperactor::runtime::signal_safe_block_on;
-use monarch_hyperactor::shape::PyShape;
 use pyo3::Bound;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::exceptions::PyValueError;
@@ -28,6 +25,10 @@ use pyo3::types::PyModule;
 use serde::Deserialize;
 use serde::Serialize;
 
+use crate::proc_mesh::PyProcMesh;
+use crate::runtime::signal_safe_block_on;
+use crate::shape::PyShape;
+
 #[pyclass(
     frozen,
     name = "WorkspaceLocation",

actor_extension/src/lib.rs

@@ -12,9 +12,22 @@
 //! It is imported by `monarch` as `monarch._src.actor._extension`.
 use pyo3::prelude::*;
 
+pub mod actor;
+mod actor_mesh;
+mod alloc;
 mod blocking;
+mod bootstrap;
+mod channel;
 mod code_sync;
+mod config;
+pub mod mailbox;
+pub mod ndslice;
 mod panic;
+pub mod proc;
+pub mod proc_mesh;
+pub mod runtime;
+mod selection;
+pub mod shape;
 mod telemetry;
 
 #[cfg(fbcode_build)]
@@ -51,65 +64,26 @@ fn get_or_add_new_module<'py>(
 #[pymodule]
 #[pyo3(name = "_extension")]
 pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> {
-    monarch_hyperactor::runtime::initialize(module.py())?;
-    let runtime = monarch_hyperactor::runtime::get_tokio_runtime();
-    ::hyperactor::initialize(runtime.handle().clone());
-
-    monarch_hyperactor::shape::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_hyperactor.shape",
-    )?)?;
-
-    monarch_hyperactor::selection::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_hyperactor.selection",
-    )?)?;
-    code_sync::register_python_bindings(&get_or_add_new_module(module, "code_sync")?)?;
-    monarch_hyperactor::bootstrap::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_hyperactor.bootstrap",
-    )?)?;
-
-    monarch_hyperactor::proc::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_hyperactor.proc",
-    )?)?;
-
-    monarch_hyperactor::actor::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_hyperactor.actor",
-    )?)?;
+    crate::runtime::initialize(module.py())?;
+    let runtime = crate::runtime::get_tokio_runtime();
 
-    monarch_hyperactor::mailbox::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_hyperactor.mailbox",
-    )?)?;
-
-    monarch_hyperactor::alloc::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_hyperactor.alloc",
-    )?)?;
-    monarch_hyperactor::channel::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_hyperactor.channel",
-    )?)?;
-    monarch_hyperactor::actor_mesh::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_hyperactor.actor_mesh",
-    )?)?;
-    monarch_hyperactor::proc_mesh::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_hyperactor.proc_mesh",
-    )?)?;
-
-    monarch_hyperactor::runtime::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_hyperactor.runtime",
-    )?)?;
-    telemetry::register_python_bindings(&get_or_add_new_module(module, "telemetry")?)?;
-    crate::panic::register_python_bindings(&get_or_add_new_module(module, "panic")?)?;
+    ::hyperactor::initialize(runtime.handle().clone());
 
+    crate::actor_mesh::register_python_bindings(&get_or_add_new_module(module, "actor_mesh")?)?;
+    crate::actor::register_python_bindings(&get_or_add_new_module(module, "actor")?)?;
+    crate::alloc::register_python_bindings(&get_or_add_new_module(module, "alloc")?)?;
     crate::blocking::register_python_bindings(&get_or_add_new_module(module, "blocking")?)?;
+    crate::bootstrap::register_python_bindings(&get_or_add_new_module(module, "bootstrap")?)?;
+    crate::channel::register_python_bindings(&get_or_add_new_module(module, "channel")?)?;
+    crate::code_sync::register_python_bindings(&get_or_add_new_module(module, "code_sync")?)?;
+    crate::mailbox::register_python_bindings(&get_or_add_new_module(module, "mailbox")?)?;
+    crate::panic::register_python_bindings(&get_or_add_new_module(module, "panic")?)?;
+    crate::proc_mesh::register_python_bindings(&get_or_add_new_module(module, "proc_mesh")?)?;
+    crate::proc::register_python_bindings(&get_or_add_new_module(module, "proc")?)?;
+    crate::runtime::register_python_bindings(&get_or_add_new_module(module, "runtime")?)?;
+    crate::selection::register_python_bindings(&get_or_add_new_module(module, "selection")?)?;
+    crate::shape::register_python_bindings(&get_or_add_new_module(module, "shape")?)?;
+    crate::telemetry::register_python_bindings(&get_or_add_new_module(module, "telemetry")?)?;
 
     #[cfg(fbcode_build)]
     crate::meta::register_python_bindings(&get_or_add_new_module(module, "meta")?)?;