From 061364211600b77acd06cfab3e747c8d46386b2a Mon Sep 17 00:00:00 2001 From: suo Date: Sun, 6 Jul 2025 21:27:25 -0700 Subject: [PATCH] [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-poisoned] --- actor_extension/Cargo.toml | 17 +++ .../src/blocking.rs | 5 +- .../src/code_sync.rs | 4 +- actor_extension/src/lib.rs | 119 ++++++++++++++++++ .../src/panic.rs | 0 hyperactor_extension/src/alloc.rs | 6 +- hyperactor_extension/src/telemetry.rs | 14 +-- monarch_extension/src/lib.rs | 80 ------------ monarch_hyperactor/src/actor.rs | 10 +- monarch_hyperactor/src/actor_mesh.rs | 2 +- monarch_hyperactor/src/alloc.rs | 8 +- monarch_hyperactor/src/bootstrap.rs | 2 +- monarch_hyperactor/src/channel.rs | 4 +- monarch_hyperactor/src/mailbox.rs | 22 ++-- monarch_hyperactor/src/ndslice.rs | 2 +- monarch_hyperactor/src/proc.rs | 10 +- monarch_hyperactor/src/proc_mesh.rs | 6 +- monarch_hyperactor/src/runtime.rs | 2 +- monarch_hyperactor/src/selection.rs | 2 +- monarch_hyperactor/src/shape.rs | 4 +- python/monarch/__init__.py | 3 + .../monarch/_monarch/hyperactor/__init__.py | 10 +- python/monarch/_monarch/selection/__init__.py | 2 +- .../monarch_extension/client.pyi | 4 +- .../monarch_extension/controller.pyi | 4 +- .../monarch_extension/debugger.pyi | 3 +- .../monarch_extension/mesh_controller.pyi | 8 +- .../monarch_extension/tensor_worker.pyi | 4 +- python/monarch/_rust_bindings/old.pyi | 7 +- python/monarch/actor/_actor_mesh.py | 18 +-- python/monarch/actor/_allocator.py | 7 +- python/monarch/actor/_bootstrap_main.py | 2 +- python/monarch/actor/_code_sync/__init__.py | 2 +- python/monarch/actor/_debugger.py | 3 +- python/monarch/actor/_extension/__init__.pyi | 9 ++ .../_extension}/blocking.pyi | 0 .../_extension}/code_sync.pyi | 4 +- .../hyperactor_extension/alloc.pyi | 0 .../hyperactor_extension/telemetry.pyi | 0 .../_extension}/monarch_hyperactor/actor.pyi | 6 +- .../monarch_hyperactor/actor_mesh.pyi | 8 +- .../_extension}/monarch_hyperactor/alloc.pyi | 2 +- .../monarch_hyperactor/bootstrap.pyi | 0 .../monarch_hyperactor/channel.pyi | 0 .../monarch_hyperactor/mailbox.pyi | 6 +- .../_extension}/monarch_hyperactor/proc.pyi | 4 +- .../monarch_hyperactor/proc_mesh.pyi | 11 +- .../monarch_hyperactor/runtime.pyi | 0 .../monarch_hyperactor/selection.pyi | 0 .../_extension}/monarch_hyperactor/shape.pyi | 0 .../_extension}/panic.pyi | 0 python/monarch/actor/_pdb_wrapper.py | 2 +- python/monarch/actor/_proc_mesh.py | 20 +-- python/monarch/actor/_shape.py | 2 +- python/monarch/actor/_telemetry/__init__.py | 4 +- .../actor/_telemetry/rust_span_tracing.py | 2 +- python/monarch/common/invocation.py | 2 +- python/monarch/common/selection.py | 2 +- python/monarch/controller/controller.py | 2 +- python/monarch/controller/history.py | 2 +- .../controller/rust_backend/controller.py | 6 +- python/monarch/mesh_controller.py | 12 +- python/monarch/rdma.py | 4 +- python/monarch/rust_backend_mesh.py | 2 +- python/monarch/rust_local_mesh.py | 8 +- python/monarch/sim_mesh.py | 2 +- python/monarch/simulator/mock_controller.py | 2 +- python/monarch/simulator/simulator.py | 2 +- python/monarch/tools/mesh_spec.py | 2 +- python/tests/_monarch/test_actor.py | 2 +- python/tests/_monarch/test_client.py | 2 +- python/tests/_monarch/test_controller.py | 2 +- python/tests/_monarch/test_hyperactor.py | 12 +- python/tests/_monarch/test_mailbox.py | 10 +- python/tests/_monarch/test_ndslice.py | 4 +- python/tests/_monarch/test_worker.py | 2 +- python/tests/error_test_binary.py | 6 +- python/tests/sleep_binary.py | 2 +- python/tests/test_actor_error.py | 2 +- python/tests/test_alloc.py | 2 +- python/tests/test_allocator.py | 18 +-- python/tests/test_future.py | 2 +- 82 files changed, 329 insertions(+), 261 deletions(-) create mode 100644 actor_extension/Cargo.toml rename {monarch_extension => actor_extension}/src/blocking.rs (90%) rename {monarch_extension => actor_extension}/src/code_sync.rs (96%) create mode 100644 actor_extension/src/lib.rs rename {monarch_extension => actor_extension}/src/panic.rs (100%) create mode 100644 python/monarch/actor/_extension/__init__.pyi rename python/monarch/{_rust_bindings/monarch_extension => actor/_extension}/blocking.pyi (100%) rename python/monarch/{_rust_bindings/monarch_extension => actor/_extension}/code_sync.pyi (87%) rename python/monarch/{_rust_bindings => actor/_extension}/hyperactor_extension/alloc.pyi (100%) rename python/monarch/{_rust_bindings => actor/_extension}/hyperactor_extension/telemetry.pyi (100%) rename python/monarch/{_rust_bindings => actor/_extension}/monarch_hyperactor/actor.pyi (95%) rename python/monarch/{_rust_bindings => actor/_extension}/monarch_hyperactor/actor_mesh.pyi (77%) rename python/monarch/{_rust_bindings => actor/_extension}/monarch_hyperactor/alloc.pyi (97%) rename python/monarch/{_rust_bindings => actor/_extension}/monarch_hyperactor/bootstrap.pyi (100%) rename python/monarch/{_rust_bindings => actor/_extension}/monarch_hyperactor/channel.pyi (100%) rename python/monarch/{_rust_bindings => actor/_extension}/monarch_hyperactor/mailbox.pyi (96%) rename python/monarch/{_rust_bindings => actor/_extension}/monarch_hyperactor/proc.pyi (95%) rename python/monarch/{_rust_bindings => actor/_extension}/monarch_hyperactor/proc_mesh.pyi (88%) rename python/monarch/{_rust_bindings => actor/_extension}/monarch_hyperactor/runtime.pyi (100%) rename python/monarch/{_rust_bindings => actor/_extension}/monarch_hyperactor/selection.pyi (100%) rename python/monarch/{_rust_bindings => actor/_extension}/monarch_hyperactor/shape.pyi (100%) rename python/monarch/{_rust_bindings/monarch_extension => actor/_extension}/panic.pyi (100%) diff --git a/actor_extension/Cargo.toml b/actor_extension/Cargo.toml new file mode 100644 index 00000000..456c302f --- /dev/null +++ b/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"] } diff --git a/monarch_extension/src/blocking.rs b/actor_extension/src/blocking.rs similarity index 90% rename from monarch_extension/src/blocking.rs rename to actor_extension/src/blocking.rs index 9b106e91..b411884c 100644 --- a/monarch_extension/src/blocking.rs +++ b/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(()) } diff --git a/monarch_extension/src/code_sync.rs b/actor_extension/src/code_sync.rs similarity index 96% rename from monarch_extension/src/code_sync.rs rename to actor_extension/src/code_sync.rs index 19ec8f1d..32a736a4 100644 --- a/monarch_extension/src/code_sync.rs +++ b/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>, diff --git a/actor_extension/src/lib.rs b/actor_extension/src/lib.rs new file mode 100644 index 00000000..838e64bd --- /dev/null +++ b/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> { + 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(()) +} diff --git a/monarch_extension/src/panic.rs b/actor_extension/src/panic.rs similarity index 100% rename from monarch_extension/src/panic.rs rename to actor_extension/src/panic.rs diff --git a/hyperactor_extension/src/alloc.rs b/hyperactor_extension/src/alloc.rs index bcc7b53d..c076ec02 100644 --- a/hyperactor_extension/src/alloc.rs +++ b/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.actor._extension.hyperactor_extension.alloc" )] pub struct PyAlloc { pub inner: Arc>>, @@ -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, diff --git a/hyperactor_extension/src/telemetry.rs b/hyperactor_extension/src/telemetry.rs index 32f9b4dc..d8327601 100644 --- a/hyperactor_extension/src/telemetry.rs +++ b/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.actor._extension.hyperactor_extension.telemetry" )] struct PySpan { span: tracing::span::EnteredSpan, @@ -142,7 +142,7 @@ 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)?; @@ -150,35 +150,35 @@ pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> { 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)?; diff --git a/monarch_extension/src/lib.rs b/monarch_extension/src/lib.rs index 5ca058f6..a85f0dfe 100644 --- a/monarch_extension/src/lib.rs +++ b/monarch_extension/src/lib.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)?)?; diff --git a/monarch_hyperactor/src/actor.rs b/monarch_hyperactor/src/actor.rs index 2a7bdcb9..4183360e 100644 --- a/monarch_hyperactor/src/actor.rs +++ b/monarch_hyperactor/src/actor.rs @@ -52,7 +52,7 @@ use crate::proc::PySerialized; use crate::runtime::signal_safe_block_on; use crate::shape::PyShape; -#[pyclass(frozen, module = "monarch._rust_bindings.monarch_hyperactor.actor")] +#[pyclass(frozen, module = "monarch.actor._extension.monarch_hyperactor.actor")] #[derive(Serialize, Deserialize, Named)] pub struct PickledMessage { sender_actor_id: ActorId, @@ -96,7 +96,7 @@ impl PickledMessage { } } -#[pyclass(module = "monarch._rust_bindings.monarch_hyperactor.actor")] +#[pyclass(module = "monarch.actor._extension.monarch_hyperactor.actor")] pub struct PickledMessageClientActor { instance: Arc>>, } @@ -171,7 +171,7 @@ impl PickledMessageClientActor { } } -#[pyclass(frozen, module = "monarch._rust_bindings.monarch_hyperactor.actor")] +#[pyclass(frozen, module = "monarch.actor._extension.monarch_hyperactor.actor")] #[derive(Default, Clone, Serialize, Deserialize, Named, PartialEq)] pub struct PythonMessage { pub(crate) method: String, @@ -260,7 +260,7 @@ impl PythonMessage { } } -#[pyclass(module = "monarch._rust_bindings.monarch_hyperactor.actor")] +#[pyclass(module = "monarch.actor._extension.monarch_hyperactor.actor")] pub(super) struct PythonActorHandle { pub(super) inner: ActorHandle, } @@ -385,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._rust_bindings.monarch_hyperactor.actor")] +#[pyclass(module = "monarch.actor._extension.monarch_hyperactor.actor")] struct PanicFlag { sender: Option>, } diff --git a/monarch_hyperactor/src/actor_mesh.rs b/monarch_hyperactor/src/actor_mesh.rs index 88e8ca68..3113aae3 100644 --- a/monarch_hyperactor/src/actor_mesh.rs +++ b/monarch_hyperactor/src/actor_mesh.rs @@ -25,7 +25,7 @@ use crate::shape::PyShape; #[pyclass( name = "PythonActorMesh", - module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh" + module = "monarch.actor._extension.monarch_hyperactor.actor_mesh" )] pub struct PythonActorMesh { pub(super) inner: SharedCell>, diff --git a/monarch_hyperactor/src/alloc.rs b/monarch_hyperactor/src/alloc.rs index 1757a54f..f5f19c8c 100644 --- a/monarch_hyperactor/src/alloc.rs +++ b/monarch_hyperactor/src/alloc.rs @@ -36,7 +36,7 @@ use crate::runtime::signal_safe_block_on; #[pyclass( name = "LocalAllocatorBase", - module = "monarch._rust_bindings.monarch_hyperactor.alloc", + module = "monarch.actor._extension.monarch_hyperactor.alloc", subclass )] pub struct PyLocalAllocator; @@ -83,7 +83,7 @@ impl PyLocalAllocator { #[pyclass( name = "ProcessAllocatorBase", - module = "monarch._rust_bindings.monarch_hyperactor.alloc", + module = "monarch.actor._extension.monarch_hyperactor.alloc", subclass )] pub struct PyProcessAllocator { @@ -151,7 +151,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._rust_bindings.monarch_hyperactor.alloc.RemoteProcessAllocInitializer`. + // instance of a Python subclass of `monarch.actor._extension.monarch_hyperactor.alloc.RemoteProcessAllocInitializer`. py_inner: Py, // allocation constraints passed onto the allocator's allocate call and passed along to python initializer. @@ -229,7 +229,7 @@ impl RemoteProcessAllocInitializer for PyRemoteProcessAllocInitializer { #[pyclass( name = "RemoteAllocatorBase", - module = "monarch._rust_bindings.monarch_hyperactor.alloc", + module = "monarch.actor._extension.monarch_hyperactor.alloc", subclass )] pub struct PyRemoteAllocator { diff --git a/monarch_hyperactor/src/bootstrap.rs b/monarch_hyperactor/src/bootstrap.rs index 22f0c594..2dfa7fa1 100644 --- a/monarch_hyperactor/src/bootstrap.rs +++ b/monarch_hyperactor/src/bootstrap.rs @@ -40,7 +40,7 @@ pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResul let f = wrap_pyfunction!(bootstrap_main, hyperactor_mod)?; f.setattr( "__module__", - "monarch._rust_bindings.monarch_hyperactor.bootstrap", + "monarch.actor._extension.monarch_hyperactor.bootstrap", )?; hyperactor_mod.add_function(f)?; diff --git a/monarch_hyperactor/src/channel.rs b/monarch_hyperactor/src/channel.rs index ef7522a2..59925a52 100644 --- a/monarch_hyperactor/src/channel.rs +++ b/monarch_hyperactor/src/channel.rs @@ -16,7 +16,7 @@ use pyo3::prelude::*; /// Python binding for [`hyperactor::channel::ChannelTransport`] #[pyclass( name = "ChannelTransport", - module = "monarch._rust_bindings.monarch_hyperactor.channel", + module = "monarch.actor._extension.monarch_hyperactor.channel", eq )] #[derive(PartialEq, Clone, Copy, Debug)] @@ -30,7 +30,7 @@ pub enum PyChannelTransport { #[pyclass( name = "ChannelAddr", - module = "monarch._rust_bindings.monarch_hyperactor.channel" + module = "monarch.actor._extension.monarch_hyperactor.channel" )] pub struct PyChannelAddr { inner: ChannelAddr, diff --git a/monarch_hyperactor/src/mailbox.rs b/monarch_hyperactor/src/mailbox.rs index e281c220..59b14f5e 100644 --- a/monarch_hyperactor/src/mailbox.rs +++ b/monarch_hyperactor/src/mailbox.rs @@ -50,7 +50,7 @@ use crate::shape::PyShape; #[derive(Clone, Debug)] #[pyclass( name = "Mailbox", - module = "monarch._rust_bindings.monarch_hyperactor.mailbox" + module = "monarch.actor._extension.monarch_hyperactor.mailbox" )] pub(super) struct PyMailbox { pub(super) inner: Mailbox, @@ -188,7 +188,7 @@ impl PyMailbox { #[pyclass( frozen, name = "PortId", - module = "monarch._rust_bindings.monarch_hyperactor.mailbox" + module = "monarch.actor._extension.monarch_hyperactor.mailbox" )] #[derive(Clone)] pub struct PyPortId { @@ -270,7 +270,7 @@ impl std::fmt::Debug for PyPortId { #[derive(Clone, Debug)] #[pyclass( name = "PortHandle", - module = "monarch._rust_bindings.monarch_hyperactor.mailbox" + module = "monarch.actor._extension.monarch_hyperactor.mailbox" )] pub(super) struct PythonPortHandle { inner: PortHandle, @@ -295,7 +295,7 @@ impl PythonPortHandle { #[derive(Clone, Debug)] #[pyclass( name = "UndeliverablePortHandle", - module = "monarch._rust_bindings.monarch_hyperactor.mailbox" + module = "monarch.actor._extension.monarch_hyperactor.mailbox" )] pub(super) struct PythonUndeliverablePortHandle { inner: PortHandle>, @@ -311,7 +311,7 @@ impl PythonUndeliverablePortHandle { #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] #[pyclass( name = "PortRef", - module = "monarch._rust_bindings.monarch_hyperactor.mailbox" + module = "monarch.actor._extension.monarch_hyperactor.mailbox" )] pub(super) struct PythonPortRef { pub(crate) inner: PortRef, @@ -345,7 +345,7 @@ impl From> for PythonPortRef { #[derive(Debug)] #[pyclass( name = "PortReceiver", - module = "monarch._rust_bindings.monarch_hyperactor.mailbox" + module = "monarch.actor._extension.monarch_hyperactor.mailbox" )] pub(super) struct PythonPortReceiver { inner: Arc>>, @@ -374,7 +374,7 @@ impl PythonPortReceiver { #[derive(Debug)] #[pyclass( name = "UndeliverableMessageEnvelope", - module = "monarch._rust_bindings.monarch_hyperactor.mailbox" + module = "monarch.actor._extension.monarch_hyperactor.mailbox" )] pub(crate) struct PythonUndeliverableMessageEnvelope { #[allow(dead_code)] // At this time, field `inner` isn't read. @@ -384,7 +384,7 @@ pub(crate) struct PythonUndeliverableMessageEnvelope { #[derive(Debug)] #[pyclass( name = "UndeliverablePortReceiver", - module = "monarch._rust_bindings.monarch_hyperactor.mailbox" + module = "monarch.actor._extension.monarch_hyperactor.mailbox" )] pub(super) struct PythonUndeliverablePortReceiver { inner: Arc>>>, @@ -420,7 +420,7 @@ impl PythonUndeliverablePortReceiver { #[derive(Debug)] #[pyclass( name = "OncePortHandle", - module = "monarch._rust_bindings.monarch_hyperactor.mailbox" + module = "monarch.actor._extension.monarch_hyperactor.mailbox" )] pub(super) struct PythonOncePortHandle { inner: Option>, @@ -451,7 +451,7 @@ impl PythonOncePortHandle { #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] #[pyclass( name = "OncePortRef", - module = "monarch._rust_bindings.monarch_hyperactor.mailbox" + module = "monarch.actor._extension.monarch_hyperactor.mailbox" )] pub(crate) struct PythonOncePortRef { pub(crate) inner: Option>, @@ -492,7 +492,7 @@ impl From> for PythonOncePortRef { #[pyclass( name = "OncePortReceiver", - module = "monarch._rust_bindings.monarch_hyperactor.mailbox" + module = "monarch.actor._extension.monarch_hyperactor.mailbox" )] pub(super) struct PythonOncePortReceiver { inner: std::sync::Mutex>>, diff --git a/monarch_hyperactor/src/ndslice.rs b/monarch_hyperactor/src/ndslice.rs index 285dce7c..e8d63e3f 100644 --- a/monarch_hyperactor/src/ndslice.rs +++ b/monarch_hyperactor/src/ndslice.rs @@ -28,7 +28,7 @@ use pyo3::types::PyTuple; #[pyclass( name = "Slice", frozen, - module = "monarch._rust_bindings.monarch_hyperactor.shape" + module = "monarch.actor._extension.monarch_hyperactor.shape" )] #[derive(Clone)] pub struct PySlice { diff --git a/monarch_hyperactor/src/proc.rs b/monarch_hyperactor/src/proc.rs index 141c94fa..e9c511e0 100644 --- a/monarch_hyperactor/src/proc.rs +++ b/monarch_hyperactor/src/proc.rs @@ -14,7 +14,7 @@ /// wheel and the rest of the codebase especially while things are in flux. Plus we are also /// building everything in hyperactor_python into this wheel already (i.e. hyperactor deps). /// 2. In order to support autoreload in bento, potentially pickling in the future etc we need to -/// have a well defined module for these deps which needs to be monarch._rust_bindings.monarch_hyperactor.proc and +/// have a well defined module for these deps which needs to be monarch.actor._extension.monarch_hyperactor.proc and /// and making that the module of the classes in hyperactor python is weird. use std::collections::HashMap; use std::hash::DefaultHasher; @@ -71,7 +71,7 @@ use crate::runtime::signal_safe_block_on; #[derive(Clone, Debug)] #[pyclass( name = "Proc", - module = "monarch._rust_bindings.monarch_hyperactor.proc" + module = "monarch.actor._extension.monarch_hyperactor.proc" )] pub struct PyProc { pub(super) inner: Proc, @@ -255,7 +255,7 @@ pub fn init_proc( #[pyclass( frozen, name = "ActorId", - module = "monarch._rust_bindings.monarch_hyperactor.proc" + module = "monarch.actor._extension.monarch_hyperactor.proc" )] #[derive(Clone)] pub struct PyActorId { @@ -371,7 +371,7 @@ enum InstanceStatus { #[pyclass( frozen, name = "Serialized", - module = "monarch._rust_bindings.monarch_hyperactor.proc" + module = "monarch.actor._extension.monarch_hyperactor.proc" )] #[derive(Debug)] pub struct PySerialized { @@ -683,7 +683,7 @@ pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResul let f = wrap_pyfunction!(init_proc, hyperactor_mod)?; f.setattr( "__module__", - "monarch._rust_bindings.monarch_hyperactor.proc", + "monarch.actor._extension.monarch_hyperactor.proc", )?; hyperactor_mod.add_function(f)?; diff --git a/monarch_hyperactor/src/proc_mesh.rs b/monarch_hyperactor/src/proc_mesh.rs index c457e130..66a9b68a 100644 --- a/monarch_hyperactor/src/proc_mesh.rs +++ b/monarch_hyperactor/src/proc_mesh.rs @@ -108,7 +108,7 @@ impl TrackedProcMesh { #[pyclass( name = "ProcMesh", - module = "monarch._rust_bindings.monarch_hyperactor.proc_mesh" + module = "monarch.actor._extension.monarch_hyperactor.proc_mesh" )] pub struct PyProcMesh { pub inner: SharedCell, @@ -380,7 +380,7 @@ impl Drop for KeepaliveState { #[pyclass( name = "ProcMeshMonitor", - module = "monarch._rust_bindings.monarch_hyperactor.proc_mesh" + module = "monarch.actor._extension.monarch_hyperactor.proc_mesh" )] pub struct PyProcMeshMonitor { proc_events: SharedCell>, @@ -419,7 +419,7 @@ impl PyProcMeshMonitor { #[pyclass( name = "ProcEvent", - module = "monarch._rust_bindings.monarch_hyperactor.proc_mesh" + module = "monarch.actor._extension.monarch_hyperactor.proc_mesh" )] pub enum PyProcEvent { /// The proc of the given rank was stopped with the provided reason. diff --git a/monarch_hyperactor/src/runtime.rs b/monarch_hyperactor/src/runtime.rs index 9bfddb21..e80ace3b 100644 --- a/monarch_hyperactor/src/runtime.rs +++ b/monarch_hyperactor/src/runtime.rs @@ -126,7 +126,7 @@ pub fn register_python_bindings(runtime_mod: &Bound<'_, PyModule>) -> PyResult<( wrap_pyfunction!(sleep_indefinitely_for_unit_tests, runtime_mod.py())?; sleep_indefinitely_fn.setattr( "__module__", - "monarch._rust_bindings.monarch_hyperactor.runtime", + "monarch.actor._extension.monarch_hyperactor.runtime", )?; runtime_mod.add_function(sleep_indefinitely_fn)?; Ok(()) diff --git a/monarch_hyperactor/src/selection.rs b/monarch_hyperactor/src/selection.rs index 07d5f5bb..de2dae98 100644 --- a/monarch_hyperactor/src/selection.rs +++ b/monarch_hyperactor/src/selection.rs @@ -13,7 +13,7 @@ use pyo3::types::PyType; #[pyclass( name = "Selection", - module = "monarch._rust_bindings.monarch_hyperactor.selection", + module = "monarch.actor._extension.monarch_hyperactor.selection", frozen )] pub struct PySelection { diff --git a/monarch_hyperactor/src/shape.rs b/monarch_hyperactor/src/shape.rs index b01d391c..30b39366 100644 --- a/monarch_hyperactor/src/shape.rs +++ b/monarch_hyperactor/src/shape.rs @@ -18,7 +18,7 @@ use crate::ndslice::PySlice; #[pyclass( name = "Shape", - module = "monarch._rust_bindings.monarch_hyperactor.shape", + module = "monarch.actor._extension.monarch_hyperactor.shape", frozen )] pub struct PyShape { @@ -126,7 +126,7 @@ impl From for PyShape { #[pyclass( name = "Point", - module = "monarch._rust_bindings.monarch_hyperactor.shape", + module = "monarch.actor._extension.monarch_hyperactor.shape", subclass, frozen )] diff --git a/python/monarch/__init__.py b/python/monarch/__init__.py index a54ecc2f..bdd9cd6c 100644 --- a/python/monarch/__init__.py +++ b/python/monarch/__init__.py @@ -9,6 +9,9 @@ from importlib import import_module as _import_module from typing import TYPE_CHECKING +# Import this to initialize actor runtime before anything else. +import monarch.actor._extension # noqa: F401 + # Import before monarch to pre-load torch DSOs as, in exploded wheel flows, # our RPATHs won't correctly find them. import torch # noqa: F401 diff --git a/python/monarch/_monarch/hyperactor/__init__.py b/python/monarch/_monarch/hyperactor/__init__.py index 0adbd36e..f9dff551 100644 --- a/python/monarch/_monarch/hyperactor/__init__.py +++ b/python/monarch/_monarch/hyperactor/__init__.py @@ -6,15 +6,15 @@ # pyre-strict -from monarch._rust_bindings.monarch_hyperactor.actor import PythonMessage +from monarch.actor._extension.monarch_hyperactor.actor import PythonMessage -from monarch._rust_bindings.monarch_hyperactor.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.monarch_hyperactor.alloc import ( # @manual=//monarch/actor_extension:actor_extension LocalAllocatorBase, ) -from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox, PortId +from monarch.actor._extension.monarch_hyperactor.mailbox import Mailbox, PortId -from monarch._rust_bindings.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/actor_extension:actor_extension ActorId, Alloc, AllocConstraints, @@ -24,7 +24,7 @@ Serialized, ) -from monarch._rust_bindings.monarch_hyperactor.shape import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.monarch_hyperactor.shape import ( # @manual=//monarch/actor_extension:actor_extension Shape, ) diff --git a/python/monarch/_monarch/selection/__init__.py b/python/monarch/_monarch/selection/__init__.py index 376604e9..0b69a84c 100644 --- a/python/monarch/_monarch/selection/__init__.py +++ b/python/monarch/_monarch/selection/__init__.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from monarch._rust_bindings.monarch_hyperactor.selection import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.monarch_hyperactor.selection import ( # @manual=//monarch/monarch_extension:monarch_extension Selection, ) diff --git a/python/monarch/_rust_bindings/monarch_extension/client.pyi b/python/monarch/_rust_bindings/monarch_extension/client.pyi index 6c69af03..c3284aab 100644 --- a/python/monarch/_rust_bindings/monarch_extension/client.pyi +++ b/python/monarch/_rust_bindings/monarch_extension/client.pyi @@ -7,9 +7,9 @@ from typing import Any, ClassVar, Dict, final, List, NamedTuple, Union from monarch._rust_bindings.monarch_extension.tensor_worker import Ref -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId, Proc, Serialized -from monarch._rust_bindings.monarch_hyperactor.shape import Slice as NDSlice from monarch._rust_bindings.monarch_messages.debugger import DebuggerActionType +from monarch.actor._extension.monarch_hyperactor.proc import ActorId, Proc, Serialized +from monarch.actor._extension.monarch_hyperactor.shape import Slice as NDSlice class Exception: """ diff --git a/python/monarch/_rust_bindings/monarch_extension/controller.pyi b/python/monarch/_rust_bindings/monarch_extension/controller.pyi index d13a37af..153d9818 100644 --- a/python/monarch/_rust_bindings/monarch_extension/controller.pyi +++ b/python/monarch/_rust_bindings/monarch_extension/controller.pyi @@ -9,9 +9,9 @@ from typing import Any, final, List, Optional, Tuple, Union from monarch._rust_bindings.monarch_extension.tensor_worker import Ref, WorkerMessage -from monarch._rust_bindings.monarch_hyperactor.proc import Serialized +from monarch.actor._extension.monarch_hyperactor.proc import Serialized -from monarch._rust_bindings.monarch_hyperactor.shape import Slice +from monarch.actor._extension.monarch_hyperactor.shape import Slice @final class Node: diff --git a/python/monarch/_rust_bindings/monarch_extension/debugger.pyi b/python/monarch/_rust_bindings/monarch_extension/debugger.pyi index c0aafa36..f5d4033e 100644 --- a/python/monarch/_rust_bindings/monarch_extension/debugger.pyi +++ b/python/monarch/_rust_bindings/monarch_extension/debugger.pyi @@ -6,12 +6,13 @@ from typing import final, Optional, Union -from monarch._rust_bindings.monarch_hyperactor.proc import Serialized from monarch._rust_bindings.monarch_messages.debugger import ( DebuggerAction, DebuggerActionType, ) +from monarch.actor._extension.monarch_hyperactor.proc import Serialized + @final class DebuggerMessage: """A message for debugger communication between worker and client.""" diff --git a/python/monarch/_rust_bindings/monarch_extension/mesh_controller.pyi b/python/monarch/_rust_bindings/monarch_extension/mesh_controller.pyi index dbdba70e..61490683 100644 --- a/python/monarch/_rust_bindings/monarch_extension/mesh_controller.pyi +++ b/python/monarch/_rust_bindings/monarch_extension/mesh_controller.pyi @@ -8,11 +8,11 @@ from traceback import FrameSummary from typing import List, NamedTuple, Sequence, Tuple, Union from monarch._rust_bindings.monarch_extension import client -from monarch._rust_bindings.monarch_hyperactor.mailbox import PortId -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId -from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcMesh +from monarch.actor._extension.monarch_hyperactor.mailbox import PortId +from monarch.actor._extension.monarch_hyperactor.proc import ActorId +from monarch.actor._extension.monarch_hyperactor.proc_mesh import ProcMesh -from monarch._rust_bindings.monarch_hyperactor.shape import Slice as NDSlice +from monarch.actor._extension.monarch_hyperactor.shape import Slice as NDSlice class _Controller: def __init__(self) -> None: ... diff --git a/python/monarch/_rust_bindings/monarch_extension/tensor_worker.pyi b/python/monarch/_rust_bindings/monarch_extension/tensor_worker.pyi index 9936090b..ccba106d 100644 --- a/python/monarch/_rust_bindings/monarch_extension/tensor_worker.pyi +++ b/python/monarch/_rust_bindings/monarch_extension/tensor_worker.pyi @@ -7,8 +7,8 @@ from typing import Callable, final, Optional, Sequence, Tuple import torch -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId -from monarch._rust_bindings.monarch_hyperactor.shape import Slice +from monarch.actor._extension.monarch_hyperactor.proc import ActorId +from monarch.actor._extension.monarch_hyperactor.shape import Slice @final class Ref: diff --git a/python/monarch/_rust_bindings/old.pyi b/python/monarch/_rust_bindings/old.pyi index dbca00eb..2b01caf4 100644 --- a/python/monarch/_rust_bindings/old.pyi +++ b/python/monarch/_rust_bindings/old.pyi @@ -8,10 +8,11 @@ from typing import Dict, final, List, Optional, Protocol, Type -from monarch._rust_bindings.hyperactor_extension.alloc import Alloc, AllocSpec -from monarch._rust_bindings.monarch_hyperactor.selection import Selection +from monarch.actor._extension.hyperactor_extension.alloc import Alloc, AllocSpec -from monarch._rust_bindings.monarch_hyperactor.shape import Shape +from monarch.actor._extension.monarch_hyperactor.selection import Selection + +from monarch.actor._extension.monarch_hyperactor.shape import Shape def init_proc( *, diff --git a/python/monarch/actor/_actor_mesh.py b/python/monarch/actor/_actor_mesh.py index 41eb4702..ac4e3c70 100644 --- a/python/monarch/actor/_actor_mesh.py +++ b/python/monarch/actor/_actor_mesh.py @@ -38,20 +38,24 @@ TypeVar, ) -from monarch._rust_bindings.hyperactor_extension.telemetry import enter_span, exit_span +from monarch.actor._allocator import LocalAllocator, ProcessAllocator + +from monarch.actor._extension.hyperactor_extension.telemetry import ( + enter_span, + exit_span, +) -from monarch._rust_bindings.monarch_hyperactor.actor import PanicFlag, PythonMessage -from monarch._rust_bindings.monarch_hyperactor.actor_mesh import PythonActorMesh -from monarch._rust_bindings.monarch_hyperactor.mailbox import ( +from monarch.actor._extension.monarch_hyperactor.actor import PanicFlag, PythonMessage +from monarch.actor._extension.monarch_hyperactor.actor_mesh import PythonActorMesh +from monarch.actor._extension.monarch_hyperactor.mailbox import ( Mailbox, OncePortReceiver, OncePortRef, PortReceiver as HyPortReceiver, PortRef, ) -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId -from monarch._rust_bindings.monarch_hyperactor.shape import Point as HyPoint, Shape -from monarch.actor._allocator import LocalAllocator, ProcessAllocator +from monarch.actor._extension.monarch_hyperactor.proc import ActorId +from monarch.actor._extension.monarch_hyperactor.shape import Point as HyPoint, Shape from monarch.actor._future import Future from monarch.actor._pdb_wrapper import remote_breakpointhook diff --git a/python/monarch/actor/_allocator.py b/python/monarch/actor/_allocator.py index ce609701..84299011 100644 --- a/python/monarch/actor/_allocator.py +++ b/python/monarch/actor/_allocator.py @@ -10,12 +10,9 @@ import logging from typing import final, Optional -from monarch._rust_bindings.hyperactor_extension.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension - Alloc, - AllocSpec, -) +from monarch.actor._extension.hyperactor_extension.alloc import Alloc, AllocSpec -from monarch._rust_bindings.monarch_hyperactor.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.monarch_hyperactor.alloc import ( LocalAllocatorBase, ProcessAllocatorBase, RemoteAllocatorBase, diff --git a/python/monarch/actor/_bootstrap_main.py b/python/monarch/actor/_bootstrap_main.py index 89c791d9..3dd06b47 100644 --- a/python/monarch/actor/_bootstrap_main.py +++ b/python/monarch/actor/_bootstrap_main.py @@ -22,7 +22,7 @@ async def main(): - from monarch._rust_bindings.monarch_hyperactor.bootstrap import bootstrap_main + from monarch.actor._extension.monarch_hyperactor.bootstrap import bootstrap_main await bootstrap_main() diff --git a/python/monarch/actor/_code_sync/__init__.py b/python/monarch/actor/_code_sync/__init__.py index 4cadbf08..b666b6df 100644 --- a/python/monarch/actor/_code_sync/__init__.py +++ b/python/monarch/actor/_code_sync/__init__.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from monarch._rust_bindings.monarch_extension.code_sync import ( # noqa: F401 +from monarch.actor._extension.code_sync import ( # noqa: F401 RsyncMeshClient, WorkspaceLocation, ) diff --git a/python/monarch/actor/_debugger.py b/python/monarch/actor/_debugger.py index dcb34a5b..86d03f3b 100644 --- a/python/monarch/actor/_debugger.py +++ b/python/monarch/actor/_debugger.py @@ -10,8 +10,9 @@ from dataclasses import dataclass from typing import Dict, List, Tuple, Union -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId from monarch.actor._actor_mesh import Actor, ActorMeshRef, endpoint + +from monarch.actor._extension.monarch_hyperactor.proc import ActorId from monarch.actor._pdb_wrapper import DebuggerWrite from monarch.actor._proc_mesh import local_proc_mesh from tabulate import tabulate diff --git a/python/monarch/actor/_extension/__init__.pyi b/python/monarch/actor/_extension/__init__.pyi new file mode 100644 index 00000000..8fd5e7e9 --- /dev/null +++ b/python/monarch/actor/_extension/__init__.pyi @@ -0,0 +1,9 @@ +# 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. + +""" +Typestubs for the Monarch actor Rust extension. +""" diff --git a/python/monarch/_rust_bindings/monarch_extension/blocking.pyi b/python/monarch/actor/_extension/blocking.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_extension/blocking.pyi rename to python/monarch/actor/_extension/blocking.pyi diff --git a/python/monarch/_rust_bindings/monarch_extension/code_sync.pyi b/python/monarch/actor/_extension/code_sync.pyi similarity index 87% rename from python/monarch/_rust_bindings/monarch_extension/code_sync.pyi rename to python/monarch/actor/_extension/code_sync.pyi index 83e6f11d..4dbfa074 100644 --- a/python/monarch/_rust_bindings/monarch_extension/code_sync.pyi +++ b/python/monarch/actor/_extension/code_sync.pyi @@ -7,9 +7,9 @@ from pathlib import Path from typing import final -from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcMesh +from monarch.actor._extension.monarch_hyperactor.proc_mesh import ProcMesh -from monarch._rust_bindings.monarch_hyperactor.shape import Shape +from monarch.actor._extension.monarch_hyperactor.shape import Shape class WorkspaceLocation: """ diff --git a/python/monarch/_rust_bindings/hyperactor_extension/alloc.pyi b/python/monarch/actor/_extension/hyperactor_extension/alloc.pyi similarity index 100% rename from python/monarch/_rust_bindings/hyperactor_extension/alloc.pyi rename to python/monarch/actor/_extension/hyperactor_extension/alloc.pyi diff --git a/python/monarch/_rust_bindings/hyperactor_extension/telemetry.pyi b/python/monarch/actor/_extension/hyperactor_extension/telemetry.pyi similarity index 100% rename from python/monarch/_rust_bindings/hyperactor_extension/telemetry.pyi rename to python/monarch/actor/_extension/hyperactor_extension/telemetry.pyi diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/actor.pyi b/python/monarch/actor/_extension/monarch_hyperactor/actor.pyi similarity index 95% rename from python/monarch/_rust_bindings/monarch_hyperactor/actor.pyi rename to python/monarch/actor/_extension/monarch_hyperactor/actor.pyi index 2732494d..376634d4 100644 --- a/python/monarch/_rust_bindings/monarch_hyperactor/actor.pyi +++ b/python/monarch/actor/_extension/monarch_hyperactor/actor.pyi @@ -10,13 +10,13 @@ import abc from typing import final, List, Optional, Protocol -from monarch._rust_bindings.monarch_hyperactor.mailbox import ( +from monarch.actor._extension.monarch_hyperactor.mailbox import ( Mailbox, OncePortRef, PortRef, ) -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId, Proc, Serialized -from monarch._rust_bindings.monarch_hyperactor.shape import Shape +from monarch.actor._extension.monarch_hyperactor.proc import ActorId, Proc, Serialized +from monarch.actor._extension.monarch_hyperactor.shape import Shape @final class PickledMessage: diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/actor_mesh.pyi b/python/monarch/actor/_extension/monarch_hyperactor/actor_mesh.pyi similarity index 77% rename from python/monarch/_rust_bindings/monarch_hyperactor/actor_mesh.pyi rename to python/monarch/actor/_extension/monarch_hyperactor/actor_mesh.pyi index f485821b..406ca47e 100644 --- a/python/monarch/_rust_bindings/monarch_hyperactor/actor_mesh.pyi +++ b/python/monarch/actor/_extension/monarch_hyperactor/actor_mesh.pyi @@ -8,12 +8,12 @@ from typing import final -from monarch._rust_bindings.monarch_hyperactor.actor import PythonMessage -from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox +from monarch.actor._extension.monarch_hyperactor.actor import PythonMessage +from monarch.actor._extension.monarch_hyperactor.mailbox import Mailbox -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId +from monarch.actor._extension.monarch_hyperactor.proc import ActorId -from monarch._rust_bindings.monarch_hyperactor.shape import Shape +from monarch.actor._extension.monarch_hyperactor.shape import Shape @final class PythonActorMesh: diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/alloc.pyi b/python/monarch/actor/_extension/monarch_hyperactor/alloc.pyi similarity index 97% rename from python/monarch/_rust_bindings/monarch_hyperactor/alloc.pyi rename to python/monarch/actor/_extension/monarch_hyperactor/alloc.pyi index a32272c2..e17332db 100644 --- a/python/monarch/_rust_bindings/monarch_hyperactor/alloc.pyi +++ b/python/monarch/actor/_extension/monarch_hyperactor/alloc.pyi @@ -9,7 +9,7 @@ from datetime import timedelta from typing import Optional -from monarch._rust_bindings.hyperactor_extension.alloc import Alloc, AllocSpec +from monarch.actor._extension.hyperactor_extension.alloc import Alloc, AllocSpec from typing_extensions import Self class ProcessAllocatorBase: diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/bootstrap.pyi b/python/monarch/actor/_extension/monarch_hyperactor/bootstrap.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_hyperactor/bootstrap.pyi rename to python/monarch/actor/_extension/monarch_hyperactor/bootstrap.pyi diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/channel.pyi b/python/monarch/actor/_extension/monarch_hyperactor/channel.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_hyperactor/channel.pyi rename to python/monarch/actor/_extension/monarch_hyperactor/channel.pyi diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/mailbox.pyi b/python/monarch/actor/_extension/monarch_hyperactor/mailbox.pyi similarity index 96% rename from python/monarch/_rust_bindings/monarch_hyperactor/mailbox.pyi rename to python/monarch/actor/_extension/monarch_hyperactor/mailbox.pyi index 35743679..4a466458 100644 --- a/python/monarch/_rust_bindings/monarch_hyperactor/mailbox.pyi +++ b/python/monarch/actor/_extension/monarch_hyperactor/mailbox.pyi @@ -8,14 +8,14 @@ from typing import final, Generic, Protocol -from monarch._rust_bindings.monarch_hyperactor.actor import ( +from monarch.actor._extension.monarch_hyperactor.actor import ( PythonMessage, UndeliverableMessageEnvelope, ) -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId +from monarch.actor._extension.monarch_hyperactor.proc import ActorId -from monarch._rust_bindings.monarch_hyperactor.shape import Shape +from monarch.actor._extension.monarch_hyperactor.shape import Shape @final class PortId: diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/proc.pyi b/python/monarch/actor/_extension/monarch_hyperactor/proc.pyi similarity index 95% rename from python/monarch/_rust_bindings/monarch_hyperactor/proc.pyi rename to python/monarch/actor/_extension/monarch_hyperactor/proc.pyi index 0f7c5abe..577bf08e 100644 --- a/python/monarch/_rust_bindings/monarch_hyperactor/proc.pyi +++ b/python/monarch/actor/_extension/monarch_hyperactor/proc.pyi @@ -8,8 +8,8 @@ from typing import final, Optional, Type -from monarch._rust_bindings.monarch_hyperactor.actor import Actor, PythonActorHandle -from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox +from monarch.actor._extension.monarch_hyperactor.actor import Actor, PythonActorHandle +from monarch.actor._extension.monarch_hyperactor.mailbox import Mailbox def init_proc( *, diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/proc_mesh.pyi b/python/monarch/actor/_extension/monarch_hyperactor/proc_mesh.pyi similarity index 88% rename from python/monarch/_rust_bindings/monarch_hyperactor/proc_mesh.pyi rename to python/monarch/actor/_extension/monarch_hyperactor/proc_mesh.pyi index 3ee5d8ce..6a25e629 100644 --- a/python/monarch/_rust_bindings/monarch_hyperactor/proc_mesh.pyi +++ b/python/monarch/actor/_extension/monarch_hyperactor/proc_mesh.pyi @@ -8,11 +8,12 @@ from typing import AsyncIterator, final, Type -from monarch._rust_bindings.hyperactor_extension.alloc import Alloc -from monarch._rust_bindings.monarch_hyperactor.actor import Actor -from monarch._rust_bindings.monarch_hyperactor.actor_mesh import PythonActorMesh -from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox -from monarch._rust_bindings.monarch_hyperactor.shape import Shape +from monarch.actor._extension.hyperactor_extension.alloc import Alloc + +from monarch.actor._extension.monarch_hyperactor.actor import Actor +from monarch.actor._extension.monarch_hyperactor.actor_mesh import PythonActorMesh +from monarch.actor._extension.monarch_hyperactor.mailbox import Mailbox +from monarch.actor._extension.monarch_hyperactor.shape import Shape @final class ProcMesh: diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/runtime.pyi b/python/monarch/actor/_extension/monarch_hyperactor/runtime.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_hyperactor/runtime.pyi rename to python/monarch/actor/_extension/monarch_hyperactor/runtime.pyi diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/selection.pyi b/python/monarch/actor/_extension/monarch_hyperactor/selection.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_hyperactor/selection.pyi rename to python/monarch/actor/_extension/monarch_hyperactor/selection.pyi diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/shape.pyi b/python/monarch/actor/_extension/monarch_hyperactor/shape.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_hyperactor/shape.pyi rename to python/monarch/actor/_extension/monarch_hyperactor/shape.pyi diff --git a/python/monarch/_rust_bindings/monarch_extension/panic.pyi b/python/monarch/actor/_extension/panic.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_extension/panic.pyi rename to python/monarch/actor/_extension/panic.pyi diff --git a/python/monarch/actor/_pdb_wrapper.py b/python/monarch/actor/_pdb_wrapper.py index dea6cba3..b3574cdb 100644 --- a/python/monarch/actor/_pdb_wrapper.py +++ b/python/monarch/actor/_pdb_wrapper.py @@ -14,7 +14,7 @@ from typing import Dict, TYPE_CHECKING -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId +from monarch.actor._extension.monarch_hyperactor.proc import ActorId if TYPE_CHECKING: from monarch.actor._debugger import DebugClient diff --git a/python/monarch/actor/_proc_mesh.py b/python/monarch/actor/_proc_mesh.py index ccfd4799..da566346 100644 --- a/python/monarch/actor/_proc_mesh.py +++ b/python/monarch/actor/_proc_mesh.py @@ -22,23 +22,23 @@ TypeVar, ) -from monarch._rust_bindings.hyperactor_extension.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._actor_mesh import _Actor, _ActorMeshRefImpl, Actor, ActorMeshRef +from monarch.actor._allocator import LocalAllocator, ProcessAllocator +from monarch.actor._code_sync import RsyncMeshClient, WorkspaceLocation +from monarch.actor._code_sync.auto_reload import AutoReloadActor +from monarch.actor._device_utils import _local_device_count + +from monarch.actor._extension.hyperactor_extension.alloc import ( # @manual=//monarch/actor_extension:actor_extension Alloc, AllocConstraints, AllocSpec, ) -from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox -from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ( +from monarch.actor._extension.monarch_hyperactor.mailbox import Mailbox +from monarch.actor._extension.monarch_hyperactor.proc_mesh import ( ProcMesh as HyProcMesh, ProcMeshMonitor, ) -from monarch._rust_bindings.monarch_hyperactor.shape import Shape, Slice -from monarch.actor._actor_mesh import _Actor, _ActorMeshRefImpl, Actor, ActorMeshRef -from monarch.actor._allocator import LocalAllocator, ProcessAllocator -from monarch.actor._code_sync import RsyncMeshClient, WorkspaceLocation -from monarch.actor._code_sync.auto_reload import AutoReloadActor - -from monarch.actor._device_utils import _local_device_count +from monarch.actor._extension.monarch_hyperactor.shape import Shape, Slice from monarch.actor._future import Future from monarch.actor._shape import MeshTrait diff --git a/python/monarch/actor/_shape.py b/python/monarch/actor/_shape.py index 41afd867..486f43bc 100644 --- a/python/monarch/actor/_shape.py +++ b/python/monarch/actor/_shape.py @@ -10,7 +10,7 @@ from typing import Dict, Generator, Sequence, Tuple, Union -from monarch._rust_bindings.monarch_hyperactor.shape import Shape, Slice +from monarch.actor._extension.monarch_hyperactor.shape import Shape, Slice from typing_extensions import Self diff --git a/python/monarch/actor/_telemetry/__init__.py b/python/monarch/actor/_telemetry/__init__.py index 35406d2e..ea25d9a1 100644 --- a/python/monarch/actor/_telemetry/__init__.py +++ b/python/monarch/actor/_telemetry/__init__.py @@ -9,9 +9,7 @@ import logging -from monarch._rust_bindings.hyperactor_extension.telemetry import ( # @manual=//monarch/monarch_extension:monarch_extension - forward_to_tracing, -) +from monarch.actor._extension.hyperactor_extension.telemetry import forward_to_tracing class TracingForwarder(logging.Handler): diff --git a/python/monarch/actor/_telemetry/rust_span_tracing.py b/python/monarch/actor/_telemetry/rust_span_tracing.py index d0b243e2..d2c0604e 100644 --- a/python/monarch/actor/_telemetry/rust_span_tracing.py +++ b/python/monarch/actor/_telemetry/rust_span_tracing.py @@ -12,7 +12,7 @@ import opentelemetry.util.types as types # @manual=fbsource//third-party/pypi/opentelemetry-api:opentelemetry-api -from monarch._rust_bindings.hyperactor_extension.telemetry import ( +from monarch.actor._extension.hyperactor_extension.telemetry import ( get_current_span_id, PySpan, ) diff --git a/python/monarch/common/invocation.py b/python/monarch/common/invocation.py index d39d1e8e..6c5b74bc 100644 --- a/python/monarch/common/invocation.py +++ b/python/monarch/common/invocation.py @@ -8,7 +8,7 @@ import traceback from typing import Any, List, Optional, Tuple -from monarch._rust_bindings.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, ) diff --git a/python/monarch/common/selection.py b/python/monarch/common/selection.py index 6af54ff2..802ce69b 100644 --- a/python/monarch/common/selection.py +++ b/python/monarch/common/selection.py @@ -4,6 +4,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from monarch._rust_bindings.monarch_hyperactor.selection import Selection +from monarch.actor._extension.monarch_hyperactor.selection import Selection __all__ = ["Selection"] diff --git a/python/monarch/controller/controller.py b/python/monarch/controller/controller.py index 9e30405d..9c753442 100644 --- a/python/monarch/controller/controller.py +++ b/python/monarch/controller/controller.py @@ -15,7 +15,7 @@ WorldState, ) -from monarch._rust_bindings.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, ) diff --git a/python/monarch/controller/history.py b/python/monarch/controller/history.py index d710d64c..7fb75dd2 100644 --- a/python/monarch/controller/history.py +++ b/python/monarch/controller/history.py @@ -8,7 +8,7 @@ from collections import deque from typing import Generator, Sequence, TYPE_CHECKING -from monarch._rust_bindings.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, ) diff --git a/python/monarch/controller/rust_backend/controller.py b/python/monarch/controller/rust_backend/controller.py index ad73850c..a5a7533c 100644 --- a/python/monarch/controller/rust_backend/controller.py +++ b/python/monarch/controller/rust_backend/controller.py @@ -23,13 +23,13 @@ SystemSnapshotFilter, WorldState, ) -from monarch._rust_bindings.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension + +from monarch._rust_bindings.monarch_messages.debugger import DebuggerAction +from monarch.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, Proc, ) -from monarch._rust_bindings.monarch_messages.debugger import DebuggerAction - from monarch.actor._shape import NDSlice from monarch.common.controller_api import LogMessage, MessageResult from monarch.common.device_mesh import no_mesh diff --git a/python/monarch/mesh_controller.py b/python/monarch/mesh_controller.py index 015a5e50..e169a566 100644 --- a/python/monarch/mesh_controller.py +++ b/python/monarch/mesh_controller.py @@ -29,11 +29,11 @@ WorldState, ) from monarch._rust_bindings.monarch_extension.mesh_controller import _Controller -from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox -from monarch._rust_bindings.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._actor_mesh import Port, PortTuple +from monarch.actor._extension.monarch_hyperactor.mailbox import Mailbox +from monarch.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, ) -from monarch.actor._actor_mesh import Port, PortTuple from monarch.actor._shape import NDSlice from monarch.common import messages from monarch.common.controller_api import TController @@ -42,12 +42,12 @@ from monarch.common.tensor import Tensor if TYPE_CHECKING: - from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ( + from monarch.actor import ProcMesh + from monarch.actor._extension.monarch_hyperactor.proc_mesh import ( ProcMesh as HyProcMesh, ) - from monarch.actor import ProcMesh -from monarch._rust_bindings.monarch_hyperactor.shape import Point +from monarch.actor._extension.monarch_hyperactor.shape import Point from monarch.common.client import Client from monarch.common.controller_api import LogMessage, MessageResult diff --git a/python/monarch/rdma.py b/python/monarch/rdma.py index 2763f0cd..098e162d 100644 --- a/python/monarch/rdma.py +++ b/python/monarch/rdma.py @@ -10,8 +10,6 @@ from typing import cast, Dict, Optional, Tuple import torch - -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId from monarch.actor._actor_mesh import ( _ActorMeshRefImpl, Actor, @@ -20,6 +18,8 @@ MonarchContext, ) +from monarch.actor._extension.monarch_hyperactor.proc import ActorId + @dataclass class LocalRDMARecord: diff --git a/python/monarch/rust_backend_mesh.py b/python/monarch/rust_backend_mesh.py index dd876cd4..3afd3b5d 100644 --- a/python/monarch/rust_backend_mesh.py +++ b/python/monarch/rust_backend_mesh.py @@ -15,7 +15,7 @@ SystemSnapshotFilter, ) -from monarch._rust_bindings.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, init_proc, Proc, diff --git a/python/monarch/rust_local_mesh.py b/python/monarch/rust_local_mesh.py index b3fe87b8..3700344c 100644 --- a/python/monarch/rust_local_mesh.py +++ b/python/monarch/rust_local_mesh.py @@ -45,15 +45,15 @@ RunCommand, ) -from monarch._rust_bindings.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension - ActorId, -) - from monarch._rust_bindings.monarch_tensor_worker.bootstrap import ( WorkerServerRequest, WorkerServerResponse, ) +from monarch.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension + ActorId, +) + from monarch.common.device_mesh import DeviceMesh from monarch.common.fake import fake_call from monarch.common.invocation import DeviceException, RemoteException diff --git a/python/monarch/sim_mesh.py b/python/monarch/sim_mesh.py index 6727c78e..b3d6aab5 100644 --- a/python/monarch/sim_mesh.py +++ b/python/monarch/sim_mesh.py @@ -35,7 +35,7 @@ SimulatorClient, ) -from monarch._rust_bindings.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, init_proc, Proc, diff --git a/python/monarch/simulator/mock_controller.py b/python/monarch/simulator/mock_controller.py index 227f6788..4783793f 100644 --- a/python/monarch/simulator/mock_controller.py +++ b/python/monarch/simulator/mock_controller.py @@ -22,7 +22,7 @@ from monarch._rust_bindings.monarch_extension.client import ( # @manual=//monarch/monarch_extension:monarch_extension WorldState, ) -from monarch._rust_bindings.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, ) from monarch.actor._shape import iter_ranks, NDSlice, Slices as Ranks diff --git a/python/monarch/simulator/simulator.py b/python/monarch/simulator/simulator.py index 0fac4439..11eaa035 100644 --- a/python/monarch/simulator/simulator.py +++ b/python/monarch/simulator/simulator.py @@ -40,7 +40,7 @@ import numpy as np import torch -from monarch._rust_bindings.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, ) from monarch.actor._shape import iter_ranks, NDSlice diff --git a/python/monarch/tools/mesh_spec.py b/python/monarch/tools/mesh_spec.py index 7c1a5168..19217bc6 100644 --- a/python/monarch/tools/mesh_spec.py +++ b/python/monarch/tools/mesh_spec.py @@ -34,7 +34,7 @@ class MeshSpec: num_hosts: int host_type: str = _UNSET_STR gpus: int = _UNSET_INT - # NOTE: using str over monarch._rust_bindings.monarch_hyperactor.channel.ChannelTransport enum + # NOTE: using str over monarch.actor._extension.monarch_hyperactor.channel.ChannelTransport enum # b/c the rust binding doesn't have Python enum semantics, hence doesn't serialize well transport: str = "tcp" port: int = DEFAULT_REMOTE_ALLOCATOR_PORT diff --git a/python/tests/_monarch/test_actor.py b/python/tests/_monarch/test_actor.py index f2d249d6..2d876325 100644 --- a/python/tests/_monarch/test_actor.py +++ b/python/tests/_monarch/test_actor.py @@ -8,7 +8,7 @@ import time -from monarch._rust_bindings.monarch_hyperactor.actor import PythonMessage +from monarch.actor._extension.monarch_hyperactor.actor import PythonMessage def test_python_message() -> None: diff --git a/python/tests/_monarch/test_client.py b/python/tests/_monarch/test_client.py index 7bbc5b38..48bb0c88 100644 --- a/python/tests/_monarch/test_client.py +++ b/python/tests/_monarch/test_client.py @@ -11,7 +11,7 @@ import torch from monarch._rust_bindings.monarch_extension import client -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId +from monarch.actor._extension.monarch_hyperactor.proc import ActorId from pyre_extensions import none_throws diff --git a/python/tests/_monarch/test_controller.py b/python/tests/_monarch/test_controller.py index 531a6b44..74c2f0ea 100644 --- a/python/tests/_monarch/test_controller.py +++ b/python/tests/_monarch/test_controller.py @@ -12,7 +12,7 @@ controller, tensor_worker, ) -from monarch._rust_bindings.monarch_hyperactor import shape +from monarch.actor._extension.monarch_hyperactor import shape class TestController(TestCase): diff --git a/python/tests/_monarch/test_hyperactor.py b/python/tests/_monarch/test_hyperactor.py index 602b1867..c9632db8 100644 --- a/python/tests/_monarch/test_hyperactor.py +++ b/python/tests/_monarch/test_hyperactor.py @@ -13,16 +13,16 @@ import monarch -from monarch._rust_bindings.hyperactor_extension.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.hyperactor_extension.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension AllocConstraints, AllocSpec, ) -from monarch._rust_bindings.monarch_hyperactor.actor import PanicFlag, PythonMessage -from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId -from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcMesh -from monarch._rust_bindings.monarch_hyperactor.shape import Shape +from monarch.actor._extension.monarch_hyperactor.actor import PanicFlag, PythonMessage +from monarch.actor._extension.monarch_hyperactor.mailbox import Mailbox +from monarch.actor._extension.monarch_hyperactor.proc import ActorId +from monarch.actor._extension.monarch_hyperactor.proc_mesh import ProcMesh +from monarch.actor._extension.monarch_hyperactor.shape import Shape class MyActor: diff --git a/python/tests/_monarch/test_mailbox.py b/python/tests/_monarch/test_mailbox.py index 0bbda8d3..2b323c2e 100644 --- a/python/tests/_monarch/test_mailbox.py +++ b/python/tests/_monarch/test_mailbox.py @@ -12,20 +12,20 @@ import monarch -from monarch._rust_bindings.hyperactor_extension.alloc import ( +from monarch.actor._extension.hyperactor_extension.alloc import ( AllocConstraints, AllocSpec, ) -from monarch._rust_bindings.monarch_hyperactor.actor import PanicFlag, PythonMessage +from monarch.actor._extension.monarch_hyperactor.actor import PanicFlag, PythonMessage -from monarch._rust_bindings.monarch_hyperactor.mailbox import ( +from monarch.actor._extension.monarch_hyperactor.mailbox import ( Mailbox, PortReceiver, PortRef, ) -from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcMesh -from monarch._rust_bindings.monarch_hyperactor.shape import Shape +from monarch.actor._extension.monarch_hyperactor.proc_mesh import ProcMesh +from monarch.actor._extension.monarch_hyperactor.shape import Shape S = TypeVar("S") diff --git a/python/tests/_monarch/test_ndslice.py b/python/tests/_monarch/test_ndslice.py index b007edbb..8ca4febd 100644 --- a/python/tests/_monarch/test_ndslice.py +++ b/python/tests/_monarch/test_ndslice.py @@ -10,9 +10,9 @@ import random from unittest import TestCase -from monarch._rust_bindings.monarch_hyperactor.selection import Selection +from monarch.actor._extension.monarch_hyperactor.selection import Selection -from monarch._rust_bindings.monarch_hyperactor.shape import Shape, Slice +from monarch.actor._extension.monarch_hyperactor.shape import Shape, Slice class TestNdslice(TestCase): diff --git a/python/tests/_monarch/test_worker.py b/python/tests/_monarch/test_worker.py index 1a8d5436..30027ea9 100644 --- a/python/tests/_monarch/test_worker.py +++ b/python/tests/_monarch/test_worker.py @@ -13,7 +13,7 @@ import cloudpickle from monarch._rust_bindings.monarch_extension import tensor_worker -from monarch._rust_bindings.monarch_hyperactor import shape +from monarch.actor._extension.monarch_hyperactor import shape from pyre_extensions import none_throws diff --git a/python/tests/error_test_binary.py b/python/tests/error_test_binary.py index 524567d4..d32deeea 100644 --- a/python/tests/error_test_binary.py +++ b/python/tests/error_test_binary.py @@ -9,12 +9,12 @@ import sys import click -from monarch._rust_bindings.monarch_extension.blocking import blocking_function - -from monarch._rust_bindings.monarch_extension.panic import panicking_function from monarch.actor import Actor, endpoint, proc_mesh, send +from monarch.actor._extension.blocking import blocking_function +from monarch.actor._extension.panic import panicking_function + class ErrorActor(Actor): """An actor that has endpoints cause segfaults.""" diff --git a/python/tests/sleep_binary.py b/python/tests/sleep_binary.py index 25d0535e..8e57b9a1 100644 --- a/python/tests/sleep_binary.py +++ b/python/tests/sleep_binary.py @@ -14,7 +14,7 @@ import sys -from monarch._rust_bindings.monarch_hyperactor.runtime import ( # @manual +from monarch.actor._extension.monarch_hyperactor.runtime import ( # @manual sleep_indefinitely_for_unit_tests, ) diff --git a/python/tests/test_actor_error.py b/python/tests/test_actor_error.py index b51a471b..5c8df510 100644 --- a/python/tests/test_actor_error.py +++ b/python/tests/test_actor_error.py @@ -9,8 +9,8 @@ import subprocess import pytest -from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcEvent from monarch.actor import Actor, ActorError, endpoint, local_proc_mesh, proc_mesh +from monarch.actor._extension.monarch_hyperactor.proc_mesh import ProcEvent class ExceptionActor(Actor): diff --git a/python/tests/test_alloc.py b/python/tests/test_alloc.py index 53843267..8230d6b8 100644 --- a/python/tests/test_alloc.py +++ b/python/tests/test_alloc.py @@ -9,7 +9,7 @@ from unittest import IsolatedAsyncioTestCase from monarch import ProcessAllocator -from monarch._rust_bindings.hyperactor_extension.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.hyperactor_extension.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension AllocConstraints, AllocSpec, ) diff --git a/python/tests/test_allocator.py b/python/tests/test_allocator.py index 4623f56a..bcf0b5c7 100644 --- a/python/tests/test_allocator.py +++ b/python/tests/test_allocator.py @@ -23,15 +23,6 @@ import torch import torch.distributed as dist import torch.nn.functional as F - -from monarch._rust_bindings.hyperactor_extension.alloc import ( - AllocConstraints, - AllocSpec, -) -from monarch._rust_bindings.monarch_hyperactor.channel import ( - ChannelAddr, - ChannelTransport, -) from monarch.actor import ( Actor, current_rank, @@ -46,6 +37,15 @@ StaticRemoteAllocInitializer, TorchXRemoteAllocInitializer, ) + +from monarch.actor._extension.hyperactor_extension.alloc import ( + AllocConstraints, + AllocSpec, +) +from monarch.actor._extension.monarch_hyperactor.channel import ( + ChannelAddr, + ChannelTransport, +) from monarch.tools.mesh_spec import MeshSpec, ServerSpec from monarch.tools.network import get_sockaddr diff --git a/python/tests/test_future.py b/python/tests/test_future.py index e57cc29c..5ff6ff83 100644 --- a/python/tests/test_future.py +++ b/python/tests/test_future.py @@ -10,7 +10,7 @@ import pytest from monarch import Future, RemoteException -from monarch._rust_bindings.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, ) from monarch.common import future