monarch_hyperactor/
channel.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::str::FromStr;
10
11use hyperactor::channel::BindSpec;
12use hyperactor::channel::ChannelAddr;
13use hyperactor::channel::ChannelTransport;
14use hyperactor::channel::TcpMode;
15use hyperactor::channel::TlsAddr;
16use hyperactor::channel::TlsMode;
17use pyo3::IntoPyObjectExt;
18use pyo3::exceptions::PyRuntimeError;
19use pyo3::exceptions::PyTypeError;
20use pyo3::exceptions::PyValueError;
21use pyo3::prelude::*;
22
23/// Python binding for [`hyperactor::channel::ChannelTransport`]
24///
25/// This enum represents the basic transport types that can be represented
26/// as simple enum variants. For explicit addresses, use `PyBindSpec`.
27#[pyclass(
28    name = "ChannelTransport",
29    module = "monarch._rust_bindings.monarch_hyperactor.channel",
30    eq
31)]
32#[derive(PartialEq, Clone, Copy, Debug)]
33pub enum PyChannelTransport {
34    TcpWithLocalhost,
35    TcpWithHostname,
36    MetaTlsWithHostname,
37    MetaTlsWithIpV6,
38    Tls,
39    Quic,
40    MetaQuicWithHostname,
41    MetaQuicWithIpV6,
42    Local,
43    Unix,
44    // Sim(/*transport:*/ ChannelTransport), TODO kiuk@ add support
45}
46
47#[pymethods]
48impl PyChannelTransport {
49    fn get(&self) -> Self {
50        *self
51    }
52
53    fn __reduce__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)> {
54        let getattr_fn = py.import("builtins")?.getattr("getattr")?;
55        let variant_name = match self {
56            PyChannelTransport::TcpWithLocalhost => "TcpWithLocalhost",
57            PyChannelTransport::TcpWithHostname => "TcpWithHostname",
58            PyChannelTransport::MetaTlsWithHostname => "MetaTlsWithHostname",
59            PyChannelTransport::MetaTlsWithIpV6 => "MetaTlsWithIpV6",
60            PyChannelTransport::Tls => "Tls",
61            PyChannelTransport::Quic => "Quic",
62            PyChannelTransport::MetaQuicWithHostname => "MetaQuicWithHostname",
63            PyChannelTransport::MetaQuicWithIpV6 => "MetaQuicWithIpV6",
64            PyChannelTransport::Local => "Local",
65            PyChannelTransport::Unix => "Unix",
66        };
67        let cls = py
68            .import("monarch._rust_bindings.monarch_hyperactor.channel")?
69            .getattr("ChannelTransport")?;
70        let args = (cls, variant_name).into_bound_py_any(py)?;
71        Ok((getattr_fn, args))
72    }
73}
74
75impl TryFrom<ChannelTransport> for PyChannelTransport {
76    type Error = PyErr;
77
78    fn try_from(transport: ChannelTransport) -> PyResult<Self> {
79        match transport {
80            ChannelTransport::Tcp(TcpMode::Localhost) => Ok(PyChannelTransport::TcpWithLocalhost),
81            ChannelTransport::Tcp(TcpMode::Hostname) => Ok(PyChannelTransport::TcpWithHostname),
82            ChannelTransport::MetaTls(TlsMode::Hostname) => {
83                Ok(PyChannelTransport::MetaTlsWithHostname)
84            }
85            ChannelTransport::MetaTls(TlsMode::IpV6) => Ok(PyChannelTransport::MetaTlsWithIpV6),
86            ChannelTransport::Tls => Ok(PyChannelTransport::Tls),
87            ChannelTransport::Quic => Ok(PyChannelTransport::Quic),
88            ChannelTransport::MetaQuic(TlsMode::Hostname) => {
89                Ok(PyChannelTransport::MetaQuicWithHostname)
90            }
91            ChannelTransport::MetaQuic(TlsMode::IpV6) => Ok(PyChannelTransport::MetaQuicWithIpV6),
92            ChannelTransport::Local => Ok(PyChannelTransport::Local),
93            ChannelTransport::Unix => Ok(PyChannelTransport::Unix),
94        }
95    }
96}
97
98/// Python binding for [`hyperactor::channel::BindSpec`]
99#[pyclass(
100    name = "BindSpec",
101    module = "monarch._rust_bindings.monarch_hyperactor.channel"
102)]
103#[derive(Clone, Debug, PartialEq)]
104pub struct PyBindSpec {
105    inner: BindSpec,
106}
107
108#[pymethods]
109impl PyBindSpec {
110    /// Create a new PyBindSpec from a ChannelTransport enum, a string representation,
111    /// or another PyBindSpec object.
112    ///
113    /// Examples:
114    ///     PyBindSpec(ChannelTransport.Unix)
115    ///     PyBindSpec("tcp://127.0.0.1:8080")
116    ///     PyBindSpec(PyBindSpec(ChannelTransport.Unix))
117    #[new]
118    pub fn new(spec: &Bound<'_, PyAny>) -> PyResult<Self> {
119        // First try to extract as PyBindSpec (for when passing an existing spec)
120        if let Ok(bind_spec) = spec.extract::<PyBindSpec>() {
121            return Ok(bind_spec);
122        }
123
124        // Then try to extract as PyChannelTransport enum
125        if let Ok(py_transport) = spec.extract::<PyChannelTransport>() {
126            let transport: ChannelTransport = py_transport.into();
127            return Ok(PyBindSpec {
128                inner: BindSpec::Any(transport),
129            });
130        }
131
132        // Then try to extract as a string and parse it as a BindSpec
133        if let Ok(spec_str) = spec.extract::<String>() {
134            let bind_spec = BindSpec::from_str(&spec_str).map_err(|e| {
135                PyValueError::new_err(format!("invalid str for BindSpec '{}': {}", spec_str, e))
136            })?;
137            return Ok(PyBindSpec { inner: bind_spec });
138        }
139
140        Err(PyTypeError::new_err(
141            "expected ChannelTransport enum, BindSpec, or str",
142        ))
143    }
144
145    fn __str__(&self) -> String {
146        self.inner.to_string()
147    }
148
149    fn __repr__(&self) -> String {
150        format!("PyBindSpec({:?})", self.inner)
151    }
152
153    fn __eq__(&self, other: &Self) -> bool {
154        self.inner == other.inner
155    }
156}
157
158impl From<PyBindSpec> for BindSpec {
159    fn from(spec: PyBindSpec) -> Self {
160        spec.inner
161    }
162}
163
164impl From<BindSpec> for PyBindSpec {
165    fn from(spec: BindSpec) -> Self {
166        PyBindSpec { inner: spec }
167    }
168}
169
170#[pyclass(
171    name = "ChannelAddr",
172    module = "monarch._rust_bindings.monarch_hyperactor.channel"
173)]
174pub struct PyChannelAddr {
175    inner: ChannelAddr,
176}
177
178impl FromStr for PyChannelAddr {
179    type Err = anyhow::Error;
180    fn from_str(addr: &str) -> anyhow::Result<Self> {
181        let inner = ChannelAddr::from_str(addr)?;
182        Ok(Self { inner })
183    }
184}
185
186#[pymethods]
187impl PyChannelAddr {
188    /// Returns an "any" address for the given transport type.
189    /// Primarily used to bind servers. Returned string form of the address
190    /// is of the format `{transport}!{address}`. For example:
191    /// `tcp![::]:0`, `unix!@a0b1c2d3`, `metatls!devgpu001.pci.facebook.com:0`
192    #[staticmethod]
193    pub fn any(transport: PyChannelTransport) -> PyResult<String> {
194        Ok(ChannelAddr::any(transport.into()).to_string())
195    }
196
197    #[staticmethod]
198    pub fn parse(addr: &str) -> PyResult<Self> {
199        Ok(PyChannelAddr::from_str(addr)?)
200    }
201
202    /// Returns the port number (if any) of this channel address,
203    /// `0` for transports for which unix ports do not apply (e.g. `unix`, `local`)
204    pub fn get_port(&self) -> PyResult<u16> {
205        match &self.inner {
206            ChannelAddr::Tcp(socket_addr) => Ok(socket_addr.port()),
207            ChannelAddr::MetaTls(TlsAddr { port, .. })
208            | ChannelAddr::Tls(TlsAddr { port, .. })
209            | ChannelAddr::Quic(TlsAddr { port, .. })
210            | ChannelAddr::MetaQuic(TlsAddr { port, .. }) => Ok(*port),
211            ChannelAddr::Unix(_) | ChannelAddr::Local(_) => Ok(0),
212            _ => Err(PyRuntimeError::new_err(format!(
213                "unsupported transport: `{:?}` for channel address: `{}`",
214                self.inner.transport(),
215                self.inner
216            ))),
217        }
218    }
219
220    /// Returns the channel transport of this channel address.
221    pub fn get_transport(&self) -> PyResult<PyChannelTransport> {
222        let transport = self.inner.transport();
223        match transport {
224            ChannelTransport::Tcp(mode) => match mode {
225                TcpMode::Localhost => Ok(PyChannelTransport::TcpWithLocalhost),
226                TcpMode::Hostname => Ok(PyChannelTransport::TcpWithHostname),
227            },
228            ChannelTransport::MetaTls(mode) => match mode {
229                TlsMode::Hostname => Ok(PyChannelTransport::MetaTlsWithHostname),
230                TlsMode::IpV6 => Ok(PyChannelTransport::MetaTlsWithIpV6),
231            },
232            ChannelTransport::Tls => Ok(PyChannelTransport::Tls),
233            ChannelTransport::Quic => Ok(PyChannelTransport::Quic),
234            ChannelTransport::MetaQuic(mode) => match mode {
235                TlsMode::Hostname => Ok(PyChannelTransport::MetaQuicWithHostname),
236                TlsMode::IpV6 => Ok(PyChannelTransport::MetaQuicWithIpV6),
237            },
238            ChannelTransport::Local => Ok(PyChannelTransport::Local),
239            ChannelTransport::Unix => Ok(PyChannelTransport::Unix),
240        }
241    }
242}
243
244impl From<PyChannelTransport> for ChannelTransport {
245    fn from(val: PyChannelTransport) -> Self {
246        match val {
247            PyChannelTransport::TcpWithLocalhost => ChannelTransport::Tcp(TcpMode::Localhost),
248            PyChannelTransport::TcpWithHostname => ChannelTransport::Tcp(TcpMode::Hostname),
249            PyChannelTransport::MetaTlsWithHostname => ChannelTransport::MetaTls(TlsMode::Hostname),
250            PyChannelTransport::MetaTlsWithIpV6 => ChannelTransport::MetaTls(TlsMode::IpV6),
251            PyChannelTransport::Tls => ChannelTransport::Tls,
252            PyChannelTransport::Quic => ChannelTransport::Quic,
253            PyChannelTransport::MetaQuicWithHostname => {
254                ChannelTransport::MetaQuic(TlsMode::Hostname)
255            }
256            PyChannelTransport::MetaQuicWithIpV6 => ChannelTransport::MetaQuic(TlsMode::IpV6),
257            PyChannelTransport::Local => ChannelTransport::Local,
258            PyChannelTransport::Unix => ChannelTransport::Unix,
259        }
260    }
261}
262
263pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
264    hyperactor_mod.add_class::<PyChannelTransport>()?;
265    hyperactor_mod.add_class::<PyBindSpec>()?;
266    hyperactor_mod.add_class::<PyChannelAddr>()?;
267    Ok(())
268}
269
270#[cfg(test)]
271mod tests {
272    use super::*;
273
274    #[test]
275    // TODO: OSS: failed to retrieve ipv6 address
276    #[cfg_attr(not(fbcode_build), ignore)]
277    fn test_channel_any_and_parse() -> PyResult<()> {
278        // just make sure any() and parse() calls work for all transports
279        for transport in [
280            PyChannelTransport::TcpWithLocalhost,
281            PyChannelTransport::TcpWithHostname,
282            PyChannelTransport::Unix,
283            PyChannelTransport::MetaTlsWithHostname,
284            PyChannelTransport::MetaTlsWithIpV6,
285            PyChannelTransport::Tls,
286            PyChannelTransport::Quic,
287            PyChannelTransport::MetaQuicWithHostname,
288            PyChannelTransport::MetaQuicWithIpV6,
289            PyChannelTransport::Local,
290        ] {
291            let address = PyChannelAddr::any(transport)?;
292            let _ = PyChannelAddr::parse(&address)?;
293        }
294        Ok(())
295    }
296
297    #[test]
298    fn test_channel_addr_get_port() -> PyResult<()> {
299        assert_eq!(PyChannelAddr::parse("tcp![::]:26600")?.get_port()?, 26600);
300        assert_eq!(
301            PyChannelAddr::parse("metatls!devgpu1.pci.facebook.com:26600")?.get_port()?,
302            26600
303        );
304        assert_eq!(
305            PyChannelAddr::parse("quic!devgpu1.pci.facebook.com:26600")?.get_port()?,
306            26600
307        );
308        assert_eq!(
309            PyChannelAddr::parse("metaquic!devgpu1.pci.facebook.com:26600")?.get_port()?,
310            26600
311        );
312        assert_eq!(PyChannelAddr::parse("local!12345")?.get_port()?, 0);
313        assert_eq!(PyChannelAddr::parse("unix!@1a2b3c")?.get_port()?, 0);
314        Ok(())
315    }
316
317    #[test]
318    fn test_channel_addr_get_transport() -> PyResult<()> {
319        assert_eq!(
320            PyChannelAddr::parse("tcp![::1]:26600")?.get_transport()?,
321            PyChannelTransport::TcpWithLocalhost,
322        );
323        assert_eq!(
324            PyChannelAddr::parse("tcp![::]:26600")?.get_transport()?,
325            PyChannelTransport::TcpWithHostname,
326        );
327        assert_eq!(
328            PyChannelAddr::parse("metatls!devgpu001.pci.facebook.com:26600")?.get_transport()?,
329            PyChannelTransport::MetaTlsWithHostname
330        );
331        assert_eq!(
332            PyChannelAddr::parse("metatls!::1:26600")?.get_transport()?,
333            PyChannelTransport::MetaTlsWithIpV6
334        );
335        assert_eq!(
336            // IpV4 will fallback to hostname
337            PyChannelAddr::parse("metatls!127.0.0.1:26600")?.get_transport()?,
338            PyChannelTransport::MetaTlsWithHostname
339        );
340        assert_eq!(
341            PyChannelAddr::parse("quic!devgpu001.pci.facebook.com:26600")?.get_transport()?,
342            PyChannelTransport::Quic
343        );
344        assert_eq!(
345            PyChannelAddr::parse("metaquic!devgpu001.pci.facebook.com:26600")?.get_transport()?,
346            PyChannelTransport::MetaQuicWithHostname
347        );
348        assert_eq!(
349            PyChannelAddr::parse("metaquic!::1:26600")?.get_transport()?,
350            PyChannelTransport::MetaQuicWithIpV6
351        );
352        assert_eq!(
353            PyChannelAddr::parse("local!12345")?.get_transport()?,
354            PyChannelTransport::Local
355        );
356        assert_eq!(
357            PyChannelAddr::parse("unix!@1a2b3c")?.get_transport()?,
358            PyChannelTransport::Unix
359        );
360        Ok(())
361    }
362}
monarch_hyperactor/channel.rs

monarch_hyperactor/
channel.rs
