Cache subprocess rbc invocations via RBC_STYXCACHE_DIR (#312)

* Wire styxcache into setup_runner for subprocess rbc invocations

tests/integration/test_all.py spawns `rbc all` via subprocess.run;
that fresh process hits setup_runner() rather than the pytest autouse
fixture, so it was running the full pipeline uncached. Cache diagnose
confirmed 0 new entries during warm runs and revealed `setup` of
those two fixtures took ~81 min combined.

setup_runner now wraps the configured runner with styxcache's
CachingRunner when RBC_STYXCACHE_DIR is set, matching the behavior
we already had in tests/conftest.py. The wrapping logic (attribute
proxy + policy) now lives in rbc.core.niwrap.maybe_wrap_with_cache;
tests/conftest.py imports and reuses it.

styxcache moves from dev to runtime dependency since rbc.core now
imports it unconditionally. Drop the try/except fallback and the
deptry DEP004 ignore.

* Clear RBC_STYXCACHE_DIR for setup_runner isinstance tests, evict corrupt gz cache entries

- tests/unit/test_niwrap.py::TestSetupRunner checks the concrete runner
  type; ensure the env var is unset so setup_runner returns the unwrapped
  runner.
- Prepare styx cache now runs gzip -t on every .styxcache.stdout.gz and
  .styxcache.stderr.gz, evicting any entry with a truncated stream.
  Prompted by an EOFError from a half-written gz in an older cache entry
  that survived into a warm run.

* Extend corrupt-gz eviction to all .gz files in cache shards

Earlier scan only looked at .styxcache.stdout.gz / .stderr.gz, but nibabel
also reads cached .nii.gz outputs directly from shard paths. A truncated
cached BOLD blew up single_session_qc with the same EOFError signature.

* Move rbc scratch dirs off cache shards via independent scratch root

Two problems with the prior generate_exec_folder pattern:

1. It hijacked the niwrap runner's data_dir + uid + execution_counter,
   forcing _CacheProxyingRunner to forward those attributes through to
   the base runner on every access.
2. Several rbc sites then wrote sibling files via `input_file.parent /
   ...`. When input_file was served from a styxcache shard (which is
   the norm under caching), those writes landed inside the shard,
   racing other workers' reads and silently truncating cached files.
   The test_single_session_qc EOFError was a cached .nii.gz truncated
   this way.

generate_exec_folder now owns its own root (RBC_SCRATCH_DIR override,
else a process-unique tempfile.mkdtemp). Every touched call site was
migrated off `input_file.parent / ...` to `generate_exec_folder(...) /
...`. Sites audited:
  - core/functional/resampling.py:107, 201
  - core/functional/nuisance.py:112
  - core/longitudinal/transform.py:145
  - core/longitudinal/freesurfer.py:142
  - core/metrics/reho.py:191
  - core/metrics/alff.py:259
  - core/metrics/standardization.py:74
  - core/metrics/timeseries.py:159

distortion.py's topup.parent is a read-only lookup into the tool's own
output set (no write), left as-is.

Unit tests for generate_exec_folder rewritten for the new semantics:
no more uid/counter-based naming, just unique scratch dirs.

* Unify niwrap data_dir and rbc scratch under a single RBC_WORK_DIR root

Previously generate_exec_folder had its own RBC_SCRATCH_DIR while
niwrap's data_dir defaulted to the system tmp. One knob is cleaner:
RBC_WORK_DIR now parents both (niwrap/, scratch/). CI points it at
JOB_TMP so the per-run cleanup reclaims both trees in one sweep.

* CLI --tmp-dir now parents rbc scratch too, not just niwrap data_dir

Author: nx10

.github/workflows/test_full.yaml

@@ -50,6 +50,27 @@ jobs:
         # Purge stale staging dirs from crashed previous runs (safe: atomic-rename commits)
         rm -rf "$STYXCACHE_DIR/.incoming" 2>/dev/null || true
 
+        # Drop entries that contain any truncated .gz (styxcache's own
+        # stdout/stderr streams or cached .nii.gz outputs). Atomic commits
+        # should prevent these, but crashes before commit have been observed.
+        corrupt=0
+        while IFS= read -r entry; do
+          bad=0
+          for gz in $(find "$entry" -type f -name '*.gz'); do
+            if ! gzip -t "$gz" 2>/dev/null; then
+              bad=1
+              break
+            fi
+          done
+          if [ "$bad" = "1" ]; then
+            rm -rf "$entry"
+            corrupt=$((corrupt + 1))
+          fi
+        done < <(find "$STYXCACHE_DIR" -mindepth 2 -maxdepth 2 -type d 2>/dev/null)
+        if [ "$corrupt" -gt 0 ]; then
+          echo "evicted $corrupt entries with corrupt gzip streams"
+        fi
+
         size_kb=$(du -sk "$STYXCACHE_DIR" 2>/dev/null | awk '{print $1}')
         size_kb=${size_kb:-0}
         cap_kb=$(( STYXCACHE_MAX_GB * 1024 * 1024 ))
@@ -85,6 +106,9 @@ jobs:
         export PYTHONPYCACHEPREFIX="$JOB_TMP/__pycache__"
         export COVERAGE_FILE="$JOB_TMP/.coverage"
         export RBC_STYXCACHE_DIR="$STYXCACHE_DIR"
+        # Pin rbc's scratch + niwrap data_dir under JOB_TMP so the Cleanup step
+        # reclaims them automatically. Default is a system tmpdir otherwise.
+        export RBC_WORK_DIR="$JOB_TMP/rbc_work"
 
         uv run pytest \
           -n 8 \

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
   "numpy>=2.4.2",
   "polars>=1.38.1",
   "scipy>=1.17.0",
+  "styxcache>=0.2.0,<0.3",
   "styxpodman>=0.1.1",
   "tqdm>=4.67.3"
 ]
