feat(workflow): prefetch segmentation model for offline GPU nodes

docs/guide/snakemake.md

@@ -223,10 +223,12 @@ prologue. The simplest path is a single shared env that the compute nodes see.
 | Symptom | Fix |
 |---------|-----|
 | `snakemake: command not found` | use `python -m snakemake` |
-| Segment jobs pend forever | wrong `slurm_partition`/`slurm_extra` GPU flag for your cluster |
+| Segment jobs pend forever | wrong `slurm_partition`/GPU request; on scicore use `gres: "gpu:1"` |
+| Segment dies, `Network is unreachable` | offline GPU nodes — the `fetch_model` localrule caches the model on the submit host first; if it still fails, your submit host has no network either (pre-download manually) |
 | `cellpose is not installed` in a job | the job's env lacks `patchworks[cellpose]` |
 | Reading the input fails | install the matching reader (`patchworks[imaris]`/`[bioio]` + a `bioio-*`) |
 | Out of GPU memory | smaller `tile_shape`, or `do_3D: false` |
+| A job fails with an empty SLURM log | read `logs/segment/<index>.log` (per tile) or `logs/steps.log` — the real traceback is there |
 | Very slow | confirm GPU is used (`nvidia-smi`); try 2-D or a lower `level` |
 
 ## How it works (for the curious)

workflow/Snakefile

@@ -19,6 +19,11 @@ include: "rules/segment.smk"
 include: "rules/merge.smk"
 
 
+# Runs on the submit host (has network) — never submitted to an offline GPU node.
+localrules:
+    fetch_model,
+
+
 rule all:
     input:
         f"{WORK}/labels.done",

workflow/pixi.toml

@@ -7,7 +7,7 @@
 # On a cluster, put this project on a shared filesystem the compute nodes can
 # read: the SLURM jobs re-launch snakemake from this env's interpreter.
 
-[project]
+[workspace]
 name = "patchworks-workflow"
 channels = ["conda-forge"]
 platforms = ["linux-64"]

workflow/rules/common.smk

@@ -20,6 +20,10 @@ STAGE_OK = f"{STAGE}.done"
 LOGS = f"{WORK}/logs"
 STEPLOG = f"{LOGS}/steps.log"
 
+# Marker that the segmentation model is cached locally. Produced by a local
+# rule (runs on the networked submit host) so offline GPU nodes never download.
+MODEL_OK = f"{WORK}/model.ready"
+
 
 def occupied_done(wildcards):
     """Per-tile markers for the occupied tiles (resolved after the checkpoint)."""

workflow/rules/segment.smk

@@ -1,5 +1,20 @@
 # Plan tiles (checkpoint) and segment each tile on a GPU.
 
+
+rule fetch_model:
+    """Cache the segmentation model on the (networked) submit host.
+
+    Declared local (see ``localrules`` in the Snakefile) so it never runs on an
+    offline GPU node — Cellpose downloads its weights here, into shared $HOME.
+    """
+    output:
+        touch(MODEL_OK),
+    log:
+        f"{LOGS}/fetch_model.log",
+    script:
+        "../scripts/fetch_model.py"
+
+
 checkpoint prepare:
     input:
         IMAGE_OK,
@@ -18,6 +33,7 @@ rule segment:
         tiles=TILES,
         stage=STAGE_OK,
         image=IMAGE_OK,
+        model=MODEL_OK,
     output:
         f"{WORK}/seg/{{index}}.done",
     log:

workflow/scripts/fetch_model.py

@@ -0,0 +1,21 @@
+"""Snakemake (local) script: cache the segmentation model before segmenting.
+
+Runs on the submit host, which has network access, so the offline GPU nodes
+never try to download Cellpose weights at run time (they read the shared
+``$HOME/.cellpose`` cache instead).
+"""
+
+from _pw import start_log
+
+start_log(snakemake.log[0])  # noqa: F821
+cfg = snakemake.config  # noqa: F821
+
+if cfg.get("method", "cellpose") == "cellpose":
+    # _get_model downloads + caches the weights keyed by (model, gpu).
+    from patchworks.plugins.cellpose import _get_model
+
+    model = cfg["cellpose"]["model"]
+    _get_model({"model": model, "gpu": False})
+    print(f"[patchworks] cached segmentation model: {model}")
+else:
+    print(f"[patchworks] method={cfg.get('method')!r}; no model to prefetch")