diff --git a/actor_extension/Cargo.toml b/actor_extension/Cargo.toml index 0dc4fa80..1ceddbb2 100644 --- a/actor_extension/Cargo.toml +++ b/actor_extension/Cargo.toml @@ -16,11 +16,12 @@ crate-type = ["cdylib"] 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" } +hyperactor_telemetry = { version = "0.0.0", path = "../hyperactor_telemetry" } 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"] } +tracing = { version = "0.1.41", features = ["attributes", "valuable"] } diff --git a/actor_extension/src/lib.rs b/actor_extension/src/lib.rs index f6b47beb..365396f7 100644 --- a/actor_extension/src/lib.rs +++ b/actor_extension/src/lib.rs @@ -15,6 +15,7 @@ use pyo3::prelude::*; mod blocking; mod code_sync; mod panic; +mod telemetry; #[cfg(fbcode_build)] mod meta; @@ -105,14 +106,7 @@ pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> { 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", - )?)?; + telemetry::register_python_bindings(&get_or_add_new_module(module, "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")?)?; diff --git a/hyperactor_extension/src/telemetry.rs b/actor_extension/src/telemetry.rs similarity index 85% rename from hyperactor_extension/src/telemetry.rs rename to actor_extension/src/telemetry.rs index d3eb6b78..03f18127 100644 --- a/hyperactor_extension/src/telemetry.rs +++ b/actor_extension/src/telemetry.rs @@ -113,7 +113,7 @@ pub fn use_sim_clock() -> PyResult<()> { #[pyclass( unsendable, subclass, - module = "monarch._src.actor._extension.hyperactor_extension.telemetry" + module = "monarch._src.actor._extension.telemetry" )] struct PySpan { span: tracing::span::EnteredSpan, @@ -140,46 +140,28 @@ use pyo3::types::PyModule; pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> { // Register the forward_to_tracing function let f = wrap_pyfunction!(forward_to_tracing, module)?; - f.setattr( - "__module__", - "monarch._src.actor._extension.hyperactor_extension.telemetry", - )?; + f.setattr("__module__", "monarch._src.actor._extension.telemetry")?; module.add_function(f)?; // Register the span-related functions let enter_span_fn = wrap_pyfunction!(enter_span, module)?; - enter_span_fn.setattr( - "__module__", - "monarch._src.actor._extension.hyperactor_extension.telemetry", - )?; + enter_span_fn.setattr("__module__", "monarch._src.actor._extension.telemetry")?; module.add_function(enter_span_fn)?; let exit_span_fn = wrap_pyfunction!(exit_span, module)?; - exit_span_fn.setattr( - "__module__", - "monarch._src.actor._extension.hyperactor_extension.telemetry", - )?; + exit_span_fn.setattr("__module__", "monarch._src.actor._extension.telemetry")?; module.add_function(exit_span_fn)?; let get_current_span_id_fn = wrap_pyfunction!(get_current_span_id, module)?; - get_current_span_id_fn.setattr( - "__module__", - "monarch._src.actor._extension.hyperactor_extension.telemetry", - )?; + get_current_span_id_fn.setattr("__module__", "monarch._src.actor._extension.telemetry")?; module.add_function(get_current_span_id_fn)?; let use_real_clock_fn = wrap_pyfunction!(use_real_clock, module)?; - use_real_clock_fn.setattr( - "__module__", - "monarch._src.actor._extension.hyperactor_extension.telemetry", - )?; + use_real_clock_fn.setattr("__module__", "monarch._src.actor._extension.telemetry")?; module.add_function(use_real_clock_fn)?; let use_sim_clock_fn = wrap_pyfunction!(use_sim_clock, module)?; - use_sim_clock_fn.setattr( - "__module__", - "monarch._src.actor._extension.hyperactor_extension.telemetry", - )?; + use_sim_clock_fn.setattr("__module__", "monarch._src.actor._extension.telemetry")?; module.add_function(use_sim_clock_fn)?; module.add_class::()?; diff --git a/hyperactor_extension/Cargo.toml b/hyperactor_extension/Cargo.toml deleted file mode 100644 index 490b1862..00000000 --- a/hyperactor_extension/Cargo.toml +++ /dev/null @@ -1,17 +0,0 @@ -# @generated by autocargo from //monarch/hyperactor_extension:hyperactor_extension - -[package] -name = "hyperactor_extension" -version = "0.0.0" -authors = ["Meta"] -edition = "2021" -license = "BSD-3-Clause" - -[dependencies] -async-trait = "0.1.86" -hyperactor = { version = "0.0.0", path = "../hyperactor" } -hyperactor_mesh = { version = "0.0.0", path = "../hyperactor_mesh" } -hyperactor_telemetry = { version = "0.0.0", path = "../hyperactor_telemetry" } -ndslice = { version = "0.0.0", path = "../ndslice" } -pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods"] } -tracing = { version = "0.1.41", features = ["attributes", "valuable"] } diff --git a/hyperactor_extension/src/alloc.rs b/hyperactor_extension/src/alloc.rs deleted file mode 100644 index cef50424..00000000 --- a/hyperactor_extension/src/alloc.rs +++ /dev/null @@ -1,169 +0,0 @@ -/* - * 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. - */ - -use std::collections::HashMap; -use std::sync::Arc; -use std::sync::Mutex; - -use async_trait::async_trait; -use hyperactor::WorldId; -use hyperactor::channel::ChannelTransport; -use hyperactor_mesh::alloc::Alloc; -use hyperactor_mesh::alloc::AllocConstraints; -use hyperactor_mesh::alloc::AllocSpec; -use hyperactor_mesh::alloc::AllocatorError; -use hyperactor_mesh::alloc::ProcState; -use hyperactor_mesh::shape::Shape; -use ndslice::Slice; -use pyo3::exceptions::PyValueError; -use pyo3::prelude::*; -use pyo3::types::PyDict; - -/// 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.hyperactor_extension.alloc" -)] -pub struct PyAlloc { - pub inner: Arc>>, -} - -impl PyAlloc { - /// Create a new PyAlloc with provided boxed trait. - pub fn new(inner: Box) -> Self { - Self { - inner: Arc::new(Mutex::new(Some(PyAllocWrapper { inner }))), - } - } - - /// Take the internal Alloc object. - pub fn take(&self) -> Option { - self.inner.lock().unwrap().take() - } -} - -#[pymethods] -impl PyAlloc { - fn __repr__(&self) -> PyResult { - let data = self.inner.lock().unwrap(); - match &*data { - None => Ok("Alloc(None)".to_string()), - Some(wrapper) => Ok(format!("Alloc({})", wrapper.shape())), - } - } -} - -/// Internal wrapper to translate from a dyn Alloc to an impl Alloc. Used -/// to support polymorphism in the Python bindings. -pub struct PyAllocWrapper { - inner: Box, -} - -#[async_trait] -impl Alloc for PyAllocWrapper { - async fn next(&mut self) -> Option { - self.inner.next().await - } - - fn shape(&self) -> &Shape { - self.inner.shape() - } - - fn world_id(&self) -> &WorldId { - self.inner.world_id() - } - - fn transport(&self) -> ChannelTransport { - self.inner.transport() - } - - async fn stop(&mut self) -> Result<(), AllocatorError> { - self.inner.stop().await - } -} - -#[pyclass( - name = "AllocConstraints", - module = "monarch._src.actor._extension.hyperactor_extension.alloc" -)] -pub struct PyAllocConstraints { - inner: AllocConstraints, -} - -#[pymethods] -impl PyAllocConstraints { - #[new] - #[pyo3(signature = (match_labels=None))] - fn new(match_labels: Option>) -> PyResult { - let mut constraints = AllocConstraints::default(); - if let Some(match_lables) = match_labels { - constraints.match_labels = match_lables; - } - - Ok(Self { inner: constraints }) - } -} - -#[pyclass( - name = "AllocSpec", - module = "monarch._src.actor._extension.hyperactor_extension.alloc" -)] -pub struct PyAllocSpec { - pub inner: AllocSpec, -} - -#[pymethods] -impl PyAllocSpec { - #[new] - #[pyo3(signature = (constraints, **kwargs))] - fn new(constraints: &PyAllocConstraints, kwargs: Option<&Bound<'_, PyAny>>) -> PyResult { - let Some(kwargs) = kwargs else { - return Err(PyValueError::new_err( - "Shape must have at least one dimension", - )); - }; - let shape_dict = kwargs.downcast::()?; - - let mut keys = Vec::new(); - let mut values = Vec::new(); - for (key, value) in shape_dict { - keys.push(key.clone()); - values.push(value.clone()); - } - - let shape = Shape::new( - keys.into_iter() - .map(|key| key.extract::()) - .collect::>>()?, - Slice::new_row_major( - values - .into_iter() - .map(|key| key.extract::()) - .collect::>>()?, - ), - ) - .map_err(|e| PyValueError::new_err(format!("Invalid shape: {:?}", e)))?; - - Ok(Self { - inner: AllocSpec { - shape, - constraints: constraints.inner.clone(), - }, - }) - } -} - -pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_class::()?; - module.add_class::()?; - module.add_class::()?; - - Ok(()) -} diff --git a/hyperactor_extension/src/lib.rs b/hyperactor_extension/src/lib.rs deleted file mode 100644 index 7fe6be48..00000000 --- a/hyperactor_extension/src/lib.rs +++ /dev/null @@ -1,12 +0,0 @@ -/* - * 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. - */ - -#![allow(unsafe_op_in_unsafe_fn)] - -pub mod alloc; -pub mod telemetry; diff --git a/monarch_hyperactor/Cargo.toml b/monarch_hyperactor/Cargo.toml index af9962b3..1cce78d4 100644 --- a/monarch_hyperactor/Cargo.toml +++ b/monarch_hyperactor/Cargo.toml @@ -15,7 +15,6 @@ clap = { version = "4.5.38", features = ["derive", "env", "string", "unicode", " 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_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" } hyperactor_telemetry = { version = "0.0.0", path = "../hyperactor_telemetry" } diff --git a/monarch_hyperactor/src/alloc.rs b/monarch_hyperactor/src/alloc.rs index 52ff9969..b36c87db 100644 --- a/monarch_hyperactor/src/alloc.rs +++ b/monarch_hyperactor/src/alloc.rs @@ -9,6 +9,7 @@ use std::collections::HashMap; use std::str::FromStr; use std::sync::Arc; +use std::sync::Mutex; use std::time::Duration; use anyhow::anyhow; @@ -16,19 +17,23 @@ use async_trait::async_trait; use hyperactor::WorldId; use hyperactor::channel::ChannelAddr; use hyperactor::channel::ChannelTransport; -use hyperactor_extension::alloc::PyAlloc; -use hyperactor_extension::alloc::PyAllocSpec; +use hyperactor_mesh::alloc::Alloc; use hyperactor_mesh::alloc::AllocConstraints; use hyperactor_mesh::alloc::AllocSpec; use hyperactor_mesh::alloc::Allocator; use hyperactor_mesh::alloc::AllocatorError; use hyperactor_mesh::alloc::LocalAllocator; +use hyperactor_mesh::alloc::ProcState; use hyperactor_mesh::alloc::ProcessAllocator; use hyperactor_mesh::alloc::remoteprocess::RemoteProcessAlloc; use hyperactor_mesh::alloc::remoteprocess::RemoteProcessAllocHost; use hyperactor_mesh::alloc::remoteprocess::RemoteProcessAllocInitializer; +use ndslice::Shape; +use ndslice::Slice; use pyo3::exceptions::PyRuntimeError; +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; +use pyo3::types::PyDict; use tokio::process::Command; use crate::channel::PyChannelAddr; @@ -327,10 +332,149 @@ 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" +)] +pub struct PyAlloc { + pub inner: Arc>>, +} + +impl PyAlloc { + /// Create a new PyAlloc with provided boxed trait. + pub fn new(inner: Box) -> Self { + Self { + inner: Arc::new(Mutex::new(Some(PyAllocWrapper { inner }))), + } + } + + /// Take the internal Alloc object. + pub fn take(&self) -> Option { + self.inner.lock().unwrap().take() + } +} + +#[pymethods] +impl PyAlloc { + fn __repr__(&self) -> PyResult { + let data = self.inner.lock().unwrap(); + match &*data { + None => Ok("Alloc(None)".to_string()), + Some(wrapper) => Ok(format!("Alloc({})", wrapper.shape())), + } + } +} + +/// Internal wrapper to translate from a dyn Alloc to an impl Alloc. Used +/// to support polymorphism in the Python bindings. +pub struct PyAllocWrapper { + inner: Box, +} + +#[async_trait] +impl Alloc for PyAllocWrapper { + async fn next(&mut self) -> Option { + self.inner.next().await + } + + fn shape(&self) -> &Shape { + self.inner.shape() + } + + fn world_id(&self) -> &WorldId { + self.inner.world_id() + } + + fn transport(&self) -> ChannelTransport { + self.inner.transport() + } + + async fn stop(&mut self) -> Result<(), AllocatorError> { + self.inner.stop().await + } +} + +#[pyclass( + name = "AllocConstraints", + module = "monarch._src.actor._extension.monarch_hyperactor.alloc" +)] +pub struct PyAllocConstraints { + inner: AllocConstraints, +} + +#[pymethods] +impl PyAllocConstraints { + #[new] + #[pyo3(signature = (match_labels=None))] + fn new(match_labels: Option>) -> PyResult { + let mut constraints = AllocConstraints::default(); + if let Some(match_lables) = match_labels { + constraints.match_labels = match_lables; + } + + Ok(Self { inner: constraints }) + } +} + +#[pyclass( + name = "AllocSpec", + module = "monarch._src.actor._extension.monarch_hyperactor.alloc" +)] +pub struct PyAllocSpec { + pub inner: AllocSpec, +} + +#[pymethods] +impl PyAllocSpec { + #[new] + #[pyo3(signature = (constraints, **kwargs))] + fn new(constraints: &PyAllocConstraints, kwargs: Option<&Bound<'_, PyAny>>) -> PyResult { + let Some(kwargs) = kwargs else { + return Err(PyValueError::new_err( + "Shape must have at least one dimension", + )); + }; + let shape_dict = kwargs.downcast::()?; + + let mut keys = Vec::new(); + let mut values = Vec::new(); + for (key, value) in shape_dict { + keys.push(key.clone()); + values.push(value.clone()); + } + + let shape = Shape::new( + keys.into_iter() + .map(|key| key.extract::()) + .collect::>>()?, + Slice::new_row_major( + values + .into_iter() + .map(|key| key.extract::()) + .collect::>>()?, + ), + ) + .map_err(|e| PyValueError::new_err(format!("Invalid shape: {:?}", e)))?; + + Ok(Self { + inner: AllocSpec { + shape, + constraints: constraints.inner.clone(), + }, + }) + } +} + pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> { hyperactor_mod.add_class::()?; hyperactor_mod.add_class::()?; hyperactor_mod.add_class::()?; + hyperactor_mod.add_class::()?; + hyperactor_mod.add_class::()?; + hyperactor_mod.add_class::()?; Ok(()) } diff --git a/monarch_hyperactor/src/proc_mesh.rs b/monarch_hyperactor/src/proc_mesh.rs index b0f9e239..20f47576 100644 --- a/monarch_hyperactor/src/proc_mesh.rs +++ b/monarch_hyperactor/src/proc_mesh.rs @@ -17,7 +17,6 @@ use hyperactor::RemoteMessage; use hyperactor::WorldId; use hyperactor::actor::RemoteActor; use hyperactor::proc::Proc; -use hyperactor_extension::alloc::PyAlloc; use hyperactor_mesh::RootActorMesh; use hyperactor_mesh::alloc::Alloc; use hyperactor_mesh::alloc::ProcStopReason; @@ -40,6 +39,7 @@ use tokio::sync::Mutex; use tokio::sync::mpsc; use crate::actor_mesh::PythonActorMesh; +use crate::alloc::PyAlloc; use crate::mailbox::PyMailbox; use crate::runtime::signal_safe_block_on; use crate::shape::PyShape; diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/alloc.pyi b/python/monarch/_rust_bindings/monarch_hyperactor/alloc.pyi deleted file mode 100644 index 0b0d854b..00000000 --- a/python/monarch/_rust_bindings/monarch_hyperactor/alloc.pyi +++ /dev/null @@ -1,107 +0,0 @@ -# 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/old.pyi b/python/monarch/_rust_bindings/old.pyi index d90fc9f3..f045ee7a 100644 --- a/python/monarch/_rust_bindings/old.pyi +++ b/python/monarch/_rust_bindings/old.pyi @@ -8,7 +8,7 @@ from typing import Dict, final, List, Optional, Protocol, Type -from monarch._src.actor._extension.hyperactor_extension.alloc import Alloc, AllocSpec +from monarch._src.actor._extension.monarch_hyperactor.alloc import Alloc, AllocSpec from monarch._src.actor._extension.monarch_hyperactor.selection import Selection diff --git a/python/monarch/_src/actor/_extension/hyperactor_extension/alloc.pyi b/python/monarch/_src/actor/_extension/hyperactor_extension/alloc.pyi deleted file mode 100644 index fc4d86d4..00000000 --- a/python/monarch/_src/actor/_extension/hyperactor_extension/alloc.pyi +++ /dev/null @@ -1,38 +0,0 @@ -# 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 typing import Dict, final, Optional - -class Alloc: - """ - An alloc represents an allocation of procs. Allocs are returned by - one of the allocator implementations, such as `ProcessAllocator` or - `LocalAllocator`. - """ - -@final -class AllocConstraints: - def __init__(self, match_labels: Optional[Dict[str, str]] = None) -> None: - """ - Create a new alloc constraints. - - Arguments: - - `match_labels`: A dictionary of labels to match. If a label is present - in the dictionary, the alloc must have that label and its value - must match the value in the dictionary. - """ - ... - -@final -class AllocSpec: - def __init__(self, constraints: AllocConstraints, **kwargs: int) -> None: - """ - Initialize a shape with the provided dimension-size pairs. - For example, `AllocSpec(constraints, replica=2, host=3, gpu=8)` creates a - shape with 2 replicas with 3 hosts each, each of which in turn - has 8 GPUs. - """ - ... diff --git a/python/monarch/_src/actor/_extension/monarch_hyperactor/alloc.pyi b/python/monarch/_src/actor/_extension/monarch_hyperactor/alloc.pyi index 0b0d854b..576be754 100644 --- a/python/monarch/_src/actor/_extension/monarch_hyperactor/alloc.pyi +++ b/python/monarch/_src/actor/_extension/monarch_hyperactor/alloc.pyi @@ -7,11 +7,41 @@ # pyre-strict from datetime import timedelta -from typing import Optional +from typing import final, Optional -from monarch._src.actor._extension.hyperactor_extension.alloc import Alloc, AllocSpec from typing_extensions import Self +class Alloc: + """ + An alloc represents an allocation of procs. Allocs are returned by + one of the allocator implementations, such as `ProcessAllocator` or + `LocalAllocator`. + """ + +@final +class AllocConstraints: + def __init__(self, match_labels: Optional[dict[str, str]] = None) -> None: + """ + Create a new alloc constraints. + + Arguments: + - `match_labels`: A dictionary of labels to match. If a label is present + in the dictionary, the alloc must have that label and its value + must match the value in the dictionary. + """ + ... + +@final +class AllocSpec: + def __init__(self, constraints: AllocConstraints, **kwargs: int) -> None: + """ + Initialize a shape with the provided dimension-size pairs. + For example, `AllocSpec(constraints, replica=2, host=3, gpu=8)` creates a + shape with 2 replicas with 3 hosts each, each of which in turn + has 8 GPUs. + """ + ... + class ProcessAllocatorBase: def __init__( self, diff --git a/python/monarch/_src/actor/_extension/monarch_hyperactor/proc_mesh.pyi b/python/monarch/_src/actor/_extension/monarch_hyperactor/proc_mesh.pyi index 5e91715c..aa072dd8 100644 --- a/python/monarch/_src/actor/_extension/monarch_hyperactor/proc_mesh.pyi +++ b/python/monarch/_src/actor/_extension/monarch_hyperactor/proc_mesh.pyi @@ -8,10 +8,10 @@ from typing import AsyncIterator, final, Type -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.alloc import Alloc from monarch._src.actor._extension.monarch_hyperactor.mailbox import Mailbox from monarch._src.actor._extension.monarch_hyperactor.shape import Shape diff --git a/python/monarch/_src/actor/_extension/hyperactor_extension/telemetry.pyi b/python/monarch/_src/actor/_extension/telemetry.pyi similarity index 100% rename from python/monarch/_src/actor/_extension/hyperactor_extension/telemetry.pyi rename to python/monarch/_src/actor/_extension/telemetry.pyi diff --git a/python/monarch/_src/actor/actor_mesh.py b/python/monarch/_src/actor/actor_mesh.py index db650495..b1b5fb8e 100644 --- a/python/monarch/_src/actor/actor_mesh.py +++ b/python/monarch/_src/actor/actor_mesh.py @@ -38,11 +38,6 @@ TypeVar, ) -from monarch._src.actor._extension.hyperactor_extension.telemetry import ( - enter_span, - exit_span, -) - from monarch._src.actor._extension.monarch_hyperactor.actor import ( PanicFlag, PythonMessage, @@ -61,6 +56,8 @@ Shape, ) +from monarch._src.actor._extension.telemetry import enter_span, exit_span + from monarch._src.actor.allocator import LocalAllocator, ProcessAllocator from monarch._src.actor.future import Future from monarch._src.actor.pdb_wrapper import remote_breakpointhook diff --git a/python/monarch/_src/actor/allocator.py b/python/monarch/_src/actor/allocator.py index 4261cb23..116b3648 100644 --- a/python/monarch/_src/actor/allocator.py +++ b/python/monarch/_src/actor/allocator.py @@ -10,9 +10,9 @@ import logging from typing import final, Optional -from monarch._src.actor._extension.hyperactor_extension.alloc import Alloc, AllocSpec - from monarch._src.actor._extension.monarch_hyperactor.alloc import ( + Alloc, + AllocSpec, LocalAllocatorBase, ProcessAllocatorBase, RemoteAllocatorBase, diff --git a/python/monarch/_src/actor/proc_mesh.py b/python/monarch/_src/actor/proc_mesh.py index dcde3309..66872879 100644 --- a/python/monarch/_src/actor/proc_mesh.py +++ b/python/monarch/_src/actor/proc_mesh.py @@ -22,7 +22,7 @@ TypeVar, ) -from monarch._src.actor._extension.hyperactor_extension.alloc import ( # @manual=//monarch/actor_extension:actor_extension +from monarch._src.actor._extension.monarch_hyperactor.alloc import ( # @manual=//monarch/actor_extension:actor_extension Alloc, AllocConstraints, AllocSpec, diff --git a/python/monarch/_src/actor/telemetry/__init__.py b/python/monarch/_src/actor/telemetry/__init__.py index 7ffe0979..b80bc5cf 100644 --- a/python/monarch/_src/actor/telemetry/__init__.py +++ b/python/monarch/_src/actor/telemetry/__init__.py @@ -9,9 +9,7 @@ import logging -from monarch._src.actor._extension.hyperactor_extension.telemetry import ( - forward_to_tracing, -) +from monarch._src.actor._extension.telemetry import forward_to_tracing class TracingForwarder(logging.Handler): diff --git a/python/monarch/_src/actor/telemetry/rust_span_tracing.py b/python/monarch/_src/actor/telemetry/rust_span_tracing.py index 41ca30d9..66c39e2e 100644 --- a/python/monarch/_src/actor/telemetry/rust_span_tracing.py +++ b/python/monarch/_src/actor/telemetry/rust_span_tracing.py @@ -12,10 +12,7 @@ import opentelemetry.util.types as types # @manual=fbsource//third-party/pypi/opentelemetry-api:opentelemetry-api -from monarch._src.actor._extension.hyperactor_extension.telemetry import ( - get_current_span_id, - PySpan, -) +from monarch._src.actor._extension.telemetry import get_current_span_id, PySpan from opentelemetry import ( # @manual=fbsource//third-party/pypi/opentelemetry-api:opentelemetry-api trace, diff --git a/python/tests/_monarch/test_hyperactor.py b/python/tests/_monarch/test_hyperactor.py index ef3a767e..25631f13 100644 --- a/python/tests/_monarch/test_hyperactor.py +++ b/python/tests/_monarch/test_hyperactor.py @@ -13,15 +13,15 @@ import monarch -from monarch._src.actor._extension.hyperactor_extension.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension - AllocConstraints, - AllocSpec, -) - from monarch._src.actor._extension.monarch_hyperactor.actor import ( PanicFlag, PythonMessage, ) + +from monarch._src.actor._extension.monarch_hyperactor.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension + AllocConstraints, + AllocSpec, +) 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 diff --git a/python/tests/_monarch/test_mailbox.py b/python/tests/_monarch/test_mailbox.py index c025f4e1..60f66e2c 100644 --- a/python/tests/_monarch/test_mailbox.py +++ b/python/tests/_monarch/test_mailbox.py @@ -12,16 +12,16 @@ import monarch -from monarch._src.actor._extension.hyperactor_extension.alloc import ( - AllocConstraints, - AllocSpec, -) - from monarch._src.actor._extension.monarch_hyperactor.actor import ( PanicFlag, PythonMessage, ) +from monarch._src.actor._extension.monarch_hyperactor.alloc import ( + AllocConstraints, + AllocSpec, +) + from monarch._src.actor._extension.monarch_hyperactor.mailbox import ( Mailbox, PortReceiver, diff --git a/python/tests/test_alloc.py b/python/tests/test_alloc.py index e8d97f07..d90a36c3 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._src.actor._extension.hyperactor_extension.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension +from monarch._src.actor._extension.monarch_hyperactor.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 d1acedc0..c18680c0 100644 --- a/python/tests/test_allocator.py +++ b/python/tests/test_allocator.py @@ -24,7 +24,7 @@ import torch.distributed as dist import torch.nn.functional as F -from monarch._src.actor._extension.hyperactor_extension.alloc import ( +from monarch._src.actor._extension.monarch_hyperactor.alloc import ( AllocConstraints, AllocSpec, )