monarch_hyperactor/code_sync/
auto_reload.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::sync::Arc;
10
11use anyhow::Result;
12use async_trait::async_trait;
13use hyperactor as reference;
14use hyperactor::Actor;
15use hyperactor::Context;
16use hyperactor::Endpoint as _;
17use hyperactor::Handler;
18use hyperactor::RemoteSpawn;
19use hyperactor_config::Flattrs;
20use monarch_types::SerializablePyErr;
21use pyo3::prelude::*;
22use serde::Deserialize;
23use serde::Serialize;
24use typeuri::Named;
25
26use crate::runtime::GilSite;
27use crate::runtime::monarch_with_gil_blocking;
28
29/// Message to trigger module reloading
30#[derive(Debug, Clone, Named, Serialize, Deserialize)]
31pub struct AutoReloadMessage {
32    pub result: reference::PortRef<Result<(), String>>,
33}
34wirevalue::register_type!(AutoReloadMessage);
35
36/// Parameters for creating an AutoReloadActor
37#[derive(Debug, Clone, Named, Serialize, Deserialize)]
38pub struct AutoReloadParams {}
39wirevalue::register_type!(AutoReloadParams);
40
41/// Simple Rust Actor that wraps the Python AutoReloader class via pyo3
42#[derive(Debug)]
43#[hyperactor::export(handlers = [AutoReloadMessage])]
44#[hyperactor::spawnable]
45pub struct AutoReloadActor {
46    state: Result<(Arc<Py<PyAny>>, Py<PyAny>), SerializablePyErr>,
47}
48
49impl Actor for AutoReloadActor {}
50
51#[async_trait]
52impl RemoteSpawn for AutoReloadActor {
53    type Params = AutoReloadParams;
54
55    async fn new(Self::Params {}: Self::Params, _environment: Flattrs) -> Result<Self> {
56        AutoReloadActor::new().await
57    }
58}
59
60impl AutoReloadActor {
61    pub(crate) async fn new() -> Result<Self, anyhow::Error> {
62        Ok(Self {
63            state: tokio::task::spawn_blocking(move || {
64                monarch_with_gil_blocking(GilSite::CodeSync, |py| {
65                    Self::create_state(py).map_err(SerializablePyErr::from_fn(py))
66                })
67            })
68            .await?,
69        })
70    }
71
72    fn create_state(py: Python) -> PyResult<(Arc<Py<PyAny>>, Py<PyAny>)> {
73        // Import the Python AutoReloader class
74        let auto_reload_module = py.import("monarch._src.actor.code_sync.auto_reload")?;
75        let auto_reloader_class = auto_reload_module.getattr("AutoReloader")?;
76
77        let reloader = auto_reloader_class.call0()?;
78
79        // Install the audit import hook: SysAuditImportHook.install(reloader.import_callback)
80        let sys_audit_import_hook_class = auto_reload_module.getattr("SysAuditImportHook")?;
81        let import_callback = reloader.getattr("import_callback")?;
82        let hook_guard = sys_audit_import_hook_class.call_method1("install", (import_callback,))?;
83
84        Ok((Arc::new(reloader.into()), hook_guard.into()))
85    }
86
87    fn reload(py: Python, py_reloader: &Py<PyAny>) -> PyResult<()> {
88        let reloader = py_reloader.bind(py);
89        let changed_modules: Vec<String> = reloader.call_method0("reload_changes")?.extract()?;
90        if !changed_modules.is_empty() {
91            eprintln!("reloaded modules: {:?}", changed_modules);
92        }
93        Ok(())
94    }
95}
96
97#[async_trait]
98impl Handler<AutoReloadMessage> for AutoReloadActor {
99    async fn handle(
100        &mut self,
101        cx: &Context<Self>,
102        AutoReloadMessage { result }: AutoReloadMessage,
103    ) -> Result<()> {
104        // Call the Python reloader's reload_changes method
105        let res = async {
106            let py_reloader: Arc<_> = self.state.as_ref().map_err(Clone::clone)?.0.clone();
107            tokio::task::spawn_blocking(move || {
108                monarch_with_gil_blocking(GilSite::CodeSync, |py| {
109                    Self::reload(py, py_reloader.as_ref()).map_err(SerializablePyErr::from_fn(py))
110                })
111            })
112            .await??;
113            anyhow::Ok(())
114        }
115        .await;
116        result.post(cx, res.map_err(|e| format!("{:#?}", e)));
117        Ok(())
118    }
119}
monarch_hyperactor/code_sync/auto_reload.rs

monarch_hyperactor/code_sync/
auto_reload.rs
