|
| 1 | +"""Invoke a script under a sibling conda/mamba env, isolated from ARC's env. |
| 2 | +
|
| 3 | +ARC runs inside ``arc_env``. Several adapters (AutoTST, GCN, TorchANI) |
| 4 | +shell out to scripts that live in their *own* envs (``tst_env``, |
| 5 | +``ts_gcn``, ``tani_env``). Running the target env's ``python`` |
| 6 | +binary directly leaves ARC's exported activation vars (``BABEL_LIBDIR``, |
| 7 | +``LD_LIBRARY_PATH``, ``CONDA_PREFIX``, ...) bound to ``arc_env``'s |
| 8 | +paths in the child, which causes ABI-mismatch crashes when shared |
| 9 | +libraries in the child resolve plugins against the wrong env's tree. |
| 10 | +
|
| 11 | +Routing through a launcher's ``run`` subcommand makes the launcher |
| 12 | +deactivate the caller env and re-activate the target, so the target |
| 13 | +env's own ``activate.d`` hooks fire and bind those vars to its paths. |
| 14 | +
|
| 15 | +Three launchers are supported, in preference order: |
| 16 | +
|
| 17 | +1. ``conda`` — needs ``--no-capture-output`` to avoid buffering child |
| 18 | + stdio. |
| 19 | +2. ``mamba`` — same parser as conda for ``run``; also needs |
| 20 | + ``--no-capture-output``. |
| 21 | +3. ``micromamba`` — independent C++ reimplementation; streams stdio by |
| 22 | + default and **rejects** ``--no-capture-output``, so the flag must be |
| 23 | + omitted. |
| 24 | +
|
| 25 | +Buffering matters: without the right flag, conda/mamba hold the child's |
| 26 | +stdout until exit, hiding tracebacks and progress. |
| 27 | +
|
| 28 | +The launcher is detected at call time, with the active one (per |
| 29 | +``CONDA_EXE`` / ``MAMBA_EXE``) preferred when available. |
| 30 | +""" |
| 31 | + |
| 32 | +import os |
| 33 | +import shutil |
| 34 | +import subprocess |
| 35 | +from pathlib import Path |
| 36 | + |
| 37 | +from arc.common import get_logger |
| 38 | + |
| 39 | +logger = get_logger() |
| 40 | + |
| 41 | + |
| 42 | +def env_prefix_from_python(python_executable: str) -> str: |
| 43 | + """Derive the env prefix from an interpreter path. |
| 44 | +
|
| 45 | + ARC's settings expose target Python interpreters as full paths |
| 46 | + (``AUTOTST_PYTHON``, ``TS_GCN_PYTHON``, ``TANI_PYTHON``). The env |
| 47 | + prefix passed to ``<launcher> run -p <prefix>`` is the directory two |
| 48 | + levels above the binary (``<prefix>/bin/python``). |
| 49 | +
|
| 50 | + Using a prefix path rather than ``-n <name>`` avoids assuming the |
| 51 | + env lives under a literal ``envs/`` segment — ``CONDA_ENVS_PATH`` |
| 52 | + and bare-prefix mamba/micromamba layouts (e.g. |
| 53 | + ``/scratch/conda_envs/<env>/bin/python``) are both fine. |
| 54 | +
|
| 55 | + Validation is lexical, NOT through ``Path.resolve()``: in real |
| 56 | + conda/mamba/micromamba envs ``<prefix>/bin/python`` is a symlink to |
| 57 | + ``python3.X``, so resolving first would replace the basename with |
| 58 | + ``python3.12`` (or similar) and trip the name check. The launcher |
| 59 | + follows its own interpreter, so all we need here is the prefix |
| 60 | + string the caller already gave us. |
| 61 | + """ |
| 62 | + path = Path(python_executable) |
| 63 | + if path.name != "python" or path.parent.name != "bin": |
| 64 | + raise ValueError( |
| 65 | + f"Cannot derive an env prefix from {python_executable!r}; " |
| 66 | + "expected a path of the form '<prefix>/bin/python'." |
| 67 | + ) |
| 68 | + return str(path.parent.parent) |
| 69 | + |
| 70 | + |
| 71 | +def _run_flags_for(launcher_path: str) -> list[str]: |
| 72 | + """Return the per-launcher flags needed for ``run`` to stream stdio. |
| 73 | +
|
| 74 | + Decided by the launcher's basename rather than which env var pointed |
| 75 | + us at it, so symlinks and odd ``MAMBA_EXE``-points-at-micromamba |
| 76 | + setups still get the right flag. |
| 77 | + """ |
| 78 | + name = Path(launcher_path).name |
| 79 | + if name == "micromamba": |
| 80 | + return [] |
| 81 | + return ["--no-capture-output"] |
| 82 | + |
| 83 | + |
| 84 | +def _detect_launcher() -> tuple[str, list[str]]: |
| 85 | + """Return ``(launcher_path, extra_run_flags)``. |
| 86 | +
|
| 87 | + Preference: whichever launcher is active in the current shell |
| 88 | + (``CONDA_EXE`` / ``MAMBA_EXE``), then conda → mamba → micromamba on |
| 89 | + PATH. |
| 90 | + """ |
| 91 | + for env_var in ("CONDA_EXE", "MAMBA_EXE"): |
| 92 | + path = os.environ.get(env_var) |
| 93 | + if path and os.path.isfile(path): |
| 94 | + return path, _run_flags_for(path) |
| 95 | + for name in ("conda", "mamba", "micromamba"): |
| 96 | + found = shutil.which(name) |
| 97 | + if found: |
| 98 | + return found, _run_flags_for(found) |
| 99 | + raise FileNotFoundError( |
| 100 | + "No conda-family launcher (conda / mamba / micromamba) found on " |
| 101 | + "PATH. ARC's cross-env adapters (AutoTST/GCN/TorchANI) need one " |
| 102 | + "of these to launch their subprocess scripts in isolated envs." |
| 103 | + ) |
| 104 | + |
| 105 | + |
| 106 | +def run_in_conda_env( |
| 107 | + python_executable: str, |
| 108 | + script_path: str, |
| 109 | + *script_args: str, |
| 110 | + check: bool = False, |
| 111 | +) -> subprocess.CompletedProcess: |
| 112 | + """Run ``python script_path *script_args`` inside the env that owns |
| 113 | + ``python_executable``, isolated from ARC's process env. |
| 114 | +
|
| 115 | + stdout and stderr are captured and logged centrally — debug on |
| 116 | + success, warning (with both streams and the return code) on |
| 117 | + non-zero exit — so call sites don't each re-implement capture and |
| 118 | + error reporting. The captured streams are also exposed on the |
| 119 | + returned :class:`subprocess.CompletedProcess` (``.stdout`` / |
| 120 | + ``.stderr``) for callers that need to inspect them. ``check=True`` |
| 121 | + raises ``CalledProcessError`` on non-zero exit. Args are passed as |
| 122 | + a list, so no shell quoting concerns. |
| 123 | + """ |
| 124 | + env_prefix = env_prefix_from_python(python_executable) |
| 125 | + launcher, extra_flags = _detect_launcher() |
| 126 | + argv = [ |
| 127 | + launcher, "run", *extra_flags, |
| 128 | + "-p", env_prefix, |
| 129 | + "python", script_path, |
| 130 | + *script_args, |
| 131 | + ] |
| 132 | + result = subprocess.run(argv, check=check, capture_output=True, text=True) |
| 133 | + if result.returncode: |
| 134 | + logger.warning( |
| 135 | + "env-run: %s exited with %d\ncmd: %s\nstdout:\n%s\nstderr:\n%s", |
| 136 | + script_path, result.returncode, " ".join(argv), |
| 137 | + result.stdout, result.stderr, |
| 138 | + ) |
| 139 | + else: |
| 140 | + logger.debug( |
| 141 | + "env-run: %s exited 0\ncmd: %s\nstdout:\n%s\nstderr:\n%s", |
| 142 | + script_path, " ".join(argv), result.stdout, result.stderr, |
| 143 | + ) |
| 144 | + return result |
0 commit comments