Add support in enable_transport so the user can set an explicit TCP address (#2142)

Summary:

As explained in D89190087, Lightning needs a way to set and explicit TCP address in `enable_transport`. This diff is to implement that.

Differential Revision: D89185959

Author: pzhan9

monarch_hyperactor/src/channel.rs

@@ -14,10 +14,15 @@ use hyperactor::channel::MetaTlsAddr;
 use hyperactor::channel::TcpMode;
 use hyperactor::channel::TlsMode;
 use pyo3::exceptions::PyRuntimeError;
+use pyo3::exceptions::PyTypeError;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 
 /// Python binding for [`hyperactor::channel::ChannelTransport`]
+///
+/// This enum represents the basic transport types that can be represented
+/// as simple enum variants. For more complex transports like `Explicit`,
+/// use string-based configuration via `PyChannelTransportConfig`.
 #[pyclass(
     name = "ChannelTransport",
     module = "monarch._rust_bindings.monarch_hyperactor.channel",
@@ -62,6 +67,94 @@ impl TryFrom<ChannelTransport> for PyChannelTransport {
     }
 }
 
+/// A wrapper for ChannelTransport that can be created from either a
+/// PyChannelTransport enum or a string (for explicit transport).
+///
+/// We need this wrapper because Python's enum type does not support attaching
+/// data to enum variants like Rust does.
+#[pyclass(
+    name = "ChannelTransportConfig",
+    module = "monarch._rust_bindings.monarch_hyperactor.channel"
+)]
+#[derive(Clone, Debug, PartialEq)]
+pub struct PyChannelTransportConfig {
+    inner: ChannelTransport,
+}
+
+#[pymethods]
+impl PyChannelTransportConfig {
+    /// Create a new PyChannelTransportConfig from either a ChannelTransport enum
+    /// or a string representation.
+    ///
+    /// Examples:
+    ///     PyChannelTransportConfig(ChannelTransport.Unix)
+    ///     PyChannelTransportConfig("explicit:tcp://127.0.0.1:8080")
+    #[new]
+    pub fn new(transport: &Bound<'_, PyAny>) -> PyResult<Self> {
+        // First try to extract as PyChannelTransportConfig (for when passing an existing config)
+        if let Ok(config) = transport.extract::<PyChannelTransportConfig>() {
+            return Ok(config);
+        }
+
+        // Then try to extract as PyChannelTransport enum
+        if let Ok(py_transport) = transport.extract::<PyChannelTransport>() {
+            return Ok(PyChannelTransportConfig {
+                inner: py_transport.into(),
+            });
+        }
+
+        // Then try to extract as a string and parse it
+        if let Ok(transport_str) = transport.extract::<String>() {
+            if !transport_str.starts_with("explicit:") {
+                return Err(PyValueError::new_err(format!(
+                    "string argument only supports explicit transport with \
+                    address in the zmq url format (e.g., 'explicit:tcp://127.0.0.1:8080'); \
+                    but got: {}",
+                    transport_str,
+                )));
+            }
+            let addr_str = transport_str.strip_prefix("explicit:").unwrap();
+            let addr = ChannelAddr::from_zmq_url(addr_str).map_err(|e| {
+                PyValueError::new_err(format!(
+                    "invalid address string used for explicit transport '{}': {}",
+                    addr_str, e
+                ))
+            })?;
+            return Ok(PyChannelTransportConfig {
+                inner: ChannelTransport::Explicit(addr),
+            });
+        }
+
+        Err(PyTypeError::new_err(
+            "expected ChannelTransport enum, ChannelTransportConfig, or str",
+        ))
+    }
+
+    fn __str__(&self) -> String {
+        self.inner.to_string()
+    }
+
+    fn __repr__(&self) -> String {
+        format!("PyChannelTransportConfig({:?})", self.inner)
+    }
+
+    fn __eq__(&self, other: &Self) -> bool {
+        self.inner == other.inner
+    }
+}
+
+impl From<PyChannelTransportConfig> for ChannelTransport {
+    fn from(config: PyChannelTransportConfig) -> Self {
+        config.inner
+    }
+}
+
+impl From<ChannelTransport> for PyChannelTransportConfig {
+    fn from(transport: ChannelTransport) -> Self {
+        PyChannelTransportConfig { inner: transport }
+    }
+}
+
 #[pyclass(
     name = "ChannelAddr",
     module = "monarch._rust_bindings.monarch_hyperactor.channel"
