bioimage_py.runner

 71class Runner(ABC):
 72    """Abstract runner. Subclasses implement :meth:`_execute` for a specific backend."""
 73
 74    def __init__(self, config: Optional[RunnerConfig] = None):
 75        self.config = config or RunnerConfig()
 76
 77    def run(
 78        self,
 79        function: ComputeFn,
 80        inputs: Sequence[SourceLike],
 81        outputs: Sequence[SourceLike] = (),
 82        *,
 83        block_shape: Optional[Tuple[int, ...]] = None,
 84        halo: Optional[Sequence[int]] = None,
 85        mask: Optional[SourceLike] = None,
 86        num_workers: int = 1,
 87        block_ids: Optional[Sequence[int]] = None,
 88        has_return_val: bool = False,
 89        name: str = "",
 90        roi: Optional[Tuple[slice, ...]] = None,
 91        pre_cleanup: Optional[Callable[[str], None]] = None,
 92        resume_from: Optional[str] = None,
 93    ) -> Optional[list]:
 94        """Run ``function`` block-wise over the inputs/outputs.
 95
 96        Args:
 97            function: Per-block function ``function(block, inputs, outputs, mask)``.
 98            inputs: Input source-like objects (read).
 99            outputs: Output source-like objects (written in place).
100            block_shape: Block shape; defaults to the domain source's chunks.
101            halo: Per-axis halo; if given, ``function`` receives a ``BlockWithHalo``.
102            mask: Optional binary mask source.
103            num_workers: Number of parallel workers / tasks.
104            block_ids: Restrict processing to these blocks (for re-running failures).
105            has_return_val: Whether ``function`` returns a value to collect.
106            name: A short name for progress display.
107            roi: Region of interest to restrict the blocking to.
108            pre_cleanup: Optional callback ``pre_cleanup(tmp_folder)`` invoked on the
109                orchestrating process with the job temp folder right before it is deleted
110                (distributed backends only, success path only). Use it to read out anything
111                worth keeping from the temp folder (e.g. the per-task timing files under
112                ``tmp_folder/timings/``) before cleanup. Ignored by the local runner, which
113                has no temp folder.
114            resume_from: Distributed backends only. Path to the preserved temp folder of a
115                failed run (``RunnerError.tmp_folder``). Re-runs only the blocks that did not
116                complete and merges them with the already-completed ones, so the result is
117                correct and complete. The run is resumed from the serialized payload, so
118                ``function``/``inputs``/``outputs``/``block_shape``/... from this call are
119                **ignored** -- pass ``resume_from`` to *finish the same call*, not to start a
120                new one. Mutually exclusive with ``block_ids``.
121
122        Returns:
123            The list of per-block return values (in ``block_ids`` order) if
124            ``has_return_val``, else ``None``.
125        """
126        if resume_from is not None:
127            if block_ids is not None:
128                raise ValueError("resume_from and block_ids are mutually exclusive; resume_from "
129                                 "re-runs the original partition's un-done blocks.")
130            return self._resume_entry(resume_from, name=name, pre_cleanup=pre_cleanup)
131
132        inputs = [as_source(i) for i in inputs]
133        outputs = [as_source(o) for o in outputs]
134        mask_source = as_source(mask) if mask is not None else None
135
136        domain = inputs[0] if inputs else (outputs[0] if outputs else None)
137        if domain is None:
138            raise ValueError("run() requires at least one input or output source.")
139
140        # Shape consistency: all inputs and the mask must match the domain shape.
141        dom_shape = tuple(domain.shape)
142        for src in inputs + ([mask_source] if mask_source is not None else []):
143            if tuple(src.shape) != dom_shape:
144                raise ValueError(
145                    f"Shape mismatch: source with shape {src.shape} does not match the "
146                    f"domain shape {domain.shape}."
147                )
148        # Outputs may carry a leading channel axis, but their trailing spatial dims must match
149        # the domain so the per-block roi indexes them consistently.
150        for out in outputs:
151            out_shape = tuple(out.shape)
152            if out_shape[-len(dom_shape):] != dom_shape:
153                raise ValueError(
154                    f"Output shape {out_shape} is incompatible with the domain shape "
155                    f"{dom_shape}: its trailing dimensions must match the domain."
156                )
157
158        block_shape = derive_block_shape(domain, block_shape)
159        halo_n = normalize_halo(halo, domain.ndim) if halo is not None else None
160        self._validate_write_safety(outputs, block_shape)
161
162        blocking = get_blocking(domain.shape, block_shape, roi)
163        if block_ids is None:
164            block_ids = list(range(int(blocking.number_of_blocks)))
165        else:
166            block_ids = [int(b) for b in block_ids]
167
168        results = self._execute(
169            function=function, inputs=inputs, outputs=outputs, mask=mask_source,
170            blocking=blocking, block_ids=block_ids, halo=halo_n,
171            has_return_val=has_return_val, num_workers=num_workers, name=name,
172            shape=tuple(domain.shape), block_shape=block_shape, roi=roi,
173            pre_cleanup=pre_cleanup,
174        )
175        return results if has_return_val else None
176
177    def map(
178        self,
179        function: Callable[[int], Any],
180        n_items: Optional[int] = None,
181        *,
182        item_ids: Optional[Sequence[int]] = None,
183        num_workers: int = 1,
184        has_return_val: bool = True,
185        name: str = "",
186        pre_cleanup: Optional[Callable[[str], None]] = None,
187        resume_from: Optional[str] = None,
188    ) -> Optional[list]:
189        """Map ``function(index)`` over item indices in parallel, across any backend.
190
191        Unlike :meth:`run`, this is not block-wise: there is no domain, blocking, sources or
192        mask. ``function`` takes a single integer index and returns its result; it must carry
193        whatever data it needs in its (cloudpickled) closure — e.g. a `SourceSpec` it reopens
194        and a file path it reads. This is the per-item counterpart used by per-object
195        workflows.
196
197        Args:
198            function: The per-item function ``function(index) -> result``.
199            n_items: The number of items; indices ``0 .. n_items - 1`` are processed. Ignored
200                if ``item_ids`` is given.
201            item_ids: Explicit item indices to process (e.g. to re-run failures). Defaults to
202                ``range(n_items)``.
203            num_workers: Number of parallel workers / tasks.
204            has_return_val: Whether ``function`` returns a value to collect.
205            name: A short name for progress display.
206            pre_cleanup: Optional ``pre_cleanup(tmp_folder)`` callback (distributed backends
207                only); see :meth:`run`.
208            resume_from: Distributed backends only; the preserved temp folder of a failed run
209                (see :meth:`run`). Re-runs only the incomplete items and merges with those
210                already done. Mutually exclusive with ``item_ids``.
211
212        Returns:
213            The list of per-item return values (in ``item_ids`` order) if ``has_return_val``,
214            else ``None``.
215
216        Raises:
217            ValueError: If neither ``n_items`` nor ``item_ids`` is given.
218        """
219        if resume_from is not None:
220            if item_ids is not None:
221                raise ValueError("resume_from and item_ids are mutually exclusive; resume_from "
222                                 "re-runs the original partition's un-done items.")
223            return self._resume_entry(resume_from, name=name, pre_cleanup=pre_cleanup)
224
225        if item_ids is None:
226            if n_items is None:
227                raise ValueError("map() requires either n_items or item_ids.")
228            item_ids = list(range(int(n_items)))
229        else:
230            item_ids = [int(i) for i in item_ids]
231
232        results = self._execute_map(
233            function=function, item_ids=item_ids, has_return_val=has_return_val,
234            num_workers=num_workers, name=name, pre_cleanup=pre_cleanup,
235        )
236        return results if has_return_val else None
237
238    def _resume_entry(self, tmp_folder: str, *, name: str,
239                      pre_cleanup: Optional[Callable[[str], None]]) -> Optional[list]:
240        """Resume a failed run from its temp folder; overridden by distributed runners.
241
242        The local runner keeps no temp folder, so resuming is not possible here.
243        """
244        raise ValueError(
245            "resume_from is only valid for distributed backends (subprocess/slurm); the local "
246            "runner keeps no temp folder. Re-run the operation to recompute in-process "
247            "(optionally with block_ids=err.failed_block_ids for a subset)."
248        )
249
250    @staticmethod
251    def _validate_write_safety(outputs: Sequence[Source], block_shape: Sequence[int]) -> None:
252        """Conservative guard: chunked output write-blocks must be a multiple of chunks.
253
254        This prevents two blocks from concurrently writing the same chunk (which would
255        corrupt it). Auto-derivation of a safe block shape is a flagged TODO.
256
257        Sharded outputs (``out.shards is not None``) are exempt: for them the atomic write
258        unit is the shard, and they are made safe by shard-exclusive routing (each shard's
259        blocks go to one worker, run sequentially) rather than by constraining the block
260        shape — see :func:`bioimage_py.util.group_blocks_by_shard`.
261
262        ``block_shape`` is spatial-only, but an output may carry a leading channel axis (its
263        ``chunks`` then have one extra leading entry); the block shape is aligned against the
264        trailing (spatial) chunk axes, since the channel axis is fully written by every block.
265        """
266        for out in outputs:
267            if out.shards is not None:
268                continue
269            chunks = out.chunks
270            if chunks is None or len(chunks) < len(block_shape):
271                continue
272            spatial_chunks = tuple(chunks[-len(block_shape):])
273            for bs, ch in zip(block_shape, spatial_chunks):
274                if bs % ch != 0:
275                    raise ValueError(
276                        f"Unsafe block shape for writing: {tuple(block_shape)} is not a multiple "
277                        f"of the output (spatial) chunk shape {spatial_chunks}. Concurrent writes "
278                        "could corrupt shared chunks; use a block shape that is a chunk multiple."
279                    )
280
281    @abstractmethod
282    def _execute(
283        self,
284        *,
285        function: ComputeFn,
286        inputs: Sequence[Source],
287        outputs: Sequence[Source],
288        mask: Optional[Source],
289        blocking: Blocking,
290        block_ids: Sequence[int],
291        halo: Optional[Sequence[int]],
292        has_return_val: bool,
293        num_workers: int,
294        name: str,
295        shape: Tuple[int, ...],
296        block_shape: Tuple[int, ...],
297        roi: Optional[Tuple[slice, ...]],
298        pre_cleanup: Optional[Callable[[str], None]] = None,
299    ) -> List[Any]:
300        """Execute the per-block function over ``block_ids`` and return ordered results."""
301        ...
302
303    @abstractmethod
304    def _execute_map(
305        self,
306        *,
307        function: Callable[[int], Any],
308        item_ids: Sequence[int],
309        has_return_val: bool,
310        num_workers: int,
311        name: str,
312        pre_cleanup: Optional[Callable[[str], None]] = None,
313    ) -> List[Any]:
314        """Execute ``function(index)`` over ``item_ids`` and return ordered results."""
315        ...

