diff --git a/tests/test_general.py b/tests/test_general.py new file mode 100644 index 0000000..07ea1c4 --- /dev/null +++ b/tests/test_general.py @@ -0,0 +1,266 @@ +import os +import subprocess +import sys + +import pytest + +import xobjects as xo +from xobjects.general import Print +from xobjects.test_helpers import allow_kernel_compilation + + +def test_print_mode_default(capsys): + printer = Print() + + printer('visible') + + assert capsys.readouterr().out == 'visible\n' + + +def test_print_mode_suppressed(capsys): + printer = Print() + with xo.settings.override(print_mode='suppress'): + printer('hidden') + + assert capsys.readouterr().out == '' + + +def test_settings_control_xsuite_printer(capsys): + with xo.settings.override(print_mode='suppress'): + xo._print('hidden') + + assert capsys.readouterr().out == '' + + +def test_python_setting_overrides_environment_default(capsys): + with xo.settings.override(print_mode='print'): + xo._print('visible') + + assert capsys.readouterr().out == 'visible\n' + + +def test_invalid_print_mode_setting(): + with pytest.raises(ValueError, match='XSUITE_PRINT_MODE.*print.*suppress'): + xo.settings.print_mode = 'invalid' + + +def test_settings_override_restores_value_after_error(): + original = xo.settings.print_mode + + with pytest.raises(RuntimeError): + with xo.settings.override(print_mode='suppress'): + assert xo.settings.print_mode == 'suppress' + raise RuntimeError + + assert xo.settings.print_mode == original + + +def test_settings_are_discoverable(): + expected_settings = { + 'print_mode', + 'progress_indicator', + 'allow_kernel_compilation', + 'force_kernel_compilation', + 'show_kernel_diagnostics', + 'cffi_forbid_compile', + 'cffi_keep_build_files', + 'cuda_backend', + 'cuda_fast_compile', + 'cuda_compiler', + } + + assert expected_settings <= set(dir(xo.settings)) + for name in expected_settings: + assert f'{name}=' in repr(xo.settings) + assert name in type(xo.settings).__doc__ + + +def test_print_mode_environment_variable_is_startup_default(): + environment = os.environ.copy() + environment['XSUITE_PRINT_MODE'] = 'suppress' + code = ( + 'import xobjects as xo; ' + 'assert xo.settings.print_mode == "suppress"; ' + 'xo._print("hidden")') + + completed = subprocess.run( + [sys.executable, '-c', code], + env=environment, + capture_output=True, + text=True, + check=True, + ) + + assert completed.stdout == '' + + +def test_print_mode_environment_variable_can_be_overridden_in_python(): + environment = os.environ.copy() + environment['XSUITE_PRINT_MODE'] = 'suppress' + code = ( + 'import xobjects as xo; ' + 'xo.settings.print_mode = "print"; ' + 'assert xo.settings.print_mode == "print"; ' + 'xo._print("visible")') + + completed = subprocess.run( + [sys.executable, '-c', code], + env=environment, + capture_output=True, + text=True, + check=True, + ) + + assert completed.stdout == 'visible\n' + + +def test_environment_change_after_import_does_not_override_python_setting( + capsys, monkeypatch): + with xo.settings.override(print_mode='print'): + monkeypatch.setenv('XSUITE_PRINT_MODE', 'suppress') + xo._print('visible') + + assert capsys.readouterr().out == 'visible\n' + + +@pytest.mark.parametrize('value', ['1', 'true', 'YES', 'on']) +def test_boolean_environment_true_values(value): + environment = os.environ.copy() + environment['XSUITE_ALLOW_KERNEL_COMPILATION'] = value + code = ( + 'import xobjects as xo; ' + 'assert xo.settings.allow_kernel_compilation is True; ' + 'xo.settings.allow_kernel_compilation = False; ' + 'assert xo.settings.allow_kernel_compilation is False') + + subprocess.run( + [sys.executable, '-c', code], + env=environment, + capture_output=True, + text=True, + check=True, + ) + + +@pytest.mark.parametrize('value', ['0', 'false', 'NO', 'off']) +def test_boolean_environment_false_values(value): + environment = os.environ.copy() + environment['XSUITE_FORCE_KERNEL_COMPILATION'] = value + code = ( + 'import xobjects as xo; ' + 'assert xo.settings.force_kernel_compilation is False') + + subprocess.run( + [sys.executable, '-c', code], + env=environment, + capture_output=True, + text=True, + check=True, + ) + + +def test_invalid_boolean_environment_value(): + environment = os.environ.copy() + environment['XSUITE_CFFI_FORBID_COMPILE'] = 'sometimes' + + completed = subprocess.run( + [sys.executable, '-c', 'import xobjects'], + env=environment, + capture_output=True, + text=True, + ) + + assert completed.returncode != 0 + assert 'Invalid boolean value' in completed.stderr + assert 'xobjects.settings.cffi_forbid_compile' in completed.stderr + assert 'XSUITE_CFFI_FORBID_COMPILE' in completed.stderr + + +def test_runtime_settings_environment_defaults(): + environment = os.environ.copy() + environment.update({ + 'XSUITE_PROGRESS_INDICATOR': 'text', + 'XSUITE_FORCE_KERNEL_COMPILATION': 'yes', + 'XSUITE_SHOW_KERNEL_DIAGNOSTICS': 'on', + 'XSUITE_CFFI_FORBID_COMPILE': 'true', + 'XSUITE_CFFI_KEEP_BUILD_FILES': '1', + 'XSUITE_CUDA_BACKEND': 'clang', + 'XSUITE_CUDA_FAST_COMPILE': 'false', + 'XSUITE_CUDA_COMPILER': '/path/to/clang++', + }) + code = ( + 'import xobjects as xo; ' + 'assert xo.settings.progress_indicator == "text"; ' + 'assert xo.settings.force_kernel_compilation is True; ' + 'assert xo.settings.show_kernel_diagnostics is True; ' + 'assert xo.settings.cffi_forbid_compile is True; ' + 'assert xo.settings.cffi_keep_build_files is True; ' + 'assert xo.settings.cuda_backend == "clang"; ' + 'assert xo.settings.cuda_fast_compile is False; ' + 'assert xo.settings.cuda_compiler == "/path/to/clang++"') + + subprocess.run( + [sys.executable, '-c', code], + env=environment, + capture_output=True, + text=True, + check=True, + ) + + +@pytest.mark.parametrize( + 'allow, force, compilation_allowed', + [ + (False, False, False), + (True, False, True), + (False, True, True), + (True, True, True), + ], +) +def test_kernel_compilation_settings(allow, force, compilation_allowed): + with xo.settings.override( + allow_kernel_compilation=allow, + force_kernel_compilation=force, + ): + assert xo.context_cpu.kernel_compilation_allowed( + xo.ContextCpu()) is compilation_allowed + + +def test_user_context_environment_variable(monkeypatch): + monkeypatch.setenv('XOBJECTS_USER_CONTEXT', 'ContextCpu:auto') + context = xo.get_user_context() + + assert context.openmp_enabled + + +def test_test_contexts_environment_variable(monkeypatch): + monkeypatch.setenv( + 'XOBJECTS_TEST_CONTEXTS', + 'ContextCpu;ContextCpu:auto', + ) + + contexts = list(xo.context.get_test_contexts()) + + assert len(contexts) == 2 + assert contexts[0].openmp_enabled is False + assert contexts[1].openmp_enabled is True + + +def test_cffi_forbid_compile_setting(): + with xo.settings.override(cffi_forbid_compile=True): + with pytest.raises(RuntimeError) as err: + xo.ContextCpu().build_kernels({}) + + message = str(err.value) + assert 'xobjects.settings.cffi_forbid_compile' in message + assert 'XSUITE_CFFI_FORBID_COMPILE' in message + + +def test_allow_kernel_compilation_decorator_restores_state(): + @allow_kernel_compilation(skip_when_forbid_compile=False) + def decorated(): + assert xo.settings.allow_kernel_compilation is True + + with xo.settings.override(allow_kernel_compilation=False): + decorated() + assert xo.settings.allow_kernel_compilation is False diff --git a/tests/test_struct.py b/tests/test_struct.py index 76508a2..3ca8d3e 100644 --- a/tests/test_struct.py +++ b/tests/test_struct.py @@ -6,7 +6,11 @@ import numpy as np import xobjects as xo -from xobjects.test_helpers import for_all_test_contexts, requires_context +from xobjects.test_helpers import ( + allow_kernel_compilation, + for_all_test_contexts, + requires_context, +) def test_static_struct_def(): @@ -299,6 +303,7 @@ class MyStruct(xo.Struct): @requires_context("ContextCpu") +@allow_kernel_compilation def test_compile_kernels_only_if_needed(tmp_path, mocker): """Test the use case of xtrack. @@ -348,6 +353,7 @@ def myfun(self): @requires_context("ContextCpu") +@allow_kernel_compilation def test_thisclass_placeholder_on_struct(): test_context = xo.ContextCpu() diff --git a/xobjects/__init__.py b/xobjects/__init__.py index 087b1e7..6b50a6d 100644 --- a/xobjects/__init__.py +++ b/xobjects/__init__.py @@ -35,6 +35,7 @@ from .linkedarray import BypassLinked from .general import _print +from .settings import settings from .general import assert_allclose diff --git a/xobjects/context.py b/xobjects/context.py index bff5d46..bd3d65f 100644 --- a/xobjects/context.py +++ b/xobjects/context.py @@ -207,6 +207,7 @@ def __getattr__(self, attr): class XContext(ABC): minimum_alignment = 1 allow_prebuilt_kernels = False + allow_kernel_compilation = False def __init__(self): self._kernels = KernelDict() @@ -776,8 +777,8 @@ def get_test_contexts(): def get_user_context(): """ - Get the context specfied by the enviroment variable XOBJECTS_USER_CONTEXT. - If not present use ContextCpu(). + Get the context specified by the environment variable + ``XOBJECTS_USER_CONTEXT``. If it is not set, use ``ContextCpu()``. Examples: ContextPyopencl:0.0 -> ContextPyopencl(device="0.0") diff --git a/xobjects/context_cpu.py b/xobjects/context_cpu.py index e26383c..53bf29f 100644 --- a/xobjects/context_cpu.py +++ b/xobjects/context_cpu.py @@ -13,42 +13,42 @@ import weakref from .general import _print +from .settings import settings import numpy as np import scipy as sp -_forbid_compile = False _suppress_warnings = False -allow_no_prebuilt_kernel = False -def _class_allows_no_prebuilt_kernel(cls): - return ( - getattr(cls, "allow_no_prebuilt_kernel", False) - or getattr( - getattr(cls, "_DressingClass", None), - "allow_no_prebuilt_kernel", - False, - ) - or getattr( - getattr(cls, "_XoStruct", None), "allow_no_prebuilt_kernel", False - ) - ) +def _class_allows_kernel_compilation(cls): + for candidate in ( + cls, + getattr(cls, "_DressingClass", None), + getattr(cls, "_XoStruct", None), + ): + if candidate is None: + continue + if getattr(candidate, "allow_kernel_compilation", False): + return True + # Compatibility with classes from packages that have not migrated yet. + if getattr(candidate, "allow_no_prebuilt_kernel", False): + return True + return False -def allow_no_prebuilt_kernel_enabled(context=None, classes=()): +def kernel_compilation_allowed(context=None, classes=()): if classes is None: classes = () elif isinstance(classes, type): classes = (classes,) - if os.environ.get("XSUITE_ALLOW_NO_PREBUILT_KERNELS") is not None: - return True - if allow_no_prebuilt_kernel: + if (settings.allow_kernel_compilation + or settings.force_kernel_compilation): return True - if any(_class_allows_no_prebuilt_kernel(cls) for cls in classes): + if any(_class_allows_kernel_compilation(cls) for cls in classes): return True - return getattr(context, "allow_no_prebuilt_kernel", False) + return getattr(context, "allow_kernel_compilation", False) def _is_serial_cpu_context(context): @@ -58,20 +58,20 @@ def _is_serial_cpu_context(context): def require_prebuilt_kernel(context=None, classes=()): - return not allow_no_prebuilt_kernel_enabled( + return not kernel_compilation_allowed( context, classes=classes ) and _is_serial_cpu_context(context) -def no_prebuilt_kernel_jit_message(): +def kernel_compilation_help_message(): return ( "To allow just-in-time compilation instead, as in older Xsuite " - "versions, set the environment variable " - "`XSUITE_ALLOW_NO_PREBUILT_KERNELS`, set " - "`xobjects.context_cpu.allow_no_prebuilt_kernel = True`, or set " - "`context.allow_no_prebuilt_kernel = True`. Classes that require " + "versions, set `xobjects.settings.allow_kernel_compilation = True`, " + "or equivalently set the environment variable " + "`XSUITE_ALLOW_KERNEL_COMPILATION=1`, or set " + "`context.allow_kernel_compilation = True`. Classes that require " "just-in-time compilation can also define " - "`allow_no_prebuilt_kernel = True` as a class attribute. Using " + "`allow_kernel_compilation = True` as a class attribute. Using " "just-in-time compilation instead of prebuilt kernels may require " "lengthy compilation whenever a different kernel is needed." ) @@ -369,13 +369,11 @@ def build_kernels( cdefs = "\n".join(cls._gen_c_decl({}) for cls in classes) cdefs += "\n" + extra_cdef - if _forbid_compile: - raise RuntimeError("Compilation is forbidden") - - if os.environ.get("XOBJECTS_FORBID_COMPILE"): + if settings.cffi_forbid_compile: raise RuntimeError( - "Compilation is forbidden by the environment variable " - "XOBJECTS_FORBID_COMPILE" + "CFFI compilation is forbidden by " + "xobjects.settings.cffi_forbid_compile or equivalently " + "the environment variable XSUITE_CFFI_FORBID_COMPILE." ) so_file = self.compile_kernel( @@ -536,7 +534,7 @@ def compile_kernel( return Path(output_file) finally: # Clean temp files - if "XOBJECTS_KEEP_BUILD_FILES" not in os.environ: + if not settings.cffi_keep_build_files: files_to_remove = [ module_name + ".c", module_name + ".o", @@ -961,8 +959,8 @@ def __init__(self, data, axes, threads=0): direction="FFTW_BACKWARD", flags=("FFTW_MEASURE",), ) - print(f"fftw simd_aligned={self.fftw.simd_aligned}") - print(f"ifftw simd_aligned={self.fftw.simd_aligned}") + _print(f"fftw simd_aligned={self.fftw.simd_aligned}") + _print(f"ifftw simd_aligned={self.fftw.simd_aligned}") else: # I perform one fft to have numpy cache the plan _ = np.fft.ifftn(np.fft.fftn(data, axes=axes), axes=axes) diff --git a/xobjects/context_cupy.py b/xobjects/context_cupy.py index 943f6ef..4111b6c 100644 --- a/xobjects/context_cupy.py +++ b/xobjects/context_cupy.py @@ -25,18 +25,11 @@ sources_from_classes, ) from .linkedarray import BaseLinkedArray +from .settings import settings from .specialize_source import specialize_source log = logging.getLogger(__name__) -no_fast_compile = False -"""Disable NVRTC fast compile tuning when building CUDA kernels. - -When set to ``True``, ``ContextCupy`` does not pass ``--Ofast-compile=min`` -to NVRTC. The ``XO_CUDA_NO_FAST_COMPILE`` environment variable provides the -same behavior. -""" - try: import cupy import cupyx.scipy @@ -399,12 +392,12 @@ class ContextCupy(XContext): Creates a Cupy Context object, that allows performing the computations on nVidia GPUs. - The module-level flag ``xobjects.context_cupy.no_fast_compile`` controls - whether NVRTC fast compile tuning is disabled. By default it is ``False``, - so CUDA kernels built with NVRTC >= 12.9 use ``--Ofast-compile=min`` to - reduce compilation time and memory usage, at the cost of some runtime - performance. Set it to ``True`` to disable this option. The environment - variable ``XO_CUDA_NO_FAST_COMPILE`` also disables it. + ``xobjects.settings.cuda_fast_compile`` controls whether NVRTC fast compile + tuning is enabled. By default it is ``True``, so CUDA kernels built with + NVRTC >= 12.9 use ``--Ofast-compile=min`` to reduce compilation time and + memory usage, at the cost of some runtime performance. Set + ``xobjects.settings.cuda_fast_compile = False``, or equivalently the + environment variable ``XSUITE_CUDA_FAST_COMPILE=0``, to disable it. Args: default_block_size (int): CUDA thread size that is used by default @@ -439,8 +432,8 @@ def __init__( self.default_block_size = default_block_size self.default_shared_mem_size_bytes = default_shared_mem_size_bytes - if not backend: - backend = os.environ.get("XO_CUDA_BACKEND", "nvrtc") + if backend is None: + backend = settings.cuda_backend if backend not in ["nvrtc", "clang"]: raise ValueError( @@ -516,9 +509,7 @@ def build_kernels( if self.backend == "nvrtc": # NVRTC (default): add NVRTC-specific flags nvrtc_args = (*extra_compile_args,) - fast_compile = not ( - no_fast_compile or os.environ.get("XO_CUDA_NO_FAST_COMPILE") - ) + fast_compile = settings.cuda_fast_compile if nvrtc and nvrtc.getVersion() >= (12, 9): # If supported, skip prohibitively heavy optimisations (e.g. # involving cloning). This it at the expense of <20% @@ -558,7 +549,7 @@ def build_kernels( return out_kernels def _find_clang(self): - override = os.environ.get("XO_CUDA_CLANG") + override = settings.cuda_compiler if override: return override @@ -567,8 +558,11 @@ def _find_clang(self): return found raise RuntimeError( - "clang++ for the CUDA context not found. Either install clang so that 'clang++' is on PATH," - "or set the XO_CUDA_CLANG variable to the desired clang++ executable." + "clang++ for the CUDA context not found. Either install clang so " + "that 'clang++' is on PATH, or set " + "xobjects.settings.cuda_compiler, or equivalently the environment " + "variable XSUITE_CUDA_COMPILER, to the desired clang++ " + "executable." ) def _build_module_with_clang(self, source, extra_compile_args=()): diff --git a/xobjects/context_pyopencl.py b/xobjects/context_pyopencl.py index d94dd63..38cff24 100644 --- a/xobjects/context_pyopencl.py +++ b/xobjects/context_pyopencl.py @@ -21,6 +21,7 @@ ) from .linkedarray import BaseLinkedArray from .specialize_source import specialize_source +from .general import _print log = logging.getLogger(__name__) @@ -100,9 +101,9 @@ def get_devices(cls): @classmethod def print_devices(cls): for ip, platform in enumerate(cl.get_platforms()): - print(f"Platform {ip} : {platform.name}") + _print(f"Platform {ip} : {platform.name}") for id, device in enumerate(platform.get_devices()): - print(f"Device {ip}.{id}: {device.name}") + _print(f"Device {ip}.{id}: {device.name}") def __init__( self, device=None, patch_pyopencl_array=True, minimum_alignment=None diff --git a/xobjects/general.py b/xobjects/general.py index 5ac08b2..a8d57cc 100644 --- a/xobjects/general.py +++ b/xobjects/general.py @@ -5,13 +5,21 @@ from numpy.testing import assert_allclose as np_assert_allclose import numpy as np +from .settings import settings + class Print: - suppress = False + """Configurable wrapper around :func:`print` used by Xsuite. + + The behavior is controlled by ``xobjects.settings.print_mode``, or + equivalently the environment variable ``XSUITE_PRINT_MODE``. + """ def __call__(self, *args, **kwargs): - if not self.suppress: - print(*args, **kwargs) + if settings.print_mode == 'suppress': + return + + print(*args, **kwargs) _print = Print() diff --git a/xobjects/settings.py b/xobjects/settings.py new file mode 100644 index 0000000..e781327 --- /dev/null +++ b/xobjects/settings.py @@ -0,0 +1,308 @@ +import os +from contextlib import contextmanager + + +class Settings: + """Process-wide settings shared by the Xsuite packages. + + Available settings and their equivalent environment variables are: + + ``print_mode`` — ``XSUITE_PRINT_MODE`` + Values: ``print`` (default), ``suppress`` + Controls informational output produced by Xsuite. + + ``progress_indicator`` — ``XSUITE_PROGRESS_INDICATOR`` + Values: ``tqdm`` (default), ``text``, ``suppress`` + Selects how progress is displayed when printing is enabled. + + ``allow_kernel_compilation`` — ``XSUITE_ALLOW_KERNEL_COMPILATION`` + Values: ``False`` (default), ``True`` + Allows compilation when no compatible prebuilt kernel is available. + + ``force_kernel_compilation`` — ``XSUITE_FORCE_KERNEL_COMPILATION`` + Values: ``False`` (default), ``True`` + Always compiles kernels and skips prebuilt-kernel lookup. Intended + primarily for Xsuite development and debugging. + + ``show_kernel_diagnostics`` — ``XSUITE_SHOW_KERNEL_DIAGNOSTICS`` + Values: ``False`` (default), ``True`` + During prebuilt-kernel selection, reports why each candidate kernel is + selected or rejected. + + ``cffi_forbid_compile`` — ``XSUITE_CFFI_FORBID_COMPILE`` + Values: ``False`` (default), ``True`` + Prevents CFFI from compiling CPU kernels. Used primarily for debugging + and testing. + + ``cffi_keep_build_files`` — ``XSUITE_CFFI_KEEP_BUILD_FILES`` + Values: ``False`` (default), ``True`` + Keeps the intermediate files generated by CFFI. Used primarily for + debugging and testing. + + ``cuda_backend`` — ``XSUITE_CUDA_BACKEND`` + Values: ``nvrtc`` (default), ``clang`` + Selects the compiler backend used for CUDA kernels. + + ``cuda_fast_compile`` — ``XSUITE_CUDA_FAST_COMPILE`` + Values: ``True`` (default), ``False`` + Enables faster NVRTC compilation when supported. + + ``cuda_compiler`` — ``XSUITE_CUDA_COMPILER`` + Values: ``None`` (default), executable path + Selects the clang++ executable used by the CUDA clang backend. + + Notes + ----- + The same object is exposed as ``xobjects.settings`` and + ``xtrack.settings``. Environment variables provide its initial values when + Xobjects is imported. Later Python assignments take precedence and can be + applied temporarily with :meth:`override`. + + Boolean environment variables accept ``1``/``0``, ``true``/``false``, + ``yes``/``no``, and ``on``/``off``, case-insensitively. + + ``force_kernel_compilation=True`` skips prebuilt-kernel lookup and takes + precedence over ``allow_kernel_compilation``. When forcing is disabled, + ``allow_kernel_compilation=True`` permits compilation only if a compatible + prebuilt kernel is unavailable. Compilation can also be enabled for one + context or element class by setting its ``allow_kernel_compilation`` + attribute to ``True``. The legacy class attribute + ``allow_no_prebuilt_kernel`` is temporarily recognized for compatibility + with packages that have not migrated yet. + + Examples + -------- + Allow Xsuite to compile a kernel when no compatible prebuilt kernel is + available: + + >>> import xtrack as xt + >>> xt.settings.allow_kernel_compilation = True + + Equivalently, set ``XSUITE_ALLOW_KERNEL_COMPILATION=1`` before starting + Python. + """ + + def __init__(self): + object.__setattr__(self, '_definitions', {}) + object.__setattr__(self, '_values', {}) + + def _register( + self, + name, + *, + default, + environment_variable=None, + choices=None, + environment_parser=None, + value_type=None, + ): + if name in self._definitions: + raise ValueError(f'Setting {name!r} is already registered.') + + definition = { + 'environment_variable': environment_variable, + 'choices': choices, + 'value_type': value_type, + } + self._definitions[name] = definition + + value = default + if (environment_variable is not None + and environment_variable in os.environ): + environment_value = os.environ[environment_variable] + try: + value = (environment_parser(environment_value) + if environment_parser else environment_value) + except (TypeError, ValueError) as err: + raise ValueError( + f'Invalid value for ' + f'{self._setting_description(name, definition)}: ' + f'{err}') from err + self._set(name, value) + + def _set(self, name, value): + try: + definition = self._definitions[name] + except KeyError as err: + raise AttributeError(f'Unknown Xsuite setting {name!r}.') from err + + choices = definition['choices'] + value_type = definition['value_type'] + setting_description = self._setting_description(name, definition) + if value_type is not None and not isinstance(value, value_type): + raise TypeError( + f'Invalid value {value!r} for {setting_description}; ' + f'expected {self._type_name(value_type)}.') + if choices is not None and value not in choices: + expected = ', '.join(repr(choice) for choice in choices) + raise ValueError( + f'Invalid value {value!r} for {setting_description}; ' + f'expected one of {expected}.') + + self._values[name] = value + + def __getattr__(self, name): + try: + return self._values[name] + except KeyError as err: + raise AttributeError(f'Unknown Xsuite setting {name!r}.') from err + + def __setattr__(self, name, value): + self._set(name, value) + + @contextmanager + def override(self, **kwargs): + """Temporarily override settings and restore them on exit.""" + previous = {} + for name, value in kwargs.items(): + if name not in self._definitions: + raise AttributeError(f'Unknown Xsuite setting {name!r}.') + previous[name] = getattr(self, name) + + # Validate every value before changing any setting. + for name, value in kwargs.items(): + definition = self._definitions[name] + setting_description = self._setting_description(name, definition) + value_type = definition['value_type'] + if value_type is not None and not isinstance(value, value_type): + raise TypeError( + f'Invalid value {value!r} for {setting_description}; ' + f'expected {self._type_name(value_type)}.') + choices = definition['choices'] + if choices is not None and value not in choices: + expected = ', '.join(repr(choice) for choice in choices) + raise ValueError( + f'Invalid value {value!r} for {setting_description}; ' + f'expected one of {expected}.') + + try: + for name, value in kwargs.items(): + self._set(name, value) + yield self + finally: + for name, value in previous.items(): + self._set(name, value) + + def __repr__(self): + values = ', '.join( + f'{name}={getattr(self, name)!r}' for name in self._values) + return f'Settings({values})' + + def __dir__(self): + return sorted(set(super().__dir__()) | set(self._definitions)) + + @staticmethod + def _type_name(value_type): + if isinstance(value_type, tuple): + return ' or '.join(tt.__name__ for tt in value_type) + return value_type.__name__ + + @staticmethod + def _setting_description(name, definition): + description = f'Python setting xobjects.settings.{name}' + environment_variable = definition['environment_variable'] + if environment_variable is not None: + description += ( + f' or equivalently the environment variable ' + f'{environment_variable}' + ) + return description + + +def _parse_boolean(value): + normalized = value.strip().lower() + if normalized in ('1', 'true', 'yes', 'on'): + return True + if normalized in ('0', 'false', 'no', 'off'): + return False + raise ValueError( + f'Invalid boolean value {value!r}; expected one of 1, 0, true, ' + 'false, yes, no, on, or off.') + + +def _parse_choice(value): + return value.strip().lower() + + +def _parse_optional_string(value): + value = value.strip() + return value if value else None + + +settings = Settings() +settings._register( + 'print_mode', + default='print', + environment_variable='XSUITE_PRINT_MODE', + choices=('print', 'suppress'), + environment_parser=_parse_choice, +) +settings._register( + 'progress_indicator', + default='tqdm', + environment_variable='XSUITE_PROGRESS_INDICATOR', + choices=('tqdm', 'text', 'suppress'), + environment_parser=_parse_choice, +) +settings._register( + 'allow_kernel_compilation', + default=False, + environment_variable='XSUITE_ALLOW_KERNEL_COMPILATION', + choices=(False, True), + environment_parser=_parse_boolean, + value_type=bool, +) +settings._register( + 'force_kernel_compilation', + default=False, + environment_variable='XSUITE_FORCE_KERNEL_COMPILATION', + choices=(False, True), + environment_parser=_parse_boolean, + value_type=bool, +) +settings._register( + 'show_kernel_diagnostics', + default=False, + environment_variable='XSUITE_SHOW_KERNEL_DIAGNOSTICS', + choices=(False, True), + environment_parser=_parse_boolean, + value_type=bool, +) +settings._register( + 'cffi_forbid_compile', + default=False, + environment_variable='XSUITE_CFFI_FORBID_COMPILE', + choices=(False, True), + environment_parser=_parse_boolean, + value_type=bool, +) +settings._register( + 'cffi_keep_build_files', + default=False, + environment_variable='XSUITE_CFFI_KEEP_BUILD_FILES', + choices=(False, True), + environment_parser=_parse_boolean, + value_type=bool, +) +settings._register( + 'cuda_backend', + default='nvrtc', + environment_variable='XSUITE_CUDA_BACKEND', + choices=('nvrtc', 'clang'), + environment_parser=_parse_choice, +) +settings._register( + 'cuda_fast_compile', + default=True, + environment_variable='XSUITE_CUDA_FAST_COMPILE', + choices=(False, True), + environment_parser=_parse_boolean, + value_type=bool, +) +settings._register( + 'cuda_compiler', + default=None, + environment_variable='XSUITE_CUDA_COMPILER', + environment_parser=_parse_optional_string, + value_type=(str, type(None)), +) diff --git a/xobjects/struct.py b/xobjects/struct.py index f576856..be36f69 100644 --- a/xobjects/struct.py +++ b/xobjects/struct.py @@ -59,7 +59,7 @@ default_conf, ) -from .general import Print +from .settings import settings from .scalar import Int64 from .array import Array from .context import Source, Arg, Kernel @@ -518,14 +518,13 @@ def compile_class_kernels( extra_classes=(), extra_compile_args=(), ): - if context.allow_prebuilt_kernels: - _print_state = Print.suppress - Print.suppress = True - try: + if (context.allow_prebuilt_kernels + and not settings.force_kernel_compilation): + with settings.override(print_mode="suppress"): try: from xsuite import ( get_suitable_kernel, - XSK_PREBUILT_KERNELS_LOCATION, + PREBUILT_KERNELS_LOCATION, ) kernel_info = get_suitable_kernel( @@ -546,14 +545,12 @@ def compile_class_kernels( "Xsuite is required to load prebuilt kernels but could " "not be imported. Please install it with " f"`pip install xsuite`. " - f"{context_cpu.no_prebuilt_kernel_jit_message()}" + f"{context_cpu.kernel_compilation_help_message()}" ) from err - finally: - Print.suppress = _print_state if kernel_info: kernels = context.kernels_from_file( module_name=kernel_info["module_name"], - containing_dir=XSK_PREBUILT_KERNELS_LOCATION, + containing_dir=PREBUILT_KERNELS_LOCATION, kernel_descriptions=cls._kernels, ) context.kernels.update(kernels) diff --git a/xobjects/test_helpers.py b/xobjects/test_helpers.py index fce38ab..0828dff 100644 --- a/xobjects/test_helpers.py +++ b/xobjects/test_helpers.py @@ -5,11 +5,11 @@ from functools import wraps from typing import Callable, Iterable, Union -import os import pytest from .context import get_context_from_string, get_test_contexts +from .settings import settings def _for_all_test_contexts_excluding( @@ -108,13 +108,13 @@ def wrapper(*args, **kwargs): return decorator -def allow_no_prebuilt_kernels( +def allow_kernel_compilation( test_function=None, *, skip_when_forbid_compile=True ): """Allow JIT compilation for tests that intentionally compile kernels. By default, the wrapped test is skipped when compilation is forbidden by - ``XOBJECTS_FORBID_COMPILE``. Use ``skip_when_forbid_compile=False`` when + ``XSUITE_CFFI_FORBID_COMPILE``. Use ``skip_when_forbid_compile=False`` when the test has more specific ``skip_if_forbid_compile()`` guards inside the test. """ @@ -124,15 +124,8 @@ def decorator(test_function): def wrapper(*args, **kwargs): if skip_when_forbid_compile: skip_if_forbid_compile() - old_value = os.environ.get("XSUITE_ALLOW_NO_PREBUILT_KERNELS") - os.environ["XSUITE_ALLOW_NO_PREBUILT_KERNELS"] = "1" - try: + with settings.override(allow_kernel_compilation=True): return test_function(*args, **kwargs) - finally: - if old_value is None: - del os.environ["XSUITE_ALLOW_NO_PREBUILT_KERNELS"] - else: - os.environ["XSUITE_ALLOW_NO_PREBUILT_KERNELS"] = old_value return wrapper @@ -142,8 +135,9 @@ def wrapper(*args, **kwargs): def skip_if_forbid_compile(): - if os.environ.get("XOBJECTS_FORBID_COMPILE"): + if settings.cffi_forbid_compile: pytest.skip( - "Compilation is forbidden by the environment variable " - "XOBJECTS_FORBID_COMPILE" + "CFFI compilation is forbidden by " + "xobjects.settings.cffi_forbid_compile or equivalently the " + "environment variable XSUITE_CFFI_FORBID_COMPILE." )