@@ -149,6 +242,7 @@ impl From<PyChannelTransport> for ChannelTransport {
 #[pymodule]
 pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
     hyperactor_mod.add_class::<PyChannelTransport>()?;
+    hyperactor_mod.add_class::<PyChannelTransportConfig>()?;
     hyperactor_mod.add_class::<PyChannelAddr>()?;
     Ok(())
 }

monarch_hyperactor/src/config.rs

@@ -30,12 +30,13 @@ use hyperactor_config::attrs::Attrs;
 use hyperactor_config::attrs::ErasedKey;
 use hyperactor_config::attrs::declare_attrs;
 use hyperactor_config::global::Source;
+use pyo3::conversion::IntoPyObject;
 use pyo3::conversion::IntoPyObjectExt;
 use pyo3::exceptions::PyTypeError;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 
-use crate::channel::PyChannelTransport;
+use crate::channel::PyChannelTransportConfig;
 
 /// Python wrapper for Duration, using humantime format strings.
 ///
@@ -289,9 +290,9 @@ inventory::collect!(PythonConfigTypeInfo);
 /// like `String` that are convertible directly to/from PyObjects,
 /// you can just use `declare_py_config_type!(String)`. For types
 /// that must first be converted to/from a rust python wrapper
-/// (e.g., keys with type `ChannelTransport` must use `PyChannelTransport`
+/// (e.g., keys with type `ChannelTransport` must use `PyChannelTransportConfig`
 /// as an intermediate step), the usage is
