Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 10 additions & 5 deletions src/magnetics/_slcontour/omfit_compat.py
Original file line number Diff line number Diff line change
Expand Up @@ -348,16 +348,21 @@ def load_sensor_table(device="DIII-D"):
return sensor_geometry(device)


def load_wall(device="DIII-D"):
def load_wall(device="DIII-D", shot=None):
"""Return the ``(r, z)`` first-wall outline arrays from the device JSON.

Returns ``(None, None)`` if no device file or wall section is found.
Resolves the shot-segmented ``wall`` (``{segments: [...]}``) through the device
SSOT resolver :func:`magnetics.data.devices.feature_at`, which also tolerates a
legacy flat ``{r, z}`` record. A ``None`` shot (device-level / plotting) takes
the latest segment, since the wall changes rarely between campaigns.

Returns ``(None, None)`` if no device file or wall segment is found.
"""
try:
dev = load_device(device)
except (FileNotFoundError, OSError):
return None, None
wall = dev.get("wall")
if not wall or "r" not in wall or "z" not in wall:
seg = _devices.feature_at(dev, "wall", shot if shot is not None else 2_000_000)
if not seg or "r" not in seg or "z" not in seg:
return None, None
return np.array(wall["r"], dtype=float), np.array(wall["z"], dtype=float)
return np.array(seg["r"], dtype=float), np.array(seg["z"], dtype=float)
19 changes: 19 additions & 0 deletions src/magnetics/data/fetch/remote.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@

from __future__ import annotations

import re
import shlex
import subprocess
import sys
Expand All @@ -63,6 +64,23 @@
# (its activate.d scripts set nothing PTDATA needs; direct is ~4-5s faster per pull).
DEFAULT_PYTHON = "/fusion/projects/codes/conda/omega/envs_public/toksearch_env/bin/python"

# ``remote_dir`` is interpolated into remote shell commands (``mkdir -p``, ``cd``) and
# rsync remote specs, so it must not carry spaces or shell metacharacters. Allow a
# leading ``~`` (home) then POSIX-safe path chars only. This is an external-input /
# security boundary (the value can originate from the GUI/API), so a bad value is
# rejected loudly rather than quoted-around or silently accepted.
_SAFE_REMOTE_DIR = re.compile(r"^~?[\w./\-]+$")


def _validate_remote_dir(remote_dir: str) -> str:
if not _SAFE_REMOTE_DIR.match(remote_dir) or remote_dir.startswith("-"):
raise ValueError(
f"unsafe remote_dir {remote_dir!r}: only '~', letters, digits, '_', '.', "
"'-' and '/' are allowed (no spaces or shell metacharacters), and relative "
"paths must not start with '-'"
)
return remote_dir


def _log(msg: str) -> None:
sys.stderr.write(msg + "\n")
Expand Down Expand Up @@ -98,6 +116,7 @@ def run_remote(
the network, no hop needed); pass an explicit host[:port] to force one, or an
empty string to force none (e.g. an ssh-config alias that carries its own
ProxyJump)."""
remote_dir = _validate_remote_dir(remote_dir)
login = cluster_login(device)
host = host or login["host"] or DEFAULT_HOST
port = login["port"] # cluster SSH port (22 unless the device overrides it)
Expand Down
10 changes: 8 additions & 2 deletions src/magnetics/data/fetch/toksearch.py
Original file line number Diff line number Diff line change
Expand Up @@ -626,8 +626,14 @@ def _fetch_toksearch(shot, pointnames, *, tmin, tmax, stride, progress, tree_sig

# Resolve EFIT-tree signals: first candidate key with usable data wins. No
# decimation (EFIT time bases are already coarse), window-trim only.
for name, keys in tree_keys.items():
ch = Channel(name, ok=False, error="no data")
# Every requested tree signal is recorded, fetched or not: when MdsSignal is
# unavailable tree_keys is empty, so iterate the full request set and mark the
# unfetched ones missing (parity with the mdsthin path — never silently drop).
for name in tree_signals:
keys = tree_keys.get(name, [])
ch = Channel(
name, ok=False, error="MdsSignal unavailable" if MdsSignal is None else "no data"
)
for key, tree, node in keys:
sig = None if key in errors else (rec[key] if key in rec.keys() else None)
if not sig or "data" not in sig:
Expand Down
6 changes: 5 additions & 1 deletion src/magnetics/service/mock.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@

from __future__ import annotations

import zlib

import numpy as np

from ..data import diiid
Expand Down Expand Up @@ -126,7 +128,9 @@ def qs_fit_frames(machine: str, params: dict) -> list[tuple[float, dict]]:
# ── spectrogram (Burgess) — coarse → fine, mode that locks or rotates ────────
def spectrogram_frames(machine: str, params: dict) -> list[tuple[float, dict]]:
p = _profile(machine)
rng = np.random.default_rng(abs(hash(machine)) % 2**32)
# Stable across process restarts: the builtin str hash() is per-process salted
# (PYTHONHASHSEED), which would make the mock non-deterministic between runs.
rng = np.random.default_rng(zlib.crc32(machine.encode()))
t0, t1 = p["t_ms"]
resolutions = [(60, 45), (120, 90), (240, 180)] # (nt, nf), → real 360×180
frames = []
Expand Down
2 changes: 1 addition & 1 deletion src/magnetics/service/nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -1234,7 +1234,7 @@ def _sensor_map_rz(shot, params=None) -> dict:

from .._slcontour.omfit_compat import load_wall

r_wall, z_wall = load_wall(device)
r_wall, z_wall = load_wall(device, int(shot))

series = []
for c in channels:
Expand Down
39 changes: 39 additions & 0 deletions tests/test_data_fetch_hardening.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
"""Regression tests for the Crucible-confirmed data-layer fixes.

P2 load_wall read the old flat wall schema after the segmented migration → (None,None)
P5 remote_dir flowed unquoted / as an option into cluster shell commands
"""

from __future__ import annotations

import pytest

from magnetics._slcontour.omfit_compat import load_wall
from magnetics.data.fetch import remote


# ── P2: load_wall resolves the segmented wall schema ────────────────────────
def test_load_wall_resolves_segmented_schema():
"""After the wall was migrated to {segments:[...]}, load_wall read the old flat
keys and returned (None, None), dropping the vessel outline from the sensor map."""
r, z = load_wall("DIII-D", 184927)
assert r is not None and z is not None
assert len(r) == len(z) > 0
# a None shot (device-level / plotting) takes the latest segment, not (None, None)
r0, z0 = load_wall("DIII-D")
assert r0 is not None and len(r0) > 0


# ── P5: remote_dir shell/option-injection boundary ──────────────────────────
@pytest.mark.parametrize("safe", ["~/magnetics_fetch", "/scratch/u/fetch", "rel/dir", "~/a.b-c_d"])
def test_validate_remote_dir_accepts_safe_paths(safe):
assert remote._validate_remote_dir(safe) == safe


@pytest.mark.parametrize(
"bad",
["~/my fetch", "$(touch x)", "a;b", "a&&b", "`id`", "a|b", "a>b", "-tmp", "--help"],
)
def test_validate_remote_dir_rejects_metacharacters_and_options(bad):
with pytest.raises(ValueError):
remote._validate_remote_dir(bad)
Loading