Revert D77847588: add actor_extension

Differential Revision:
D77847588

Original commit changeset: ec1f4bfba9ec

Original Phabricator Diff: D77847588

fbshipit-source-id: 87f1e8849535e0f04917f230f3e3b3b1163e0ea9

Author: suo

Cargo.toml

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

actor_extension/Cargo.toml

@@ -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._src.actor._extension.hyperactor_extension.alloc"
+    module = "monarch._rust_bindings.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._src.actor._extension.hyperactor_extension.alloc"
+    module = "monarch._rust_bindings.hyperactor_extension.alloc"
 )]
 pub struct PyAllocConstraints {
     inner: AllocConstraints,
@@ -113,7 +113,7 @@ impl PyAllocConstraints {
 
 #[pyclass(
     name = "AllocSpec",
-    module = "monarch._src.actor._extension.hyperactor_extension.alloc"
+    module = "monarch._rust_bindings.hyperactor_extension.alloc"
 )]
 pub struct PyAllocSpec {
     pub inner: AllocSpec,

actor_extension/src/lib.rs

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

hyperactor_extension/src/alloc.rs

@@ -16,11 +16,14 @@ 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 }
@@ -30,6 +33,7 @@ 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 }

hyperactor_extension/src/telemetry.rs

@@ -31,7 +31,10 @@ 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._src.actor._extension.blocking")?;
+    f.setattr(
+        "__module__",
+        "monarch._rust_bindings.monarch_extension.blocking",
+    )?;
     module.add_function(f)?;
     Ok(())
 }

monarch_extension/Cargo.toml

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

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

@@ -10,6 +10,7 @@
 
 #[cfg(feature = "tensor_engine")]
 mod client;
+pub mod code_sync;
 #[cfg(feature = "tensor_engine")]
 mod controller;
 #[cfg(feature = "tensor_engine")]
@@ -23,6 +24,8 @@ mod simulator_client;
 #[cfg(feature = "tensor_engine")]
 mod tensor_worker;
 
+mod blocking;
+mod panic;
 use pyo3::prelude::*;
 
 #[pyfunction]
@@ -61,6 +64,20 @@ 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(
@@ -108,6 +125,69 @@ 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)?)?;