monarch_types/
lib.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
9#![feature(assert_matches)]
10
11pub mod nccl_types;
12mod pyobject;
13mod python;
14mod pytree;
15
16pub use nccl_types::NcclUniqueId;
17pub use nccl_types::ReduceOp;
18pub use nccl_types::UniqueId;
19use pyo3::PyErr;
20use pyo3::exceptions::PyValueError;
21pub use pyobject::PickledPyObject;
22pub use python::SerializablePyErr;
23pub use python::TryIntoPyObjectUnsafe;
24pub use pytree::PyTree;
25
26/// Macro to generate a Python object lookup function with caching
27///
28/// # Arguments
29/// * `$fn_name` - Name of the Rust function to generate
30/// * `$python_path` - Path to the Python object as a string (e.g., "module.submodule.function")
31#[macro_export]
32macro_rules! py_global {
33    ($fn_name:ident, $python_module:literal, $python_class:literal) => {
34        fn $fn_name<'py>(py: ::pyo3::Python<'py>) -> ::pyo3::Bound<'py, ::pyo3::PyAny> {
35            static CACHE: ::pyo3::sync::PyOnceLock<::pyo3::Py<::pyo3::PyAny>> =
36                ::pyo3::sync::PyOnceLock::new();
37            CACHE
38                .import(py, $python_module, $python_class)
39                .unwrap()
40                .clone()
41        }
42    };
43}
44
45/// Macro to register a function to a Python module.
46#[macro_export]
47macro_rules! py_module_add_function {
48    ($mod:ident, $mod_name:literal, $fn:ident) => {
49        let f = pyo3::wrap_pyfunction!($fn, $mod)?;
50        f.setattr("__module__", $mod_name)?;
51        $mod.add_function(f)?;
52    };
53}
54
55pub trait MapPyErr<T> {
56    fn map_pyerr(self) -> Result<T, PyErr>;
57}
58impl<T, E> MapPyErr<T> for Result<T, E>
59where
60    E: ToString,
61{
62    fn map_pyerr(self) -> Result<T, PyErr> {
63        self.map_err(|err| PyErr::new::<PyValueError, _>(err.to_string()))
64    }
65}
monarch_types/lib.rs

monarch_types/
lib.rs
