[monarch] add actor_extension

Pull Request resolved: #440

Per the proposed structure, we will have a separate native extension for the actor api and the tensor_api.

This is the starting point for the actor extension.

Eventually the following binding-related crates will either be merged in or call from here:
- hyperactor_extension
- hyperactor_meta
- monarch_hyperactor
- monarch_meta_extension
ghstack-source-id: 294788182
@exported-using-ghexport

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

Author: suo

Cargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 resolver = "2"
 members = [
+    "actor_extension",
     "controller",
     "hyper",
     "hyperactor",

actor_extension/Cargo.toml

@@ -0,0 +1,26 @@
+# @generated by autocargo from //monarch/actor_extension:actor_extension-lib
+
+[package]
+name = "_extension"
+version = "0.0.0"
+authors = ["Meta"]
+edition = "2021"
+license = "BSD-3-Clause"
+
+[lib]
+test = false
+doctest = false
+crate-type = ["cdylib"]
+
+[dependencies]
+anyhow = "1.0.98"
+bincode = "1.3.3"
+hyperactor = { version = "0.0.0", path = "../hyperactor" }
+hyperactor_extension = { version = "0.0.0", path = "../hyperactor_extension" }
+hyperactor_mesh = { version = "0.0.0", path = "../hyperactor_mesh" }
+libc = "0.2.139"
+monarch_hyperactor = { version = "0.0.0", path = "../monarch_hyperactor" }
+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"] }
+tokio = { version = "1.45.0", features = ["full", "test-util", "tracing"] }

monarch_extension/src/blocking.rs renamed to actor_extension/src/blocking.rs

@@ -31,10 +31,7 @@ pub fn blocking_function() {
 /// Register Python bindings for the blocking module.
 pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
     let f = wrap_pyfunction!(blocking_function, module)?;
-    f.setattr(
-        "__module__",
-        "monarch._rust_bindings.monarch_extension.blocking",
-    )?;
+    f.setattr("__module__", "monarch._src.actor._extension.blocking")?;
     module.add_function(f)?;
     Ok(())
 }

monarch_extension/src/code_sync.rs renamed to actor_extension/src/code_sync.rs

@@ -31,7 +31,7 @@ use serde::Serialize;
 #[pyclass(
     frozen,
     name = "WorkspaceLocation",
-    module = "monarch._rust_bindings.monarch_extension.code_sync"
+    module = "monarch._src.actor._extension.code_sync"
 )]
 #[derive(Clone, Debug, Serialize, Deserialize)]
 enum PyWorkspaceLocation {
@@ -75,7 +75,7 @@ impl PyWorkspaceLocation {
 #[pyclass(
     frozen,
     name = "RsyncMeshClient",
-    module = "monarch._rust_bindings.monarch_extension.code_sync"
+    module = "monarch._src.actor._extension.code_sync"
 )]
 pub struct RsyncMeshClient {
     actor_mesh: SharedCell<RootActorMesh<'static, rsync::RsyncActor>>,

actor_extension/src/lib.rs

@@ -0,0 +1,117 @@
+/*
+ * 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.
+ */
+
+//! This module is used to expose Rust bindings for code supporting the
+//! `monarch.actor` module.
+//!
+//! It is imported by `monarch` as `monarch._src.actor._extension`.
+use pyo3::prelude::*;
+
+mod blocking;
+mod code_sync;
+mod panic;
+
+fn get_or_add_new_module<'py>(
+    module: &Bound<'py, PyModule>,
+    module_name: &str,
+) -> PyResult<Bound<'py, pyo3::types::PyModule>> {
+    let mut current_module = module.clone();
+    let mut parts = Vec::new();
+    for part in module_name.split(".") {
+        parts.push(part);
+        let submodule = current_module.getattr(part).ok();
+        if let Some(submodule) = submodule {
+            current_module = submodule.extract()?;
+        } else {
+            let new_module = PyModule::new(current_module.py(), part)?;
+            current_module.add_submodule(&new_module)?;
+            current_module
+                .py()
+                .import("sys")?
+                .getattr("modules")?
+                .set_item(
+                    format!("monarch._src.actor._extension.{}", parts.join(".")),
+                    new_module.clone(),
+                )?;
+            current_module = new_module;
+        }
+    }
+    Ok(current_module)
+}
+
+#[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",
+    )?)?;
+
+    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",
+    )?)?;
+    hyperactor_extension::alloc::register_python_bindings(&get_or_add_new_module(
+        module,
+        "hyperactor_extension.alloc",
+    )?)?;
+    hyperactor_extension::telemetry::register_python_bindings(&get_or_add_new_module(
+        module,
+        "hyperactor_extension.telemetry",
+    )?)?;
+    crate::panic::register_python_bindings(&get_or_add_new_module(module, "panic")?)?;
+
+    crate::blocking::register_python_bindings(&get_or_add_new_module(module, "blocking")?)?;
+    Ok(())
+}

monarch_extension/src/panic.rs renamed to actor_extension/src/panic.rs

@@ -18,10 +18,7 @@ pub fn panicking_function() {
 /// Register Python bindings for the panic module.
 pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
     let f = wrap_pyfunction!(panicking_function, module)?;
-    f.setattr(
-        "__module__",
-        "monarch._rust_bindings.monarch_extension.panic",
-    )?;
+    f.setattr("__module__", "monarch._src.actor._extension.panic")?;
     module.add_function(f)?;
     Ok(())
 }

