# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from typing import Any, Literal, overload
import kvikio._lib.defaults
from kvikio.utils import call_once
# Once the MULTI_POLL reactor pool has started (i.e. the first remote I/O has been
# issued), these properties are fixed for the remaining process lifetime: there is no
# value they can be reverted to, since even a plain (non-context-manager) `set()` call
# would fail at that point too. They are therefore not among the properties
# `ConfigContextManager` knows how to get/set, and cannot be used with `with`.
_PROCESS_LIFETIME_PROPERTIES = frozenset(
{
"remote_io_num_reactors",
"remote_io_reactor_dispatch",
"remote_io_max_concurrent_requests",
}
)
[docs]
class ConfigContextManager:
"""Context manager allowing the KvikIO configurations to be set upon entering a
`with` block, and automatically reset upon leaving the block.
"""
[docs]
def __init__(self, config: dict[str, str]):
(
self._property_getters,
self._property_setters,
) = self._property_getter_and_setter()
self._old_properties = {}
for key, value in config.items():
self._old_properties[key] = self._get_property(key)
self._set_property(key, value)
def __enter__(self):
return None
def __exit__(self, type_unused, value, traceback_unused):
for key, value in self._old_properties.items():
self._set_property(key, value)
def _get_property(self, property: str) -> Any:
if property == "num_threads":
property = "thread_pool_nthreads"
func = self._property_getters[property]
return func()
def _set_property(self, property: str, value: Any):
if property == "num_threads":
property = "thread_pool_nthreads"
func = self._property_setters[property]
func(value)
@call_once
def _property_getter_and_setter(self) -> tuple[dict[str, Any], dict[str, Any]]:
module_dict = vars(kvikio._lib.defaults)
property_getter_names = [
"compat_mode",
"thread_pool_nthreads",
"task_size",
"gds_threshold",
"bounce_buffer_size",
"http_max_attempts",
"http_status_codes",
"http_timeout",
"auto_direct_io_read",
"auto_direct_io_write",
"remote_io_backend",
]
property_getters = {}
property_setters = {}
for name in property_getter_names:
property_getters[name] = module_dict[name]
property_setters[name] = module_dict["set_" + name]
return property_getters, property_setters
@overload
def set(config: dict[str, Any], /) -> ConfigContextManager: ...
@overload
def set(
key: Literal[
"remote_io_num_reactors",
"remote_io_reactor_dispatch",
"remote_io_max_concurrent_requests",
],
value: Any,
/,
) -> None: ...
@overload
def set(key: str, value: Any, /) -> ConfigContextManager: ...
[docs]
def set(*config) -> ConfigContextManager | None:
"""Set KvikIO configurations.
Examples:
- To set one or more properties
.. code-block:: python
# Set the property globally.
kvikio.defaults.set({"prop1": value1, "prop2": value2})
# Set the property with a context manager.
# The property automatically reverts to its old value
# after leaving the `with` block.
with kvikio.defaults.set({"prop1": value1, "prop2": value2}):
...
- To set a single property
.. code-block:: python
# Set the property globally.
kvikio.defaults.set("prop", value)
# Set the property with a context manager.
# The property automatically reverts to its old value
# after leaving the `with` block.
with kvikio.defaults.set("prop", value):
...
Parameters
----------
config
The configurations. Can either be a single parameter (dict) consisting of one
or more properties, or two parameters key (string) and value (Any)
indicating a single property.
Valid configuration names are:
- ``"compat_mode"``
- ``"num_threads"``
- ``"task_size"``
- ``"gds_threshold"``
- ``"bounce_buffer_size"``
- ``"http_max_attempts"``
- ``"http_status_codes"``
- ``"http_timeout"``
- ``"auto_direct_io_read"``
- ``"auto_direct_io_write"``
- ``"remote_io_backend"``
The following are fixed for the remaining process lifetime once the first
remote I/O has been issued, so they must be set with
``kvikio.defaults.set(key, value)`` (not the dict form) and cannot be used
with a `with` block:
- ``"remote_io_num_reactors"``
- ``"remote_io_reactor_dispatch"``
- ``"remote_io_max_concurrent_requests"``
Returns
-------
ConfigContextManager or None
A context manager. If used in a `with` statement, the configuration will revert
to its old value upon leaving the block. `None` when setting one of the
process-lifetime configurations listed above, which cannot be reverted and thus
must be set with the two-argument form, not the dict form.
"""
err_msg = (
"Valid arguments are kvikio.defaults.set(config: dict) or "
"kvikio.defaults.set(key: str, value: Any)"
)
if len(config) == 1:
if not isinstance(config[0], dict):
raise ValueError(err_msg)
config_dict = config[0]
process_lifetime_keys = _PROCESS_LIFETIME_PROPERTIES.intersection(config_dict)
if process_lifetime_keys:
raise ValueError(
f"{sorted(process_lifetime_keys)} are fixed for the process's "
"remaining lifetime once the first remote I/O has been issued, so "
"they cannot be reverted and must be set individually with "
"kvikio.defaults.set(key, value), not as part of a dict."
)
return ConfigContextManager(config_dict)
elif len(config) == 2:
if not isinstance(config[0], str):
raise ValueError(err_msg)
key, value = config
if key in _PROCESS_LIFETIME_PROPERTIES:
getattr(kvikio._lib.defaults, "set_" + key)(value)
return None
return ConfigContextManager({key: value})
else:
raise ValueError(err_msg)
[docs]
def get(config_name: str) -> Any:
"""Get KvikIO configurations.
Parameters
----------
config_name: str
The name of the configuration.
Valid configuration names are:
- ``"compat_mode"``
- ``"num_threads"``
- ``"task_size"``
- ``"gds_threshold"``
- ``"bounce_buffer_size"``
- ``"http_max_attempts"``
- ``"http_status_codes"``
- ``"http_timeout"``
- ``"auto_direct_io_read"``
- ``"auto_direct_io_write"``
- ``"remote_io_backend"``
- ``"remote_io_num_reactors"``
- ``"remote_io_reactor_dispatch"``
- ``"remote_io_max_concurrent_requests"``
Returns
-------
Any
The value of the configuration.
"""
if config_name in _PROCESS_LIFETIME_PROPERTIES:
return getattr(kvikio._lib.defaults, config_name)()
context_manager = ConfigContextManager({})
return context_manager._get_property(config_name)
def is_compat_mode_preferred() -> bool:
return kvikio._lib.defaults.is_compat_mode_preferred()
