Gpu Only Scheduling

GPU-Only Scheduling - Research Notes

Experiment Overview

Context

The original _build_cycle_assignment() in the Snakefile used a dual-pool architecture: GPU slots filled first, then overflow cycles were assigned to CPU. With 5 GPU slots and 9 cycles, cycles 6-9 would run on CPU. Investigation of CX_19-003_lymph-node_R1 revealed CPU jobs were taking 2+ hours per cycle vs ~8 minutes on GPU for stitching alone.

Performance Data (CX_19-003_lymph-node_R1, 9x7 tiles, 20 z-planes, 4 channels)

Stitching (per cycle)

Full Pipeline (stitch + decon + EDF per cycle)

Root Cause: BaSiC Illumination Correction

• BaSiC fit() is called per z-plane (not per channel) — 80 calls per cycle (4 channels × 20 z-planes)
• Each fit() runs 500 iterative DCT operations
• On GPU (CuPy): FFT/DCT is massively parallel → milliseconds per iteration
• On CPU (SciPy): Sequential DCT → 1-2.5 minutes per z-plane
• BaSiC caching (compute once per channel) was evaluated and REJECTED — causes 15-20% intensity errors for sparse markers (see basic-caching-evaluation skill)

Verified Workflow

GPU-Only Cycle Assignment (Snakefile)

def _build_cycle_assignment():
    """Pre-assign each cycle to an account and mode (gpu-only).

    All cycles are assigned to GPU, round-robin across accounts proportional
    to each account's GPU slot count.  Overflow cycles queue in SLURM until
    a GPU slot frees up — this is dramatically faster than CPU fallback
    (GPU ~8 min vs CPU ~2 hours per cycle for stitching alone).
    """
    assignment = {}
    gpu_queue = []
    for acct in ACCOUNTS:
        gpu_queue.extend([{
            "account": acct["name"],
            "mode": "gpu",
            "partition": acct.get("partition_gpu", "hpg-b200,hpg-turin"),
        }] * acct.get("gpu_slots", 0))

    gpu_accounts = [acct for acct in ACCOUNTS if acct.get("gpu_slots", 0) > 0]
    for i, cyc in enumerate(CYCLES):
        cyc_key = cyc_fmt(cyc)
        if i < len(gpu_queue):
            assignment[cyc_key] = gpu_queue[i]
        else:
            acct = gpu_accounts[i % len(gpu_accounts)] if gpu_accounts else ACCOUNTS[0]
            assignment[cyc_key] = {
                "account": acct["name"],
                "mode": "gpu",
                "partition": acct.get("partition_gpu", "hpg-b200,hpg-turin"),
            }
    return assignment

CLI: Set -j to GPU Slots Only

# In workflow run command:
j_val = pool["total_gpu_avail"]  # NOT pool["total_avail"]

Dynamic Worker Counts (Wrapper Scripts)

# stitch.py and edf.py — use SLURM allocation, not hardcoded 4
CPUS = int(getattr(snakemake.resources, "cpus_per_task", 4))
# GPU jobs get 4 cores, CPU jobs would get 8

Failed Attempts (Critical)

Final Parameters

# Snakefile: _build_cycle_assignment()
# Mode: GPU-only for ALL cycles (no CPU fallback)
# Overflow: round-robin across GPU accounts, queue in SLURM

# CLI: workflow run
# -j = total_gpu_avail (5 for clive+maigan)

# Wrapper scripts: dynamic worker counts
# CPUS = snakemake.resources.cpus_per_task (4 for GPU, 8 for CPU)

Key Insights

• Measured per-step speedups (CX_19-003, 9x7 grid, 4ch, 20z): stitch 25x, decon 5x, EDF 15x (~13x full cycle)
• GPU ~22 min vs CPU ~282 min per cycle: stitch ~8 vs ~200 min, decon ~12 vs ~60 min, EDF ~1.5 vs ~22 min
• Queuing is faster than CPU: Even if a cycle waits 12 min for a GPU slot, total time is still ~44 min vs ~18 hours with CPU
• Per-cycle pipeline (stitch→decon→edf) means GPU slots free up incrementally — overflow cycles start as soon as one cycle's stitch finishes
• BaSiC fit() is the bottleneck, not stitch_images() or blending — it runs 500 DCT iterations per z-plane, 80 z-planes per cycle
• Never cache BaSiC flatfields — confirmed by basic-caching-evaluation skill (15-20% intensity errors)
• Dynamic worker counts matter — max_workers=4 wasted 50% of CPU job allocation
• Target specific Snakemake rules to avoid GPU QC contention — Running bare snakemake triggers QC rules (qc_stitch, qc_decon, qc_edf, qc_registration) which don't need GPU but compete for SLURM scheduling slots. Use snakemake registration --configfile config.yaml to run only registration. QC can run separately without GPU allocation (16 GB RAM, 2 CPUs, 30 min)
• Account distribution for registration — _registration_assignment() picks the first GPU account in config.yaml resources.accounts list. Reorder accounts to control which account runs registration jobs. Useful for load balancing when one account has queued jobs

When to Apply

• Configuring SLURM scheduling for KINTSUGI batch processing
• Investigating slow CPU stitching or deconvolution jobs
• Deciding whether to add CPU fallback for overflow cycles
• Tuning -j parameter for Snakemake SLURM submissions

References

• basic-caching-evaluation skill — Why BaSiC caching fails
• snakemake-workflow-architecture skill — Snakefile design decisions
• gpu-parallel-scheduling skill — GPU queue pattern for notebooks
• slurm-concurrent-processing skill — Multi-account SLURM architecture