diff --git a/pytensor/link/c/cmodule.py b/pytensor/link/c/cmodule.py index acfc32fe46..3d6cd55a4c 100644 --- a/pytensor/link/c/cmodule.py +++ b/pytensor/link/c/cmodule.py @@ -10,6 +10,7 @@ import pickle import platform import re +import shlex import shutil import stat import subprocess @@ -19,7 +20,7 @@ import textwrap import time import warnings -from collections.abc import Callable +from collections.abc import Callable, Collection, Sequence from contextlib import AbstractContextManager, nullcontext from io import BytesIO, StringIO from pathlib import Path @@ -2104,49 +2105,48 @@ def compile_args(march_flags=True): ) detect_march = False - if detect_march: - GCC_compiler.march_flags = [] + def get_lines(cmd: list[str], parse: bool = True) -> list[str] | None: + p = subprocess_Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + stdin=subprocess.PIPE, + ) + # For mingw64 with GCC >= 4.7, passing os.devnull + # as stdin (which is the default) results in the process + # waiting forever without returning. For that reason, + # we use a pipe, and use the empty string as input. + (stdout, stderr) = p.communicate(input=b"") + if p.returncode != 0: + return None - def get_lines(cmd, parse=True): - p = subprocess_Popen( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - stdin=subprocess.PIPE, - shell=True, - ) - # For mingw64 with GCC >= 4.7, passing os.devnull - # as stdin (which is the default) results in the process - # waiting forever without returning. For that reason, - # we use a pipe, and use the empty string as input. - (stdout, stderr) = p.communicate(input=b"") - if p.returncode != 0: - return None - - lines = BytesIO(stdout + stderr).readlines() - lines = (l.decode() for l in lines) - if parse: - selected_lines = [] - for line in lines: - if ( - "COLLECT_GCC_OPTIONS=" in line - or "CFLAGS=" in line - or "CXXFLAGS=" in line - or "-march=native" in line - ): - continue - selected_lines.extend( - line.strip() - for reg in ("-march=", "-mtune=", "-target-cpu", "-mabi=") - if reg in line - ) - lines = list(set(selected_lines)) # to remove duplicate + lines_bytes = BytesIO(stdout + stderr).readlines() + lines = [l.decode() for l in lines_bytes] + if parse: + selected_lines: list[str] = [] + for line in lines: + if ( + "COLLECT_GCC_OPTIONS=" in line + or "CFLAGS=" in line + or "CXXFLAGS=" in line + or "-march=native" in line + ): + continue + selected_lines.extend( + line.strip() + for reg in ("-march=", "-mtune=", "-target-cpu", "-mabi=") + if reg in line + ) + lines = list(set(selected_lines)) # to remove duplicate - return lines + return lines + + if detect_march: + GCC_compiler.march_flags = [] # The '-' at the end is needed. Otherwise, g++ do not output # enough information. - native_lines = get_lines(f"{config.cxx} -march=native -E -v -") + native_lines = get_lines([config.cxx, "-march=native", "-E", "-v", "-"]) if native_lines is None: _logger.info( "Call to 'g++ -march=native' failed, not setting -march flag" @@ -2161,7 +2161,7 @@ def get_lines(cmd, parse=True): # That means we did not select the right lines, so # we have to report all the lines instead reported_lines = get_lines( - f"{config.cxx} -march=native -E -v -", parse=False + [config.cxx, "-march=native", "-E", "-v", "-"], parse=False ) else: reported_lines = native_lines @@ -2174,10 +2174,12 @@ def get_lines(cmd, parse=True): f" problem:\n {reported_lines}" ) else: - default_lines = get_lines(f"{config.cxx} -E -v -") + default_lines = get_lines([config.cxx, "-E", "-v", "-"]) _logger.info(f"g++ default lines: {default_lines}") if len(default_lines) < 1: - reported_lines = get_lines(f"{config.cxx} -E -v -", parse=False) + reported_lines = get_lines( + [config.cxx, "-E", "-v", "-"], parse=False + ) warnings.warn( "PyTensor was not able to find the " "default g++ parameters. This is needed to tune " @@ -2610,7 +2612,7 @@ def compile_str( cmd.append(f"{path_wrapper}{cppfilename}{path_wrapper}") cmd.extend(GCC_compiler.linking_patch(lib_dirs, libs)) # print >> sys.stderr, 'COMPILING W CMD', cmd - _logger.debug(f"Running cmd: {' '.join(cmd)}") + _logger.debug(f"Running cmd: {shlex.join(cmd)}") def print_command_line_error(): # Print command line when a problem occurred. @@ -2618,7 +2620,7 @@ def print_command_line_error(): ("Problem occurred during compilation with the command line below:"), file=sys.stderr, ) - print(" ".join(cmd), file=sys.stderr) + print(shlex.join(cmd), file=sys.stderr) try: p_out = output_subprocess_Popen(cmd) @@ -2734,6 +2736,96 @@ def check_mkl_openmp(): ) +def _check_required_file( + paths: Collection[Path], + required_regexs: Collection[str | re.Pattern[str]], +) -> list[tuple[str, str]]: + """Select path parents for each required pattern.""" + libs: list[tuple[str, str]] = [] + for req in required_regexs: + found = False + for path in paths: + m = re.search(req, path.name) + if m: + libs.append((str(path.parent), m.string[slice(*m.span())])) + found = True + break + if not found: + _logger.debug("Required file '%s' not found", req) + raise RuntimeError(f"Required file {req} not found") + return libs + + +def _get_cxx_library_dirs() -> list[str]: + """Query C++ search dirs and return those the existing ones.""" + cmd = [config.cxx, "-print-search-dirs"] + p = subprocess_Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + stdin=subprocess.PIPE, + ) + (stdout, stderr) = p.communicate(input=b"") + if p.returncode != 0: + warnings.warn( + "Pytensor cxx failed to communicate its search dirs. As a consequence, " + "it might not be possible to automatically determine the blas link flags to use.\n" + f"Command that was run: {config.cxx} -print-search-dirs\n" + f"Output printed to stderr: {stderr.decode(sys.stderr.encoding)}" + ) + return [] + + maybe_lib_dirs = [ + [Path(p).resolve() for p in line[len("libraries: =") :].split(":")] + for line in stdout.decode(sys.getdefaultencoding()).splitlines() + if line.startswith("libraries: =") + ] + if not maybe_lib_dirs: + return [] + return [str(d) for d in maybe_lib_dirs[0] if d.exists() and d.is_dir()] + + +def _check_libs( + all_libs: Collection[Path], + required_libs: Collection[str | re.Pattern], + extra_compile_flags: Sequence[str] = (), + cxx_library_dirs: Sequence[str] = (), +) -> str: + """Assembly library paths and try BLAS flags, returning the flags on success.""" + found_libs = _check_required_file( + all_libs, + required_libs, + ) + path_quote = '"' if sys.platform == "win32" else "" + libdir_ldflags = list( + dict.fromkeys( + [ + f"-L{path_quote}{lib_path}{path_quote}" + for lib_path, _ in found_libs + if lib_path not in cxx_library_dirs + ] + ) + ) + + flags = ( + libdir_ldflags + + [f"-l{lib_name}" for _, lib_name in found_libs] + + list(extra_compile_flags) + ) + res = try_blas_flag(flags) + if not res: + _logger.debug("Supplied flags '%s' failed to compile", res) + raise RuntimeError(f"Supplied flags {flags} failed to compile") + + if any("mkl" in flag for flag in flags): + try: + check_mkl_openmp() + except Exception as e: + _logger.debug(e) + _logger.debug("The following blas flags will be used: '%s'", res) + return res + + def default_blas_ldflags() -> str: """Look for an available BLAS implementation in the system. @@ -2761,89 +2853,6 @@ def default_blas_ldflags() -> str: """ - def check_required_file(paths, required_regexs): - libs = [] - for req in required_regexs: - found = False - for path in paths: - m = re.search(req, path.name) - if m: - libs.append((str(path.parent), m.string[slice(*m.span())])) - found = True - break - if not found: - _logger.debug("Required file '%s' not found", req) - raise RuntimeError(f"Required file {req} not found") - return libs - - def get_cxx_library_dirs(): - cmd = f"{config.cxx} -print-search-dirs" - p = subprocess_Popen( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - stdin=subprocess.PIPE, - shell=True, - ) - (stdout, stderr) = p.communicate(input=b"") - if p.returncode != 0: - warnings.warn( - "Pytensor cxx failed to communicate its search dirs. As a consequence, " - "it might not be possible to automatically determine the blas link flags to use.\n" - f"Command that was run: {config.cxx} -print-search-dirs\n" - f"Output printed to stderr: {stderr.decode(sys.stderr.encoding)}" - ) - return [] - - maybe_lib_dirs = [ - [Path(p).resolve() for p in line[len("libraries: =") :].split(":")] - for line in stdout.decode(sys.getdefaultencoding()).splitlines() - if line.startswith("libraries: =") - ] - if len(maybe_lib_dirs) > 0: - maybe_lib_dirs = maybe_lib_dirs[0] - return [str(d) for d in maybe_lib_dirs if d.exists() and d.is_dir()] - - def check_libs( - all_libs, required_libs, extra_compile_flags=None, cxx_library_dirs=None - ) -> str: - if cxx_library_dirs is None: - cxx_library_dirs = [] - if extra_compile_flags is None: - extra_compile_flags = [] - found_libs = check_required_file( - all_libs, - required_libs, - ) - path_quote = '"' if sys.platform == "win32" else "" - libdir_ldflags = list( - dict.fromkeys( - [ - f"-L{path_quote}{lib_path}{path_quote}" - for lib_path, _ in found_libs - if lib_path not in cxx_library_dirs - ] - ) - ) - - flags = ( - libdir_ldflags - + [f"-l{lib_name}" for _, lib_name in found_libs] - + extra_compile_flags - ) - res = try_blas_flag(flags) - if res: - if any("mkl" in flag for flag in flags): - try: - check_mkl_openmp() - except Exception as e: - _logger.debug(e) - _logger.debug("The following blas flags will be used: '%s'", res) - return res - else: - _logger.debug("Supplied flags '%s' failed to compile", res) - raise RuntimeError(f"Supplied flags {flags} failed to compile") - # If no compiler is available we default to empty ldflags if not config.cxx: return "" @@ -2853,7 +2862,7 @@ def check_libs( else: rpath = None - cxx_library_dirs = get_cxx_library_dirs() + cxx_library_dirs = _get_cxx_library_dirs() searched_library_dirs = cxx_library_dirs + _std_lib_dirs if sys.platform == "win32": # Conda on Windows saves MKL libraries under CONDA_PREFIX\Library\bin @@ -2883,7 +2892,7 @@ def check_libs( try: # 1. Try to use MKL with INTEL OpenMP threading _logger.debug("Checking MKL flags with intel threading") - return check_libs( + return _check_libs( all_libs, required_libs=[ "mkl_core", @@ -2900,7 +2909,7 @@ def check_libs( try: # 2. Try to use MKL with GNU OpenMP threading _logger.debug("Checking MKL flags with GNU OpenMP threading") - return check_libs( + return _check_libs( all_libs, required_libs=["mkl_core", "mkl_rt", "mkl_gnu_thread", "gomp", "pthread"], extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [], @@ -2923,7 +2932,7 @@ def check_libs( try: _logger.debug("Checking Lapack + blas") # 4. Try to use LAPACK + BLAS - return check_libs( + return _check_libs( all_libs, required_libs=["lapack", "blas", "cblas", "m"], extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [], @@ -2934,7 +2943,7 @@ def check_libs( try: # 5. Try to use BLAS alone _logger.debug("Checking blas alone") - return check_libs( + return _check_libs( all_libs, required_libs=["blas", "cblas"], extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [], @@ -2945,7 +2954,7 @@ def check_libs( try: # 6. Try to use openblas _logger.debug("Checking openblas") - return check_libs( + return _check_libs( all_libs, required_libs=["openblas", "gfortran", "gomp", "m"], extra_compile_flags=["-fopenmp", f"-Wl,-rpath,{rpath}"] diff --git a/pytensor/utils.py b/pytensor/utils.py index 01eb06f2e2..c81fb74f56 100644 --- a/pytensor/utils.py +++ b/pytensor/utils.py @@ -123,7 +123,7 @@ def maybe_add_to_os_environ_pathlist(var: str, newpath: Path | str) -> None: pass -def subprocess_Popen(command: str | list[str], **params): +def subprocess_Popen(command: list[str], **params) -> subprocess.Popen: """ Utility function to work around windows behavior that open windows. @@ -137,37 +137,17 @@ def subprocess_Popen(command: str | list[str], **params): except AttributeError: startupinfo.dwFlags |= subprocess._subprocess.STARTF_USESHOWWINDOW # type: ignore[attr-defined] - # Anaconda for Windows does not always provide .exe files - # in the PATH, they also have .bat files that call the corresponding - # executable. For instance, "g++.bat" is in the PATH, not "g++.exe" - # Unless "shell=True", "g++.bat" is not executed when trying to - # execute "g++" without extensions. - # (Executing "g++.bat" explicitly would also work.) - params["shell"] = True # "If shell is True, it is recommended to pass args as a string rather than as a sequence." (cite taken from https://docs.python.org/2/library/subprocess.html#frequently-used-arguments) # In case when command arguments have spaces, passing a command as a list will result in incorrect arguments break down, and consequently # in "The filename, directory name, or volume label syntax is incorrect" error message. # Passing the command as a single string solves this problem. if isinstance(command, list): - command = " ".join(command) + command = " ".join(command) # type: ignore[assignment] - # Using the dummy file descriptors below is a workaround for a - # crash experienced in an unusual Python 2.4.4 Windows environment - # with the default None values. - stdin = None - if "stdin" not in params: - stdin = Path(os.devnull).open() - params["stdin"] = stdin.fileno() - - try: - proc = subprocess.Popen(command, startupinfo=startupinfo, **params) - finally: - if stdin is not None: - stdin.close() - return proc + return subprocess.Popen(command, startupinfo=startupinfo, **params) -def call_subprocess_Popen(command, **params): +def call_subprocess_Popen(command: list[str], **params) -> int: """ Calls subprocess_Popen and discards the output, returning only the exit code. @@ -185,13 +165,17 @@ def call_subprocess_Popen(command, **params): return returncode -def output_subprocess_Popen(command, **params): +def output_subprocess_Popen(command: list[str], **params) -> tuple[bytes, bytes, int]: """ Calls subprocess_Popen, returning the output, error and exit code in a tuple. """ if "stdout" in params or "stderr" in params: raise TypeError("don't use stderr or stdout with output_subprocess_Popen") + if "encoding" in params: + raise TypeError( + "adjust the output_subprocess_Popen type annotation to support str" + ) params["stdout"] = subprocess.PIPE params["stderr"] = subprocess.PIPE p = subprocess_Popen(command, **params) diff --git a/tests/compile/function/test_function.py b/tests/compile/function/test_function.py index d1f94dd689..b4748e78c5 100644 --- a/tests/compile/function/test_function.py +++ b/tests/compile/function/test_function.py @@ -1,5 +1,4 @@ import pickle -import re import shutil import tempfile from pathlib import Path @@ -50,8 +49,7 @@ def test_function_name(): x = vector("x") func = function([x], x + 1.0) - regex = re.compile(f".*{__file__}c?") - assert regex.match(func.name) is not None + assert __file__ in func.name def test_trust_input(): diff --git a/tests/link/c/test_cmodule.py b/tests/link/c/test_cmodule.py index 46533fef35..212a2d8181 100644 --- a/tests/link/c/test_cmodule.py +++ b/tests/link/c/test_cmodule.py @@ -128,7 +128,7 @@ def test_cache_versioning(): z = my_add(x) z_v = my_add_ver(x) - with tempfile.TemporaryDirectory() as dir_name: + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as dir_name: cache = ModuleCache(dir_name) lnk = CLinker().accept(FunctionGraph(outputs=[z])) @@ -228,23 +228,19 @@ def cxx_search_dirs(blas_libs, mock_system): ) -@pytest.fixture( - scope="function", params=[True, False], ids=["Working_CXX", "Broken_CXX"] +@pytest.mark.parametrize( + "working_cxx", [True, False], ids=["Working_CXX", "Broken_CXX"] ) -def cxx_search_dirs_status(request): - return request.param - - @patch("pytensor.link.c.cmodule.std_lib_dirs", return_value=[]) @patch("pytensor.link.c.cmodule.check_mkl_openmp", return_value=None) def test_default_blas_ldflags( - mock_std_lib_dirs, mock_check_mkl_openmp, cxx_search_dirs, cxx_search_dirs_status + mock_std_lib_dirs, mock_check_mkl_openmp, cxx_search_dirs, working_cxx ): cxx_search_dirs, expected_blas_ldflags, enabled_accelerate_framework = ( cxx_search_dirs ) mock_process = MagicMock() - if cxx_search_dirs_status: + if working_cxx: error_message = "" mock_process.communicate = lambda *args, **kwargs: (cxx_search_dirs, b"") mock_process.returncode = 0 @@ -273,7 +269,7 @@ def wrapped(test_code, tmp_prefix, flags, try_run, output): "try_compile_tmp", new_callable=patched_compile_tmp, ): - if cxx_search_dirs_status: + if working_cxx: assert set(default_blas_ldflags().split(" ")) == set( expected_blas_ldflags.split(" ") ) diff --git a/tests/link/c/test_op.py b/tests/link/c/test_op.py index 5ddf6443a4..f25cadb7e8 100644 --- a/tests/link/c/test_op.py +++ b/tests/link/c/test_op.py @@ -1,7 +1,7 @@ import os +import string import subprocess import sys -import tempfile from pathlib import Path import numpy as np @@ -37,7 +37,7 @@ class QuadraticCOpFunc(ExternalCOp): def __init__(self, a, b, c): super().__init__( - "{test_dir}/c_code/test_quadratic_function.c", "APPLY_SPECIFIC(compute_quadratic)" + "{str(test_dir).replace(os.sep, "/")}/c_code/test_quadratic_function.c", "APPLY_SPECIFIC(compute_quadratic)" ) self.a = a self.b = b @@ -215,9 +215,10 @@ def get_hash(modname, seed=None): def test_ExternalCOp_c_code_cache_version(): """Make sure the C cache versions produced by `ExternalCOp` don't depend on `hash` seeding.""" - with tempfile.NamedTemporaryFile(dir=".", suffix=".py") as tmp: - tmp.write(externalcop_test_code.encode()) - tmp.seek(0) + tmp = Path() / ("".join(np.random.choice(list(string.ascii_letters), 8)) + ".py") + tmp.write_bytes(externalcop_test_code.encode()) + + try: modname = tmp.name out_1, err1, returncode1 = get_hash(modname, seed=428) out_2, err2, returncode2 = get_hash(modname, seed=3849) @@ -225,9 +226,11 @@ def test_ExternalCOp_c_code_cache_version(): assert returncode2 == 0 assert err1 == err2 - hash_1, msg, _ = out_1.decode().split("\n") + hash_1, msg, _ = out_1.decode().split(os.linesep) assert msg == "__success__" - hash_2, msg, _ = out_2.decode().split("\n") + hash_2, msg, _ = out_2.decode().split(os.linesep) assert msg == "__success__" assert hash_1 == hash_2 + finally: + tmp.unlink()