@@ -28,8 +29,7 @@ dev = [
   "pytest-cov>=7.0.0",
   "ruff>=0.8.1",
   "deptry>=0.23.0",
-  "pytest-xdist[psutil]>=3.8.0",
-  "styxcache>=0.2.0,<0.3"
+  "pytest-xdist[psutil]>=3.8.0"
 ]
 docs = ["pdoc>=15.0.0"]
 
@@ -44,11 +44,6 @@ markers = [
   "full_pipeline: End-to-end workflow tests"
 ]
 
-[tool.deptry.per_rule_ignores]
-# styxcache is optional — imported via try/except so the dependency is only
-# needed on CI where the cache is wired up. src/rbc runs fine without it.
-DEP004 = ["styxcache"]
-
 [tool.coverage.report]
 omit = ["src/rbc_resources/"]
 

src/rbc/core/common.py

@@ -7,29 +7,20 @@
 
 from __future__ import annotations
 
-import contextlib
 from typing import TYPE_CHECKING
 
 import nibabel as nib
+import styxcache
 from niwrap import afni
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator, Sequence
+    from collections.abc import Sequence
     from pathlib import Path
 
 from rbc.core.fileops import file_tmp_copy
 from rbc.core.nifti import strip_afni_volatile_metadata
 from rbc.core.niwrap import generate_exec_folder
 
-try:
-    from styxcache import bypass as _styxcache_bypass
-except ImportError:  # styxcache is optional — only used on CI
-
-    @contextlib.contextmanager
-    def _styxcache_bypass() -> Iterator[None]:
-        yield
-
-
 __all__ = ["deoblique_and_reorient", "merge_3d_to_4d", "split_4d"]
 
 
