monarch_hyperactor/

actor_mesh.rs

/*
 * 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::future::Future;
use std::ops::Deref;
use std::pin::Pin;
use std::sync::Arc;
use std::thread;
use std::time::Duration;

use async_trait::async_trait;
use futures::future;
use futures::future::FutureExt;
use futures::future::Shared;
use hyperactor::Instance;
use hyperactor::supervision::ActorSupervisionEvent;
use hyperactor_mesh::actor_mesh::ActorMesh;
use hyperactor_mesh::actor_mesh::ActorMeshRef;
use monarch_types::py_global;
use monarch_types::py_module_add_function;
use ndslice::view::Ranked;
use ndslice::view::RankedSliceable;
use pyo3::IntoPyObjectExt;
use pyo3::exceptions::PyNotImplementedError;
use pyo3::exceptions::PyRuntimeError;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::PyBytes;
use pyo3::types::PyTuple;
use tokio::sync::mpsc::UnboundedSender;
use tokio::sync::mpsc::unbounded_channel;
use tracing::Instrument;

use crate::actor::PythonActor;
use crate::actor::PythonMessage;
use crate::actor::PythonMessageKind;
use crate::context::PyInstance;
use crate::pickle::PendingMessage;
use crate::proc::PyActorAddr;
use crate::pytokio::PyPythonTask;
use crate::runtime::GilSite;
use crate::runtime::get_tokio_runtime;
use crate::runtime::monarch_with_gil;
use crate::runtime::monarch_with_gil_blocking;
use crate::shape::PyRegion;
use crate::supervision::Supervisable;
use crate::supervision::SupervisionError;

py_global!(_pickle, "monarch._src.actor.actor_mesh", "_pickle");

py_global!(
    shared_class,
    "monarch._rust_bindings.monarch_hyperactor.pytokio",
    "Shared"
);

/// Closed actor-mesh selection surface used by public send APIs.
///
/// Public Python APIs expose only `"all"` and `"choose"`. Keep that
/// invariant explicit here so arbitrary `ndslice::Selection` values stay out
/// of the public mesh-casting surface.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum AllOrChoose {
    All,
    Choose,
}

impl AllOrChoose {
    pub(crate) fn as_str(self) -> &'static str {
        match self {
            Self::All => "all",
            Self::Choose => "choose",
        }
    }
}

/// Trait defining the common interface for actor mesh, mesh ref and actor mesh implementations.
/// This corresponds to the Python ActorMeshProtocol ABC.
pub(crate) trait ActorMeshProtocol: Send + Sync {
    /// Cast a message to actors selected by the given selection using the specified mailbox.
    fn cast(
        &self,
        message: PythonMessage,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
    ) -> PyResult<()>;

    /// Cast a message, merging caller-supplied envelope headers into
    /// the outbound request. Implementations that reach the real
    /// envelope emission site override this to thread `caller_headers`
    /// through `hyperactor_mesh::ActorMeshRef::cast_with_headers`;
    /// the default collapses to the non-headers path for impls that
    /// have no envelope access.
    fn cast_with_headers(
        &self,
        message: PythonMessage,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
        _caller_headers: hyperactor_config::Flattrs,
    ) -> PyResult<()> {
        self.cast(message, selection, instance)
    }

    /// Cast a pending message (which may contain unresolved async values) to actors.
    ///
    /// The default implementation blocks on resolving the message and then calls cast.
    /// AsyncActorMesh overrides this with an optimized async implementation.
    fn cast_unresolved(
        &self,
        message: PendingMessage,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
    ) -> PyResult<()> {
        let message = get_tokio_runtime().block_on(message.resolve())?;
        self.cast(message, selection, instance)
    }

    /// Async counterpart of `cast_with_headers`. The default
    /// resolves the pending message synchronously and delegates;
    /// `AsyncActorMesh` overrides this to resolve asynchronously
    /// and route through `cast_with_headers`.
    fn cast_unresolved_with_headers(
        &self,
        message: PendingMessage,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
        caller_headers: hyperactor_config::Flattrs,
    ) -> PyResult<()> {
        let message = get_tokio_runtime().block_on(message.resolve())?;
        self.cast_with_headers(message, selection, instance, caller_headers)
    }

    fn __reduce__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)>;

    /// The serializable reference for this mesh, for the out-of-band `refs`
    /// table. Required so every impl chooses: hold a resolved ref and return
    /// it, or (a pending mesh) have none and error explicitly.
    fn mesh_ref(&self) -> PyResult<ActorMeshRef<PythonActor>>;

    /// Stop the actor mesh asynchronously.
    /// Default implementation raises NotImplementedError for types that don't support stopping.
    fn stop(&self, _instance: &PyInstance, _reason: String) -> PyResult<PyPythonTask> {
        Err(PyNotImplementedError::new_err(format!(
            "stop() is not supported for {}",
            std::any::type_name::<Self>()
        )))
    }

    /// Initialize the actor mesh asynchronously.
    /// Default implementation returns None (no initialization needed).
    fn initialized(&self) -> PyResult<PyPythonTask> {
        PyPythonTask::new(async { Ok(None::<()>) })
    }

    /// The name of the mesh.
    fn name(&self) -> PyResult<PyPythonTask>;
}

pub(crate) trait SupervisableActorMesh: ActorMeshProtocol + Supervisable {
    fn new_with_region(&self, region: &PyRegion) -> PyResult<Box<dyn SupervisableActorMesh>>;
}

/// This just forwards to the rust trait that can implement these bindings
#[pyclass(
    name = "PythonActorMesh",
    module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh"
)]
#[derive(Clone)]
pub(crate) struct PythonActorMesh {
    inner: Arc<dyn SupervisableActorMesh>,
}

impl PythonActorMesh {
    pub(crate) fn new<F>(f: F, supervised: bool) -> Self
    where
        F: Future<Output = PyResult<Box<dyn SupervisableActorMesh>>> + Send + 'static,
    {
        let f = async move { Ok(Arc::from(f.await?)) }.boxed().shared();
        PythonActorMesh {
            inner: Arc::new(AsyncActorMesh::new_queue(f, supervised)),
        }
    }

    pub(crate) fn from_impl(inner: Arc<dyn SupervisableActorMesh>) -> Self {
        PythonActorMesh { inner }
    }

    pub(crate) fn get_inner(&self) -> Arc<dyn SupervisableActorMesh> {
        self.inner.clone()
    }
}

pub(crate) fn to_all_or_choose(selection: &str) -> PyResult<AllOrChoose> {
    match selection {
        "choose" => Ok(AllOrChoose::Choose),
        "all" => Ok(AllOrChoose::All),
        _ => Err(PyErr::new::<PyValueError, _>(format!(
            "Invalid selection: {}",
            selection
        ))),
    }
}

#[pymethods]
impl PythonActorMesh {
    #[tracing::instrument(level = "debug", skip_all)]
    #[pyo3(name = "cast")]
    fn py_cast(
        &self,
        message: &PythonMessage,
        selection: &str,
        instance: &PyInstance,
    ) -> PyResult<()> {
        let sel = to_all_or_choose(selection)?;
        self.inner.cast(message.clone(), sel, instance.deref())
    }

    #[hyperactor::instrument]
    pub(crate) fn cast_unresolved(
        &self,
        message: &mut PendingMessage,
        selection: &str,
        instance: &PyInstance,
    ) -> PyResult<()> {
        let sel = to_all_or_choose(selection)?;
        let message = message.take()?;
        self.inner.cast_unresolved(message, sel, instance)
    }

    fn new_with_region(&self, region: &PyRegion) -> PyResult<PythonActorMesh> {
        let inner = self.inner.new_with_region(region)?;
        Ok(PythonActorMesh {
            inner: Arc::from(inner),
        })
    }

    fn stop(&self, instance: &PyInstance, reason: String) -> PyResult<PyPythonTask> {
        self.inner.stop(instance, reason)
    }

    fn initialized(&self) -> PyResult<PyPythonTask> {
        self.inner.initialized()
    }

    fn name(&self) -> PyResult<PyPythonTask> {
        self.inner.name()
    }

    fn __reduce__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)> {
        self.inner.__reduce__(py)
    }
}

#[derive(Debug)]
pub(crate) struct ClonePyErr {
    inner: PyErr,
}

impl From<ClonePyErr> for PyErr {
    fn from(value: ClonePyErr) -> PyErr {
        value.inner
    }
}
impl From<PyErr> for ClonePyErr {
    fn from(inner: PyErr) -> ClonePyErr {
        ClonePyErr { inner }
    }
}

impl Clone for ClonePyErr {
    fn clone(&self) -> Self {
        monarch_with_gil_blocking(GilSite::Convert, |py| self.inner.clone_ref(py).into())
    }
}

type ActorMeshResult = Result<Arc<dyn SupervisableActorMesh>, ClonePyErr>;
type ActorMeshFut = Shared<Pin<Box<dyn Future<Output = ActorMeshResult> + Send + 'static>>>;

pub(crate) struct AsyncActorMesh {
    mesh: ActorMeshFut,
    queue: UnboundedSender<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
    supervised: bool,
}

impl AsyncActorMesh {
    pub(crate) fn new_queue(f: ActorMeshFut, supervised: bool) -> AsyncActorMesh {
        let (queue, mut recv) = unbounded_channel();

        get_tokio_runtime().spawn(async move {
            loop {
                let r = recv.recv().await;
                if let Some(r) = r {
                    r.await;
                } else {
                    return;
                }
            }
        });

        let mesh = AsyncActorMesh::new(queue, supervised, f);
        // Eagerly trigger the mesh initialization by pushing an init task onto
        // the queue. This ensures actors are spawned immediately rather than
        // waiting for the first endpoint call, which is critical for:
        // 1. Tests/code that wait for supervision events from actor __init__
        //    failures without making any endpoint calls.
        // 2. Ensuring all meshes on a proc are spawned before any errors occur,
        //    preventing spawn rejections due to stale supervision events.
        let f = mesh.mesh.clone();
        mesh.push(async move {
            let _ = f.await;
        });
        mesh
    }

    fn new(
        queue: UnboundedSender<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
        supervised: bool,
        f: ActorMeshFut,
    ) -> AsyncActorMesh {
        AsyncActorMesh {
            mesh: f,
            queue,
            supervised,
        }
    }

    fn push<F>(&self, f: F)
    where
        F: Future<Output = ()> + Send + 'static,
    {
        self.queue.send(f.boxed()).unwrap();
    }

    pub(crate) fn from_impl(mesh: Arc<dyn SupervisableActorMesh>) -> Self {
        let fut = future::ready(Ok::<Arc<dyn SupervisableActorMesh>, ClonePyErr>(mesh))
            .boxed()
            .shared();
        // Poll the future so that its result can be observed without blocking the tokio runtime.
        let _ = futures::executor::block_on(fut.clone());
        Self::new_queue(fut, true)
    }
}

impl ActorMeshProtocol for AsyncActorMesh {
    fn cast(
        &self,
        _message: PythonMessage,
        _selection: AllOrChoose,
        _instance: &Instance<PythonActor>,
    ) -> PyResult<()> {
        panic!("not implemented")
    }

    fn cast_unresolved(
        &self,
        message: PendingMessage,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
    ) -> PyResult<()> {
        self.cast_unresolved_with_headers(
            message,
            selection,
            instance,
            hyperactor_config::Flattrs::new(),
        )
    }

    fn cast_unresolved_with_headers(
        &self,
        message: PendingMessage,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
        caller_headers: hyperactor_config::Flattrs,
    ) -> PyResult<()> {
        let mesh = self.mesh.clone();
        let instance = instance.clone_for_py();
        let port = match &message.kind {
            PythonMessageKind::CallMethod { response_port, .. } => response_port.clone(),
            _ => None,
        };
        self.push(
            async move {
                let result = async {
                    let resolved = message.resolve().await?;
                    mesh.await?
                        .cast_with_headers(resolved, selection, &instance, caller_headers)
                }
                .await;
                if let (Some(mut port_ref), Err(pyerr)) = (port, result) {
                    let _ = monarch_with_gil(GilSite::Traceback, |py: Python<'_>| {
                        let exception_str = crate::logging::format_traceback(py, &pyerr);
                        tracing::error!(
                            actor_id = instance.self_addr().to_string(),
                            "error occurred during cast unresolved: {}",
                            exception_str
                        );

                        // Endpoint calls create a response port: the
                        // PortRef is sent to the remote worker (to send
                        // results back), and collect_valuemesh owns the
                        // PortReceiver. If mesh.cast() fails here, we try
                        // to send the exception back to the caller via
                        // the PortRef ourselves. But a supervision event
                        // can cause collect_valuemesh to drop the
                        // PortReceiver (removing the port from the
                        // mailbox) before we get here. Disable
                        // return-undeliverable so a delivery failure
                        // doesn't bounce back and crash the root client.
                        //
                        // TODO: Tie the lifetime of this queued work to
                        // the PortReceiver (e.g. a cancellation token set
                        // on drop) so we can distinguish
                        // supervision-caused failures — where the caller
                        // already knows — from other cast errors where
                        // the caller actually needs this exception.

                        port_ref.set_return_undeliverable(false);

                        let mut state = crate::pickle::pickle(
                            py,
                            pyerr.into_value(py).into_any(),
                            false,
                            false,
                        )?;
                        let _ = port_ref.post(
                            &instance,
                            PythonMessage::new_from_buf(
                                PythonMessageKind::Exception { rank: Some(0) },
                                state.take_inner()?.take_buffer(),
                            ),
                        );

                        Ok::<_, PyErr>(())
                    })
                    .await;
                }
            }
            .instrument(tracing::debug_span!("AsyncActorMesh::cast")),
        );
        Ok(())
    }

    fn __reduce__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)> {
        let fut = self.mesh.clone();
        match fut.peek().cloned() {
            Some(mesh) => mesh?.__reduce__(py),
            None => {
                let shared =
                    PyPythonTask::new(async move { Ok(PythonActorMesh::from_impl(fut.await?)) })?
                        .spawn_abortable()?;
                let shared = Py::new(py, shared)?;
                if crate::pickle::reserve_mesh_reference_if_active(shared.clone_ref(py)) {
                    let pop_fn = py
                        .import("monarch._rust_bindings.monarch_hyperactor.pickle")?
                        .getattr("pop_mesh_reference")?;
                    return Ok((pop_fn, PyTuple::empty(py).into_any()));
                }
                // Get Shared.block_on as an unbound method
                let block_on = shared_class(py).getattr("block_on")?;
                let args = PyTuple::new(py, [shared])?;
                Ok((block_on, args.into_any()))
            }
        }
    }

    fn mesh_ref(&self) -> PyResult<ActorMeshRef<PythonActor>> {
        // A pending mesh has no serializable ref of its own; the reserve/fill
        // slot carries it out-of-band. This is a backstop: the happy path never
        // asks a pending mesh for a ref.
        Err(pyo3::exceptions::PyRuntimeError::new_err(
            "pending actor mesh has no serializable ref; it is carried via the reserve/fill slot",
        ))
    }

    fn stop(&self, instance: &PyInstance, reason: String) -> PyResult<PyPythonTask> {
        let mesh = self.mesh.clone();
        let instance = monarch_with_gil_blocking(GilSite::Stop, |_py| instance.clone());
        let (tx, rx) = tokio::sync::oneshot::channel();
        self.push(async move {
            let result =
                async move { mesh.await?.stop(&instance, reason)?.take_task()?.await }.await;
            if tx.send(result).is_err() {
                panic!("oneshot failed");
            }
        });
        PyPythonTask::new(async move { rx.await.map_err(anyhow::Error::from)? })
    }

    fn initialized<'py>(&self) -> PyResult<PyPythonTask> {
        let mesh = self.mesh.clone();
        PyPythonTask::new(async {
            mesh.await?;
            Ok(None::<()>)
        })
    }

    fn name(&self) -> PyResult<PyPythonTask> {
        let mesh = self.mesh.clone();
        let (tx, rx) = tokio::sync::oneshot::channel();
        self.push(async move {
            let result = async move { mesh.await?.name()?.take_task()?.await }.await;
            if tx.send(result).is_err() {
                panic!("oneshot failed");
            }
        });
        PyPythonTask::new(async move { rx.await.map_err(anyhow::Error::from)? })
    }
}

#[async_trait]
impl Supervisable for AsyncActorMesh {
    async fn supervision_event(&self, instance: &Instance<PythonActor>) -> Option<PyErr> {
        if !self.supervised {
            return None;
        }
        let mesh = self.mesh.clone();
        match mesh.await {
            Ok(mesh) => mesh.supervision_event(instance).await,
            Err(e) => Some(e.into()),
        }
    }
}

impl SupervisableActorMesh for AsyncActorMesh {
    fn new_with_region(&self, region: &PyRegion) -> PyResult<Box<dyn SupervisableActorMesh>> {
        let mesh = self.mesh.clone();
        let region = region.clone();
        Ok(Box::new(AsyncActorMesh::new(
            self.queue.clone(),
            self.supervised,
            async move { Ok(Arc::from(mesh.await?.new_with_region(&region)?)) }
                .boxed()
                .shared(),
        )))
    }
}

#[derive(Debug, Clone)]
#[pyclass(
    name = "PyActorMesh",
    module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh"
)]
pub(crate) struct PyActorMesh {
    mesh: ActorMesh<PythonActor>,
}

#[derive(Debug, Clone)]
#[pyclass(
    name = "PyActorMeshRef",
    module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh"
)]
pub(crate) struct PyActorMeshRef {
    mesh: ActorMeshRef<PythonActor>,
}

#[derive(Debug, Clone)]
#[pyclass(
    name = "PythonActorMeshImpl",
    module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh"
)]
#[expect(
    clippy::large_enum_variant,
    reason = "PyO3 #[pyclass] enum; Box wrapping interacts with PyO3 codegen and Python interop — separate diff"
)]
pub(crate) enum PythonActorMeshImpl {
    Owned(PyActorMesh),
    Ref(PyActorMeshRef),
}

impl PythonActorMeshImpl {
    /// Get a new owned [`PythonActorMeshImpl`].
    pub(crate) fn new_owned(inner: ActorMesh<PythonActor>) -> Self {
        PythonActorMeshImpl::Owned(PyActorMesh { mesh: inner })
    }

    /// Get a new ref-based [`PythonActorMeshImpl`].
    pub(crate) fn new_ref(inner: ActorMeshRef<PythonActor>) -> Self {
        PythonActorMeshImpl::Ref(PyActorMeshRef { mesh: inner })
    }

    fn mesh_ref(&self) -> &ActorMeshRef<PythonActor> {
        match self {
            PythonActorMeshImpl::Owned(inner) => &inner.mesh,
            PythonActorMeshImpl::Ref(inner) => &inner.mesh,
        }
    }
}

#[async_trait]
impl Supervisable for PythonActorMeshImpl {
    async fn supervision_event(&self, instance: &Instance<PythonActor>) -> Option<PyErr> {
        let mesh = self.mesh_ref();
        match mesh.next_supervision_event(instance).await {
            Ok(supervision_failure) => Some(SupervisionError::new_err_from(supervision_failure)),
            Err(e) => Some(PyValueError::new_err(e.to_string())),
        }
    }
}

impl ActorMeshProtocol for PythonActorMeshImpl {
    fn cast(
        &self,
        message: PythonMessage,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
    ) -> PyResult<()> {
        <ActorMeshRef<PythonActor> as ActorMeshProtocol>::cast(
            self.mesh_ref(),
            message,
            selection,
            instance,
        )
    }

    fn cast_with_headers(
        &self,
        message: PythonMessage,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
        caller_headers: hyperactor_config::Flattrs,
    ) -> PyResult<()> {
        <ActorMeshRef<PythonActor> as ActorMeshProtocol>::cast_with_headers(
            self.mesh_ref(),
            message,
            selection,
            instance,
            caller_headers,
        )
    }

    fn stop(&self, instance: &PyInstance, reason: String) -> PyResult<PyPythonTask> {
        let (slf, instance) =
            monarch_with_gil_blocking(GilSite::Stop, |_py| (self.clone(), instance.clone()));
        match slf {
            PythonActorMeshImpl::Owned(mut mesh) => PyPythonTask::new(async move {
                mesh.mesh
                    .stop(instance.deref(), reason)
                    .await
                    .map_err(|err| PyValueError::new_err(err.to_string()))
            }),
            PythonActorMeshImpl::Ref(_) => Err(PyNotImplementedError::new_err(
                "Cannot call stop on an ActorMeshRef, requires an owned ActorMesh",
            )),
        }
    }

    fn __reduce__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)> {
        self.mesh_ref().__reduce__(py)
    }

    fn mesh_ref(&self) -> PyResult<ActorMeshRef<PythonActor>> {
        // `self.mesh_ref()` resolves to the inherent borrow-returning method.
        Ok(PythonActorMeshImpl::mesh_ref(self).clone())
    }

    fn name(&self) -> PyResult<PyPythonTask> {
        let name = self.mesh_ref().id().to_string();
        PyPythonTask::new(async move { Ok(name) })
    }
}

impl SupervisableActorMesh for PythonActorMeshImpl {
    fn new_with_region(&self, region: &PyRegion) -> PyResult<Box<dyn SupervisableActorMesh>> {
        assert!(region.as_inner().is_subset(self.mesh_ref().region()));
        Ok(Box::new(PythonActorMeshImpl::new_ref(
            self.mesh_ref().sliced(region.as_inner().clone()),
        )))
    }
}

// Convert a hyperactor_mesh::Error to a Python exception. hyperactor_mesh::Error::Supervision becomes a SupervisionError,
// all others become a RuntimeError.
fn cast_error_to_py_error(err: hyperactor_mesh::Error) -> PyErr {
    if let hyperactor_mesh::Error::Supervision(failure) = err {
        SupervisionError::new_err_from(*failure)
    } else {
        PyRuntimeError::new_err(err.to_string())
    }
}

impl ActorMeshProtocol for ActorMeshRef<PythonActor> {
    fn cast(
        &self,
        message: PythonMessage,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
    ) -> PyResult<()> {
        <Self as ActorMeshProtocol>::cast_with_headers(
            self,
            message,
            selection,
            instance,
            hyperactor_config::Flattrs::new(),
        )
    }

    fn cast_with_headers(
        &self,
        message: PythonMessage,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
        caller_headers: hyperactor_config::Flattrs,
    ) -> PyResult<()> {
        match selection {
            AllOrChoose::All => ActorMeshRef::<PythonActor>::cast_with_headers(
                self,
                instance,
                &caller_headers,
                message,
            )
            .map_err(cast_error_to_py_error),
            AllOrChoose::Choose => ActorMeshRef::<PythonActor>::cast_choose_with_headers(
                self,
                instance,
                &caller_headers,
                message,
            )
            .map_err(cast_error_to_py_error),
        }
    }

    /// Stop the actor mesh asynchronously.
    fn stop(&self, _instance: &PyInstance, _reason: String) -> PyResult<PyPythonTask> {
        Err(PyNotImplementedError::new_err(
            "This cannot be used on ActorMeshRef, only on owned ActorMesh",
        ))
    }

    fn __reduce__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)> {
        if crate::pickle::push_mesh_reference_if_active(crate::actor::MeshRef::Actor(Box::new(
            self.clone(),
        ))) {
            let pop_fn = py
                .import("monarch._rust_bindings.monarch_hyperactor.pickle")?
                .getattr("pop_mesh_reference")?;
            return Ok((pop_fn, pyo3::types::PyTuple::empty(py).into_any()));
        }
        let bytes = bincode::serde::encode_to_vec(self, bincode::config::legacy())
            .map_err(|e| PyValueError::new_err(e.to_string()))?;
        let py_bytes = (PyBytes::new(py, &bytes),).into_bound_py_any(py).unwrap();
        let module = py
            .import("monarch._rust_bindings.monarch_hyperactor.actor_mesh")
            .unwrap();
        let from_bytes = module.getattr("py_actor_mesh_from_bytes").unwrap();
        Ok((from_bytes, py_bytes))
    }

    fn mesh_ref(&self) -> PyResult<ActorMeshRef<PythonActor>> {
        Ok(self.clone())
    }

    fn name(&self) -> PyResult<PyPythonTask> {
        let name = self.id().to_string();
        PyPythonTask::new(async move { Ok(name) })
    }
}

#[pymethods]
impl PythonActorMeshImpl {
    fn get(&self, rank: usize) -> PyResult<Option<PyActorAddr>> {
        Ok(self
            .mesh_ref()
            .get(rank)
            .map(|r| hyperactor::ActorRef::into_actor_addr(r.clone()))
            .map(PyActorAddr::from))
    }

    fn __repr__(&self) -> String {
        format!("PythonActorMeshImpl({:?})", self.mesh_ref())
    }
}

#[pyfunction]
fn py_actor_mesh_from_bytes(bytes: &Bound<'_, PyBytes>) -> PyResult<PythonActorMesh> {
    let r: PyResult<ActorMeshRef<PythonActor>> =
        bincode::serde::decode_from_slice(bytes.as_bytes(), bincode::config::legacy())
            .map(|(v, _)| v)
            .map_err(|e| PyValueError::new_err(e.to_string()));
    r.map(|r| AsyncActorMesh::from_impl(Arc::new(PythonActorMeshImpl::new_ref(r))))
        .map(|r| PythonActorMesh::from_impl(Arc::from(r)))
}

#[pyclass(
    name = "ActorSupervisionEvent",
    module = "monarch._rust_bindings.monarch_hyperactor.actor_mesh"
)]
#[derive(Debug)]
pub struct PyActorSupervisionEvent {
    inner: ActorSupervisionEvent,
}

#[pymethods]
impl PyActorSupervisionEvent {
    pub(crate) fn __repr__(&self) -> PyResult<String> {
        Ok(format!("<PyActorSupervisionEvent: {}>", self.inner))
    }

    #[getter]
    pub(crate) fn actor_id(&self) -> PyResult<PyActorAddr> {
        Ok(PyActorAddr::from(self.inner.actor_id.clone()))
    }

    #[getter]
    pub(crate) fn actor_status(&self) -> PyResult<String> {
        Ok(self.inner.actor_status.to_string())
    }
}

impl From<ActorSupervisionEvent> for PyActorSupervisionEvent {
    fn from(event: ActorSupervisionEvent) -> Self {
        PyActorSupervisionEvent { inner: event }
    }
}

#[pyfunction]
fn py_identity(obj: Py<PyAny>) -> PyResult<Py<PyAny>> {
    Ok(obj)
}

/// Holds the GIL for the specified number of seconds without releasing it.
///
/// This is a test utility function that spawns a background thread which
/// acquires the GIL using Rust's Python::attach and holds it for the
/// specified duration using thread::sleep. Unlike Python code which
/// periodically releases the GIL, this function holds it continuously.
///
/// We intentionally use `std::thread::sleep` here (not `Clock::sleep` or async sleep)
/// because the purpose is to simulate a blocking operation that holds the GIL without
/// releasing it. Using an async sleep would release the GIL periodically, defeating
/// the purpose of this test utility.
///
/// Args:
///     delay_secs: Seconds to wait before acquiring the GIL
///     hold_secs: Seconds to hold the GIL
#[pyfunction]
#[pyo3(name = "hold_gil_for_test", signature = (delay_secs, hold_secs))]
pub fn hold_gil_for_test(delay_secs: f64, hold_secs: f64) {
    thread::spawn(move || {
        // Wait before grabbing the GIL (blocking sleep is fine here, we're in a spawned thread)
        thread::sleep(Duration::from_secs_f64(delay_secs));
        // Acquire and hold the GIL - MUST use blocking sleep to keep GIL held
        monarch_with_gil_blocking(GilSite::Test, |_py| {
            tracing::info!("start holding the gil...");
            thread::sleep(Duration::from_secs_f64(hold_secs));
            tracing::info!("end holding the gil...");
        });
    });
}

pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
    py_module_add_function!(
        hyperactor_mod,
        "monarch._rust_bindings.monarch_hyperactor.actor_mesh",
        py_identity
    );
    py_module_add_function!(
        hyperactor_mod,
        "monarch._rust_bindings.monarch_hyperactor.actor_mesh",
        py_actor_mesh_from_bytes
    );
    py_module_add_function!(
        hyperactor_mod,
        "monarch._rust_bindings.monarch_hyperactor.actor_mesh",
        hold_gil_for_test
    );
    hyperactor_mod.add_class::<PythonActorMesh>()?;
    hyperactor_mod.add_class::<PythonActorMeshImpl>()?;
    hyperactor_mod.add_class::<PyActorSupervisionEvent>()?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use std::sync::OnceLock;
    use std::time::Duration;

    use async_trait::async_trait;
    use hyperactor::Actor;
    use hyperactor::Context;
    use hyperactor::Endpoint as _;
    use hyperactor::Handler;
    use hyperactor::Instance;
    use hyperactor::Proc;
    use hyperactor::actor::Signal;
    use hyperactor::channel::ChannelTransport;
    use hyperactor::mailbox;
    use hyperactor::proc::ActorWorkReceiver;
    use hyperactor::supervision::ActorSupervisionEvent;
    use hyperactor_mesh::host_mesh::HostMesh;
    use hyperactor_mesh::mesh_controller::GetSubscriberCount;
    use hyperactor_mesh::supervision::MeshFailure;
    use monarch_types::PickledPyObject;
    use ndslice::extent;
    use tokio::sync::mpsc;

    use super::*;
    use crate::actor::PythonActor;
    use crate::actor::PythonActorParams;
    use crate::config::ACTOR_QUEUE_DISPATCH;

    /// Minimal root-client actor for test infrastructure.
    /// Handles MeshFailure by panicking (test failure).
    #[derive(Debug)]
    struct TestClient {
        signal_rx: mpsc::UnboundedReceiver<Signal>,
        supervision_rx: mpsc::UnboundedReceiver<ActorSupervisionEvent>,
        work_rx: ActorWorkReceiver<Self>,
    }

    impl Actor for TestClient {}

    #[async_trait]
    impl Handler<MeshFailure> for TestClient {
        async fn handle(
            &mut self,
            _cx: &Context<Self>,
            msg: MeshFailure,
        ) -> Result<(), anyhow::Error> {
            panic!("unexpected supervision failure in test: {}", msg);
        }
    }

    impl TestClient {
        fn run(mut self, instance: &'static Instance<Self>) {
            tokio::spawn(async move {
                loop {
                    tokio::select! {
                        work = self.work_rx.recv() => {
                            match work {
                                Some(work) => {
                                    let _ = work.handle(&mut self, instance).await;
                                }
                                None => break,
                            }
                        }
                        Some(_) = self.signal_rx.recv() => {}
                        Some(event) = self.supervision_rx.recv() => {
                            let _ = instance
                                .handle_supervision_event(&mut self, event)
                                .await;
                        }
                    }
                }
            });
        }
    }

    fn init_test_instance() -> &'static Instance<TestClient> {
        static INSTANCE: OnceLock<Instance<TestClient>> = OnceLock::new();
        let proc = Proc::direct(ChannelTransport::Unix.any(), "test_proc".to_string()).unwrap();
        let ai = proc.actor_instance("test_client").unwrap();

        INSTANCE
            .set(ai.instance)
            .map_err(|_| "already initialized")
            .unwrap();
        let instance = INSTANCE.get().unwrap();

        TestClient {
            signal_rx: ai.signal,
            supervision_rx: ai.supervision,
            work_rx: ai.work,
        }
        .run(instance);

        instance
    }

    fn test_instance() -> &'static Instance<TestClient> {
        static INSTANCE: OnceLock<&'static Instance<TestClient>> = OnceLock::new();
        INSTANCE.get_or_init(init_test_instance)
    }

    /// Verify that calling `supervision_event` repeatedly through a
    /// [`PythonActorMesh`] does not increase the subscriber count on the
    /// controller.  This guards against a regression where each call
    /// would create a new supervision subscriber.
    #[tokio::test]
    async fn test_subscriber_count_stable_across_supervision_calls() {
        crate::pytokio::ensure_python();

        let instance = test_instance();

        let mut host_mesh = HostMesh::local_in_process().await.unwrap();
        let proc_mesh = host_mesh
            .spawn(instance, "test", extent!(replicas = 2), None, None)
            .await
            .unwrap();

        // Create a minimal Python class and pickle it so we can spawn
        // PythonActor instances (mirroring PyProcMesh::spawn_async).
        // The class must live in __main__'s globals for pickle to find it.
        let pickled_type = monarch_with_gil_blocking(GilSite::Test, |py| {
            py.run(c"class MinimalActor: pass", None, None).unwrap();

            PickledPyObject::pickle(
                &py.import("__main__")
                    .unwrap()
                    .getattr("MinimalActor")
                    .unwrap(),
            )
            .unwrap()
        });

        let actor_params = PythonActorParams::new(pickled_type, None, None);
        let config_lock = hyperactor_config::global::lock();
        let _queue_dispatch_guard = config_lock.override_key(ACTOR_QUEUE_DISPATCH, false);
        let actor_mesh = proc_mesh
            .spawn::<PythonActor, _>(instance, "test_actors", &actor_params)
            .await
            .unwrap();
        drop(_queue_dispatch_guard);
        drop(config_lock);

        let controller = actor_mesh.controller().as_ref().unwrap().clone();

        // Wrap using the production code path from PyProcMesh::spawn_async.
        let mesh_impl =
            async move { Ok::<_, PyErr>(Box::new(PythonActorMeshImpl::new_owned(actor_mesh))) };
        let python_actor_mesh = PythonActorMesh::new(
            async move {
                let mesh_impl: Box<dyn SupervisableActorMesh> = mesh_impl.await?;
                Ok(mesh_impl)
            },
            true,
        );

        // Instance<PythonActor> required by the Supervisable trait
        // signature. Only used for subscription routing inside
        // next_supervision_event.
        let py_ai = Proc::direct(ChannelTransport::Unix.any(), "py_proc".to_string())
            .unwrap()
            .actor_instance::<PythonActor>("py_client")
            .unwrap();
        let py_instance = py_ai.instance;

        // Query the subscriber count from the controller.
        let (port, mut rx) = mailbox::open_port::<usize>(instance);
        controller.post(instance, GetSubscriberCount(port.bind()));
        let initial_count = tokio::time::timeout(Duration::from_secs(5), rx.recv())
            .await
            .expect("timed out waiting for subscriber count")
            .expect("channel closed");
        assert_eq!(initial_count, 0, "should have 0 subscribers initially");

        // Call supervision_event through the PythonActorMesh multiple
        // times, racing against a short timeout each time.  The mesh is
        // healthy so no event fires; we just want to trigger the lazy
        // subscriber initialization repeatedly.
        for _ in 0..5 {
            tokio::select! {
                _ = python_actor_mesh.inner.supervision_event(&py_instance) => {
                    panic!("unexpected supervision event on healthy mesh");
                }
                _ = tokio::time::sleep(Duration::from_millis(200)) => {}
            }
        }

        // After 5 calls from the same context, there should be exactly 1
        // subscriber (created lazily on the first call, reused thereafter).
        let (port, mut rx) = mailbox::open_port::<usize>(instance);
        controller.post(instance, GetSubscriberCount(port.bind()));
        let after_count = tokio::time::timeout(Duration::from_secs(5), rx.recv())
            .await
            .expect("timed out waiting for subscriber count")
            .expect("channel closed");
        assert_eq!(
            after_count, 1,
            "subscriber count should be exactly 1, not growing with each call"
        );

        // Do 5 more calls to confirm it stays stable.
        for _ in 0..5 {
            tokio::select! {
                _ = python_actor_mesh.inner.supervision_event(&py_instance) => {
                    panic!("unexpected supervision event on healthy mesh");
                }
                _ = tokio::time::sleep(Duration::from_millis(200)) => {}
            }
        }

        let (port, mut rx) = mailbox::open_port::<usize>(instance);
        controller.post(instance, GetSubscriberCount(port.bind()));
        let final_count = tokio::time::timeout(Duration::from_secs(5), rx.recv())
            .await
            .expect("timed out waiting for subscriber count")
            .expect("channel closed");
        assert_eq!(
            final_count, 1,
            "subscriber count should still be 1 after repeated calls"
        );

        let _ = host_mesh.shutdown(instance).await;
    }
}