# 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
import sys
from typing import Any, Callable, Dict, Iterator, TYPE_CHECKING, TypedDict
from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport
from monarch._rust_bindings.monarch_hyperactor.config import (
clear_runtime_config as _clear_runtime_config,
configure as _configure,
Encoding,
get_global_config as _get_global_config,
get_runtime_config as _get_runtime_config,
)
__all__ = [
"clear_runtime_config",
"configure",
"configured",
"Encoding",
"get_global_config",
"get_runtime_config",
"parametrize_config",
"parametrize_config_pointwise",
]
if TYPE_CHECKING:
if sys.version_info >= (3, 12):
from typing import NotRequired, Unpack
class ConfigureArgs(TypedDict):
default_transport: NotRequired[ChannelTransport | str]
enable_log_forwarding: NotRequired[bool]
enable_file_capture: NotRequired[bool]
tail_log_lines: NotRequired[int]
codec_max_frame_length: NotRequired[int]
message_delivery_timeout: NotRequired[str]
host_spawn_ready_timeout: NotRequired[str]
mesh_proc_spawn_max_idle: NotRequired[str]
process_exit_timeout: NotRequired[str]
message_ack_time_interval: NotRequired[str]
message_ack_every_n_messages: NotRequired[int]
message_ttl_default: NotRequired[int]
split_max_buffer_size: NotRequired[int]
split_max_buffer_age: NotRequired[str]
stop_actor_timeout: NotRequired[str]
cleanup_timeout: NotRequired[str]
default_encoding: NotRequired[Encoding]
channel_net_rx_buffer_full_check_interval: NotRequired[str]
message_latency_sampling_rate: NotRequired[float]
enable_dest_actor_reordering_buffer: NotRequired[bool]
mesh_bootstrap_enable_pdeathsig: NotRequired[bool]
mesh_terminate_concurrency: NotRequired[int]
mesh_terminate_timeout: NotRequired[str]
small_write_threshold: NotRequired[int]
max_cast_dimension_size: NotRequired[int]
remote_alloc_bind_to_inaddr_any: NotRequired[bool]
remote_alloc_bootstrap_addr: NotRequired[str]
remote_alloc_allowed_port_range: NotRequired[slice]
read_log_buffer: NotRequired[int]
force_file_log: NotRequired[bool]
prefix_with_rank: NotRequired[bool]
actor_spawn_max_idle: NotRequired[str]
get_actor_state_max_idle: NotRequired[str]
supervision_watchdog_timeout: NotRequired[str]
proc_stop_max_idle: NotRequired[str]
get_proc_state_max_idle: NotRequired[str]
actor_queue_dispatch: NotRequired[bool]
mesh_admin_addr: NotRequired[str]
mesh_attach_config_timeout: NotRequired[str]
mesh_orphan_timeout: NotRequired[str]
rdma_allow_tcp_fallback: NotRequired[bool]
rdma_disable_ibverbs: NotRequired[bool]
rdma_max_chunk_size_mb: NotRequired[int]
# pyrefly: ignore [invalid-annotation]
ConfigureKwargsType = Unpack[ConfigureArgs]
else:
ConfigureKwargsType = object
[docs]def get_global_config() -> Dict[str, Any]:
"""Return a merged view of all configuration layers.
The resulting dict includes defaults, environment overrides, file-based
settings, and the current Runtime layer. Mutating the returned dict does
*not* change the active configuration; use :func:`configure` instead.
"""
return _get_global_config()
[docs]def get_runtime_config() -> Dict[str, Any]:
"""Return a snapshot of just the Runtime layer configuration.
Useful for snapshot/restore flows (see :func:`configured`) or for
inspecting which keys were last set via Python.
"""
return _get_runtime_config()
[docs]def clear_runtime_config() -> None:
"""Remove every key from the Runtime configuration layer.
Environment variables, config files, and defaults are untouched. This is
typically paired with :func:`configure` to reset overrides in long-lived
processes.
"""
_clear_runtime_config()
def parametrize_config(
**config_options: set[Any],
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
"""Create a pytest parametrize decorator for configuration cross-products.
This decorator runs the test function under every combination of the
specified configuration values. Each test invocation wraps the test body
in a `configured(...)` context manager with the corresponding settings.
Args:
**config_options: Configuration keys mapped to sets of values to test.
Each key should be a valid argument to `configure()`.
Returns:
A decorator that parametrizes and wraps the test function.
Example:
>>> from monarch.config import parametrize_config
>>>
>>> @parametrize_config(
... actor_queue_dispatch={True, False},
... prefix_with_rank={True, False},
... )
... async def test_actor_feature():
... # Test runs 4 times: all combinations of the two bool options
... pass
"""
import asyncio
import functools
import inspect
import itertools
import pytest
if not config_options:
raise ValueError("parametrize_config requires at least one config option")
keys = list(config_options.keys())
value_lists = [list(config_options[k]) for k in keys]
combinations = list(itertools.product(*value_lists))
# Create parameter IDs for clearer test output
param_ids = [
"-".join(f"{k}={v}" for k, v in zip(keys, combo)) for combo in combinations
]
# Create the config dicts for each combination
config_dicts = [dict(zip(keys, combo)) for combo in combinations]
def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
# Get original function's signature and add _config_overrides as first param
orig_sig = inspect.signature(fn)
new_params = [
inspect.Parameter(
"_config_overrides", inspect.Parameter.POSITIONAL_OR_KEYWORD
)
] + list(orig_sig.parameters.values())
new_sig = orig_sig.replace(parameters=new_params)
if asyncio.iscoroutinefunction(fn):
async def async_wrapper(
_config_overrides: Dict[str, Any], *args: Any, **kwargs: Any
) -> Any:
with configured(**_config_overrides):
return await fn(*args, **kwargs)
functools.update_wrapper(async_wrapper, fn)
async_wrapper.__signature__ = new_sig # type: ignore[attr-defined]
wrapped = async_wrapper
else:
def sync_wrapper(
_config_overrides: Dict[str, Any], *args: Any, **kwargs: Any
) -> Any:
with configured(**_config_overrides):
return fn(*args, **kwargs)
functools.update_wrapper(sync_wrapper, fn)
sync_wrapper.__signature__ = new_sig # type: ignore[attr-defined]
wrapped = sync_wrapper
return pytest.mark.parametrize(
"_config_overrides", config_dicts, ids=param_ids
)(wrapped)
return decorator
def parametrize_config_pointwise(
**config_options: Any,
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
"""Create a pytest parametrize decorator for pointwise configuration sweeps.
Like :func:`parametrize_config`, but iterates the input sequences in
lockstep (``zip``) rather than computing their cartesian product. Every
sequence must have the same length; the test is invoked once per index
with one value drawn from each sequence at that position.
Args:
**config_options: Configuration keys mapped to ordered sequences of
values. Each key should be a valid argument to ``configure()``.
Returns:
A decorator that parametrizes and wraps the test function.
Example:
>>> from monarch.config import parametrize_config_pointwise
>>>
>>> @parametrize_config_pointwise(
... actor_queue_dispatch=[True, False],
... prefix_with_rank=[True, False],
... )
... async def test_actor_feature():
... # Runs 2 times:
... # (actor_queue_dispatch=True, prefix_with_rank=True)
... # (actor_queue_dispatch=False, prefix_with_rank=False)
... pass
"""
import asyncio
import functools
import inspect
import pytest
if not config_options:
raise ValueError(
"parametrize_config_pointwise requires at least one config option"
)
keys = list(config_options.keys())
value_lists = [list(config_options[k]) for k in keys]
lengths = {len(v) for v in value_lists}
if len(lengths) != 1:
raise ValueError(
"parametrize_config_pointwise requires every option to have the "
f"same number of values; got lengths {dict(zip(keys, [len(v) for v in value_lists]))}"
)
combinations = list(zip(*value_lists))
param_ids = [
"-".join(f"{k}={v}" for k, v in zip(keys, combo)) for combo in combinations
]
config_dicts = [dict(zip(keys, combo)) for combo in combinations]
def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
orig_sig = inspect.signature(fn)
new_params = [
inspect.Parameter(
"_config_overrides", inspect.Parameter.POSITIONAL_OR_KEYWORD
)
] + list(orig_sig.parameters.values())
new_sig = orig_sig.replace(parameters=new_params)
if asyncio.iscoroutinefunction(fn):
async def async_wrapper(
_config_overrides: Dict[str, Any], *args: Any, **kwargs: Any
) -> Any:
with configured(**_config_overrides):
return await fn(*args, **kwargs)
functools.update_wrapper(async_wrapper, fn)
async_wrapper.__signature__ = new_sig # type: ignore[attr-defined]
wrapped = async_wrapper
else:
def sync_wrapper(
_config_overrides: Dict[str, Any], *args: Any, **kwargs: Any
) -> Any:
with configured(**_config_overrides):
return fn(*args, **kwargs)
functools.update_wrapper(sync_wrapper, fn)
sync_wrapper.__signature__ = new_sig # type: ignore[attr-defined]
wrapped = sync_wrapper
return pytest.mark.parametrize(
"_config_overrides", config_dicts, ids=param_ids
)(wrapped)
return decorator