318class LocalRunner(Runner):
319    """Run blocks locally with a thread pool."""
320
321    def _execute(
322        self,
323        *,
324        function: ComputeFn,
325        inputs: Sequence[Source],
326        outputs: Sequence[Source],
327        mask: Optional[Source],
328        blocking: Blocking,
329        block_ids: Sequence[int],
330        halo: Optional[Sequence[int]],
331        has_return_val: bool,
332        num_workers: int,
333        name: str,
334        shape: Tuple[int, ...],
335        block_shape: Tuple[int, ...],
336        roi: Optional[Tuple[slice, ...]],
337        pre_cleanup: Optional[Callable[[str], None]] = None,
338    ) -> List[Any]:
339        """Run the blocks in a thread pool, collecting results and re-raising failures.
340
341        ``pre_cleanup`` is accepted for interface parity but ignored: the local runner has
342        no temp folder (and no per-worker concept) to read out before returning.
343        """
344        def call_one(bid: int) -> Any:
345            return run_block(function, blocking, bid, inputs, outputs, mask, halo)
346
347        # For sharded outputs, group blocks so each shard is written by a single thread
348        # (a group runs sequentially) and never corrupted by concurrent writes; otherwise
349        # each block is its own group, reproducing the plain one-future-per-block path.
350        groups = group_blocks_by_shard(blocking, outputs, block_ids)
351        if groups is None:
352            groups = [[int(b)] for b in block_ids]
353        else:
354            maybe_warn_imbalance([len(g) for g in groups], num_workers, len(groups), name)
355        return self._run_pool(groups, call_one, num_workers, name, unit="block")
356
357    def _execute_map(
358        self,
359        *,
360        function: Callable[[int], Any],
361        item_ids: Sequence[int],
362        has_return_val: bool,
363        num_workers: int,
364        name: str,
365        pre_cleanup: Optional[Callable[[str], None]] = None,
366    ) -> List[Any]:
367        """Run ``function(index)`` over ``item_ids`` in a thread pool (``pre_cleanup`` ignored)."""
368        groups = [[int(i)] for i in item_ids]
369        return self._run_pool(groups, lambda i: function(int(i)), num_workers, name, unit="item")
370
371    @staticmethod
372    def _run_pool(groups: Sequence[Sequence[int]], call_one: Callable[[int], Any],
373                  num_workers: int, name: str, *, unit: str = "block") -> List[Any]:
374        """Run ``call_one(id)`` for each id in a thread pool, ordered, re-raising failures.
375
376        The schedulable unit is a *group*: the ids in a group are run sequentially within one
377        worker thread, while distinct groups run concurrently. Singleton groups reproduce the
378        one-future-per-id behavior; multi-id groups serialize same-shard writes (see
379        :func:`bioimage_py.util.group_blocks_by_shard`).
380
381        Args:
382            groups: The work groups; each is a list of ids (block ids or item indices) run
383                sequentially. Results are returned in flattened ``groups`` order.
384            call_one: The per-id callable returning that id's result.
385            num_workers: Number of worker threads.
386            name: A short name for the progress bar (disabled when empty).
387            unit: The noun used in the failure message ("block" or "item").
388
389        Returns:
390            The per-id results in flattened ``groups`` order.
391
392        Raises:
393            RunnerError: If any id fails; the failed ids are attached for re-running. When an
394                id in a group fails, the remaining (un-run) ids of that group are reported as
395                failed too, since later same-shard writes cannot safely proceed.
396        """
397        groups = [list(g) for g in groups]
398        flat_ids = [i for g in groups for i in g]
399        result_by_id: Dict[int, Any] = {}
400        failed: List[int] = []
401        first_error: Optional[BaseException] = None
402
403        def _run_group(group: List[int]):
404            # Limit nested BLAS/OpenMP parallelism to 1 within this worker thread, reusing the
405            # process-wide controller (see _TP_CONTROLLER) so no per-thread dl_iterate_phdr scan
406            # is performed here -- that scan would deadlock against concurrent first-time imports.
407            with _TP_CONTROLLER.limit(limits=1):
408                local: Dict[int, Any] = {}
409                local_failed: List[int] = []
410                err: Optional[BaseException] = None
411                for k, bid in enumerate(group):
412                    try:
413                        local[bid] = call_one(bid)
414                    except Exception as error:  # noqa: BLE001 - we re-raise as RunnerError
415                        err = error
416                        local_failed = list(group[k:])
417                        break
418                return local, local_failed, err
419
420        with futures.ThreadPoolExecutor(max(1, int(num_workers))) as tp:
421            fut_to_group = {tp.submit(_run_group, g): g for g in groups}
422            with tqdm(total=len(flat_ids), desc=name or None, disable=not name) as pbar:
423                for fut in futures.as_completed(fut_to_group):
424                    group = fut_to_group[fut]
425                    local, local_failed, err = fut.result()
426                    result_by_id.update(local)
427                    if local_failed:
428                        failed.extend(local_failed)
429                        if first_error is None:
430                            first_error = err
431                    pbar.update(len(group))
432
433        if failed:
434            failed = sorted(set(failed))
435            raise RunnerError(
436                f"{len(failed)} {unit}(s) failed in '{name or 'run'}': "
437                f"{failed[:10]}. First error: {first_error!r}",
438                failed_block_ids=failed,
439            )
440        return [result_by_id[i] for i in flat_ids]

