diff --git a/python/monarch/__init__.py b/python/monarch/__init__.py index a7c420edf..c5db3d454 100644 --- a/python/monarch/__init__.py +++ b/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", diff --git a/python/monarch/_src/actor/__init__.py b/python/monarch/_src/actor/__init__.py index 02efd2014..7c7db10a4 100644 --- a/python/monarch/_src/actor/__init__.py +++ b/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 diff --git a/python/monarch/config/__init__.py b/python/monarch/config/__init__.py new file mode 100644 index 000000000..7da0a91d6 --- /dev/null +++ b/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) diff --git a/python/tests/test_config.py b/python/tests/test_config.py index 673b8c010..4157a306f 100644 --- a/python/tests/test_config.py +++ b/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: diff --git a/python/tests/test_python_actors.py b/python/tests/test_python_actors.py index 15f3fe3a0..8d4bbc453 100644 --- a/python/tests/test_python_actors.py +++ b/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().