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

monarch_hyperactor/
proc_mesh.rs
