[monarch] add actor_extension

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

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

ghstack-source-id: 294539856
Pull Request resolved: #440

Author: suo

actor_extension/Cargo.toml

@@ -0,0 +1,17 @@
+# @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]
+pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods"] }
+pyo3-async-runtimes = { git = "https://github.com/PyO3/pyo3-async-runtimes", rev = "c5a3746f110b4d246556b0f6c29f5f555919eee3", features = ["attributes", "tokio-runtime"] }

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.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.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.actor._extension.code_sync"
 )]
 pub struct RsyncMeshClient {
     actor_mesh: SharedCell<RootActorMesh<'static, rsync::RsyncActor>>,

actor_extension/src/lib.rs

@@ -0,0 +1,119 @@
+/*
+ * 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.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.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")?)?;
+
+    println!("hello I am imported :)");
+    Ok(())
+}

monarch_extension/src/panic.rs renamed to actor_extension/src/panic.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.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.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.actor._extension.hyperactor_extension.alloc"
 )]
 pub struct PyAllocSpec {
     pub inner: AllocSpec,

hyperactor_extension/src/alloc.rs

@@ -113,7 +113,7 @@ pub fn use_sim_clock() -> PyResult<()> {
 #[pyclass(
     unsendable,
     subclass,
-    module = "monarch._rust_bindings.hyperactor_extension.telemetry"
+    module = "monarch.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.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.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.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.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.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.actor._extension.hyperactor_extension.telemetry",
     )?;
     module.add_function(use_sim_clock_fn)?;
 

hyperactor_extension/src/telemetry.rs

@@ -10,7 +10,6 @@
 
 #[cfg(feature = "tensor_engine")]
 mod client;
-pub mod code_sync;
 #[cfg(feature = "tensor_engine")]
 mod controller;
 #[cfg(feature = "tensor_engine")]
@@ -24,8 +23,6 @@ mod simulator_client;
 #[cfg(feature = "tensor_engine")]
 mod tensor_worker;
 
-mod blocking;
-mod panic;
 use pyo3::prelude::*;
 
 #[pyfunction]
@@ -64,20 +61,6 @@ fn get_or_add_new_module<'py>(
 #[pymodule]
 #[pyo3(name = "_rust_bindings")]
 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",
-    )?)?;
-
     #[cfg(feature = "tensor_engine")]
     {
         client::register_python_bindings(&get_or_add_new_module(
@@ -125,69 +108,6 @@ pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> {
         module,
         "monarch_extension.simulation_tools",
     )?)?;
-    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",
-    )?)?;
-    code_sync::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_extension.code_sync",
-    )?)?;
-
-    crate::panic::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_extension.panic",
-    )?)?;
-
-    crate::blocking::register_python_bindings(&get_or_add_new_module(
-        module,
-        "monarch_extension.blocking",
-    )?)?;
 
     // Add feature detection function
     module.add_function(wrap_pyfunction!(has_tensor_engine, module)?)?;