Skip to content

[monarch/ext] fold monarch_hyperactor into actor_extension #465

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: gh/suo/54/base
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions actor_extension/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# @generated by autocargo from //monarch/actor_extension:actor_extension-lib
# @generated by autocargo from //monarch/actor_extension:[actor_extension-lib,process_allocator-oss]

[package]
name = "_extension"
name = "actor_extension_lib"
version = "0.0.0"
authors = ["Meta"]
edition = "2021"
Expand All @@ -14,14 +14,23 @@ 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"] }
erased-serde = "0.3.27"
fbinit = { version = "0.2.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
hyperactor = { version = "0.0.0", path = "../hyperactor" }
hyperactor_mesh = { version = "0.0.0", path = "../hyperactor_mesh" }
hyperactor_multiprocess = { version = "0.0.0", path = "../hyperactor_multiprocess" }
hyperactor_telemetry = { version = "0.0.0", path = "../hyperactor_telemetry" }
inventory = "0.3.8"
libc = "0.2.139"
monarch_hyperactor = { version = "0.0.0", path = "../monarch_hyperactor" }
monarch_types = { version = "0.0.0", path = "../monarch_types" }
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"] }
serde_bytes = "0.11"
thiserror = "2.0.12"
tokio = { version = "1.45.0", features = ["full", "test-util", "tracing"] }
tracing = { version = "0.1.41", features = ["attributes", "valuable"] }
16 changes: 5 additions & 11 deletions monarch_hyperactor/src/actor.rs → actor_extension/src/actor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,10 +52,7 @@ use crate::proc::PySerialized;
use crate::runtime::signal_safe_block_on;
use crate::shape::PyShape;

#[pyclass(
frozen,
module = "monarch._src.actor._extension.monarch_hyperactor.actor"
)]
#[pyclass(frozen, module = "monarch._src.actor._extension.actor")]
#[derive(Serialize, Deserialize, Named)]
pub struct PickledMessage {
sender_actor_id: ActorId,
Expand Down Expand Up @@ -99,7 +96,7 @@ impl PickledMessage {
}
}

#[pyclass(module = "monarch._src.actor._extension.monarch_hyperactor.actor")]
#[pyclass(module = "monarch._src.actor._extension.actor")]
pub struct PickledMessageClientActor {
instance: Arc<Mutex<InstanceWrapper<PickledMessage>>>,
}
Expand Down Expand Up @@ -174,10 +171,7 @@ impl PickledMessageClientActor {
}
}

#[pyclass(
frozen,
module = "monarch._src.actor._extension.monarch_hyperactor.actor"
)]
#[pyclass(frozen, module = "monarch._src.actor._extension.actor")]
#[derive(Default, Clone, Serialize, Deserialize, Named, PartialEq)]
pub struct PythonMessage {
pub(crate) method: String,
Expand Down Expand Up @@ -266,7 +260,7 @@ impl PythonMessage {
}
}