hyperactor_extension/src/alloc.rs

@@ -29,7 +29,7 @@ use pyo3::types::PyDict;
 /// It ensures that the Alloc is only used once (i.e. moved) in rust.
 #[pyclass(
     name = "Alloc",
-    module = "monarch._rust_bindings.hyperactor_extension.alloc"
+    module = "monarch._src.actor._extension.hyperactor_extension.alloc"
 )]
 pub struct PyAlloc {
     pub inner: Arc<Mutex<Option<PyAllocWrapper>>>,
@@ -91,7 +91,7 @@ impl Alloc for PyAllocWrapper {
 
 #[pyclass(
     name = "AllocConstraints",
-    module = "monarch._rust_bindings.hyperactor_extension.alloc"
+    module = "monarch._src.actor._extension.hyperactor_extension.alloc"
 )]
 pub struct PyAllocConstraints {
     inner: AllocConstraints,
@@ -113,7 +113,7 @@ impl PyAllocConstraints {
 
 #[pyclass(
     name = "AllocSpec",
-    module = "monarch._rust_bindings.hyperactor_extension.alloc"
+    module = "monarch._src.actor._extension.hyperactor_extension.alloc"
 )]
 pub struct PyAllocSpec {
     pub inner: AllocSpec,

hyperactor_extension/src/telemetry.rs

@@ -113,7 +113,7 @@ pub fn use_sim_clock() -> PyResult<()> {
 #[pyclass(
     unsendable,
     subclass,
-    module = "monarch._rust_bindings.hyperactor_extension.telemetry"
+    module = "monarch._src.actor._extension.hyperactor_extension.telemetry"
 )]
 struct PySpan {
     span: tracing::span::EnteredSpan,
@@ -142,43 +142,43 @@ pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
     let f = wrap_pyfunction!(forward_to_tracing, module)?;
     f.setattr(
         "__module__",
-        "monarch._rust_bindings.hyperactor_extension.telemetry",
+        "monarch._src.actor._extension.hyperactor_extension.telemetry",
     )?;
     module.add_function(f)?;
 
     // Register the span-related functions
     let enter_span_fn = wrap_pyfunction!(enter_span, module)?;
     enter_span_fn.setattr(
         "__module__",
-        "monarch._rust_bindings.hyperactor_extension.telemetry",
+        "monarch._src.actor._extension.hyperactor_extension.telemetry",
     )?;
     module.add_function(enter_span_fn)?;
 
     let exit_span_fn = wrap_pyfunction!(exit_span, module)?;
     exit_span_fn.setattr(
         "__module__",
-        "monarch._rust_bindings.hyperactor_extension.telemetry",
+        "monarch._src.actor._extension.hyperactor_extension.telemetry",
     )?;
     module.add_function(exit_span_fn)?;
 
     let get_current_span_id_fn = wrap_pyfunction!(get_current_span_id, module)?;
     get_current_span_id_fn.setattr(
         "__module__",
-        "monarch._rust_bindings.hyperactor_extension.telemetry",
+        "monarch._src.actor._extension.hyperactor_extension.telemetry",
     )?;
     module.add_function(get_current_span_id_fn)?;
 
     let use_real_clock_fn = wrap_pyfunction!(use_real_clock, module)?;
     use_real_clock_fn.setattr(
         "__module__",
-        "monarch._rust_bindings.hyperactor_extension.telemetry",
+        "monarch._src.actor._extension.hyperactor_extension.telemetry",
     )?;
     module.add_function(use_real_clock_fn)?;
 
     let use_sim_clock_fn = wrap_pyfunction!(use_sim_clock, module)?;
     use_sim_clock_fn.setattr(
         "__module__",
-        "monarch._rust_bindings.hyperactor_extension.telemetry",
+        "monarch._src.actor._extension.hyperactor_extension.telemetry",
     )?;
     module.add_function(use_sim_clock_fn)?;
 

monarch_extension/Cargo.toml

@@ -16,14 +16,11 @@ 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"] }
 controller = { version = "0.0.0", path = "../controller", optional = true }
 hyperactor = { version = "0.0.0", path = "../hyperactor" }
-hyperactor_extension = { version = "0.0.0", path = "../hyperactor_extension" }
 hyperactor_mesh = { version = "0.0.0", path = "../hyperactor_mesh" }
 hyperactor_multiprocess = { version = "0.0.0", path = "../hyperactor_multiprocess" }
-libc = "0.2.139"
 monarch_hyperactor = { version = "0.0.0", path = "../monarch_hyperactor" }
 monarch_messages = { version = "0.0.0", path = "../monarch_messages", optional = true }
 monarch_simulator_lib = { version = "0.0.0", path = "../monarch_simulator", optional = true }
@@ -33,7 +30,6 @@ nccl-sys = { path = "../nccl-sys", optional = true }
 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"] }
 tokio = { version = "1.45.0", features = ["full", "test-util", "tracing"] }
 torch-sys = { version = "0.0.0", path = "../torch-sys", optional = true }
 torch-sys-cuda = { version = "0.0.0", path = "../torch-sys-cuda", optional = true }