feat(nextflow): integrate virtual-stain into mantis-v2 pipeline

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_minutes = 360
+    echo_resources(num_cpus, mem_gb, time_minutes)
 
     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_minutes,
         "slurm_partition": "cpu",
     }
 

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,35 @@
 logger = logging.getLogger(__name__)
 
 
+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
+    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_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_minutes": int(time_minutes)}
+    click.echo("RESOURCES:" + json.dumps(payload))
+
+
 def get_submitit_cluster(
     local: bool = False,
     cluster: str | None = None,

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_minutes = 60
+    echo_resources(num_cpus, mem_gb, time_minutes)
 
     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_minutes,
         "slurm_partition": "preempted",
     }
     if sbatch_filepath:

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_minutes = 60
+    echo_resources(num_cpus, mem_gb, time_minutes)
 
     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_minutes,
         "slurm_partition": "cpu",
     }
 

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,30 +330,31 @@ 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_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)")
         return
 
     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_per_cpu": 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_minutes,
         "slurm_partition": "gpu",
     }
 

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

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_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_minutes: res.time_minutes as int]
 }
 
 def slurm_log_dir(step_name) {

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_minutes * task.attempt} min" }
     maxRetries 1
     errorStrategy 'retry'
 

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_minutes * task.attempt} min" }
     maxRetries 1
     errorStrategy 'retry'
 

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_minutes * task.attempt} min" }
     maxRetries 1
     errorStrategy 'retry'