monarch_hyperactor/
proc_mesh.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use std::fmt::Debug;
10use std::ops::Deref;
11
12use hyperactor::id::Label;
13use hyperactor_mesh::ProcMesh;
14use hyperactor_mesh::ProcMeshRef;
15use hyperactor_mesh::mesh_id::ActorMeshId;
16use hyperactor_mesh::shared_cell::SharedCell;
17use monarch_types::PickledPyObject;
18use monarch_types::py_module_add_function;
19use ndslice::View;
20use ndslice::view::RankedSliceable;
21use pyo3::IntoPyObjectExt;
22use pyo3::exceptions::PyRuntimeError;
23use pyo3::exceptions::PyValueError;
24use pyo3::prelude::*;
25use pyo3::types::PyBytes;
26use pyo3::types::PyType;
27
28use crate::actor::ActorKind;
29use crate::actor::PythonActorParams;
30use crate::actor_mesh::PythonActorMesh;
31use crate::actor_mesh::PythonActorMeshImpl;
32use crate::actor_mesh::SupervisableActorMesh;
33use crate::context::PyInstance;
34use crate::handle::PyHandle;
35use crate::pickle::PendingMessage;
36use crate::pytokio::PyShared;
37use crate::runtime::GilSite;
38use crate::runtime::get_tokio_runtime;
39use crate::runtime::monarch_with_gil;
40use crate::runtime::monarch_with_gil_blocking;
41use crate::shape::PyRegion;
42
43#[pyclass(
44    name = "ProcMesh",
45    module = "monarch._rust_bindings.monarch_hyperactor.proc_mesh"
46)]
47#[expect(
48    clippy::large_enum_variant,
49    reason = "PyO3 #[pyclass] enum; Box wrapping interacts with PyO3 codegen and Python interop — separate diff"
50)]
51pub enum PyProcMesh {
52    Owned(PyProcMeshImpl),
53    Ref(PyProcMeshRefImpl),
54}
55
56impl PyProcMesh {
57    pub fn new_owned(inner: ProcMesh) -> Self {
58        Self::Owned(PyProcMeshImpl(inner.into()))
59    }
60
61    pub(crate) fn new_ref(inner: ProcMeshRef) -> Self {
62        Self::Ref(PyProcMeshRefImpl(inner))
63    }
64
65    pub fn mesh_ref(&self) -> PyResult<ProcMeshRef> {
66        match self {
67            PyProcMesh::Owned(inner) => Ok(inner
68                .0
69                .borrow()
70                .map_err(|_| PyRuntimeError::new_err("`ProcMesh` has already been stopped"))?
71                .clone()),
72            PyProcMesh::Ref(inner) => Ok(inner.0.clone()),
73        }
74    }
75}
76
77#[pymethods]
78impl PyProcMesh {
79    #[staticmethod]
80    #[pyo3(signature = (proc_mesh, instance, mesh_base_name, actor, init_message, kind, emulated, supervision_display_name = None))]
81    fn spawn_async(
82        proc_mesh: &mut PyShared,
83        instance: &PyInstance,
84        mesh_base_name: String,
85        actor: Py<PyType>,
86        init_message: &mut PendingMessage,
87        kind: ActorKind,
88        emulated: bool,
89        supervision_display_name: Option<String>,
90    ) -> PyResult<Py<PyAny>> {
91        let init_message = init_message.take()?;
92        let task = proc_mesh.task()?.take_task()?;
93        let instance = instance.clone();
94        let mesh_impl = async move {
95            let proc_mesh = task.await?;
96
97            let init_message = init_message.resolve().await?;
98
99            let (proc_mesh, params) =
100                monarch_with_gil(GilSite::ActorConstruct, |py| -> PyResult<_> {
101                    let slf: Bound<PyProcMesh> = proc_mesh.extract(py)?;
102                    let slf = slf.borrow();
103                    let pickled_type = PickledPyObject::pickle(actor.bind(py).as_any())?;
104                    Ok((
105                        slf.mesh_ref()?.clone(),
106                        PythonActorParams::new(pickled_type, Some(init_message), kind),
107                    ))
108                })
109                .await?;
110
111            let mesh_name = ActorMeshId::instance(Label::strip(&mesh_base_name));
112            let actor_mesh = proc_mesh
113                .spawn_with_name(
114                    instance.deref(),
115                    mesh_name,
116                    &params,
117                    supervision_display_name,
118                    false,
119                )
120                .await
121                .map_err(anyhow::Error::from)?;
122            Ok::<_, PyErr>(Box::new(PythonActorMeshImpl::new_owned(actor_mesh)))
123        };
124        if emulated {
125            // we give up on doing mesh spawn async for the emulated old version
126            // it is too complicated to make both work.
127            let r = get_tokio_runtime().block_on(mesh_impl)?;
128            monarch_with_gil_blocking(GilSite::Convert, |py| r.into_py_any(py))
129        } else {
130            let r = PythonActorMesh::new(
131                async move {
132                    let mesh_impl: Box<dyn SupervisableActorMesh> = mesh_impl.await?;
133                    Ok(mesh_impl)
134                },
135                true,
136            );
137            monarch_with_gil_blocking(GilSite::Convert, |py| r.into_py_any(py))
138        }
139    }
140
141    fn __repr__(&self) -> PyResult<String> {
142        match self {
143            PyProcMesh::Owned(inner) => Ok(format!("<ProcMesh: {:?}>", inner.__repr__()?)),
144            PyProcMesh::Ref(inner) => Ok(format!("<ProcMesh: {:?}>", inner.__repr__()?)),
145        }
146    }
147
148    fn __reduce__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)> {
149        let mesh_ref = self.mesh_ref()?;
150        if crate::pickle::push_mesh_reference_if_active(crate::actor::MeshRef::Proc(Box::new(
151            mesh_ref.clone(),
152        ))) {
153            let pop_fn = PyModule::import(py, "monarch._rust_bindings.monarch_hyperactor.pickle")?
154                .getattr("pop_mesh_reference")?;
155            return Ok((pop_fn, pyo3::types::PyTuple::empty(py).into_any()));
156        }
157        let bytes = bincode::serde::encode_to_vec(&mesh_ref, bincode::config::legacy())
158            .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))?;
159        let py_bytes = (PyBytes::new(py, &bytes),).into_bound_py_any(py).unwrap();
160        let from_bytes =
161            PyModule::import(py, "monarch._rust_bindings.monarch_hyperactor.proc_mesh")?
162                .getattr("py_proc_mesh_from_bytes")?;
163        Ok((from_bytes, py_bytes))
164    }
165
166    #[getter]
167    fn region(&self) -> PyResult<PyRegion> {
168        Ok(self.mesh_ref()?.region().into())
169    }
170
171    fn stop_nonblocking(&self, instance: &PyInstance, reason: String) -> PyResult<PyHandle> {
172        // Clone the necessary fields from self to avoid capturing self in the async block
173        let (owned_inner, instance) = monarch_with_gil_blocking(GilSite::Stop, |_py| {
174            let owned_inner = match self {
175                PyProcMesh::Owned(inner) => inner.clone(),
176                PyProcMesh::Ref(_) => {
177                    return Err(PyValueError::new_err(
178                        "ProcMesh is not owned; must be stopped by an owner",
179                    ));
180                }
181            };
182
183            let instance = instance.clone();
184            Ok((owned_inner, instance))
185        })?;
186        Ok(PyHandle::spawn(async move {
187            let mesh = owned_inner.0.take().await;
188            match mesh {
189                Ok(mut mesh) => mesh
190                    .stop(instance.deref(), reason)
191                    .await
192                    .map_err(|e| PyValueError::new_err(format!("error stopping mesh: {}", e))),
193                Err(e) => {
194                    // Don't return an exception, silently ignore the stop request
195                    // because it was already done.
196                    tracing::info!("proc mesh already stopped: {}", e);
197                    Ok(())
198                }
199            }
200        }))
201    }
202
203    fn sliced(&self, region: &PyRegion) -> PyResult<Self> {
204        Ok(Self::new_ref(
205            self.mesh_ref()?.sliced(region.as_inner().clone()),
206        ))
207    }
208}
209
210#[derive(Clone)]
211#[pyclass(
212    name = "ProcMeshImpl",
213    module = "monarch._rust_bindings.monarch_hyperactor.proc_mesh"
214)]
215pub struct PyProcMeshImpl(SharedCell<ProcMesh>);
216
217impl PyProcMeshImpl {
218    fn __repr__(&self) -> PyResult<String> {
219        Ok(format!(
220            "<ProcMeshImpl {:?}>",
221            *self.0.borrow().map_err(anyhow::Error::from)?
222        ))
223    }
224}
225
226#[derive(Debug, Clone)]
227#[pyclass(
228    name = "ProcMeshRefImpl",
229    module = "monarch._rust_bindings.monarch_hyperactor.proc_mesh"
230)]
231pub struct PyProcMeshRefImpl(ProcMeshRef);
232
233impl PyProcMeshRefImpl {
234    fn __repr__(&self) -> PyResult<String> {
235        Ok(format!("<ProcMeshRefImpl {:?}>", self.0))
236    }
237}
238
239#[pyfunction]
240fn py_proc_mesh_from_bytes(bytes: &Bound<'_, PyBytes>) -> PyResult<PyProcMesh> {
241    let r: PyResult<ProcMeshRef> =
242        bincode::serde::decode_from_slice(bytes.as_bytes(), bincode::config::legacy())
243            .map(|(v, _)| v)
244            .map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()));
245    r.map(PyProcMesh::new_ref)
246}
247
248pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
249    hyperactor_mod.add_class::<PyProcMesh>()?;
250    py_module_add_function!(
251        hyperactor_mod,
252        "monarch._rust_bindings.monarch_hyperactor.proc_mesh",
253        py_proc_mesh_from_bytes
254    );
255    Ok(())
256}
monarch_hyperactor/proc_mesh.rs

monarch_hyperactor/
proc_mesh.rs
