: new module 'config'

python/monarch/__init__.py

@@ -53,6 +53,13 @@
     from monarch.common.selection import Selection
     from monarch.common.stream import get_active_stream, Stream
     from monarch.common.tensor import reduce, reduce_, Tensor
+    from monarch.config import (  # noqa
+        clear_runtime_config,
+        configure,
+        configured,
+        get_global_config,
+        get_runtime_config,
+    )
     from monarch.fetch import fetch_shard, inspect, show
     from monarch.gradient_generator import grad_function, grad_generator
     from monarch.simulator.config import set_meta  # noqa
@@ -62,6 +69,11 @@
 
 _public_api = {
     "coalescing": ("monarch.common._coalescing", "coalescing"),
+    "clear_runtime_config": ("monarch.config", "clear_runtime_config"),
+    "configure": ("monarch.config", "configure"),
+    "configured": ("monarch.config", "configured"),
+    "get_global_config": ("monarch.config", "get_global_config"),
+    "get_runtime_config": ("monarch.config", "get_runtime_config"),
     "remote": ("monarch.common.remote", "remote"),
     "get_active_mesh": ("monarch.common.device_mesh", "get_active_mesh"),
     "no_mesh": ("monarch.common.device_mesh", "no_mesh"),
@@ -117,6 +129,11 @@ def __getattr__(name):
 # otherwise tools think the imports are unused
 __all__ = [
     "coalescing",
+    "clear_runtime_config",
+    "configure",
+    "configured",
+    "get_global_config",
+    "get_runtime_config",
     "get_active_mesh",
     "no_mesh",
     "remote",

python/monarch/_src/actor/__init__.py

@@ -10,7 +10,7 @@
 Monarch Actor API
 """
 
-from monarch._rust_bindings.monarch_hyperactor.config import configure
+from monarch.config import configure
 
 # Detect if we're running in IPython/Jupyter
 _in_ipython = False

python/monarch/config/__init__.py

@@ -0,0 +1,87 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+"""Configuration utilities for Monarch.
+
+This module provides utilities for managing Monarch's runtime
+configuration, particularly useful for testing and temporary
+configuration overrides.
+"""
+
+import contextlib
+from typing import Any, Dict, Iterator
+
+from monarch._rust_bindings.monarch_hyperactor.config import (
+    clear_runtime_config,
+    configure,
+    get_global_config,
+    get_runtime_config,
+)
+
+
+__all__ = [
+    "clear_runtime_config",
+    "configure",
+    "configured",
+    "get_global_config",
+    "get_runtime_config",
+]
+
+
+@contextlib.contextmanager
+def configured(**overrides) -> Iterator[Dict[str, Any]]:
+    """Temporarily apply Python-side config overrides for this
+    process.
+
+    This context manager:
+      * snapshots the current **Runtime** configuration layer
+        (`get_runtime_config()`),
+      * applies the given `overrides` via `configure(**overrides)`,
+        and
+      * yields the **merged** view of config (`get_global_config()`),
+        including defaults, env, file, and Runtime.
+
+    On exit it restores the previous Runtime layer by:
+      * clearing all Runtime entries, and
+      * re-applying the saved snapshot.
+
+    `configured` alters the global configuration; thus other threads
+    will be subject to the overriden configuration while the context
+    manager is active.
+
+    Thus: this is intended for tests, which run as single threads;
+    per-test overrides do not leak into other tests.
+
+    Args:
+        **overrides: Configuration key-value pairs to override for the
+            duration of the context.
+
+    Yields:
+        Dict[str, Any]: The merged global configuration including all
+            layers (defaults, environment, file, and runtime).
+
+    Example:
+        >>> from monarch.config import configured
+        >>> with configured(enable_log_forwarding=True, tail_log_lines=100):
+        ...     # Configuration is temporarily overridden
+        ...     assert get_global_config()["enable_log_forwarding"] is True
+        >>> # Configuration is automatically restored after the context
+
+    """
+    # Retrieve runtime
+    prev = get_runtime_config()
+    try:
+        # Merge overrides into runtime
+        configure(**overrides)
+
+        # Snapshot of merged config (all layers)
+        yield get_global_config()
+    finally:
+        # Restore previous runtime
+        clear_runtime_config()
+        configure(**prev)

python/tests/test_config.py

@@ -6,52 +6,10 @@
 
 # pyre-unsafe
 
-import contextlib
-from typing import Any, Dict, Iterator
-
 import pytest
 from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport
-from monarch._rust_bindings.monarch_hyperactor.config import (
-    clear_runtime_config,
-    configure,
-    get_global_config,
-    get_runtime_config,
-)
-
-
-@contextlib.contextmanager
-def configured(**overrides) -> Iterator[Dict[str, Any]]:
-    """Temporarily apply Python-side config overrides for this
-    process.
-
-    This context manager:
-      * snapshots the current **Runtime** configuration layer
-        (`get_runtime_configuration()`),
-      * applies the given `overrides` via `configure(**overrides)`,
-        and
-      * yields the **merged** view of config (`get_configuration()`),
-        including defaults, env, file, and Runtime.
-
-    On exit it restores the previous Runtime layer by:
-      * clearing all Runtime entries, and
-      * re-applying the saved snapshot.
-
-    This is intended for tests, so per-test overrides do not leak into
-    other tests.
-
-    """
-    # Retrieve runtime
-    prev = get_runtime_config()
-    try:
-        # Merge overrides into runtime
-        configure(**overrides)
-
-        # Snapshot of merged config (all layers)
-        yield get_global_config()
-    finally:
-        # Restore previous runtime
-        clear_runtime_config()
-        configure(**prev)
+from monarch._rust_bindings.monarch_hyperactor.config import get_global_config
+from monarch.config import configured
 
 
 def test_get_set_transport() -> None:

python/tests/test_python_actors.py

@@ -34,12 +34,6 @@
     PythonMessageKind,
 )
 from monarch._rust_bindings.monarch_hyperactor.alloc import Alloc, AllocSpec
-from monarch._rust_bindings.monarch_hyperactor.config import (
-    clear_runtime_config,
-    configure,
-    get_global_config,
-    get_runtime_config,
-)
 from monarch._rust_bindings.monarch_hyperactor.mailbox import (
     PortId,
     PortRef,
@@ -74,6 +68,7 @@
     endpoint,
     ProcMesh,
 )
+from monarch.config import configured
 from monarch.tools.config import defaults
 from typing_extensions import assert_type
 
@@ -461,41 +456,6 @@ def _handle_undeliverable_message(
         return True
 
 
-@contextlib.contextmanager
-def configured(**overrides) -> Iterator[Dict[str, Any]]:
-    """Temporarily apply Python-side config overrides for this
-    process.
-
-    This context manager:
-      * snapshots the current **Runtime** configuration layer
-        (`get_runtime_config()`),
-      * applies the given `overrides` via `configure(**overrides)`,
-        and
-      * yields the **merged** view of config (`get_global_config()`),
-        including defaults, env, file, and Runtime.
-
-    On exit it restores the previous Runtime layer by:
-      * clearing all Runtime entries, and
-      * re-applying the saved snapshot.
-
-    This is intended for tests, so per-test overrides do not leak into
-    other tests.
-
-    """
-    # Retrieve runtime
-    prev = get_runtime_config()
-    try:
-        # Merge overrides into runtime
-        configure(**overrides)
-
-        # Snapshot of merged config (all layers)
-        yield get_global_config()
-    finally:
-        # Restore previous runtime
-        clear_runtime_config()
-        configure(**prev)
-
-
 class RedirectedPaths(NamedTuple):
     """Filesystem paths for captured stdio from redirected_stdio().