From 92f3455439525edc8efb662cda7b857fc754dec0 Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Tue, 16 Jun 2026 16:23:47 -0700
Subject: [PATCH 1/7] feat(nextflow): integrate virtual-stain into mantis-v2
 pipeline
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Recreate the Nextflow virtual-stain integration against the current
in-process cytoland CLI (post-#267), replacing the #259 approach which
was built on the old `viscy predict` + `--copy` temp-zarr flow.

- New `nextflow/modules/virtual_stain.nf`: init → preprocess → fan-out
  per-position GPU prediction. Per-position work is a single
  `biahub virtual-stain --cluster debug` call (no temp zarr, no --copy).
- `viscy preprocess` runs over the whole input plate to compute the
  normalization statistics the model reads via read_norm_meta.
- Both `biahub virtual-stain` and `viscy` run under biahub's `stain`
  extra (cytoland → viscy-utils provides the `viscy` console script).
- Add a `gpu` process label (queue=gpu) to nextflow.config.
- Wire virtual_stain_wf after reconstruct in mantis-v2.nf.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 nextflow/mantis-v2.nf             |  33 +++++--
 nextflow/modules/virtual_stain.nf | 154 ++++++++++++++++++++++++++++++
 nextflow/nextflow.config          |   8 ++
 3 files changed, 186 insertions(+), 9 deletions(-)
 create mode 100644 nextflow/modules/virtual_stain.nf

diff --git a/nextflow/mantis-v2.nf b/nextflow/mantis-v2.nf
index f1a8b0a7..b5963e5f 100644
--- a/nextflow/mantis-v2.nf
+++ b/nextflow/mantis-v2.nf
@@ -16,23 +16,25 @@ nextflow.enable.dsl = 2
 //  (in some pipelines the first step converts raw input to zarr). To reorder
 //  steps, change where a step reads from here; the modules stay untouched.
 //
-//  Flat-field → deskew → reconstruct is wired today. The remaining steps
-//  (virtual-stain, track, assemble) arrive with their own PRs — follow the
+//  Flat-field → deskew → reconstruct → virtual-stain is wired today. The
+//  remaining steps (track, assemble) arrive with their own PRs — follow the
 //  chaining below for the pattern.
 // ---------------------------------------------------------------------------
 
-params.input              = null   // raw source — may not be a zarr store
-params.output             = null   // output directory for all step zarrs
-params.deskew_config      = null
-params.flat_field_config  = null
+params.input = null   // raw source — may not be a zarr store
+params.output = null   // output directory for all step zarrs
+params.deskew_config = null
+params.flat_field_config = null
 params.reconstruct_config = null
-params.biahub_project     = null
-params.max_positions      = 0
+params.virtual_stain_config = null
+params.biahub_project = null
+params.max_positions = 0
 
 include { collect_positions; dataset_name } from './modules/common'
 include { deskew_wf } from './modules/deskew'
 include { flat_field_wf } from './modules/flat_field'
 include { reconstruct_wf } from './modules/reconstruct'
+include { virtual_stain_wf } from './modules/virtual_stain'
 
 // Output directory layout for the reconstruction steps — single source of
 // truth. Each entry is a subdirectory under params.output where that step
@@ -44,7 +46,7 @@ DIRECTORY_LAYOUT = [
     flat_field    : '0-flatfield',
     deskew        : '1-deskew',
     reconstruct   : '2-reconstruct',
-    // virtual_stain : '3-virtual-stain',
+    virtual_stain : '3-virtual-stain',
     // track         : '4-track',
     // assemble      : '5-assemble',
 ]
@@ -56,6 +58,7 @@ workflow {
     if (!params.flat_field_config)  error "Provide --flat_field_config"
     if (!params.deskew_config)      error "Provide --deskew_config"
     if (!params.reconstruct_config) error "Provide --reconstruct_config"
+    if (!params.virtual_stain_config) error "Provide --virtual_stain_config"
 
     def ds  = dataset_name()
     def out = params.output
@@ -91,4 +94,16 @@ workflow {
     reconstruct_output  = "${out}/${DIRECTORY_LAYOUT.reconstruct}/${ds}.zarr"
 
     reconstruct_done = reconstruct_wf(all_positions, reconstruct_input, reconstruct_output, params.reconstruct_config, reconstruct_trigger)
+
+    // ----- Virtual stain ----------------------------------------------------
+    // Virtual staining runs cytoland (VisCy) prediction on the reconstructed
+    // output and waits on reconstruct_done. A `viscy preprocess` step inside the
+    // subworkflow computes the normalization statistics the model needs; which
+    // source/target channels are used is set by the virtual-stain config, not
+    // here.
+    virtual_stain_trigger = reconstruct_done.done
+    virtual_stain_input   = reconstruct_output
+    virtual_stain_output  = "${out}/${DIRECTORY_LAYOUT.virtual_stain}/${ds}.zarr"
+
+    virtual_stain_done = virtual_stain_wf(all_positions, virtual_stain_input, virtual_stain_output, params.virtual_stain_config, virtual_stain_trigger)
 }
diff --git a/nextflow/modules/virtual_stain.nf b/nextflow/modules/virtual_stain.nf
new file mode 100644
index 00000000..372d7d00
--- /dev/null
+++ b/nextflow/modules/virtual_stain.nf
@@ -0,0 +1,154 @@
+// Virtual-stain subworkflow: init + preprocess → fan-out run × N positions.
+//
+// This subworkflow is PATH-AGNOSTIC. Callers pass the input zarr, output zarr,
+// and config explicitly; the module has no idea where it sits in the pipeline
+// directory layout. The orchestrating pipeline (see mantis-v2.nf) owns the
+// layout and the order of steps; this module just virtually stains whatever
+// it's handed.
+//
+// Since PR #267 the `biahub virtual-stain` CLI runs cytoland (modular VisCy)
+// prediction IN-PROCESS, so per-position work is a single `biahub virtual-stain
+// --cluster debug` call — no temp per-position zarr and no `--copy` merge step
+// (the old #259 flow). `--cluster debug` makes submitit's DebugExecutor run the
+// work synchronously inside the Nextflow task; Nextflow already handles
+// per-position fan-out and resource scheduling, so the CLI must NOT submit its
+// own SLURM jobs. See:
+// examples/submitit_debug_nextflow/2026-05-27-submitit-debug-nextflow-concerns.md
+//
+// Three-phase pattern:
+// 1. init_virtual_stain: validates the config, creates the output plate with
+//    the predicted channels, emits RESOURCES:
+// 2. run_virtual_stain_preprocess: `viscy preprocess` over the whole input
+//    plate. virtual_stain_position reads precomputed normalization statistics
+//    from the input store (viscy_data.read_norm_meta) and errors if they are
+//    missing, so this must run before fan-out. NOTE: this MUTATES the input
+//    store by writing normalization metadata into it.
+// 3. run_virtual_stain: per-position GPU prediction using RESOURCES.
+//
+// Both `biahub virtual-stain` and `viscy` live in biahub's optional `stain`
+// extra (cytoland → viscy-utils provides the `viscy` console script), so the
+// tasks here run in that extra's environment rather than the plain biahub env.
+
+include { parse_resources; slurm_logs; slurm_log_dir } from './common'
+
+// Command prefix for tools that require biahub's `stain` extra. Both
+// `biahub virtual-stain` (it imports cytoland) and `viscy preprocess` need it.
+// Falls back to the bare tool on the active environment when biahub_project is
+// unset (assumes that env already has the stain extra installed).
+def stain_cmd(tool) {
+    return params.biahub_project ?
+        "uv run --project ${params.biahub_project} --extra stain ${tool}" : tool
+}
+
+
+process init_virtual_stain {
+    label 'cpu_local'
+
+    input:
+    val input_zarr
+    val output_zarr
+    val config
+    val trigger
+
+    output:
+    stdout
+
+    script:
+    """
+    mkdir -p "${slurm_log_dir('virtual_stain')}"
+    ${stain_cmd('biahub')} virtual-stain --init \
+        -i "${input_zarr}"/*/*/* \
+        -o "${output_zarr}" \
+        -c "${config}"
+    """
+}
+
+process run_virtual_stain_preprocess {
+    label 'cpu'
+    clusterOptions { slurm_logs('virtual_stain') }
+    cpus 16
+    memory { "${64 * task.attempt} GB" }
+    time '1h'
+    maxRetries 1
+    errorStrategy 'retry'
+
+    input:
+    val input_zarr
+    val trigger
+
+    output:
+    val true
+
+    script:
+    """
+    ${stain_cmd('viscy')} preprocess \
+        --data_path "${input_zarr}" \
+        --channel_names -1 \
+        --num_workers ${task.cpus} \
+        --block_size 32
+    """
+}
+
+process run_virtual_stain {
+    tag "${position}"
+    label 'gpu'
+    clusterOptions { "--gres=gpu:1 " + slurm_logs('virtual_stain') }
+    maxForks 30
+    cpus { meta.cpus }
+    memory { "${meta.mem_gb} GB" }
+    time { task.attempt == 1 ? '8h' : '12h' }
+    maxRetries 2
+    errorStrategy 'retry'
+
+    input:
+    tuple val(position), val(meta)
+    val input_zarr
+    val output_zarr
+    val config
+
+    output:
+    val position
+
+    script:
+    """
+    ${stain_cmd('biahub')} virtual-stain --cluster debug \
+        -i "${input_zarr}/${position}" \
+        -o "${output_zarr}" \
+        -c "${config}"
+    """
+}
+
+
+// take:
+//   positions    collected channel of position keys (e.g. ['A/1/0', 'B/1/0'])
+//   input_zarr   path to the input plate.zarr (reconstruct output)
+//   output_zarr  path to the virtual-stain output plate.zarr
+//   config       path to the virtual-stain (viscy predict) settings YAML
+//   prev_done    gating channel — virtual stain starts once this emits
+workflow virtual_stain_wf {
+    take:
+    positions
+    input_zarr
+    output_zarr
+    config
+    prev_done
+
+    main:
+    init_out = init_virtual_stain(input_zarr, output_zarr, config, prev_done.map { 'done' })
+    resources = init_out.map { parse_resources(it) }
+
+    // Preprocess the whole plate in parallel with init; both gate the fan-out.
+    vs_preprocess = run_virtual_stain_preprocess(input_zarr, prev_done.map { 'done' })
+
+    ready = resources.combine(vs_preprocess)
+
+    pos_meta = positions
+        .flatMap { it }
+        .combine(ready)
+        .map { pos, meta, preprocess_done -> [pos, meta] }
+
+    vs_done = run_virtual_stain(pos_meta, input_zarr, output_zarr, config) | collect
+
+    emit:
+    done = vs_done
+}
diff --git a/nextflow/nextflow.config b/nextflow/nextflow.config
index ec7c0468..1e91ed59 100644
--- a/nextflow/nextflow.config
+++ b/nextflow/nextflow.config
@@ -32,6 +32,11 @@ process {
         memory = '32 GB'
         time   = '2h'
     }
+    withLabel: 'gpu' {
+        cpus   = 16
+        memory = '64 GB'
+        time   = '8h'
+    }
 }
 
 profiles {
@@ -56,6 +61,9 @@ profiles {
             withLabel: 'cpu' {
                 queue = 'cpu'
             }
+            withLabel: 'gpu' {
+                queue = 'gpu'
+            }
         }
     }
 }