-/// `declare_py_config_type!(PyChannelTransport as ChannelTransport)`.
+/// `declare_py_config_type!(PyChannelTransportConfig as ChannelTransport)`.
 macro_rules! declare_py_config_type {
     ($($ty:ty),+ $(,)?) => {
         hyperactor::paste! {
@@ -341,7 +342,7 @@ macro_rules! declare_py_config_type {
     };
 }
 
-declare_py_config_type!(PyChannelTransport as ChannelTransport);
+declare_py_config_type!(PyChannelTransportConfig as ChannelTransport);
 declare_py_config_type!(PyDuration as Duration);
 declare_py_config_type!(
     i8, i16, i32, i64, u8, u16, u32, u64, usize, f32, f64, bool, String
@@ -367,9 +368,20 @@ declare_py_config_type!(
 fn configure(py: Python<'_>, kwargs: Option<HashMap<String, PyObject>>) -> PyResult<()> {
     kwargs
         .map(|kwargs| {
-            kwargs
-                .into_iter()
-                .try_for_each(|(key, val)| configure_kwarg(py, &key, val))
+            kwargs.into_iter().try_for_each(|(key, val)| {
+                // Special handling for default_transport: convert ChannelTransport
+                // enum or explicit transport string to PyChannelTransportConfig
+                // before processing
+                let val = if key == "default_transport" {
+                    PyChannelTransportConfig::new(val.bind(py))?
+                        .into_pyobject(py)?
+                        .into_any()
+                        .unbind()
+                } else {
+                    val
+                };
+                configure_kwarg(py, &key, val)
+            })
         })
         .transpose()?;
     Ok(())

python/monarch/_rust_bindings/monarch_hyperactor/channel.pyi

@@ -9,6 +9,11 @@
 from enum import Enum
 
 class ChannelTransport(Enum):
+    """
+    Enum representing basic transport types for channels.
+    For the explicit transport type, use ChannelTransportConfig instead.
+    """
+
     TcpWithLocalhost = "tcp(localhost)"
     TcpWithHostname = "tcp(hostname)"
     MetaTlsWithHostname = "metatls(hostname)"
@@ -17,6 +22,25 @@ class ChannelTransport(Enum):
     Unix = "unix"
     # Sim  # TODO add support
 
+class ChannelTransportConfig:
+    """
+    Internal wrapper for ChannelTransport that accepts either a ChannelTransport enum
+    or a string for complex transports.
+
+    Note: This class is for internal use only. Users should pass ChannelTransport
+    enum values or strings directly to enable_transport().
+    """
+
+    def __init__(self, transport: ChannelTransport | str) -> None: ...
+    """
+        Basic transport types supported by ChannelTransport should be used directly as enum values.
+        For the explicit transport type, use ChannelTransportConfig instead.
+        - "explicit:<addr>": Use a specific channel address with zmq url format, e.g. "explicit:tcp://127.0.0.1:8080"
+    """
+    def __str__(self) -> str: ...
+    def __repr__(self) -> str: ...
+    def __eq__(self, other: object) -> bool: ...
+
 class ChannelAddr:
     @staticmethod
     def any(transport: ChannelTransport) -> str:

python/monarch/_rust_bindings/monarch_hyperactor/config.pyi

@@ -12,7 +12,10 @@ Type hints for the monarch_hyperactor.config Rust bindings.
 
 from typing import Any, Dict
 
-from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport
+from monarch._rust_bindings.monarch_hyperactor.channel import (
+    ChannelTransport,
+    ChannelTransportConfig,
+)
 
 def reload_config_from_env() -> None:
     """
@@ -32,7 +35,7 @@ def reset_config_to_defaults() -> None:
     ...
 
 def configure(
-    default_transport: ChannelTransport = ...,
+    default_transport: ChannelTransportConfig | ChannelTransport | str = ...,
     enable_log_forwarding: bool = ...,
     enable_file_capture: bool = ...,
     tail_log_lines: int = ...,
@@ -50,7 +53,10 @@ def configure(
     plus any additional CONFIG-marked keys passed via **kwargs.
 
     Args:
-        default_transport: Default channel transport for communication
+        default_transport: Default channel transport for communication. Can be:
+            - A ChannelTransportConfig object
+            - A ChannelTransport enum value (e.g., ChannelTransport.Unix)
+            - A string for explicit transport (e.g., "explicit::tcp://127.0.0.1:8080")
         enable_log_forwarding: Whether to forward logs from actors
         enable_file_capture: Whether to capture file output
         tail_log_lines: Number of log lines to tail

python/monarch/_src/actor/actor_mesh.py

@@ -53,7 +53,10 @@
 )
 from monarch._rust_bindings.monarch_hyperactor.actor_mesh import PythonActorMesh
 from monarch._rust_bindings.monarch_hyperactor.buffers import FrozenBuffer
-from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport
+from monarch._rust_bindings.monarch_hyperactor.channel import (
+    ChannelTransport,
+    ChannelTransportConfig,
+)
 from monarch._rust_bindings.monarch_hyperactor.config import configure
 from monarch._rust_bindings.monarch_hyperactor.context import Instance as HyInstance
 from monarch._rust_bindings.monarch_hyperactor.mailbox import (
@@ -410,7 +413,7 @@ def context() -> Context:
     return c
 
 
-_transport: Optional[ChannelTransport] = None
+_transport: Optional[ChannelTransportConfig] = None
 _transport_lock = threading.Lock()
 
 
@@ -424,17 +427,33 @@ def enable_transport(transport: "ChannelTransport | str") -> None:
     Currently only one transport type may be enabled at one time.
     In the future we may allow multiple to be enabled.
 
+    Supported transport values:
+        - ChannelTransport enum: ChannelTransport.Unix, ChannelTransport.TcpWithHostname, etc.
+        - string short cuts for the ChannelTransport enum:
+            - "tcp": ChannelTransport.TcpWithHostname
+            - "ipc": ChannelTransport.Unix
+            - "metatls": ChannelTransport.MetaTlsWithIpV6
+            - "metatls-hostname": ChannelTransport.MetaTlsWithHostname
+        - string for explicit transport. e.g.:
+            - "explicit:tcp://127.0.0.1:8080"
+
     For Meta usage, use metatls-hostname
     """
     if isinstance(transport, str):
-        transport = {
+        # Handle string shortcuts for the ChannelTransport enum,
+        resolved = {
             "tcp": ChannelTransport.TcpWithHostname,
             "ipc": ChannelTransport.Unix,
             "metatls": ChannelTransport.MetaTlsWithIpV6,
             "metatls-hostname": ChannelTransport.MetaTlsWithHostname,
         }.get(transport)
-        if transport is None:
-            raise ValueError(f"unknown transport: {transport}")
+        if resolved is not None:
+            transport_config = ChannelTransportConfig(resolved)
+        else:
+            transport_config = ChannelTransportConfig(transport)
+    else:
+        # ChannelTransport enum
+        transport_config = ChannelTransportConfig(transport)
 
     if _context.get(None) is not None:
         raise RuntimeError(
@@ -445,14 +464,14 @@ def enable_transport(transport: "ChannelTransport | str") -> None:
 
     global _transport
     with _transport_lock:
-        if _transport is not None and _transport != transport:
+        if _transport is not None and _transport != transport_config:
             raise RuntimeError(
                 f"Only one transport type may be enabled at one time. "
                 f"Currently enabled transport type is `{_transport}`. "
-                f"Attempted to enable transport type `{transport}`."
+                f"Attempted to enable transport type `{transport_config}`."
             )
-        _transport = transport
-    configure(default_transport=transport)
+        _transport = transport_config
+    configure(default_transport=transport_config)
 
 
 @dataclass