[monarch] Typed global configuration from python integrated with the Attrs system

Pull Request resolved: #1414

This diff introduces a typed global configuration system in python for monarch, integrated with the hyperactor `Attrs` system and the `hyperactor::global::config` module. To make an attribute configurable via python, first add the meta attribute `PYTHON_CONFIG_KEY`, like so:
```
@meta(
    CONFIG_ENV_VAR = "HYPERACTOR_MESH_DEFAULT_TRANSPORT".to_string(),
    PYTHON_CONFIG_KEY = "default_transport".to_string(),
)
pub attr DEFAULT_TRANSPORT: ChannelTransport = ChannelTransport::Unix;
```
This makes it so that the `DEFAULT_TRANSPORT` attr can be configured using `monarch.configure(default_transport=...)`. In order to ensure that attrs with type `ChannelTransport` can be configured by passing `PyChannelTransport` values to `monarch.configure`, in `monarch_hyperactor/src/config.rs`, add the macro invocation:
```
declare_py_config_type!(PyChannelTransport as ChannelTransport);
```
For attrs with rust types that can be passed directly to and from python (like `String`), simply add:
```
declare_py_config_type!(String);
```
These macro invocations only need to be added once per unique attr type we want to support -- *not* once per actual attr.

From python, the `DEFAULT_TRANSPORT` attr is then configurable in the global config by calling
```
monarch._rust_bindings.monarch_hyperactor.config.configure(
    default_transport=ChannelTransport.MetaTlsWithHostname,
    ...
)
```
You can then get the currently configured values with
```
monarch._rust_bindings.monarch_hyperactor.config.get_configuration()
```
which returns
```
{
    "default_transport": ChannelTransport.MetaTlsWithHostname,
    ...
}
```
ghstack-source-id: 313996274

Differential Revision: [D83701581](https://our.internmc.facebook.com/intern/diff/D83701581/)

Author: samlurye

hyperactor/src/attrs.rs

@@ -135,6 +135,9 @@ pub struct AttrKeyInfo {
     pub parse: fn(&str) -> Result<Box<dyn SerializableValue>, anyhow::Error>,
     /// Default value for the attribute, if any.
     pub default: Option<&'static dyn SerializableValue>,
+    /// A reference to the relevant key object with the associated
+    /// type parameter erased. Can be downcast to a concrete Key<T>.
+    pub erased: &'static dyn ErasedKey,
 }
 
 inventory::collect!(AttrKeyInfo);
@@ -199,6 +202,39 @@ impl<T: 'static> Clone for Key<T> {
 
 impl<T: 'static> Copy for Key<T> {}
 
+/// A trait for type-erased keys.
+pub trait ErasedKey: Any + Send + Sync + 'static {
+    /// The name of the key.
+    fn name(&self) -> &'static str;
+
+    /// The typehash of the key's associated type.
+    fn typehash(&self) -> u64;
+
+    /// The typename of the key's associated type.
+    fn typename(&self) -> &'static str;
+}
+
+impl dyn ErasedKey {
+    /// Downcast a type-erased key to a specific key type.
+    pub fn downcast_ref<T: Named + 'static>(&'static self) -> Option<&'static Key<T>> {
+        (self as &dyn Any).downcast_ref::<Key<T>>()
+    }
+}
+
+impl<T: AttrValue> ErasedKey for Key<T> {
+    fn name(&self) -> &'static str {
+        self.name
+    }
+
+    fn typehash(&self) -> u64 {
+        T::typehash()
+    }
+
+    fn typename(&self) -> &'static str {
+        T::typename()
+    }
+}
+
 // Enable attr[key] syntax.
 impl<T: AttrValue> Index<Key<T>> for Attrs {
     type Output = T;
@@ -772,6 +808,7 @@ macro_rules! declare_attrs {
                     Ok(Box::new(value) as Box<dyn $crate::attrs::SerializableValue>)
                 },
                 default: Some($crate::paste! { &[<$name _DEFAULT>] }),
+                erased: &$name,
             }
         }
     };
@@ -824,6 +861,7 @@ macro_rules! declare_attrs {
                     Ok(Box::new(value) as Box<dyn $crate::attrs::SerializableValue>)
                 },
                 default: None,
+                erased: &$name,
             }
         }
     };

hyperactor/src/config.rs