From 07ae4efe1dc3f38e249c257b7bf13c81962e0743 Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Tue, 16 Jun 2026 16:35:05 -0700
Subject: [PATCH 2/7] fix biahub virtual-stain mem request

---
 biahub/virtual_stain.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/biahub/virtual_stain.py b/biahub/virtual_stain.py
index 8210805e..f6b03945 100644
--- a/biahub/virtual_stain.py
+++ b/biahub/virtual_stain.py
@@ -349,7 +349,7 @@ def virtual_stain(
     slurm_args = {
         "slurm_job_name": "virtual-stain",
         "slurm_gres": "gpu:1",
-        "slurm_mem_per_cpu": f"{gb_ram}G",
+        "slurm_mem": f"{gb_ram}G",
         "slurm_cpus_per_task": num_cpus,
         "slurm_array_parallelism": 20,  # process up to 20 positions at a time
         "slurm_time": slurm_time,

From cc1361fb590415777e27d943c440c9c6a95f3336 Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Tue, 16 Jun 2026 17:10:03 -0700
Subject: [PATCH 3/7] refactor: unify per-position resource requests across
 step CLIs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Every step CLI (deskew, flat-field, apply-inv-tf, virtual-stain) now emits
its per-position resource request through a single shared helper,
biahub.cli.utils.echo_resources, as a JSON payload:

    RESOURCES:{"cpus": 16, "mem_gb": 64, "time_min": 480}

The same (cpus, total mem_gb, time_min) values feed both the CLI's own
slurm_* submission args and, via parse_resources in nextflow, the Nextflow
per-position task directives — so the SLURM fan-out and the Nextflow fan-out
can no longer request different resources.

- Add echo_resources() and parse the JSON in common.nf::parse_resources,
  which now returns time_min alongside cpus/mem_gb.
- Emit memory as a TOTAL (mem_gb) and request it via slurm_mem rather than
  slurm_mem_per_cpu, fixing virtual-stain's per-cpu over-request.
- Wire each CLI's slurm_time from time_min; drop the hardcoded `time`
  literals in the run_* processes for `meta.time_min * task.attempt`.
- virtual-stain computes time_min before the init_only return so --init
  emits it.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 biahub/apply_inverse_transfer_function.py |  9 +++++---
 biahub/cli/utils.py                       | 28 +++++++++++++++++++++++
 biahub/deskew.py                          | 13 +++++++----
 biahub/flat_field_correction.py           | 13 +++++++----
 biahub/virtual_stain.py                   | 24 ++++++++++---------
 nextflow/modules/common.nf                |  8 +++++--
 nextflow/modules/deskew.nf                |  2 +-
 nextflow/modules/flat_field.nf            |  2 +-
 nextflow/modules/reconstruct.nf           |  2 +-
 nextflow/modules/virtual_stain.nf         |  4 ++--
 10 files changed, 76 insertions(+), 29 deletions(-)

diff --git a/biahub/apply_inverse_transfer_function.py b/biahub/apply_inverse_transfer_function.py
index 813df05b..6aed201b 100644
--- a/biahub/apply_inverse_transfer_function.py
+++ b/biahub/apply_inverse_transfer_function.py
@@ -24,6 +24,7 @@
     sbatch_to_submitit,
 )
 from biahub.cli.utils import (
+    echo_resources,
     get_submitit_cluster,
     yaml_to_model,
 )
@@ -112,7 +113,9 @@ def apply_inverse_transfer_function(
 
     max_num_cpus = 16
     num_cpus, mem_per_cpu = wo_estimate_resources(list(input_shape), settings, max_num_cpus)
-    click.echo(f"RESOURCES:{num_cpus} {num_cpus * mem_per_cpu}")
+    mem_gb = num_cpus * mem_per_cpu
+    time_min = 360
+    echo_resources(num_cpus, mem_gb, time_min)
 
     if init_only:
         click.echo(
@@ -130,9 +133,9 @@ def apply_inverse_transfer_function(
     # examples/submitit_debug_nextflow/2026-05-27-submitit-debug-nextflow-concerns.md
     slurm_args = {
         "slurm_job_name": "apply-inverse-transfer-function",
-        "slurm_mem_per_cpu": f"{mem_per_cpu}G",
+        "slurm_mem": f"{mem_gb}G",
         "slurm_cpus_per_task": num_cpus,
-        "slurm_time": 60,
+        "slurm_time": time_min,
         "slurm_partition": "cpu",
     }
 
diff --git a/biahub/cli/utils.py b/biahub/cli/utils.py
index a252093e..01451eb7 100644
--- a/biahub/cli/utils.py
+++ b/biahub/cli/utils.py
@@ -1,9 +1,11 @@
+import json
 import logging
 import os
 
 from pathlib import Path
 from typing import Literal
 
+import click
 import numpy as np
 import yaml
 
@@ -14,6 +16,32 @@
 logger = logging.getLogger(__name__)
 
 
+def echo_resources(num_cpus: int, mem_gb: int, time_min: int) -> None:
+    """Emit the per-position resource request consumed by the Nextflow pipeline.
+
+    Every step CLI calls this from its ``--init`` path so there is a single
+    source of truth for per-position CPU, memory, and wall-clock time. The
+    Nextflow ``init_*`` process captures this line on stdout and
+    ``parse_resources`` (``nextflow/modules/common.nf``) reads the JSON payload
+    to set the per-position task's ``cpus``/``memory``/``time`` directives. The
+    same values also feed the CLI's own ``slurm_*`` submission args, so the
+    SLURM fan-out and the Nextflow fan-out request identical resources.
+
+    A single JSON payload keeps the contract order-independent and extensible
+    (new fields can be added without breaking the positional parsing).
+
+    Parameters
+    ----------
+    num_cpus : int
+        CPUs per position.
+    mem_gb : int
+        TOTAL memory per position in GB (not per-CPU).
+    time_min : int
+        Wall-clock budget per position in minutes.
+    """
+    click.echo("RESOURCES:" + json.dumps({"cpus": num_cpus, "mem_gb": mem_gb, "time_min": time_min}))
+
+
 def get_submitit_cluster(
     local: bool = False,
     cluster: str | None = None,
diff --git a/biahub/deskew.py b/biahub/deskew.py
index ddbfee21..68119933 100644
--- a/biahub/deskew.py
+++ b/biahub/deskew.py
@@ -27,6 +27,7 @@
     sbatch_to_submitit,
 )
 from biahub.cli.utils import (
+    echo_resources,
     estimate_resources,
     get_submitit_cluster,
     resolve_ome_zarr_version,
@@ -673,8 +674,12 @@ def deskew(
     _warn_pixel_size_mismatch(settings, input_position_dirpaths[0])
     input_shape, _ = _init_output_plate(input_position_dirpaths, output_dirpath, settings)
 
-    num_cpus, gb_ram = estimate_resources(shape=input_shape, ram_multiplier=8, max_num_cpus=16)
-    click.echo(f"RESOURCES:{num_cpus} {num_cpus * gb_ram}")
+    num_cpus, gb_ram_per_cpu = estimate_resources(
+        shape=input_shape, ram_multiplier=8, max_num_cpus=16
+    )
+    mem_gb = num_cpus * gb_ram_per_cpu
+    time_min = 60
+    echo_resources(num_cpus, mem_gb, time_min)
 
     if init_only:
         click.echo(f"Initialized {output_dirpath} ({len(input_position_dirpaths)} positions)")
@@ -694,10 +699,10 @@ def deskew(
 
     slurm_args = {
         "slurm_job_name": "deskew",
-        "slurm_mem_per_cpu": f"{gb_ram}G",
+        "slurm_mem": f"{mem_gb}G",
         "slurm_cpus_per_task": num_cpus,
         "slurm_array_parallelism": 100,  # process up to 100 positions at a time
-        "slurm_time": 60,
+        "slurm_time": time_min,
         "slurm_partition": "preempted",
     }
     if sbatch_filepath:
diff --git a/biahub/flat_field_correction.py b/biahub/flat_field_correction.py
index 996aa44f..62f81072 100644
--- a/biahub/flat_field_correction.py
+++ b/biahub/flat_field_correction.py
@@ -20,6 +20,7 @@
     sbatch_to_submitit,
 )
 from biahub.cli.utils import (
+    echo_resources,
     estimate_resources,
     get_submitit_cluster,
     resolve_ome_zarr_version,
@@ -161,8 +162,12 @@ def flat_field(
     )
 
     T, C, Z, Y, X = input_shape
-    num_cpus, gb_ram = estimate_resources(shape=input_shape, ram_multiplier=8, max_num_cpus=16)
-    click.echo(f"RESOURCES:{num_cpus} {num_cpus * gb_ram}")
+    num_cpus, gb_ram_per_cpu = estimate_resources(
+        shape=input_shape, ram_multiplier=8, max_num_cpus=16
+    )
+    mem_gb = num_cpus * gb_ram_per_cpu
+    time_min = 60
+    echo_resources(num_cpus, mem_gb, time_min)
 
     if init_only:
         click.echo(f"Initialized {output_dirpath} ({len(input_position_dirpaths)} positions)")
@@ -178,10 +183,10 @@ def flat_field(
 
     slurm_args = {
         "slurm_job_name": "flat-field",
-        "slurm_mem_per_cpu": f"{gb_ram}G",
+        "slurm_mem": f"{mem_gb}G",
         "slurm_cpus_per_task": num_cpus,
         "slurm_array_parallelism": 100,
-        "slurm_time": 360,
+        "slurm_time": time_min,
         "slurm_partition": "cpu",
     }
 
diff --git a/biahub/virtual_stain.py b/biahub/virtual_stain.py
index f6b03945..c75a29b8 100644
--- a/biahub/virtual_stain.py
+++ b/biahub/virtual_stain.py
@@ -24,6 +24,7 @@
     sbatch_to_submitit,
 )
 from biahub.cli.utils import (
+    echo_resources,
     get_submitit_cluster,
     resolve_ome_zarr_version,
 )
@@ -329,8 +330,16 @@ def virtual_stain(
 
     # Timepoints are processed sequentially on a single GPU, so resource needs
     # are independent of dataset size.
-    num_cpus, gb_ram = 16, 64
-    click.echo(f"RESOURCES:{num_cpus} {gb_ram}")
+    num_cpus, mem_gb = 16, 64
+    # Generous wall-clock budget (minutes), assuming median TTA (4 rotations).
+    # Each timepoint runs ~Z sliding windows along Z. Measured ~1.4 s/window
+    # with TTA on this model; budget ~5 s/window (~3-4x margin for slower GPUs,
+    # larger FOVs, and compute-bound runs) with a 60-minute floor. Computed
+    # before the init_only return so --init emits it for the Nextflow pipeline.
+    T, Z = input_shape[0], input_shape[2]
+    seconds_per_window = 5
+    time_min = int(np.ceil(max(60, T * Z * seconds_per_window / 60)))
+    echo_resources(num_cpus, mem_gb, time_min)
 
     if init_only:
         click.echo(f"Initialized {output_dirpath} ({len(input_position_dirpaths)} positions)")
@@ -338,21 +347,14 @@ def virtual_stain(
 
     output_position_paths = utils.get_output_paths(input_position_dirpaths, output_dirpath)
 
-    T, Z = input_shape[0], input_shape[2]
-    # Generous wall-clock budget (minutes), assuming median TTA (4 rotations).
-    # Each timepoint runs ~Z sliding windows along Z. Measured ~1.4 s/window
-    # with TTA on this model; budget ~5 s/window (~3-4x margin for slower GPUs,
-    # larger FOVs, and compute-bound runs) with a 60-minute floor.
-    seconds_per_window = 5
-    slurm_time = int(np.ceil(max(60, T * Z * seconds_per_window / 60)))
     # Prepare SLURM arguments
     slurm_args = {
         "slurm_job_name": "virtual-stain",
         "slurm_gres": "gpu:1",
-        "slurm_mem": f"{gb_ram}G",
+        "slurm_mem": f"{mem_gb}G",
         "slurm_cpus_per_task": num_cpus,
         "slurm_array_parallelism": 20,  # process up to 20 positions at a time
-        "slurm_time": slurm_time,
+        "slurm_time": time_min,
         "slurm_partition": "gpu",
     }
 
diff --git a/nextflow/modules/common.nf b/nextflow/modules/common.nf
index c00689b6..6d235d23 100644
--- a/nextflow/modules/common.nf
+++ b/nextflow/modules/common.nf
@@ -8,8 +8,12 @@ def parse_resources(stdout_text, prefix = 'RESOURCES:') {
     if (!matching) {
         error "Expected a '${prefix}' line in command output but none was found. The underlying CLI may have failed."
     }
-    def parts = matching.last().replace(prefix, '').trim().split(/\s+/)
-    return [cpus: parts[0].toInteger(), mem_gb: parts[1].toInteger()]
+    // The CLI emits a JSON payload (see biahub.cli.utils.echo_resources): cpus,
+    // total mem_gb, and per-position time_min. Parsing JSON keeps the contract
+    // order-independent and extensible.
+    def payload = matching.last().replace(prefix, '').trim()
+    def res = new groovy.json.JsonSlurper().parseText(payload)
+    return [cpus: res.cpus as int, mem_gb: res.mem_gb as int, time_min: res.time_min as int]
 }
 
 def slurm_log_dir(step_name) {
diff --git a/nextflow/modules/deskew.nf b/nextflow/modules/deskew.nf
index 1910a574..88518d1c 100644
--- a/nextflow/modules/deskew.nf
+++ b/nextflow/modules/deskew.nf
@@ -46,7 +46,7 @@ process run_deskew {
     maxForks 30
     cpus { meta.cpus }
     memory { "${meta.mem_gb} GB" }
-    time { task.attempt == 1 ? '1h' : '2h' }
+    time { "${meta.time_min * task.attempt} min" }
     maxRetries 1
     errorStrategy 'retry'
 
diff --git a/nextflow/modules/flat_field.nf b/nextflow/modules/flat_field.nf
index 242f8344..e33ac4c1 100644
--- a/nextflow/modules/flat_field.nf
+++ b/nextflow/modules/flat_field.nf
@@ -45,7 +45,7 @@ process run_flat_field {
     maxForks 30
     cpus { meta.cpus }
     memory { "${meta.mem_gb} GB" }
-    time '1h'
+    time { "${meta.time_min * task.attempt} min" }
     maxRetries 1
     errorStrategy 'retry'
 
diff --git a/nextflow/modules/reconstruct.nf b/nextflow/modules/reconstruct.nf
index 0bffabeb..25e53643 100644
--- a/nextflow/modules/reconstruct.nf
+++ b/nextflow/modules/reconstruct.nf
@@ -92,7 +92,7 @@ process run_apply_inv_tf {
     maxForks 30
     cpus { meta.cpus }
     memory { "${meta.mem_gb} GB" }
-    time '6h'
+    time { "${meta.time_min * task.attempt} min" }
     maxRetries 1
     errorStrategy 'retry'
 
diff --git a/nextflow/modules/virtual_stain.nf b/nextflow/modules/virtual_stain.nf
index 372d7d00..ca58fbc3 100644
--- a/nextflow/modules/virtual_stain.nf
+++ b/nextflow/modules/virtual_stain.nf
@@ -96,8 +96,8 @@ process run_virtual_stain {
     maxForks 30
     cpus { meta.cpus }
     memory { "${meta.mem_gb} GB" }
-    time { task.attempt == 1 ? '8h' : '12h' }
-    maxRetries 2
+    time { "${meta.time_min * task.attempt} min" }
+    maxRetries 1
     errorStrategy 'retry'
 
     input:

From c2bc54d874d1deaba18880543b9c14ee4aea941a Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Tue, 16 Jun 2026 17:29:19 -0700
Subject: [PATCH 4/7] payload bugfix

---
 biahub/cli/utils.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/biahub/cli/utils.py b/biahub/cli/utils.py
index 01451eb7..29659605 100644
--- a/biahub/cli/utils.py
+++ b/biahub/cli/utils.py
@@ -39,7 +39,10 @@ def echo_resources(num_cpus: int, mem_gb: int, time_min: int) -> None:
     time_min : int
         Wall-clock budget per position in minutes.
     """
-    click.echo("RESOURCES:" + json.dumps({"cpus": num_cpus, "mem_gb": mem_gb, "time_min": time_min}))
+    # Coerce to plain int: estimators may return numpy integers, which json
+    # cannot serialize.
+    payload = {"cpus": int(num_cpus), "mem_gb": int(mem_gb), "time_min": int(time_min)}
+    click.echo("RESOURCES:" + json.dumps(payload))
 
 
 def get_submitit_cluster(

From 34a4e482e0930aeaf33f900dd06394263f9c9ddd Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Tue, 16 Jun 2026 19:40:14 -0700
Subject: [PATCH 5/7] debug viscy predict

---
 nextflow/modules/virtual_stain.nf | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/nextflow/modules/virtual_stain.nf b/nextflow/modules/virtual_stain.nf
index ca58fbc3..9c166f95 100644
--- a/nextflow/modules/virtual_stain.nf
+++ b/nextflow/modules/virtual_stain.nf
@@ -79,13 +79,26 @@ process run_virtual_stain_preprocess {
     output:
     val true
 
+    // `--trainer.logger false` disables the viscy CLI's default WandbLogger.
+    // The VisCy LightningCLI sets trainer.logger to a lazy WandbLogger for every
+    // subcommand (viscy_utils/cli.py); preprocess needs no logger, and W&B isn't
+    // in the `stain` extra, so instantiating it fails with a missing-wandb error.
+    //
+    // `unset SLURM_NTASKS`: sbatch exports the submit environment, so when this
+    // pipeline is launched from inside a SLURM allocation, the submit shell's
+    // SLURM_NTASKS leaks into the job. Lightning's Trainer then auto-detects a
+    // SLURMEnvironment and rejects SLURM_NTASKS>1 (it expects --ntasks-per-node).
+    // preprocess is a single-process CPU job, so clearing it lets Lightning fall
+    // back to LightningEnvironment. The dataloader uses --num_workers, not tasks.
     script:
     """
+    unset SLURM_NTASKS
     ${stain_cmd('viscy')} preprocess \
         --data_path "${input_zarr}" \
         --channel_names -1 \
         --num_workers ${task.cpus} \
-        --block_size 32
+        --block_size 32 \
+        --trainer.logger false
     """
 }
 

From 1cc22a103634213ed0cdc409f65db975f096c31c Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Tue, 16 Jun 2026 19:49:36 -0700
Subject: [PATCH 6/7] refactor: rename time_min to time_minutes for clarity

Avoids confusion with "minimum time"; renames the resource payload key
and Nextflow meta field consistently across both languages.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 biahub/apply_inverse_transfer_function.py | 6 +++---
 biahub/cli/utils.py                       | 6 +++---
 biahub/deskew.py                          | 6 +++---
 biahub/flat_field_correction.py           | 6 +++---
 biahub/virtual_stain.py                   | 6 +++---
 nextflow/modules/common.nf                | 4 ++--
 nextflow/modules/deskew.nf                | 2 +-
 nextflow/modules/flat_field.nf            | 2 +-
 nextflow/modules/reconstruct.nf           | 2 +-
 nextflow/modules/virtual_stain.nf         | 2 +-
 10 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/biahub/apply_inverse_transfer_function.py b/biahub/apply_inverse_transfer_function.py
index 6aed201b..902aadeb 100644
--- a/biahub/apply_inverse_transfer_function.py
+++ b/biahub/apply_inverse_transfer_function.py
@@ -114,8 +114,8 @@ def apply_inverse_transfer_function(
     max_num_cpus = 16
     num_cpus, mem_per_cpu = wo_estimate_resources(list(input_shape), settings, max_num_cpus)
     mem_gb = num_cpus * mem_per_cpu
-    time_min = 360
-    echo_resources(num_cpus, mem_gb, time_min)
+    time_minutes = 360
+    echo_resources(num_cpus, mem_gb, time_minutes)
 
     if init_only:
         click.echo(
@@ -135,7 +135,7 @@ def apply_inverse_transfer_function(
         "slurm_job_name": "apply-inverse-transfer-function",
         "slurm_mem": f"{mem_gb}G",
         "slurm_cpus_per_task": num_cpus,
-        "slurm_time": time_min,
+        "slurm_time": time_minutes,
         "slurm_partition": "cpu",
     }
 
diff --git a/biahub/cli/utils.py b/biahub/cli/utils.py
index 29659605..011d6f0c 100644
--- a/biahub/cli/utils.py
+++ b/biahub/cli/utils.py
@@ -16,7 +16,7 @@
 logger = logging.getLogger(__name__)
 
 
-def echo_resources(num_cpus: int, mem_gb: int, time_min: int) -> None:
+def echo_resources(num_cpus: int, mem_gb: int, time_minutes: int) -> None:
     """Emit the per-position resource request consumed by the Nextflow pipeline.
 
     Every step CLI calls this from its ``--init`` path so there is a single
@@ -36,12 +36,12 @@ def echo_resources(num_cpus: int, mem_gb: int, time_min: int) -> None:
         CPUs per position.
     mem_gb : int
         TOTAL memory per position in GB (not per-CPU).
-    time_min : int
+    time_minutes : int
         Wall-clock budget per position in minutes.
     """
     # Coerce to plain int: estimators may return numpy integers, which json
     # cannot serialize.
-    payload = {"cpus": int(num_cpus), "mem_gb": int(mem_gb), "time_min": int(time_min)}
+    payload = {"cpus": int(num_cpus), "mem_gb": int(mem_gb), "time_minutes": int(time_minutes)}
     click.echo("RESOURCES:" + json.dumps(payload))
 
 
diff --git a/biahub/deskew.py b/biahub/deskew.py
index 68119933..77dfa7b3 100644
--- a/biahub/deskew.py
+++ b/biahub/deskew.py
@@ -678,8 +678,8 @@ def deskew(
         shape=input_shape, ram_multiplier=8, max_num_cpus=16
     )
     mem_gb = num_cpus * gb_ram_per_cpu
-    time_min = 60
-    echo_resources(num_cpus, mem_gb, time_min)
+    time_minutes = 60
+    echo_resources(num_cpus, mem_gb, time_minutes)
 
     if init_only:
         click.echo(f"Initialized {output_dirpath} ({len(input_position_dirpaths)} positions)")
@@ -702,7 +702,7 @@ def deskew(
         "slurm_mem": f"{mem_gb}G",
         "slurm_cpus_per_task": num_cpus,
         "slurm_array_parallelism": 100,  # process up to 100 positions at a time
-        "slurm_time": time_min,
+        "slurm_time": time_minutes,
         "slurm_partition": "preempted",
     }
     if sbatch_filepath:
diff --git a/biahub/flat_field_correction.py b/biahub/flat_field_correction.py
index 62f81072..600b81e2 100644
--- a/biahub/flat_field_correction.py
+++ b/biahub/flat_field_correction.py
@@ -166,8 +166,8 @@ def flat_field(
         shape=input_shape, ram_multiplier=8, max_num_cpus=16
     )
     mem_gb = num_cpus * gb_ram_per_cpu
-    time_min = 60
-    echo_resources(num_cpus, mem_gb, time_min)
+    time_minutes = 60
+    echo_resources(num_cpus, mem_gb, time_minutes)
 
     if init_only:
         click.echo(f"Initialized {output_dirpath} ({len(input_position_dirpaths)} positions)")
@@ -186,7 +186,7 @@ def flat_field(
         "slurm_mem": f"{mem_gb}G",
         "slurm_cpus_per_task": num_cpus,
         "slurm_array_parallelism": 100,
-        "slurm_time": time_min,
+        "slurm_time": time_minutes,
         "slurm_partition": "cpu",
     }
 
diff --git a/biahub/virtual_stain.py b/biahub/virtual_stain.py
index c75a29b8..5d96ddd5 100644
--- a/biahub/virtual_stain.py
+++ b/biahub/virtual_stain.py
@@ -338,8 +338,8 @@ def virtual_stain(
     # before the init_only return so --init emits it for the Nextflow pipeline.
     T, Z = input_shape[0], input_shape[2]
     seconds_per_window = 5
-    time_min = int(np.ceil(max(60, T * Z * seconds_per_window / 60)))
-    echo_resources(num_cpus, mem_gb, time_min)
+    time_minutes = int(np.ceil(max(60, T * Z * seconds_per_window / 60)))
+    echo_resources(num_cpus, mem_gb, time_minutes)
 
     if init_only:
         click.echo(f"Initialized {output_dirpath} ({len(input_position_dirpaths)} positions)")
@@ -354,7 +354,7 @@ def virtual_stain(
         "slurm_mem": f"{mem_gb}G",
         "slurm_cpus_per_task": num_cpus,
         "slurm_array_parallelism": 20,  # process up to 20 positions at a time
-        "slurm_time": time_min,
+        "slurm_time": time_minutes,
         "slurm_partition": "gpu",
     }
 
diff --git a/nextflow/modules/common.nf b/nextflow/modules/common.nf
index 6d235d23..a873f3b5 100644
--- a/nextflow/modules/common.nf
+++ b/nextflow/modules/common.nf
@@ -9,11 +9,11 @@ def parse_resources(stdout_text, prefix = 'RESOURCES:') {
         error "Expected a '${prefix}' line in command output but none was found. The underlying CLI may have failed."
     }
     // The CLI emits a JSON payload (see biahub.cli.utils.echo_resources): cpus,
-    // total mem_gb, and per-position time_min. Parsing JSON keeps the contract
+    // total mem_gb, and per-position time_minutes. Parsing JSON keeps the contract
     // order-independent and extensible.
     def payload = matching.last().replace(prefix, '').trim()
     def res = new groovy.json.JsonSlurper().parseText(payload)
-    return [cpus: res.cpus as int, mem_gb: res.mem_gb as int, time_min: res.time_min as int]
+    return [cpus: res.cpus as int, mem_gb: res.mem_gb as int, time_minutes: res.time_minutes as int]
 }
 
 def slurm_log_dir(step_name) {
diff --git a/nextflow/modules/deskew.nf b/nextflow/modules/deskew.nf
index 88518d1c..a53ea984 100644
--- a/nextflow/modules/deskew.nf
+++ b/nextflow/modules/deskew.nf
@@ -46,7 +46,7 @@ process run_deskew {
     maxForks 30
     cpus { meta.cpus }
     memory { "${meta.mem_gb} GB" }
-    time { "${meta.time_min * task.attempt} min" }
+    time { "${meta.time_minutes * task.attempt} min" }
     maxRetries 1
     errorStrategy 'retry'
 
diff --git a/nextflow/modules/flat_field.nf b/nextflow/modules/flat_field.nf
index e33ac4c1..50bdd57b 100644
--- a/nextflow/modules/flat_field.nf
+++ b/nextflow/modules/flat_field.nf
@@ -45,7 +45,7 @@ process run_flat_field {
     maxForks 30
     cpus { meta.cpus }
     memory { "${meta.mem_gb} GB" }
-    time { "${meta.time_min * task.attempt} min" }
+    time { "${meta.time_minutes * task.attempt} min" }
     maxRetries 1
     errorStrategy 'retry'
 
diff --git a/nextflow/modules/reconstruct.nf b/nextflow/modules/reconstruct.nf
index 25e53643..a8b080ff 100644
--- a/nextflow/modules/reconstruct.nf
+++ b/nextflow/modules/reconstruct.nf
@@ -92,7 +92,7 @@ process run_apply_inv_tf {
     maxForks 30
     cpus { meta.cpus }
     memory { "${meta.mem_gb} GB" }
-    time { "${meta.time_min * task.attempt} min" }
+    time { "${meta.time_minutes * task.attempt} min" }
     maxRetries 1
     errorStrategy 'retry'
 
diff --git a/nextflow/modules/virtual_stain.nf b/nextflow/modules/virtual_stain.nf
index 9c166f95..30094c17 100644
--- a/nextflow/modules/virtual_stain.nf
+++ b/nextflow/modules/virtual_stain.nf
@@ -109,7 +109,7 @@ process run_virtual_stain {
     maxForks 30
     cpus { meta.cpus }
     memory { "${meta.mem_gb} GB" }
-    time { "${meta.time_min * task.attempt} min" }
+    time { "${meta.time_minutes * task.attempt} min" }
     maxRetries 1
     errorStrategy 'retry'
 

From 0046ee39067500f812acc2015e97c577d5b1e189 Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Tue, 23 Jun 2026 11:48:01 -0700
Subject: [PATCH 7/7] update iohub dep

---
 pyproject.toml |  6 +-----
 uv.lock        | 11 ++++++++---
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 39a7bdbb..ec6d15ac 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,7 +20,7 @@ classifiers = [
 
 dependencies = [
   "click",
-  "iohub>=0.3.7",
+  "iohub>=0.3.8",
   "matplotlib",
   "natsort",
   "numpy",
@@ -84,10 +84,6 @@ index-strategy = "unsafe-best-match"
 # Pin cytoland to the commit from VisCy PR #465 (non-square TTA and helpers).
 # Revert to the PyPI release once that PR is merged and published.
 cytoland = { git = "https://github.com/mehta-lab/VisCy", subdirectory = "applications/cytoland", rev = "7fb3bc1e78c0e0bd26e1845caba860f76223ba02" }
-# Pin iohub to the branch that stores per-step provenance as top-level
-# zattrs keys (e.g. biahub-flat_field, biahub-deskew). Revert to the PyPI
-# release once those changes are merged and published.
-iohub = { git = "https://github.com/czbiohub-sf/iohub.git", branch = "feat/flat-extra-metadata" }
 
 [tool.hatch.metadata]
 allow-direct-references = true
diff --git a/uv.lock b/uv.lock
index d50bb4ac..b68fc221 100644
--- a/uv.lock
+++ b/uv.lock
@@ -366,7 +366,7 @@ requires-dist = [
     { name = "dask", extras = ["array"] },
     { name = "humanize" },
     { name = "imageio-ffmpeg" },
-    { name = "iohub", git = "https://github.com/czbiohub-sf/iohub.git?branch=feat%2Fflat-extra-metadata" },
+    { name = "iohub", specifier = ">=0.3.8" },
     { name = "largestinteriorrectangle" },
     { name = "llvmlite", specifier = ">=0.41.0" },
     { name = "markdown" },
@@ -1766,8 +1766,8 @@ wheels = [
 
 [[package]]
 name = "iohub"
-version = "0.3.8.dev7+46eb69c"
-source = { git = "https://github.com/czbiohub-sf/iohub.git?branch=feat%2Fflat-extra-metadata#46eb69c9d985f6cca99961aef1c6cede6ad47187" }
+version = "0.3.8"
+source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "blosc2" },
     { name = "dask", extra = ["array"] },
@@ -1784,6 +1784,10 @@ dependencies = [
     { name = "zarr" },
     { name = "zarrs" },
 ]
+sdist = { url = "https://files.pythonhosted.org/packages/ce/72/c5941fb52dc71ba78d08c05c90d4a7e15a4d7fc25000f7a5b4af20d96a72/iohub-0.3.8.tar.gz", hash = "sha256:db9e7f941058573ea6eb72028081bc830332dc28c7d59c74f7111cf4730cf793", size = 842230, upload-time = "2026-06-22T18:54:19.801Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/5a/55/8ddc8152eb3d38f94d3d5e15a95d210e044b70917071c31c97fd03148b76/iohub-0.3.8-py3-none-any.whl", hash = "sha256:4d34317104a635e5f23fb002352880292ecc04b0ea5286ee86347aaa649066c4", size = 97164, upload-time = "2026-06-22T18:54:18.582Z" },
+]
 
 [[package]]
 name = "ipykernel"
@@ -5772,6 +5776,7 @@ dependencies = [
     { name = "torch" },
     { name = "wget" },
 ]
+sdist = { url = "https://files.pythonhosted.org/packages/71/1d/96208c489ac53b48f024031191f4435985496d8bea4cf14e0132617a8fef/waveorder-3.0.4.tar.gz", hash = "sha256:e72e819bb5652e69f91203e21b309e6da15869a84016039f15479dd479794c39", size = 67110755, upload-time = "2026-06-23T03:37:29.238Z" }
 wheels = [
     { url = "https://files.pythonhosted.org/packages/1a/bc/3f3657bdf1461fcaea05958278bfe53c2f2f3c0084e1a51e5517150c96da/waveorder-3.0.4-py3-none-any.whl", hash = "sha256:6d075b5eeeed898f9f6cf5e00b4043b12e4fd7085e25b35cbbc769b1312c6e05", size = 286705, upload-time = "2026-04-29T23:03:00.171Z" },
 ]