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    fn extract_bound(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
56        Ok(WireValue::PyObject(PickledPyObject::pickle(obj)?))
57    }
58}
59
60impl<'py> IntoPyObject<'py> for WireValue {
61    type Target = PyAny;
62    type Output = Bound<'py, PyAny>;
63    type Error = PyErr;
64
65    fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
66        match self {
67            WireValue::Ref(ref_) => ref_.into_bound_py_any(py),
68            WireValue::RefList(ref_list) => ref_list.clone().into_bound_py_any(py),
69            WireValue::Int(int) => int.into_bound_py_any(py),
70            WireValue::IntList(int_list) => int_list.clone().into_bound_py_any(py),
71            WireValue::Double(double) => double.into_bound_py_any(py),
72            WireValue::Bool(bool_) => bool_.into_bound_py_any(py),
73            WireValue::String(string) => string.into_bound_py_any(py),
74            WireValue::Device(device) => device.into_bound_py_any(py),
75            WireValue::Layout(val) => val.into_bound_py_any(py),
76            WireValue::ScalarType(val) => val.into_bound_py_any(py),
77            WireValue::MemoryFormat(val) => val.into_bound_py_any(py),
78            WireValue::None(()) => PyNone::get(py).into_bound_py_any(py),
79            WireValue::PyObject(val) => val.unpickle(py),
80        }
81    }
82}
83
84impl From<Py<PyAny>> for WireValue {
85    fn from(obj: Py<PyAny>) -> Self {
86        monarch_with_gil_blocking(GilSite::Convert, |py| {
87            WireValue::PyObject(PickledPyObject::pickle(obj.bind(py)).unwrap())
88        })
89    }
90}
monarch_messages/wire_value.rs

monarch_messages/
wire_value.rs
