"""Functionality for single-cell analysis"""
from __future__ import annotations
import contextlib
import json
import logging
import multiprocessing
import os
import shutil
import warnings
from collections.abc import Iterable
from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait
from functools import reduce
from itertools import starmap
from multiprocessing import Pool
from pathlib import Path
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from matplotlib.figure import Figure
import numpy as np
import pandas as pd
import papermill as pm
import shapely
from pint._typing import UnitLike
from tqdm.auto import tqdm
from acia import Q_, U_
from acia.analysis._scale_progress import _WorkerBars
from acia.analysis._scale_progress import register_engine as register_progress_engine
from acia.analysis._stage_io import STAGE_NOTEBOOK_ENV, STAGE_PARAMS_ENV
from acia.analysis.growth_rate import AggMode as AggMode
from acia.analysis.growth_rate import GrowthRateResult as GrowthRateResult
from acia.analysis.growth_rate import estimate_growth_rate as estimate_growth_rate
from acia.analysis.stage import CALIBRATION_SCHEMA as CALIBRATION_SCHEMA
from acia.analysis.stage import DEFAULT_KEY_PATTERN as DEFAULT_KEY_PATTERN
from acia.analysis.stage import IO_SCHEMA as IO_SCHEMA
from acia.analysis.stage import MANIFEST_NAME as MANIFEST_NAME
from acia.analysis.stage import STAGE_SCHEMA as STAGE_SCHEMA
from acia.analysis.stage import StageContext as StageContext
from acia.analysis.stage import check_stale as check_stale
from acia.analysis.stage import population_id_of as population_id_of
from acia.analysis.stage import read_manifest as read_manifest
from acia.analysis.stage import stage_graph as stage_graph
from acia.analysis.stage import stage_table as stage_table
from acia.analysis.stage import stages_run as stages_run
from acia.analysis.units import UNIT_ATTR, units_in_header
from acia.analysis.units import attach_units as attach_units
from acia.analysis.units import from_header as from_header
from acia.analysis.units import read_units_csv as read_units_csv
from acia.analysis.units import strip_units as strip_units
from acia.analysis.units import write_units_csv as write_units_csv
from acia.base import BaseImage, ImageSequenceSource, Overlay
from acia.segm.rasterize import contour_pixels
logger = logging.getLogger(__name__)
DEFAULT_UNIT_LENGTH = "micrometer"
DEFAULT_UNIT_AREA = "micrometer ** 2"
def _polygons(overlay: Overlay) -> list:
"""Every contour's polygon, in overlay order."""
return [cont.polygon for cont in overlay]
def _rect_edge_lengths(polygons: list) -> np.ndarray:
"""``(n, 4)`` edge lengths of each polygon's minimum rotated rectangle.
Computes the rectangles for the whole overlay in one shapely call rather
than one per contour per extractor, which had ``LengthEx`` and ``WidthEx``
each deriving the same rectangle independently.
A degenerate outline (collinear, zero-extent, or absent) has a rectangle
that collapses to a ``LineString``/``Point`` with no ``exterior``, so it has
no length or width to report and measures ``nan``. Reporting 0 instead would
be a claim -- and a 0-length cell passes any open-below bound, letting junk
through -- whereas ``nan`` states that the property is undefined, and the
filters drop what they cannot measure.
"""
if not polygons:
return np.empty((0, 4), dtype=float)
rectangles = shapely.oriented_envelope(np.array(polygons, dtype=object))
# a rectangle ring is 5 coordinates (closed); anything else is degenerate
edges = np.full((len(polygons), 4), np.nan, dtype=float)
intact = shapely.get_num_coordinates(rectangles) == 5
if intact.any():
coords = shapely.get_coordinates(rectangles[intact]).reshape(-1, 5, 2)
edges[intact] = np.linalg.norm(np.diff(coords, axis=1), axis=2)
return edges
def _value_frame(overlay: Overlay, name: str, values: np.ndarray) -> pd.DataFrame:
"""``id``-indexed single-column frame, float-typed even when empty."""
return pd.DataFrame(
{"id": [cont.id for cont in overlay], name: np.asarray(values, dtype=float)}
).set_index("id")
def _id_indexed(records: list[dict], columns: list[str]) -> pd.DataFrame:
"""Build an ``id``-indexed DataFrame from per-contour records, empty-safe.
``pd.DataFrame([])`` has no columns, so ``.set_index("id")`` raises for an
empty overlay; passing explicit ``columns`` keeps ``id`` and the value
column(s) present (as an empty frame) instead.
"""
return pd.DataFrame(records, columns=["id", *columns]).set_index("id")
[docs]
class AreaEx(PropertyExtractor):
"""Extract area for every contour"""
_dim = 2
[docs]
def __init__(
self,
input_unit: UnitLike | None = None,
output_unit: UnitLike | None = None,
):
self._auto_unit = input_unit is None
PropertyExtractor.__init__(
self,
"area",
input_unit=input_unit if input_unit is not None else DEFAULT_UNIT_AREA,
output_unit=output_unit if output_unit is not None else DEFAULT_UNIT_AREA,
)
[docs]
class PerimeterEx(PropertyExtractor):
"""Extract area for every contour"""
_dim = 1
[docs]
def __init__(
self,
input_unit: UnitLike | None = None,
output_unit: UnitLike | None = None,
):
self._auto_unit = input_unit is None
PropertyExtractor.__init__(
self,
"perimeter",
input_unit=input_unit if input_unit is not None else DEFAULT_UNIT_LENGTH,
output_unit=output_unit if output_unit is not None else DEFAULT_UNIT_LENGTH,
)
[docs]
class CircularityEx(PropertyExtractor):
"""Extract area for every contour"""
[docs]
def __init__(
self,
input_unit: UnitLike | None = "1",
output_unit: UnitLike | None = "1",
):
PropertyExtractor.__init__(
self, "circularity", input_unit=input_unit, output_unit=output_unit
)
[docs]
class BoundaryClosenessEx(PropertyExtractor):
"""Distance from each cell's bounding box to the nearest image border.
The backing property for
:class:`~acia.segm.filter.BoundaryClosenessFilter`, which filters on how
close a cell sits to the edge of the field of view (a cell that is partly
outside it has unreliable size and shape). Extracting it as a column means
the filter reads it like every other filter reads its own property, and it
becomes plottable next to them in
:func:`~acia.analysis.properties.plot_property_histograms`.
The frame extent comes from the source (``size_w`` / ``size_h``), so a cell
whose bounding box touches a border measures 0.
"""
_dim = 1
[docs]
def __init__(
self,
input_unit: UnitLike | None = None,
output_unit: UnitLike | None = None,
):
self._auto_unit = input_unit is None
PropertyExtractor.__init__(
self,
"boundary_closeness",
input_unit=input_unit if input_unit is not None else DEFAULT_UNIT_LENGTH,
output_unit=output_unit if output_unit is not None else DEFAULT_UNIT_LENGTH,
)
[docs]
class LengthEx(PropertyExtractor):
"""Extracts width of cells based on the shorter edge of a minimum rotated bbox approximation"""
_dim = 1
[docs]
def __init__(
self,
input_unit: UnitLike | None = None,
output_unit: UnitLike | None = None,
):
self._auto_unit = input_unit is None
PropertyExtractor.__init__(
self,
"length",
input_unit=input_unit if input_unit is not None else DEFAULT_UNIT_LENGTH,
output_unit=output_unit if output_unit is not None else DEFAULT_UNIT_LENGTH,
)
[docs]
class WidthEx(PropertyExtractor):
"""Extracts width of cells based on the shorter edge of a minimum rotated bbox approximation"""
_dim = 1
[docs]
def __init__(
self,
input_unit: UnitLike | None = None,
output_unit: UnitLike | None = None,
):
self._auto_unit = input_unit is None
PropertyExtractor.__init__(
self,
"width",
input_unit=input_unit if input_unit is not None else DEFAULT_UNIT_LENGTH,
output_unit=output_unit if output_unit is not None else DEFAULT_UNIT_LENGTH,
)
[docs]
class LengthWidthEx(PropertyExtractor):
"""Extracts length and width of cells based on the shorter edge of a minimum rotated bbox approximation"""
_dim = 1
[docs]
def __init__(
self,
prefix="",
input_unit: UnitLike | None = None,
output_unit: UnitLike | None = None,
):
self.prefix = prefix
self._auto_unit = input_unit is None
PropertyExtractor.__init__(
self,
f"{prefix}length-width",
input_unit=input_unit if input_unit is not None else DEFAULT_UNIT_LENGTH,
output_unit=output_unit if output_unit is not None else DEFAULT_UNIT_LENGTH,
)
[docs]
class FrameEx(PropertyExtractor):
"""Extract the frame information for every contour"""
[docs]
def __init__(self):
super().__init__("frame", 1)
[docs]
class IdEx(PropertyExtractor):
"""Extract single-cell id for every contour"""
[docs]
def __init__(self):
super().__init__("id", 1)
[docs]
class LabelEx(PropertyExtractor):
"""Extract single-cell label (from tracking) for every contour"""
[docs]
def __init__(self):
super().__init__("label", 1)
[docs]
class TimeEx(PropertyExtractor):
"""Extract time information for every contour.
If no ``input_unit`` is given, the per-frame timepoints are taken from the
image source (or the overlay) calibration -- so a source loaded with a
``frame_interval`` (or sliced) yields correct, automatically-updated times.
Passing ``input_unit`` keeps the legacy ``frame * interval`` behavior.
"""
[docs]
def __init__(
self, input_unit: UnitLike | None = None, output_unit: UnitLike | None = "hour"
):
self._auto_unit = input_unit is None
super().__init__(
"time", input_unit if input_unit is not None else "second", output_unit
)
[docs]
class DynamicTimeEx(PropertyExtractor):
"""Extract time information for every contour when timepoints are not equi-distant"""
[docs]
def __init__(
self,
timepoints: list,
relative=True,
input_unit: UnitLike = "second",
output_unit: UnitLike | None = "hour",
):
super().__init__("time", input_unit, output_unit)
if len(timepoints) == 0:
raise ValueError("Need non-empty timepoint list")
self.timepoints = np.array(timepoints)
if relative:
self.timepoints -= self.timepoints[0]
[docs]
class PositionEx(PropertyExtractor):
"""Extract cell center information from image RoI detections"""
_dim = 1
[docs]
def __init__(
self,
input_unit: UnitLike | None = None,
output_unit: UnitLike | None = DEFAULT_UNIT_LENGTH,
):
self._auto_unit = input_unit is None
super().__init__(
"position",
input_unit=input_unit if input_unit is not None else DEFAULT_UNIT_LENGTH,
output_unit=output_unit,
)
[docs]
class FluorescenceEx(PropertyExtractor):
"""Extracting fluorescence properties from image sequence and RoI detections"""
[docs]
def __init__(
self,
channels,
channel_names,
summarize_operator=np.median,
input_unit: UnitLike = "1",
output_unit: UnitLike | None = "",
parallel=6,
):
super().__init__("Fluorescence", input_unit=input_unit, output_unit=output_unit)
self.channels = channels
self.channel_names = channel_names
self.summarize_operator = summarize_operator
self.parallel = parallel
assert len(self.channels) == len(self.channel_names), (
"Number of channels and number of channel names must comply"
)
[docs]
def default_execution_naming(source) -> str:
"""Source-aware default folder name for one scaled execution.
* ``int`` (e.g. an OMERO image id) -> ``"execution_<id>"``
* ``str`` (a file path or fsspec URL) -> the file *stem*, i.e. the file name
without directory and extension (``smb://host/share/pos1.tif`` -> ``pos1``)
For other item types (e.g. a parameter ``dict``) the name cannot be inferred;
pass an explicit ``execution_naming`` to :func:`scale` in that case.
"""
if isinstance(source, bool): # bool is an int subclass; treat as unsupported
raise ValueError(f"Cannot derive an execution name from {source!r}.")
if isinstance(source, int):
return f"execution_{source}"
if isinstance(source, str):
# strip a possible query string, then take the file stem (works for URLs)
return Path(source.split("?", 1)[0]).stem
raise ValueError(
f"Cannot derive a default execution name from {type(source).__name__}. "
"Please pass an explicit `execution_naming` function to scale()."
)
def _source_label(image_id, name: str) -> str:
"""Short, human-readable label for one scaled source, for progress output.
Path/URL entries show the input **file name** (``/data/pos1_roi2.tiff`` ->
``pos1_roi2.tiff``); ids and parameter dicts fall back to the execution folder
name, which is the only thing that identifies them.
"""
if isinstance(image_id, str):
return Path(image_id.split("?", 1)[0]).name or name
return name
# jupyter_client raises a bare RuntimeError with one of these messages when a
# kernel never becomes ready (KernelClient.wait_for_ready). All three mean the
# handshake failed *before* any cell ran, so re-executing the notebook repeats no
# work and has no side effect to duplicate -- which is what makes retrying safe.
# A kernel that dies mid-notebook raises DeadKernelError instead and is NOT
# retried here, because cells have already run.
_KERNEL_START_FAILURES = (
"Kernel died before replying to kernel_info",
"Kernel didn't respond in",
"Kernel didn't respond to heartbeats in",
)
def _is_kernel_start_failure(exc: BaseException) -> bool:
"""Did this exception come from a kernel that never started?
Matched on the message because jupyter_client raises a plain ``RuntimeError``
for all three cases rather than a dedicated exception class.
"""
if not isinstance(exc, RuntimeError):
return False
return any(msg in str(exc) for msg in _KERNEL_START_FAILURES)
def _execute_notebook(*args, kernel_start_retries: int = 1, **kwargs) -> None:
"""``papermill.execute_notebook`` that retries a kernel which never started.
Starting a kernel is occasionally flaky under load -- observed in CI as
``RuntimeError: Kernel died before replying to kernel_info`` on one image of a
batch. Nothing has executed at that point, so one more attempt is cheap and
turns a lost stage into a slow one.
"""
for attempt in range(kernel_start_retries + 1):
try:
pm.execute_notebook(*args, **kwargs)
return
except Exception as exc:
if attempt >= kernel_start_retries or not _is_kernel_start_failure(exc):
raise
logger.warning(
"The kernel for %s failed to start (%s); retrying (attempt %d/%d).",
args[0] if args else kwargs.get("input_path"),
exc,
attempt + 2,
kernel_start_retries + 1,
)
def _record_failure(failed_ids: list, job: dict) -> None:
"""Isolate one source's failure from the rest of the batch.
Every exception is isolated, not only ``PapermillExecutionError``. That one
means "a cell raised"; a kernel that fails to start raises a bare
``RuntimeError``, and letting it escape aborted the whole batch -- the
remaining sources were never attempted, which is exactly what running each
source independently is supposed to prevent.
The exception is logged here because ``failed_ids`` only carries the id: the
summary at the end of :func:`scale` can say how many sources failed, never why.
"""
logger.warning(
"Scaling failed for %s; continuing with the remaining sources.",
job["label"],
exc_info=True,
)
failed_ids.append(job["image_id"])
def _scale_execute_one(
job,
additional_parameters,
exist_ok,
exist_skip,
kernel_name,
storage_parameter_name,
on_stage=None,
stage_progress="keep",
progress_queue=None,
):
"""Run every analysis script for one image entry.
Module-level (not a closure) so it is picklable for the ``spawn`` process pool.
``job`` is a dict with ``image_id``, ``output_parent`` (str), ``source_parameters``
(dict) and ``scripts`` (list of ``(src, dst)`` path-string pairs). Returns the
list of execution records; raises ``PapermillExecutionError`` if a notebook fails.
``on_stage`` (if given) is called with the name of each stage notebook just
before it is executed, so the caller can report what is running. Only used on
the sequential path -- the parallel one runs this in a child process, from
which the parent's progress bar is not reachable.
How the per-stage cell progress is displayed depends on where we run:
* no ``progress_queue`` (sequential) -- papermill draws its own bar in this
process, labelled with the stage notebook and the source it runs on.
* with a ``progress_queue`` (parallel) -- nothing is drawn here at all; cell
progress is reported to the parent, which renders one bar per pool worker.
Children must not draw: they share one stderr with no shared cursor, so
their bars would overwrite each other (see :mod:`acia.analysis._scale_progress`).
``stage_progress`` is :func:`scale`'s parameter of the same name.
"""
output_parent = Path(job["output_parent"])
os.makedirs(output_parent, exist_ok=exist_ok)
engine_name = None
progress_key = None
if progress_queue is not None and stage_progress != "off":
engine_name = register_progress_engine()
progress_key = os.getpid()
else:
progress_queue = None
executions = []
for src, dst in job["scripts"]:
dst = Path(dst)
# the notebook already exists and we should skip it
if dst.exists() and exist_skip:
continue
if on_stage is not None:
on_stage(dst.name)
if progress_queue is not None:
# report before the copy + kernel start, so the worker's bar shows what
# it is about to run instead of staying blank for the kernel boot
progress_queue.put(
("stage_start", progress_key, job["label"], dst.name),
)
shutil.copy(src, dst)
# parameters to integrate into notebook -- inject the execution folder
# under `storage_parameter_name` unless it is None (some notebooks derive
# their output location differently and don't declare `storage_folder`,
# which papermill would otherwise warn about as an unknown parameter).
parameters = {}
if storage_parameter_name is not None:
parameters[storage_parameter_name] = str(dst.parent.absolute())
parameters.update(job["source_parameters"])
parameters.update(additional_parameters)
# execute the notebook (its own kernel subprocess; cwd is this process's,
# which is safe because each worker is a separate process)
extra: dict[str, Any] = {}
if progress_queue is not None:
extra = {
"progress_bar": False,
"engine_name": engine_name,
"progress_queue": progress_queue,
"progress_key": progress_key,
}
elif stage_progress == "off":
extra = {"progress_bar": False}
else:
# papermill takes a dict of tqdm kwargs; indent the label so the bar
# reads as "cells of this stage" under scale()'s own source bar
extra = {
"progress_bar": {
"desc": f" ↳ {dst.name} | {job['label']}",
"leave": stage_progress == "keep",
}
}
# Tell the kernel which notebook it is running, so StageContext can record
# the code that produced a result (a kernel cannot know this otherwise).
# An env var rather than a papermill parameter: it needs no `parameters`
# cell in the notebook and raises no "unknown parameter" warning.
#
# The *source* notebook, not the copy being executed: papermill writes
# outputs back into `dst` as it runs, so every execution folder's copy has
# different bytes. Hashing those would make the digest differ per image and
# destroy the one question it exists to answer -- did this batch all run the
# same code?
os.environ[STAGE_NOTEBOOK_ENV] = str(src)
# ... and which parameters it was given, so StageContext records the run's
# real settings without every notebook having to repeat them into
# log_params(). Same reasoning as the line above: an env var needs no
# `parameters` cell and raises no "unknown parameter" warning.
with contextlib.suppress(TypeError, ValueError):
os.environ[STAGE_PARAMS_ENV] = json.dumps(parameters, default=str)
try:
_execute_notebook(
dst,
dst,
parameters=parameters,
cwd=dst.parent,
kernel_name=kernel_name,
**extra,
)
finally:
os.environ.pop(STAGE_PARAMS_ENV, None)
# papermill writes the executed notebook back into `dst` only after it
# returns, so this cannot be done from inside the notebook -- the last
# cell never sees the finished document
_export_notebook_html(dst)
if progress_queue is not None:
# papermill closes the reporter itself, but not when it fails
# before the notebook manager exists -- make sure the parent's
# worker bar is always released
progress_queue.put(("stage_end", progress_key))
executions.append(dict(parameters=parameters, storage_folder=dst.parent))
return executions
def _export_notebook_html(notebook: Path) -> Path | None:
"""Write an HTML rendering of an executed notebook beside it.
A run's own record of what it looked like, readable without Jupyter. Entirely
best-effort: ``nbconvert`` is an optional dependency and a failed export must
never be the reason a completed analysis reports failure.
"""
try:
from nbconvert import HTMLExporter # noqa: PLC0415 -- optional dependency
body, _ = HTMLExporter().from_filename(str(notebook))
target = notebook.with_suffix(".html")
target.write_text(body, encoding="utf-8")
return target
except Exception:
logger.debug("could not export %s to HTML", notebook, exc_info=True)
return None
[docs]
def scale(
output_path: Path,
analysis_script: Path | list[Path],
image_ids: list[int | str | dict],
additional_parameters=None,
exist_ok=False,
execution_naming=None,
exist_skip=False,
kernel_name=None,
parameter_name: str = "image_id",
max_workers: int = 1,
storage_parameter_name: str | None = "storage_folder",
stage_progress: str = "keep",
):
"""Scale an analysis notebook to several image sources.
Each entry in ``image_ids`` identifies one image source and triggers one
notebook execution. An entry may be:
* an ``int`` -- e.g. an OMERO image id (default folder ``execution_<id>``),
* a ``str`` -- a local path or fsspec URL such as ``smb://host/share/x.tif``
(default folder name is the file stem, e.g. ``x``),
* a ``dict`` -- arbitrary parameters merged into the notebook; provide an
explicit ``execution_naming`` for these.
The identifier is injected into the notebook under ``parameter_name``
(default ``"image_id"``), so existing notebooks keep working. The notebook is
responsible for turning it into a concrete source (e.g.
``OmeroSequenceSource(image_id)`` or ``SambaSequenceSource.from_url(image_id)``).
**Hint:** the analysis script should only use absolute paths as the file is copied and executed in another folder.
Args:
output_path (Path): the general output path to the storage
analysis_script (Path): the template script
image_ids (list[int | str | dict]): image sources to scale over (ids,
paths/URLs, or parameter dicts).
additional_parameters (dict): Parameters to be inserted into the jupyter script
exist_ok (Bool): True when it is okay that the directory exists, False will throw an error when the directory exists.
execution_naming (Callable): maps an entry to its output folder name. By
default :func:`default_execution_naming` is used, which dispatches on
the entry type (id -> ``execution_<id>``, path -> file stem).
exist_skip (Bool): If true existing executions are skipped.
kernel_name (str): specifies the notebook kernel to be used. None is the default kernel.
parameter_name (str): name of the notebook parameter the identifier is
injected as (ignored for ``dict`` entries, which are merged as-is).
storage_parameter_name (str | None): the notebook parameter the per-run
output/execution folder (absolute) is injected under. Defaults to
``"storage_folder"``. Set it to match the notebook's own output
parameter, or to ``None`` to not inject it at all (avoids papermill's
"Passed unknown parameter" warning for notebooks that derive their
output location some other way).
max_workers (int): how many notebooks to execute concurrently. ``1``
(default) runs them sequentially, exactly as before. Values > 1 run
that many notebooks in parallel using a **process** pool started with
the ``"spawn"`` method (not threads: papermill sets the working
directory with a process-global ``os.chdir`` that threads would race
on; and not ``fork``: spawning fresh processes avoids duplicating a
CUDA-initialised parent kernel -- the classic Jupyter crash). Each
execution is still its own kernel subprocess, so a worker whose kernel
dies only fails its own image. On a **single GPU** keep this small
(2-3): every concurrent run loads its own model, so throughput is
bounded by GPU memory, not CPU cores.
stage_progress (str): how the per-notebook cell progress is shown --
``"keep"`` (default) leaves the finished bars on screen as a per-stage
timing log, ``"collapse"`` removes each bar once its stage is done, and
``"off"`` shows only the source bar. See below for what the two levels
look like.
**Progress output.** There are always two levels: one bar counting *sources*,
and per-notebook bars counting *cells*. Running sequentially, the source bar is
labelled with what is running right now and papermill's own bars sit underneath
it::
01_Segment.ipynb | pos001_roi001.tiff: 33%|███| 1/3 [04:41<09:23, 281.63s/source]
↳ 01_Segment.ipynb | pos001_roi001.tiff: 100%|███| 21/21 [01:21<00:00, 3.88s/cell]
With ``max_workers > 1`` there is no single "current" source, so the source bar
reports the one that just finished and each worker gets its own bar instead::
Sources: 33%|███| 1/3 [04:41<09:23, 281.63s/source]
[w1] pos001_roi001.tiff | 02_Track.ipynb: 57%|███| 12/21 [00:32<00:35, 3.91s/cell]
[w2] pos002_roi001.tiff | 01_Segment.ipynb: 19%|█ | 4/21 [00:00<00:03, 4.88cell/s]
Those worker bars are drawn by *this* process from progress the workers report
over a queue -- children never write bars themselves, because they share one
stderr with no shared cursor and would overwrite each other.
"""
if execution_naming is None:
execution_naming = default_execution_naming
if isinstance(analysis_script, str):
# if this is just a single string, then we make it a list of a single path
analysis_script = [Path(analysis_script)]
elif isinstance(analysis_script, Path):
analysis_script = [analysis_script]
elif isinstance(analysis_script, Iterable):
analysis_script = list(map(Path, analysis_script))
for script in analysis_script:
if not script.exists():
raise ValueError(f"Analysis script {script} does not exist!")
if additional_parameters is None:
additional_parameters = {}
experiment_executions = []
failed_ids: list[int | str | dict] = []
# warn (don't fail) if two entries map to the same output folder, since later
# executions would otherwise silently overwrite earlier ones
seen_names: dict[str, object] = {}
for entry in image_ids:
if isinstance(entry, dict):
continue # naming for dicts is user-defined; skip the heuristic check
name = execution_naming(entry)
if name in seen_names:
logger.warning(
"Execution name %r maps to multiple sources (%r and %r); their "
"outputs will collide. Pass a unique `execution_naming` to avoid this.",
name,
seen_names[name],
entry,
)
seen_names[name] = entry
if max_workers < 1:
raise ValueError(f"max_workers must be >= 1, got {max_workers}")
if stage_progress not in ("keep", "collapse", "off"):
raise ValueError(
"stage_progress must be one of 'keep', 'collapse', 'off'; "
f"got {stage_progress!r}"
)
# Resolve naming + per-entry parameters up front (in this process, so a lambda
# `execution_naming` still works). Each job is fully picklable for the pool.
jobs: list[dict[str, Any]] = []
for image_id in image_ids:
# per-entry notebook parameters: dict entries are merged as-is,
# everything else is injected under `parameter_name`.
if isinstance(image_id, dict):
source_parameters = dict(image_id)
else:
source_parameters = {parameter_name: image_id}
name = execution_naming(image_id)
output_parent = output_path / name
jobs.append(
{
"image_id": image_id,
"label": _source_label(image_id, name),
"output_parent": str(output_parent),
"source_parameters": source_parameters,
"scripts": [
(str(s), str(output_parent / s.name)) for s in analysis_script
],
}
)
if max_workers == 1:
# sequential (default) -- the bar reports the notebook currently running
# and the source it runs on, e.g. "02_Track.ipynb | pos001_roi002.tiff".
with tqdm(jobs, unit="source") as pbar:
for job in pbar:
pbar.set_description(job["label"])
try:
experiment_executions.extend(
_scale_execute_one(
job,
additional_parameters,
exist_ok,
exist_skip,
kernel_name,
storage_parameter_name,
on_stage=lambda stage, label=job["label"]: (
pbar.set_description(f"{stage} | {label}")
),
stage_progress=stage_progress,
)
)
except Exception:
_record_failure(failed_ids, job)
else:
# Run several notebooks in parallel. A *process* pool (spawn) is required,
# not threads: papermill sets the working directory with a process-global
# os.chdir, which threads would race on (mislocating each run's relative
# outputs). spawn also avoids duplicating a CUDA-initialised parent kernel.
# Failures are collected here in the parent as each future completes.
ctx = multiprocessing.get_context("spawn")
with contextlib.ExitStack() as stack:
# The workers report their progress here instead of drawing bars of
# their own (see the "Progress output" note above). A Manager queue is
# used because a plain mp.Queue cannot be passed as a submit() argument.
progress_queue = None
if stage_progress != "off":
try:
manager = stack.enter_context(ctx.Manager())
progress_queue = manager.Queue()
except Exception: # pragma: no cover - environment dependent
logger.warning(
"Could not start the progress manager; falling back to the "
"source bar only.",
exc_info=True,
)
# several sources run at once, so "currently running" is not a single
# thing on the source bar: it reports the one that just finished (and
# whether it failed). What each worker does right now is on its own bar.
bars = _WorkerBars(total_sources=len(jobs), mode=stage_progress)
pool = stack.enter_context(
ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx)
)
futures = {
pool.submit(
_scale_execute_one,
job,
additional_parameters,
exist_ok,
exist_skip,
kernel_name,
storage_parameter_name,
None,
stage_progress,
progress_queue,
): job
for job in jobs
}
try:
pending = set(futures)
while pending:
# a short timeout instead of as_completed() so the worker bars
# are drained and redrawn while nothing completes
done, pending = wait(
pending, timeout=0.2, return_when=FIRST_COMPLETED
)
if progress_queue is not None:
bars.drain(progress_queue)
for future in done:
job = futures[future]
status = "done"
try:
experiment_executions.extend(future.result())
except Exception:
_record_failure(failed_ids, job)
status = "FAILED"
bars.advance_total(f"{status} {job['label']}")
bars.tick()
finally:
bars.close()
if len(failed_ids) > 0:
error_ratio = len(failed_ids) / len(image_ids) * 100
logging.warning(
"The scaling failed in %d/%d (%.3f%%) executions. Please report failes with the link to the script and the image id to your administrator in order to further improve the software.",
len(failed_ids),
len(image_ids),
error_ratio,
)
if error_ratio > 10:
# error rates of more than 10% are definitively acceptable
logging.error("Such a high error rate is not acceptable!")
return experiment_executions