diff --git a/Cargo.toml b/Cargo.toml index 9c623529..e687802d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,7 @@ [workspace] resolver = "2" members = [ + "actor_extension", "controller", "hyper", "hyperactor", diff --git a/actor_extension/Cargo.toml b/actor_extension/Cargo.toml new file mode 100644 index 00000000..0dc4fa80 --- /dev/null +++ b/actor_extension/Cargo.toml @@ -0,0 +1,26 @@ +# @generated by autocargo from //monarch/actor_extension:actor_extension-lib + +[package] +name = "_extension" +version = "0.0.0" +authors = ["Meta"] +edition = "2021" +license = "BSD-3-Clause" + +[lib] +test = false +doctest = false +crate-type = ["cdylib"] + +[dependencies] +anyhow = "1.0.98" +bincode = "1.3.3" +hyperactor = { version = "0.0.0", path = "../hyperactor" } +hyperactor_extension = { version = "0.0.0", path = "../hyperactor_extension" } +hyperactor_mesh = { version = "0.0.0", path = "../hyperactor_mesh" } +libc = "0.2.139" +monarch_hyperactor = { version = "0.0.0", path = "../monarch_hyperactor" } +pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods"] } +pyo3-async-runtimes = { version = "0.24", features = ["attributes", "tokio-runtime"] } +serde = { version = "1.0.185", features = ["derive", "rc"] } +tokio = { version = "1.45.0", features = ["full", "test-util", "tracing"] } 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..6c2c57de 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._src.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 d1652cc1..eb8f6226 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._src.actor._extension.code_sync" )] #[derive(Clone, Debug, Serialize, Deserialize)] enum PyWorkspaceLocation { @@ -75,7 +75,7 @@ impl PyWorkspaceLocation { #[pyclass( frozen, name = "RsyncMeshClient", - module = "monarch._rust_bindings.monarch_extension.code_sync" + module = "monarch._src.actor._extension.code_sync" )] pub struct RsyncMeshClient { actor_mesh: SharedCell>, diff --git a/actor_extension/src/lib.rs b/actor_extension/src/lib.rs new file mode 100644 index 00000000..13ebfbdf --- /dev/null +++ b/actor_extension/src/lib.rs @@ -0,0 +1,117 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +//! This module is used to expose Rust bindings for code supporting the +//! `monarch.actor` module. +//! +//! It is imported by `monarch` as `monarch._src.actor._extension`. +use pyo3::prelude::*; + +mod blocking; +mod code_sync; +mod panic; + +fn get_or_add_new_module<'py>( + module: &Bound<'py, PyModule>, + module_name: &str, +) -> PyResult> { + let mut current_module = module.clone(); + let mut parts = Vec::new(); + for part in module_name.split(".") { + parts.push(part); + let submodule = current_module.getattr(part).ok(); + if let Some(submodule) = submodule { + current_module = submodule.extract()?; + } else { + let new_module = PyModule::new(current_module.py(), part)?; + current_module.add_submodule(&new_module)?; + current_module + .py() + .import("sys")? + .getattr("modules")? + .set_item( + format!("monarch._src.actor._extension.{}", parts.join(".")), + new_module.clone(), + )?; + current_module = new_module; + } + } + Ok(current_module) +} + +#[pymodule] +#[pyo3(name = "_extension")] +pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> { + monarch_hyperactor::runtime::initialize(module.py())?; + let runtime = monarch_hyperactor::runtime::get_tokio_runtime(); + ::hyperactor::initialize(runtime.handle().clone()); + + monarch_hyperactor::shape::register_python_bindings(&get_or_add_new_module( + module, + "monarch_hyperactor.shape", + )?)?; + + monarch_hyperactor::selection::register_python_bindings(&get_or_add_new_module( + module, + "monarch_hyperactor.selection", + )?)?; + code_sync::register_python_bindings(&get_or_add_new_module(module, "code_sync")?)?; + monarch_hyperactor::bootstrap::register_python_bindings(&get_or_add_new_module( + module, + "monarch_hyperactor.bootstrap", + )?)?; + + monarch_hyperactor::proc::register_python_bindings(&get_or_add_new_module( + module, + "monarch_hyperactor.proc", + )?)?; + + monarch_hyperactor::actor::register_python_bindings(&get_or_add_new_module( + module, + "monarch_hyperactor.actor", + )?)?; + + monarch_hyperactor::mailbox::register_python_bindings(&get_or_add_new_module( + module, + "monarch_hyperactor.mailbox", + )?)?; + + monarch_hyperactor::alloc::register_python_bindings(&get_or_add_new_module( + module, + "monarch_hyperactor.alloc", + )?)?; + monarch_hyperactor::channel::register_python_bindings(&get_or_add_new_module( + module, + "monarch_hyperactor.channel", + )?)?; + monarch_hyperactor::actor_mesh::register_python_bindings(&get_or_add_new_module( + module, + "monarch_hyperactor.actor_mesh", + )?)?; + monarch_hyperactor::proc_mesh::register_python_bindings(&get_or_add_new_module( + module, + "monarch_hyperactor.proc_mesh", + )?)?; + + monarch_hyperactor::runtime::register_python_bindings(&get_or_add_new_module( + module, + "monarch_hyperactor.runtime", + )?)?; + hyperactor_extension::alloc::register_python_bindings(&get_or_add_new_module( + module, + "hyperactor_extension.alloc", + )?)?; + hyperactor_extension::telemetry::register_python_bindings(&get_or_add_new_module( + module, + "hyperactor_extension.telemetry", + )?)?; + crate::panic::register_python_bindings(&get_or_add_new_module(module, "panic")?)?; + + crate::blocking::register_python_bindings(&get_or_add_new_module(module, "blocking")?)?; + Ok(()) +} diff --git a/monarch_extension/src/panic.rs b/actor_extension/src/panic.rs similarity index 87% rename from monarch_extension/src/panic.rs rename to actor_extension/src/panic.rs index 5df12b9b..dd865ae4 100644 --- a/monarch_extension/src/panic.rs +++ b/actor_extension/src/panic.rs @@ -18,10 +18,7 @@ pub fn panicking_function() { /// Register Python bindings for the panic module. pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> { let f = wrap_pyfunction!(panicking_function, module)?; - f.setattr( - "__module__", - "monarch._rust_bindings.monarch_extension.panic", - )?; + f.setattr("__module__", "monarch._src.actor._extension.panic")?; module.add_function(f)?; Ok(()) } diff --git a/hyperactor_extension/src/alloc.rs b/hyperactor_extension/src/alloc.rs index bcc7b53d..cef50424 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._src.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._src.actor._extension.hyperactor_extension.alloc" )] pub struct PyAllocConstraints { inner: AllocConstraints, @@ -113,7 +113,7 @@ impl PyAllocConstraints { #[pyclass( name = "AllocSpec", - module = "monarch._rust_bindings.hyperactor_extension.alloc" + module = "monarch._src.actor._extension.hyperactor_extension.alloc" )] pub struct PyAllocSpec { pub inner: AllocSpec, diff --git a/hyperactor_extension/src/telemetry.rs b/hyperactor_extension/src/telemetry.rs index 32f9b4dc..d3eb6b78 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._src.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._src.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._src.actor._extension.hyperactor_extension.telemetry", )?; module.add_function(enter_span_fn)?; let exit_span_fn = wrap_pyfunction!(exit_span, module)?; exit_span_fn.setattr( "__module__", - "monarch._rust_bindings.hyperactor_extension.telemetry", + "monarch._src.actor._extension.hyperactor_extension.telemetry", )?; module.add_function(exit_span_fn)?; let get_current_span_id_fn = wrap_pyfunction!(get_current_span_id, module)?; get_current_span_id_fn.setattr( "__module__", - "monarch._rust_bindings.hyperactor_extension.telemetry", + "monarch._src.actor._extension.hyperactor_extension.telemetry", )?; module.add_function(get_current_span_id_fn)?; let use_real_clock_fn = wrap_pyfunction!(use_real_clock, module)?; use_real_clock_fn.setattr( "__module__", - "monarch._rust_bindings.hyperactor_extension.telemetry", + "monarch._src.actor._extension.hyperactor_extension.telemetry", )?; module.add_function(use_real_clock_fn)?; let use_sim_clock_fn = wrap_pyfunction!(use_sim_clock, module)?; use_sim_clock_fn.setattr( "__module__", - "monarch._rust_bindings.hyperactor_extension.telemetry", + "monarch._src.actor._extension.hyperactor_extension.telemetry", )?; module.add_function(use_sim_clock_fn)?; diff --git a/monarch_extension/Cargo.toml b/monarch_extension/Cargo.toml index e4952771..3500628d 100644 --- a/monarch_extension/Cargo.toml +++ b/monarch_extension/Cargo.toml @@ -16,14 +16,11 @@ crate-type = ["cdylib"] [dependencies] anyhow = "1.0.98" async-trait = "0.1.86" -bincode = "1.3.3" clap = { version = "4.5.38", features = ["derive", "env", "string", "unicode", "wrap_help"] } controller = { version = "0.0.0", path = "../controller", optional = true } hyperactor = { version = "0.0.0", path = "../hyperactor" } -hyperactor_extension = { version = "0.0.0", path = "../hyperactor_extension" } hyperactor_mesh = { version = "0.0.0", path = "../hyperactor_mesh" } hyperactor_multiprocess = { version = "0.0.0", path = "../hyperactor_multiprocess" } -libc = "0.2.139" monarch_hyperactor = { version = "0.0.0", path = "../monarch_hyperactor" } monarch_messages = { version = "0.0.0", path = "../monarch_messages", optional = true } monarch_simulator_lib = { version = "0.0.0", path = "../monarch_simulator", optional = true } @@ -33,7 +30,6 @@ nccl-sys = { path = "../nccl-sys", optional = true } ndslice = { version = "0.0.0", path = "../ndslice" } pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods"] } pyo3-async-runtimes = { version = "0.24", features = ["attributes", "tokio-runtime"] } -serde = { version = "1.0.185", features = ["derive", "rc"] } tokio = { version = "1.45.0", features = ["full", "test-util", "tracing"] } torch-sys = { version = "0.0.0", path = "../torch-sys", optional = true } torch-sys-cuda = { version = "0.0.0", path = "../torch-sys-cuda", optional = true } 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..c066aa48 100644 --- a/monarch_hyperactor/src/actor.rs +++ b/monarch_hyperactor/src/actor.rs @@ -52,7 +52,10 @@ 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._src.actor._extension.monarch_hyperactor.actor" +)] #[derive(Serialize, Deserialize, Named)] pub struct PickledMessage { sender_actor_id: ActorId, @@ -96,7 +99,7 @@ impl PickledMessage { } } -#[pyclass(module = "monarch._rust_bindings.monarch_hyperactor.actor")] +#[pyclass(module = "monarch._src.actor._extension.monarch_hyperactor.actor")] pub struct PickledMessageClientActor { instance: Arc>>, } @@ -171,7 +174,10 @@ impl PickledMessageClientActor { } } -#[pyclass(frozen, module = "monarch._rust_bindings.monarch_hyperactor.actor")] +#[pyclass( + frozen, + module = "monarch._src.actor._extension.monarch_hyperactor.actor" +)] #[derive(Default, Clone, Serialize, Deserialize, Named, PartialEq)] pub struct PythonMessage { pub(crate) method: String, @@ -260,7 +266,7 @@ impl PythonMessage { } } -#[pyclass(module = "monarch._rust_bindings.monarch_hyperactor.actor")] +#[pyclass(module = "monarch._src.actor._extension.monarch_hyperactor.actor")] pub(super) struct PythonActorHandle { pub(super) inner: ActorHandle, } @@ -385,7 +391,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._src.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 acac417e..357e0bd6 100644 --- a/monarch_hyperactor/src/actor_mesh.rs +++ b/monarch_hyperactor/src/actor_mesh.rs @@ -26,7 +26,7 @@ use crate::shape::PyShape; #[pyclass( name = "PythonActorMesh", - module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh" + module = "monarch._src.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..52ff9969 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._src.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._src.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._src.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._src.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..315f2318 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._src.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..e27d100d 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._src.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._src.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..3024ec50 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._src.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._src.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._src.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._src.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._src.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._src.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._src.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._src.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._src.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._src.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._src.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..681f3c1c 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._src.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..8a0b12e0 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._src.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._src.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._src.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._src.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._src.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..b0f9e239 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._src.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._src.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._src.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..a1ffc226 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._src.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 f4bb8ec3..3d73697c 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._src.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..8b73d727 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._src.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._src.actor._extension.monarch_hyperactor.shape", subclass, frozen )] diff --git a/python/monarch/__init__.py b/python/monarch/__init__.py index 7a0b1330..11b649b6 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._src.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/selection/__init__.py b/python/monarch/_monarch/selection/__init__.py new file mode 100644 index 00000000..fb4fd5cc --- /dev/null +++ b/python/monarch/_monarch/selection/__init__.py @@ -0,0 +1,13 @@ +# 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. + +from monarch._src.actor._extension.monarch_hyperactor.selection import ( # @manual=//monarch/monarch_extension:monarch_extension + Selection, +) + +__all__ = [ + "Selection", +] diff --git a/python/monarch/_rust_bindings/monarch_extension/client.pyi b/python/monarch/_rust_bindings/monarch_extension/client.pyi index 6c69af03..a38a466c 100644 --- a/python/monarch/_rust_bindings/monarch_extension/client.pyi +++ b/python/monarch/_rust_bindings/monarch_extension/client.pyi @@ -7,9 +7,13 @@ 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._src.actor._extension.monarch_hyperactor.proc import ( + ActorId, + Proc, + Serialized, +) +from monarch._src.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..6f6248a5 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._src.actor._extension.monarch_hyperactor.proc import Serialized -from monarch._rust_bindings.monarch_hyperactor.shape import Slice +from monarch._src.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..50284854 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._src.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 a256c556..7f2b70d3 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._src.actor._extension.monarch_hyperactor.mailbox import PortId +from monarch._src.actor._extension.monarch_hyperactor.proc import ActorId +from monarch._src.actor._extension.monarch_hyperactor.proc_mesh import ProcMesh -from monarch._rust_bindings.monarch_hyperactor.shape import Slice as NDSlice +from monarch._src.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..a84806a5 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._src.actor._extension.monarch_hyperactor.proc import ActorId +from monarch._src.actor._extension.monarch_hyperactor.shape import Slice @final class Ref: diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/alloc.pyi b/python/monarch/_rust_bindings/monarch_hyperactor/alloc.pyi index 30840ac9..0b0d854b 100644 --- a/python/monarch/_rust_bindings/monarch_hyperactor/alloc.pyi +++ b/python/monarch/_rust_bindings/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._src.actor._extension.hyperactor_extension.alloc import Alloc, AllocSpec from typing_extensions import Self class ProcessAllocatorBase: diff --git a/python/monarch/_rust_bindings/old.pyi b/python/monarch/_rust_bindings/old.pyi index dbca00eb..d90fc9f3 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._src.actor._extension.hyperactor_extension.alloc import Alloc, AllocSpec -from monarch._rust_bindings.monarch_hyperactor.shape import Shape +from monarch._src.actor._extension.monarch_hyperactor.selection import Selection + +from monarch._src.actor._extension.monarch_hyperactor.shape import Shape def init_proc( *, diff --git a/python/monarch/_src/actor/_extension/__init__.pyi b/python/monarch/_src/actor/_extension/__init__.pyi new file mode 100644 index 00000000..8fd5e7e9 --- /dev/null +++ b/python/monarch/_src/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/_src/actor/_extension/blocking.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_extension/blocking.pyi rename to python/monarch/_src/actor/_extension/blocking.pyi diff --git a/python/monarch/_rust_bindings/monarch_extension/code_sync.pyi b/python/monarch/_src/actor/_extension/code_sync.pyi similarity index 86% rename from python/monarch/_rust_bindings/monarch_extension/code_sync.pyi rename to python/monarch/_src/actor/_extension/code_sync.pyi index 83e6f11d..bc6d6dce 100644 --- a/python/monarch/_rust_bindings/monarch_extension/code_sync.pyi +++ b/python/monarch/_src/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._src.actor._extension.monarch_hyperactor.proc_mesh import ProcMesh -from monarch._rust_bindings.monarch_hyperactor.shape import Shape +from monarch._src.actor._extension.monarch_hyperactor.shape import Shape class WorkspaceLocation: """ diff --git a/python/monarch/_rust_bindings/hyperactor_extension/alloc.pyi b/python/monarch/_src/actor/_extension/hyperactor_extension/alloc.pyi similarity index 100% rename from python/monarch/_rust_bindings/hyperactor_extension/alloc.pyi rename to python/monarch/_src/actor/_extension/hyperactor_extension/alloc.pyi diff --git a/python/monarch/_rust_bindings/hyperactor_extension/telemetry.pyi b/python/monarch/_src/actor/_extension/hyperactor_extension/telemetry.pyi similarity index 100% rename from python/monarch/_rust_bindings/hyperactor_extension/telemetry.pyi rename to python/monarch/_src/actor/_extension/hyperactor_extension/telemetry.pyi diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/actor.pyi b/python/monarch/_src/actor/_extension/monarch_hyperactor/actor.pyi similarity index 94% rename from python/monarch/_rust_bindings/monarch_hyperactor/actor.pyi rename to python/monarch/_src/actor/_extension/monarch_hyperactor/actor.pyi index 2732494d..c41a152e 100644 --- a/python/monarch/_rust_bindings/monarch_hyperactor/actor.pyi +++ b/python/monarch/_src/actor/_extension/monarch_hyperactor/actor.pyi @@ -10,13 +10,17 @@ import abc from typing import final, List, Optional, Protocol -from monarch._rust_bindings.monarch_hyperactor.mailbox import ( +from monarch._src.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._src.actor._extension.monarch_hyperactor.proc import ( + ActorId, + Proc, + Serialized, +) +from monarch._src.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/_src/actor/_extension/monarch_hyperactor/actor_mesh.pyi similarity index 73% rename from python/monarch/_rust_bindings/monarch_hyperactor/actor_mesh.pyi rename to python/monarch/_src/actor/_extension/monarch_hyperactor/actor_mesh.pyi index 90371dab..db8df5e0 100644 --- a/python/monarch/_rust_bindings/monarch_hyperactor/actor_mesh.pyi +++ b/python/monarch/_src/actor/_extension/monarch_hyperactor/actor_mesh.pyi @@ -8,11 +8,11 @@ from typing import final -from monarch._rust_bindings.monarch_hyperactor.actor import 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.selection import Selection -from monarch._rust_bindings.monarch_hyperactor.shape import Shape +from monarch._src.actor._extension.monarch_hyperactor.actor import PythonMessage +from monarch._src.actor._extension.monarch_hyperactor.mailbox import Mailbox +from monarch._src.actor._extension.monarch_hyperactor.proc import ActorId +from monarch._src.actor._extension.monarch_hyperactor.selection import Selection +from monarch._src.actor._extension.monarch_hyperactor.shape import Shape @final class PythonActorMesh: diff --git a/python/monarch/_src/actor/_extension/monarch_hyperactor/alloc.pyi b/python/monarch/_src/actor/_extension/monarch_hyperactor/alloc.pyi new file mode 100644 index 00000000..0b0d854b --- /dev/null +++ b/python/monarch/_src/actor/_extension/monarch_hyperactor/alloc.pyi @@ -0,0 +1,107 @@ +# 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. + +# pyre-strict + +from datetime import timedelta +from typing import Optional + +from monarch._src.actor._extension.hyperactor_extension.alloc import Alloc, AllocSpec +from typing_extensions import Self + +class ProcessAllocatorBase: + def __init__( + self, + program: str, + args: Optional[list[str]] = None, + envs: Optional[dict[str, str]] = None, + ) -> None: + """ + Create a new process allocator. + + Arguments: + - `program`: The program for each process to run. Must be a hyperactor + bootstrapped program. + - `args`: The arguments to pass to the program. + - `envs`: The environment variables to set for the program. + """ + ... + + async def allocate_nonblocking(self, spec: AllocSpec) -> Alloc: + """ + Allocate a process according to the provided spec. + + Arguments: + - `spec`: The spec to allocate according to. + """ + ... + + def allocate_blocking(self, spec: AllocSpec) -> Alloc: + """ + Allocate a process according to the provided spec, blocking until an + alloc is returned. + + Arguments: + - `spec`: The spec to allocate according to. + """ + ... + +class LocalAllocatorBase: + async def allocate_nonblocking(self, spec: AllocSpec) -> Alloc: + """ + Allocate a process according to the provided spec. + + Arguments: + - `spec`: The spec to allocate according to. + """ + ... + + def allocate_blocking(self, spec: AllocSpec) -> Alloc: + """ + Allocate a process according to the provided spec, blocking until an + alloc is returned. + + Arguments: + - `spec`: The spec to allocate according to. + """ + ... + +class RemoteAllocatorBase: + def __new__( + cls, + world_id: str, + initializer: "monarch._src.actor.allocator.RemoteAllocInitializer", # pyre-ignore[11] + heartbeat_interval: timedelta = timedelta(seconds=5), + ) -> Self: + """ + Create a new (client-side) allocator instance that submits allocation requests to + remote hosts that are running hyperactor's RemoteProcessAllocator. + + Arguments: + - `world_id`: The world id to use for the remote allocator. + - `initializer`: Returns the server addresses to send allocation requests to. + - `heartbeat_interval`: Heartbeat interval used to maintain health status of remote hosts. + """ + ... + + async def allocate_nonblocking(self, spec: AllocSpec) -> Alloc: + """ + Allocate a process according to the provided spec. + + Arguments: + - `spec`: The spec to allocate according to. + """ + ... + + def allocate_blocking(self, spec: AllocSpec) -> Alloc: + """ + Allocate a process according to the provided spec, blocking until an + alloc is returned. + + Arguments: + - `spec`: The spec to allocate according to. + """ + ... diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/bootstrap.pyi b/python/monarch/_src/actor/_extension/monarch_hyperactor/bootstrap.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_hyperactor/bootstrap.pyi rename to python/monarch/_src/actor/_extension/monarch_hyperactor/bootstrap.pyi diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/channel.pyi b/python/monarch/_src/actor/_extension/monarch_hyperactor/channel.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_hyperactor/channel.pyi rename to python/monarch/_src/actor/_extension/monarch_hyperactor/channel.pyi diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/mailbox.pyi b/python/monarch/_src/actor/_extension/monarch_hyperactor/mailbox.pyi similarity index 96% rename from python/monarch/_rust_bindings/monarch_hyperactor/mailbox.pyi rename to python/monarch/_src/actor/_extension/monarch_hyperactor/mailbox.pyi index 35743679..0c057587 100644 --- a/python/monarch/_rust_bindings/monarch_hyperactor/mailbox.pyi +++ b/python/monarch/_src/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._src.actor._extension.monarch_hyperactor.actor import ( PythonMessage, UndeliverableMessageEnvelope, ) -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId +from monarch._src.actor._extension.monarch_hyperactor.proc import ActorId -from monarch._rust_bindings.monarch_hyperactor.shape import Shape +from monarch._src.actor._extension.monarch_hyperactor.shape import Shape @final class PortId: diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/proc.pyi b/python/monarch/_src/actor/_extension/monarch_hyperactor/proc.pyi similarity index 95% rename from python/monarch/_rust_bindings/monarch_hyperactor/proc.pyi rename to python/monarch/_src/actor/_extension/monarch_hyperactor/proc.pyi index 0f7c5abe..42f01c94 100644 --- a/python/monarch/_rust_bindings/monarch_hyperactor/proc.pyi +++ b/python/monarch/_src/actor/_extension/monarch_hyperactor/proc.pyi @@ -8,8 +8,11 @@ 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._src.actor._extension.monarch_hyperactor.actor import ( + Actor, + PythonActorHandle, +) +from monarch._src.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/_src/actor/_extension/monarch_hyperactor/proc_mesh.pyi similarity index 87% rename from python/monarch/_rust_bindings/monarch_hyperactor/proc_mesh.pyi rename to python/monarch/_src/actor/_extension/monarch_hyperactor/proc_mesh.pyi index 3ee5d8ce..5e91715c 100644 --- a/python/monarch/_rust_bindings/monarch_hyperactor/proc_mesh.pyi +++ b/python/monarch/_src/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._src.actor._extension.hyperactor_extension.alloc import Alloc + +from monarch._src.actor._extension.monarch_hyperactor.actor import Actor +from monarch._src.actor._extension.monarch_hyperactor.actor_mesh import PythonActorMesh +from monarch._src.actor._extension.monarch_hyperactor.mailbox import Mailbox +from monarch._src.actor._extension.monarch_hyperactor.shape import Shape @final class ProcMesh: diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/runtime.pyi b/python/monarch/_src/actor/_extension/monarch_hyperactor/runtime.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_hyperactor/runtime.pyi rename to python/monarch/_src/actor/_extension/monarch_hyperactor/runtime.pyi diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/selection.pyi b/python/monarch/_src/actor/_extension/monarch_hyperactor/selection.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_hyperactor/selection.pyi rename to python/monarch/_src/actor/_extension/monarch_hyperactor/selection.pyi diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/shape.pyi b/python/monarch/_src/actor/_extension/monarch_hyperactor/shape.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_hyperactor/shape.pyi rename to python/monarch/_src/actor/_extension/monarch_hyperactor/shape.pyi diff --git a/python/monarch/_rust_bindings/monarch_extension/panic.pyi b/python/monarch/_src/actor/_extension/panic.pyi similarity index 100% rename from python/monarch/_rust_bindings/monarch_extension/panic.pyi rename to python/monarch/_src/actor/_extension/panic.pyi diff --git a/python/monarch/_src/actor/actor_mesh.py b/python/monarch/_src/actor/actor_mesh.py index 687add3e..db650495 100644 --- a/python/monarch/_src/actor/actor_mesh.py +++ b/python/monarch/_src/actor/actor_mesh.py @@ -38,19 +38,29 @@ TypeVar, ) -from monarch._rust_bindings.hyperactor_extension.telemetry import enter_span, exit_span +from monarch._src.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._src.actor._extension.monarch_hyperactor.actor import ( + PanicFlag, + PythonMessage, +) +from monarch._src.actor._extension.monarch_hyperactor.actor_mesh import PythonActorMesh +from monarch._src.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._src.actor._extension.monarch_hyperactor.proc import ActorId +from monarch._src.actor._extension.monarch_hyperactor.shape import ( + Point as HyPoint, + Shape, +) + from monarch._src.actor.allocator import LocalAllocator, ProcessAllocator from monarch._src.actor.future import Future from monarch._src.actor.pdb_wrapper import remote_breakpointhook @@ -93,7 +103,7 @@ def get() -> "MonarchContext": _context: contextvars.ContextVar[MonarchContext] = contextvars.ContextVar( - "monarch.actor_mesh._context" + "monarch._src.actor.actor_mesh._context" ) diff --git a/python/monarch/_src/actor/allocator.py b/python/monarch/_src/actor/allocator.py index 552fba71..4261cb23 100644 --- a/python/monarch/_src/actor/allocator.py +++ b/python/monarch/_src/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._src.actor._extension.hyperactor_extension.alloc import Alloc, AllocSpec -from monarch._rust_bindings.monarch_hyperactor.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch._src.actor._extension.monarch_hyperactor.alloc import ( LocalAllocatorBase, ProcessAllocatorBase, RemoteAllocatorBase, diff --git a/python/monarch/_src/actor/bootstrap_main.py b/python/monarch/_src/actor/bootstrap_main.py index 5b377ac2..a5f50898 100644 --- a/python/monarch/_src/actor/bootstrap_main.py +++ b/python/monarch/_src/actor/bootstrap_main.py @@ -22,7 +22,9 @@ async def main(): - from monarch._rust_bindings.monarch_hyperactor.bootstrap import bootstrap_main + from monarch._src.actor._extension.monarch_hyperactor.bootstrap import ( + bootstrap_main, + ) await bootstrap_main() diff --git a/python/monarch/_src/actor/code_sync/__init__.py b/python/monarch/_src/actor/code_sync/__init__.py index 4cadbf08..4ec64790 100644 --- a/python/monarch/_src/actor/code_sync/__init__.py +++ b/python/monarch/_src/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._src.actor._extension.code_sync import ( # noqa: F401 RsyncMeshClient, WorkspaceLocation, ) diff --git a/python/monarch/_src/actor/debugger.py b/python/monarch/_src/actor/debugger.py index 2906c63f..31d0dcbb 100644 --- a/python/monarch/_src/actor/debugger.py +++ b/python/monarch/_src/actor/debugger.py @@ -10,7 +10,7 @@ from dataclasses import dataclass from typing import Dict, List, Tuple, Union -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId +from monarch._src.actor._extension.monarch_hyperactor.proc import ActorId from monarch._src.actor.actor_mesh import Actor, ActorMeshRef, endpoint from monarch._src.actor.pdb_wrapper import DebuggerWrite from monarch._src.actor.proc_mesh import local_proc_mesh diff --git a/python/monarch/_src/actor/pdb_wrapper.py b/python/monarch/_src/actor/pdb_wrapper.py index 87031dd4..a79d995c 100644 --- a/python/monarch/_src/actor/pdb_wrapper.py +++ b/python/monarch/_src/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._src.actor._extension.monarch_hyperactor.proc import ActorId if TYPE_CHECKING: from monarch._src.actor.debugger import DebugClient diff --git a/python/monarch/_src/actor/proc_mesh.py b/python/monarch/_src/actor/proc_mesh.py index d0c3fc79..dcde3309 100644 --- a/python/monarch/_src/actor/proc_mesh.py +++ b/python/monarch/_src/actor/proc_mesh.py @@ -22,17 +22,17 @@ TypeVar, ) -from monarch._rust_bindings.hyperactor_extension.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension # @manual=//monarch/monarch_extension:monarch_extension +from monarch._src.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._src.actor._extension.monarch_hyperactor.mailbox import Mailbox +from monarch._src.actor._extension.monarch_hyperactor.proc_mesh import ( ProcMesh as HyProcMesh, ProcMeshMonitor, ) -from monarch._rust_bindings.monarch_hyperactor.shape import Shape, Slice +from monarch._src.actor._extension.monarch_hyperactor.shape import Shape, Slice from monarch._src.actor.actor_mesh import _Actor, _ActorMeshRefImpl, Actor, ActorMeshRef from monarch._src.actor.allocator import LocalAllocator, ProcessAllocator from monarch._src.actor.code_sync import RsyncMeshClient, WorkspaceLocation diff --git a/python/monarch/_src/actor/shape.py b/python/monarch/_src/actor/shape.py index 41afd867..0e26fdc0 100644 --- a/python/monarch/_src/actor/shape.py +++ b/python/monarch/_src/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._src.actor._extension.monarch_hyperactor.shape import Shape, Slice from typing_extensions import Self diff --git a/python/monarch/_src/actor/telemetry/__init__.py b/python/monarch/_src/actor/telemetry/__init__.py index 35406d2e..7ffe0979 100644 --- a/python/monarch/_src/actor/telemetry/__init__.py +++ b/python/monarch/_src/actor/telemetry/__init__.py @@ -9,7 +9,7 @@ import logging -from monarch._rust_bindings.hyperactor_extension.telemetry import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch._src.actor._extension.hyperactor_extension.telemetry import ( forward_to_tracing, ) diff --git a/python/monarch/_src/actor/telemetry/rust_span_tracing.py b/python/monarch/_src/actor/telemetry/rust_span_tracing.py index d0b243e2..41ca30d9 100644 --- a/python/monarch/_src/actor/telemetry/rust_span_tracing.py +++ b/python/monarch/_src/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._src.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..1a01e492 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._src.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..5d7d4c77 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._src.actor._extension.monarch_hyperactor.selection import Selection __all__ = ["Selection"] diff --git a/python/monarch/controller/controller.py b/python/monarch/controller/controller.py index d96b2ac5..a591c6d8 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._src.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..61ca7508 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._src.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 f6dfecf4..919d7fb3 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._src.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._src.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 04225fe1..201d9c95 100644 --- a/python/monarch/mesh_controller.py +++ b/python/monarch/mesh_controller.py @@ -30,8 +30,8 @@ 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._src.actor._extension.monarch_hyperactor.mailbox import Mailbox +from monarch._src.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, ) from monarch._src.actor.actor_mesh import Port, PortTuple @@ -45,12 +45,12 @@ from monarch.tensor_worker_main import _set_trace if TYPE_CHECKING: - from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ( + from monarch._src.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._src.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 b9cc771a..cad4131a 100644 --- a/python/monarch/rdma.py +++ b/python/monarch/rdma.py @@ -11,7 +11,7 @@ import torch -from monarch._rust_bindings.monarch_hyperactor.proc import ActorId +from monarch._src.actor._extension.monarch_hyperactor.proc import ActorId from monarch._src.actor.actor_mesh import ( _ActorMeshRefImpl, Actor, diff --git a/python/monarch/rust_backend_mesh.py b/python/monarch/rust_backend_mesh.py index 665ccb2b..260e26f1 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._src.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..9b9cf2cf 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._src.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 b91361bf..ffd447fb 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._src.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 605ac58b..ab1db332 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._src.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, ) from monarch._src.actor.shape import iter_ranks, NDSlice, Slices as Ranks diff --git a/python/monarch/simulator/simulator.py b/python/monarch/simulator/simulator.py index 9a7c79c9..f6568b7f 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._src.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, ) from monarch._src.actor.shape import iter_ranks, NDSlice diff --git a/python/monarch/tools/mesh_spec.py b/python/monarch/tools/mesh_spec.py index 7c1a5168..7fd32a49 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._src.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..8b36a998 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._src.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..ffd983a5 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._src.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..4080076e 100644 --- a/python/tests/_monarch/test_controller.py +++ b/python/tests/_monarch/test_controller.py @@ -8,11 +8,11 @@ from unittest import TestCase -from monarch._rust_bindings.monarch_extension import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch._rust_bindings.monarch_extension import ( # @manual=//monarch/python/monarch/_rust_bindings:_rust_bindings controller, tensor_worker, ) -from monarch._rust_bindings.monarch_hyperactor import shape +from monarch._src.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..ef3a767e 100644 --- a/python/tests/_monarch/test_hyperactor.py +++ b/python/tests/_monarch/test_hyperactor.py @@ -13,16 +13,19 @@ import monarch -from monarch._rust_bindings.hyperactor_extension.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch._src.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._src.actor._extension.monarch_hyperactor.actor import ( + PanicFlag, + PythonMessage, +) +from monarch._src.actor._extension.monarch_hyperactor.mailbox import Mailbox +from monarch._src.actor._extension.monarch_hyperactor.proc import ActorId +from monarch._src.actor._extension.monarch_hyperactor.proc_mesh import ProcMesh +from monarch._src.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 8027f9e9..c025f4e1 100644 --- a/python/tests/_monarch/test_mailbox.py +++ b/python/tests/_monarch/test_mailbox.py @@ -12,21 +12,24 @@ import monarch -from monarch._rust_bindings.hyperactor_extension.alloc import ( +from monarch._src.actor._extension.hyperactor_extension.alloc import ( AllocConstraints, AllocSpec, ) -from monarch._rust_bindings.monarch_hyperactor.actor import PanicFlag, PythonMessage +from monarch._src.actor._extension.monarch_hyperactor.actor import ( + PanicFlag, + PythonMessage, +) -from monarch._rust_bindings.monarch_hyperactor.mailbox import ( +from monarch._src.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.selection import Selection -from monarch._rust_bindings.monarch_hyperactor.shape import Shape +from monarch._src.actor._extension.monarch_hyperactor.proc_mesh import ProcMesh +from monarch._src.actor._extension.monarch_hyperactor.selection import Selection +from monarch._src.actor._extension.monarch_hyperactor.shape import Shape S = TypeVar("S") U = TypeVar("U") diff --git a/python/tests/_monarch/test_ndslice.py b/python/tests/_monarch/test_ndslice.py index b007edbb..5757ac8e 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._src.actor._extension.monarch_hyperactor.selection import Selection -from monarch._rust_bindings.monarch_hyperactor.shape import Shape, Slice +from monarch._src.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..38827469 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._src.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 c3a58df4..9c28a5da 100644 --- a/python/tests/error_test_binary.py +++ b/python/tests/error_test_binary.py @@ -9,9 +9,9 @@ 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._src.actor._extension.blocking import blocking_function +from monarch._src.actor._extension.panic import panicking_function from monarch.actor import Actor, endpoint, proc_mesh, send diff --git a/python/tests/sleep_binary.py b/python/tests/sleep_binary.py index 25d0535e..ab9f2973 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._src.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 0f5c523e..add0fc21 100644 --- a/python/tests/test_actor_error.py +++ b/python/tests/test_actor_error.py @@ -9,7 +9,7 @@ import subprocess import pytest -from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcEvent +from monarch._src.actor._extension.monarch_hyperactor.proc_mesh import ProcEvent from monarch.actor import Actor, ActorError, endpoint, local_proc_mesh, proc_mesh diff --git a/python/tests/test_alloc.py b/python/tests/test_alloc.py index 53843267..e8d97f07 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._src.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 e875826e..d1acedc0 100644 --- a/python/tests/test_allocator.py +++ b/python/tests/test_allocator.py @@ -24,11 +24,11 @@ import torch.distributed as dist import torch.nn.functional as F -from monarch._rust_bindings.hyperactor_extension.alloc import ( +from monarch._src.actor._extension.hyperactor_extension.alloc import ( AllocConstraints, AllocSpec, ) -from monarch._rust_bindings.monarch_hyperactor.channel import ( +from monarch._src.actor._extension.monarch_hyperactor.channel import ( ChannelAddr, ChannelTransport, ) diff --git a/python/tests/test_future.py b/python/tests/test_future.py index e57cc29c..e2dc9847 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._src.actor._extension.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension ActorId, ) from monarch.common import future diff --git a/setup.py b/setup.py index 4e0ed2f6..85cf3ed9 100644 --- a/setup.py +++ b/setup.py @@ -121,16 +121,24 @@ def run(self): rust_extensions = [ RustExtension( - "monarch._rust_bindings", + "monarch._src.actor._extension", binding=Binding.PyO3, - path="monarch_extension/Cargo.toml", + path="actor_extension/Cargo.toml", debug=False, - features=["tensor_engine"] if USE_TENSOR_ENGINE else [], - args=[] if USE_TENSOR_ENGINE else ["--no-default-features"], ), ] if USE_TENSOR_ENGINE: + rust_extensions.append( + RustExtension( + "monarch._rust_bindings", + binding=Binding.PyO3, + path="monarch_extension/Cargo.toml", + debug=False, + features=["tensor_engine"] if USE_TENSOR_ENGINE else [], + args=[] if USE_TENSOR_ENGINE else ["--no-default-features"], + ) + ) rust_extensions.append( RustExtension( {"controller_bin": "monarch.monarch_controller"},