feat(workflow): per-step and per-tile log files

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: <partition>-<duration> 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

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

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"

workflow/rules/merge.smk

@@ -5,5 +5,7 @@ rule merge:
         occupied_done,
     output:
         touch(f"{WORK}/labels.done"),
+    log:
+        STEPLOG,
     script:
         "../scripts/merge.py"

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"

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.
 

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(

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")

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"])

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"]