monarch_types/
pytree.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 pyo3::Borrowed;
10use pyo3::Bound;
11use pyo3::FromPyObject;
12use pyo3::IntoPyObject;
13use pyo3::IntoPyObjectExt;
14use pyo3::PyAny;
15use pyo3::PyErr;
16use pyo3::Python;
17use pyo3::exceptions::PyRuntimeError;
18use pyo3::exceptions::PyTypeError;
19use pyo3::prelude::PyAnyMethods;
20use pyo3::prelude::PyResult;
21use pyo3::prelude::PyTupleMethods;
22use pyo3::types::PyBool;
23use pyo3::types::PyBoolMethods;
24use pyo3::types::PyList;
25use pyo3::types::PyModule;
26use pyo3::types::PyTuple;
27use serde::Deserialize;
28use serde::Serialize;
29use typeuri::Named;
30
31use crate::PickledPyObject;
32use crate::python::TryIntoPyObjectUnsafe;
33
34#[derive(Clone, Debug, Serialize, Deserialize)]
35pub enum TreeSpec {
36    Leaf,
37    Tree(PickledPyObject),
38}
39
40/// A Rust wrapper around A PyTorch pytree, which can be serialized and sent
41/// across the wire.
42/// https://github.com/pytorch/pytorch/blob/main/torch/utils/_pytree.py
43// NOTE: We have runtime deps on torch's pytree module and `pickle`, which the
44// user must ensure is available.
45#[derive(Clone, Debug, Serialize, Deserialize, Named)]
46pub struct PyTree<T> {
47    /// A wrapper around the tree spec.
48    // NOTE: This is currently just the pickled bytes of the tree spec.  We
49    // could also just deserialize to a `PyObject`, but this would mean
50    // acquiring the GIL when deserializing the message.
51    treespec: TreeSpec,
52    /// The deserialized leaves of the pytree.
53    leaves: Vec<T>,
54}
55
56impl<T> PyTree<T> {
57    pub fn is_leaf(&self) -> bool {
58        matches!(self.treespec, TreeSpec::Leaf)
59    }
60
61    pub fn into_leaf(mut self) -> Option<T> {
62        if self.is_leaf() {
63            self.leaves.pop()
64        } else {
65            None
66        }
67    }
68
69    pub fn leaves(&self) -> &[T] {
70        &self.leaves
71    }
72
73    pub fn into_leaves(self) -> Vec<T> {
74        self.leaves
75    }
76
77    pub fn for_each<F>(&self, func: F)
78    where
79        F: FnMut(&T),
80    {
81        self.leaves.iter().for_each(func)
82    }
83
84    pub fn iter(&self) -> std::slice::Iter<'_, T> {
85        self.leaves.iter()
86    }
87
88    pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, T> {
89        self.leaves.iter_mut()
90    }
91
92    /// Map leaf values with the given func.
93    pub fn into_map<U, F>(self, func: F) -> PyTree<U>
94    where
95        F: FnMut(T) -> U,
96    {
97        PyTree {
98            treespec: self.treespec,
99            leaves: self.leaves.into_iter().map(func).collect(),
100        }
101    }
102
103    /// Map the leaves of the pytree with the given fallible callback.
104    // NOTE: This ends up copying the serialized treespec.
105    pub fn try_map<U, F, E>(&self, mut func: F) -> Result<PyTree<U>, E>
106    where
107        F: FnMut(&T) -> Result<U, E>,
108    {
109        let mut leaves = vec![];
110        for leaf in self.leaves.iter() {
111            leaves.push(func(leaf)?);
112        }
113        Ok(PyTree {
114            treespec: self.treespec.clone(),
115            leaves,
116        })
117    }
118
119    /// Map the leaves of the pytree with the given fallible callback.
120    pub fn try_into_map<U, F, E>(self, mut func: F) -> Result<PyTree<U>, E>
121    where
122        F: FnMut(T) -> Result<U, E>,
123    {
124        let mut leaves = vec![];
125        for leaf in self.leaves.into_iter() {
126            leaves.push(func(leaf)?);
127        }
128        Ok(PyTree {
129            treespec: self.treespec,
130            leaves,
131        })
132    }
133
134    fn unflatten_impl<'a>(
135        py: Python<'a>,
136        treespec: &TreeSpec,
137        mut leaves: impl Iterator<Item = PyResult<Bound<'a, PyAny>>>,
138    ) -> PyResult<Bound<'a, PyAny>> {
139        if let TreeSpec::Tree(tree) = treespec {
140            // Call into pytorch's unflatten.
141            let module = PyModule::import(py, "torch.utils._pytree")?;
142            let function = module.getattr("tree_unflatten")?;
143            let leaves = leaves.collect::<Result<Vec<_>, _>>()?;
144            let leaves = PyList::new(py, &leaves)?;
145            let args = PyTuple::new(py, vec![leaves.as_any(), &tree.unpickle(py)?])?;
146            let result = function.call(args, None)?;
147            Ok(result)
148        } else {
149            leaves.next().ok_or(PyRuntimeError::new_err(
150                "Pytree leaf unexpectedly had no value",
151            ))?
152        }
153    }
154}
155
156impl<T> From<T> for PyTree<T> {
157    fn from(value: T) -> Self {
158        PyTree {
159            treespec: TreeSpec::Leaf,
160            leaves: vec![value],
161        }
162    }
163}
164
165impl<'py, T> IntoPyObject<'py> for PyTree<T>
166where
167    T: IntoPyObject<'py>,
168{
169    type Target = PyAny;
170    type Output = Bound<'py, Self::Target>;
171    type Error = PyErr;
172
173    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
174        PyTree::<T>::unflatten_impl(
175            py,
176            &self.treespec,
177            self.leaves.into_iter().map(|l| l.into_bound_py_any(py)),
178        )
179    }
180}
181
182impl<'a, 'py, T> IntoPyObject<'py> for &'a PyTree<T>
183where
184    &'a T: IntoPyObject<'py>,
185    T: 'a,
186{
187    type Target = PyAny;
188    type Output = Bound<'py, Self::Target>;
189    type Error = PyErr;
190
191    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
192        PyTree::<T>::unflatten_impl(
193            py,
194            &self.treespec,
195            self.leaves.iter().map(|l| l.into_bound_py_any(py)),
196        )
197    }
198}
199
200/// Serialize into a `PyObject`.
201impl<'a, 'py, T> TryIntoPyObjectUnsafe<'py, PyAny> for &'a PyTree<T>
202where
203    &'a T: TryIntoPyObjectUnsafe<'py, PyAny>,
204    T: 'a,
205{
206    unsafe fn try_to_object_unsafe(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
207        PyTree::<T>::unflatten_impl(
208            py,
209            &self.treespec,
210            self.leaves
211                .iter()
212                // SAFETY: Safety requirements are propagated via the `unsafe`
213                // tag on this method.
214                .map(|l| unsafe { l.try_to_object_unsafe(py) }),
215        )
216    }
217}
218
219impl<'py, T: for<'a> FromPyObject<'a, 'py>> PyTree<T> {
220    pub fn flatten(tree: &Bound<'py, PyAny>) -> PyResult<Self> {
221        let py = tree.py();
222
223        // Call into pytorch's flatten.
224        let pytree_module = PyModule::import(py, "torch.utils._pytree")?;
225        let tree_flatten = pytree_module.getattr("tree_flatten")?;
226        let res = tree_flatten.call1((tree,))?;
227
228        // Convert leaves to Rust objects.
229        let (leaves, treespec) = match res.cast::<PyTuple>()?.as_slice() {
230            [leaves, treespec] => {
231                let mut out = vec![];
232                for leaf in leaves.try_iter()? {
233                    out.push(leaf?.extract::<T>().map_err(Into::into)?);
234                }
235                (out, treespec)
236            }
237            _ => return Err(PyTypeError::new_err("unexpected result from tree_flatten")),
238        };
239
240        if treespec
241            .call_method0("is_leaf")?
242            .cast::<PyBool>()?
243            .is_true()
244        {
245            Ok(Self {
246                treespec: TreeSpec::Leaf,
247                leaves,
248            })
249        } else {
250            Ok(Self {
251                treespec: TreeSpec::Tree(PickledPyObject::pickle(treespec)?),
252                leaves,
253            })
254        }
255    }
256}
257
258/// Deserialize from a `PyObject`.
259impl<'a, 'py, T: for<'b> FromPyObject<'b, 'py>> FromPyObject<'a, 'py> for PyTree<T> {
260    type Error = PyErr;
261
262    fn extract(obj: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
263        Self::flatten(&obj)
264    }
265}
266
267#[cfg(all(test, fbcode_build))]
268mod tests {
269    use anyhow::Result;
270    use monarch_gil::GilSite;
271    use monarch_gil::monarch_with_gil_blocking;
272    use pyo3::IntoPyObject;
273    use pyo3::Python;
274    use pyo3::ffi::c_str;
275    use pyo3::py_run;
276
277    use super::PyTree;
278
279    #[cfg_attr(not(feature = "tensor_engine"), ignore = "requires tensor_engine")]
280    #[test]
281    fn flatten_unflatten() -> Result<()> {
282        Python::initialize();
283        monarch_with_gil_blocking(GilSite::Test, |py| {
284            let tree = py.eval(c_str!("[1, 2]"), None, None)?;
285            let tree: PyTree<u64> = PyTree::flatten(&tree)?;
286            assert_eq!(tree.leaves, vec![1u64, 2u64]);
287            let list = tree.into_pyobject(py)?;
288            py_run!(py, list, "assert list == [1, 2]");
289            anyhow::Ok(())
290        })?;
291        Ok(())
292    }
293
294    #[cfg_attr(not(feature = "tensor_engine"), ignore = "requires tensor_engine")]
295    #[test]
296    fn try_map() -> Result<()> {
297        Python::initialize();
298        monarch_with_gil_blocking(GilSite::Test, |py| {
299            let tree = py.eval(c_str!("[1, 2]"), None, None)?;
300            let tree: PyTree<u64> = PyTree::flatten(&tree)?;
301            let tree = tree.try_map(|v| anyhow::Ok(v + 1))?;
302            assert_eq!(tree.leaves, vec![2u64, 3u64]);
303            anyhow::Ok(())
304        })?;
305        Ok(())
306    }
307}
monarch_types/pytree.rs

monarch_types/
pytree.rs
