monarch_messages/
wire_value.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 derive_more::From;
10use derive_more::TryInto;
11use monarch_gil::GilSite;
12use monarch_gil::monarch_with_gil_blocking;
13use monarch_types::PickledPyObject;
14use pyo3::IntoPyObjectExt;
15use pyo3::prelude::*;
16use pyo3::types::PyNone;
17use serde::Deserialize;
18use serde::Serialize;
19use torch_sys2::Device;
20use torch_sys2::Layout;
21use torch_sys2::MemoryFormat;
22use torch_sys2::ScalarType;
23use typeuri::Named;
24
25use crate::worker::Ref;
26
27/// A value used as an input to CallFunction.
28// TODO, this is basically the same as RValue, but with TensorIndices swapped
29// out for refs. And IValue is the same as RValue, but with real tensors and
30// C++ types. I wonder if there is a nicer way to express this relationship.
31// TODO extend this to support other types of values, like bytes, dicts etc.
32#[derive(Serialize, Deserialize, Debug, Clone, TryInto, Named, From)]
33pub enum WireValue {
34    // Make sure boolean goes ealier than int as bool is a subclass of int.
35    // Otherwise, bool will be converted to int.
36    Bool(bool),
37    Int(i64),
38    Double(f64),
39    String(String),
40    Ref(Ref),
41    IntList(Vec<i64>),
42    RefList(Vec<Ref>),
43    Device(Device),
44    Layout(#[serde(with = "torch_sys2::LayoutDef")] Layout),
45    ScalarType(#[serde(with = "torch_sys2::ScalarTypeDef")] ScalarType),
46    MemoryFormat(#[serde(with = "torch_sys2::MemoryFormatDef")] MemoryFormat),
47    // Make this wrap the unit type, as `pyo3::FromPyObject` doesn't work with
48    // empty enum variants.
49    None(()),
50    PyObject(PickledPyObject),
51}
52wirevalue::register_type!(WireValue);
53
54impl FromPyObject<'_, '_> for WireValue {
55    type Error = PyErr;
56
57    fn extract(obj: pyo3::Borrowed<'_, '_, PyAny>) -> PyResult<Self> {
58        Ok(WireValue::PyObject(PickledPyObject::pickle(&obj)?))
59    }
60}
61
62impl<'py> IntoPyObject<'py> for WireValue {
63    type Target = PyAny;
64    type Output = Bound<'py, PyAny>;
65    type Error = PyErr;
66
67    fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
68        match self {
69            WireValue::Ref(ref_) => ref_.into_bound_py_any(py),
70            WireValue::RefList(ref_list) => ref_list.clone().into_bound_py_any(py),
71            WireValue::Int(int) => int.into_bound_py_any(py),
72            WireValue::IntList(int_list) => int_list.clone().into_bound_py_any(py),
73            WireValue::Double(double) => double.into_bound_py_any(py),
74            WireValue::Bool(bool_) => bool_.into_bound_py_any(py),
75            WireValue::String(string) => string.into_bound_py_any(py),
76            WireValue::Device(device) => device.into_bound_py_any(py),
77            WireValue::Layout(val) => val.into_bound_py_any(py),
78            WireValue::ScalarType(val) => val.into_bound_py_any(py),
79            WireValue::MemoryFormat(val) => val.into_bound_py_any(py),
80            WireValue::None(()) => PyNone::get(py).into_bound_py_any(py),
81            WireValue::PyObject(val) => val.unpickle(py),
82        }
83    }
84}
85
86impl From<Py<PyAny>> for WireValue {
87    fn from(obj: Py<PyAny>) -> Self {
88        monarch_with_gil_blocking(GilSite::Convert, |py| {
89            WireValue::PyObject(PickledPyObject::pickle(obj.bind(py)).unwrap())
90        })
91    }
92}
monarch_messages/wire_value.rs

monarch_messages/
wire_value.rs
