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]

548class SubprocessRunner(_DistributedRunner):
549    """Distributed runner that launches each task as a local subprocess.
550
551    Exercises the full distributed protocol (cloudpickle payload, generated harness,
552    result/sentinel files, ``block_ids`` re-run) without a scheduler.
553    """
554
555    def _launch_and_wait(self, tmp: str, n_tasks: int, num_workers: int, name: str,
556                         task_ids: Optional[Sequence[int]] = None) -> None:
557        """Run each task as a local subprocess, up to ``num_workers`` concurrently.
558
559        The progress bar counts processed *blocks* (summed from the per-task done-logs) rather
560        than tasks; a background thread polls the logs while the tasks run. ``task_ids``
561        restricts the launch to a subset (resume); the bar still spans all tasks.
562        """
563        ids = list(range(n_tasks)) if task_ids is None else list(task_ids)
564        python = self.config.python_executable or sys.executable
565        cmd_base = [python, "-m", "bioimage_py.runner._harness", tmp]
566
567        def _write_error(task_id: int, summary: str, stdout: Optional[str] = None,
568                         stderr: Optional[str] = None) -> None:
569            # Synthesize error/<id>.txt for a failure the harness could not report itself. The
570            # human summary is written *last* so it is the line _failure_message surfaces
571            # (it reads lines[-1]). Never clobber an error file the harness already wrote.
572            err_path = os.path.join(tmp, "error", f"{task_id}.txt")
573            if os.path.exists(err_path):
574                return
575            with open(err_path, "w") as f:
576                if stdout:
577                    f.write(f"--- stdout ---\n{stdout}\n")
578                if stderr:
579                    f.write(f"--- stderr ---\n{stderr}\n")
580                f.write(summary + "\n")
581
582        def _run_task(task_id: int):
583            # A timeout kill or a launch failure raises here; both are turned into an error
584            # file + a normal return so _launch_and_wait always returns and _finalize can raise
585            # the standard RunnerError (with failed_block_ids + preserved tmp) instead of an
586            # escaping exception. Blocks completed before a timeout kill are already in the
587            # done-log, so _finalize reports only the un-done ones and resume skips the rest.
588            try:
589                proc = subprocess.run(cmd_base + [str(task_id)], capture_output=True, text=True,
590                                      timeout=self.config.task_timeout)
591            except subprocess.TimeoutExpired as err:
592                _write_error(task_id, f"TimeoutError: worker for task {task_id} exceeded "
593                             f"{self.config.task_timeout}s and was killed.", err.stdout, err.stderr)
594                return None
595            except OSError as err:
596                _write_error(task_id, f"OSError: failed to launch worker for task {task_id}: {err!r}")
597                return None
598            # The harness writes its own error/<id>.txt on a caught exception. But a failure
599            # *before* that try (e.g. an import error launching the module) would otherwise be
600            # silent, so capture the subprocess output as a fallback error file.
601            if proc.returncode != 0:
602                _write_error(task_id, f"Worker for task {task_id} exited with code {proc.returncode}.",
603                             proc.stdout, proc.stderr)
604            return proc
605
606        # Drive a block-counting progress bar from the done-logs (single source of truth, so no
607        # double-counting); clamp to the total in case a resume re-reads prior lines.
608        n_blocks = _total_blocks(tmp, n_tasks)
609        stop = threading.Event()
610        bar_thread = None
611        if name:
612            def _poll_bar() -> None:
613                counter = _DoneLogCounter(tmp, n_tasks)
614                with tqdm(total=n_blocks, desc=name, unit="block") as pbar:
615                    while not stop.wait(0.5):
616                        pbar.n = min(counter.count(), n_blocks)
617                        pbar.refresh()
618                    pbar.n = min(counter.count(), n_blocks)
619                    pbar.refresh()
620            bar_thread = threading.Thread(target=_poll_bar, daemon=True)
621            bar_thread.start()
622        try:
623            with futures.ThreadPoolExecutor(max(1, int(num_workers))) as tp:
624                list(tp.map(_run_task, ids))
625        finally:
626            stop.set()
627            if bar_thread is not None:
628                bar_thread.join()

 642class SlurmRunner(_DistributedRunner):
 643    """Distributed runner that submits one sbatch array job and polls it with ``sacct``.
 644
 645    Reuses the full distributed protocol from :class:`_DistributedRunner` (cloudpickle
 646    payload, generated work-lists, per-task result + ``.success`` sentinel files, failure
 647    reporting and ``block_ids`` re-run) and overrides only how tasks are launched and
 648    awaited. The per-task sentinel file remains the ground truth for success; ``sacct`` is
 649    queried only to detect tasks that died without writing a sentinel. A manifest is written
 650    at submission time so an interrupted run can be picked back up with :meth:`reattach`.
 651    """
 652
 653    def __init__(self, config: Optional[RunnerConfig] = None):
 654        """Create the runner, requiring a :class:`SlurmConfig`.
 655
 656        Args:
 657            config: The slurm configuration. ``None`` loads the user defaults from the config
 658                file via :meth:`SlurmConfig.load` (honoring ``BIOIMAGE_PY_NO_CONFIG`` /
 659                ``BIOIMAGE_PY_CONFIG``); ``tmp_root`` must still be set, here or in the file,
 660                before running.
 661
 662        Raises:
 663            TypeError: If ``config`` is a non-slurm ``RunnerConfig``.
 664        """
 665        if config is None:
 666            config = SlurmConfig.load()
 667        if not isinstance(config, SlurmConfig):
 668            raise TypeError(
 669                f"SlurmRunner requires a SlurmConfig, got {type(config).__name__}. "
 670                "Pass job_config=SlurmConfig(...) (it carries partition/account/time/etc.)."
 671            )
 672        super().__init__(config)
 673        self._max_array_cache: Optional[int] = None
 674
 675    def _max_tasks(self) -> int:
 676        """The cluster's ``MaxArraySize`` -- the hard cap on the number of tasks per array job.
 677
 678        Uses ``config.max_array_size`` when set, else queries ``scontrol`` once and memoizes it
 679        (a cluster property, stable for this runner's lifetime). This bounds
 680        :meth:`_resolve_n_tasks` so over-partitioning degrades gracefully to the array limit
 681        instead of failing at submit.
 682        """
 683        if self.config.max_array_size is not None:
 684            return int(self.config.max_array_size)
 685        if self._max_array_cache is None:
 686            self._max_array_cache = self._max_array_size()
 687        return self._max_array_cache
 688
 689    def _launch_and_wait(self, tmp: str, n_tasks: int, num_workers: int, name: str,
 690                         task_ids: Optional[Sequence[int]] = None) -> None:
 691        """Submit an sbatch array job for the tasks and poll until they all finish.
 692
 693        Args:
 694            tmp: The job temp folder (must live on a shared filesystem).
 695            n_tasks: The total number of tasks the run was partitioned into.
 696            num_workers: The array throttle (max tasks running concurrently).
 697            name: A short name used for the job name and progress display.
 698            task_ids: Restrict the submitted array to this subset of task indices (used by
 699                :meth:`resume` to resubmit only the incomplete tasks); ``None`` submits all
 700                ``0 .. n_tasks - 1``.
 701        """
 702        launch_ids = list(range(n_tasks)) if task_ids is None else sorted(set(int(t) for t in task_ids))
 703        is_resume = task_ids is not None
 704
 705        def _guard_fail(message: str) -> None:
 706            # On a resume we must never remove the user's preserved temp folder.
 707            if not is_resume:
 708                shutil.rmtree(tmp, ignore_errors=True)
 709            raise ValueError(message)
 710
 711        if self.config.tmp_root is None:
 712            _guard_fail(
 713                "SlurmRunner requires config.tmp_root to be set to a shared filesystem "
 714                "visible to all compute nodes (node-local /tmp is not usable)."
 715            )
 716
 717        max_array = self._max_tasks()
 718        if len(launch_ids) > max_array:
 719            _guard_fail(
 720                f"Run partitioned into {len(launch_ids)} tasks exceeds the maximum array size "
 721                f"{max_array}. Lower num_workers/tasks_per_worker or use a larger block_shape."
 722            )
 723
 724        os.makedirs(os.path.join(tmp, "logs"), exist_ok=True)
 725        throttle = max(1, min(int(num_workers), len(launch_ids)))
 726        script_path = os.path.join(tmp, "submit.sh")
 727        with open(script_path, "w") as f:
 728            f.write(self._build_script(tmp, launch_ids, throttle, name))
 729
 730        # Unlike the tmp_root / max_array guards above, a submission failure deliberately does NOT
 731        # remove the temp folder: the generated submit.sh, payload, and per-task block lists are
 732        # exactly what's needed to diagnose why sbatch rejected the job. Re-raise naming the folder
 733        # so the user knows where to look.
 734        try:
 735            job_id = self._submit(script_path)
 736        except RuntimeError as err:
 737            raise RuntimeError(f"{err} Temp folder preserved for debugging: {tmp}.") from err
 738        manifest = {
 739            "job_id": job_id,
 740            "n_tasks": n_tasks,
 741            "launch_ids": launch_ids,
 742            "throttle": throttle,
 743            "name": name,
 744            "tmp": tmp,
 745            "script": script_path,
 746            "python_executable": self.config.python_executable or sys.executable,
 747            "submit_time": time.strftime("%Y-%m-%d %H:%M:%S"),
 748        }
 749        manifest_path = os.path.join(tmp, "manifest.json")
 750        if is_resume and os.path.exists(manifest_path):  # keep the prior job id for forensics
 751            try:
 752                with open(manifest_path) as f:
 753                    manifest["resumed_from_job_id"] = json.load(f).get("job_id")
 754            except (OSError, ValueError):
 755                pass
 756        with open(manifest_path, "w") as f:
 757            json.dump(manifest, f, indent=2)
 758
 759        self._poll(job_id, n_tasks, tmp, name, task_ids=launch_ids)
 760
 761    @staticmethod
 762    def _format_array_indices(task_ids: Sequence[int], throttle: int) -> str:
 763        """Compress task ids into an sbatch ``--array`` spec, e.g. ``0,3,7-9%4``."""
 764        ids = sorted(set(int(t) for t in task_ids))
 765        parts: List[str] = []
 766        i = 0
 767        while i < len(ids):
 768            j = i
 769            while j + 1 < len(ids) and ids[j + 1] == ids[j] + 1:
 770                j += 1
 771            parts.append(str(ids[i]) if i == j else f"{ids[i]}-{ids[j]}")
 772            i = j + 1
 773        return ",".join(parts) + f"%{throttle}"
 774
 775    def _build_script(self, tmp: str, task_ids: Sequence[int], throttle: int, name: str) -> str:
 776        """Render the sbatch array script for the given task indices."""
 777        cfg = self.config
 778        shebang, preamble = "#!/bin/bash", ""
 779        if cfg.shebang:
 780            lines = cfg.shebang.splitlines()
 781            if lines and lines[0].startswith("#!"):
 782                shebang, preamble = lines[0], "\n".join(lines[1:])
 783            else:
 784                preamble = cfg.shebang
 785
 786        # Collapse whitespace/newlines so the name cannot break or inject directives.
 787        job_name = "_".join((name or "").split()) or "bioimage_py"
 788        directives = [
 789            f"--job-name={job_name}",
 790            f"--array={self._format_array_indices(task_ids, throttle)}",
 791            f"--cpus-per-task={int(cfg.cpus_per_task)}",
 792            f"--output={os.path.join(tmp, 'logs', 'slurm-%A_%a.out')}",
 793            f"--error={os.path.join(tmp, 'logs', 'slurm-%A_%a.err')}",
 794        ]
 795        if cfg.partition is not None:
 796            directives.append(f"--partition={cfg.partition}")
 797        if cfg.time is not None:
 798            directives.append(f"--time={cfg.time}")
 799        if cfg.mem is not None:
 800            directives.append(f"--mem={cfg.mem}")
 801        if int(cfg.gpus) > 0:
 802            directives.append(f"--gpus={int(cfg.gpus)}")
 803        if cfg.account is not None:
 804            directives.append(f"--account={cfg.account}")
 805        if cfg.qos is not None:
 806            directives.append(f"--qos={cfg.qos}")
 807        if cfg.constraint is not None:
 808            directives.append(f"--constraint={cfg.constraint}")
 809
 810        python = shlex.quote(cfg.python_executable or sys.executable)
 811        command = f'{python} -m bioimage_py.runner._harness {shlex.quote(tmp)} "${{SLURM_ARRAY_TASK_ID}}"'
 812        lines = [shebang]
 813        lines += [f"#SBATCH {d}" for d in directives]
 814        if preamble:
 815            lines.append(preamble)
 816        lines.append(command)
 817        return "\n".join(lines) + "\n"
 818
 819    @staticmethod
 820    def _submit(script_path: str) -> str:
 821        """Submit ``script_path`` with ``sbatch --parsable`` and return the job id."""
 822        sbatch = shutil.which("sbatch")
 823        if sbatch is None:
 824            raise RuntimeError("sbatch not found on PATH; the slurm CLI must be available.")
 825        proc = subprocess.run([sbatch, "--parsable", script_path],
 826                              capture_output=True, text=True)
 827        if proc.returncode != 0:
 828            raise RuntimeError(f"sbatch submission failed (exit {proc.returncode}): "
 829                               f"{proc.stderr.strip() or proc.stdout.strip()}")
 830        job_id = proc.stdout.strip().split(";")[0].strip()
 831        if not job_id.isdigit():
 832            raise RuntimeError(f"Could not parse job id from sbatch output: {proc.stdout!r}")
 833        return job_id
 834
 835    @staticmethod
 836    def _max_array_size() -> int:
 837        """Return the cluster's ``MaxArraySize`` (or a safe fallback)."""
 838        scontrol = shutil.which("scontrol")
 839        if scontrol is None:
 840            return _DEFAULT_MAX_ARRAY
 841        try:
 842            proc = subprocess.run([scontrol, "show", "config"], capture_output=True, text=True)
 843        except OSError:
 844            return _DEFAULT_MAX_ARRAY
 845        match = re.search(r"MaxArraySize\s*=\s*(\d+)", proc.stdout)
 846        return int(match.group(1)) if match else _DEFAULT_MAX_ARRAY
 847
 848    @staticmethod
 849    def _parse_array_range(spec: str) -> List[int]:
 850        """Expand a pending-collapse range like ``[2-9,11%4]`` into its task indices."""
 851        body = spec.strip("[]").split("%", 1)[0]
 852        indices: List[int] = []
 853        for part in body.split(","):
 854            part = part.strip()
 855            if not part:
 856                continue
 857            if "-" in part:
 858                lo, hi = part.split("-", 1)
 859                indices.extend(range(int(lo), int(hi) + 1))
 860            else:
 861                indices.append(int(part))
 862        return indices
 863
 864    def _sacct_states(self, job_id: str) -> Optional[Dict[int, str]]:
 865        """Return ``{array_index: STATE}`` for the array job, or ``None`` on a poll error.
 866
 867        ``None`` (a transient ``sacct`` failure) means *skip this poll*; an empty dict means
 868        the job is simply not registered with the scheduler yet. A task absent from the
 869        result is treated as pending, never as dead.
 870        """
 871        sacct = shutil.which("sacct")
 872        if sacct is None:
 873            raise RuntimeError("sacct not found on PATH; the slurm CLI must be available.")
 874        try:
 875            proc = subprocess.run(
 876                [sacct, "-X", "-n", "-P", "--format=JobID,State", "-j", str(job_id)],
 877                capture_output=True, text=True,
 878            )
 879        except OSError:
 880            return None
 881        if proc.returncode != 0:
 882            return None
 883
 884        states: Dict[int, str] = {}
 885        for line in proc.stdout.splitlines():
 886            line = line.strip()
 887            if not line or "|" not in line:
 888                continue
 889            jid, _, raw_state = line.partition("|")
 890            jid = jid.split(";", 1)[0]
 891            if "." in jid or "_" not in jid:  # step rows (defensive; -X already excludes them)
 892                continue
 893            # Take the first token: normalises e.g. "CANCELLED by 12345" -> "CANCELLED".
 894            tokens = raw_state.split()
 895            state = tokens[0].upper() if tokens else ""
 896            suffix = jid.split("_", 1)[1]
 897            if suffix.startswith("["):
 898                for idx in self._parse_array_range(suffix):
 899                    states[idx] = state
 900            else:
 901                try:
 902                    states[int(suffix)] = state
 903                except ValueError:
 904                    continue
 905        return states
 906
 907    def _job_known(self, job_id: str, attempts: int = 3) -> bool:
 908        """Whether the job is known to ``sacct``, retrying to tolerate post-submit lag.
 909
 910        A transient ``sacct`` error (``None``) or any returned row counts as known; only a
 911        sustained empty result across ``attempts`` polls is treated as unknown.
 912        """
 913        for attempt in range(attempts):
 914            states = self._sacct_states(job_id)
 915            if states is None or states:
 916                return True
 917            if attempt + 1 < attempts:
 918                time.sleep(self.config.poll_interval)
 919        return False
 920
 921    def _poll(self, job_id: str, n_tasks: int, tmp: str, name: str,
 922              task_ids: Optional[Sequence[int]] = None) -> None:
 923        """Poll ``sacct`` until every task has a visible sentinel or is confirmed dead.
 924
 925        The scheduler ``State`` is not subject to NFS lag, but the ``.success`` sentinels the
 926        compute nodes write can take up to the mount's attribute-cache timeout to become
 927        visible here. So a ``COMPLETED`` task (its harness exited 0, hence wrote a sentinel)
 928        is given ``config.latency_wait`` for that sentinel to appear; any other terminal
 929        state means the harness did not succeed and the task is declared dead after a short
 930        confirmation grace. Tasks absent from ``sacct`` are pending, never dead.
 931
 932        Args:
 933            job_id: The submitted array job id.
 934            n_tasks: The total number of tasks (spans the block-counting progress bar).
 935            tmp: The job temp folder (where sentinels are written).
 936            name: A short name for the progress bar (disables it when empty).
 937            task_ids: The subset of task indices this job actually runs (a resume submits only
 938                the incomplete tasks); resolution is over this subset, ``None`` means all tasks.
 939        """
 940        poll_ids = list(range(n_tasks)) if task_ids is None else sorted(set(int(t) for t in task_ids))
 941
 942        def has_sentinel(t: int) -> bool:
 943            return os.path.exists(os.path.join(tmp, "success", f"{t}.success"))
 944
 945        latency_wait = max(float(self.config.latency_wait), self.config.poll_interval)
 946        fail_grace = max(self.config.poll_interval, 5.0)
 947        terminal_since: Dict[int, float] = {}
 948        terminal_count: Dict[int, int] = {}
 949        resolved: set = set()
 950        # The bar counts processed blocks across ALL tasks (a resume credits prior progress).
 951        n_blocks = _total_blocks(tmp, n_tasks)
 952        counter = _DoneLogCounter(tmp, n_tasks)
 953        with tqdm(total=n_blocks, desc=name or None, disable=not name, unit="block") as pbar:
 954            while len(resolved) < len(poll_ids):
 955                states = self._sacct_states(job_id)
 956                if states is None:  # transient sacct error: skip this poll.
 957                    time.sleep(self.config.poll_interval)
 958                    continue
 959
 960                now = time.monotonic()
 961                ok = {t for t in poll_ids if has_sentinel(t)}
 962                running = sum(1 for s in states.values() if s == "RUNNING")
 963                dead = set()
 964                for t in poll_ids:
 965                    if t in ok:
 966                        terminal_since.pop(t, None)
 967                        terminal_count.pop(t, None)
 968                        continue
 969                    state = states.get(t)
 970                    if state in _TERMINAL_STATES:
 971                        terminal_since.setdefault(t, now)
 972                        terminal_count[t] = terminal_count.get(t, 0) + 1
 973                        # COMPLETED -> sentinel was written, just wait it out over NFS; any
 974                        # other terminal state -> the task will never produce a sentinel.
 975                        grace = latency_wait if state == "COMPLETED" else fail_grace
 976                        if (terminal_count[t] >= 2 and now - terminal_since[t] >= grace
 977                                and not has_sentinel(t)):
 978                            dead.add(t)
 979                    else:  # pending/running/requeued: reset the dead countdown.
 980                        terminal_since.pop(t, None)
 981                        terminal_count.pop(t, None)
 982
 983                resolved = ok | dead
 984                pbar.n = min(counter.count(), n_blocks)
 985                pbar.set_postfix(ok=len(ok), failed=len(dead), run=running,
 986                                 pending=max(0, len(poll_ids) - len(resolved) - running), refresh=False)
 987                pbar.refresh()
 988                if len(resolved) >= len(poll_ids):
 989                    break
 990                try:
 991                    time.sleep(self.config.poll_interval)
 992                except KeyboardInterrupt:
 993                    print(f"\nInterrupted while waiting on slurm job {job_id}. The job was left "
 994                          f"running; reattach with SlurmRunner(...).reattach({tmp!r}).")
 995                    raise
 996
 997    def reattach(self, tmp_folder: str, name: str = "reattach",
 998                 pre_cleanup: Optional[Callable[[str], None]] = None) -> Optional[list]:
 999        """Reattach to a previously submitted run and finalize it.
1000
1001        Picks a run back up from its manifest (e.g. after the orchestrating login-node
1002        process was interrupted) instead of resubmitting. Only ``poll_interval`` is read
1003        from this runner's config, so a freshly constructed ``SlurmRunner`` can reattach.
1004
1005        Args:
1006            tmp_folder: The job temp folder containing ``manifest.json`` and ``payload.pkl``.
1007            name: A short name for the progress display.
1008            pre_cleanup: Optional ``pre_cleanup(tmp)`` callback invoked right before the temp
1009                folder is removed (forwarded to :meth:`_finalize`).
1010
1011        Returns:
1012            The per-block return values (if the run collected any), else ``None``.
1013
1014        Raises:
1015            RunnerError: If any task failed (sentinel missing).
1016            RuntimeError: If the manifest's job is unknown to slurm and the run did not
1017                already complete.
1018        """
1019        with open(os.path.join(tmp_folder, "manifest.json")) as f:
1020            manifest = json.load(f)
1021        job_id, n_tasks = str(manifest["job_id"]), int(manifest["n_tasks"])
1022        with open(os.path.join(tmp_folder, "payload.pkl"), "rb") as f:
1023            payload = cloudpickle.load(f)
1024        # A reattaching orchestrator unpickles/reduces the per-block results locally, so guard
1025        # its environment against skew from the submitting one just as the workers do.
1026        _check_versions(payload.get("versions", {"python": payload.get("python")}), role="orchestrator")
1027        has_return_val = bool(payload["has_return_val"])
1028
1029        # Reconstruct the partition in numeric task order (never glob: it sorts lexically).
1030        tasks: List[List[int]] = []
1031        for task_id in range(n_tasks):
1032            with open(os.path.join(tmp_folder, "blocks", f"{task_id}.json")) as f:
1033                tasks.append(json.load(f))
1034        block_ids = [b for task in tasks for b in task]
1035
1036        all_done = all(os.path.exists(os.path.join(tmp_folder, "success", f"{t}.success"))
1037                       for t in range(n_tasks))
1038        if not all_done:
1039            # Only a job that stays unknown to sacct across retries (not registration lag
1040            # right after submit, nor a transient error) is treated as unrecoverable.
1041            if not self._job_known(job_id):
1042                raise RuntimeError(
1043                    f"Slurm job {job_id} is not known to the scheduler and the run did not "
1044                    f"complete. Inspect {tmp_folder} or resubmit."
1045                )
1046            self._poll(job_id, n_tasks, tmp_folder, name)
1047
1048        results = self._finalize(tmp_folder, n_tasks, tasks, block_ids, has_return_val, name,
1049                                 pre_cleanup=pre_cleanup)
1050        return results if has_return_val else None