[monarch/ext] fold monarch_hyperactor into actor_extension

Now we have achieved a good end-state for the actor native layer.
- One rust extension (`actor_extension`) for defining bindings from actor stuff -> Python.
- Exposed names in the extension directly map to the names in the `actor_extension` module hierarchy.

Differential Revision: [D77949213](https://our.internmc.facebook.com/intern/diff/D77949213/)

ghstack-source-id: 294914658
Pull Request resolved: #465

Author: suo

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/bin/process_allocator/common.rs renamed to actor_extension/bin/process_allocator/common.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/bin/process_allocator/main.rs renamed to actor_extension/bin/process_allocator/main.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/actor.rs renamed to actor_extension/src/actor.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/actor_mesh.rs renamed to actor_extension/src/actor_mesh.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/alloc.rs renamed to actor_extension/src/alloc.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,
 }

monarch_hyperactor/src/bootstrap.rs renamed to actor_extension/src/bootstrap.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",