@@ -35,6 +35,10 @@ declare_attrs! {
     /// key.
     pub attr CONFIG_ENV_VAR: String;
 
+    /// This is a meta-attribute specifying the name of the kwarg to pass to monarch.configure()
+    /// to set the attribute value in the global config.
+    pub attr PYTHON_CONFIG_KEY: String;
+
     /// Maximum frame length for codec
     @meta(CONFIG_ENV_VAR = "HYPERACTOR_CODEC_MAX_FRAME_LENGTH".to_string())
     pub attr CODEC_MAX_FRAME_LENGTH: usize = 10 * 1024 * 1024 * 1024; // 10 GiB
@@ -244,6 +248,12 @@ pub mod global {
         CONFIG.read().unwrap().get(key).unwrap().clone()
     }
 
+    /// Get a key from the global configuration by cloning the value,
+    /// if it exists. Returns None if the key is not present.
+    pub fn try_get_cloned<T: AttrValue>(key: Key<T>) -> Option<T> {
+        CONFIG.read().unwrap().get(key).cloned()
+    }
+
     /// Get the global attrs
     pub fn attrs() -> Attrs {
         CONFIG.read().unwrap().clone()
@@ -258,6 +268,12 @@ pub mod global {
         *config = Attrs::new();
     }
 
+    /// Set a key in the global configuration.
+    pub fn set<T: AttrValue>(key: Key<T>, value: T) {
+        let mut config = CONFIG.write().unwrap();
+        config.insert_value(key, Box::new(value));
+    }
+
     /// A guard that holds the global configuration lock and provides override functionality.
     ///
     /// This struct acts as both a lock guard (preventing other tests from modifying global config)

hyperactor_mesh/src/proc_mesh.rs

@@ -33,6 +33,7 @@ use hyperactor::channel::ChannelAddr;
 use hyperactor::channel::ChannelTransport;
 use hyperactor::config;
 use hyperactor::config::CONFIG_ENV_VAR;
+use hyperactor::config::PYTHON_CONFIG_KEY;
 use hyperactor::context;
 use hyperactor::declare_attrs;
 use hyperactor::mailbox;
@@ -85,10 +86,12 @@ use std::sync::RwLock;
 
 declare_attrs! {
     /// Default transport type to use across the application.
-    @meta(CONFIG_ENV_VAR = "HYPERACTOR_MESH_DEFAULT_TRANSPORT".to_string())
-    attr DEFAULT_TRANSPORT: ChannelTransport = ChannelTransport::Unix;
+    @meta(
+        CONFIG_ENV_VAR = "HYPERACTOR_MESH_DEFAULT_TRANSPORT".to_string(),
+        PYTHON_CONFIG_KEY = "default_transport".to_string(),
+    )
+    pub attr DEFAULT_TRANSPORT: ChannelTransport = ChannelTransport::Unix;
 }
-
 /// Get the default transport type to use across the application.
 pub fn default_transport() -> ChannelTransport {
     config::global::get_cloned(DEFAULT_TRANSPORT)

monarch_hyperactor/src/channel.rs

@@ -13,6 +13,7 @@ use hyperactor::channel::ChannelTransport;
 use hyperactor::channel::MetaTlsAddr;
 use hyperactor::channel::TlsMode;
 use pyo3::exceptions::PyRuntimeError;
+use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 
 /// Python binding for [`hyperactor::channel::ChannelTransport`]
@@ -31,6 +32,33 @@ pub enum PyChannelTransport {
     // Sim(/*transport:*/ ChannelTransport), TODO kiuk@ add support
 }
 
+#[pymethods]
+impl PyChannelTransport {
+    fn get(&self) -> Self {
+        self.clone()
+    }
+}
+
+impl TryFrom<ChannelTransport> for PyChannelTransport {
+    type Error = PyErr;
+
+    fn try_from(transport: ChannelTransport) -> PyResult<Self> {
+        match transport {
+            ChannelTransport::Tcp => Ok(PyChannelTransport::Tcp),
+            ChannelTransport::MetaTls(TlsMode::Hostname) => {
+                Ok(PyChannelTransport::MetaTlsWithHostname)
+            }
+            ChannelTransport::MetaTls(TlsMode::IpV6) => Ok(PyChannelTransport::MetaTlsWithIpV6),
+            ChannelTransport::Local => Ok(PyChannelTransport::Local),
+            ChannelTransport::Unix => Ok(PyChannelTransport::Unix),
+            _ => Err(PyValueError::new_err(format!(
+                "unsupported transport: {}",
+                transport
+            ))),
+        }
+    }
+}
+
 #[pyclass(
     name = "ChannelAddr",
     module = "monarch._rust_bindings.monarch_hyperactor.channel"

monarch_hyperactor/src/config.rs

@@ -11,9 +11,23 @@
 //! This module provides monarch-specific configuration attributes that extend
 //! the base hyperactor configuration system.
 
+use std::collections::HashMap;
+use std::fmt::Debug;
+
+use hyperactor::AttrValue;
+use hyperactor::Named;
+use hyperactor::attrs::AttrKeyInfo;
+use hyperactor::attrs::ErasedKey;
 use hyperactor::attrs::declare_attrs;
+use hyperactor::channel::ChannelTransport;
+use hyperactor::config::PYTHON_CONFIG_KEY;
+use pyo3::conversion::IntoPyObjectExt;
+use pyo3::exceptions::PyTypeError;
+use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 
+use crate::channel::PyChannelTransport;
+
 // Declare monarch-specific configuration keys
 declare_attrs! {
     /// Use a single asyncio runtime for all Python actors, rather than one per actor
@@ -30,6 +44,183 @@ pub fn reload_config_from_env() -> PyResult<()> {
     Ok(())
 }
 
+/// Map from name of the kwarg that will be passed to `monarch.configure(...)`
+/// to the `Key<T>`` associated with that kwarg. This contains all of the
+/// attribute keys with meta-attribute `PYTHON_CONFIG_KEY`.
+static KEY_BY_NAME: std::sync::LazyLock<HashMap<&'static str, &'static dyn ErasedKey>> =
+    std::sync::LazyLock::new(|| {
+        inventory::iter::<AttrKeyInfo>()
+            .filter_map(|info| {
+                info.meta
+                    .get(PYTHON_CONFIG_KEY)
+                    .map(|py_name| (py_name.as_str(), info.erased))
+            })
+            .collect()
+    });
+
+/// Map from typehash to an info struct that can be used to downcast an `ErasedKey`
+/// to a concrete `Key<T>` and use it to get/set values in the global configl
+static TYPEHASH_TO_INFO: std::sync::LazyLock<HashMap<u64, &'static PythonConfigTypeInfo>> =
+    std::sync::LazyLock::new(|| {
+        inventory::iter::<PythonConfigTypeInfo>()
+            .map(|info| ((info.typehash)(), info))
+            .collect()
+    });
+
+/// Given a key, get the associated `T`-typed value from the global config, then
+/// convert it to a `P`-typed object that can be converted to PyObject, and
+/// return that PyObject.
+fn get_global_config<'py, P, T>(
+    py: Python<'py>,
+    key: &'static dyn ErasedKey,
+) -> PyResult<Option<PyObject>>
+where
+    T: AttrValue + TryInto<P>,
+    P: IntoPyObjectExt<'py>,
+    PyErr: From<<T as TryInto<P>>::Error>,
+{
+    // Well, it can't fail unless there's a bug in the code in this file.
+    let key = key.downcast_ref::<T>().expect("cannot fail");
+    let val: Option<P> = hyperactor::config::global::try_get_cloned(key.clone())
+        .map(|v| v.try_into())
+        .transpose()?;
+    val.map(|v| v.into_py_any(py)).transpose()
+}
+
+fn set_global_config<T: AttrValue + Debug>(key: &'static dyn ErasedKey, value: T) -> PyResult<()> {
+    // Again, can't fail unless there's a bug in the code in this file.
+    let key = key.downcast_ref().expect("cannot fail");
+    hyperactor::config::global::set(*key, value.clone());
+    Ok(())
+}
+
+fn set_global_config_from_py_obj(py: Python<'_>, name: &str, val: PyObject) -> PyResult<()> {
+    // Get the `ErasedKey` from the kwarg `name` passed to `monarch.configure(...)`.
+    let key = match KEY_BY_NAME.get(name) {
+        None => {
+            return Err(PyValueError::new_err(format!(
+                "invalid configuration key: `{}`",
+                name
+            )));
+        }
+        Some(key) => *key,
+    };
+
+    // Using the typehash from the erased key, get/call the function that can downcast
+    // the key and set the value on the global config.
+    match TYPEHASH_TO_INFO.get(&key.typehash()) {
+        None => Err(PyTypeError::new_err(format!(
+            "configuration key `{}` has type `{}`, but configuring with values of this type from Python is not supported.",
+            name,
+            key.typename()
+        ))),
+        Some(info) => (info.set_global_config)(py, key, val),
+    }
+}
+
+/// Struct to associate a typehash with functions for getting/setting
+/// values in the global config with keys of type `Key<T>`, where
+/// `T::typehash() == PythonConfigTypeInfo::typehash()`.
+struct PythonConfigTypeInfo {
+    typehash: fn() -> u64,
+    set_global_config:
+        fn(py: Python<'_>, key: &'static dyn ErasedKey, val: PyObject) -> PyResult<()>,
+    get_global_config:
+        fn(py: Python<'_>, key: &'static dyn ErasedKey) -> PyResult<Option<PyObject>>,
+}
+
+inventory::collect!(PythonConfigTypeInfo);
+
+/// Macro to declare that keys of this type can be configured
+/// from python using `monarch.configure(...)`. For types
+/// 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`
+/// as an intermediate step), the usage is
+/// `declare_py_config_type!(PyChannelTransport as ChannelTransport)`.
+macro_rules! declare_py_config_type {
+    ($($ty:ty),+ $(,)?) => {
+        hyperactor::paste! {
+            $(
+                hyperactor::submit! {
+                    PythonConfigTypeInfo {
+                        typehash: $ty::typehash,
+                        set_global_config: |py, key, val| {
+                            let val: $ty = val.extract::<$ty>(py).map_err(|err| PyTypeError::new_err(format!(
+                                "invalid value `{}` for configuration key `{}` ({})",
+                                val, key.name(), err
+                            )))?;
+                            set_global_config(key, val)
+                        },
+                        get_global_config: |py, key| {
+                            get_global_config::<$ty, $ty>(py, key)
+                        }
+                    }
+                }
+            )+
+        }
+    };
+    ($py_ty:ty as $ty:ty) => {
+        hyperactor::paste! {
+            hyperactor::submit! {
+                PythonConfigTypeInfo {
+                    typehash: $ty::typehash,
+                    set_global_config: |py, key, val| {
+                        let val: $ty = val.extract::<$py_ty>(py).map_err(|err| PyTypeError::new_err(format!(
+                            "invalid value `{}` for configuration key `{}` ({})",
+                            val, key.name(), err
+                        )))?.into();
+                        set_global_config(key, val)
+                    },
+                    get_global_config: |py, key| {
+                        get_global_config::<$py_ty, $ty>(py, key)
+                    }
+                }
+            }
+        }
+    };
+}
+
+declare_py_config_type!(PyChannelTransport as ChannelTransport);
+declare_py_config_type!(
+    i8, i16, i32, i64, u8, u16, u32, u64, usize, f32, f64, bool, String
+);
+
+/// Iterate over each key-value pair. Attempt to retrieve the `Key<T>`
+/// associated with the key and convert the value to `T`, then set
+/// them on the global config. The association between kwarg and
+/// `Key<T>` is specified using the `PYTHON_CONFIG_KEY` meta-attribute.
+#[pyfunction]
+#[pyo3(signature = (**kwargs))]
+fn configure(py: Python<'_>, kwargs: Option<HashMap<String, PyObject>>) -> PyResult<()> {
+    kwargs
+        .map(|kwargs| {
+            kwargs
+                .into_iter()
+                .try_for_each(|(key, val)| set_global_config_from_py_obj(py, &key, val))
+        })
+        .transpose()?;
+    Ok(())
+}
+
+/// For all attribute keys with meta-attribute `PYTHON_CONFIG_KEY` defined, return the
+/// current associated value in the global config. The key will not be present in the
+/// result of it has no value in the global config.
+#[pyfunction]
+fn get_configuration(py: Python<'_>) -> PyResult<HashMap<String, PyObject>> {
+    KEY_BY_NAME
+        .iter()
+        .filter_map(|(name, key)| match TYPEHASH_TO_INFO.get(&key.typehash()) {
+            None => None,
+            Some(info) => match (info.get_global_config)(py, *key) {
+                Err(err) => Some(Err(err)),
+                Ok(val) => val.map(|val| Ok(((*name).into(), val))),
+            },
+        })
+        .collect()
+}
+
 /// Register Python bindings for the config module
 pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
     let reload = wrap_pyfunction!(reload_config_from_env, module)?;
@@ -38,5 +229,20 @@ pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
         "monarch._rust_bindings.monarch_hyperactor.config",
     )?;
     module.add_function(reload)?;
+
+    let configure = wrap_pyfunction!(configure, module)?;
+    configure.setattr(
+        "__module__",
+        "monarch._rust_bindings.monarch_hyperactor.config",
+    )?;
+    module.add_function(configure)?;
+
+    let get_configuration = wrap_pyfunction!(get_configuration, module)?;
+    get_configuration.setattr(
+        "__module__",
+        "monarch._rust_bindings.monarch_hyperactor.config",
+    )?;
+    module.add_function(get_configuration)?;
+
     Ok(())
 }