435class SubprocessRunner(_DistributedRunner):
436    """Distributed runner that launches each task as a local subprocess.
437
438    Exercises the full distributed protocol (cloudpickle payload, generated harness,
439    result/sentinel files, ``block_ids`` re-run) without a scheduler.
440    """
441
442    def _launch_and_wait(self, tmp: str, n_tasks: int, num_workers: int, name: str,
443                         task_ids: Optional[Sequence[int]] = None) -> None:
444        """Run each task as a local subprocess, up to ``num_workers`` concurrently.
445
446        The progress bar counts processed *blocks* (summed from the per-task done-logs) rather
447        than tasks; a background thread polls the logs while the tasks run. ``task_ids``
448        restricts the launch to a subset (resume); the bar still spans all tasks.
449        """
450        ids = list(range(n_tasks)) if task_ids is None else list(task_ids)
451        python = self.config.python_executable or sys.executable
452        cmd_base = [python, "-m", "bioimage_py.runner._harness", tmp]
453
454        def _run_task(task_id: int):
455            proc = subprocess.run(cmd_base + [str(task_id)], capture_output=True, text=True)
456            # The harness writes its own error/<id>.txt on a caught exception. But a failure
457            # *before* that try (e.g. an import error launching the module) would otherwise be
458            # silent, so capture the subprocess output as a fallback error file.
459            if proc.returncode != 0:
460                err_path = os.path.join(tmp, "error", f"{task_id}.txt")
461                if not os.path.exists(err_path):
462                    with open(err_path, "w") as f:
463                        f.write(f"Worker for task {task_id} exited with code {proc.returncode}.\n")
464                        if proc.stdout:
465                            f.write(f"--- stdout ---\n{proc.stdout}\n")
466                        if proc.stderr:
467                            f.write(f"--- stderr ---\n{proc.stderr}\n")
468            return proc
469
470        # Drive a block-counting progress bar from the done-logs (single source of truth, so no
471        # double-counting); clamp to the total in case a resume re-reads prior lines.
472        n_blocks = _total_blocks(tmp, n_tasks)
473        stop = threading.Event()
474        bar_thread = None
475        if name:
476            def _poll_bar() -> None:
477                with tqdm(total=n_blocks, desc=name, unit="block") as pbar:
478                    while not stop.wait(0.5):
479                        pbar.n = min(_count_done_blocks(tmp, n_tasks), n_blocks)
480                        pbar.refresh()
481                    pbar.n = min(_count_done_blocks(tmp, n_tasks), n_blocks)
482                    pbar.refresh()
483            bar_thread = threading.Thread(target=_poll_bar, daemon=True)
484            bar_thread.start()
485        try:
486            with futures.ThreadPoolExecutor(max(1, int(num_workers))) as tp:
487                list(tp.map(_run_task, ids))
488        finally:
489            stop.set()
490            if bar_thread is not None:
491                bar_thread.join()