#[pyclass(module = "monarch._src.actor._extension.monarch_hyperactor.actor")]
#[pyclass(module = "monarch._src.actor._extension.actor")]
pub(super) struct PythonActorHandle {
pub(super) inner: ActorHandle<PythonActor>,
}
Expand Down Expand Up @@ -391,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._src.actor._extension.monarch_hyperactor.actor")]
#[pyclass(module = "monarch._src.actor._extension.actor")]
struct PanicFlag {
sender: Option<tokio::sync::oneshot::Sender<PyObject>>,
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ use crate::shape::PyShape;

#[pyclass(
name = "PythonActorMesh",
module = "monarch._src.actor._extension.monarch_hyperactor.actor_mesh"
module = "monarch._src.actor._extension.actor_mesh"
)]
pub struct PythonActorMesh {
pub(super) inner: SharedCell<RootActorMesh<'static, PythonActor>>,
Expand Down
20 changes: 7 additions & 13 deletions monarch_hyperactor/src/alloc.rs → actor_extension/src/alloc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ use crate::runtime::signal_safe_block_on;

#[pyclass(
name = "LocalAllocatorBase",
module = "monarch._src.actor._extension.monarch_hyperactor.alloc",
module = "monarch._src.actor._extension.alloc",
subclass
)]
pub struct PyLocalAllocator;
Expand Down Expand Up @@ -88,7 +88,7 @@ impl PyLocalAllocator {

#[pyclass(
name = "ProcessAllocatorBase",
module = "monarch._src.actor._extension.monarch_hyperactor.alloc",
module = "monarch._src.actor._extension.alloc",
subclass
)]
pub struct PyProcessAllocator {
Expand Down Expand Up @@ -156,7 +156,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._src.actor._extension.monarch_hyperactor.alloc.RemoteProcessAllocInitializer`.
// instance of a Python subclass of `monarch._src.actor._extension.alloc.RemoteProcessAllocInitializer`.
py_inner: Py<PyAny>,

// allocation constraints passed onto the allocator's allocate call and passed along to python initializer.
Expand Down Expand Up @@ -234,7 +234,7 @@ impl RemoteProcessAllocInitializer for PyRemoteProcessAllocInitializer {

#[pyclass(
name = "RemoteAllocatorBase",
module = "monarch._src.actor._extension.monarch_hyperactor.alloc",
module = "monarch._src.actor._extension.alloc",
subclass
)]
pub struct PyRemoteAllocator {
Expand Down Expand Up @@ -335,10 +335,7 @@ impl PyRemoteAllocator {
/// A python class that wraps a Rust Alloc trait object. It represents what
/// is shown on the python side. Internals are not exposed.
/// It ensures that the Alloc is only used once (i.e. moved) in rust.
#[pyclass(
name = "Alloc",
module = "monarch._src.actor._extension.monarch_hyperactor.alloc"
)]
#[pyclass(name = "Alloc", module = "monarch._src.actor._extension.alloc")]
pub struct PyAlloc {
pub inner: Arc<Mutex<Option<PyAllocWrapper>>>,
}
Expand Down Expand Up @@ -399,7 +396,7 @@ impl Alloc for PyAllocWrapper {

#[pyclass(
name = "AllocConstraints",
module = "monarch._src.actor._extension.monarch_hyperactor.alloc"
module = "monarch._src.actor._extension.alloc"
)]
pub struct PyAllocConstraints {
inner: AllocConstraints,
Expand All @@ -419,10 +416,7 @@ impl PyAllocConstraints {
}
}

#[pyclass(
name = "AllocSpec",
module = "monarch._src.actor._extension.monarch_hyperactor.alloc"
)]
#[pyclass(name = "AllocSpec", module = "monarch._src.actor._extension.alloc")]
pub struct PyAllocSpec {
pub inner: AllocSpec,
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,10 +38,7 @@ pub fn bootstrap_main(py: Python) -> PyResult<Bound<PyAny>> {

pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
let f = wrap_pyfunction!(bootstrap_main, hyperactor_mod)?;
f.setattr(
"__module__",
"monarch._src.actor._extension.monarch_hyperactor.bootstrap",
)?;
f.setattr("__module__", "monarch._src.actor._extension.bootstrap")?;
hyperactor_mod.add_function(f)?;

Ok(())
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ use pyo3::prelude::*;
/// Python binding for [`hyperactor::channel::ChannelTransport`]
#[pyclass(
name = "ChannelTransport",
module = "monarch._src.actor._extension.monarch_hyperactor.channel",
module = "monarch._src.actor._extension.channel",
eq
)]
#[derive(PartialEq, Clone, Copy, Debug)]
Expand All @@ -28,10 +28,7 @@ pub enum PyChannelTransport {
// Sim(/*proxy address:*/ ChannelAddr), TODO kiuk@ add support
}

#[pyclass(
name = "ChannelAddr",
module = "monarch._src.actor._extension.monarch_hyperactor.channel"
)]
#[pyclass(name = "ChannelAddr", module = "monarch._src.actor._extension.channel")]
pub struct PyChannelAddr {
inner: ChannelAddr,
}
Expand Down
7 changes: 4 additions & 3 deletions actor_extension/src/code_sync.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,6 @@ use hyperactor_mesh::code_sync::WorkspaceLocation;
use hyperactor_mesh::code_sync::rsync;
use hyperactor_mesh::shape::Shape;
use hyperactor_mesh::shared_cell::SharedCell;
use monarch_hyperactor::proc_mesh::PyProcMesh;
use monarch_hyperactor::runtime::signal_safe_block_on;
use monarch_hyperactor::shape::PyShape;
use pyo3::Bound;
use pyo3::exceptions::PyRuntimeError;
use pyo3::exceptions::PyValueError;
Expand All @@ -28,6 +25,10 @@ use pyo3::types::PyModule;
use serde::Deserialize;
use serde::Serialize;

use crate::proc_mesh::PyProcMesh;
use crate::runtime::signal_safe_block_on;
use crate::shape::PyShape;

#[pyclass(
frozen,
name = "WorkspaceLocation",
Expand Down
File renamed without changes.
86 changes: 30 additions & 56 deletions actor_extension/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,22 @@
//! It is imported by `monarch` as `monarch._src.actor._extension`.
use pyo3::prelude::*;

pub mod actor;
mod actor_mesh;
mod alloc;
mod blocking;
mod bootstrap;
mod channel;
mod code_sync;
mod config;
pub mod mailbox;
pub mod ndslice;
mod panic;
pub mod proc;
pub mod proc_mesh;
pub mod runtime;
mod selection;
pub mod shape;
mod telemetry;

#[cfg(fbcode_build)]
Expand Down Expand Up @@ -51,65 +64,26 @@ fn get_or_add_new_module<'py>(
#[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",
)?)?;
crate::runtime::initialize(module.py())?;
let runtime = crate::runtime::get_tokio_runtime();

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",
)?)?;
telemetry::register_python_bindings(&get_or_add_new_module(module, "telemetry")?)?;
crate::panic::register_python_bindings(&get_or_add_new_module(module, "panic")?)?;
::hyperactor::initialize(runtime.handle().clone());

crate::actor_mesh::register_python_bindings(&get_or_add_new_module(module, "actor_mesh")?)?;
crate::actor::register_python_bindings(&get_or_add_new_module(module, "actor")?)?;
crate::alloc::register_python_bindings(&get_or_add_new_module(module, "alloc")?)?;
crate::blocking::register_python_bindings(&get_or_add_new_module(module, "blocking")?)?;
crate::bootstrap::register_python_bindings(&get_or_add_new_module(module, "bootstrap")?)?;
crate::channel::register_python_bindings(&get_or_add_new_module(module, "channel")?)?;
crate::code_sync::register_python_bindings(&get_or_add_new_module(module, "code_sync")?)?;
crate::mailbox::register_python_bindings(&get_or_add_new_module(module, "mailbox")?)?;
crate::panic::register_python_bindings(&get_or_add_new_module(module, "panic")?)?;
crate::proc_mesh::register_python_bindings(&get_or_add_new_module(module, "proc_mesh")?)?;
crate::proc::register_python_bindings(&get_or_add_new_module(module, "proc")?)?;
crate::runtime::register_python_bindings(&get_or_add_new_module(module, "runtime")?)?;
crate::selection::register_python_bindings(&get_or_add_new_module(module, "selection")?)?;
crate::shape::register_python_bindings(&get_or_add_new_module(module, "shape")?)?;
crate::telemetry::register_python_bindings(&get_or_add_new_module(module, "telemetry")?)?;

#[cfg(fbcode_build)]
crate::meta::register_python_bindings(&get_or_add_new_module(module, "meta")?)?;
Expand Down
Loading
Loading