From cdc9169c742ce6fb02e30fced319833324e9f9d8 Mon Sep 17 00:00:00 2001 From: Laurent Guerard Date: Thu, 25 Jun 2026 10:10:32 +0200 Subject: [PATCH] =?UTF-8?q?feat(workflow):=20=E2=9C=A8=20prefetch=20the=20?= =?UTF-8?q?model=20on=20the=20submit=20host=20for=20offline=20nodes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compute nodes without internet made Cellpose fail at run time while downloading its weights ("Network is unreachable"). Add a fetch_model localrule that runs on the (networked) submit host and caches the model into shared $HOME before any GPU job; the segment rule now depends on it. Its own log file avoids racing the shared steps.log. Also fix the pixi [project] -> [workspace] deprecation warning and document the offline-cluster and empty-SLURM-log cases. Co-Authored-By: Claude Opus 4.8 --- docs/guide/snakemake.md | 4 +++- workflow/Snakefile | 5 +++++ workflow/pixi.toml | 2 +- workflow/rules/common.smk | 4 ++++ workflow/rules/segment.smk | 16 ++++++++++++++++ workflow/scripts/fetch_model.py | 21 +++++++++++++++++++++ 6 files changed, 50 insertions(+), 2 deletions(-) create mode 100644 workflow/scripts/fetch_model.py diff --git a/docs/guide/snakemake.md b/docs/guide/snakemake.md index 0cdd786..c633f02 100644 --- a/docs/guide/snakemake.md +++ b/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/.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) diff --git a/workflow/Snakefile b/workflow/Snakefile index d59a8ed..99eecde 100644 --- a/workflow/Snakefile +++ b/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", diff --git a/workflow/pixi.toml b/workflow/pixi.toml index 98c1c7e..e84d7f3 100644 --- a/workflow/pixi.toml +++ b/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"] diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 69f9fd2..eedba11 100644 --- a/workflow/rules/common.smk +++ b/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).""" diff --git a/workflow/rules/segment.smk b/workflow/rules/segment.smk index f3e9976..4713bcc 100644 --- a/workflow/rules/segment.smk +++ b/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: diff --git a/workflow/scripts/fetch_model.py b/workflow/scripts/fetch_model.py new file mode 100644 index 0000000..0b1ce4c --- /dev/null +++ b/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")