diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7ae1b71..8dadc1c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -57,13 +57,17 @@ jobs: run: uv sync --python ${{ matrix.python }} # ruff lint/format are Python-version-agnostic; run once, on the floor. + # PIN the ruff version: bare `uvx ruff` resolves the latest release, so a + # new ruff (0.16.0 broadened the default rule set — E501, N802, S-codes) + # would spontaneously fail CI on unchanged code. Bump this deliberately, + # fixing findings in the same PR. Keep it in step with the dev-group cap. - name: Lint (ruff) if: ${{ matrix.python == '3.12' }} - run: uvx ruff check . + run: uvx ruff@0.15.22 check . - name: Format check (ruff format) if: ${{ matrix.python == '3.12' }} - run: uvx ruff format --check . + run: uvx ruff@0.15.22 format --check . # Coverage floor: a regression guard, not a target. Today's suite sits at # ~79% line coverage; 75 leaves headroom for legitimately-untestable code diff --git a/CLAUDE.md b/CLAUDE.md index 797985a..6f9b7b5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -10,7 +10,11 @@ The full **fetch → process → service → GUI** path runs end-to-end for **bo - **Fetch:** `magnetics.data.fetch.toksearch` (mdsthin via the `cybele` ssh-config alias, or a cluster-side `python -m` run orchestrated by `fetch/remote.py`) writes one HDF5 per shot to `data/datafile/` (gitignored); read back via `magnetics.data.h5source`. The GUI can trigger a - pull from the left rail (`PullControl` → `POST /api/fetch`). + pull from the left rail (`PullControl` → `POST /api/fetch`). The `remote` backend is + location-aware: `network.on_cluster_host()` detects when the server is already running ON the + device's cluster (load-balanced submit nodes included) and fetches in-process — via the cluster + interpreter's toksearch (`remote.run_on_cluster`), else mdsthin — instead of SSHing to the host + it is running on (which hangs on an unanswerable host-key prompt). - **Process:** `core/spectral.py` (MODESPEC) is real and pure. The SLCONTOUR quasi-stationary fit runs end-to-end via the shim-free `core/qs_*` modules (`qs_io_data` → `qs_prep` → `qs_fit`, adapted to nodes by `qs_bridge`; the former `_slcontour` translation, promoted into core) — real @@ -66,6 +70,12 @@ The Python project **is the repo root** (a uv project, served as a webapp). `src - `data/` — sources + `fetch/` (toksearch/mdsthin pulls, cluster orchestration); device configs in `data/device/*.json`. - `service/` — FastAPI; the built GUI is bundled at `service/webapp/` and served here. +- `connect.py` — the `magnetics-connect` remote-GUI launcher (server on a cluster node, + browser local, one SSH tunnel). Self-bootstraps the remote: existing `magnetics` → else + install uv → else `uvx magnetics` (targets PyPI; provisions its own Python). Deliberately + **stdlib-only and standalone** so it runs as a + bare `python3 connect.py` on gateway nodes with no magnetics install — keep it free of + package imports and of syntax newer than ~Python 3.9. Tests in `tests/`, maintainer scripts in `scripts/`. `gui/web/` — React + Vite + TypeScript frontend (its `dist/` is staged into `service/webapp/` for the wheel). @@ -76,10 +86,13 @@ frontend (its `dist/` is staged into `service/webapp/` for the wheel). tests (GA gateway / PPPL flux) are env-gated and skip unless `MAGNETICS_GA_USER` / `MAGNETICS_FLUX_USER` is set; they are manual-only, never wired into CI. - **Frontend:** `cd gui/web && npm run test` (vitest, one-shot; `npm run test:watch` to iterate). -- **Everything CI runs:** `uv run ruff format --check .` + `uvx ruff check .` + `uv run pytest` +- **Everything CI runs:** `uv run ruff format --check .` + `uv run ruff check .` + `uv run pytest` + `uv run ty check src/magnetics` (Python), and `npm run lint` + `npm run typecheck` + `npm run test` + `npm run build` (in `gui/web/`). Run these locally before committing — a plain `pytest` + `tsc` pass does NOT cover everything CI checks. + Use `uv run ruff` (not a bare `uvx ruff`, which resolves the *latest* release): the dev group + caps ruff at `<0.16`, matching the version CI pins, so local lint can't drift from CI when a + new ruff broadens its default rules. ## Conventions - Physics lives in `src/magnetics/core` (pure, device-agnostic, testable); **no physics in the diff --git a/README.md b/README.md index 8269e29..0c7976e 100644 --- a/README.md +++ b/README.md @@ -23,11 +23,12 @@ magnetics --no-browser # headless (e.g. on a server) magnetics --data-dir PATH # where shot data lives (see below) ``` -Two more console scripts come with the install: +Three more console scripts come with the install: ```sh magnetics-fetch --shot 184927 # pull a shot to the local data dir magnetics-service # start just the API/service (honors HOST/PORT) +magnetics-connect omega # run the server on a remote host, browse it locally ``` ### Where shot data lives @@ -47,6 +48,47 @@ usually quota'd and shot files are large) — e.g. in your shell profile: export MAGNETICS_DATA_DIR=$SCRATCH/magnetics ``` +## Remote use — run on a cluster, browse locally + +The server should run where the data lives (a cluster work node), but that is +rarely where your browser lives — you're on a laptop off-site, or on a NoMachine +desktop on the site's login/gateway node with a firewall between you and the +work node. `magnetics-connect` does the standard SSH-tunnel dance in one command: +it opens one authenticated connection to the host, asks it for a free port and +its real node name, starts the server there **bound to loopback only**, carries +a `-L` port-forward on the same connection, and opens your local browser once +the server answers through the tunnel. Only HTTP crosses the wire — fetches and +fits run next to the data, plots stream back. Ctrl-C stops both ends. + +```sh +magnetics-connect omega # ~/.ssh/config alias (keys, ProxyJump) +magnetics-connect me@omega.gat.com -J me@cybele.gat.com:2039 +magnetics-connect omega --data-dir /cscratch/$USER/magnetics +magnetics-connect flux --remote-cmd \ + 'module load magnetics && magnetics --no-browser --port {port}' +``` + +**Zero-install on the remote.** By default the launcher bootstraps the server +for you: it runs an already-installed `magnetics` if the host has one; otherwise +it installs `uv` (the official root-free installer, into `~/.local/bin`) and +launches with `uvx magnetics`, which fetches the package **and provisions its own +Python** — so a bare cluster node with only an old system Python needs nothing +installed by hand. The first run downloads + provisions (slower; watch the +streamed `[host]` log); later runs are cached and fast. + +- `--install-from SOURCE` — install from a wheel URL/path (any uv source) instead + of PyPI. This is the shim until `magnetics` is published to PyPI. +- `--remote-cmd '…{port}…'` — full override of the launch line (bypasses the + bootstrap), e.g. for an eventual `module load magnetics`. + +The launcher itself is **stdlib-only and standalone** — on a gateway node with +only system `python3`, copy the single file and run it directly: + +```sh +scp src/magnetics/connect.py gateway: +python3 connect.py omega +``` + ## Development The Python project is the repo root (a uv project); the React GUI is in `gui/web/`. diff --git a/pyproject.toml b/pyproject.toml index 3e10d35..78a4772 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,10 @@ build-backend = "uv_build" [dependency-groups] dev = [ "pytest>=9.1.1", - "ruff>=0.15.18", + # Cap below 0.16: that minor broadened ruff's default rule set and would fail + # the existing tree. Raise the cap deliberately (fixing findings) in step with + # the pinned `uvx ruff@…` in .github/workflows/ci.yml. + "ruff>=0.15.18,<0.16", # notebook tooling (examples + clean notebook commits) "ipykernel>=7.3.0", "jupyter>=1.1.0", @@ -51,6 +54,10 @@ magnetics = "magnetics.cli:main" magnetics-service = "magnetics.service.app:main" # Fetch a shot to the local data dir (the toksearch/mdsthin CLI). magnetics-fetch = "magnetics.data.fetch.toksearch:main" +# Run the server on a remote host (cluster) over an SSH tunnel, browse locally. +# Stdlib-only and standalone: src/magnetics/connect.py also runs as a bare +# `python3 connect.py` on hosts where magnetics is not installed. +magnetics-connect = "magnetics.connect:main" [tool.ruff] line-length = 100 # target version is inferred from requires-python diff --git a/src/magnetics/connect.py b/src/magnetics/connect.py new file mode 100644 index 0000000..02d9d2a --- /dev/null +++ b/src/magnetics/connect.py @@ -0,0 +1,350 @@ +#!/usr/bin/env python3 +"""One-command remote GUI: run the magnetics server on a cluster, browse it locally. + +The problem this solves: the server must run where the data lives (omega, a PPPL +cluster, ...), but the browser lives somewhere else -- a NoMachine desktop on the +site gateway, or a laptop off-site -- and 127.0.0.1 there is not 127.0.0.1 here. +Rather than X-forwarding a browser across the firewall, this does the standard +SSH tunnel dance in one command: + + 1. open one authenticated SSH master to HOST (ProxyJump'd when given); + 2. ask the remote for a free loopback port and its real node name (submit + nodes are load-balanced -- "omega" may land you on omega-a or omega-b); + 3. start the server there, bound to 127.0.0.1 only, with the tunnel + (``-L local:127.0.0.1:remote``) carried by the same connection; + 4. poll through the tunnel until the server answers, then open the local + browser at ``http://localhost:``. + +Only HTTP crosses the wire -- fetches and fits run on the cluster, plots stream +back. Ctrl-C tears everything down (the forced remote pty SIGHUPs the server, +so nothing is orphaned on the shared node). + +This file is deliberately STDLIB-ONLY and self-contained so it also runs where +magnetics is not installed (e.g. a gateway node with only system python3): + + scp src/magnetics/connect.py gateway: + python3 connect.py omega + +With the package installed it is simply ``magnetics-connect``. + +Examples: + magnetics-connect omega # ~/.ssh/config alias, key auth + magnetics-connect pharrm@omega.gat.com -J pharrm@cybele.gat.com:2039 + magnetics-connect omega --data-dir /cscratch/$USER/magnetics + magnetics-connect flux --remote-cmd \\ + 'module load magnetics && magnetics --no-browser --port {port}' +""" + +from __future__ import annotations + +import argparse +import shlex +import signal +import socket +import subprocess +import sys +import threading +import time +import urllib.error +import urllib.request +import webbrowser + +# The PyPI package + console-script the server ships as. Once published, +# `uvx magnetics` fetches it and provisions its own Python — nothing to install +# by hand. Until then, the bootstrap's fast path uses an already-installed +# `magnetics` (or point --install-from at a wheel). +PACKAGE = "magnetics" +# The root-free, official uv installer. omega/cybele have curl + outbound net; +# it drops uv in ~/.local/bin without touching the system. +UV_INSTALLER = "https://astral.sh/uv/install.sh" + +# One remote python3 -c: print a free loopback port and the node's real name. +_PROBE_PY = ( + 'import socket; s = socket.socket(); s.bind(("127.0.0.1", 0)); ' + "print(s.getsockname()[1], socket.getfqdn())" +) + + +def _ctl_opts(sock: str) -> list[str]: + """ControlMaster options: authenticate once, reuse for the serve connection.""" + return ["-o", "ControlMaster=auto", "-o", f"ControlPath={sock}", "-o", "ControlPersist=60"] + + +def control_sock(target: str) -> str: + """Short per-target ControlPath in /tmp (macOS socket-path length limit).""" + tag = "".join(c if (c.isalnum() or c in "._-") else "_" for c in target) + return f"/tmp/mc-{tag}.sock" + + +def _serve_args(data_dir: str | None) -> str: + """The `magnetics` server flags shared by every launch path. `{port}` is a + placeholder resolved later (it appears more than once in the bootstrap).""" + args = "--no-browser --port {port}" + if data_dir: + args += " --data-dir " + shlex.quote(data_dir) + return args + + +def default_remote_cmd(data_dir: str | None = None, install_from: str | None = None) -> str: + """A self-bootstrapping serve command (POSIX sh, run via ``bash -lc``). + + In order: run an already-installed ``magnetics`` if the remote has one (fast + path, no install); else ensure ``uv`` (fetch the root-free installer if it is + missing); else launch with ``uvx``, which downloads the package and provisions + a suitable Python — omega's system Python is far too old, and uvx sidesteps + that. ``install_from`` (a wheel URL / path / any uv source) overrides the PyPI + default via ``uvx --from`` — the shim until the package is published. + """ + serve = _serve_args(data_dir) + from_opt = ("--from " + shlex.quote(install_from) + " ") if install_from else "" + return ( + 'export PATH="$HOME/.local/bin:$PATH"; ' + f"if command -v {PACKAGE} >/dev/null 2>&1; then exec {PACKAGE} {serve}; fi; " + "if ! command -v uvx >/dev/null 2>&1; then " + "echo 'magnetics-connect: installing uv (one-time, no root)…' >&2; " + f"curl -LsSf {UV_INSTALLER} | sh || exit 1; " + 'export PATH="$HOME/.local/bin:$PATH"; fi; ' + "echo 'magnetics-connect: launching via uvx (first run provisions Python + deps)…' >&2; " + f"exec uvx {from_opt}{PACKAGE} {serve}" + ) + + +def resolve_remote_cmd(template: str, port: int) -> str: + """Substitute the chosen remote port into the serve command (every + ``{port}`` occurrence; the bootstrap names it more than once).""" + if "{port}" in template: + return template.replace("{port}", str(port)) + return f"{template} --port {port}" + + +def build_probe_cmd(target: str, sock: str, jump: str | None, ssh_opts: list[str]) -> list[str]: + cmd = ["ssh", *_ctl_opts(sock)] + for opt in ssh_opts: + cmd += ["-o", opt] + if jump: + cmd += ["-J", jump] + cmd += [target, f"python3 -c {shlex.quote(_PROBE_PY)}"] + return cmd + + +def build_serve_cmd( + target: str, + sock: str, + local_port: int, + remote_port: int, + remote_cmd: str, + jump: str | None, + ssh_opts: list[str], +) -> list[str]: + # -tt forces a remote pty even with our stdout piped, so when this ssh dies + # (Ctrl-C, laptop lid) the server gets SIGHUP -- no orphans on a shared node. + cmd = ["ssh", *_ctl_opts(sock), "-tt", "-o", "ExitOnForwardFailure=yes"] + for opt in ssh_opts: + cmd += ["-o", opt] + if jump: + cmd += ["-J", jump] + cmd += ["-L", f"{local_port}:127.0.0.1:{remote_port}"] + # A login shell picks up module-/profile-managed PATHs that a bare ssh + # command shell would miss (module load in ~/.bash_profile etc.). + cmd += [target, f"bash -lc {shlex.quote(remote_cmd)}"] + return cmd + + +def pick_local_port(preferred: int | None = None) -> int: + """`preferred` if it is free locally (same number both ends reads nicer), + else any OS-assigned free port.""" + if preferred: + with socket.socket() as s: + try: + s.bind(("127.0.0.1", preferred)) + return preferred + except OSError: + pass + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _wait_ready(url: str, child: subprocess.Popen, timeout: float) -> bool: + """Poll `url` through the tunnel until the server answers (any HTTP status) + or `child` (the ssh carrying it) exits / `timeout` passes.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if child.poll() is not None: + return False + try: + with urllib.request.urlopen(url, timeout=2): + return True + except urllib.error.HTTPError: + return True # an HTTP response IS the server answering + except Exception: + time.sleep(0.5) + return False + + +def _pump(stream, prefix: str) -> None: + """Mirror the remote server's output locally, line by line as it arrives.""" + for line in stream: + sys.stdout.write(f"{prefix}{line.rstrip()}\n") + sys.stdout.flush() + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser( + prog="magnetics-connect", + description="Start the magnetics GUI on a remote host and open it in your local browser.", + ) + ap.add_argument("host", help="ssh destination: a ~/.ssh/config alias or user@host") + ap.add_argument( + "-J", + "--jump", + default=None, + help="ProxyJump gateway (user@host[:port]); omit when the alias carries its own", + ) + ap.add_argument( + "--remote-cmd", + default=None, + help="full override for the command that starts the server on HOST, with " + "{port} substituted; run through `bash -lc`. Bypasses the uv/uvx " + "bootstrap — use when you have your own launch line (e.g. a module load).", + ) + ap.add_argument( + "--install-from", + default=None, + metavar="SOURCE", + help="install the server from SOURCE instead of PyPI (a wheel URL/path or " + "any uv source), passed to `uvx --from`. The shim until the package is " + "on PyPI; ignored if the remote already has `magnetics` installed.", + ) + ap.add_argument( + "--data-dir", + default=None, + help="remote shot-data directory, forwarded to the launch command " + "(on a cluster, point at scratch/project space)", + ) + ap.add_argument( + "--remote-port", + type=int, + default=None, + help="fixed remote port (default: ask the remote for a free one)", + ) + ap.add_argument( + "--local-port", + type=int, + default=None, + help="fixed local port (default: mirror the remote port when free)", + ) + ap.add_argument("--no-browser", action="store_true", help="print the URL, don't open it") + ap.add_argument( + "--timeout", + type=float, + default=300.0, + help="seconds to wait for the remote server to answer (default: 300; the " + "first uvx run provisions Python + deps and can be slow)", + ) + ap.add_argument( + "-o", + "--ssh-opt", + action="append", + default=[], + metavar="OPT", + help="extra ssh -o option (repeatable), e.g. -o ConnectTimeout=30", + ) + args = ap.parse_args(argv) + + sock = control_sock(args.host) + remote_cmd_template = args.remote_cmd or default_remote_cmd(args.data_dir, args.install_from) + + # Ctrl-C signals the whole foreground process group, so ssh dies with us and + # SIGHUPs the remote server. A bare `kill ` does not: it would drop us + # without running any cleanup, orphaning the server on a shared cluster node. + # Route SIGTERM into the same KeyboardInterrupt path so both tear down. + def _on_sigterm(_signum, _frame): + raise KeyboardInterrupt + + signal.signal(signal.SIGTERM, _on_sigterm) + + # 1) probe: authenticates the master (prompts go to the tty, not our pipe) + # and reports a free remote loopback port + the actual node we landed on. + print(f"Connecting to {args.host}" + (f" via {args.jump}" if args.jump else "") + " ...") + probe = subprocess.run( + build_probe_cmd(args.host, sock, args.jump, args.ssh_opt), + stdout=subprocess.PIPE, + text=True, + ) + child: subprocess.Popen | None = None # the finally tears this down + try: + if probe.returncode != 0: + print( + "error: could not reach the remote host (check the ssh alias / " + "gateway / credentials).", + file=sys.stderr, + ) + return 1 + probed_port, _, node = probe.stdout.strip().partition(" ") + remote_port = args.remote_port or int(probed_port) + local_port = args.local_port or pick_local_port(remote_port) + url = f"http://localhost:{local_port}" + + # 2) one connection carries both the tunnel and the server process. + remote_cmd = resolve_remote_cmd(remote_cmd_template, remote_port) + print(f"Starting on {node or args.host}: {remote_cmd}") + child = subprocess.Popen( + build_serve_cmd( + args.host, sock, local_port, remote_port, remote_cmd, args.jump, args.ssh_opt + ), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + threading.Thread( + target=_pump, args=(child.stdout, f" [{args.host}] "), daemon=True + ).start() + + # 3) the server is up when it answers through the tunnel. + if not _wait_ready(url, child, args.timeout): + print( + f"error: the remote server never answered on {url} within " + f"{args.timeout:.0f}s (last output above). If it was still " + "installing, retry with a larger --timeout; otherwise check the " + "remote can reach PyPI, or pass --install-from / --remote-cmd.", + file=sys.stderr, + ) + return 1 + + print() + print(f" magnetics GUI: {url}") + print(f" server: {node or args.host}, port {remote_port} (loopback only)") + print(" Ctrl-C to stop both ends.") + print() + if not args.no_browser: + webbrowser.open(url) + + try: + return child.wait() + except KeyboardInterrupt: + print("\nStopped.") + return 0 + finally: + # Always take the tunnel down with us: while this ssh lives, so does the + # remote server it carries. Covers every exit — Ctrl-C, SIGTERM, a failed + # readiness wait, an exception — so nothing is left on a shared node. + if child is not None and child.poll() is None: + child.terminate() + try: + child.wait(timeout=10) + except subprocess.TimeoutExpired: + # ssh ignored SIGTERM (a stuck remote pipeline: mid-bootstrap, a + # wedged fetch). Escalate rather than leave it holding the server. + child.kill() + child.wait() + subprocess.run( + ["ssh", "-o", f"ControlPath={sock}", "-O", "exit", args.host], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/magnetics/data/fetch/network.py b/src/magnetics/data/fetch/network.py index 648f517..15bb298 100644 --- a/src/magnetics/data/fetch/network.py +++ b/src/magnetics/data/fetch/network.py @@ -66,6 +66,34 @@ def on_site_network(device: str) -> bool: return socket.getfqdn().lower().endswith(domain) +def on_cluster_host(device: str) -> bool: + """True when THIS process is already running on the device's compute cluster, + so the ``remote`` backend's SSH round trip would dial the very host it is + running on. + + That self-SSH is not merely wasteful: on a node whose ``known_hosts`` lacks + the cluster's own entry, ssh stops at the host-key confirmation prompt, and a + server process has no terminal to answer it -- the pull hangs at 0% forever + rather than failing. Callers use this to fetch in-process instead. + + Cluster submit nodes are load-balanced, so the FQDN we land on + (``omega-a.gat.com``) is rarely the configured cluster address + (``omega.gat.com``); we compare short names by prefix, gated on already being + inside the site domain. ``MAGNETICS_ON_CLUSTER`` (``0``/``1``) forces it. + """ + override = os.environ.get("MAGNETICS_ON_CLUSTER") + if override is not None: + return override.strip().lower() not in ("", "0", "false", "no") + host = str(cluster_login(device).get("host") or "").strip().lower() + if not host or not on_site_network(device): + return False + cluster_short = host.split(".")[0] + me_short = socket.getfqdn().lower().split(".")[0] + if not cluster_short or not me_short: + return False + return me_short.startswith(cluster_short) + + def _hostport(block, default_port: int) -> str | None: """Render a ``{host, port}`` sub-block as ``host:port`` (None if no host).""" block = block or {} diff --git a/src/magnetics/data/fetch/remote.py b/src/magnetics/data/fetch/remote.py index d064ad2..c0337eb 100644 --- a/src/magnetics/data/fetch/remote.py +++ b/src/magnetics/data/fetch/remote.py @@ -37,10 +37,13 @@ from __future__ import annotations +import os import re import shlex +import shutil import subprocess import sys +import tempfile from pathlib import Path from .. import h5source @@ -91,6 +94,123 @@ def _log(msg: str) -> None: sys.stderr.flush() +def _fetch_argv( + python, shot, analysis, device, out, tmin, tmax, decimate, sensor_set, raw_pointnames +): + """The `python -m magnetics.data.fetch.toksearch --backend toksearch …` argv. + + Shared by the SSH path (run_remote) and the on-cluster path + (run_on_cluster) so both invoke the native-PTDATA reader identically. + """ + argv = [ + python, + "-m", + "magnetics.data.fetch.toksearch", + "--backend", + "toksearch", + "--shot", + str(shot), + "--device", + device, + "--out", + out, + ] + if raw_pointnames: + argv += ["--pointnames", ",".join(raw_pointnames)] + elif sensor_set: + argv += ["--sensor-set", sensor_set] + else: + argv += ["--analysis", analysis] + if tmin is not None: + argv += ["--tmin", str(tmin)] + if tmax is not None: + argv += ["--tmax", str(tmax)] + if decimate and decimate > 1: + argv += ["--decimate", str(decimate)] + return argv + + +def cluster_python_has_toksearch(python: str | None) -> bool: + """Whether `python` can import toksearch + the d3d PTDATA plugin. + + The probe itself is the arbiter (a bad interpreter just returns non-zero), + so resolve bare command names through PATH the way subprocess would rather + than pre-rejecting them: `Path("python3").exists()` is false for anything + not sitting in the cwd, which would silently downgrade an on-cluster pull to + mdsthin (~5-7x slower) for a caller who passed `--remote-python python3`. + """ + if not python: + return False + resolved = python if os.path.isabs(python) else (shutil.which(python) or python) + try: + probe = subprocess.run( + [resolved, "-c", "import toksearch, toksearch_d3d"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + except OSError: # not found / not executable + return False + return probe.returncode == 0 + + +def run_on_cluster( + shot, + analysis="both", + *, + python=None, + device="diiid", + tmin=None, + tmax=None, + decimate=1, + sensor_set=None, + raw_pointnames=None, + out=None, + progress=None, +) -> str: + """Run the toksearch pull HERE, under the cluster's env interpreter. + + We are already ON the cluster: there is nothing to ssh to and nothing to copy + back. But `toksearch` lives in the site conda env, not in whatever venv runs + the GUI, so a plain in-process fetch would fall back to mdsthin -- the mdsip + path, which is ~5-7x slower than the native PTDATA read and is the *reason* + the remote backend exists. Re-invoking this same module under the cluster + interpreter keeps the fast reader without any network hop. + + This is exactly the command `run_remote` sends over SSH, minus the SSH. + """ + python = python or DEFAULT_PYTHON + out = str(out) + argv = _fetch_argv( + python, shot, analysis, device, out, tmin, tmax, decimate, sensor_set, raw_pointnames + ) + # Expose ONLY the magnetics package to the cluster interpreter -- never our + # whole venv site-packages. That directory carries numpy/scipy/etc. built for + # THIS interpreter (e.g. 3.14); on the conda env's 3.11 those C extensions + # shadow its own working copies and break `import toksearch` (which pulls in + # numpy). A temp dir holding just a `magnetics` symlink resolves `-m + # magnetics.…` while the core's deps come from the conda env itself -- exactly + # what the SSH path achieves by rsync'ing only the package. + stage = tempfile.mkdtemp(prefix="magnetics_pkg_") + try: + os.symlink(PKG_ROOT, os.path.join(stage, PKG_NAME)) + env = dict(os.environ) + # Fully override (don't append): a leaked outer PYTHONPATH could re-add the + # venv. TOKSEARCH_INDEX_DIR only matters for SQL shot *discovery* (we pass + # an explicit shot), so point it at the workdir to silence its warning. + env["PYTHONPATH"] = stage + env.setdefault("TOKSEARCH_INDEX_DIR", str(Path(out).parent)) + _log(f"Fetching on this node with {python}: {shlex.join(argv)}") + if progress: + progress(0.5, f"pulling shot {shot} on this node (toksearch)") + if subprocess.run(argv, env=env).returncode != 0: + raise RuntimeError("on-cluster toksearch fetch failed") + finally: + shutil.rmtree(stage, ignore_errors=True) + if progress: + progress(1.0, f"done: {Path(out).name}") + return out + + def run_remote( shot, analysis="both", @@ -201,31 +321,18 @@ def run(cmd, **kw): # an ABSOLUTE path in the /tmp stage, so it lands there regardless of the # `cd {remote_dir}` the inner command does for the package imports. remote_out = f"{out_stage}/shot_{shot}.h5" - fetch = [ + fetch = _fetch_argv( python, - "-m", - "magnetics.data.fetch.toksearch", - "--backend", - "toksearch", - "--shot", - str(shot), - "--device", + shot, + analysis, device, - "--out", remote_out, - ] - if raw_pointnames: - fetch += ["--pointnames", ",".join(raw_pointnames)] - elif sensor_set: - fetch += ["--sensor-set", sensor_set] - else: - fetch += ["--analysis", analysis] - if tmin is not None: - fetch += ["--tmin", str(tmin)] - if tmax is not None: - fetch += ["--tmax", str(tmax)] - if decimate and decimate > 1: - fetch += ["--decimate", str(decimate)] + tmin, + tmax, + decimate, + sensor_set, + raw_pointnames, + ) # TOKSEARCH_INDEX_DIR only matters for SQL/index shot *discovery*; we pass an # explicit shotlist, so point it at the workdir to silence ~1 warning/signal. inner = ( diff --git a/src/magnetics/data/fetch/toksearch.py b/src/magnetics/data/fetch/toksearch.py index 584291f..294d5b8 100644 --- a/src/magnetics/data/fetch/toksearch.py +++ b/src/magnetics/data/fetch/toksearch.py @@ -66,7 +66,13 @@ # Connection endpoints (mdsip / gateway) + automatic on-site detection, resolved # from the device file's `network` block so the hop count is picked for the user. -from .network import gateway_address, mdsip_address, on_site_network +from .network import ( + cluster_login, + gateway_address, + mdsip_address, + on_cluster_host, + on_site_network, +) # All fetched shot files land in the runtime data dir (data/datafile/) — the same # place h5source reads back from ($MAGNETICS_DATA_DIR or the repo's data/ dir). @@ -1497,6 +1503,45 @@ def fetch_shot( sys.stderr.write(f"All {n_skipped} requested signals already in {out}; nothing to fetch.\n") return out + # Already ON the cluster? Then "remote" has nothing to reach: its SSH would + # dial this very host, and on a node without the cluster in known_hosts it + # blocks on a host-key prompt no server process can answer (pull pinned at + # 0%). The data is local here -- but the POINT of the remote backend is + # toksearch's native PTDATA read (~5-7x mdsip), and toksearch lives in the + # site conda env, not in the venv running the GUI. So prefer, in order: + # 1. toksearch importable right here -> plain in-process fetch + # 2. the cluster interpreter has it -> subprocess, no ssh + # 3. neither -> mdsthin (slow mdsip) + # Falling straight to (3) is what made an on-cluster pull SLOWER than the + # laptop's remote pull, which had been quietly using toksearch all along. + if backend == "remote" and not _tree_transport and on_cluster_host(device): + from . import remote as remote_run + + try: + import toksearch # noqa: F401 # ty: ignore[unresolved-import] + + backend = "toksearch" + progress(0.0, "already on the cluster; remote→toksearch") + except ImportError: + cluster_py = remote_python or cluster_login(device).get("python") + if remote_run.cluster_python_has_toksearch(cluster_py): + progress(0.0, "already on the cluster; toksearch via the cluster env") + return remote_run.run_on_cluster( + shot, + analysis, + python=cluster_py, + device=device, + tmin=tmin, + tmax=tmax, + decimate=decimate, + sensor_set=sensor_set, + raw_pointnames=raw_pointnames, + out=out, + progress=progress, + ) + backend = "mdsthin" + progress(0.0, "already on the cluster; no toksearch → mdsthin (slow)") + if backend == "remote" and not _tree_transport: # Orchestrate a pull on the cluster from here; remote side runs this same # script with --backend toksearch and writes the file we copy back. When a diff --git a/tests/test_connect_launcher.py b/tests/test_connect_launcher.py new file mode 100644 index 0000000..3fa61c2 --- /dev/null +++ b/tests/test_connect_launcher.py @@ -0,0 +1,278 @@ +"""magnetics-connect: the remote-GUI tunnel launcher, with ssh/subprocess faked. + +The cluster round trip (NoMachine desktop -> gateway -> work node, or laptop -> +cluster) can't run offline, but everything the launcher does is command +construction + orchestration: the ControlMaster probe, the free-port handshake, +the single ssh that carries both the -L tunnel and the server process, readiness +polling through the tunnel, and teardown. These tests script all of it. +""" + +from __future__ import annotations + +import socket +import subprocess +import urllib.request +from types import SimpleNamespace + +import pytest + +from magnetics import connect + +PROBED = "43211 omega-b.gat.com\n" + + +class _Run: + """Record subprocess.run calls; script the probe's stdout/returncode.""" + + def __init__(self, probe_stdout=PROBED, probe_rc=0): + self.calls: list[list[str]] = [] + self.probe_stdout = probe_stdout + self.probe_rc = probe_rc + + def __call__(self, cmd, **kw): + self.calls.append(list(cmd)) + joined = " ".join(str(c) for c in cmd) + if "python3 -c" in joined: # the free-port/node probe + return SimpleNamespace(returncode=self.probe_rc, stdout=self.probe_stdout, stderr="") + return SimpleNamespace(returncode=0, stdout="", stderr="") + + def joined(self): + return [" ".join(str(c) for c in call) for call in self.calls] + + +class _Popen: + """The ssh carrying tunnel + server: stays 'alive' unless told otherwise.""" + + instances: list["_Popen"] = [] + + def __init__(self, cmd, **kw): + self.cmd = list(cmd) + self.stdout = iter(["INFO: Uvicorn running\r\n"]) + self.terminated = False + self._poll_rc = None + _Popen.instances.append(self) + + def poll(self): + return self._poll_rc + + def wait(self, timeout=None): + return 0 + + def terminate(self): + self.terminated = True + + +class _Resp: + """Context-manager fake for urlopen (dunders must live on the TYPE).""" + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + +@pytest.fixture() +def harness(monkeypatch): + """Fake run/Popen/urlopen/webbrowser; return the recorders.""" + _Popen.instances = [] + run = _Run() + opened: list[str] = [] + monkeypatch.setattr(subprocess, "run", run) + monkeypatch.setattr(subprocess, "Popen", _Popen) + monkeypatch.setattr(urllib.request, "urlopen", lambda url, timeout=0: _Resp()) + monkeypatch.setattr(connect.webbrowser, "open", opened.append) + # deterministic local port regardless of what's free on the test machine + monkeypatch.setattr(connect, "pick_local_port", lambda preferred=None: 55555) + return SimpleNamespace(run=run, opened=opened) + + +def test_happy_path_builds_probe_tunnel_and_serve(harness): + rc = connect.main(["omega"]) + assert rc == 0 + cmds = harness.run.joined() + + # probe: ControlMaster established once; asks the remote for port + node + probe = cmds[0] + assert "ControlMaster=auto" in probe and "omega" in probe + assert "python3 -c" in probe + + # serve: ONE ssh carries the forced pty, the tunnel, and the server command + serve = " ".join(_Popen.instances[0].cmd) + assert "-tt" in serve + assert "ExitOnForwardFailure=yes" in serve + assert "-L 55555:127.0.0.1:43211" in serve # local port -> probed remote port + assert "bash -lc" in serve + assert "--port 43211" in serve # substituted into the default remote cmd + assert "--no-browser" in serve # the REMOTE server must not open a browser + + # local browser opened on the tunnel entrance; master torn down at exit + assert harness.opened == ["http://localhost:55555"] + assert any("-O exit" in c for c in cmds) + + +def test_jump_flag_reaches_both_ssh_invocations(harness): + connect.main(["u@omega.gat.com", "-J", "u@cybele.gat.com:2039"]) + assert "-J u@cybele.gat.com:2039" in harness.run.joined()[0] + assert "-J u@cybele.gat.com:2039" in " ".join(_Popen.instances[0].cmd) + + +def test_no_browser_and_data_dir(harness): + connect.main(["omega", "--no-browser", "--data-dir", "/cscratch/me/mag data"]) + assert harness.opened == [] + # the path is forwarded and survives re-quoting through `bash -lc` (spaces intact) + assert "--data-dir" in _Popen.instances[0].cmd[-1] + assert "mag data" in _Popen.instances[0].cmd[-1] + # the helper quotes it so the remote shell sees ONE argument + default = connect.default_remote_cmd("/cscratch/me/mag data") + assert default.endswith("--data-dir '/cscratch/me/mag data'") + + +def test_custom_remote_cmd_port_placeholder(harness): + connect.main(["flux", "--remote-cmd", "module load magnetics && magnetics --port {port}"]) + serve = " ".join(_Popen.instances[0].cmd) + assert "module load magnetics && magnetics --port 43211" in serve + assert serve.count("--port") == 1 # placeholder consumed, nothing appended + + +def test_probe_failure_exits_before_any_server(monkeypatch): + _Popen.instances = [] + run = _Run(probe_rc=255) + monkeypatch.setattr(subprocess, "run", run) + monkeypatch.setattr(subprocess, "Popen", _Popen) + assert connect.main(["omega"]) == 1 + assert _Popen.instances == [] # never tried to start the server + + +def _refuse_connections(monkeypatch): + def _refused(url, timeout=0): + raise OSError("connection refused") + + monkeypatch.setattr(urllib.request, "urlopen", _refused) + + +def test_remote_command_died_reports_failure(harness, monkeypatch): + """ssh already exited (e.g. the launch command was not found): report the + failure and open no browser. Nothing to terminate — it's gone.""" + _refuse_connections(monkeypatch) + + class _DeadPopen(_Popen): + def __init__(self, cmd, **kw): + super().__init__(cmd, **kw) + self._poll_rc = 127 + + monkeypatch.setattr(subprocess, "Popen", _DeadPopen) + assert connect.main(["omega"]) == 1 + assert harness.opened == [] + assert not _Popen.instances[0].terminated # already dead; no signal needed + + +def test_server_never_ready_still_tears_the_tunnel_down(harness, monkeypatch): + """The timeout path with ssh still alive: it must not be left holding the + tunnel (and the remote server) open on a shared node.""" + _refuse_connections(monkeypatch) + # _Popen's default poll() is None → still running + assert connect.main(["omega", "--timeout", "0.1"]) == 1 + assert _Popen.instances[0].terminated + assert harness.opened == [] + + +def test_ctrl_c_escalates_to_kill_when_ssh_ignores_sigterm(harness, monkeypatch): + """A stuck remote (mid-bootstrap, wedged fetch) is exactly when Ctrl-C + matters. TimeoutExpired is not a KeyboardInterrupt, so an uncaught one would + turn the clean 'Stopped.' exit into a traceback.""" + + class _StubbornPopen(_Popen): + def __init__(self, cmd, **kw): + super().__init__(cmd, **kw) + self.killed = False + self._waits = 0 + + def wait(self, timeout=None): + self._waits += 1 + if self._waits == 1: + raise KeyboardInterrupt # the user's Ctrl-C + if timeout is not None and not self.killed: + raise subprocess.TimeoutExpired(self.cmd, timeout) # ignores SIGTERM + return 0 + + def kill(self): + self.killed = True + + monkeypatch.setattr(subprocess, "Popen", _StubbornPopen) + assert connect.main(["omega"]) == 0 # clean exit, no traceback + child = _Popen.instances[0] + assert child.terminated and child.killed # SIGTERM first, then escalate + + +class TestBootstrap: + """The default remote command self-installs: fast path if magnetics is + present, else ensure uv, else `uvx magnetics` (PyPI, provisions its own + Python). This is the zero-touch hand-off for a fresh GA machine.""" + + def test_three_tiers_in_order(self): + cmd = connect.default_remote_cmd() + # 1) fast path: use an installed magnetics before anything else + fast = cmd.index("command -v magnetics") + exec_installed = cmd.index("exec magnetics") + # 2) ensure uv only if uvx is missing (root-free installer) + ensure_uv = cmd.index("command -v uvx") + assert "astral.sh/uv/install.sh" in cmd + # 3) last resort: uvx provisions + runs from PyPI + uvx = cmd.index("exec uvx magnetics") + assert fast < exec_installed < ensure_uv < uvx + + def test_prepends_local_bin_to_path(self): + # uv installs into ~/.local/bin; it must be on PATH for the uvx step + assert "$HOME/.local/bin" in connect.default_remote_cmd() + + def test_install_from_becomes_uvx_from(self): + url = "https://example.org/magnetics-0.1.0-py3-none-any.whl" + cmd = connect.default_remote_cmd(install_from=url) + assert f"uvx --from {url} magnetics" in cmd + # no --from when unset (plain PyPI) + assert "--from" not in connect.default_remote_cmd() + + def test_every_port_placeholder_is_resolved(self, harness): + # {port} appears in BOTH the fast-path and uvx exec lines + connect.main(["omega"]) + serve = " ".join(_Popen.instances[0].cmd) + assert "{port}" not in serve + assert serve.count("--port 43211") == 2 # installed path + uvx path + + def test_bootstrap_is_the_default_but_remote_cmd_bypasses_it(self, harness): + connect.main(["omega", "--remote-cmd", "magnetics --port {port}"]) + serve = " ".join(_Popen.instances[0].cmd) + assert "uvx" not in serve and "astral.sh" not in serve + + +class TestHelpers: + def test_resolve_remote_cmd_appends_when_no_placeholder(self): + assert connect.resolve_remote_cmd("magnetics --no-browser", 8123).endswith("--port 8123") + + def test_control_sock_sanitizes_target(self): + sock = connect.control_sock("user@host.gat.com:2039") + assert sock.startswith("/tmp/mc-") and ":" not in sock and "@" not in sock + + def test_pick_local_port_mirrors_free_preferred(self): + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + free = s.getsockname()[1] + assert connect.pick_local_port(free) == free + + def test_pick_local_port_falls_back_when_busy(self): + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + busy = s.getsockname()[1] + got = connect.pick_local_port(busy) + assert got != busy and got > 0 + + def test_wait_ready_true_on_http_error_response(self, monkeypatch): + import urllib.error + + def _forbidden(url, timeout=0): + raise urllib.error.HTTPError(url, 403, "forbidden", {}, None) + + monkeypatch.setattr(urllib.request, "urlopen", _forbidden) + child = SimpleNamespace(poll=lambda: None) + assert connect._wait_ready("http://localhost:1", child, timeout=5) diff --git a/tests/test_on_cluster_detection.py b/tests/test_on_cluster_detection.py new file mode 100644 index 0000000..7660a96 --- /dev/null +++ b/tests/test_on_cluster_detection.py @@ -0,0 +1,286 @@ +"""On-cluster detection: don't SSH to the host you're already running on. + +Running the GUI on a cluster work node (the NoMachine → gateway → omega flow) +sent the "remote" backend's ssh to omega.gat.com *from omega*. Where the node's +known_hosts lacks that entry, ssh blocks on the host-key prompt, and a server +process has no tty to answer it — the pull sat at 0% forever instead of failing. + +These pin the detector's boundaries (a load-balanced submit node still counts; +the site gateway does NOT) and the dispatch downgrade that avoids the self-SSH. +""" + +from __future__ import annotations + +import os +import socket +from types import SimpleNamespace + +import pytest + +from magnetics.data.fetch import network, toksearch + + +@pytest.fixture(autouse=True) +def _no_env_override(monkeypatch): + """The env forces the answer; clear it so tests see the real logic.""" + monkeypatch.delenv("MAGNETICS_ON_CLUSTER", raising=False) + monkeypatch.delenv("MAGNETICS_ON_NETWORK", raising=False) + + +def _as_host(monkeypatch, fqdn: str) -> None: + monkeypatch.setattr(socket, "getfqdn", lambda: fqdn) + + +class TestOnClusterHost: + @pytest.mark.parametrize( + "fqdn", + [ + "omega.gat.com", # the configured address itself + "omega-a.gat.com", # load-balanced submit nodes — what you actually land on + "omega-b.gat.com", + "omega25.gat.com", # worker nodes (no dash) + ], + ) + def test_cluster_nodes_are_detected(self, monkeypatch, fqdn): + _as_host(monkeypatch, fqdn) + assert network.on_cluster_host("diiid") + + def test_site_gateway_is_not_the_cluster(self, monkeypatch): + # cybele is inside gat.com and reaches omega by ssh — from there the + # remote backend is correct and must NOT be downgraded. + _as_host(monkeypatch, "cybele.gat.com") + assert not network.on_cluster_host("diiid") + + def test_off_site_laptop_is_not_the_cluster(self, monkeypatch): + _as_host(monkeypatch, "my-laptop.local") + assert not network.on_cluster_host("diiid") + + def test_same_short_name_off_site_does_not_count(self, monkeypatch): + # on_site gating: a personal box named "omega" elsewhere is not the cluster + _as_host(monkeypatch, "omega.home.example") + assert not network.on_cluster_host("diiid") + + def test_env_override_forces_both_answers(self, monkeypatch): + _as_host(monkeypatch, "my-laptop.local") + monkeypatch.setenv("MAGNETICS_ON_CLUSTER", "1") + assert network.on_cluster_host("diiid") + monkeypatch.setenv("MAGNETICS_ON_CLUSTER", "0") + _as_host(monkeypatch, "omega-a.gat.com") + assert not network.on_cluster_host("diiid") + + def test_device_without_a_cluster_block_is_never_on_cluster(self, monkeypatch): + _as_host(monkeypatch, "omega-a.gat.com") + monkeypatch.setattr(network, "cluster_login", lambda d: {"host": None, "port": 22}) + assert not network.on_cluster_host("diiid") + + +class TestDispatchDowngrade: + """`backend="remote"` on the cluster must fetch in-process, never via ssh.""" + + def _spy(self, monkeypatch): + calls = {"remote": 0} + + def _boom(*a, **k): + calls["remote"] += 1 + raise AssertionError("run_remote must not be called from the cluster") + + from magnetics.data.fetch import remote as remote_run + + monkeypatch.setattr(remote_run, "run_remote", _boom) + return calls + + def test_on_cluster_remote_never_sshes_and_falls_back_to_mdsthin(self, monkeypatch, tmp_path): + """Without toksearch importable (the uv-tool install on omega), the pull + downgrades to mdsthin rather than dialing itself over ssh.""" + self._spy(monkeypatch) + monkeypatch.setattr(toksearch, "on_cluster_host", lambda device: True) + + seen = {} + + def _fake_mdsthin(shot, pointnames, **kw): + seen["called"] = True + return [] + + monkeypatch.setattr(toksearch, "_fetch_mdsthin", _fake_mdsthin) + + def _fake_write(out, *a, **k): + seen["wrote"] = out + # (got, missing) — one fetched channel, or the empty-result guard trips + return ["MPI66M307D"], [] + + monkeypatch.setattr(toksearch, "_write_h5", _fake_write) + + notes: list[str] = [] + toksearch.fetch_shot( + 184927, + analysis="rotating", + backend="remote", + out=str(tmp_path / "shot.h5"), + progress=lambda f, m: notes.append(m), + ) + assert seen.get("called"), "should have fetched in-process via mdsthin" + assert any("already on the cluster" in n for n in notes) + + def test_on_cluster_prefers_toksearch_via_the_cluster_interpreter(self, monkeypatch, tmp_path): + """The whole point: when the site conda env has toksearch, route the pull + through it (native PTDATA) rather than the slow mdsthin/mdsip fallback. + Regression for the fix where an on-cluster pull quietly used mdsip and ran + ~17x slower than the laptop's remote pull, which had toksearch all along.""" + self._spy(monkeypatch) # run_remote (ssh) must not fire + monkeypatch.setattr(toksearch, "on_cluster_host", lambda device: True) + + from magnetics.data.fetch import remote as remote_run + + monkeypatch.setattr(remote_run, "cluster_python_has_toksearch", lambda py: True) + # mdsthin must NOT be reached — that would be the slow fallback + monkeypatch.setattr( + toksearch, + "_fetch_mdsthin", + lambda *a, **k: pytest.fail("mdsthin used despite cluster toksearch"), + ) + hit = {} + + def _fake_on_cluster(shot, analysis, **kw): + hit["py"] = kw.get("python") + return str(tmp_path / "o.h5") + + monkeypatch.setattr(remote_run, "run_on_cluster", _fake_on_cluster) + + out = toksearch.fetch_shot( + 184927, + analysis="rotating", + backend="remote", + out=str(tmp_path / "o.h5"), + progress=lambda f, m: None, + ) + assert out == str(tmp_path / "o.h5") + # it passed the device's configured cluster interpreter through + assert hit.get("py") and "toksearch_env" in hit["py"] + + def test_off_cluster_remote_still_orchestrates_over_ssh(self, monkeypatch, tmp_path): + """The laptop path is untouched — remote still means remote.""" + monkeypatch.setattr(toksearch, "on_cluster_host", lambda device: False) + hit = {} + + def _fake_remote(shot, analysis, **kw): + hit["shot"] = shot + return str(tmp_path / "shot_184927.h5") + + from magnetics.data.fetch import remote as remote_run + + monkeypatch.setattr(remote_run, "run_remote", _fake_remote) + toksearch.fetch_shot( + 184927, + analysis="rotating", + backend="remote", + out=str(tmp_path / "shot_184927.h5"), + progress=lambda f, m: None, + ) + assert hit.get("shot") == 184927 + + +class TestRunOnClusterIsolation: + """run_on_cluster must expose ONLY the magnetics package to the cluster + interpreter — never the venv site-packages, whose numpy/scipy are built for a + different Python and shadow the conda env's copies (breaking `import + toksearch`). This is the bug that made the first on-cluster attempt crash.""" + + def test_pythonpath_is_an_isolated_package_symlink(self, monkeypatch, tmp_path): + from magnetics.data.fetch import remote as remote_run + + captured = {} + + def _fake_run(argv, env=None, **kw): + captured["env"] = env + captured["pp"] = env.get("PYTHONPATH", "") + # the symlink must exist AT RUN TIME (cleaned up in finally after) + entries = os.listdir(captured["pp"]) + captured["entries"] = entries + target = os.path.realpath(os.path.join(captured["pp"], "magnetics")) + captured["target"] = target + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(remote_run.subprocess, "run", _fake_run) + remote_run.run_on_cluster( + 184927, "rotating", python="/some/conda/env/bin/python", out=str(tmp_path / "o.h5") + ) + + pp = captured["pp"] + # exactly the staging dir — NOT our site-packages (no numpy/scipy alongside) + assert captured["entries"] == ["magnetics"], captured["entries"] + assert "site-packages" not in pp + # the symlink resolves to the real installed package + assert captured["target"] == os.path.realpath(str(remote_run.PKG_ROOT)) + # and the stage is removed afterward (finally: shutil.rmtree) + assert not os.path.exists(pp) + + def test_failure_raises_and_still_cleans_up(self, monkeypatch, tmp_path): + from magnetics.data.fetch import remote as remote_run + + seen = {} + + def _fail_run(argv, env=None, **kw): + seen["pp"] = env.get("PYTHONPATH") + return SimpleNamespace(returncode=1) + + monkeypatch.setattr(remote_run.subprocess, "run", _fail_run) + with pytest.raises(RuntimeError, match="on-cluster toksearch fetch failed"): + remote_run.run_on_cluster(1, "both", python="/x/py", out=str(tmp_path / "o.h5")) + assert not os.path.exists(seen["pp"]) # temp stage gone even on failure + + +class TestClusterPythonProbe: + """The probe decides toksearch-vs-mdsthin, so a false negative silently costs + ~5-7x on every on-cluster pull. The subprocess is the arbiter — don't + pre-reject interpreters it would happily run.""" + + def test_bare_command_name_is_resolved_through_path(self, monkeypatch): + from magnetics.data.fetch import remote as remote_run + + # `Path("python3").exists()` is False anywhere but the cwd; PATH lookup + # is what subprocess would do, so the probe must get a chance to run. + monkeypatch.setattr(remote_run.shutil, "which", lambda p: "/opt/conda/bin/" + p) + seen = {} + + def _probe(argv, **kw): + seen["argv"] = argv + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(remote_run.subprocess, "run", _probe) + assert remote_run.cluster_python_has_toksearch("python3") + assert seen["argv"][0] == "/opt/conda/bin/python3" + + def test_absolute_path_is_passed_through_untouched(self, monkeypatch): + from magnetics.data.fetch import remote as remote_run + + monkeypatch.setattr( + remote_run.shutil, "which", lambda p: pytest.fail("which() on an absolute path") + ) + seen = {} + + def _probe(argv, **kw): + seen["argv"] = argv + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(remote_run.subprocess, "run", _probe) + assert remote_run.cluster_python_has_toksearch("/fusion/env/bin/python") + assert seen["argv"][0] == "/fusion/env/bin/python" + + def test_missing_interpreter_is_false_not_a_crash(self, monkeypatch): + from magnetics.data.fetch import remote as remote_run + + monkeypatch.setattr(remote_run.shutil, "which", lambda p: None) + + def _enoent(argv, **kw): + raise FileNotFoundError(argv[0]) + + monkeypatch.setattr(remote_run.subprocess, "run", _enoent) + assert not remote_run.cluster_python_has_toksearch("nope-python") + + def test_none_short_circuits_without_spawning(self, monkeypatch): + from magnetics.data.fetch import remote as remote_run + + monkeypatch.setattr( + remote_run.subprocess, "run", lambda *a, **k: pytest.fail("spawned for None") + ) + assert not remote_run.cluster_python_has_toksearch(None) diff --git a/uv.lock b/uv.lock index 472a7d1..8160845 100644 --- a/uv.lock +++ b/uv.lock @@ -11,7 +11,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-07-03T05:43:21.277622Z" +exclude-newer = "2026-07-17T17:03:56.236499Z" exclude-newer-span = "P1W" [[package]] @@ -1218,7 +1218,7 @@ dev = [ { name = "nbconvert", specifier = ">=7.16" }, { name = "nbstripout", specifier = ">=0.9.1" }, { name = "pytest", specifier = ">=9.1.1" }, - { name = "ruff", specifier = ">=0.15.18" }, + { name = "ruff", specifier = ">=0.15.18,<0.16" }, { name = "ty", specifier = ">=0.0.52" }, ]