505class SlurmRunner(_DistributedRunner):
506    """Distributed runner that submits one sbatch array job and polls it with ``sacct``.
507
508    Reuses the full distributed protocol from :class:`_DistributedRunner` (cloudpickle
509    payload, generated work-lists, per-task result + ``.success`` sentinel files, failure
510    reporting and ``block_ids`` re-run) and overrides only how tasks are launched and
511    awaited. The per-task sentinel file remains the ground truth for success; ``sacct`` is
512    queried only to detect tasks that died without writing a sentinel. A manifest is written
513    at submission time so an interrupted run can be picked back up with :meth:`reattach`.
514    """
515
516    def __init__(self, config: Optional[RunnerConfig] = None):
517        """Create the runner, requiring a :class:`SlurmConfig`.
518
519        Args:
520            config: The slurm configuration. ``None`` loads the user defaults from the config
521                file via :meth:`SlurmConfig.load` (honoring ``BIOIMAGE_PY_NO_CONFIG`` /
522                ``BIOIMAGE_PY_CONFIG``); ``tmp_root`` must still be set, here or in the file,
523                before running.
524
525        Raises:
526            TypeError: If ``config`` is a non-slurm ``RunnerConfig``.
527        """
528        if config is None:
529            config = SlurmConfig.load()
530        if not isinstance(config, SlurmConfig):
531            raise TypeError(
532                f"SlurmRunner requires a SlurmConfig, got {type(config).__name__}. "
533                "Pass job_config=SlurmConfig(...) (it carries partition/account/time/etc.)."
534            )
535        super().__init__(config)
536
537    def _launch_and_wait(self, tmp: str, n_tasks: int, num_workers: int, name: str,
538                         task_ids: Optional[Sequence[int]] = None) -> None:
539        """Submit an sbatch array job for the tasks and poll until they all finish.
540
541        Args:
542            tmp: The job temp folder (must live on a shared filesystem).
543            n_tasks: The total number of tasks the run was partitioned into.
544            num_workers: The array throttle (max tasks running concurrently).
545            name: A short name used for the job name and progress display.
546            task_ids: Restrict the submitted array to this subset of task indices (used by
547                :meth:`resume` to resubmit only the incomplete tasks); ``None`` submits all
548                ``0 .. n_tasks - 1``.
549        """
550        launch_ids = list(range(n_tasks)) if task_ids is None else sorted(set(int(t) for t in task_ids))
551        is_resume = task_ids is not None
552
553        def _guard_fail(message: str) -> None:
554            # On a resume we must never remove the user's preserved temp folder.
555            if not is_resume:
556                shutil.rmtree(tmp, ignore_errors=True)
557            raise ValueError(message)
558
559        if self.config.tmp_root is None:
560            _guard_fail(
561                "SlurmRunner requires config.tmp_root to be set to a shared filesystem "
562                "visible to all compute nodes (node-local /tmp is not usable)."
563            )
564
565        max_array = (self.config.max_array_size if self.config.max_array_size is not None
566                     else self._max_array_size())
567        if len(launch_ids) > max_array:
568            _guard_fail(
569                f"Run partitioned into {len(launch_ids)} tasks exceeds the maximum array size "
570                f"{max_array}. Lower num_workers or use a larger block_shape."
571            )
572
573        os.makedirs(os.path.join(tmp, "logs"), exist_ok=True)
574        throttle = max(1, min(int(num_workers), len(launch_ids)))
575        script_path = os.path.join(tmp, "submit.sh")
576        with open(script_path, "w") as f:
577            f.write(self._build_script(tmp, launch_ids, throttle, name))
578
579        # Unlike the tmp_root / max_array guards above, a submission failure deliberately does NOT
580        # remove the temp folder: the generated submit.sh, payload, and per-task block lists are
581        # exactly what's needed to diagnose why sbatch rejected the job. Re-raise naming the folder
582        # so the user knows where to look.
583        try:
584            job_id = self._submit(script_path)
585        except RuntimeError as err:
586            raise RuntimeError(f"{err} Temp folder preserved for debugging: {tmp}.") from err
587        manifest = {
588            "job_id": job_id,
589            "n_tasks": n_tasks,
590            "launch_ids": launch_ids,
591            "throttle": throttle,
592            "name": name,
593            "tmp": tmp,
594            "script": script_path,
595            "python_executable": self.config.python_executable or sys.executable,
596            "submit_time": time.strftime("%Y-%m-%d %H:%M:%S"),
597        }
598        manifest_path = os.path.join(tmp, "manifest.json")
599        if is_resume and os.path.exists(manifest_path):  # keep the prior job id for forensics
600            try:
601                with open(manifest_path) as f:
602                    manifest["resumed_from_job_id"] = json.load(f).get("job_id")
603            except (OSError, ValueError):
604                pass
605        with open(manifest_path, "w") as f:
606            json.dump(manifest, f, indent=2)
607
608        self._poll(job_id, n_tasks, tmp, name, task_ids=launch_ids)
609
610    @staticmethod
611    def _format_array_indices(task_ids: Sequence[int], throttle: int) -> str:
612        """Compress task ids into an sbatch ``--array`` spec, e.g. ``0,3,7-9%4``."""
613        ids = sorted(set(int(t) for t in task_ids))
614        parts: List[str] = []
615        i = 0
616        while i < len(ids):
617            j = i
618            while j + 1 < len(ids) and ids[j + 1] == ids[j] + 1:
619                j += 1
620            parts.append(str(ids[i]) if i == j else f"{ids[i]}-{ids[j]}")
621            i = j + 1
622        return ",".join(parts) + f"%{throttle}"
623
624    def _build_script(self, tmp: str, task_ids: Sequence[int], throttle: int, name: str) -> str:
625        """Render the sbatch array script for the given task indices."""
626        cfg = self.config
627        shebang, preamble = "#!/bin/bash", ""
628        if cfg.shebang:
629            lines = cfg.shebang.splitlines()
630            if lines and lines[0].startswith("#!"):
631                shebang, preamble = lines[0], "\n".join(lines[1:])
632            else:
633                preamble = cfg.shebang
634
635        # Collapse whitespace/newlines so the name cannot break or inject directives.
636        job_name = "_".join((name or "").split()) or "bioimage_py"
637        directives = [
638            f"--job-name={job_name}",
639            f"--array={self._format_array_indices(task_ids, throttle)}",
640            f"--cpus-per-task={int(cfg.cpus_per_task)}",
641            f"--output={os.path.join(tmp, 'logs', 'slurm-%A_%a.out')}",
642            f"--error={os.path.join(tmp, 'logs', 'slurm-%A_%a.err')}",
643        ]
644        if cfg.partition is not None:
645            directives.append(f"--partition={cfg.partition}")
646        if cfg.time is not None:
647            directives.append(f"--time={cfg.time}")
648        if cfg.mem is not None:
649            directives.append(f"--mem={cfg.mem}")
650        if int(cfg.gpus) > 0:
651            directives.append(f"--gpus={int(cfg.gpus)}")
652        if cfg.account is not None:
653            directives.append(f"--account={cfg.account}")
654        if cfg.qos is not None:
655            directives.append(f"--qos={cfg.qos}")
656        if cfg.constraint is not None:
657            directives.append(f"--constraint={cfg.constraint}")
658
659        python = shlex.quote(cfg.python_executable or sys.executable)
660        command = f'{python} -m bioimage_py.runner._harness {shlex.quote(tmp)} "${{SLURM_ARRAY_TASK_ID}}"'
661        lines = [shebang]
662        lines += [f"#SBATCH {d}" for d in directives]
663        if preamble:
664            lines.append(preamble)
665        lines.append(command)
666        return "\n".join(lines) + "\n"
667
668    @staticmethod
669    def _submit(script_path: str) -> str:
670        """Submit ``script_path`` with ``sbatch --parsable`` and return the job id."""
671        sbatch = shutil.which("sbatch")
672        if sbatch is None:
673            raise RuntimeError("sbatch not found on PATH; the slurm CLI must be available.")
674        proc = subprocess.run([sbatch, "--parsable", script_path],
675                              capture_output=True, text=True)
676        if proc.returncode != 0:
677            raise RuntimeError(f"sbatch submission failed (exit {proc.returncode}): "
678                               f"{proc.stderr.strip() or proc.stdout.strip()}")
679        job_id = proc.stdout.strip().split(";")[0].strip()
680        if not job_id.isdigit():
681            raise RuntimeError(f"Could not parse job id from sbatch output: {proc.stdout!r}")
682        return job_id
683
684    @staticmethod
685    def _max_array_size() -> int:
686        """Return the cluster's ``MaxArraySize`` (or a safe fallback)."""
687        scontrol = shutil.which("scontrol")
688        if scontrol is None:
689            return _DEFAULT_MAX_ARRAY
690        try:
691            proc = subprocess.run([scontrol, "show", "config"], capture_output=True, text=True)
692        except OSError:
693            return _DEFAULT_MAX_ARRAY
694        match = re.search(r"MaxArraySize\s*=\s*(\d+)", proc.stdout)
695        return int(match.group(1)) if match else _DEFAULT_MAX_ARRAY
696
697    @staticmethod
698    def _parse_array_range(spec: str) -> List[int]:
699        """Expand a pending-collapse range like ``[2-9,11%4]`` into its task indices."""
700        body = spec.strip("[]").split("%", 1)[0]
701        indices: List[int] = []
702        for part in body.split(","):
703            part = part.strip()
704            if not part:
705                continue
706            if "-" in part:
707                lo, hi = part.split("-", 1)
708                indices.extend(range(int(lo), int(hi) + 1))
709            else:
710                indices.append(int(part))
711        return indices
712
713    def _sacct_states(self, job_id: str) -> Optional[Dict[int, str]]:
714        """Return ``{array_index: STATE}`` for the array job, or ``None`` on a poll error.
715
716        ``None`` (a transient ``sacct`` failure) means *skip this poll*; an empty dict means
717        the job is simply not registered with the scheduler yet. A task absent from the
718        result is treated as pending, never as dead.
719        """
720        sacct = shutil.which("sacct")
721        if sacct is None:
722            raise RuntimeError("sacct not found on PATH; the slurm CLI must be available.")
723        try:
724            proc = subprocess.run(
725                [sacct, "-X", "-n", "-P", "--format=JobID,State", "-j", str(job_id)],
726                capture_output=True, text=True,
727            )
728        except OSError:
729            return None
730        if proc.returncode != 0:
731            return None
732
733        states: Dict[int, str] = {}
734        for line in proc.stdout.splitlines():
735            line = line.strip()
736            if not line or "|" not in line:
737                continue
738            jid, _, raw_state = line.partition("|")
739            jid = jid.split(";", 1)[0]
740            if "." in jid or "_" not in jid:  # step rows (defensive; -X already excludes them)
741                continue
742            # Take the first token: normalises e.g. "CANCELLED by 12345" -> "CANCELLED".
743            tokens = raw_state.split()
744            state = tokens[0].upper() if tokens else ""
745            suffix = jid.split("_", 1)[1]
746            if suffix.startswith("["):
747                for idx in self._parse_array_range(suffix):
748                    states[idx] = state
749            else:
750                try:
751                    states[int(suffix)] = state
752                except ValueError:
753                    continue
754        return states
755
756    def _job_known(self, job_id: str, attempts: int = 3) -> bool:
757        """Whether the job is known to ``sacct``, retrying to tolerate post-submit lag.
758
759        A transient ``sacct`` error (``None``) or any returned row counts as known; only a
760        sustained empty result across ``attempts`` polls is treated as unknown.
761        """
762        for attempt in range(attempts):
763            states = self._sacct_states(job_id)
764            if states is None or states:
765                return True
766            if attempt + 1 < attempts:
767                time.sleep(self.config.poll_interval)
768        return False
769
770    def _poll(self, job_id: str, n_tasks: int, tmp: str, name: str,
771              task_ids: Optional[Sequence[int]] = None) -> None:
772        """Poll ``sacct`` until every task has a visible sentinel or is confirmed dead.
773
774        The scheduler ``State`` is not subject to NFS lag, but the ``.success`` sentinels the
775        compute nodes write can take up to the mount's attribute-cache timeout to become
776        visible here. So a ``COMPLETED`` task (its harness exited 0, hence wrote a sentinel)
777        is given ``config.latency_wait`` for that sentinel to appear; any other terminal
778        state means the harness did not succeed and the task is declared dead after a short
779        confirmation grace. Tasks absent from ``sacct`` are pending, never dead.
780
781        Args:
782            job_id: The submitted array job id.
783            n_tasks: The total number of tasks (spans the block-counting progress bar).
784            tmp: The job temp folder (where sentinels are written).
785            name: A short name for the progress bar (disables it when empty).
786            task_ids: The subset of task indices this job actually runs (a resume submits only
787                the incomplete tasks); resolution is over this subset, ``None`` means all tasks.
788        """
789        poll_ids = list(range(n_tasks)) if task_ids is None else sorted(set(int(t) for t in task_ids))
790
791        def has_sentinel(t: int) -> bool:
792            return os.path.exists(os.path.join(tmp, "success", f"{t}.success"))
793
794        latency_wait = max(float(self.config.latency_wait), self.config.poll_interval)
795        fail_grace = max(self.config.poll_interval, 5.0)
796        terminal_since: Dict[int, float] = {}
797        terminal_count: Dict[int, int] = {}
798        resolved: set = set()
799        # The bar counts processed blocks across ALL tasks (a resume credits prior progress).
800        n_blocks = _total_blocks(tmp, n_tasks)
801        with tqdm(total=n_blocks, desc=name or None, disable=not name, unit="block") as pbar:
802            while len(resolved) < len(poll_ids):
803                states = self._sacct_states(job_id)
804                if states is None:  # transient sacct error: skip this poll.
805                    time.sleep(self.config.poll_interval)
806                    continue
807
808                now = time.monotonic()
809                ok = {t for t in poll_ids if has_sentinel(t)}
810                running = sum(1 for s in states.values() if s == "RUNNING")
811                dead = set()
812                for t in poll_ids:
813                    if t in ok:
814                        terminal_since.pop(t, None)
815                        terminal_count.pop(t, None)
816                        continue
817                    state = states.get(t)
818                    if state in _TERMINAL_STATES:
819                        terminal_since.setdefault(t, now)
820                        terminal_count[t] = terminal_count.get(t, 0) + 1
821                        # COMPLETED -> sentinel was written, just wait it out over NFS; any
822                        # other terminal state -> the task will never produce a sentinel.
823                        grace = latency_wait if state == "COMPLETED" else fail_grace
824                        if (terminal_count[t] >= 2 and now - terminal_since[t] >= grace
825                                and not has_sentinel(t)):
826                            dead.add(t)
827                    else:  # pending/running/requeued: reset the dead countdown.
828                        terminal_since.pop(t, None)
829                        terminal_count.pop(t, None)
830
831                resolved = ok | dead
832                pbar.n = min(_count_done_blocks(tmp, n_tasks), n_blocks)
833                pbar.set_postfix(ok=len(ok), failed=len(dead), run=running,
834                                 pending=max(0, len(poll_ids) - len(resolved) - running), refresh=False)
835                pbar.refresh()
836                if len(resolved) >= len(poll_ids):
837                    break
838                try:
839                    time.sleep(self.config.poll_interval)
840                except KeyboardInterrupt:
841                    print(f"\nInterrupted while waiting on slurm job {job_id}. The job was left "
842                          f"running; reattach with SlurmRunner(...).reattach({tmp!r}).")
843                    raise
844
845    def reattach(self, tmp_folder: str, name: str = "reattach",
846                 pre_cleanup: Optional[Callable[[str], None]] = None) -> Optional[list]:
847        """Reattach to a previously submitted run and finalize it.
848
849        Picks a run back up from its manifest (e.g. after the orchestrating login-node
850        process was interrupted) instead of resubmitting. Only ``poll_interval`` is read
851        from this runner's config, so a freshly constructed ``SlurmRunner`` can reattach.
852
853        Args:
854            tmp_folder: The job temp folder containing ``manifest.json`` and ``payload.pkl``.
855            name: A short name for the progress display.
856            pre_cleanup: Optional ``pre_cleanup(tmp)`` callback invoked right before the temp
857                folder is removed (forwarded to :meth:`_finalize`).
858
859        Returns:
860            The per-block return values (if the run collected any), else ``None``.
861
862        Raises:
863            RunnerError: If any task failed (sentinel missing).
864            RuntimeError: If the manifest's job is unknown to slurm and the run did not
865                already complete.
866        """
867        with open(os.path.join(tmp_folder, "manifest.json")) as f:
868            manifest = json.load(f)
869        job_id, n_tasks = str(manifest["job_id"]), int(manifest["n_tasks"])
870        with open(os.path.join(tmp_folder, "payload.pkl"), "rb") as f:
871            has_return_val = bool(cloudpickle.load(f)["has_return_val"])
872
873        # Reconstruct the partition in numeric task order (never glob: it sorts lexically).
874        tasks: List[List[int]] = []
875        for task_id in range(n_tasks):
876            with open(os.path.join(tmp_folder, "blocks", f"{task_id}.json")) as f:
877                tasks.append(json.load(f))
878        block_ids = [b for task in tasks for b in task]
879
880        all_done = all(os.path.exists(os.path.join(tmp_folder, "success", f"{t}.success"))
881                       for t in range(n_tasks))
882        if not all_done:
883            # Only a job that stays unknown to sacct across retries (not registration lag
884            # right after submit, nor a transient error) is treated as unrecoverable.
885            if not self._job_known(job_id):
886                raise RuntimeError(
887                    f"Slurm job {job_id} is not known to the scheduler and the run did not "
888                    f"complete. Inspect {tmp_folder} or resubmit."
889                )
890            self._poll(job_id, n_tasks, tmp_folder, name)
891
892        results = self._finalize(tmp_folder, n_tasks, tasks, block_ids, has_return_val, name,
893                                 pre_cleanup=pre_cleanup)
894        return results if has_return_val else None