@@ -52,11 +43,11 @@ def deoblique_and_reorient(
         AFNI 3dresample outputs (use ``.out_file`` for the reoriented image).
     """
     with file_tmp_copy(in_file) as tmp_file:
-        # 3drefit mutates in place, and styxcache 0.2.0 does not replay
+        # 3drefit mutates in place, and styxcache does not replay
         # mutable-input mutations on cache hits. Bypass it so it always runs,
         # then strip AFNI's non-deterministic extension (timestamps + random
         # UUID) so the downstream cached 3dresample call keys on stable bytes.
-        with _styxcache_bypass():
+        with styxcache.bypass():
             afni.v_3drefit(in_file=tmp_file, deoblique=True)
         strip_afni_volatile_metadata(tmp_file)
         return afni.v_3dresample(

src/rbc/core/functional/nuisance.py

@@ -109,7 +109,7 @@ def bandpass_regressor_file(
         f_data[~freq_mask] = 0.0
         filtered[:, col] = np.real(ifft(f_data))[:n_tp]
 
-    out_path = regressor_file.parent / "regressors_filtered.1D"
+    out_path = generate_exec_folder("regressors_filtered") / "regressors_filtered.1D"
     with out_path.open("w") as f:
         for line in header_lines:
             f.write(line)

src/rbc/core/functional/resampling.py

@@ -104,7 +104,7 @@ def apply_motion_transforms(
         )
         transformed_vols.append(result.output.output_image_outfile)
 
-    out_path = transformed_vols[0].parent / "preproc_bold.nii.gz"
+    out_path = generate_exec_folder("preproc_bold_merge") / "preproc_bold.nii.gz"
     merged = merge_3d_to_4d(transformed_vols, out_path)
 
     # antsApplyTransforms writes the reference image's pixdim into the output
@@ -198,7 +198,10 @@ def resample_bold_to_template(
         )
         transformed_vols.append(result.output.output_image_outfile)
 
-    out_path = transformed_vols[0].parent / "bold_to_template_resampled.nii.gz"
+    out_path = (
+        generate_exec_folder("bold_to_template_merge")
+        / "bold_to_template_resampled.nii.gz"
+    )
     merged = merge_3d_to_4d(transformed_vols, out_path)
 
     # antsApplyTransforms writes the reference (template) image's pixdim into

src/rbc/core/longitudinal/freesurfer.py

@@ -16,6 +16,7 @@
 from niwrap import freesurfer
 
 from rbc.core.fsl2itk import mat_to_itk
+from rbc.core.niwrap import generate_exec_folder
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -139,7 +140,7 @@ def fs_to_itk_xfm(
     for ses, source, in_xfm in zip(sessions, sources, in_xfms, strict=True):
         fsl_fname = in_xfm.with_suffix(".mat").name
         lta = freesurfer.lta_convert(in_lta=in_xfm, out_fsl=fsl_fname)
-        itk_path = in_xfm.parent / itk_filename(sub, ses)
+        itk_path = generate_exec_folder("itk_xfm") / itk_filename(sub, ses)
         mat_to_itk(
             mat=lta.root / fsl_fname,
             reference=reference,

src/rbc/core/longitudinal/transform.py

@@ -8,6 +8,7 @@
 
 from rbc.core.common import merge_3d_to_4d, split_4d
 from rbc.core.functional.resampling import _restore_tr
+from rbc.core.niwrap import generate_exec_folder
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -142,7 +143,10 @@ def _transform_4d_chunked(in_file: Path, template: Path, xfm: Path) -> Path:
         )
         transformed_vols.append(result.output.output_image_outfile)
 
-    out_path = transformed_vols[0].parent / "bold_to_longitudinal.nii.gz"
+    out_path = (
+        generate_exec_folder("bold_to_longitudinal_merge")
+        / "bold_to_longitudinal.nii.gz"
+    )
     merged = merge_3d_to_4d(transformed_vols, out_path)
     _restore_tr(merged, in_file)
     return merged

src/rbc/core/metrics/alff.py

@@ -17,6 +17,7 @@
 import numpy as np
 
 from rbc.core.nifti import Volume
+from rbc.core.niwrap import generate_exec_folder
 
 if TYPE_CHECKING:
     from typing import Literal
@@ -256,7 +257,8 @@ def compute_alff(
 
     stem = in_file.name.split(".nii")[0]
     if out_file is None:
-        alff_path = in_file.parent / f"{stem}_alff.nii.gz"
+        out_dir = generate_exec_folder("alff")
+        alff_path = out_dir / f"{stem}_alff.nii.gz"
     else:
         alff_path = Path(out_file)
 

src/rbc/core/metrics/reho.py

@@ -15,6 +15,7 @@
 from scipy.stats import rankdata
 
 from rbc.core.nifti import Volume
+from rbc.core.niwrap import generate_exec_folder
 
 if TYPE_CHECKING:
     from typing import Literal
@@ -188,7 +189,7 @@ def compute_reho(
 
     if out_file is None:
         stem = in_file.name.split(".nii")[0]
-        out_file = in_file.parent / f"{stem}_reho.nii.gz"
+        out_file = generate_exec_folder("reho") / f"{stem}_reho.nii.gz"
     out_file = Path(out_file)
 
     bold.derive(reho_map).save(out_file)

src/rbc/core/metrics/standardization.py

@@ -11,6 +11,8 @@
 
 import numpy as np
 
+from rbc.core.niwrap import generate_exec_folder
+
 
 def zscore(data: np.ndarray, mask: np.ndarray) -> np.ndarray:
     """Z-score a 3D map within a brain mask.
@@ -71,7 +73,7 @@ def compute_zscore(
 
     if out_file is None:
         stem = in_file.name.split(".nii")[0]
-        out_file = in_file.parent / f"{stem}_zscored.nii.gz"
+        out_file = generate_exec_folder("zscore") / f"{stem}_zscored.nii.gz"
     out_file = Path(out_file)
 
     nib.nifti1.Nifti1Image(zscored, img.affine, img.header).to_filename(str(out_file))