From 6a647a4f1838ad056e619faa7bdf1ac60757a27a Mon Sep 17 00:00:00 2001 From: Laurent Guerard Date: Thu, 25 Jun 2026 09:46:45 +0200 Subject: [PATCH] =?UTF-8?q?feat(workflow):=20=E2=9C=A8=20per-step=20and=20?= =?UTF-8?q?per-tile=20log=20files?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a shared logs/steps.log for the sequential CPU steps (convert/prepare/merge) and a logs/segment/.log per GPU tile. A line-buffered tee in _pw.start_log mirrors stdout/stderr and library logging into the file, so a Python traceback (or output up to an OOM/ walltime kill) is preserved even when the SLURM job log is empty — which is exactly the blind spot we hit debugging segment failures. Also revert the segment mem_mb bump: sacct showed MaxRSS ~1 GB, so the failure was a Python error (exit 1), not OOM. Co-Authored-By: Claude Opus 4.8 --- workflow/profile/slurm/config.yaml | 14 ++++---- workflow/rules/common.smk | 6 ++++ workflow/rules/convert.smk | 2 ++ workflow/rules/merge.smk | 2 ++ workflow/rules/segment.smk | 4 +++ workflow/scripts/_pw.py | 51 ++++++++++++++++++++++++++++++ workflow/scripts/convert.py | 3 ++ workflow/scripts/merge.py | 3 +- workflow/scripts/prepare_tiles.py | 3 +- workflow/scripts/segment_tile.py | 3 +- 10 files changed, 82 insertions(+), 9 deletions(-) diff --git a/workflow/profile/slurm/config.yaml b/workflow/profile/slurm/config.yaml index 7038ed0..22373b8 100644 --- a/workflow/profile/slurm/config.yaml +++ b/workflow/profile/slurm/config.yaml @@ -17,7 +17,7 @@ keep-going: true # Defaults for every rule (CPU jobs). default-resources: - slurm_partition: "cpu" + slurm_partition: "scicore" # slurm_account: "my_account" mem_mb: 16000 cpus_per_task: 4 @@ -32,12 +32,14 @@ set-resources: prepare: mem_mb: 32000 runtime: 120 - segment: # one GPU per tile — this is what spreads Cellpose across GPUs - slurm_partition: "gpu" - slurm_extra: "'--gres=gpu:1'" - mem_mb: 32000 + segment: + # one GPU per tile — this is what spreads Cellpose across GPUs + slurm_partition: "rtx4090" + gpu: 1 + qos: "rtx4090-6hours" # scicore: - QOS + mem_mb: 32000 # plenty — a tile used ~1G; the failure was a Python error cpus_per_task: 4 - runtime: 120 + runtime: 360 # 6 hours — must match the QOS, NOT 120 (=2h → killed early) merge: mem_mb: 128000 cpus_per_task: 8 diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 03f6c3e..69f9fd2 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -15,6 +15,12 @@ STAGE = f"{WORK}/stage.zarr" STAGE_OK = f"{STAGE}.done" +# Logs: one shared file for the sequential CPU steps (convert/prepare/merge), +# one file per tile for the GPU segment jobs. +LOGS = f"{WORK}/logs" +STEPLOG = f"{LOGS}/steps.log" + + def occupied_done(wildcards): """Per-tile markers for the occupied tiles (resolved after the checkpoint).""" tiles = checkpoints.prepare.get().output.tiles diff --git a/workflow/rules/convert.smk b/workflow/rules/convert.smk index 70f05f1..f2a08d9 100644 --- a/workflow/rules/convert.smk +++ b/workflow/rules/convert.smk @@ -4,5 +4,7 @@ rule convert: output: # marker file inside the store; existence => skip re-conversion. IMAGE_OK, + log: + STEPLOG, script: "../scripts/convert.py" diff --git a/workflow/rules/merge.smk b/workflow/rules/merge.smk index 249e171..5353ec2 100644 --- a/workflow/rules/merge.smk +++ b/workflow/rules/merge.smk @@ -5,5 +5,7 @@ rule merge: occupied_done, output: touch(f"{WORK}/labels.done"), + log: + STEPLOG, script: "../scripts/merge.py" diff --git a/workflow/rules/segment.smk b/workflow/rules/segment.smk index 2111868..f3e9976 100644 --- a/workflow/rules/segment.smk +++ b/workflow/rules/segment.smk @@ -6,6 +6,8 @@ checkpoint prepare: output: tiles=TILES, stage=touch(STAGE_OK), + log: + STEPLOG, script: "../scripts/prepare_tiles.py" @@ -18,5 +20,7 @@ rule segment: image=IMAGE_OK, output: f"{WORK}/seg/{{index}}.done", + log: + f"{LOGS}/segment/{{index}}.log", script: "../scripts/segment_tile.py" diff --git a/workflow/scripts/_pw.py b/workflow/scripts/_pw.py index 39222bf..08c83a0 100644 --- a/workflow/scripts/_pw.py +++ b/workflow/scripts/_pw.py @@ -8,11 +8,62 @@ from __future__ import annotations import json +import logging +import sys from pathlib import Path from patchworks import load_ome_zarr +class _Tee: + """Write to several streams at once (e.g. the SLURM log and a file).""" + + def __init__(self, *streams): + self._streams = streams + + def write(self, data): + for stream in self._streams: + stream.write(data) + + def flush(self): + for stream in self._streams: + stream.flush() + + +def start_log(path, *, append=True): + """Tee stdout/stderr (and logging) into ``path``. + + Captures prints, tracebacks and library logging into a file in the work + directory, independent of the (often empty) SLURM job log. Line-buffered, + so output up to a crash or OOM kill is preserved. + + Parameters + ---------- + path : str or Path + Log file to write. Parent directories are created. + append : bool, optional + Append to an existing log (keep retry history) instead of truncating. + Default True. + + Returns + ------- + TextIO + The open log file (kept open for the lifetime of the process). + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + handle = open(path, "a" if append else "w", buffering=1) + sys.stdout = _Tee(sys.__stdout__, handle) + sys.stderr = _Tee(sys.__stderr__, handle) + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + handlers=[logging.StreamHandler(sys.stdout)], + force=True, + ) + return handle + + def open_image(work_dir, channel, level): """Open the converted image for segmentation. diff --git a/workflow/scripts/convert.py b/workflow/scripts/convert.py index 19e3533..d9ab62e 100644 --- a/workflow/scripts/convert.py +++ b/workflow/scripts/convert.py @@ -8,6 +8,9 @@ from patchworks.plugins.ome_zarr import to_ome_zarr +from _pw import start_log + +start_log(snakemake.log[0]) # noqa: F821 cfg = snakemake.config # noqa: F821 (injected by Snakemake) chunks = cfg.get("convert_chunks") to_ome_zarr( diff --git a/workflow/scripts/merge.py b/workflow/scripts/merge.py index 4e3e712..c69aac3 100644 --- a/workflow/scripts/merge.py +++ b/workflow/scripts/merge.py @@ -11,8 +11,9 @@ from patchworks import merge_tile_labels from patchworks.plugins.ome_zarr import write_labels -from _pw import stage_path +from _pw import stage_path, start_log +start_log(snakemake.log[0]) # noqa: F821 cfg = snakemake.config # noqa: F821 work_dir = cfg["work_dir"] image_store = str(Path(work_dir) / "image.zarr") diff --git a/workflow/scripts/prepare_tiles.py b/workflow/scripts/prepare_tiles.py index ea49062..abe1652 100644 --- a/workflow/scripts/prepare_tiles.py +++ b/workflow/scripts/prepare_tiles.py @@ -11,8 +11,9 @@ spatial_tiles, ) -from _pw import open_image, stage_path +from _pw import open_image, stage_path, start_log +start_log(snakemake.log[0]) # noqa: F821 cfg = snakemake.config # noqa: F821 work_dir = cfg["work_dir"] image = open_image(work_dir, cfg["channel"], cfg["level"]) diff --git a/workflow/scripts/segment_tile.py b/workflow/scripts/segment_tile.py index 0398d52..be6f158 100644 --- a/workflow/scripts/segment_tile.py +++ b/workflow/scripts/segment_tile.py @@ -6,8 +6,9 @@ from patchworks import stage_tile -from _pw import build_fn, load_tiles_json, open_image, stage_path +from _pw import build_fn, load_tiles_json, open_image, stage_path, start_log +start_log(snakemake.log[0]) # noqa: F821 cfg = snakemake.config # noqa: F821 index = int(snakemake.wildcards.index) # noqa: F821 work_dir = cfg["work_dir"]