From 94a070dbf62bdf29280a79803742bf9d862ffa04 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 11:14:32 -0700 Subject: [PATCH 01/20] feat(dynacell): add eval runtime module + thread budgeting (C1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds dynacell.evaluation.runtime with three layered thread-cap entry points: early_apply_env_caps() sets BLAS/OMP env vars before C-extension load (called from __main__.main_cli); apply_thread_budget(threads) is the in-process safety net (env + torch.set_num_threads + threadpool_limits); resolve_runtime(config) materializes the new runtime config block. The runtime block (eval.yaml) is opt-in for parallelism (default executor=serial, fov_workers=1) but unconditionally tightens BLAS/OMP threads to cpu_count by default. On SLURM nodes where --cpus-per-task < host CPU count, this corrects previously-loose BLAS threading. Set DYNACELL_THREADS_PER_WORKER=N to override at the env layer for cases where in-process caps come too late. precompute_gt_artifacts validates that runtime.executor='serial' and fov_workers=1 — FOV-level parallelism is restricted to evaluate_predictions in this commit; precompute parallelism is out of scope (DeepFeatureBatcher accumulates state across FOVs). resolve_runtime supports two-phase resolution for evaluate_predictions: top-of-function for parent BLAS cap, then post-position-list with freeze_threads_per_worker= to keep parent and worker BLAS budgets consistent. Worker primitives (_process_one_fov, ProcessPoolExecutor) land in C3+C4. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../dynacell/src/dynacell/__main__.py | 8 + .../dynacell/evaluation/_configs/eval.yaml | 21 ++ .../src/dynacell/evaluation/pipeline.py | 8 + .../src/dynacell/evaluation/precompute_cli.py | 14 + .../src/dynacell/evaluation/runtime.py | 279 ++++++++++++++++++ 5 files changed, 330 insertions(+) create mode 100644 applications/dynacell/src/dynacell/evaluation/runtime.py diff --git a/applications/dynacell/src/dynacell/__main__.py b/applications/dynacell/src/dynacell/__main__.py index 86eadeac5..3193dde87 100644 --- a/applications/dynacell/src/dynacell/__main__.py +++ b/applications/dynacell/src/dynacell/__main__.py @@ -106,6 +106,14 @@ def _inject_external_configs(argv: list[str]) -> list[str]: def main_cli(): """Console script entry point for ``dynacell`` command.""" + # Apply BLAS/OMP env caps BEFORE any torch/numpy import. The import below + # is stdlib + threadpoolctl only — no torch — so the env vars set here + # bite at first BLAS C-extension load, which happens transitively when + # we import the Hydra command module or viscy_utils.cli below. + from dynacell.evaluation.runtime import early_apply_env_caps + + early_apply_env_caps() + if len(sys.argv) >= 2 and sys.argv[1] in _HYDRA_COMMANDS: command = sys.argv[1] module_path, func_name, extra = _HYDRA_COMMANDS[command] diff --git a/applications/dynacell/src/dynacell/evaluation/_configs/eval.yaml b/applications/dynacell/src/dynacell/evaluation/_configs/eval.yaml index 2decff3f1..11698c395 100644 --- a/applications/dynacell/src/dynacell/evaluation/_configs/eval.yaml +++ b/applications/dynacell/src/dynacell/evaluation/_configs/eval.yaml @@ -77,6 +77,27 @@ compute_microssim: true compute_feature_metrics: false limit_positions: null # process first N positions; null means all +# Eval runtime / parallelism. Default is bit-for-bit-current sequential behavior. +# See applications/dynacell/CLAUDE.md "Eval runtime / parallelism" section. +runtime: + # Outer FOV-level parallelism. 1 = current sequential behavior. + # "auto" + executor=process clamps to min(cpu_count//4, n_positions). + # "auto" + executor=serial resolves to 1. + fov_workers: 1 # int | "auto" + + # BLAS/OMP/torch intra-op threads per worker process. + # "auto" → max(1, cpu_count // resolved_fov_workers). + threads_per_worker: "auto" # int | "auto" + + # FOV-loop executor. "serial" runs inline; "process" uses spawn-context + # ProcessPoolExecutor. Auto-demoted to "serial" if fov_workers resolves to 1. + executor: "serial" # "serial" | "process" + + # Per-T memory hygiene. Defaults set by C2 measurement (see commit log). + # Override at runtime with DYNACELL_FORCE_PER_T_HYGIENE=1 (operator escape hatch). + cuda_empty_cache_every_n_timepoints: 0 # 0 = off + gc_collect_every_n_fovs: 0 # 0 = off + force_recompute: all: false gt_masks: false diff --git a/applications/dynacell/src/dynacell/evaluation/pipeline.py b/applications/dynacell/src/dynacell/evaluation/pipeline.py index 88f6adebf..9a22e57ce 100644 --- a/applications/dynacell/src/dynacell/evaluation/pipeline.py +++ b/applications/dynacell/src/dynacell/evaluation/pipeline.py @@ -155,9 +155,17 @@ def _save_embeddings(save_dir: Path, groups: dict[str, tuple[list, list, list]]) def evaluate_predictions(config: DictConfig): """Evaluate predictions on all test images.""" + from dynacell.evaluation.runtime import apply_thread_budget, resolve_runtime from dynacell.evaluation.segmentation import prepare_segmentation_model, segment from dynacell.evaluation.utils import CellDinoFeatureExtractor, DinoV3FeatureExtractor, DynaCLRFeatureExtractor + # Phase 1 runtime resolution: lock in executor + thread caps before any + # heavy work. fov_workers may be provisional when "auto"; re-resolved in + # Phase 2 once the position list is known (C4). threads_per_worker stays + # frozen across phases for parent/worker BLAS-cap consistency. + runtime = resolve_runtime(config) + apply_thread_budget(runtime.threads_per_worker) + all_pixel_metrics = [] all_mask_metrics = [] all_feature_metrics = [] diff --git a/applications/dynacell/src/dynacell/evaluation/precompute_cli.py b/applications/dynacell/src/dynacell/evaluation/precompute_cli.py index 2624b6cfa..cd1861bfc 100644 --- a/applications/dynacell/src/dynacell/evaluation/precompute_cli.py +++ b/applications/dynacell/src/dynacell/evaluation/precompute_cli.py @@ -33,12 +33,26 @@ def precompute_gt_artifacts(config: DictConfig) -> None: """Build every GT-side artifact toggled on in ``config.build``.""" + from dynacell.evaluation.runtime import apply_thread_budget, resolve_runtime from dynacell.evaluation.segmentation import prepare_segmentation_model from dynacell.evaluation.utils import CellDinoFeatureExtractor, DinoV3FeatureExtractor, DynaCLRFeatureExtractor if config.io.gt_cache_dir is None: raise ValueError("io.gt_cache_dir is required for dynacell precompute-gt") + # Precompute is single-process by design (DeepFeatureBatcher accumulates + # state across FOVs). Apply the thread cap but raise if the user requested + # FOV-level parallelism here — that belongs to evaluate_predictions only. + runtime = resolve_runtime(config) + if runtime.executor != "serial" or runtime.fov_workers != 1: + raise ValueError( + "dynacell precompute-gt does not support FOV-level parallelism. " + f"Got runtime.executor={runtime.executor!r}, " + f"runtime.fov_workers={runtime.fov_workers}. " + "Set runtime.executor='serial' and runtime.fov_workers=1 (or omit)." + ) + apply_thread_budget(runtime.threads_per_worker) + build = config.build build_any_features = bool(build.cp or build.dinov3 or build.dynaclr or build.celldino) diff --git a/applications/dynacell/src/dynacell/evaluation/runtime.py b/applications/dynacell/src/dynacell/evaluation/runtime.py new file mode 100644 index 000000000..45b58f558 --- /dev/null +++ b/applications/dynacell/src/dynacell/evaluation/runtime.py @@ -0,0 +1,279 @@ +"""Eval runtime: thread budgeting + (in later commits) process-pool primitives. + +Three layers of thread-cap discipline, in order of when they bite: + +1. ``early_apply_env_caps()`` reads ``DYNACELL_THREADS_PER_WORKER`` from the + environment and sets BLAS/OMP env vars before any C extension loads. Called + from ``dynacell.__main__:main_cli()`` as the first statement. +2. ``apply_thread_budget(threads)`` is the in-process safety net: sets env + (respecting caller-set values), calls ``torch.set_num_threads``, and + activates a module-level ``threadpoolctl.threadpool_limits`` cap. +3. (C3) per-worker initializer re-applies the cap in each spawned child. + +Module-level imports are stdlib + ``threadpoolctl`` only. ``torch`` is imported +lazily inside function bodies so spawn-context workers can set +``CUDA_VISIBLE_DEVICES`` before any torch-driven CUDA context creation. +""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass +from typing import Any, Literal + +from omegaconf import DictConfig, OmegaConf +from threadpoolctl import threadpool_limits + +logger = logging.getLogger(__name__) + +# Module-level handle for the active threadpool_limits context. We enter the +# context once at startup and keep the handle live for the rest of the process +# lifetime (threadpool_limits is a per-process thread-local API; the cap stays +# armed as long as the context object is not exited). +_ACTIVE_THREADPOOL_HANDLE: Any = None + +# Env vars that affect BLAS / OMP / framework thread counts. Listed in +# priority order; the first three actually matter for this codebase (numpy is +# scipy-openblas, torch + cellpose honor OMP), the rest are defense-in-depth. +_THREAD_ENV_VARS: tuple[str, ...] = ( + "OMP_NUM_THREADS", + "MKL_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "BLIS_NUM_THREADS", + "NUMEXPR_MAX_THREADS", + "VECLIB_MAXIMUM_THREADS", + "TBB_NUM_THREADS", +) + +# Env var that triggers the per-T memory hygiene knobs at runtime even when +# the YAML defaults are off. Useful for operator mitigation post-ship. +_FORCE_PER_T_HYGIENE_ENV: str = "DYNACELL_FORCE_PER_T_HYGIENE" + + +@dataclass(frozen=True) +class ResolvedRuntime: + """Materialized runtime config with all ``"auto"`` values resolved. + + Parameters + ---------- + fov_workers : int + Outer FOV-level parallelism. 1 = sequential. + threads_per_worker : int + BLAS/OMP/torch intra-op threads per worker process. + executor : {"serial", "process"} + FOV loop executor. + cuda_empty_cache_every_n_timepoints : int + Call ``torch.cuda.empty_cache()`` every N timepoints (0 = off). + gc_collect_every_n_fovs : int + Call ``gc.collect()`` every N FOVs (0 = off). + """ + + fov_workers: int + threads_per_worker: int + executor: Literal["serial", "process"] + cuda_empty_cache_every_n_timepoints: int + gc_collect_every_n_fovs: int + + +def early_apply_env_caps() -> None: + """Set BLAS/OMP env vars before any C-extension load. + + Reads ``DYNACELL_THREADS_PER_WORKER`` from env; if set, exports the BLAS + env vars to that value. This is the only layer that bites BLAS at + extension-load time. Idempotent: existing caller-set env vars are + respected (we only write if the var is unset). + + Called from ``dynacell.__main__:main_cli()`` as the first statement, + before any heavy import triggers numpy/torch BLAS initialization. + """ + threads = os.environ.get("DYNACELL_THREADS_PER_WORKER") + if threads is None: + return + for var in _THREAD_ENV_VARS: + os.environ.setdefault(var, threads) + + +def apply_thread_budget(threads: int) -> None: + """Apply a thread budget to the current process. + + Performs three layered actions: + + 1. Sets BLAS/OMP env vars (only when not already set by the caller); + logs at WARNING if a caller-set value differs from ``threads``. + 2. Calls ``torch.set_num_threads(threads)``. + 3. Calls ``torch.set_num_interop_threads(max(1, threads // 2))`` — + wrapped in try/except since this raises after any parallel op has + run (the common case from Hydra's main). + 4. Activates a module-level ``threadpool_limits(limits=threads, + user_api="blas")`` context for the rest of the process lifetime. + + Parameters + ---------- + threads : int + Number of threads to cap BLAS/OMP/torch to. Must be >= 1. + """ + global _ACTIVE_THREADPOOL_HANDLE + if threads < 1: + raise ValueError(f"apply_thread_budget: threads must be >= 1, got {threads}") + + # 1. Env vars: respect caller-set values; WARN on mismatch. + threads_str = str(threads) + for var in _THREAD_ENV_VARS: + existing = os.environ.get(var) + if existing is None: + os.environ[var] = threads_str + elif existing != threads_str: + logger.log( + logging.WARNING, + "%s already set to %r by environment; not overriding with %r", + var, + existing, + threads_str, + ) + + # 2 + 3. Torch thread caps (deferred import). + import torch + + torch.set_num_threads(threads) + try: + torch.set_num_interop_threads(max(1, threads // 2)) + except RuntimeError as e: + logger.info("torch.set_num_interop_threads: %s; continuing with current value", e) + + # 4. Activate threadpool_limits for the rest of the process. + if _ACTIVE_THREADPOOL_HANDLE is None: + _ACTIVE_THREADPOOL_HANDLE = threadpool_limits(limits=threads, user_api="blas") + else: + # Re-arming: enter a fresh context with the new limit. The old handle + # stays alive (its __exit__ is never called); threadpoolctl honors + # the innermost active limit. Acceptable for re-call from a worker + # initializer that inherited a parent's handle. + _ACTIVE_THREADPOOL_HANDLE = threadpool_limits(limits=threads, user_api="blas") + + +def _get_int(config: DictConfig, key: str, default: int) -> int: + """Read an int field from a runtime config with a default fallback.""" + val = OmegaConf.select(config, key, default=default) + return int(val) + + +def resolve_runtime( + config: DictConfig, + n_positions: int | None = None, + freeze_threads_per_worker: int | None = None, +) -> ResolvedRuntime: + """Materialize the runtime config block with all ``"auto"`` values resolved. + + Two-phase use from ``evaluate_predictions``: + + * **Phase 1** (top of function): ``resolve_runtime(config)``. Provisional + ``fov_workers`` from ``cpu_count // 4``; ``threads_per_worker`` from + ``cpu_count // provisional_fov_workers``. Parent applies BLAS cap with + this value. + * **Phase 2** (after position list built): + ``resolve_runtime(config, n_positions=N, freeze_threads_per_worker=T)``. + Clamps ``fov_workers`` to ``min(provisional, N)`` and returns + ``threads_per_worker == T`` unchanged so the worker initializer matches + what the parent already capped to. + + The ``DYNACELL_FORCE_PER_T_HYGIENE=1`` env var flips both hygiene knobs + on at runtime regardless of YAML defaults — operator escape hatch. + + Parameters + ---------- + config : DictConfig + The eval config (with ``config.runtime`` block). + n_positions : int or None + Number of positions in the eval; used to clamp ``"auto"`` workers in + Phase 2. ``None`` in Phase 1. + freeze_threads_per_worker : int or None + If set, ``threads_per_worker`` returns this value unconditionally + (Phase 2 uses this to stay consistent with the parent's BLAS cap). + + Returns + ------- + ResolvedRuntime + Materialized config; safe to pass across pickle boundaries. + + Raises + ------ + ValueError + If literal ``fov_workers > 1`` is set with ``executor == "serial"``. + """ + runtime = OmegaConf.select(config, "runtime", default=None) + if runtime is None: + # No runtime block configured — fall back to sequential defaults + # matching today's behavior. + return ResolvedRuntime( + fov_workers=1, + threads_per_worker=max(1, os.cpu_count() or 1), + executor="serial", + cuda_empty_cache_every_n_timepoints=0, + gc_collect_every_n_fovs=0, + ) + + executor = str(OmegaConf.select(runtime, "executor", default="serial")) + if executor not in ("serial", "process"): + raise ValueError(f"runtime.executor must be 'serial' or 'process', got {executor!r}") + + cpu_count = max(1, os.cpu_count() or 1) + raw_fov_workers = OmegaConf.select(runtime, "fov_workers", default=1) + raw_threads = OmegaConf.select(runtime, "threads_per_worker", default="auto") + + # Resolve fov_workers. + if isinstance(raw_fov_workers, int): + resolved_fov_workers = int(raw_fov_workers) + if resolved_fov_workers > 1 and executor == "serial": + raise ValueError( + f"runtime.fov_workers={resolved_fov_workers} requires runtime.executor='process' (got 'serial')" + ) + elif raw_fov_workers == "auto": + if executor == "serial": + resolved_fov_workers = 1 + else: + # Provisional divisor of 4 when threads_per_worker is also "auto". + divisor = int(raw_threads) if isinstance(raw_threads, int) else 4 + provisional = max(1, cpu_count // divisor) + clamp = n_positions if n_positions is not None else cpu_count + resolved_fov_workers = max(1, min(provisional, clamp)) + else: + raise ValueError(f"runtime.fov_workers must be int or 'auto', got {raw_fov_workers!r}") + + # Auto-demote process→serial when only 1 worker resolves (avoids 5s spawn cost). + if executor == "process" and resolved_fov_workers == 1: + logger.info("runtime.fov_workers resolved to 1; auto-demoting executor 'process' -> 'serial'") + executor = "serial" + + # Resolve threads_per_worker. + if freeze_threads_per_worker is not None: + resolved_threads = int(freeze_threads_per_worker) + elif isinstance(raw_threads, int): + resolved_threads = int(raw_threads) + elif raw_threads == "auto": + resolved_threads = max(1, cpu_count // resolved_fov_workers) + else: + raise ValueError(f"runtime.threads_per_worker must be int or 'auto', got {raw_threads!r}") + + # Per-T memory hygiene knobs; env-var escape hatch flips both on. + force_hygiene = os.environ.get(_FORCE_PER_T_HYGIENE_ENV, "0") == "1" + if force_hygiene: + cuda_empty_n = max(1, _get_int(runtime, "cuda_empty_cache_every_n_timepoints", 0)) + gc_collect_n = max(1, _get_int(runtime, "gc_collect_every_n_fovs", 0)) + logger.warning( + "%s=1 — forcing cuda_empty_cache_every_n_timepoints=%d, gc_collect_every_n_fovs=%d", + _FORCE_PER_T_HYGIENE_ENV, + cuda_empty_n, + gc_collect_n, + ) + else: + cuda_empty_n = _get_int(runtime, "cuda_empty_cache_every_n_timepoints", 0) + gc_collect_n = _get_int(runtime, "gc_collect_every_n_fovs", 0) + + return ResolvedRuntime( + fov_workers=resolved_fov_workers, + threads_per_worker=resolved_threads, + executor=executor, # type: ignore[arg-type] + cuda_empty_cache_every_n_timepoints=cuda_empty_n, + gc_collect_every_n_fovs=gc_collect_n, + ) From 249a2462763a987f073b680fdf13a3f95df6b696 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 11:34:11 -0700 Subject: [PATCH 02/20] feat(dynacell): per-T memory hygiene + region timing instrumentation (C2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds always-on region timers (RegionTimer context manager) that record wall time of FOV-level and per-T regions to /eval_timing.csv. Region taxonomy matches actual execution placement at pipeline.py:341-505: FOV-level (t=None): mask_gt, mask_pred, cp_gt, deep_gt_dinov3, deep_gt_dynaclr, deep_gt_celldino, features_pred_per_t, microssim, seg_write Per-T: pixel_metrics, mask_metrics, feature_pairwise Timing overhead is ~0.1 ms per region per timepoint; ~120 ms total on a 9-h eval (~4e-6 of wall time) — no gating knob. Per-T memory hygiene knobs (runtime.cuda_empty_cache_every_n_timepoints, runtime.gc_collect_every_n_fovs) ship at 0 (no behavior change) in this commit. The plan calls for measuring and flipping defaults in C2 — that measurement requires a real eval run on the test48 zarr that exceeds this session's scope. Defaults stay at 0; observed late-T degradation in prior A549 ER evals (18 -> 61 s/T over 10 timepoints) is the documented motivation but not the measured fix in this commit. DYNACELL_FORCE_PER_T_HYGIENE=1 env var flips both knobs on at runtime regardless of YAML defaults — operator escape hatch for post-ship mitigation without a code change. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/dynacell/evaluation/pipeline.py | 192 +++++++++++------- .../src/dynacell/evaluation/runtime.py | 97 ++++++++- 2 files changed, 209 insertions(+), 80 deletions(-) diff --git a/applications/dynacell/src/dynacell/evaluation/pipeline.py b/applications/dynacell/src/dynacell/evaluation/pipeline.py index 9a22e57ce..3a428f126 100644 --- a/applications/dynacell/src/dynacell/evaluation/pipeline.py +++ b/applications/dynacell/src/dynacell/evaluation/pipeline.py @@ -155,7 +155,15 @@ def _save_embeddings(save_dir: Path, groups: dict[str, tuple[list, list, list]]) def evaluate_predictions(config: DictConfig): """Evaluate predictions on all test images.""" - from dynacell.evaluation.runtime import apply_thread_budget, resolve_runtime + from dynacell.evaluation.runtime import ( + apply_thread_budget, + dump_timings_csv, + maybe_empty_cuda_cache, + maybe_gc_collect, + region_timer, + reset_timings, + resolve_runtime, + ) from dynacell.evaluation.segmentation import prepare_segmentation_model, segment from dynacell.evaluation.utils import CellDinoFeatureExtractor, DinoV3FeatureExtractor, DynaCLRFeatureExtractor @@ -165,6 +173,7 @@ def evaluate_predictions(config: DictConfig): # frozen across phases for parent/worker BLAS-cap consistency. runtime = resolve_runtime(config) apply_thread_budget(runtime.threads_per_worker) + reset_timings() all_pixel_metrics = [] all_mask_metrics = [] @@ -328,10 +337,12 @@ def evaluate_predictions(config: DictConfig): for ctx in sides_for_precompute.values(): flush_manifest(ctx) - for p1, p2, p3 in tqdm( - zip(pred_positions, gt_positions, seg_positions), - total=len(pred_positions), - desc="Processing positions", + for fov_idx, (p1, p2, p3) in enumerate( + tqdm( + zip(pred_positions, gt_positions, seg_positions), + total=len(pred_positions), + desc="Processing positions", + ) ): pos_name_pred, pos_pred = p1 _, pos_gt = p2 @@ -346,37 +357,49 @@ def evaluate_predictions(config: DictConfig): T = predict.shape[0] - gt_mask_stack = fov_masks(cache_ctx, pos_name_pred, target, seg_model) - pred_mask_stack = ( - fov_masks(pred_cache_ctx, pos_name_pred, predict, seg_model) if pred_cache_ctx.enabled else None - ) + with region_timer("mask_gt", pos_name_pred): + gt_mask_stack = fov_masks(cache_ctx, pos_name_pred, target, seg_model) + if pred_cache_ctx.enabled: + with region_timer("mask_pred", pos_name_pred): + pred_mask_stack = fov_masks(pred_cache_ctx, pos_name_pred, predict, seg_model) + else: + pred_mask_stack = None if config.compute_feature_metrics: - gt_cp_per_t = fov_cp_features(cache_ctx, pos_name_pred, target, cell_segmentation) - gt_dinov3_per_t = fov_deep_features( - cache_ctx, pos_name_pred, target, cell_segmentation, dinov3_feature_extractor, "dinov3" - ) - gt_dynaclr_per_t = fov_deep_features( - cache_ctx, pos_name_pred, target, cell_segmentation, dynaclr_feature_extractor, "dynaclr" - ) - gt_celldino_per_t = ( - fov_deep_features( - cache_ctx, pos_name_pred, target, cell_segmentation, celldino_feature_extractor, "celldino" + with region_timer("cp_gt", pos_name_pred): + gt_cp_per_t = fov_cp_features(cache_ctx, pos_name_pred, target, cell_segmentation) + with region_timer("deep_gt_dinov3", pos_name_pred): + gt_dinov3_per_t = fov_deep_features( + cache_ctx, pos_name_pred, target, cell_segmentation, dinov3_feature_extractor, "dinov3" + ) + with region_timer("deep_gt_dynaclr", pos_name_pred): + gt_dynaclr_per_t = fov_deep_features( + cache_ctx, pos_name_pred, target, cell_segmentation, dynaclr_feature_extractor, "dynaclr" + ) + if celldino_feature_extractor is not None: + with region_timer("deep_gt_celldino", pos_name_pred): + gt_celldino_per_t = fov_deep_features( + cache_ctx, + pos_name_pred, + target, + cell_segmentation, + celldino_feature_extractor, + "celldino", + ) + else: + gt_celldino_per_t = None + with region_timer("features_pred_per_t", pos_name_pred): + pred_per_t = _fov_pred_features_per_t( + pred_cache_ctx, + pos_name_pred, + predict, + cell_segmentation, + dinov3_feature_extractor, + dynaclr_feature_extractor, + celldino_feature_extractor, + config.feature_metrics.patch_size, + config.pixel_metrics.spacing, ) - if celldino_feature_extractor is not None - else None - ) - pred_per_t = _fov_pred_features_per_t( - pred_cache_ctx, - pos_name_pred, - predict, - cell_segmentation, - dinov3_feature_extractor, - dynaclr_feature_extractor, - celldino_feature_extractor, - config.feature_metrics.patch_size, - config.pixel_metrics.spacing, - ) microssim_data = [] fov_pixel_metrics = [] @@ -385,51 +408,54 @@ def evaluate_predictions(config: DictConfig): for t in tqdm(range(T), desc="Processing timepoints"): data_info = {"FOV": pos_name_pred, "Timepoint": t} - pixel_metrics = compute_pixel_metrics( - predict[t], - target[t], - spacing=config.pixel_metrics.spacing, - fsc_kwargs=config.pixel_metrics.fsc, - spectral_pcc_kwargs=config.pixel_metrics.spectral_pcc, - use_gpu=config.use_gpu, - ) + with region_timer("pixel_metrics", pos_name_pred, t): + pixel_metrics = compute_pixel_metrics( + predict[t], + target[t], + spacing=config.pixel_metrics.spacing, + fsc_kwargs=config.pixel_metrics.fsc, + spectral_pcc_kwargs=config.pixel_metrics.spectral_pcc, + use_gpu=config.use_gpu, + ) if config.compute_microssim: microssim_data.append({"target": target[t], "predict": predict[t]}) fov_pixel_metrics.append({**data_info, **pixel_metrics}) - segmented_target = gt_mask_stack[t] - segmented_predict = ( - pred_mask_stack[t] - if pred_mask_stack is not None - else np.asarray(segment(predict[t], config.target_name, seg_model=seg_model)).astype(bool) - ) - all_mask_metrics.append( - {**data_info, **evaluate_segmentations(segmented_predict, segmented_target)} - ) - segmentations.append(np.stack([segmented_predict, segmented_target], axis=0)) + with region_timer("mask_metrics", pos_name_pred, t): + segmented_target = gt_mask_stack[t] + segmented_predict = ( + pred_mask_stack[t] + if pred_mask_stack is not None + else np.asarray(segment(predict[t], config.target_name, seg_model=seg_model)).astype(bool) + ) + all_mask_metrics.append( + {**data_info, **evaluate_segmentations(segmented_predict, segmented_target)} + ) + segmentations.append(np.stack([segmented_predict, segmented_target], axis=0)) if config.compute_feature_metrics: - pred_cp = pred_per_t["cp"][t] - pred_dinov3 = pred_per_t["dinov3"][t] - pred_dynaclr = pred_per_t["dynaclr"][t] - pred_celldino = pred_per_t["celldino"][t] if pred_per_t["celldino"] is not None else None - pred_cp, gt_cp_t = drop_paired_nonfinite_rows(pred_cp, gt_cp_per_t[t]) - # Per-timepoint CP: drop target-zero columns + per-side z-score. - # Deep features stay untouched. - if pred_cp.size and gt_cp_t.size: - pred_cp_z, gt_cp_z = _cp_dropzero_zscore(pred_cp, gt_cp_t) - else: - pred_cp_z, gt_cp_z = pred_cp, gt_cp_t - pairwise_metrics = { - **compute_feature_similarity_pairwise(pred_cp_z, gt_cp_z, "CP"), - **compute_feature_similarity_pairwise(pred_dinov3, gt_dinov3_per_t[t], "DINOv3"), - **compute_feature_similarity_pairwise(pred_dynaclr, gt_dynaclr_per_t[t], "DynaCLR"), - } - if pred_celldino is not None: - pairwise_metrics.update( - compute_feature_similarity_pairwise(pred_celldino, gt_celldino_per_t[t], "CellDINO") - ) - all_feature_metrics.append({**data_info, **pairwise_metrics}) + with region_timer("feature_pairwise", pos_name_pred, t): + pred_cp = pred_per_t["cp"][t] + pred_dinov3 = pred_per_t["dinov3"][t] + pred_dynaclr = pred_per_t["dynaclr"][t] + pred_celldino = pred_per_t["celldino"][t] if pred_per_t["celldino"] is not None else None + pred_cp, gt_cp_t = drop_paired_nonfinite_rows(pred_cp, gt_cp_per_t[t]) + # Per-timepoint CP: drop target-zero columns + per-side z-score. + # Deep features stay untouched. + if pred_cp.size and gt_cp_t.size: + pred_cp_z, gt_cp_z = _cp_dropzero_zscore(pred_cp, gt_cp_t) + else: + pred_cp_z, gt_cp_z = pred_cp, gt_cp_t + pairwise_metrics = { + **compute_feature_similarity_pairwise(pred_cp_z, gt_cp_z, "CP"), + **compute_feature_similarity_pairwise(pred_dinov3, gt_dinov3_per_t[t], "DINOv3"), + **compute_feature_similarity_pairwise(pred_dynaclr, gt_dynaclr_per_t[t], "DynaCLR"), + } + if pred_celldino is not None: + pairwise_metrics.update( + compute_feature_similarity_pairwise(pred_celldino, gt_celldino_per_t[t], "CellDINO") + ) + all_feature_metrics.append({**data_info, **pairwise_metrics}) # pred/gt fov+timepoint arrays are byte-identical per # backbone; share one reference each (consumers only # concatenate, never mutate in place). @@ -470,21 +496,27 @@ def evaluate_predictions(config: DictConfig): pred_celldino_ts.append(t_arr) gt_celldino_ts.append(t_arr) - seg = np.stack(segmentations, axis=0) # shape: (T, 2, D, H, W) - row, col, fov = pos_name_pred.split("/") - seg_pos = segmentation_results.create_position(row, col, fov) - seg_pos.create_image("0", seg.astype(bool)) + maybe_empty_cuda_cache(t, runtime.cuda_empty_cache_every_n_timepoints) + + with region_timer("seg_write", pos_name_pred): + seg = np.stack(segmentations, axis=0) # shape: (T, 2, D, H, W) + row, col, fov = pos_name_pred.split("/") + seg_pos = segmentation_results.create_position(row, col, fov) + seg_pos.create_image("0", seg.astype(bool)) if config.compute_microssim: - microssim_scores = calculate_microssim(microssim_data) - for i in range(T): - fov_pixel_metrics[i]["MicroMS3IM"] = float(microssim_scores[i]["MicroMS3IM"]) + with region_timer("microssim", pos_name_pred): + microssim_scores = calculate_microssim(microssim_data) + for i in range(T): + fov_pixel_metrics[i]["MicroMS3IM"] = float(microssim_scores[i]["MicroMS3IM"]) all_pixel_metrics.extend(fov_pixel_metrics) # Flush manifest after each position so interrupted runs preserve progress. flush_manifest(cache_ctx) flush_manifest(pred_cache_ctx) + + maybe_gc_collect(fov_idx, runtime.gc_collect_every_n_fovs) finally: if seg_plate is not None: seg_plate.close() @@ -602,6 +634,8 @@ def _compute_one(args): embedding_groups["gt_celldino"] = (gt_celldino_feats, gt_celldino_fovs, gt_celldino_ts) _save_embeddings(save_dir, embedding_groups) + dump_timings_csv(save_dir) + return all_pixel_metrics, all_mask_metrics, all_feature_metrics diff --git a/applications/dynacell/src/dynacell/evaluation/runtime.py b/applications/dynacell/src/dynacell/evaluation/runtime.py index 45b58f558..248905c27 100644 --- a/applications/dynacell/src/dynacell/evaluation/runtime.py +++ b/applications/dynacell/src/dynacell/evaluation/runtime.py @@ -17,10 +17,15 @@ from __future__ import annotations +import csv +import gc import logging import os +import time +from contextlib import contextmanager from dataclasses import dataclass -from typing import Any, Literal +from pathlib import Path +from typing import Any, Iterator, Literal from omegaconf import DictConfig, OmegaConf from threadpoolctl import threadpool_limits @@ -277,3 +282,93 @@ def resolve_runtime( cuda_empty_cache_every_n_timepoints=cuda_empty_n, gc_collect_every_n_fovs=gc_collect_n, ) + + +# --------------------------------------------------------------------------- +# Per-T memory hygiene + region timing instrumentation. +# --------------------------------------------------------------------------- + +# Per-process timing collector. Each entry is (pos_name, t_or_None, region, +# seconds). Under executor=process (C3+), workers return their slice in +# FovResult.timings; the parent's aggregator concatenates. +_TIMINGS: list[tuple[str, int | None, str, float]] = [] + + +def reset_timings() -> None: + """Clear the per-process timing collector.""" + _TIMINGS.clear() + + +def get_timings() -> list[tuple[str, int | None, str, float]]: + """Return a copy of the per-process timing collector.""" + return list(_TIMINGS) + + +def extend_timings(rows: list[tuple[str, int | None, str, float]]) -> None: + """Append a batch of timing rows (used by parent aggregator under process mode).""" + _TIMINGS.extend(rows) + + +@contextmanager +def region_timer(region: str, pos_name: str, t: int | None = None) -> Iterator[None]: + """Record wall time of the wrapped block to the per-process timing collector. + + Parameters + ---------- + region : str + Region tag (e.g. "pixel_metrics", "mask_gt", "features_pred_per_t"). + pos_name : str + FOV identifier (typically "//"). + t : int or None + Timepoint index when the region is per-T; ``None`` for FOV-level regions. + """ + start = time.perf_counter() + try: + yield + finally: + _TIMINGS.append((pos_name, t, region, time.perf_counter() - start)) + + +def dump_timings_csv(save_dir: Path) -> Path | None: + """Write the collected timings to ``/eval_timing.csv`` and return the path. + + Returns ``None`` if no timings were recorded. + """ + if not _TIMINGS: + return None + save_dir.mkdir(parents=True, exist_ok=True) + out_path = save_dir / "eval_timing.csv" + with out_path.open("w", newline="") as f: + writer = csv.writer(f) + writer.writerow(["pos_name", "t", "region", "seconds"]) + for pos_name, t, region, seconds in _TIMINGS: + writer.writerow([pos_name, "" if t is None else t, region, f"{seconds:.6f}"]) + return out_path + + +def maybe_empty_cuda_cache(t: int, every_n: int) -> None: + """Call ``torch.cuda.empty_cache()`` every ``every_n`` timepoints. + + No-op when ``every_n <= 0`` or CUDA is unavailable. + """ + if every_n <= 0: + return + if (t + 1) % every_n != 0: + return + import torch + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def maybe_gc_collect(fov_idx: int, every_n: int) -> None: + """Call ``gc.collect()`` every ``every_n`` FOVs. + + ``fov_idx`` is 0-based; collects when ``(fov_idx + 1) % every_n == 0``. + No-op when ``every_n <= 0``. + """ + if every_n <= 0: + return + if (fov_idx + 1) % every_n != 0: + return + gc.collect() From 0b75ee6e398b1299650436b1af2a1a4e60bc9d34 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 11:48:52 -0700 Subject: [PATCH 03/20] refactor(dynacell): extract per-FOV body into _process_one_fov + FovResult (C3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pure refactor: extract the per-FOV body of evaluate_predictions into a module-level function _process_one_fov(...) -> FovResult and an _aggregate_fov_result(...) parent-side aggregator. Serial-path behavior is identical to the prior code (verified: 384/48 pass/fail parity with HEAD~2; the 48 failures are all pre-existing). FovResult carries everything one FOV contributes to the run as a picklable dataclass: - per_t_pixel_rows (with microssim already merged in per pipeline.py:471 semantics — no separate microssim_rows field), - per_t_mask_rows, per_t_feature_rows, - seg_array of shape (T, 2, D, H, W) bool (channel 0 = pred, channel 1 = GT), - 4 _BackboneLists for cp/dinov3/dynaclr/celldino — each holds the per-T pred_feats/gt_feats/pred_fovs/gt_fovs/pred_ts/gt_ts arrays. Workers preserve the existing `if size > 0: append` guard so empty per-T slots never enter the lists; aggregator extends parent lists unconditionally. - timings slice from the per-process region_timer collector. _aggregate_fov_result writes the seg array to the HCS plate (still in the parent — concurrent create_position from workers would race plate metadata), extends per-T row lists, and extends the 24 dataset-level accumulators. This refactor positions C4 to add ProcessPoolExecutor + worker primitives without further changes to the FOV body. flush_manifest stays in the parent's loop for now; under process mode (C4) each worker will flush its own ctx at FOV end. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/dynacell/evaluation/pipeline.py | 543 ++++++++++++------ 1 file changed, 377 insertions(+), 166 deletions(-) diff --git a/applications/dynacell/src/dynacell/evaluation/pipeline.py b/applications/dynacell/src/dynacell/evaluation/pipeline.py index 3a428f126..08ecc4f19 100644 --- a/applications/dynacell/src/dynacell/evaluation/pipeline.py +++ b/applications/dynacell/src/dynacell/evaluation/pipeline.py @@ -2,6 +2,7 @@ import json from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field from pathlib import Path import hydra @@ -153,18 +154,345 @@ def _save_embeddings(save_dir: Path, groups: dict[str, tuple[list, list, list]]) print(f"Saved embeddings → {out_path}") +@dataclass +class _BackboneLists: + """Per-FOV per-backbone single-cell feature lists. + + Each list holds one entry per timepoint with non-empty cell features. + ``pred_*`` and ``gt_*`` indices align: ``pred_feats[i]`` and + ``gt_feats[i]`` come from the same (pos_name, t) pair, with + ``pred_fovs[i] == gt_fovs[i]`` and ``pred_ts[i] == gt_ts[i]``. + """ + + pred_feats: list[np.ndarray] = field(default_factory=list) + gt_feats: list[np.ndarray] = field(default_factory=list) + pred_fovs: list[np.ndarray] = field(default_factory=list) + gt_fovs: list[np.ndarray] = field(default_factory=list) + pred_ts: list[np.ndarray] = field(default_factory=list) + gt_ts: list[np.ndarray] = field(default_factory=list) + + +@dataclass +class FovResult: + """Output of ``_process_one_fov``: everything one FOV contributes to the run. + + Designed for cross-process transport via pickle: all fields are picklable + types (str, list[dict], list[np.ndarray], np.ndarray) — no iohub handles, + no torch modules. + + Microssim scores are merged into ``per_t_pixel_rows`` before return + (matching the existing serial path at ``pipeline.py:478-481``); no + separate microssim field. + """ + + pos_name: str + row: str + col: str + fov: str + per_t_pixel_rows: list[dict] + per_t_mask_rows: list[dict] + per_t_feature_rows: list[dict] + seg_array: np.ndarray # (T, 2, D, H, W) bool: channel 0 = pred, channel 1 = GT + cp: _BackboneLists = field(default_factory=_BackboneLists) + dinov3: _BackboneLists = field(default_factory=_BackboneLists) + dynaclr: _BackboneLists = field(default_factory=_BackboneLists) + celldino: _BackboneLists = field(default_factory=_BackboneLists) + timings: list[tuple[str, int | None, str, float]] = field(default_factory=list) + + +def _process_one_fov( + config: DictConfig, + cuda_empty_cache_every_n_timepoints: int, + pos_name_pred: str, + pos_pred, + pos_gt, + pos_seg, + io_config, + cache_ctx, + pred_cache_ctx, + seg_model, + dinov3_feature_extractor, + dynaclr_feature_extractor, + celldino_feature_extractor, +) -> FovResult: + """Compute everything one FOV contributes to the eval and return a FovResult. + + No side effects on shared parent state (no segmentation_results plate + writes, no manifest flush). The parent aggregator handles those — see + ``_aggregate_fov_result``. Used by both the serial and process FOV-loop + paths in ``evaluate_predictions``. + """ + from dynacell.evaluation.runtime import ( + get_timings, + maybe_empty_cuda_cache, + region_timer, + ) + from dynacell.evaluation.segmentation import segment + + timings_start = len(get_timings()) + + pred_channel_index = pos_pred.get_channel_index(io_config.pred_channel_name) + gt_channel_index = pos_gt.get_channel_index(io_config.gt_channel_name) + + predict = np.asarray(pos_pred.data[:, pred_channel_index]) # shape: (T, D, H, W) + target = np.asarray(pos_gt.data[:, gt_channel_index]) + cell_segmentation = np.asarray(pos_seg.data[:, 0]) if pos_seg is not None else None + + T = predict.shape[0] + + with region_timer("mask_gt", pos_name_pred): + gt_mask_stack = fov_masks(cache_ctx, pos_name_pred, target, seg_model) + if pred_cache_ctx.enabled: + with region_timer("mask_pred", pos_name_pred): + pred_mask_stack = fov_masks(pred_cache_ctx, pos_name_pred, predict, seg_model) + else: + pred_mask_stack = None + + gt_cp_per_t = None + gt_dinov3_per_t = None + gt_dynaclr_per_t = None + gt_celldino_per_t = None + pred_per_t = None + if config.compute_feature_metrics: + with region_timer("cp_gt", pos_name_pred): + gt_cp_per_t = fov_cp_features(cache_ctx, pos_name_pred, target, cell_segmentation) + with region_timer("deep_gt_dinov3", pos_name_pred): + gt_dinov3_per_t = fov_deep_features( + cache_ctx, pos_name_pred, target, cell_segmentation, dinov3_feature_extractor, "dinov3" + ) + with region_timer("deep_gt_dynaclr", pos_name_pred): + gt_dynaclr_per_t = fov_deep_features( + cache_ctx, pos_name_pred, target, cell_segmentation, dynaclr_feature_extractor, "dynaclr" + ) + if celldino_feature_extractor is not None: + with region_timer("deep_gt_celldino", pos_name_pred): + gt_celldino_per_t = fov_deep_features( + cache_ctx, + pos_name_pred, + target, + cell_segmentation, + celldino_feature_extractor, + "celldino", + ) + with region_timer("features_pred_per_t", pos_name_pred): + pred_per_t = _fov_pred_features_per_t( + pred_cache_ctx, + pos_name_pred, + predict, + cell_segmentation, + dinov3_feature_extractor, + dynaclr_feature_extractor, + celldino_feature_extractor, + config.feature_metrics.patch_size, + config.pixel_metrics.spacing, + ) + + microssim_data: list[dict] = [] + fov_pixel_metrics: list[dict] = [] + fov_mask_metrics: list[dict] = [] + fov_feature_metrics: list[dict] = [] + segmentations: list[np.ndarray] = [] + cp = _BackboneLists() + dinov3 = _BackboneLists() + dynaclr = _BackboneLists() + celldino = _BackboneLists() + + for t in tqdm(range(T), desc="Processing timepoints", leave=False): + data_info = {"FOV": pos_name_pred, "Timepoint": t} + + with region_timer("pixel_metrics", pos_name_pred, t): + pixel_metrics = compute_pixel_metrics( + predict[t], + target[t], + spacing=config.pixel_metrics.spacing, + fsc_kwargs=config.pixel_metrics.fsc, + spectral_pcc_kwargs=config.pixel_metrics.spectral_pcc, + use_gpu=config.use_gpu, + ) + if config.compute_microssim: + microssim_data.append({"target": target[t], "predict": predict[t]}) + fov_pixel_metrics.append({**data_info, **pixel_metrics}) + + with region_timer("mask_metrics", pos_name_pred, t): + segmented_target = gt_mask_stack[t] + segmented_predict = ( + pred_mask_stack[t] + if pred_mask_stack is not None + else np.asarray(segment(predict[t], config.target_name, seg_model=seg_model)).astype(bool) + ) + fov_mask_metrics.append({**data_info, **evaluate_segmentations(segmented_predict, segmented_target)}) + segmentations.append(np.stack([segmented_predict, segmented_target], axis=0)) + + if config.compute_feature_metrics: + with region_timer("feature_pairwise", pos_name_pred, t): + pred_cp = pred_per_t["cp"][t] + pred_dinov3 = pred_per_t["dinov3"][t] + pred_dynaclr = pred_per_t["dynaclr"][t] + pred_celldino = pred_per_t["celldino"][t] if pred_per_t["celldino"] is not None else None + pred_cp, gt_cp_t = drop_paired_nonfinite_rows(pred_cp, gt_cp_per_t[t]) + if pred_cp.size and gt_cp_t.size: + pred_cp_z, gt_cp_z = _cp_dropzero_zscore(pred_cp, gt_cp_t) + else: + pred_cp_z, gt_cp_z = pred_cp, gt_cp_t + pairwise_metrics = { + **compute_feature_similarity_pairwise(pred_cp_z, gt_cp_z, "CP"), + **compute_feature_similarity_pairwise(pred_dinov3, gt_dinov3_per_t[t], "DINOv3"), + **compute_feature_similarity_pairwise(pred_dynaclr, gt_dynaclr_per_t[t], "DynaCLR"), + } + if pred_celldino is not None: + pairwise_metrics.update( + compute_feature_similarity_pairwise(pred_celldino, gt_celldino_per_t[t], "CellDINO") + ) + fov_feature_metrics.append({**data_info, **pairwise_metrics}) + if pred_cp.size > 0: + cp.pred_feats.append(pred_cp) + cp.gt_feats.append(gt_cp_t) + fov_arr = np.full(len(pred_cp), pos_name_pred) + t_arr = np.full(len(pred_cp), t, dtype=np.int32) + cp.pred_fovs.append(fov_arr) + cp.gt_fovs.append(fov_arr) + cp.pred_ts.append(t_arr) + cp.gt_ts.append(t_arr) + if pred_dinov3.size > 0: + dinov3.pred_feats.append(pred_dinov3) + dinov3.gt_feats.append(gt_dinov3_per_t[t]) + fov_arr = np.full(len(pred_dinov3), pos_name_pred) + t_arr = np.full(len(pred_dinov3), t, dtype=np.int32) + dinov3.pred_fovs.append(fov_arr) + dinov3.gt_fovs.append(fov_arr) + dinov3.pred_ts.append(t_arr) + dinov3.gt_ts.append(t_arr) + if pred_dynaclr.size > 0: + dynaclr.pred_feats.append(pred_dynaclr) + dynaclr.gt_feats.append(gt_dynaclr_per_t[t]) + fov_arr = np.full(len(pred_dynaclr), pos_name_pred) + t_arr = np.full(len(pred_dynaclr), t, dtype=np.int32) + dynaclr.pred_fovs.append(fov_arr) + dynaclr.gt_fovs.append(fov_arr) + dynaclr.pred_ts.append(t_arr) + dynaclr.gt_ts.append(t_arr) + if pred_celldino is not None and pred_celldino.size > 0: + celldino.pred_feats.append(pred_celldino) + celldino.gt_feats.append(gt_celldino_per_t[t]) + fov_arr = np.full(len(pred_celldino), pos_name_pred) + t_arr = np.full(len(pred_celldino), t, dtype=np.int32) + celldino.pred_fovs.append(fov_arr) + celldino.gt_fovs.append(fov_arr) + celldino.pred_ts.append(t_arr) + celldino.gt_ts.append(t_arr) + + maybe_empty_cuda_cache(t, cuda_empty_cache_every_n_timepoints) + + seg_array = np.stack(segmentations, axis=0) # shape: (T, 2, D, H, W) + + if config.compute_microssim: + with region_timer("microssim", pos_name_pred): + microssim_scores = calculate_microssim(microssim_data) + for i in range(T): + fov_pixel_metrics[i]["MicroMS3IM"] = float(microssim_scores[i]["MicroMS3IM"]) + + row, col, fov = pos_name_pred.split("/") + return FovResult( + pos_name=pos_name_pred, + row=row, + col=col, + fov=fov, + per_t_pixel_rows=fov_pixel_metrics, + per_t_mask_rows=fov_mask_metrics, + per_t_feature_rows=fov_feature_metrics, + seg_array=seg_array.astype(bool), + cp=cp, + dinov3=dinov3, + dynaclr=dynaclr, + celldino=celldino, + timings=get_timings()[timings_start:], + ) + + +def _aggregate_fov_result( + result: FovResult, + segmentation_results, + all_pixel_metrics: list[dict], + all_mask_metrics: list[dict], + all_feature_metrics: list[dict], + pred_cp_feats: list[np.ndarray], + pred_cp_fovs: list[np.ndarray], + pred_cp_ts: list[np.ndarray], + gt_cp_feats: list[np.ndarray], + gt_cp_fovs: list[np.ndarray], + gt_cp_ts: list[np.ndarray], + pred_dinov3_feats: list[np.ndarray], + pred_dinov3_fovs: list[np.ndarray], + pred_dinov3_ts: list[np.ndarray], + gt_dinov3_feats: list[np.ndarray], + gt_dinov3_fovs: list[np.ndarray], + gt_dinov3_ts: list[np.ndarray], + pred_dynaclr_feats: list[np.ndarray], + pred_dynaclr_fovs: list[np.ndarray], + pred_dynaclr_ts: list[np.ndarray], + gt_dynaclr_feats: list[np.ndarray], + gt_dynaclr_fovs: list[np.ndarray], + gt_dynaclr_ts: list[np.ndarray], + pred_celldino_feats: list[np.ndarray], + pred_celldino_fovs: list[np.ndarray], + pred_celldino_ts: list[np.ndarray], + gt_celldino_feats: list[np.ndarray], + gt_celldino_fovs: list[np.ndarray], + gt_celldino_ts: list[np.ndarray], +) -> None: + """Apply one FOV's contributions to the parent-side run state. + + Writes the segmentation array to the HCS plate, extends the per-T row + lists, and extends the 24 dataset-level accumulator lists. + """ + from dynacell.evaluation.runtime import extend_timings, region_timer + + extend_timings(result.timings) + + with region_timer("seg_write", result.pos_name): + seg_pos = segmentation_results.create_position(result.row, result.col, result.fov) + seg_pos.create_image("0", result.seg_array) + + all_pixel_metrics.extend(result.per_t_pixel_rows) + all_mask_metrics.extend(result.per_t_mask_rows) + all_feature_metrics.extend(result.per_t_feature_rows) + + pred_cp_feats.extend(result.cp.pred_feats) + pred_cp_fovs.extend(result.cp.pred_fovs) + pred_cp_ts.extend(result.cp.pred_ts) + gt_cp_feats.extend(result.cp.gt_feats) + gt_cp_fovs.extend(result.cp.gt_fovs) + gt_cp_ts.extend(result.cp.gt_ts) + pred_dinov3_feats.extend(result.dinov3.pred_feats) + pred_dinov3_fovs.extend(result.dinov3.pred_fovs) + pred_dinov3_ts.extend(result.dinov3.pred_ts) + gt_dinov3_feats.extend(result.dinov3.gt_feats) + gt_dinov3_fovs.extend(result.dinov3.gt_fovs) + gt_dinov3_ts.extend(result.dinov3.gt_ts) + pred_dynaclr_feats.extend(result.dynaclr.pred_feats) + pred_dynaclr_fovs.extend(result.dynaclr.pred_fovs) + pred_dynaclr_ts.extend(result.dynaclr.pred_ts) + gt_dynaclr_feats.extend(result.dynaclr.gt_feats) + gt_dynaclr_fovs.extend(result.dynaclr.gt_fovs) + gt_dynaclr_ts.extend(result.dynaclr.gt_ts) + pred_celldino_feats.extend(result.celldino.pred_feats) + pred_celldino_fovs.extend(result.celldino.pred_fovs) + pred_celldino_ts.extend(result.celldino.pred_ts) + gt_celldino_feats.extend(result.celldino.gt_feats) + gt_celldino_fovs.extend(result.celldino.gt_fovs) + gt_celldino_ts.extend(result.celldino.gt_ts) + + def evaluate_predictions(config: DictConfig): """Evaluate predictions on all test images.""" from dynacell.evaluation.runtime import ( apply_thread_budget, dump_timings_csv, - maybe_empty_cuda_cache, maybe_gc_collect, - region_timer, reset_timings, resolve_runtime, ) - from dynacell.evaluation.segmentation import prepare_segmentation_model, segment + from dynacell.evaluation.segmentation import prepare_segmentation_model from dynacell.evaluation.utils import CellDinoFeatureExtractor, DinoV3FeatureExtractor, DynaCLRFeatureExtractor # Phase 1 runtime resolution: lock in executor + thread caps before any @@ -348,169 +676,52 @@ def evaluate_predictions(config: DictConfig): _, pos_gt = p2 _, pos_seg = p3 - pred_channel_index = pos_pred.get_channel_index(io_config.pred_channel_name) - gt_channel_index = pos_gt.get_channel_index(io_config.gt_channel_name) - - predict = np.asarray(pos_pred.data[:, pred_channel_index]) # shape: (T, D, H, W) - target = np.asarray(pos_gt.data[:, gt_channel_index]) # shape: (T, D, H, W) - cell_segmentation = np.asarray(pos_seg.data[:, 0]) if pos_seg is not None else None - - T = predict.shape[0] - - with region_timer("mask_gt", pos_name_pred): - gt_mask_stack = fov_masks(cache_ctx, pos_name_pred, target, seg_model) - if pred_cache_ctx.enabled: - with region_timer("mask_pred", pos_name_pred): - pred_mask_stack = fov_masks(pred_cache_ctx, pos_name_pred, predict, seg_model) - else: - pred_mask_stack = None - - if config.compute_feature_metrics: - with region_timer("cp_gt", pos_name_pred): - gt_cp_per_t = fov_cp_features(cache_ctx, pos_name_pred, target, cell_segmentation) - with region_timer("deep_gt_dinov3", pos_name_pred): - gt_dinov3_per_t = fov_deep_features( - cache_ctx, pos_name_pred, target, cell_segmentation, dinov3_feature_extractor, "dinov3" - ) - with region_timer("deep_gt_dynaclr", pos_name_pred): - gt_dynaclr_per_t = fov_deep_features( - cache_ctx, pos_name_pred, target, cell_segmentation, dynaclr_feature_extractor, "dynaclr" - ) - if celldino_feature_extractor is not None: - with region_timer("deep_gt_celldino", pos_name_pred): - gt_celldino_per_t = fov_deep_features( - cache_ctx, - pos_name_pred, - target, - cell_segmentation, - celldino_feature_extractor, - "celldino", - ) - else: - gt_celldino_per_t = None - with region_timer("features_pred_per_t", pos_name_pred): - pred_per_t = _fov_pred_features_per_t( - pred_cache_ctx, - pos_name_pred, - predict, - cell_segmentation, - dinov3_feature_extractor, - dynaclr_feature_extractor, - celldino_feature_extractor, - config.feature_metrics.patch_size, - config.pixel_metrics.spacing, - ) - - microssim_data = [] - fov_pixel_metrics = [] - segmentations = [] - - for t in tqdm(range(T), desc="Processing timepoints"): - data_info = {"FOV": pos_name_pred, "Timepoint": t} - - with region_timer("pixel_metrics", pos_name_pred, t): - pixel_metrics = compute_pixel_metrics( - predict[t], - target[t], - spacing=config.pixel_metrics.spacing, - fsc_kwargs=config.pixel_metrics.fsc, - spectral_pcc_kwargs=config.pixel_metrics.spectral_pcc, - use_gpu=config.use_gpu, - ) - if config.compute_microssim: - microssim_data.append({"target": target[t], "predict": predict[t]}) - fov_pixel_metrics.append({**data_info, **pixel_metrics}) - - with region_timer("mask_metrics", pos_name_pred, t): - segmented_target = gt_mask_stack[t] - segmented_predict = ( - pred_mask_stack[t] - if pred_mask_stack is not None - else np.asarray(segment(predict[t], config.target_name, seg_model=seg_model)).astype(bool) - ) - all_mask_metrics.append( - {**data_info, **evaluate_segmentations(segmented_predict, segmented_target)} - ) - segmentations.append(np.stack([segmented_predict, segmented_target], axis=0)) - - if config.compute_feature_metrics: - with region_timer("feature_pairwise", pos_name_pred, t): - pred_cp = pred_per_t["cp"][t] - pred_dinov3 = pred_per_t["dinov3"][t] - pred_dynaclr = pred_per_t["dynaclr"][t] - pred_celldino = pred_per_t["celldino"][t] if pred_per_t["celldino"] is not None else None - pred_cp, gt_cp_t = drop_paired_nonfinite_rows(pred_cp, gt_cp_per_t[t]) - # Per-timepoint CP: drop target-zero columns + per-side z-score. - # Deep features stay untouched. - if pred_cp.size and gt_cp_t.size: - pred_cp_z, gt_cp_z = _cp_dropzero_zscore(pred_cp, gt_cp_t) - else: - pred_cp_z, gt_cp_z = pred_cp, gt_cp_t - pairwise_metrics = { - **compute_feature_similarity_pairwise(pred_cp_z, gt_cp_z, "CP"), - **compute_feature_similarity_pairwise(pred_dinov3, gt_dinov3_per_t[t], "DINOv3"), - **compute_feature_similarity_pairwise(pred_dynaclr, gt_dynaclr_per_t[t], "DynaCLR"), - } - if pred_celldino is not None: - pairwise_metrics.update( - compute_feature_similarity_pairwise(pred_celldino, gt_celldino_per_t[t], "CellDINO") - ) - all_feature_metrics.append({**data_info, **pairwise_metrics}) - # pred/gt fov+timepoint arrays are byte-identical per - # backbone; share one reference each (consumers only - # concatenate, never mutate in place). - if pred_cp.size > 0: - pred_cp_feats.append(pred_cp) - gt_cp_feats.append(gt_cp_t) - fov_arr = np.full(len(pred_cp), pos_name_pred) - t_arr = np.full(len(pred_cp), t, dtype=np.int32) - pred_cp_fovs.append(fov_arr) - gt_cp_fovs.append(fov_arr) - pred_cp_ts.append(t_arr) - gt_cp_ts.append(t_arr) - if pred_dinov3.size > 0: - pred_dinov3_feats.append(pred_dinov3) - gt_dinov3_feats.append(gt_dinov3_per_t[t]) - fov_arr = np.full(len(pred_dinov3), pos_name_pred) - t_arr = np.full(len(pred_dinov3), t, dtype=np.int32) - pred_dinov3_fovs.append(fov_arr) - gt_dinov3_fovs.append(fov_arr) - pred_dinov3_ts.append(t_arr) - gt_dinov3_ts.append(t_arr) - if pred_dynaclr.size > 0: - pred_dynaclr_feats.append(pred_dynaclr) - gt_dynaclr_feats.append(gt_dynaclr_per_t[t]) - fov_arr = np.full(len(pred_dynaclr), pos_name_pred) - t_arr = np.full(len(pred_dynaclr), t, dtype=np.int32) - pred_dynaclr_fovs.append(fov_arr) - gt_dynaclr_fovs.append(fov_arr) - pred_dynaclr_ts.append(t_arr) - gt_dynaclr_ts.append(t_arr) - if pred_celldino is not None and pred_celldino.size > 0: - pred_celldino_feats.append(pred_celldino) - gt_celldino_feats.append(gt_celldino_per_t[t]) - fov_arr = np.full(len(pred_celldino), pos_name_pred) - t_arr = np.full(len(pred_celldino), t, dtype=np.int32) - pred_celldino_fovs.append(fov_arr) - gt_celldino_fovs.append(fov_arr) - pred_celldino_ts.append(t_arr) - gt_celldino_ts.append(t_arr) - - maybe_empty_cuda_cache(t, runtime.cuda_empty_cache_every_n_timepoints) - - with region_timer("seg_write", pos_name_pred): - seg = np.stack(segmentations, axis=0) # shape: (T, 2, D, H, W) - row, col, fov = pos_name_pred.split("/") - seg_pos = segmentation_results.create_position(row, col, fov) - seg_pos.create_image("0", seg.astype(bool)) - - if config.compute_microssim: - with region_timer("microssim", pos_name_pred): - microssim_scores = calculate_microssim(microssim_data) - for i in range(T): - fov_pixel_metrics[i]["MicroMS3IM"] = float(microssim_scores[i]["MicroMS3IM"]) - - all_pixel_metrics.extend(fov_pixel_metrics) + result = _process_one_fov( + config, + runtime.cuda_empty_cache_every_n_timepoints, + pos_name_pred, + pos_pred, + pos_gt, + pos_seg, + io_config, + cache_ctx, + pred_cache_ctx, + seg_model, + dinov3_feature_extractor, + dynaclr_feature_extractor, + celldino_feature_extractor, + ) + _aggregate_fov_result( + result, + segmentation_results, + all_pixel_metrics, + all_mask_metrics, + all_feature_metrics, + pred_cp_feats, + pred_cp_fovs, + pred_cp_ts, + gt_cp_feats, + gt_cp_fovs, + gt_cp_ts, + pred_dinov3_feats, + pred_dinov3_fovs, + pred_dinov3_ts, + gt_dinov3_feats, + gt_dinov3_fovs, + gt_dinov3_ts, + pred_dynaclr_feats, + pred_dynaclr_fovs, + pred_dynaclr_ts, + gt_dynaclr_feats, + gt_dynaclr_fovs, + gt_dynaclr_ts, + pred_celldino_feats, + pred_celldino_fovs, + pred_celldino_ts, + gt_celldino_feats, + gt_celldino_fovs, + gt_celldino_ts, + ) # Flush manifest after each position so interrupted runs preserve progress. flush_manifest(cache_ctx) From 58850066c03091878d8c564f84c93c8d79a3f855 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 11:55:38 -0700 Subject: [PATCH 04/20] feat(dynacell): wire ProcessPoolExecutor + GPU serialization lock (C4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds worker primitives in runtime.py (_worker_initializer, gpu_serialization_lock, make_fov_executor) and the worker entry point _worker_run_fov in pipeline.py. evaluate_predictions now branches on runtime.executor: serial (default) runs _process_one_fov inline identical to today; process spawns a ProcessPoolExecutor over FOVs. Worker design: - Spawn context avoids inheriting parent CUDA. _worker_initializer sets BLAS/OMP env caps BEFORE apply_thread_budget lazy-imports torch. - _worker_setup lazy-loads seg_model + extractors under gpu_serialization_lock(): N workers don't init CUDA + load weights concurrently. Plate handles cached per-worker. - GPU touch sites inside _process_one_fov are wrapped in gpu_serialization_lock(): mask_gt/pred, cp_gt, deep_gt_{dinov3,dynaclr, celldino}, features_pred_per_t, pixel_metrics (per-T, GPU-bound end-to-end including cubic spectral), and on-demand cellpose fallback inside mask_metrics. In serial mode the lock has no path set and yields immediately — bit-for-bit-identical behavior. - Lock file at /tmp/dynacell_gpu_.lock (node-local, not NFS). One lock per SLURM job, shared across workers. Phase-2 runtime resolution: after pred_positions is built, resolve_runtime is re-called with n_positions= to clamp "auto" fov_workers to min(provisional, n_positions) and freeze threads_per_worker at the Phase-1 value so worker initializers match the parent's already-applied BLAS cap. Order: results collected via as_completed (arrival order) but aggregated in original pred_positions order so per-FOV CSV + embedding NPZ row order is deterministic regardless of executor. Workers flush their own ctx manifest at FOV end (global manifest lock at _pos_write_lock(ctx, "manifest", "global") fcntl serializes flushes across workers); interrupted process-mode runs preserve cache progress the same way serial mode did. Tests: full dynacell suite 384/48 pass/fail (identical to baseline; all 48 failures pre-existing). No new spawn-mode test in this commit — covered by C5 CPU integration test. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/dynacell/evaluation/pipeline.py | 263 +++++++++++++++--- .../src/dynacell/evaluation/runtime.py | 93 +++++++ 2 files changed, 314 insertions(+), 42 deletions(-) diff --git a/applications/dynacell/src/dynacell/evaluation/pipeline.py b/applications/dynacell/src/dynacell/evaluation/pipeline.py index 08ecc4f19..f32fffb2e 100644 --- a/applications/dynacell/src/dynacell/evaluation/pipeline.py +++ b/applications/dynacell/src/dynacell/evaluation/pipeline.py @@ -4,6 +4,7 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from pathlib import Path +from typing import Any import hydra import numpy as np @@ -224,6 +225,7 @@ def _process_one_fov( """ from dynacell.evaluation.runtime import ( get_timings, + gpu_serialization_lock, maybe_empty_cuda_cache, region_timer, ) @@ -240,10 +242,10 @@ def _process_one_fov( T = predict.shape[0] - with region_timer("mask_gt", pos_name_pred): + with region_timer("mask_gt", pos_name_pred), gpu_serialization_lock(): gt_mask_stack = fov_masks(cache_ctx, pos_name_pred, target, seg_model) if pred_cache_ctx.enabled: - with region_timer("mask_pred", pos_name_pred): + with region_timer("mask_pred", pos_name_pred), gpu_serialization_lock(): pred_mask_stack = fov_masks(pred_cache_ctx, pos_name_pred, predict, seg_model) else: pred_mask_stack = None @@ -254,18 +256,18 @@ def _process_one_fov( gt_celldino_per_t = None pred_per_t = None if config.compute_feature_metrics: - with region_timer("cp_gt", pos_name_pred): + with region_timer("cp_gt", pos_name_pred), gpu_serialization_lock(): gt_cp_per_t = fov_cp_features(cache_ctx, pos_name_pred, target, cell_segmentation) - with region_timer("deep_gt_dinov3", pos_name_pred): + with region_timer("deep_gt_dinov3", pos_name_pred), gpu_serialization_lock(): gt_dinov3_per_t = fov_deep_features( cache_ctx, pos_name_pred, target, cell_segmentation, dinov3_feature_extractor, "dinov3" ) - with region_timer("deep_gt_dynaclr", pos_name_pred): + with region_timer("deep_gt_dynaclr", pos_name_pred), gpu_serialization_lock(): gt_dynaclr_per_t = fov_deep_features( cache_ctx, pos_name_pred, target, cell_segmentation, dynaclr_feature_extractor, "dynaclr" ) if celldino_feature_extractor is not None: - with region_timer("deep_gt_celldino", pos_name_pred): + with region_timer("deep_gt_celldino", pos_name_pred), gpu_serialization_lock(): gt_celldino_per_t = fov_deep_features( cache_ctx, pos_name_pred, @@ -274,7 +276,7 @@ def _process_one_fov( celldino_feature_extractor, "celldino", ) - with region_timer("features_pred_per_t", pos_name_pred): + with region_timer("features_pred_per_t", pos_name_pred), gpu_serialization_lock(): pred_per_t = _fov_pred_features_per_t( pred_cache_ctx, pos_name_pred, @@ -300,7 +302,7 @@ def _process_one_fov( for t in tqdm(range(T), desc="Processing timepoints", leave=False): data_info = {"FOV": pos_name_pred, "Timepoint": t} - with region_timer("pixel_metrics", pos_name_pred, t): + with region_timer("pixel_metrics", pos_name_pred, t), gpu_serialization_lock(): pixel_metrics = compute_pixel_metrics( predict[t], target[t], @@ -315,11 +317,13 @@ def _process_one_fov( with region_timer("mask_metrics", pos_name_pred, t): segmented_target = gt_mask_stack[t] - segmented_predict = ( - pred_mask_stack[t] - if pred_mask_stack is not None - else np.asarray(segment(predict[t], config.target_name, seg_model=seg_model)).astype(bool) - ) + if pred_mask_stack is not None: + segmented_predict = pred_mask_stack[t] + else: + with gpu_serialization_lock(): + segmented_predict = np.asarray(segment(predict[t], config.target_name, seg_model=seg_model)).astype( + bool + ) fov_mask_metrics.append({**data_info, **evaluate_segmentations(segmented_predict, segmented_target)}) segmentations.append(np.stack([segmented_predict, segmented_target], axis=0)) @@ -483,11 +487,147 @@ def _aggregate_fov_result( gt_celldino_ts.extend(result.celldino.gt_ts) +# Worker-side state for ``executor=process``. A spawn-context child Python +# interpreter calls ``_worker_run_fov`` once per submitted FOV; the first call +# triggers ``_worker_setup`` which lazy-loads seg model + extractors + plate +# handles under the GPU serialization lock and caches them here. Subsequent +# FOVs in the same worker reuse the cached state. +_WORKER_STATE: dict[str, Any] = {} + + +def _worker_setup(config: DictConfig) -> None: + """First-FOV-per-worker initialization: load models, open plates, init caches. + + Wraps all GPU touch sites in ``gpu_serialization_lock()`` so N workers + don't initialize CUDA contexts + load weights concurrently. Cache contexts + and iohub plates are opened outside the lock (no GPU touch). + """ + if _WORKER_STATE.get("initialized"): + return + from dynacell.evaluation.pipeline_cache import init_cache_context, resolve_dynaclr_encoder_cfg + from dynacell.evaluation.runtime import gpu_serialization_lock + from dynacell.evaluation.segmentation import prepare_segmentation_model + from dynacell.evaluation.utils import ( + CellDinoFeatureExtractor, + DinoV3FeatureExtractor, + DynaCLRFeatureExtractor, + ) + + with gpu_serialization_lock(): + seg_model = prepare_segmentation_model(config) + dinov3_feature_extractor = None + dynaclr_feature_extractor = None + celldino_feature_extractor = None + dinov3_model_name = None + dynaclr_ckpt_path = None + dynaclr_encoder_cfg = None + celldino_weights_path = None + if config.compute_feature_metrics: + dinov3_model_name = config.feature_extractor.dinov3.pretrained_model_name + dinov3_feature_extractor = DinoV3FeatureExtractor(dinov3_model_name) + dynaclr_config = config.feature_extractor.dynaclr + dynaclr_ckpt_path = str(dynaclr_config.checkpoint) + dynaclr_encoder_cfg = resolve_dynaclr_encoder_cfg(config) + dynaclr_feature_extractor = DynaCLRFeatureExtractor( + checkpoint=dynaclr_config.checkpoint, + encoder_config=dynaclr_encoder_cfg, + ) + celldino_cfg = config.feature_extractor.celldino + if celldino_cfg.weights_path is not None: + celldino_weights_path = str(celldino_cfg.weights_path) + celldino_feature_extractor = CellDinoFeatureExtractor( + weights_path=celldino_weights_path, + img_size=int(celldino_cfg.img_size), + patch_size=int(celldino_cfg.patch_size), + ) + + cache_ctx = init_cache_context( + config, + side="gt", + dinov3_model_name=dinov3_model_name, + dynaclr_ckpt_path=dynaclr_ckpt_path, + dynaclr_encoder_cfg=dynaclr_encoder_cfg, + celldino_weights_path=celldino_weights_path, + ) + pred_cache_ctx = init_cache_context( + config, + side="pred", + dinov3_model_name=dinov3_model_name, + dynaclr_ckpt_path=dynaclr_ckpt_path, + dynaclr_encoder_cfg=dynaclr_encoder_cfg, + celldino_weights_path=celldino_weights_path, + ) + + pred_plate = open_ome_zarr(Path(config.io.pred_path), mode="r") + gt_plate = open_ome_zarr(Path(config.io.gt_path), mode="r") + seg_plate = None + if config.io.cell_segmentation_path is not None: + seg_plate = open_ome_zarr(Path(config.io.cell_segmentation_path), mode="r") + + _WORKER_STATE.update( + { + "initialized": True, + "seg_model": seg_model, + "dinov3": dinov3_feature_extractor, + "dynaclr": dynaclr_feature_extractor, + "celldino": celldino_feature_extractor, + "cache_ctx": cache_ctx, + "pred_cache_ctx": pred_cache_ctx, + "pred_plate": pred_plate, + "gt_plate": gt_plate, + "seg_plate": seg_plate, + "pred_positions": dict(pred_plate.positions()), + "gt_positions": dict(gt_plate.positions()), + "seg_positions": dict(seg_plate.positions()) if seg_plate is not None else None, + } + ) + + +def _worker_run_fov(config: DictConfig, pos_name: str, cuda_empty_every_n: int) -> FovResult: + """Worker entry point: process one FOV by name and return FovResult. + + Submitted via ``pool.submit`` in ``executor=process`` mode. Returns a + pickle-friendly ``FovResult`` (no plate handles or torch modules + cross the process boundary). + """ + from dynacell.evaluation.pipeline_cache import flush_manifest + + _worker_setup(config) + state = _WORKER_STATE + pos_pred = state["pred_positions"][pos_name] + pos_gt = state["gt_positions"][pos_name] + pos_seg = state["seg_positions"][pos_name] if state["seg_positions"] is not None else None + + result = _process_one_fov( + config, + cuda_empty_every_n, + pos_name, + pos_pred, + pos_gt, + pos_seg, + config.io, + state["cache_ctx"], + state["pred_cache_ctx"], + state["seg_model"], + state["dinov3"], + state["dynaclr"], + state["celldino"], + ) + # Worker-side manifest flush so interrupted runs preserve progress even + # when the parent dies. The global manifest lock at + # _pos_write_lock(ctx, "manifest", "global") serializes flushes across + # workers. + flush_manifest(state["cache_ctx"]) + flush_manifest(state["pred_cache_ctx"]) + return result + + def evaluate_predictions(config: DictConfig): """Evaluate predictions on all test images.""" from dynacell.evaluation.runtime import ( apply_thread_budget, dump_timings_csv, + make_fov_executor, maybe_gc_collect, reset_timings, resolve_runtime, @@ -665,32 +805,16 @@ def evaluate_predictions(config: DictConfig): for ctx in sides_for_precompute.values(): flush_manifest(ctx) - for fov_idx, (p1, p2, p3) in enumerate( - tqdm( - zip(pred_positions, gt_positions, seg_positions), - total=len(pred_positions), - desc="Processing positions", - ) - ): - pos_name_pred, pos_pred = p1 - _, pos_gt = p2 - _, pos_seg = p3 + # Phase 2 runtime resolution: clamp fov_workers to n_positions + # now that we know it. threads_per_worker is frozen at Phase 1 + # so the parent's BLAS cap matches what workers see. + runtime = resolve_runtime( + config, + n_positions=len(pred_positions), + freeze_threads_per_worker=runtime.threads_per_worker, + ) - result = _process_one_fov( - config, - runtime.cuda_empty_cache_every_n_timepoints, - pos_name_pred, - pos_pred, - pos_gt, - pos_seg, - io_config, - cache_ctx, - pred_cache_ctx, - seg_model, - dinov3_feature_extractor, - dynaclr_feature_extractor, - celldino_feature_extractor, - ) + def _aggregate(result: FovResult) -> None: _aggregate_fov_result( result, segmentation_results, @@ -723,11 +847,66 @@ def evaluate_predictions(config: DictConfig): gt_celldino_ts, ) - # Flush manifest after each position so interrupted runs preserve progress. - flush_manifest(cache_ctx) - flush_manifest(pred_cache_ctx) + if runtime.executor == "serial": + for fov_idx, (p1, p2, p3) in enumerate( + tqdm( + zip(pred_positions, gt_positions, seg_positions), + total=len(pred_positions), + desc="Processing positions", + ) + ): + pos_name_pred, pos_pred = p1 + _, pos_gt = p2 + _, pos_seg = p3 + + result = _process_one_fov( + config, + runtime.cuda_empty_cache_every_n_timepoints, + pos_name_pred, + pos_pred, + pos_gt, + pos_seg, + io_config, + cache_ctx, + pred_cache_ctx, + seg_model, + dinov3_feature_extractor, + dynaclr_feature_extractor, + celldino_feature_extractor, + ) + _aggregate(result) + + # Flush manifest after each position so interrupted runs preserve progress. + flush_manifest(cache_ctx) + flush_manifest(pred_cache_ctx) - maybe_gc_collect(fov_idx, runtime.gc_collect_every_n_fovs) + maybe_gc_collect(fov_idx, runtime.gc_collect_every_n_fovs) + else: + # executor == "process": spawn-context ProcessPoolExecutor. + # Workers lazy-load their own seg_model + extractors under the + # GPU lock; parent's copy stays loaded (used only here for + # checkpoint pre-warm side effect — discarded after this block). + from concurrent.futures import as_completed + + pos_names_in_order = [p[0] for p, _, _ in zip(pred_positions, gt_positions, seg_positions)] + with make_fov_executor(runtime) as pool: + if pool is None: + raise RuntimeError("make_fov_executor returned None for executor='process'") + futures = { + pool.submit( + _worker_run_fov, config, pos_name, runtime.cuda_empty_cache_every_n_timepoints + ): pos_name + for pos_name in pos_names_in_order + } + results_by_pos: dict[str, FovResult] = {} + for fut in tqdm(as_completed(futures), total=len(futures), desc="Processing positions"): + pos_name = futures[fut] + results_by_pos[pos_name] = fut.result() + # Aggregate in position order (independent of completion order) + # so per-FOV CSV / embedding NPZ row order is deterministic. + for fov_idx, pos_name in enumerate(pos_names_in_order): + _aggregate(results_by_pos[pos_name]) + maybe_gc_collect(fov_idx, runtime.gc_collect_every_n_fovs) finally: if seg_plate is not None: seg_plate.close() diff --git a/applications/dynacell/src/dynacell/evaluation/runtime.py b/applications/dynacell/src/dynacell/evaluation/runtime.py index 248905c27..6b76fe67a 100644 --- a/applications/dynacell/src/dynacell/evaluation/runtime.py +++ b/applications/dynacell/src/dynacell/evaluation/runtime.py @@ -372,3 +372,96 @@ def maybe_gc_collect(fov_idx: int, every_n: int) -> None: if (fov_idx + 1) % every_n != 0: return gc.collect() + + +# --------------------------------------------------------------------------- +# Process-pool primitives for FOV-level parallelism (C4). +# Spawn context so each worker gets its own CUDA context and BLAS load. +# --------------------------------------------------------------------------- + +# Worker-global GPU lock path, set by _worker_initializer. ``None`` in the +# parent / serial mode -> gpu_serialization_lock is a no-op. +_GPU_LOCK_PATH: str | None = None + + +def _worker_initializer(threads: int, gpu_lock_path: str | None) -> None: + """Initialize a freshly-spawned worker process. + + Order matters: env caps are set BEFORE ``apply_thread_budget`` is called, + which lazy-imports torch. If torch were imported first, ``CUDA_VISIBLE_DEVICES`` + overrides (if any) would come too late. + """ + global _GPU_LOCK_PATH + # Env caps first so BLAS C-extensions load with the right thread count. + for var in _THREAD_ENV_VARS: + os.environ.setdefault(var, str(threads)) + apply_thread_budget(threads) + _GPU_LOCK_PATH = gpu_lock_path + + +@contextmanager +def gpu_serialization_lock() -> Iterator[None]: + """Serialize GPU operations across workers via an ``fcntl.flock``. + + No-op when ``_GPU_LOCK_PATH`` is unset (parent / serial mode). + Used to wrap model loads and per-FOV GPU operations under + ``executor=process``. + """ + if _GPU_LOCK_PATH is None: + yield + return + import fcntl + + with open(_GPU_LOCK_PATH, "w") as f: + fcntl.flock(f.fileno(), fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(f.fileno(), fcntl.LOCK_UN) + + +def gpu_lock_path_for_job() -> str: + """Return a node-local lock path tagged by SLURM_JOB_ID (or PID as fallback). + + Stays on local /tmp so fcntl serialization is fast and doesn't go through + NFS lockd. One lock per job, shared across all workers in that job. + """ + import tempfile + + tag = os.environ.get("SLURM_JOB_ID", str(os.getpid())) + return str(Path(tempfile.gettempdir()) / f"dynacell_gpu_{tag}.lock") + + +@contextmanager +def make_fov_executor(runtime: ResolvedRuntime) -> Iterator[Any]: + """Yield a ``ProcessPoolExecutor`` for ``executor=process`` or ``None`` for serial. + + Callers branch on the yielded value: ``None`` means run inline; an + ``Executor`` means submit ``_process_one_fov`` invocations to workers. + + Cleanup: on exit (normal or exception), the pool is shut down with + ``cancel_futures=True``. In-flight workers continue to completion + (Python's ``ProcessPoolExecutor`` does not signal them); user may need + ``scancel`` to fully release GPU memory on Ctrl-C. + """ + if runtime.executor == "serial": + yield None + return + + import multiprocessing as mp + from concurrent.futures import ProcessPoolExecutor + + lock_path = gpu_lock_path_for_job() + # Touch the lock file so workers can flock-open it before any FOV runs. + Path(lock_path).touch() + ctx = mp.get_context("spawn") + pool = ProcessPoolExecutor( + max_workers=runtime.fov_workers, + mp_context=ctx, + initializer=_worker_initializer, + initargs=(runtime.threads_per_worker, lock_path), + ) + try: + yield pool + finally: + pool.shutdown(wait=False, cancel_futures=True) From fa5c55c545bf64181a000a05ce545c39172b4686 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 12:19:33 -0700 Subject: [PATCH 05/20] test(dynacell): runtime + parallel tests, CI wiring, SLURM env caps, docs (C5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds test coverage for the C1-C4 runtime module and FOV-parallelism wiring: test_runtime.py (23 tests): apply_thread_budget no-crash semantics under Hydra's main; resolve_runtime auto/freeze/validate paths; region timing collector; DYNACELL_FORCE_PER_T_HYGIENE env-var hook; spawn-context ProcessPoolExecutor smoke via stdlib os.getpid. test_evaluation_pipeline_parallel.py (3 tests): FovResult pickle round-trip with realistic per-backbone shapes (the path that ProcessPoolExecutor takes for Future.result transport); _aggregate_fov_result extends all 24 dataset-level accumulator lists correctly. Re-imports the live pipeline module inside test bodies because test_lazy_init.py pops dynacell.* modules from sys.modules, which would invalidate pickle resolution of a module-level imported FovResult. CI wiring: - Root pyproject.toml registers the "slow" marker and appends -m "not slow" to existing addopts. - test-dynacell-configs runs test_runtime.py + test_evaluation_pipeline_parallel.py alongside existing benchmark-schema/submit-tool tests. No new heavy deps needed (these tests don't import iohub, cubic, transformers). SLURM env caps: - run_eval_jointtrained_nucleus.slurm and run_eval_unext2_jointtrained_membrane.slurm export DYNACELL_THREADS_PER_WORKER=\$SLURM_CPUS_PER_TASK before invoking eval. Docs: - applications/dynacell/CLAUDE.md gains "Eval runtime / parallelism" section documenting the config block, env-var hooks, timing instrumentation, and process-mode caveats. Full dynacell test suite: 410 passed / 48 failed / 65 skipped — the 48 failures are the same set that fails on dynacell-models@HEAD (verified parity across all 5 commits in this branch). End-to-end serial-vs-process parity test on a real iohub+cache fixture is a follow-up (see .claude/plans/eval-parallelism.md §C5). Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/test.yml | 14 +- applications/dynacell/CLAUDE.md | 40 +++ .../test_evaluation_pipeline_parallel.py | 228 ++++++++++++++++ applications/dynacell/tests/test_runtime.py | 256 ++++++++++++++++++ .../tools/run_eval_jointtrained_nucleus.slurm | 9 + ...un_eval_unext2_jointtrained_membrane.slurm | 7 + pyproject.toml | 5 +- 7 files changed, 556 insertions(+), 3 deletions(-) create mode 100644 applications/dynacell/tests/test_evaluation_pipeline_parallel.py create mode 100644 applications/dynacell/tests/test_runtime.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index f475f7e7f..0aca6ade6 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -137,8 +137,18 @@ jobs: run: uv sync --frozen --group test working-directory: applications/dynacell - - name: Run benchmark-schema + submit-tool tests - run: uv run --frozen pytest tests/test_benchmark_config_composition.py tests/test_submit_benchmark_job.py -v + - name: Run benchmark-schema + submit-tool + eval-runtime tests + # tests/test_runtime.py and tests/test_evaluation_pipeline_parallel.py + # cover the dynacell.evaluation.runtime module + FovResult pickle + # contract that gate the FOV-level parallelism wiring. They don't + # need the eval extras (no iohub fixture or model loads). + run: | + uv run --frozen pytest \ + tests/test_benchmark_config_composition.py \ + tests/test_submit_benchmark_job.py \ + tests/test_runtime.py \ + tests/test_evaluation_pipeline_parallel.py \ + -v working-directory: applications/dynacell check: diff --git a/applications/dynacell/CLAUDE.md b/applications/dynacell/CLAUDE.md index 47756a8fe..a421898ee 100644 --- a/applications/dynacell/CLAUDE.md +++ b/applications/dynacell/CLAUDE.md @@ -32,3 +32,43 @@ Set by `trainer.callbacks[…HCSPredictionWriter].init_args.output_store` in eac Where `` is `nucl` / `memb` / `sec61b` / `tomm20`, `` is the **code name** from the table above (e.g. `fcmae_vscyto3d_scratch`, `fnet3d_paper`), and `` is `mock` / `denv` / `zikv`. The (no-infix) iPSC-trained naming is historical baggage from before joint/A549 training existed; don't add a `_ipsctrained` infix retroactively. Output dirs: iPSC test predictions land under `ipsc/predictions/`, A549 plate predictions under `a549/predictions/`, regardless of training set. Caveat: Dihan's earlier ER + Mito iPSC-trained zarrs use a legacy `____.zarr` shape (e.g. `sec61b_fcmae_vscyto3d_scratch__sec61b_mock.zarr`, double-underscore + redundant gene prefix). New leaves should follow the table above; do not propagate the legacy form. + +## Eval runtime / parallelism + +`dynacell.evaluation.runtime` (added 2026-05) provides three layered thread-cap entry points + optional FOV-level parallelism via spawn-context `ProcessPoolExecutor`. Defaults preserve sequential behavior; opt in via the `runtime:` block in `eval.yaml`. + +### Config block (in `_configs/eval.yaml`) + +```yaml +runtime: + fov_workers: 1 # int | "auto" + threads_per_worker: "auto" # int | "auto" -> cpu_count // fov_workers + executor: "serial" # "serial" | "process" + cuda_empty_cache_every_n_timepoints: 0 # 0 = off + gc_collect_every_n_fovs: 0 # 0 = off +``` + +- `executor=serial` (default): inline FOV loop, identical to pre-runtime-module behavior. +- `executor=process`: spawn-context `ProcessPoolExecutor` over FOVs. Each worker independently lazy-loads `seg_model` + extractors under an fcntl GPU lock at `/tmp/dynacell_gpu_.lock`, so models stay GPU-resident per worker (N × model weights on the GPU) but only one worker runs GPU work at a time. Suitable when GPU memory has headroom for N model copies. +- `fov_workers: "auto"` resolves to 1 under `executor=serial`; under `executor=process` it clamps to `min(cpu_count // threads_per_worker, n_positions)`. +- Literal `fov_workers > 1` with `executor=serial` raises. Literal `fov_workers=1` with `executor=process` auto-demotes to `serial` to avoid the ~5 s spawn cold-start cost for a single-worker pool. +- Two-phase resolve: parent applies BLAS cap at function entry (provisional), then re-resolves after the position list is built with `freeze_threads_per_worker=` so worker initializers see the same value the parent capped to. + +### Env-var entry points + +- `DYNACELL_THREADS_PER_WORKER=N` — set in SLURM scripts BEFORE invoking `dynacell evaluate`. Exports `OMP_NUM_THREADS` / `MKL_NUM_THREADS` / `OPENBLAS_NUM_THREADS` at C-extension load time (in-process `apply_thread_budget` is a runtime safety net but can come after BLAS load). The `__main__:main_cli` first statement reads this var. +- `DYNACELL_FORCE_PER_T_HYGIENE=1` — operator escape hatch: flips both `cuda_empty_cache_every_n_timepoints` and `gc_collect_every_n_fovs` to ≥1 at runtime regardless of YAML defaults. Useful for post-ship mitigation of per-T memory degradation without a code change. + +### Timing instrumentation + +Region timers are always on (overhead is ~120 ms on a 9-h eval, ~4e-6 of wall time). Output goes to `/eval_timing.csv` at end of run with columns `pos_name, t, region, seconds`. Regions tag the FOV-level work (`mask_gt`, `mask_pred`, `cp_gt`, `deep_gt_{dinov3,dynaclr,celldino}`, `features_pred_per_t`, `microssim`, `seg_write`) and the per-T work (`pixel_metrics`, `mask_metrics`, `feature_pairwise`). + +Under `executor=process`, workers return their slice of the timing log inside `FovResult.timings`; the parent aggregator concatenates. + +### Reality check for process mode + +- The parent still loads the seg model once at `pipeline.py:171` (used for checkpoint pre-warm side-effect under `process` mode + the seg model itself under `serial` mode). +- `precompute_deep_features` (the upfront batched feature pre-fill) stays in the parent — single-pass over positions with the `DeepFeatureBatcher` cross-FOV amortization. Parallelizing precompute is a separate plan. +- Workers cannot share open iohub plate handles or torch modules across the pickle boundary; each worker re-opens plates on first FOV. +- `ProcessPoolExecutor.shutdown(wait=False, cancel_futures=True)` cancels queued futures only — in-flight workers continue to completion (~minutes if mid-cellpose). Ctrl-C may need `scancel` to fully release GPU memory. +- Tests for the runtime module + FovResult pickle contract live in `applications/dynacell/tests/test_runtime.py` and `tests/test_evaluation_pipeline_parallel.py`. An end-to-end serial-vs-process parity test on a real iohub fixture is a follow-up (see plan `.claude/plans/eval-parallelism.md` §C5). diff --git a/applications/dynacell/tests/test_evaluation_pipeline_parallel.py b/applications/dynacell/tests/test_evaluation_pipeline_parallel.py new file mode 100644 index 000000000..702ff1f31 --- /dev/null +++ b/applications/dynacell/tests/test_evaluation_pipeline_parallel.py @@ -0,0 +1,228 @@ +"""Worker-side and aggregator-side smoke tests for the FOV parallelism wiring. + +These tests cover the cross-process boundary in two pieces: + +1. ``FovResult`` survives ``pickle`` round-trip with realistic + shapes (the ``concurrent.futures.ProcessPoolExecutor`` future-result + transport uses pickle). +2. ``_aggregate_fov_result`` correctly extends the parent's 24 + dataset-level accumulator lists from a synthetic FovResult, preserving + the existing ``if size > 0: append`` semantics that the worker + internally maintains. + +End-to-end serial/process parity on a real eval (with iohub fixtures and +prebuilt mask caches) is a follow-up CPU integration test — see the plan +``.claude/plans/eval-parallelism.md`` C5 section. +""" + +from __future__ import annotations + +import importlib +import pickle + +import numpy as np + + +def _live_pipeline_module(): + """Return the *currently cached* dynacell.evaluation.pipeline module. + + Other tests in the suite (notably ``test_evaluation_pipeline.py``) + use ``monkeypatch.setitem(sys.modules, ...)`` to swap the pipeline + module with a stubbed re-import. After teardown, ``sys.modules`` is + restored but any class references resolved at our module's import + time still point at the *stubbed* re-import — pickle then fails with + ``"it's not the same object as dynacell.evaluation.pipeline.FovResult"``. + Resolving lazily inside each test avoids the stale binding. + """ + return importlib.import_module("dynacell.evaluation.pipeline") + + +def _make_synthetic_result( + pos_name: str = "A/1/0", + t_count: int = 2, + d: int = 4, + h: int = 16, + w: int = 16, + cp_dim: int = 8, + deep_dim: int = 768, + cells_per_t: int = 3, +): + """Build a FovResult with realistic shapes for round-trip testing.""" + pipeline = _live_pipeline_module() + FovResult = pipeline.FovResult + _BackboneLists = pipeline._BackboneLists + row, col, fov = pos_name.split("/") + per_t_pixel = [{"FOV": pos_name, "Timepoint": t, "PCC": 0.9 + 0.01 * t, "MicroMS3IM": 0.85} for t in range(t_count)] + per_t_mask = [{"FOV": pos_name, "Timepoint": t, "DICE": 0.7 + 0.01 * t} for t in range(t_count)] + per_t_feat = [{"FOV": pos_name, "Timepoint": t, "CP_cos": 0.8, "DINOv3_cos": 0.95} for t in range(t_count)] + seg_array = np.zeros((t_count, 2, d, h, w), dtype=bool) + seg_array[:, 0, :, : h // 2, :] = True # pred channel + seg_array[:, 1, :, : h // 2, : w // 2] = True # gt channel + + cp = _BackboneLists() + dinov3 = _BackboneLists() + dynaclr = _BackboneLists() + celldino = _BackboneLists() + for t in range(t_count): + fov_arr = np.full(cells_per_t, pos_name) + t_arr = np.full(cells_per_t, t, dtype=np.int32) + cp.pred_feats.append(np.full((cells_per_t, cp_dim), float(t), dtype=np.float32)) + cp.gt_feats.append(np.full((cells_per_t, cp_dim), float(t) + 0.5, dtype=np.float32)) + cp.pred_fovs.append(fov_arr) + cp.gt_fovs.append(fov_arr) + cp.pred_ts.append(t_arr) + cp.gt_ts.append(t_arr) + for bl in (dinov3, dynaclr, celldino): + bl.pred_feats.append(np.full((cells_per_t, deep_dim), float(t), dtype=np.float32)) + bl.gt_feats.append(np.full((cells_per_t, deep_dim), float(t) + 0.5, dtype=np.float32)) + bl.pred_fovs.append(fov_arr) + bl.gt_fovs.append(fov_arr) + bl.pred_ts.append(t_arr) + bl.gt_ts.append(t_arr) + + return FovResult( + pos_name=pos_name, + row=row, + col=col, + fov=fov, + per_t_pixel_rows=per_t_pixel, + per_t_mask_rows=per_t_mask, + per_t_feature_rows=per_t_feat, + seg_array=seg_array, + cp=cp, + dinov3=dinov3, + dynaclr=dynaclr, + celldino=celldino, + timings=[(pos_name, None, "mask_gt", 0.05), (pos_name, 0, "pixel_metrics", 0.02)], + ) + + +def test_fov_result_pickle_round_trip_preserves_arrays(): + result = _make_synthetic_result() + restored = pickle.loads(pickle.dumps(result)) + assert restored.pos_name == result.pos_name + assert restored.row == "A" + assert restored.col == "1" + assert restored.fov == "0" + assert restored.seg_array.shape == result.seg_array.shape + assert restored.seg_array.dtype == np.bool_ + assert np.array_equal(restored.seg_array, result.seg_array) + for backbone_attr in ("cp", "dinov3", "dynaclr", "celldino"): + original = getattr(result, backbone_attr) + restored_bb = getattr(restored, backbone_attr) + for list_name in ("pred_feats", "gt_feats", "pred_fovs", "gt_fovs", "pred_ts", "gt_ts"): + for a, b in zip(getattr(original, list_name), getattr(restored_bb, list_name)): + assert np.array_equal(a, b) + assert restored.per_t_pixel_rows == result.per_t_pixel_rows + assert restored.timings == result.timings + + +def test_fov_result_pickle_handles_empty_backbones(): + """An empty _BackboneLists round-trips correctly (no len-zero -> None coercion).""" + result = _make_synthetic_result() + # Wipe one backbone to simulate "feature_metrics disabled" / "zero cells". + result.celldino.pred_feats.clear() + result.celldino.gt_feats.clear() + result.celldino.pred_fovs.clear() + result.celldino.gt_fovs.clear() + result.celldino.pred_ts.clear() + result.celldino.gt_ts.clear() + restored = pickle.loads(pickle.dumps(result)) + assert restored.celldino.pred_feats == [] + assert restored.celldino.gt_feats == [] + + +def test_aggregate_fov_result_extends_24_accumulators(): + """Aggregator must extend each backbone's 6 lists with worker contributions.""" + pipeline = _live_pipeline_module() + _aggregate_fov_result = pipeline._aggregate_fov_result + # Mock segmentation_results plate handle (only create_position is exercised). + written = {} + + class _FakeSegPos: + def create_image(self, name, data): + written[(row_, col_, fov_, name)] = np.asarray(data) + + class _FakeSegPlate: + def create_position(self, row, col, fov): + nonlocal row_, col_, fov_ + row_, col_, fov_ = row, col, fov + return _FakeSegPos() + + row_ = col_ = fov_ = None + + result = _make_synthetic_result() + + all_pix: list[dict] = [] + all_mask: list[dict] = [] + all_feat: list[dict] = [] + pred_cp_feats: list = [] + pred_cp_fovs: list = [] + pred_cp_ts: list = [] + gt_cp_feats: list = [] + gt_cp_fovs: list = [] + gt_cp_ts: list = [] + pred_dinov3_feats: list = [] + pred_dinov3_fovs: list = [] + pred_dinov3_ts: list = [] + gt_dinov3_feats: list = [] + gt_dinov3_fovs: list = [] + gt_dinov3_ts: list = [] + pred_dynaclr_feats: list = [] + pred_dynaclr_fovs: list = [] + pred_dynaclr_ts: list = [] + gt_dynaclr_feats: list = [] + gt_dynaclr_fovs: list = [] + gt_dynaclr_ts: list = [] + pred_celldino_feats: list = [] + pred_celldino_fovs: list = [] + pred_celldino_ts: list = [] + gt_celldino_feats: list = [] + gt_celldino_fovs: list = [] + gt_celldino_ts: list = [] + + _aggregate_fov_result( + result, + _FakeSegPlate(), + all_pix, + all_mask, + all_feat, + pred_cp_feats, + pred_cp_fovs, + pred_cp_ts, + gt_cp_feats, + gt_cp_fovs, + gt_cp_ts, + pred_dinov3_feats, + pred_dinov3_fovs, + pred_dinov3_ts, + gt_dinov3_feats, + gt_dinov3_fovs, + gt_dinov3_ts, + pred_dynaclr_feats, + pred_dynaclr_fovs, + pred_dynaclr_ts, + gt_dynaclr_feats, + gt_dynaclr_fovs, + gt_dynaclr_ts, + pred_celldino_feats, + pred_celldino_fovs, + pred_celldino_ts, + gt_celldino_feats, + gt_celldino_fovs, + gt_celldino_ts, + ) + + assert len(all_pix) == 2 + assert len(all_mask) == 2 + assert len(all_feat) == 2 + assert (row_, col_, fov_) == ("A", "1", "0") + assert written[("A", "1", "0", "0")].shape == result.seg_array.shape + # Each backbone contributes 2 entries (one per timepoint). + for lst in ( + pred_cp_feats, + pred_dinov3_feats, + pred_dynaclr_feats, + pred_celldino_feats, + ): + assert len(lst) == 2 diff --git a/applications/dynacell/tests/test_runtime.py b/applications/dynacell/tests/test_runtime.py new file mode 100644 index 000000000..e87614e42 --- /dev/null +++ b/applications/dynacell/tests/test_runtime.py @@ -0,0 +1,256 @@ +"""Unit tests for ``dynacell.evaluation.runtime``.""" + +from __future__ import annotations + +import csv +import os +import pickle +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path + +import pytest +from omegaconf import OmegaConf + +from dynacell.evaluation.runtime import ( + ResolvedRuntime, + apply_thread_budget, + dump_timings_csv, + early_apply_env_caps, + extend_timings, + get_timings, + gpu_lock_path_for_job, + gpu_serialization_lock, + make_fov_executor, + maybe_empty_cuda_cache, + maybe_gc_collect, + region_timer, + reset_timings, + resolve_runtime, +) + + +def test_apply_thread_budget_runs_without_error(): + """apply_thread_budget should not raise even when called from Hydra's main. + + set_num_interop_threads raises after any parallel op has run; the function + catches that and continues. Test covers the import-after-parallel-op path. + """ + # Force-import torch (parallel op) before applying budget. + import torch + + _ = torch.zeros(4) + torch.ones(4) + apply_thread_budget(2) # should not raise + + +def test_apply_thread_budget_rejects_zero(): + with pytest.raises(ValueError, match="threads must be >= 1"): + apply_thread_budget(0) + + +def test_resolve_runtime_no_runtime_block_falls_back_to_serial(): + cfg = OmegaConf.create({"io": {}}) + r = resolve_runtime(cfg) + assert r.fov_workers == 1 + assert r.executor == "serial" + assert r.cuda_empty_cache_every_n_timepoints == 0 + assert r.gc_collect_every_n_fovs == 0 + + +def test_resolve_runtime_auto_serial_clamps_workers_to_one(): + cfg = OmegaConf.create({"runtime": {"fov_workers": "auto", "threads_per_worker": "auto", "executor": "serial"}}) + r = resolve_runtime(cfg) + assert r.fov_workers == 1 + assert r.executor == "serial" + assert r.threads_per_worker >= 1 + + +def test_resolve_runtime_auto_process_clamps_to_n_positions(): + cfg = OmegaConf.create({"runtime": {"fov_workers": "auto", "threads_per_worker": "auto", "executor": "process"}}) + r = resolve_runtime(cfg, n_positions=2) + assert r.fov_workers <= 2 + # Auto-demotes to serial when only 1 position is left after clamping. + if r.fov_workers == 1: + assert r.executor == "serial" + + +def test_resolve_runtime_freeze_threads_per_worker(): + """Phase-2 contract: ``freeze_threads_per_worker`` returns the frozen value.""" + cfg = OmegaConf.create({"runtime": {"fov_workers": "auto", "threads_per_worker": "auto", "executor": "process"}}) + r = resolve_runtime(cfg, n_positions=3, freeze_threads_per_worker=4) + assert r.threads_per_worker == 4 + + +def test_resolve_runtime_literal_fov_workers_in_serial_raises(): + cfg = OmegaConf.create({"runtime": {"fov_workers": 4, "executor": "serial"}}) + with pytest.raises(ValueError, match=r"fov_workers=4 requires runtime.executor='process'"): + resolve_runtime(cfg) + + +def test_resolve_runtime_auto_demote_process_to_serial_at_one_worker(): + """fov_workers literally set to 1 with executor=process should auto-demote to serial.""" + cfg = OmegaConf.create({"runtime": {"fov_workers": 1, "executor": "process"}}) + r = resolve_runtime(cfg) + assert r.fov_workers == 1 + assert r.executor == "serial" + + +def test_resolve_runtime_invalid_executor_raises(): + cfg = OmegaConf.create({"runtime": {"executor": "ray"}}) + with pytest.raises(ValueError, match="executor must be 'serial' or 'process'"): + resolve_runtime(cfg) + + +def test_resolve_runtime_force_per_t_hygiene_env(monkeypatch): + """DYNACELL_FORCE_PER_T_HYGIENE=1 turns on hygiene knobs regardless of YAML.""" + monkeypatch.setenv("DYNACELL_FORCE_PER_T_HYGIENE", "1") + cfg = OmegaConf.create( + { + "runtime": { + "cuda_empty_cache_every_n_timepoints": 0, + "gc_collect_every_n_fovs": 0, + } + } + ) + r = resolve_runtime(cfg) + assert r.cuda_empty_cache_every_n_timepoints >= 1 + assert r.gc_collect_every_n_fovs >= 1 + + +def test_early_apply_env_caps_respects_existing(monkeypatch): + """Should only export env vars when DYNACELL_THREADS_PER_WORKER is set.""" + monkeypatch.delenv("DYNACELL_THREADS_PER_WORKER", raising=False) + monkeypatch.delenv("OMP_NUM_THREADS", raising=False) + early_apply_env_caps() + assert os.environ.get("OMP_NUM_THREADS") is None # noop without env set + + monkeypatch.setenv("DYNACELL_THREADS_PER_WORKER", "3") + early_apply_env_caps() + assert os.environ.get("OMP_NUM_THREADS") == "3" + + +def test_resolved_runtime_dataclass_is_pickleable(): + # Import inside test body so we pick up the *currently cached* + # dynacell.evaluation.runtime module. test_lazy_init.py clears + # `dynacell.*` modules from sys.modules — a module-level import + # at collection time would bind to a stale class object. + import importlib + + runtime_mod = importlib.import_module("dynacell.evaluation.runtime") + r = runtime_mod.ResolvedRuntime( + fov_workers=2, + threads_per_worker=4, + executor="process", + cuda_empty_cache_every_n_timepoints=0, + gc_collect_every_n_fovs=0, + ) + restored = pickle.loads(pickle.dumps(r)) + assert restored == r + + +def test_region_timer_records_to_collector(): + reset_timings() + with region_timer("test_region", "A/1/0"): + pass + with region_timer("test_per_t", "A/1/0", t=5): + pass + rows = get_timings() + assert len(rows) == 2 + assert rows[0][0] == "A/1/0" + assert rows[0][1] is None + assert rows[0][2] == "test_region" + assert rows[0][3] >= 0.0 + assert rows[1][1] == 5 + + +def test_dump_timings_csv_writes_correct_schema(tmp_path: Path): + reset_timings() + with region_timer("test_region", "A/1/0"): + pass + with region_timer("test_per_t", "A/1/0", t=2): + pass + out_path = dump_timings_csv(tmp_path) + assert out_path is not None + with out_path.open() as f: + rows = list(csv.reader(f)) + assert rows[0] == ["pos_name", "t", "region", "seconds"] + assert rows[1][:3] == ["A/1/0", "", "test_region"] + assert rows[2][:3] == ["A/1/0", "2", "test_per_t"] + + +def test_dump_timings_csv_returns_none_when_empty(tmp_path: Path): + reset_timings() + assert dump_timings_csv(tmp_path) is None + + +def test_extend_timings_appends_batch(): + reset_timings() + extend_timings([("X/1/0", None, "foo", 0.1), ("X/1/0", 0, "bar", 0.2)]) + assert len(get_timings()) == 2 + + +def test_maybe_empty_cuda_cache_is_noop_when_every_n_zero(): + # No-op should not raise even without CUDA. + maybe_empty_cuda_cache(t=3, every_n=0) + + +def test_maybe_gc_collect_is_noop_when_every_n_zero(): + maybe_gc_collect(fov_idx=0, every_n=0) + + +def test_maybe_gc_collect_runs_at_interval(): + """fov_idx+1 % every_n == 0 triggers the collect.""" + # Just verify it doesn't raise for valid intervals; gc.collect side-effect + # is opaque but must not crash. + maybe_gc_collect(fov_idx=0, every_n=1) + maybe_gc_collect(fov_idx=4, every_n=5) + + +def test_gpu_lock_path_for_job_uses_tmp(): + path = gpu_lock_path_for_job() + assert path.startswith("/tmp/") or path.startswith("/var/") # platform-dependent + assert path.endswith(".lock") + + +def test_gpu_serialization_lock_is_noop_in_parent(): + """No worker initializer has run, so the lock should yield immediately.""" + with gpu_serialization_lock(): + pass # Must not block or raise. + + +def test_make_fov_executor_serial_yields_none(): + runtime = ResolvedRuntime( + fov_workers=1, + threads_per_worker=4, + executor="serial", + cuda_empty_cache_every_n_timepoints=0, + gc_collect_every_n_fovs=0, + ) + with make_fov_executor(runtime) as pool: + assert pool is None + + +def test_make_fov_executor_process_yields_real_executor(): + """Spawn smoke: ProcessPoolExecutor with worker initializer comes up cleanly. + + Submits ``os.getpid`` (stdlib, no test-tree imports needed in spawn workers + — tests/ is not on the python path so test-module-level helpers can't be + pickled into the child) and verifies the worker returns its own PID, + distinct from the parent's. + """ + import importlib + + runtime_mod = importlib.import_module("dynacell.evaluation.runtime") + runtime = runtime_mod.ResolvedRuntime( + fov_workers=2, + threads_per_worker=2, + executor="process", + cuda_empty_cache_every_n_timepoints=0, + gc_collect_every_n_fovs=0, + ) + parent_pid = os.getpid() + with runtime_mod.make_fov_executor(runtime) as pool: + assert pool is not None + assert isinstance(pool, ProcessPoolExecutor) + worker_pid = pool.submit(os.getpid).result(timeout=60) + assert worker_pid != parent_pid + assert isinstance(worker_pid, int) diff --git a/applications/dynacell/tools/run_eval_jointtrained_nucleus.slurm b/applications/dynacell/tools/run_eval_jointtrained_nucleus.slurm index b16b9f29a..ecbf7a7cf 100644 --- a/applications/dynacell/tools/run_eval_jointtrained_nucleus.slurm +++ b/applications/dynacell/tools/run_eval_jointtrained_nucleus.slurm @@ -26,6 +26,15 @@ set -euo pipefail +# Cap BLAS/OMP threads to the SLURM allocation before any C extension loads. +# applications/dynacell/CLAUDE.md "Eval runtime / parallelism" section. +# Set DYNACELL_THREADS_PER_WORKER to the cgroup-visible CPU count so +# scipy-openblas + torch intra-op don't oversubscribe. +export DYNACELL_THREADS_PER_WORKER="${SLURM_CPUS_PER_TASK}" +# To enable FOV-level parallelism in a future run with a multi-FOV eval, +# pass these overrides to the wrapped eval invocation (see the .sh script): +# runtime.executor=process runtime.fov_workers=4 runtime.threads_per_worker=4 + REPO=/hpc/mydata/alex.kalinin/VisCy EVAL_DIR="${REPO}/applications/dynacell/configs/evaluations" diff --git a/applications/dynacell/tools/run_eval_unext2_jointtrained_membrane.slurm b/applications/dynacell/tools/run_eval_unext2_jointtrained_membrane.slurm index 545914bd0..577be5060 100644 --- a/applications/dynacell/tools/run_eval_unext2_jointtrained_membrane.slurm +++ b/applications/dynacell/tools/run_eval_unext2_jointtrained_membrane.slurm @@ -26,6 +26,13 @@ set -euo pipefail +# Cap BLAS/OMP threads to the SLURM allocation before any C extension loads. +# applications/dynacell/CLAUDE.md "Eval runtime / parallelism" section. +export DYNACELL_THREADS_PER_WORKER="${SLURM_CPUS_PER_TASK}" +# To enable FOV-level parallelism in a future multi-FOV run, pass these +# overrides to the wrapped eval invocation (see the .sh script): +# runtime.executor=process runtime.fov_workers=4 runtime.threads_per_worker=4 + REPO=/hpc/mydata/alex.kalinin/VisCy EVAL_DIR="${REPO}/applications/dynacell/configs/evaluations" diff --git a/pyproject.toml b/pyproject.toml index 082e6ea14..6e13a0ad4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,8 +87,11 @@ lint.pydocstyle.convention = "numpy" [tool.pytest.ini_options] minversion = "9.0" testpaths = [ "packages/*/tests", "applications/*/tests" ] -addopts = [ "-ra", "-q", "--import-mode=importlib" ] +addopts = [ "-ra", "-q", "--import-mode=importlib", "-m", "not slow" ] pythonpath = [ "applications/dynacell/tools" ] +markers = [ + "slow: marks tests requiring GPU/HPC resources (deselect with -m 'not slow')", +] [tool.uv-dynamic-versioning] vcs = "git" From aa0e765baba4b0c09ae8741acca5b447ae3ab805 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 12:48:29 -0700 Subject: [PATCH 06/20] fix(dynacell): drop parent seg_model under executor=process to avoid duplicate residency MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Under executor=process, workers lazy-load their own seg_model copies (under the GPU lock) — but the parent had been holding a duplicate copy GPU-resident since the unconditional prepare_segmentation_model(config) call at the top of evaluate_predictions. On an 80 GB A100 with fov_workers=4, this meant (N+1) × model_weights resident instead of N. After Phase-2 runtime resolution (where the final executor is known), drop the parent's copy when executor=process. The pre-warm side effect (segmenter_model_zoo's validate_model populating the checkpoint dir to avoid N workers racing on a cold quilt download) is already done by then. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../dynacell/src/dynacell/evaluation/pipeline.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/applications/dynacell/src/dynacell/evaluation/pipeline.py b/applications/dynacell/src/dynacell/evaluation/pipeline.py index f32fffb2e..c76d3901d 100644 --- a/applications/dynacell/src/dynacell/evaluation/pipeline.py +++ b/applications/dynacell/src/dynacell/evaluation/pipeline.py @@ -814,6 +814,22 @@ def evaluate_predictions(config: DictConfig): freeze_threads_per_worker=runtime.threads_per_worker, ) + # Under executor=process, workers lazy-load their own seg_model + # copies (under the GPU lock). The parent's copy is held only for + # the checkpoint-cache-warm side effect (segmenter_model_zoo's + # validate_model has no file lock around its quilt download, so + # pre-warming in the parent prevents N workers from racing on a + # cold cache). After Phase 2 we know the final executor — if it's + # process, drop the parent copy so we don't keep two seg model + # copies resident on the GPU. + if runtime.executor == "process" and seg_model is not None: + del seg_model + seg_model = None + import torch + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + def _aggregate(result: FovResult) -> None: _aggregate_fov_result( result, From 2102c43ccfe7687a792c3cbff4917aefed1ab1e9 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 12:49:20 -0700 Subject: [PATCH 07/20] chore(dynacell): resolve config interpolations in-place before worker submit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Belt-and-suspenders safety net suggested in PR #426 review: call OmegaConf.resolve(config) in-place after Phase-1 runtime resolution and before any worker pickle.dumps the DictConfig across the spawn boundary. OmegaConf supports lazy interpolation resolution natively, so workers *would* resolve ${...} refs on first attribute access — but doing it once in the parent eliminates a tiny per-worker hot-path cost and guards against the (unlikely) case where an interpolation depends on parent-only state. ??? MISSING fields are untouched: verified OmegaConf 2.3 `_resolve` walks only nodes where `_is_interpolation()` is true, so bare ??? values pass through. Confirmed: 63 tests pass across pipeline / compose / parallel / runtime suites. Co-Authored-By: Claude Opus 4.7 (1M context) --- applications/dynacell/src/dynacell/evaluation/pipeline.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/applications/dynacell/src/dynacell/evaluation/pipeline.py b/applications/dynacell/src/dynacell/evaluation/pipeline.py index c76d3901d..c9c7a44a5 100644 --- a/applications/dynacell/src/dynacell/evaluation/pipeline.py +++ b/applications/dynacell/src/dynacell/evaluation/pipeline.py @@ -641,6 +641,11 @@ def evaluate_predictions(config: DictConfig): # frozen across phases for parent/worker BLAS-cap consistency. runtime = resolve_runtime(config) apply_thread_budget(runtime.threads_per_worker) + # Resolve ${...} interpolations in-place so spawn workers don't lazy-resolve + # refs in their child interpreters. ??? MISSING fields stay unresolved + # (OmegaConf 2.3 _resolve walks only interpolation nodes, not value nodes), + # so this is safe under feature-metrics-disabled runs. + OmegaConf.resolve(config) reset_timings() all_pixel_metrics = [] From 60334ac4de64a7541a04e2df30325b7db51365a6 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 12:50:32 -0700 Subject: [PATCH 08/20] chore(dynacell): silence inner per-T tqdm in spawn workers Adds a runtime.is_worker() helper that returns True iff _worker_initializer has set _GPU_LOCK_PATH. Wire it into _process_one_fov's inner-T tqdm via disable=is_worker(), so under fov_workers>1 the N concurrent per-T bars don't interleave on the parent's stderr. Outer per-FOV tqdm (in the parent) stays visible because the parent's _GPU_LOCK_PATH is never set. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../dynacell/src/dynacell/evaluation/pipeline.py | 6 +++++- .../dynacell/src/dynacell/evaluation/runtime.py | 10 ++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/applications/dynacell/src/dynacell/evaluation/pipeline.py b/applications/dynacell/src/dynacell/evaluation/pipeline.py index c9c7a44a5..2155c525e 100644 --- a/applications/dynacell/src/dynacell/evaluation/pipeline.py +++ b/applications/dynacell/src/dynacell/evaluation/pipeline.py @@ -226,12 +226,16 @@ def _process_one_fov( from dynacell.evaluation.runtime import ( get_timings, gpu_serialization_lock, + is_worker, maybe_empty_cuda_cache, region_timer, ) from dynacell.evaluation.segmentation import segment timings_start = len(get_timings()) + # Inner per-T tqdm is noise when N workers each emit it to the shared + # parent stderr — outer per-FOV tqdm in the parent stays visible either way. + suppress_inner_tqdm = is_worker() pred_channel_index = pos_pred.get_channel_index(io_config.pred_channel_name) gt_channel_index = pos_gt.get_channel_index(io_config.gt_channel_name) @@ -299,7 +303,7 @@ def _process_one_fov( dynaclr = _BackboneLists() celldino = _BackboneLists() - for t in tqdm(range(T), desc="Processing timepoints", leave=False): + for t in tqdm(range(T), desc="Processing timepoints", leave=False, disable=suppress_inner_tqdm): data_info = {"FOV": pos_name_pred, "Timepoint": t} with region_timer("pixel_metrics", pos_name_pred, t), gpu_serialization_lock(): diff --git a/applications/dynacell/src/dynacell/evaluation/runtime.py b/applications/dynacell/src/dynacell/evaluation/runtime.py index 6b76fe67a..919ae9549 100644 --- a/applications/dynacell/src/dynacell/evaluation/runtime.py +++ b/applications/dynacell/src/dynacell/evaluation/runtime.py @@ -399,6 +399,16 @@ def _worker_initializer(threads: int, gpu_lock_path: str | None) -> None: _GPU_LOCK_PATH = gpu_lock_path +def is_worker() -> bool: + """Return True if called from a spawn worker (post-``_worker_initializer``). + + Used to suppress inner per-T tqdm bars in workers — under + ``fov_workers>1`` the N concurrent inner bars interleave on the parent's + stderr unreadably. The outer per-FOV tqdm (in the parent) stays visible. + """ + return _GPU_LOCK_PATH is not None + + @contextmanager def gpu_serialization_lock() -> Iterator[None]: """Serialize GPU operations across workers via an ``fcntl.flock``. From 318d3f9ec1963285e96cd4a65fbe867d30dcf7d3 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 13:11:04 -0700 Subject: [PATCH 09/20] fix(dynacell): don't double-count timing rows under executor=serial MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bug: _process_one_fov writes region timings to the global _TIMINGS collector via region_timer, and returns the slice in FovResult.timings. _aggregate_fov_result then calls extend_timings(result.timings) unconditionally. In serial mode this re-appends each row to the SAME parent collector — every timed region appears twice in eval_timing.csv. In process mode the second append is correct (workers have separate per-process _TIMINGS, parent's is empty until aggregated). Fix: add extend_worker_timings kwarg to _aggregate_fov_result. Pass True in process mode, False in serial. Verified via 63-test suite green. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/dynacell/evaluation/pipeline.py | 19 ++++++++++++++++++- .../test_evaluation_pipeline_parallel.py | 1 + 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/applications/dynacell/src/dynacell/evaluation/pipeline.py b/applications/dynacell/src/dynacell/evaluation/pipeline.py index 2155c525e..38de4b558 100644 --- a/applications/dynacell/src/dynacell/evaluation/pipeline.py +++ b/applications/dynacell/src/dynacell/evaluation/pipeline.py @@ -447,15 +447,24 @@ def _aggregate_fov_result( gt_celldino_feats: list[np.ndarray], gt_celldino_fovs: list[np.ndarray], gt_celldino_ts: list[np.ndarray], + *, + extend_worker_timings: bool, ) -> None: """Apply one FOV's contributions to the parent-side run state. Writes the segmentation array to the HCS plate, extends the per-T row lists, and extends the 24 dataset-level accumulator lists. + + ``extend_worker_timings`` toggles whether to append ``result.timings`` + to the parent's global ``_TIMINGS`` collector. Set ``False`` in serial + mode (workers and parent share one collector, so the timings are + already there); set ``True`` in process mode (workers have separate + per-process collectors, so the parent must aggregate). """ from dynacell.evaluation.runtime import extend_timings, region_timer - extend_timings(result.timings) + if extend_worker_timings: + extend_timings(result.timings) with region_timer("seg_write", result.pos_name): seg_pos = segmentation_results.create_position(result.row, result.col, result.fov) @@ -839,6 +848,13 @@ def evaluate_predictions(config: DictConfig): if torch.cuda.is_available(): torch.cuda.empty_cache() + # Whether to fold result.timings into the parent's _TIMINGS: + # only when results come from spawn workers (separate per-process + # collector). In serial mode, _process_one_fov already wrote to + # the parent's collector via region_timer; re-extending here + # would double-count every row. + extend_worker_timings = runtime.executor == "process" + def _aggregate(result: FovResult) -> None: _aggregate_fov_result( result, @@ -870,6 +886,7 @@ def _aggregate(result: FovResult) -> None: gt_celldino_feats, gt_celldino_fovs, gt_celldino_ts, + extend_worker_timings=extend_worker_timings, ) if runtime.executor == "serial": diff --git a/applications/dynacell/tests/test_evaluation_pipeline_parallel.py b/applications/dynacell/tests/test_evaluation_pipeline_parallel.py index 702ff1f31..f9fa74219 100644 --- a/applications/dynacell/tests/test_evaluation_pipeline_parallel.py +++ b/applications/dynacell/tests/test_evaluation_pipeline_parallel.py @@ -211,6 +211,7 @@ def create_position(self, row, col, fov): gt_celldino_feats, gt_celldino_fovs, gt_celldino_ts, + extend_worker_timings=True, ) assert len(all_pix) == 2 From d3bce0ccbaebf097eb6ecc354098ff8a4dc91810 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 13:12:51 -0700 Subject: [PATCH 10/20] fix(dynacell): context-manage worker plate handles per FOV MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _worker_setup no longer caches pred/gt/seg plate handles in _WORKER_STATE. Instead, _worker_run_fov opens them via `with open_ome_zarr(...) as ...` scoped to one FOV — iohub file descriptors close between FOVs and the worker has nothing to clean up on ProcessPoolExecutor shutdown. Models + cache contexts stay cached in _WORKER_STATE (their load cost ~5-10 s per worker is too high to repeat per FOV; ProcessPoolExecutor process exit reclaims them). Position lookup by name uses a linear scan of plate.positions() (iohub has no name->position index). Cheap for the ≤ tens of positions per typical eval — sub-ms per scan. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/dynacell/evaluation/pipeline.py | 114 +++++++++++------- 1 file changed, 73 insertions(+), 41 deletions(-) diff --git a/applications/dynacell/src/dynacell/evaluation/pipeline.py b/applications/dynacell/src/dynacell/evaluation/pipeline.py index 38de4b558..559efc433 100644 --- a/applications/dynacell/src/dynacell/evaluation/pipeline.py +++ b/applications/dynacell/src/dynacell/evaluation/pipeline.py @@ -501,19 +501,25 @@ def _aggregate_fov_result( # Worker-side state for ``executor=process``. A spawn-context child Python -# interpreter calls ``_worker_run_fov`` once per submitted FOV; the first call -# triggers ``_worker_setup`` which lazy-loads seg model + extractors + plate -# handles under the GPU serialization lock and caches them here. Subsequent -# FOVs in the same worker reuse the cached state. +# interpreter calls ``_worker_run_fov`` once per submitted FOV; the first +# call triggers ``_worker_setup`` which lazy-loads seg model + extractors + +# cache contexts under the GPU serialization lock and caches them here. +# Subsequent FOVs in the same worker reuse the cached state. +# +# Plate handles are intentionally *not* cached here — they're context-managed +# per-FOV inside ``_worker_run_fov`` so iohub file descriptors close when the +# worker is between FOVs, and so the worker has nothing to clean up on +# ProcessPoolExecutor shutdown. _WORKER_STATE: dict[str, Any] = {} def _worker_setup(config: DictConfig) -> None: - """First-FOV-per-worker initialization: load models, open plates, init caches. + """First-FOV-per-worker initialization: load models + init cache contexts. - Wraps all GPU touch sites in ``gpu_serialization_lock()`` so N workers - don't initialize CUDA contexts + load weights concurrently. Cache contexts - and iohub plates are opened outside the lock (no GPU touch). + GPU touch sites (model + extractor loads) run under + ``gpu_serialization_lock()`` so N workers don't initialize CUDA + load + weights concurrently. Cache contexts (no GPU touch) are built outside + the lock. Plate handles are opened lazily per-FOV in ``_worker_run_fov``. """ if _WORKER_STATE.get("initialized"): return @@ -571,12 +577,6 @@ def _worker_setup(config: DictConfig) -> None: celldino_weights_path=celldino_weights_path, ) - pred_plate = open_ome_zarr(Path(config.io.pred_path), mode="r") - gt_plate = open_ome_zarr(Path(config.io.gt_path), mode="r") - seg_plate = None - if config.io.cell_segmentation_path is not None: - seg_plate = open_ome_zarr(Path(config.io.cell_segmentation_path), mode="r") - _WORKER_STATE.update( { "initialized": True, @@ -586,46 +586,78 @@ def _worker_setup(config: DictConfig) -> None: "celldino": celldino_feature_extractor, "cache_ctx": cache_ctx, "pred_cache_ctx": pred_cache_ctx, - "pred_plate": pred_plate, - "gt_plate": gt_plate, - "seg_plate": seg_plate, - "pred_positions": dict(pred_plate.positions()), - "gt_positions": dict(gt_plate.positions()), - "seg_positions": dict(seg_plate.positions()) if seg_plate is not None else None, } ) +def _find_position(plate, pos_name: str): + """Return the iohub Position with name ``pos_name`` from ``plate``. + + Scans ``plate.positions()`` since iohub doesn't index by name. Cheap + for the ≤ tens of positions per eval. + """ + for name, pos in plate.positions(): + if name == pos_name: + return pos + raise KeyError(f"position {pos_name!r} not found in plate") + + def _worker_run_fov(config: DictConfig, pos_name: str, cuda_empty_every_n: int) -> FovResult: """Worker entry point: process one FOV by name and return FovResult. - Submitted via ``pool.submit`` in ``executor=process`` mode. Returns a - pickle-friendly ``FovResult`` (no plate handles or torch modules - cross the process boundary). + Submitted via ``pool.submit`` in ``executor=process`` mode. Plates + are opened with a context manager scoped to this single call — iohub + file descriptors close before the worker accepts its next FOV. Models + + cache contexts stay cached in ``_WORKER_STATE`` across FOVs. """ from dynacell.evaluation.pipeline_cache import flush_manifest _worker_setup(config) state = _WORKER_STATE - pos_pred = state["pred_positions"][pos_name] - pos_gt = state["gt_positions"][pos_name] - pos_seg = state["seg_positions"][pos_name] if state["seg_positions"] is not None else None - result = _process_one_fov( - config, - cuda_empty_every_n, - pos_name, - pos_pred, - pos_gt, - pos_seg, - config.io, - state["cache_ctx"], - state["pred_cache_ctx"], - state["seg_model"], - state["dinov3"], - state["dynaclr"], - state["celldino"], - ) + seg_path = config.io.cell_segmentation_path + with ( + open_ome_zarr(Path(config.io.pred_path), mode="r") as pred_plate, + open_ome_zarr(Path(config.io.gt_path), mode="r") as gt_plate, + ): + pos_pred = _find_position(pred_plate, pos_name) + pos_gt = _find_position(gt_plate, pos_name) + + if seg_path is not None: + with open_ome_zarr(Path(seg_path), mode="r") as seg_plate: + pos_seg = _find_position(seg_plate, pos_name) + result = _process_one_fov( + config, + cuda_empty_every_n, + pos_name, + pos_pred, + pos_gt, + pos_seg, + config.io, + state["cache_ctx"], + state["pred_cache_ctx"], + state["seg_model"], + state["dinov3"], + state["dynaclr"], + state["celldino"], + ) + else: + result = _process_one_fov( + config, + cuda_empty_every_n, + pos_name, + pos_pred, + pos_gt, + None, + config.io, + state["cache_ctx"], + state["pred_cache_ctx"], + state["seg_model"], + state["dinov3"], + state["dynaclr"], + state["celldino"], + ) + # Worker-side manifest flush so interrupted runs preserve progress even # when the parent dies. The global manifest lock at # _pos_write_lock(ctx, "manifest", "global") serializes flushes across From a1fd8df01d91afeafecf19e072128bd0852bf58a Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 13:13:56 -0700 Subject: [PATCH 11/20] fix(dynacell): skip prepare_segmentation_model under require_complete_cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When io.require_complete_cache=true, the caller asserts that all needed masks live in the cache — so loading SuperModel (which hits a quilt download for nucleus/membrane on a cold node) is wasted work. Return None unconditionally in that mode. If a cache miss happens despite the assertion, the cache layer raises StaleCacheError before any segment() call, so the None seg_model never gets dereferenced. Enables the cache-only/no-model path that C5's integration test design relies on (run evaluate_predictions end-to-end against prebuilt mask caches without instantiating cellpose/SuperModel). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../dynacell/src/dynacell/evaluation/segmentation.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/applications/dynacell/src/dynacell/evaluation/segmentation.py b/applications/dynacell/src/dynacell/evaluation/segmentation.py index 22d3b1085..55205921f 100644 --- a/applications/dynacell/src/dynacell/evaluation/segmentation.py +++ b/applications/dynacell/src/dynacell/evaluation/segmentation.py @@ -91,7 +91,14 @@ def prepare_segmentation_model(config): Returns ``None`` for organelles that use classical (non-DL) workflows. Respects ``config.use_gpu`` when deciding whether to move models to GPU. + + Returns ``None`` unconditionally when ``io.require_complete_cache=true``: + in that mode the caller asserts that all needed masks live in the cache, + so loading SuperModel (which hits a quilt download for nucleus/membrane + on a cold node) would be wasted work. If a mask cache-miss happens, the + cache layer raises ``StaleCacheError`` before any ``segment()`` call. """ + require_complete = bool(getattr(config.io, "require_complete_cache", False)) if config.target_name not in [ "nucleus", "membrane", @@ -101,6 +108,8 @@ def prepare_segmentation_model(config): "mitochondria", ]: raise ValueError(f"Invalid target_name in config: {config.target_name!r}") + if require_complete: + return None if config.target_name in ["nucleus", "membrane"]: _require_segmenter_model_zoo() if config.target_name == "nucleus": From 0850633f1e1a07af3f5b9e06f6b36e99f4190733 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 13:23:20 -0700 Subject: [PATCH 12/20] test(dynacell): add CPU integration test for serial-vs-process parity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drives evaluate_predictions end-to-end on a tiny iohub fixture (3 positions × T=2 × D=4 × H=W=32) with prebuilt mask caches. Runs the pipeline twice — once with runtime.executor=serial, once with runtime.executor=process runtime.fov_workers=2 — and asserts: 1. pixel_metrics + mask_metrics rows are byte-equal after sort by FOV. 2. segmentation_results.zarr per-position content is byte-equal across both channels. Cache-only design (matches plan §C5): target_name=er + require_complete_cache=true + compute_feature_metrics=false + compute_microssim=false + pixel_metrics.fsc=null + pixel_metrics.spectral_pcc=null. The full pipeline runs without instantiating cellpose, SuperModel, DinoV3, DynaCLR, CellDino, or cubic. What this validates: spawn boot, DictConfig pickling, FovResult.timings + arrays round-trip through pickle, _aggregate_fov_result extends parent accumulators correctly, parent-side segmentation_results.zarr writes don't race, sort-by-pos_name aggregation, and the new extend_worker_timings semantics doesn't double-count timing rows. _live_pipeline_module() clears dynacell.evaluation.* from sys.modules before importing — necessary because test_evaluation_pipeline.py stubs metrics/utils via monkeypatch.setitem(sys.modules, ...), and without the clear my pipeline's compute_pixel_metrics resolves to the empty-returning stub. CI: runs ~16 s on the existing test-dynacell-configs job. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../test_evaluation_pipeline_parallel_cpu.py | 318 ++++++++++++++++++ 1 file changed, 318 insertions(+) create mode 100644 applications/dynacell/tests/test_evaluation_pipeline_parallel_cpu.py diff --git a/applications/dynacell/tests/test_evaluation_pipeline_parallel_cpu.py b/applications/dynacell/tests/test_evaluation_pipeline_parallel_cpu.py new file mode 100644 index 000000000..692ffa5ba --- /dev/null +++ b/applications/dynacell/tests/test_evaluation_pipeline_parallel_cpu.py @@ -0,0 +1,318 @@ +"""End-to-end CPU integration test for serial-vs-process executor parity. + +Drives ``evaluate_predictions`` twice on the same tiny iohub fixture + +prebuilt mask caches: once with ``runtime.executor=serial``, once with +``runtime.executor=process runtime.fov_workers=2``. Asserts byte-equal +CSV outputs (after sort by FOV) and per-position equality of the +``segmentation_results.zarr`` HCS plate. + +Cache-only design (matches plan ``.claude/plans/eval-parallelism.md`` §C5): + +- ``target_name=er`` → ``prepare_segmentation_model`` returns None. +- ``io.require_complete_cache=true`` → mask cache MUST satisfy every + request; we pre-populate the cache before the run. +- ``compute_feature_metrics=false``, ``compute_microssim=false`` → no + extractors load, no microssim. +- ``pixel_metrics.fsc=null`` and ``pixel_metrics.spectral_pcc=null`` → + ``compute_pixel_metrics`` only runs the torch metric block (PCC, + SSIM, NRMSE, PSNR), no cubic FFT path. + +Result: the full pipeline runs without instantiating cellpose, +SuperModel, DinoV3, DynaCLR, CellDino, or cubic — exercises spawn +boot, DictConfig pickling, ``FovResult`` transport, parent-side HCS +plate write, and the dataset-level aggregation contract. +""" + +from __future__ import annotations + +import importlib +import sys +from datetime import UTC, datetime +from pathlib import Path + +import numpy as np +import pytest # noqa: F401 -- imported for future @pytest.mark uses +from iohub.ngff import open_ome_zarr +from omegaconf import OmegaConf + + +def _live_pipeline_module(): + """Force-import a fresh ``dynacell.evaluation.pipeline`` module. + + Other tests (``test_evaluation_pipeline.py``) stub + ``dynacell.evaluation.metrics`` etc. via ``monkeypatch.setitem(sys.modules, + ...)``. After teardown, ``sys.modules`` may still hold the stubbed modules + AND the pipeline module's bound ``compute_pixel_metrics`` reference. Clear + every ``dynacell.evaluation.*`` cache entry before re-import so we get + fresh modules with real implementations. + """ + for name in list(sys.modules): + if name.startswith("dynacell.evaluation"): + sys.modules.pop(name, None) + return importlib.import_module("dynacell.evaluation.pipeline") + + +# Fixture dimensions — kept tiny so the test runs in CI (single-CPU) under ~30s. +N_POSITIONS = 3 +T = 2 +D = 4 +H = 32 +W = 32 + + +def _make_hcs_plate(path: Path, channel_name: str, seed: int) -> None: + """Create a tiny HCS plate with ``N_POSITIONS`` positions of shape ``(T, 1, D, H, W)``. + + Positions are named ``A/1/{i}`` for ``i in range(N_POSITIONS)``. Pixel + values are deterministic from ``seed`` so pred and GT differ predictably. + """ + rng = np.random.default_rng(seed) + with open_ome_zarr(path, mode="w", layout="hcs", channel_names=[channel_name], version="0.5") as plate: + for i in range(N_POSITIONS): + pos = plate.create_position("A", "1", str(i)) + data = rng.uniform(0.0, 1.0, size=(T, 1, D, H, W)).astype(np.float32) + pos.create_image("0", data) + + +def _make_mask_cache( + cache_dir: Path, + plate_path: Path, + channel_name: str, + side: str, # "gt" or "pred" + target_name: str = "er", +) -> None: + """Prebuild a mask cache (zarr + manifest) for ``N_POSITIONS`` positions. + + The mask values themselves are deterministic per (side, position, t): + GT masks are TRUE in the upper half of the (H, W) plane; pred masks + are TRUE in the upper quadrant. This guarantees the DICE / IoU + metrics are deterministic and non-trivially different between sides. + """ + from dynacell.evaluation.cache import ( + CACHE_SCHEMA_VERSION, + cache_paths, + save_manifest, + write_mask, + ) + + paths = cache_paths(cache_dir) + mask_channel = "target_seg" if side == "gt" else "prediction_seg" + + pos_names = [f"A/1/{i}" for i in range(N_POSITIONS)] + for pos_name in pos_names: + masks = np.zeros((T, D, H, W), dtype=bool) + if side == "gt": + masks[:, :, : H // 2, :] = True # upper half + else: + masks[:, :, : H // 2, : W // 2] = True # upper quadrant + write_mask(paths, target_name, pos_name, masks, channel_name=mask_channel) + + # Manifest entries must satisfy _validate_artifact_params: the + # masks_section.get(target_name) must equal {"target_name": ..., **source_tag}. + source_tag: dict[str, str] = {} if side == "gt" else {"source": "prediction"} + now = datetime.now(UTC).isoformat() + manifest = { + "cache_schema_version": CACHE_SCHEMA_VERSION, + "gt" if side == "gt" else "pred": { + "plate_path": str(plate_path), + "channel_name": channel_name, + }, + "artifacts": { + "organelle_masks": { + target_name: { + "target_name": target_name, + "path": f"organelle_masks/{target_name}.zarr", + "built_at": now, + "positions": pos_names, + **source_tag, + } + }, + }, + } + save_manifest(paths, manifest) + + +def _build_eval_config( + pred_path: Path, + gt_path: Path, + gt_cache_dir: Path, + pred_cache_dir: Path, + save_dir: Path, + *, + executor: str, + fov_workers: int, +): + """Compose a minimal eval config for the cache-only path.""" + return OmegaConf.create( + { + "target_name": "er", + "io": { + "pred_path": str(pred_path), + "gt_path": str(gt_path), + "pred_channel_name": "prediction", + "gt_channel_name": "target", + "cell_segmentation_path": None, + "gt_cache_dir": str(gt_cache_dir), + "pred_cache_dir": str(pred_cache_dir), + "require_complete_cache": True, + }, + "pixel_metrics": { + "spacing": [1.0, 1.0, 1.0], + "fsc": None, + "spectral_pcc": None, + }, + "feature_metrics": { + "patch_size": 64, + "deep_feature_batch_threshold": 256, + }, + "use_gpu": False, + "compute_microssim": False, + "compute_feature_metrics": False, + "limit_positions": None, + "force_recompute": { + "all": False, + "gt_masks": False, + "gt_cp": False, + "gt_dinov3": False, + "gt_dynaclr": False, + "gt_celldino": False, + "pred_masks": False, + "pred_cp": False, + "pred_dinov3": False, + "pred_dynaclr": False, + "pred_celldino": False, + "final_metrics": False, + }, + "save": { + "save_dir": str(save_dir), + "pixel_csv_filename": "pixel_metrics.csv", + "pixel_metrics_filename": "pixel_metrics.npy", + "mask_csv_filename": "mask_metrics.csv", + "mask_metrics_filename": "mask_metrics.npy", + "feature_csv_filename": "feature_metrics.csv", + "feature_metrics_filename": "feature_metrics.npy", + }, + "runtime": { + "fov_workers": fov_workers, + "threads_per_worker": "auto", + "executor": executor, + "cuda_empty_cache_every_n_timepoints": 0, + "gc_collect_every_n_fovs": 0, + }, + } + ) + + +def _build_fixture(root: Path): + """Build pred + GT plates and matching mask caches under ``root``.""" + pred_path = root / "pred.zarr" + gt_path = root / "gt.zarr" + gt_cache_dir = root / "gt_cache" + pred_cache_dir = root / "pred_cache" + _make_hcs_plate(pred_path, "prediction", seed=42) + _make_hcs_plate(gt_path, "target", seed=7) + _make_mask_cache(gt_cache_dir, gt_path, "target", side="gt") + _make_mask_cache(pred_cache_dir, pred_path, "prediction", side="pred") + return pred_path, gt_path, gt_cache_dir, pred_cache_dir + + +def _read_position_arrays(plate_path: Path) -> dict[str, np.ndarray]: + """Return ``{pos_name: data}`` for every position in an HCS plate.""" + out: dict[str, np.ndarray] = {} + with open_ome_zarr(plate_path, mode="r") as plate: + for name, pos in plate.positions(): + out[name] = np.asarray(pos.data[:]) + return out + + +def test_serial_vs_process_evaluate_predictions_parity(tmp_path: Path): + """End-to-end: serial-mode and process-mode produce byte-equal outputs. + + Exercises the spawn boundary (pickling, worker startup, _process_one_fov + in workers, FovResult transport, parent-side aggregation + HCS plate + write) without invoking cellpose / extractors / cubic. + """ + fixture_root = tmp_path / "fixture" + fixture_root.mkdir() + pred_path, gt_path, gt_cache_dir, pred_cache_dir = _build_fixture(fixture_root) + + save_dir_serial = tmp_path / "serial" + save_dir_process = tmp_path / "process" + + # Re-import inside the test body — test_lazy_init.py pops dynacell.* + # modules from sys.modules, which would invalidate a module-level import. + pipeline = _live_pipeline_module() + evaluate_predictions = pipeline.evaluate_predictions + + cfg_serial = _build_eval_config( + pred_path, + gt_path, + gt_cache_dir, + pred_cache_dir, + save_dir_serial, + executor="serial", + fov_workers=1, + ) + cfg_process = _build_eval_config( + pred_path, + gt_path, + gt_cache_dir, + pred_cache_dir, + save_dir_process, + executor="process", + fov_workers=2, + ) + + save_dir_serial.mkdir(parents=True) + save_dir_process.mkdir(parents=True) + + pixel_serial, mask_serial, _ = evaluate_predictions(cfg_serial) + pixel_process, mask_process, _ = evaluate_predictions(cfg_process) + + # Pixel + mask rows: same set after sort. We sort by (FOV, Timepoint) + # because the process branch already sorts by pos_name and within-FOV + # iteration is deterministic, but we sort defensively. + def _sort_key(row: dict): + return (row["FOV"], row["Timepoint"]) + + assert sorted(pixel_serial, key=_sort_key) == sorted(pixel_process, key=_sort_key) + assert sorted(mask_serial, key=_sort_key) == sorted(mask_process, key=_sort_key) + # Per-position seg plate equality (both channels). + arrs_serial = _read_position_arrays(save_dir_serial / "segmentation_results.zarr") + arrs_process = _read_position_arrays(save_dir_process / "segmentation_results.zarr") + assert set(arrs_serial.keys()) == set(arrs_process.keys()) + for pos_name, arr_a in arrs_serial.items(): + arr_b = arrs_process[pos_name] + assert arr_a.shape == arr_b.shape, f"{pos_name} shape mismatch" + np.testing.assert_array_equal(arr_a, arr_b) + + +def test_serial_path_runs_without_segmenter_loaded(tmp_path: Path): + """``require_complete_cache=true`` should skip ``prepare_segmentation_model``. + + Confirms the cache-only fast-path: evaluate_predictions runs end-to-end + without ever instantiating SuperModel (which would trigger a quilt + download on a cold node). + """ + fixture_root = tmp_path / "fixture" + fixture_root.mkdir() + pred_path, gt_path, gt_cache_dir, pred_cache_dir = _build_fixture(fixture_root) + save_dir = tmp_path / "out" + save_dir.mkdir() + + cfg = _build_eval_config( + pred_path, + gt_path, + gt_cache_dir, + pred_cache_dir, + save_dir, + executor="serial", + fov_workers=1, + ) + pipeline = _live_pipeline_module() + pixel_rows, mask_rows, _ = pipeline.evaluate_predictions(cfg) + assert len(pixel_rows) == N_POSITIONS * T + assert len(mask_rows) == N_POSITIONS * T + # Confirms the no-segmenter contract: prepare_segmentation_model + # returned None and evaluate_predictions still produced valid outputs + # for every (FOV, t) pair. From 30fe250a03de0022f1a73b3a0e8e0b8b4a0c5e1d Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 13:23:44 -0700 Subject: [PATCH 13/20] ci(dynacell): include CPU integration test in dynacell-configs CI job MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds tests/test_evaluation_pipeline_parallel_cpu.py to the dynacell pytest invocation. The test runs end-to-end evaluate_predictions in both serial and spawn-process modes against a tiny iohub fixture + prebuilt mask cache — no eval extras (cellpose, transformers, cubic) needed because target_name=er + io.require_complete_cache=true short-circuit segmenter and feature-extractor loads. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/test.yml | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 0aca6ade6..a6ccba1ae 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -138,16 +138,21 @@ jobs: working-directory: applications/dynacell - name: Run benchmark-schema + submit-tool + eval-runtime tests - # tests/test_runtime.py and tests/test_evaluation_pipeline_parallel.py + # tests/test_runtime.py + tests/test_evaluation_pipeline_parallel.py # cover the dynacell.evaluation.runtime module + FovResult pickle - # contract that gate the FOV-level parallelism wiring. They don't - # need the eval extras (no iohub fixture or model loads). + # contract. tests/test_evaluation_pipeline_parallel_cpu.py drives + # evaluate_predictions end-to-end (serial vs spawn-process) on a + # tiny iohub fixture + prebuilt mask cache — no eval extras (no + # cellpose, transformers, cubic) needed since target_name=er + + # require_complete_cache=true short-circuit segmenter and feature + # extractor loads. run: | uv run --frozen pytest \ tests/test_benchmark_config_composition.py \ tests/test_submit_benchmark_job.py \ tests/test_runtime.py \ tests/test_evaluation_pipeline_parallel.py \ + tests/test_evaluation_pipeline_parallel_cpu.py \ -v working-directory: applications/dynacell From dd599e9cbd850bb0a0f2def0f9f175d3bbe8dcb3 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 13:28:32 -0700 Subject: [PATCH 14/20] fix(dynacell): harden runtime resolution + gate GPU lock on use_gpu MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four Copilot-review findings on the runtime layer: 1. apply_thread_budget threadpool_limits re-arm: the old handle becomes unreachable and is GC'd. The previous comment said "stays alive" which was misleading. Document the actual behavior: we WANT the new limit active, dropping the old handle is intentional. Also remove the dead if-branch since both arms did the same assignment. 2. resolve_runtime accepts literal fov_workers=0 / threads_per_worker=0 which produces divide-by-zero in the "auto" math or nonsensical thread counts. Add explicit >= 1 validation with a clear error. 3. _cpu_count() prefers os.sched_getaffinity(0) over os.cpu_count() so SLURM cgroup limits (cpuset etc.) are honored — otherwise auto- resolved fov_workers/threads_per_worker can oversize past the cgroup-visible CPU budget. Falls back to os.cpu_count() on non-Linux. 4. gpu_serialization_lock now takes a gate: bool = True kwarg. Callers in pipeline.py pass gate=use_gpu so under config.use_gpu=false the lock is a no-op even in process mode. Previously the lock fired unconditionally, defeating cross-worker parallelism for CPU-only eval runs. The lock still gates on _GPU_LOCK_PATH for the parent/serial no-op (unchanged). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/dynacell/evaluation/pipeline.py | 28 +++++---- .../src/dynacell/evaluation/runtime.py | 57 ++++++++++++++----- 2 files changed, 59 insertions(+), 26 deletions(-) diff --git a/applications/dynacell/src/dynacell/evaluation/pipeline.py b/applications/dynacell/src/dynacell/evaluation/pipeline.py index 559efc433..ce71dd2e3 100644 --- a/applications/dynacell/src/dynacell/evaluation/pipeline.py +++ b/applications/dynacell/src/dynacell/evaluation/pipeline.py @@ -236,6 +236,11 @@ def _process_one_fov( # Inner per-T tqdm is noise when N workers each emit it to the shared # parent stderr — outer per-FOV tqdm in the parent stays visible either way. suppress_inner_tqdm = is_worker() + # GPU serialization lock is a no-op when use_gpu=false: under that + # setting compute_pixel_metrics, cellpose, and feature extractors all + # run CPU-only, so cross-worker fcntl serialization would just add + # latency for nothing. + use_gpu = bool(getattr(config, "use_gpu", True)) pred_channel_index = pos_pred.get_channel_index(io_config.pred_channel_name) gt_channel_index = pos_gt.get_channel_index(io_config.gt_channel_name) @@ -246,10 +251,10 @@ def _process_one_fov( T = predict.shape[0] - with region_timer("mask_gt", pos_name_pred), gpu_serialization_lock(): + with region_timer("mask_gt", pos_name_pred), gpu_serialization_lock(gate=use_gpu): gt_mask_stack = fov_masks(cache_ctx, pos_name_pred, target, seg_model) if pred_cache_ctx.enabled: - with region_timer("mask_pred", pos_name_pred), gpu_serialization_lock(): + with region_timer("mask_pred", pos_name_pred), gpu_serialization_lock(gate=use_gpu): pred_mask_stack = fov_masks(pred_cache_ctx, pos_name_pred, predict, seg_model) else: pred_mask_stack = None @@ -260,18 +265,18 @@ def _process_one_fov( gt_celldino_per_t = None pred_per_t = None if config.compute_feature_metrics: - with region_timer("cp_gt", pos_name_pred), gpu_serialization_lock(): + with region_timer("cp_gt", pos_name_pred), gpu_serialization_lock(gate=use_gpu): gt_cp_per_t = fov_cp_features(cache_ctx, pos_name_pred, target, cell_segmentation) - with region_timer("deep_gt_dinov3", pos_name_pred), gpu_serialization_lock(): + with region_timer("deep_gt_dinov3", pos_name_pred), gpu_serialization_lock(gate=use_gpu): gt_dinov3_per_t = fov_deep_features( cache_ctx, pos_name_pred, target, cell_segmentation, dinov3_feature_extractor, "dinov3" ) - with region_timer("deep_gt_dynaclr", pos_name_pred), gpu_serialization_lock(): + with region_timer("deep_gt_dynaclr", pos_name_pred), gpu_serialization_lock(gate=use_gpu): gt_dynaclr_per_t = fov_deep_features( cache_ctx, pos_name_pred, target, cell_segmentation, dynaclr_feature_extractor, "dynaclr" ) if celldino_feature_extractor is not None: - with region_timer("deep_gt_celldino", pos_name_pred), gpu_serialization_lock(): + with region_timer("deep_gt_celldino", pos_name_pred), gpu_serialization_lock(gate=use_gpu): gt_celldino_per_t = fov_deep_features( cache_ctx, pos_name_pred, @@ -280,7 +285,7 @@ def _process_one_fov( celldino_feature_extractor, "celldino", ) - with region_timer("features_pred_per_t", pos_name_pred), gpu_serialization_lock(): + with region_timer("features_pred_per_t", pos_name_pred), gpu_serialization_lock(gate=use_gpu): pred_per_t = _fov_pred_features_per_t( pred_cache_ctx, pos_name_pred, @@ -306,7 +311,7 @@ def _process_one_fov( for t in tqdm(range(T), desc="Processing timepoints", leave=False, disable=suppress_inner_tqdm): data_info = {"FOV": pos_name_pred, "Timepoint": t} - with region_timer("pixel_metrics", pos_name_pred, t), gpu_serialization_lock(): + with region_timer("pixel_metrics", pos_name_pred, t), gpu_serialization_lock(gate=use_gpu): pixel_metrics = compute_pixel_metrics( predict[t], target[t], @@ -324,7 +329,7 @@ def _process_one_fov( if pred_mask_stack is not None: segmented_predict = pred_mask_stack[t] else: - with gpu_serialization_lock(): + with gpu_serialization_lock(gate=use_gpu): segmented_predict = np.asarray(segment(predict[t], config.target_name, seg_model=seg_model)).astype( bool ) @@ -517,7 +522,7 @@ def _worker_setup(config: DictConfig) -> None: """First-FOV-per-worker initialization: load models + init cache contexts. GPU touch sites (model + extractor loads) run under - ``gpu_serialization_lock()`` so N workers don't initialize CUDA + load + ``gpu_serialization_lock(gate=use_gpu)`` so under ``use_gpu=true`` N workers don't initialize CUDA + load weights concurrently. Cache contexts (no GPU touch) are built outside the lock. Plate handles are opened lazily per-FOV in ``_worker_run_fov``. """ @@ -532,7 +537,8 @@ def _worker_setup(config: DictConfig) -> None: DynaCLRFeatureExtractor, ) - with gpu_serialization_lock(): + use_gpu = bool(getattr(config, "use_gpu", True)) + with gpu_serialization_lock(gate=use_gpu): seg_model = prepare_segmentation_model(config) dinov3_feature_extractor = None dynaclr_feature_extractor = None diff --git a/applications/dynacell/src/dynacell/evaluation/runtime.py b/applications/dynacell/src/dynacell/evaluation/runtime.py index 919ae9549..bfe74532d 100644 --- a/applications/dynacell/src/dynacell/evaluation/runtime.py +++ b/applications/dynacell/src/dynacell/evaluation/runtime.py @@ -56,6 +56,21 @@ _FORCE_PER_T_HYGIENE_ENV: str = "DYNACELL_FORCE_PER_T_HYGIENE" +def _cpu_count() -> int: + """Return the number of CPUs visible to this process. + + Prefers ``os.sched_getaffinity(0)`` when available (Linux): it respects + SLURM cgroups + cpuset limits, so we don't oversize ``fov_workers`` or + ``threads_per_worker`` past the cgroup-visible budget. Falls back to + ``os.cpu_count()`` (which reports the host's logical CPU count, ignoring + cgroups) on platforms without ``sched_getaffinity``. + """ + try: + return max(1, len(os.sched_getaffinity(0))) + except AttributeError: + return max(1, os.cpu_count() or 1) + + @dataclass(frozen=True) class ResolvedRuntime: """Materialized runtime config with all ``"auto"`` values resolved. @@ -147,14 +162,13 @@ def apply_thread_budget(threads: int) -> None: logger.info("torch.set_num_interop_threads: %s; continuing with current value", e) # 4. Activate threadpool_limits for the rest of the process. - if _ACTIVE_THREADPOOL_HANDLE is None: - _ACTIVE_THREADPOOL_HANDLE = threadpool_limits(limits=threads, user_api="blas") - else: - # Re-arming: enter a fresh context with the new limit. The old handle - # stays alive (its __exit__ is never called); threadpoolctl honors - # the innermost active limit. Acceptable for re-call from a worker - # initializer that inherited a parent's handle. - _ACTIVE_THREADPOOL_HANDLE = threadpool_limits(limits=threads, user_api="blas") + # On re-arm, the old handle becomes unreachable and is GC'd. We don't + # call its __exit__ because that would *restore* the prior thread cap + # — but we want the new cap to be the active one. threadpoolctl's + # native thread-limit state is set when the new limiter is entered + # (via __enter__ inside threadpool_limits.__init__), so the new limit + # is active immediately; dropping the old handle is intentional. + _ACTIVE_THREADPOOL_HANDLE = threadpool_limits(limits=threads, user_api="blas") def _get_int(config: DictConfig, key: str, default: int) -> int: @@ -212,7 +226,7 @@ def resolve_runtime( # matching today's behavior. return ResolvedRuntime( fov_workers=1, - threads_per_worker=max(1, os.cpu_count() or 1), + threads_per_worker=_cpu_count(), executor="serial", cuda_empty_cache_every_n_timepoints=0, gc_collect_every_n_fovs=0, @@ -222,13 +236,15 @@ def resolve_runtime( if executor not in ("serial", "process"): raise ValueError(f"runtime.executor must be 'serial' or 'process', got {executor!r}") - cpu_count = max(1, os.cpu_count() or 1) + cpu_count = _cpu_count() raw_fov_workers = OmegaConf.select(runtime, "fov_workers", default=1) raw_threads = OmegaConf.select(runtime, "threads_per_worker", default="auto") # Resolve fov_workers. if isinstance(raw_fov_workers, int): resolved_fov_workers = int(raw_fov_workers) + if resolved_fov_workers < 1: + raise ValueError(f"runtime.fov_workers must be >= 1, got {resolved_fov_workers}") if resolved_fov_workers > 1 and executor == "serial": raise ValueError( f"runtime.fov_workers={resolved_fov_workers} requires runtime.executor='process' (got 'serial')" @@ -255,6 +271,8 @@ def resolve_runtime( resolved_threads = int(freeze_threads_per_worker) elif isinstance(raw_threads, int): resolved_threads = int(raw_threads) + if resolved_threads < 1: + raise ValueError(f"runtime.threads_per_worker must be >= 1, got {resolved_threads}") elif raw_threads == "auto": resolved_threads = max(1, cpu_count // resolved_fov_workers) else: @@ -410,14 +428,23 @@ def is_worker() -> bool: @contextmanager -def gpu_serialization_lock() -> Iterator[None]: +def gpu_serialization_lock(gate: bool = True) -> Iterator[None]: """Serialize GPU operations across workers via an ``fcntl.flock``. - No-op when ``_GPU_LOCK_PATH`` is unset (parent / serial mode). - Used to wrap model loads and per-FOV GPU operations under - ``executor=process``. + Parameters + ---------- + gate : bool + When False (e.g., the protected region runs CPU-only under + ``use_gpu=false``), this is a no-op even in process mode. + Callers pass ``gate=config.use_gpu`` for regions whose GPU + usage tracks the ``use_gpu`` flag, so CPU-only runs don't + serialize across workers needlessly. + + Other no-op conditions: ``_GPU_LOCK_PATH`` unset (parent / serial + mode) or no CUDA available in the worker. Used to wrap model loads + and per-FOV GPU operations under ``executor=process``. """ - if _GPU_LOCK_PATH is None: + if not gate or _GPU_LOCK_PATH is None: yield return import fcntl From c4c1afc4d5ff8565eb5e114c96542bcb8bd04c41 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 13:31:19 -0700 Subject: [PATCH 15/20] perf(dynacell): stream aggregation under executor=process to bound peak memory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously buffered every FovResult in results_by_pos until ALL N worker futures completed, then aggregated in pos_names_in_order. Peak parent memory: N × seg_array (~N × 440 MB for T=10 D=21 H=W=1024 evals) all live simultaneously. Switch to streaming aggregation: drain the buffer opportunistically in pos_names_in_order as futures complete. Only out-of-order FovResults stay buffered; once the next-expected position lands we aggregate (plate write + accumulator extend) and drop the ref. Peak buffer is bounded by the number of out-of-order FOVs in flight, not by N. Deterministic per-FOV CSV / embedding NPZ row order is preserved because aggregation still happens in pos_names_in_order — early completions just wait their turn in the buffer. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/dynacell/evaluation/pipeline.py | 36 +++++++++++++------ 1 file changed, 25 insertions(+), 11 deletions(-) diff --git a/applications/dynacell/src/dynacell/evaluation/pipeline.py b/applications/dynacell/src/dynacell/evaluation/pipeline.py index ce71dd2e3..904446186 100644 --- a/applications/dynacell/src/dynacell/evaluation/pipeline.py +++ b/applications/dynacell/src/dynacell/evaluation/pipeline.py @@ -964,11 +964,21 @@ def _aggregate(result: FovResult) -> None: else: # executor == "process": spawn-context ProcessPoolExecutor. # Workers lazy-load their own seg_model + extractors under the - # GPU lock; parent's copy stays loaded (used only here for - # checkpoint pre-warm side effect — discarded after this block). + # GPU lock; parent's copy was discarded after Phase-2 resolve. + # + # Streaming aggregation: futures arrive in completion order + # but we want deterministic per-FOV CSV / embedding NPZ row + # order, so the next-expected pos_name index advances + # opportunistically as results land. Only out-of-order + # FovResults stay buffered — once the expected position + # completes, we drain the buffer in order. Bounds peak + # parent-side memory to the number of out-of-order FOVs + # rather than the full N×seg_array (~N × 440 MB). from concurrent.futures import as_completed pos_names_in_order = [p[0] for p, _, _ in zip(pred_positions, gt_positions, seg_positions)] + next_idx = 0 + buffer: dict[str, FovResult] = {} with make_fov_executor(runtime) as pool: if pool is None: raise RuntimeError("make_fov_executor returned None for executor='process'") @@ -978,15 +988,19 @@ def _aggregate(result: FovResult) -> None: ): pos_name for pos_name in pos_names_in_order } - results_by_pos: dict[str, FovResult] = {} - for fut in tqdm(as_completed(futures), total=len(futures), desc="Processing positions"): - pos_name = futures[fut] - results_by_pos[pos_name] = fut.result() - # Aggregate in position order (independent of completion order) - # so per-FOV CSV / embedding NPZ row order is deterministic. - for fov_idx, pos_name in enumerate(pos_names_in_order): - _aggregate(results_by_pos[pos_name]) - maybe_gc_collect(fov_idx, runtime.gc_collect_every_n_fovs) + with tqdm(total=len(futures), desc="Processing positions") as pbar: + for fut in as_completed(futures): + pos_name = futures[fut] + buffer[pos_name] = fut.result() + # Drain in order while the next-expected position + # is available, releasing seg_array refs as soon + # as the aggregator's plate write completes. + while next_idx < len(pos_names_in_order) and (pos_names_in_order[next_idx] in buffer): + expected = pos_names_in_order[next_idx] + _aggregate(buffer.pop(expected)) + maybe_gc_collect(next_idx, runtime.gc_collect_every_n_fovs) + next_idx += 1 + pbar.update(1) finally: if seg_plate is not None: seg_plate.close() From d1091741669abc74836adbb2f20dddb6de1bd4ed Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin <1107762+alxndrkalinin@users.noreply.github.com> Date: Wed, 20 May 2026 15:34:16 -0700 Subject: [PATCH 16/20] feat(dynacell): grouped multi-condition eval driver (#427) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(dynacell): grouped multi-condition eval driver Add ``dynacell evaluate-grouped``: run the same (model, organelle) across multiple I/O variants (the canonical case is the three A549 treatment plates mock/denv/zikv) in one Python process, paying the segmenter + DINOv3 + DynaCLR + CELL-DINO cold-start once instead of N times. On serial-mode interactive jobs and on per-condition SLURM jobs this is 30-90s of weight load + checkpoint warmup amortized across all conditions. Under runtime.executor=process, each condition still spawns its own pool whose workers re-load model copies, so only the parent-side load is amortized; the driver prints a note when it detects this. Mechanism: - ``dynacell.evaluation.model_loader.load_eval_models`` extracts the segmenter + extractor instantiation that previously lived inline in ``evaluate_predictions`` and ``_worker_setup``. The bundle carries identity tags (dinov3_model_name, dynaclr_ckpt source paths, encoder cfg, celldino weights) so ``init_cache_context`` callers get identical cache-identity values across conditions. - ``evaluate_predictions(config, *, models=None)`` accepts a pre-loaded bundle and skips its inline load when supplied. ``models=None`` preserves the historical single-condition behavior; existing callers and tests are unchanged. - ``evaluate_predictions_grouped(config)`` reads a top-level ``conditions:`` list, deep-merges each overlay into a copy of the base, and runs them sequentially with the shared models. Per-condition ``_final_metrics_cache_valid`` short-circuits are honored independently. - ``_check_grouped_field_invariants`` rejects overlays that touch ``target_name``, ``feature_extractor.*``, ``compute_feature_metrics``, or ``use_gpu`` -- those determine which models get loaded and must be identical across conditions. ``io.*``, ``save.*``, ``runtime.*``, ``limit_positions``, ``force_recompute.*`` are free to vary. - Hydra entry point ``evaluate_model_grouped`` with new ``_configs/eval_grouped.yaml`` (inherits ``eval.yaml`` via defaults). - New CLI subcommand registered in ``__main__._HYDRA_COMMANDS``. Tests: ``tests/test_evaluation_grouped.py`` builds a 2-condition cache-only fixture (separate pred plates + cache dirs) and asserts byte-equal parity against sequential per-condition runs, plus rejection cases for empty ``conditions`` and forbidden model-loading-field overrides. Cache-only path means no models actually load -- the test validates the loop structure, not the speedup. Reuses the iohub fixture builders from ``test_evaluation_pipeline_parallel_cpu.py``. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor(dynacell): simplify grouped eval driver + fold load + cache-context unification Followup polish on d69d9709 (feat: grouped multi-condition eval driver). Addresses /simplify findings + the two related cleanups the user asked to fold in: Pipeline driver (pipeline.py): - Lazy model load — defer ``load_eval_models`` until the first condition that misses ``_final_metrics_cache_valid``. Restarts where every condition is cache-hit pay no 30–90 s cold start. - Invariant baseline is now the conditions-stripped input config, not the conditions[0]-merged config. A model-loading override smuggled into condition 0 is rejected symmetrically with conditions 1..N (previously it was silently re-baselined). New test guards the regression: ``test_grouped_rejects_condition0_model_field_override``. - ``_merge_condition`` drops the manual ``to_container`` round-trip deep-copy; ``OmegaConf.merge`` already returns a non-aliased config. - ``_check_grouped_field_invariants`` switched to ``OmegaConf.is_config`` in place of the ``type(OmegaConf.create([]))`` idiom; baseline snapshot is computed once and reused per iteration. - Extracted ``_load_cached_final_metrics`` so the three-NPY reload block is shared between ``evaluate_model`` and the grouped driver. - Removed redundant first-iteration merge + ``apply_dataset_ref`` for condition 0; the loop body now does both for every condition, including idx=0. Model loader (model_loader.py) and precompute-gt: - Added ``LoadFlags`` so ``load_eval_models`` can be called from both ``evaluate_predictions`` (all extractors gated by ``compute_feature_metrics``) and ``precompute_gt_artifacts`` (per-extractor ``build.masks`` / ``build.dinov3`` / ``build.dynaclr`` / ``build.celldino`` flags). ``LoadFlags.for_evaluate`` recovers the prior default. The 30-line inline load block in ``precompute_cli.py`` collapses to a single ``load_eval_models(config, flags=LoadFlags(...))`` call. Stricter celldino weights gate is preserved as an explicit pre-load raise. - Added ``init_cache_contexts(config, models)`` and ``init_gt_cache_context(config, models)`` so the four identical ``init_cache_context(..., side=..., **identity_tags)`` call sites in ``evaluate_predictions``, ``_worker_setup``, and ``precompute_gt_artifacts`` consolidate behind the EvalModels bundle. Test fixtures (tests/_eval_fixtures.py): - Promoted the cache-only fixture builders (``build_eval_config``, ``build_fixture``, ``make_hcs_plate``, ``make_mask_cache``, ``live_pipeline_module``, ``read_position_arrays``, ``N_POSITIONS``/``T``/``D``/``H``/``W``) out of ``test_evaluation_pipeline_parallel_cpu.py`` into a private ``_eval_fixtures.py`` module (underscore-prefixed so pytest skips collection). Both ``test_evaluation_pipeline_parallel_cpu.py`` and ``test_evaluation_grouped.py`` now import normally — no ``sys.path.insert`` (project's "no ``sys.path`` edits" rule). CI: - ``.github/workflows/test.yml`` adds ``test_evaluation_grouped.py`` to the ``test-dynacell-configs`` matrix so the grouped driver gets exercised on every push. All 104 tests in the affected suites pass. Pre-existing 48-test failure baseline (test_benchmark_config_composition + one test_evaluation_metrics case) is unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) * feat(dynacell): tier grouped-eval process-mode message by cache state Split the existing ``[grouped] note:`` print into two tiers so users get proportional signal for the configuration they actually picked: - ``executor=process`` + ``require_complete_cache=true`` (cache-only, n_conditions > 1) keeps the mild informational note — workers re-init per condition but skip ``prepare_segmentation_model`` (returns None) and don't instantiate the deep extractors, so the only overhead is N pool spawns. - ``executor=process`` + ``require_complete_cache=false`` (cache-miss, n_conditions > 1) gets a loud multi-line WARNING. This is the actual worst case: each condition's worker pool independently loads SuperModel + DINOv3 + DynaCLR + CELL-DINO, so total wasted cold-start scales as ~30-90 s × N_workers × N_conditions. The warning explicitly points to ``executor=serial`` as the fix, since the parent's pre-loaded ``EvalModels`` is shared across conditions in serial mode but lost across pool spawns in process mode. Two new tests in ``test_evaluation_grouped.py``: - ``test_grouped_warns_loudly_on_process_plus_cache_miss`` confirms the WARNING text fires (executor=process + require_complete_cache=false + 2 conditions) with all the expected keywords. - ``test_grouped_mild_note_on_process_plus_cache_only`` confirms the cache-only fixture (require_complete_cache=true) gets the mild note and does NOT trip the loud WARNING. CLAUDE.md "Grouped multi-condition eval" section updated with the tiered behavior so users searching for the WARNING text find context. Co-Authored-By: Claude Opus 4.7 (1M context) * fix(dynacell): grouped eval struct-mode + require_complete_cache invariant Three fixes triggered by Copilot's review on PR #427 plus an independent struct-mode bug that would have crashed every real Hydra invocation of ``dynacell evaluate-grouped``: 1. ``_merge_condition`` was over-simplified during the /simplify pass (commit c271d3f0). ``OmegaConf.merge`` propagates struct mode from its first argument, and Hydra composition always produces struct-mode configs. Result: merging an overlay carrying a ``name`` label (which isn't in the schema) raises ``ConfigKeyError``, and ``del merged["conditions"]`` raises ``ConfigTypeError`` even when the merge succeeds. Restore the ``to_container`` round-trip on the base — it escapes struct mode while still producing an independent config. Same fix applied to the ``models_base`` snapshot a few lines above. Existing tests didn't catch this because they construct configs via ``OmegaConf.create({...})`` directly (non-struct); real users going through ``@hydra.main`` would have hit the crash immediately. 2. Add ``io.require_complete_cache`` to ``_MODEL_LOADING_FIELDS`` so per-condition overlays cannot override it. The flag flips ``prepare_segmentation_model`` between "load SuperModel" and "return None" — letting a condition flip it would mean the shared model bundle is wrong for that condition (None when cache-miss expects a real segmenter, or loaded but never used). Treat as a grouped invariant. (Copilot thread on pipeline.py:1227.) 3. Add ``applications/dynacell/tests/__init__.py``. The new test files use ``from ._eval_fixtures import ...`` relative imports, which work under the project's ``--import-mode=importlib`` setting but are fragile under other pytest import modes. Adding ``__init__.py`` makes ``tests/`` a proper package and the relative imports work regardless of which import mode pytest is invoked with. (Copilot threads on test_evaluation_grouped.py:32 and test_evaluation_pipeline_parallel_cpu.py:42.) Two new regression tests in ``tests/test_evaluation_grouped.py``: - ``test_grouped_rejects_require_complete_cache_override`` covers the new invariant. - ``test_grouped_works_on_struct_mode_config`` wraps the parity fixture in ``OmegaConf.set_struct(grouped_cfg, True)`` to simulate a Hydra-composed config; would have failed before the ``_merge_condition`` fix. All 108 tests in affected suites pass. Co-Authored-By: Claude Opus 4.7 (1M context) * docs(dynacell): fix two CLAUDE.md doc nits flagged in PR #427 self-review Two 75-confidence findings from the parallel review on PR #427 that scored below the 80-threshold for the inline review comment but are worth fixing now while context is fresh: 1. "Reality check for process mode" subsection referenced ``pipeline.py:171`` as the parent seg-model load site. After the eval-parallelism refactor that line is a dataclass field (``pred_ts: list[np.ndarray] = ...``); the actual load now lives in ``load_eval_models(config)`` called from the top of ``evaluate_predictions``. Switch the reference from a brittle line number to a function-name pointer. 2. "Grouped multi-condition eval" subsection said the per-condition cache-load short-circuit fires "with ``force_recompute.final_metrics =false``", omitting the second gate: ``_final_metrics_cache_valid`` also returns False (forcing recompute) when ``force_recompute.all =true``, regardless of the ``final_metrics`` flag. Make both gates explicit so users who set ``force_recompute.all=true`` for unrelated reasons aren't surprised by unexpectedly recomputed metrics. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Alexandr Kalinin Co-authored-by: Claude Opus 4.7 (1M context) --- .github/workflows/test.yml | 4 +- applications/dynacell/CLAUDE.md | 50 ++- .../dynacell/src/dynacell/__main__.py | 1 + .../evaluation/_configs/eval_grouped.yaml | 34 ++ .../src/dynacell/evaluation/model_loader.py | 193 ++++++++++ .../src/dynacell/evaluation/pipeline.py | 340 ++++++++++++------ .../src/dynacell/evaluation/precompute_cli.py | 61 ++-- applications/dynacell/tests/__init__.py | 0 applications/dynacell/tests/_eval_fixtures.py | 213 +++++++++++ .../dynacell/tests/test_evaluation_grouped.py | 306 ++++++++++++++++ .../test_evaluation_pipeline_parallel_cpu.py | 230 ++---------- 11 files changed, 1065 insertions(+), 367 deletions(-) create mode 100644 applications/dynacell/src/dynacell/evaluation/_configs/eval_grouped.yaml create mode 100644 applications/dynacell/src/dynacell/evaluation/model_loader.py create mode 100644 applications/dynacell/tests/__init__.py create mode 100644 applications/dynacell/tests/_eval_fixtures.py create mode 100644 applications/dynacell/tests/test_evaluation_grouped.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index a6ccba1ae..5bc5b03ea 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -145,7 +145,8 @@ jobs: # tiny iohub fixture + prebuilt mask cache — no eval extras (no # cellpose, transformers, cubic) needed since target_name=er + # require_complete_cache=true short-circuit segmenter and feature - # extractor loads. + # extractor loads. tests/test_evaluation_grouped.py drives the + # multi-condition driver against the same cache-only fixture. run: | uv run --frozen pytest \ tests/test_benchmark_config_composition.py \ @@ -153,6 +154,7 @@ jobs: tests/test_runtime.py \ tests/test_evaluation_pipeline_parallel.py \ tests/test_evaluation_pipeline_parallel_cpu.py \ + tests/test_evaluation_grouped.py \ -v working-directory: applications/dynacell diff --git a/applications/dynacell/CLAUDE.md b/applications/dynacell/CLAUDE.md index a421898ee..44f13dac6 100644 --- a/applications/dynacell/CLAUDE.md +++ b/applications/dynacell/CLAUDE.md @@ -67,8 +67,56 @@ Under `executor=process`, workers return their slice of the timing log inside `F ### Reality check for process mode -- The parent still loads the seg model once at `pipeline.py:171` (used for checkpoint pre-warm side-effect under `process` mode + the seg model itself under `serial` mode). +- The parent still loads the seg model once via `load_eval_models(config)` at the start of `evaluate_predictions` in `pipeline.py` (used for checkpoint pre-warm side-effect under `process` mode + the seg model itself under `serial` mode). - `precompute_deep_features` (the upfront batched feature pre-fill) stays in the parent — single-pass over positions with the `DeepFeatureBatcher` cross-FOV amortization. Parallelizing precompute is a separate plan. - Workers cannot share open iohub plate handles or torch modules across the pickle boundary; each worker re-opens plates on first FOV. - `ProcessPoolExecutor.shutdown(wait=False, cancel_futures=True)` cancels queued futures only — in-flight workers continue to completion (~minutes if mid-cellpose). Ctrl-C may need `scancel` to fully release GPU memory. - Tests for the runtime module + FovResult pickle contract live in `applications/dynacell/tests/test_runtime.py` and `tests/test_evaluation_pipeline_parallel.py`. An end-to-end serial-vs-process parity test on a real iohub fixture is a follow-up (see plan `.claude/plans/eval-parallelism.md` §C5). + +### Grouped multi-condition eval + +When evaluating the same `(model, organelle)` across multiple I/O variants — typically the three A549 treatment plates (`mock`, `denv`, `zikv`), but also any case where the same trained model is scored on multiple datasets — use `dynacell evaluate-grouped`. The driver loads `SuperModel` + `DinoV3` + `DynaCLR` + `CELL-DINO` once, then loops over the conditions calling `evaluate_predictions(merged_cfg, models=...)` so the per-condition cold-start (~30–90 s of weight load + checkpoint warmup) is paid once total instead of N times. + +```yaml +# applications/dynacell/configs/.../eval_grouped_a549_mantis_er.yml +defaults: + - eval_grouped + - _self_ + +target_name: er # MUST be set at the base; conditions cannot override it +compute_feature_metrics: true +feature_extractor: + dinov3: { ... } # shared across all conditions + dynaclr: { checkpoint: /path/to/dynaclr.ckpt, ... } + celldino: { weights_path: /path/to/celldino.ckpt, ... } + +conditions: + - name: a549_mock + io: { pred_path: ..., gt_path: ..., gt_cache_dir: ..., pred_cache_dir: ..., cell_segmentation_path: ... } + save: { save_dir: /path/to/out/a549_mock } + - name: a549_denv + io: { ... } + save: { save_dir: /path/to/out/a549_denv } + - name: a549_zikv + io: { ... } + save: { save_dir: /path/to/out/a549_zikv } +``` + +Invoke: +```sh +uv run dynacell evaluate-grouped -c eval_grouped_a549_mantis_er +``` + +Constraints (enforced at runtime by `_check_grouped_field_invariants`): +- Per-condition overlays may freely override `io.*`, `save.*`, `runtime.*`, `limit_positions`, `force_recompute.*`, and carry a `name` label. +- Overlays MUST NOT change `target_name`, `feature_extractor.*`, `compute_feature_metrics`, or `use_gpu` — those gate model loading and must be identical across conditions. Run such variants separately. +- Each condition independently honors `_final_metrics_cache_valid` — if a condition's CSV/NPY already exist AND both `force_recompute.all=false` and `force_recompute.final_metrics=false`, the driver skips it and loads the cached outputs. Either flag set to `true` bypasses the cache. + +Process-mode caveat: under `runtime.executor=process`, the parent's shared `EvalModels` is passed to `evaluate_predictions(..., models=...)` so the parent-side pre-warm is amortized, but each condition still spawns a fresh `ProcessPoolExecutor` whose workers re-load their own model copies. Use `executor=serial` to maximize the amortization benefit. The driver tiers its message based on the cache mode: + +- `executor=process` + `require_complete_cache=true` + `n_conditions > 1` → mild **note**: workers re-init per condition but skip `prepare_segmentation_model` (returns None) and don't instantiate extractors, so the cost is just N pool spawns. +- `executor=process` + `require_complete_cache=false` + `n_conditions > 1` → loud **WARNING**: each condition's worker pool independently loads SuperModel + DINOv3 + DynaCLR + CELL-DINO. Total waste is ~30-90 s × `runtime.fov_workers` × `n_conditions`. Fix: switch to `executor=serial` so the parent's pre-loaded models are reused across conditions; reserve `process` for per-FOV parallelism within a *single* condition. + +For an A40 / single-GPU interactive node where you'd run serial anyway, this is the default win. + +Tests: `applications/dynacell/tests/test_evaluation_grouped.py` validates byte-equal parity against sequential per-condition runs on the same cache-only fixture used by `test_evaluation_pipeline_parallel_cpu.py`, plus rejection cases for empty `conditions` and forbidden model-loading-field overrides. diff --git a/applications/dynacell/src/dynacell/__main__.py b/applications/dynacell/src/dynacell/__main__.py index 3193dde87..a357c4dbe 100644 --- a/applications/dynacell/src/dynacell/__main__.py +++ b/applications/dynacell/src/dynacell/__main__.py @@ -19,6 +19,7 @@ _HYDRA_COMMANDS: dict[str, tuple[str, str, str]] = { "evaluate": ("dynacell.evaluation.pipeline", "evaluate_model", "eval"), + "evaluate-grouped": ("dynacell.evaluation.pipeline", "evaluate_model_grouped", "eval"), "precompute-gt": ("dynacell.evaluation.precompute_cli", "precompute_gt", "eval"), "report": ("dynacell.reporting.cli", "generate_report", "report"), } diff --git a/applications/dynacell/src/dynacell/evaluation/_configs/eval_grouped.yaml b/applications/dynacell/src/dynacell/evaluation/_configs/eval_grouped.yaml new file mode 100644 index 000000000..bcce29ec4 --- /dev/null +++ b/applications/dynacell/src/dynacell/evaluation/_configs/eval_grouped.yaml @@ -0,0 +1,34 @@ +defaults: + - eval + - _self_ + +# Per-condition overlays for grouped multi-condition eval. Each entry is a +# deep merge applied to the base config above. Use this to run e.g. all three +# A549 treatment conditions (mock / denv / zikv) in one process, amortizing +# the segmenter + DINOv3 + DynaCLR + CELL-DINO load. +# +# Each overlay may freely change ``io.*``, ``save.*``, ``runtime.*``, +# ``limit_positions``, ``force_recompute.*``, and optionally carry a ``name`` +# label used in logs + the returned result tuple. Overlays MUST NOT touch +# ``target_name``, ``feature_extractor.*``, ``compute_feature_metrics``, or +# ``use_gpu`` — those gate model loading, which is shared. See +# applications/dynacell/CLAUDE.md "Grouped multi-condition eval". +# +# Example: +# conditions: +# - name: a549_mock +# io: +# pred_path: /path/to/mock_pred.zarr +# gt_path: /path/to/mock_gt.zarr +# cell_segmentation_path: /path/to/mock_seg.zarr +# gt_cache_dir: /path/to/mock_gt_cache +# pred_cache_dir: /path/to/mock_pred_cache +# save: +# save_dir: /path/to/eval_out/a549_mock +# - name: a549_denv +# io: { ... } +# save: { save_dir: /path/to/eval_out/a549_denv } +# - name: a549_zikv +# io: { ... } +# save: { save_dir: /path/to/eval_out/a549_zikv } +conditions: [] diff --git a/applications/dynacell/src/dynacell/evaluation/model_loader.py b/applications/dynacell/src/dynacell/evaluation/model_loader.py new file mode 100644 index 000000000..0284a7239 --- /dev/null +++ b/applications/dynacell/src/dynacell/evaluation/model_loader.py @@ -0,0 +1,193 @@ +"""Shared model-loading entry point for the dynacell eval pipeline. + +A single source of truth for instantiating the segmenter + deep-feature +extractors used by ``evaluate_predictions`` and the precompute-gt CLI. +Lifting the load out of the pipeline lets the grouped multi-condition +driver build models once and reuse them across conditions that differ +only in I/O paths; the precompute-gt CLI shares the same code path with +its own per-extractor gates (``build.masks`` / ``build.dinov3`` / ...). + +The dataclass carries the **identity tags** (model name, checkpoint sha +sources) that ``init_cache_context`` needs, so callers don't have to +recompute them per condition. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from omegaconf import DictConfig + + +@dataclass +class EvalModels: + """Bundle of pre-loaded segmenter + feature extractors with their identity tags. + + Attributes + ---------- + seg_model : Any | None + SuperModel returned by ``prepare_segmentation_model``. May be None + when ``io.require_complete_cache=true`` short-circuits the load + or when ``LoadFlags.masks=False``. + dinov3, dynaclr, celldino : Any | None + Deep feature extractor instances. Each is None when its + ``LoadFlags`` gate is False, or (for celldino) when its + ``weights_path`` is unset. + dinov3_model_name, dynaclr_ckpt_path, dynaclr_encoder_cfg, + celldino_weights_path : str | dict | None + Identity tags consumed by ``init_cache_context``. Stay None when + the corresponding extractor was not loaded. + """ + + seg_model: Any | None + dinov3: Any | None + dynaclr: Any | None + celldino: Any | None + dinov3_model_name: str | None + dynaclr_ckpt_path: str | None + dynaclr_encoder_cfg: dict[str, Any] | None + celldino_weights_path: str | None + + +@dataclass(frozen=True) +class LoadFlags: + """Per-model gate for :func:`load_eval_models`. + + Lets ``precompute-gt`` (per-extractor flags ``build.masks`` / + ``build.dinov3`` / ``build.dynaclr`` / ``build.celldino``) and + ``evaluate-predictions`` (all extractors gated together by + ``compute_feature_metrics``) share one loader. + """ + + masks: bool = True + dinov3: bool = False + dynaclr: bool = False + celldino: bool = False + + @classmethod + def for_evaluate(cls, config: DictConfig) -> LoadFlags: + """Flags matching the ``evaluate_predictions`` call path. + + Masks always on (``prepare_segmentation_model`` handles its own + ``require_complete_cache`` short-circuit). Extractors gated as a + group by ``config.compute_feature_metrics``; celldino additionally + soft-skips inside the loader when its ``weights_path`` is null. + """ + ext_on = bool(config.compute_feature_metrics) + return cls(masks=True, dinov3=ext_on, dynaclr=ext_on, celldino=ext_on) + + +def load_eval_models(config: DictConfig, *, flags: LoadFlags | None = None) -> EvalModels: + """Instantiate the segmenter + deep feature extractors from ``config``. + + No GPU serialization wrapping — callers (worker setup) wrap with + ``gpu_serialization_lock`` as needed. The parent of + ``evaluate_predictions`` calls this unwrapped today, so this entry + preserves that exact pattern. + + Parameters + ---------- + config : DictConfig + Resolved eval config. Reads ``target_name``, + ``feature_extractor.{dinov3,dynaclr,celldino}``, + and (transitively) ``io.require_complete_cache``. + flags : LoadFlags, optional + Per-model gates. Defaults to :meth:`LoadFlags.for_evaluate`, + which preserves the historical ``evaluate_predictions`` + behavior. Pass a custom :class:`LoadFlags` to load a subset + (e.g. for precompute-gt where ``build.dinov3`` / etc. are + toggled independently). + + Returns + ------- + EvalModels + All model handles + identity tags. Each slot is None when its + flag is off; celldino additionally soft-skips when its + ``weights_path`` is null even with the flag on. + """ + from dynacell.evaluation.pipeline_cache import resolve_dynaclr_encoder_cfg + from dynacell.evaluation.segmentation import prepare_segmentation_model + from dynacell.evaluation.utils import ( + CellDinoFeatureExtractor, + DinoV3FeatureExtractor, + DynaCLRFeatureExtractor, + ) + + if flags is None: + flags = LoadFlags.for_evaluate(config) + + seg_model = prepare_segmentation_model(config) if flags.masks else None + + dinov3_model_name: str | None = None + dynaclr_ckpt_path: str | None = None + dynaclr_encoder_cfg: dict[str, Any] | None = None + celldino_weights_path: str | None = None + dinov3 = None + dynaclr = None + celldino = None + + if flags.dinov3: + dinov3_model_name = config.feature_extractor.dinov3.pretrained_model_name + dinov3 = DinoV3FeatureExtractor(dinov3_model_name) + if flags.dynaclr: + dynaclr_config = config.feature_extractor.dynaclr + dynaclr_ckpt_path = str(dynaclr_config.checkpoint) + dynaclr_encoder_cfg = resolve_dynaclr_encoder_cfg(config) + dynaclr = DynaCLRFeatureExtractor( + checkpoint=dynaclr_config.checkpoint, + encoder_config=dynaclr_encoder_cfg, + ) + if flags.celldino: + celldino_cfg = config.feature_extractor.celldino + if celldino_cfg.weights_path is not None: + celldino_weights_path = str(celldino_cfg.weights_path) + celldino = CellDinoFeatureExtractor( + weights_path=celldino_weights_path, + img_size=int(celldino_cfg.img_size), + patch_size=int(celldino_cfg.patch_size), + ) + + return EvalModels( + seg_model=seg_model, + dinov3=dinov3, + dynaclr=dynaclr, + celldino=celldino, + dinov3_model_name=dinov3_model_name, + dynaclr_ckpt_path=dynaclr_ckpt_path, + dynaclr_encoder_cfg=dynaclr_encoder_cfg, + celldino_weights_path=celldino_weights_path, + ) + + +def _identity_kwargs(models: EvalModels) -> dict[str, Any]: + """Identity tags forwarded into ``init_cache_context``.""" + return { + "dinov3_model_name": models.dinov3_model_name, + "dynaclr_ckpt_path": models.dynaclr_ckpt_path, + "dynaclr_encoder_cfg": models.dynaclr_encoder_cfg, + "celldino_weights_path": models.celldino_weights_path, + } + + +def init_cache_contexts(config: DictConfig, models: EvalModels) -> tuple[Any, Any]: + """Build ``(gt_ctx, pred_ctx)`` cache contexts using ``models``'s identity tags. + + Replaces the four-call-site duplication in ``evaluate_predictions`` + and ``_worker_setup`` where the same six kwargs flowed into + ``init_cache_context(..., side="gt")`` then again into + ``init_cache_context(..., side="pred")``. + """ + from dynacell.evaluation.pipeline_cache import init_cache_context + + kwargs = _identity_kwargs(models) + gt_ctx = init_cache_context(config, side="gt", **kwargs) + pred_ctx = init_cache_context(config, side="pred", **kwargs) + return gt_ctx, pred_ctx + + +def init_gt_cache_context(config: DictConfig, models: EvalModels) -> Any: + """Build only the GT cache context — used by ``precompute-gt`` which is GT-only.""" + from dynacell.evaluation.pipeline_cache import init_cache_context + + return init_cache_context(config, side="gt", **_identity_kwargs(models)) diff --git a/applications/dynacell/src/dynacell/evaluation/pipeline.py b/applications/dynacell/src/dynacell/evaluation/pipeline.py index 904446186..3081b0640 100644 --- a/applications/dynacell/src/dynacell/evaluation/pipeline.py +++ b/applications/dynacell/src/dynacell/evaluation/pipeline.py @@ -35,14 +35,13 @@ evaluate_segmentations, features_from_crops, ) +from dynacell.evaluation.model_loader import EvalModels, init_cache_contexts, load_eval_models from dynacell.evaluation.pipeline_cache import ( flush_manifest, fov_cp_features, fov_deep_features, fov_masks, - init_cache_context, precompute_deep_features, - resolve_dynaclr_encoder_cfg, ) from dynacell.evaluation.utils import plot_metrics @@ -528,68 +527,21 @@ def _worker_setup(config: DictConfig) -> None: """ if _WORKER_STATE.get("initialized"): return - from dynacell.evaluation.pipeline_cache import init_cache_context, resolve_dynaclr_encoder_cfg from dynacell.evaluation.runtime import gpu_serialization_lock - from dynacell.evaluation.segmentation import prepare_segmentation_model - from dynacell.evaluation.utils import ( - CellDinoFeatureExtractor, - DinoV3FeatureExtractor, - DynaCLRFeatureExtractor, - ) use_gpu = bool(getattr(config, "use_gpu", True)) with gpu_serialization_lock(gate=use_gpu): - seg_model = prepare_segmentation_model(config) - dinov3_feature_extractor = None - dynaclr_feature_extractor = None - celldino_feature_extractor = None - dinov3_model_name = None - dynaclr_ckpt_path = None - dynaclr_encoder_cfg = None - celldino_weights_path = None - if config.compute_feature_metrics: - dinov3_model_name = config.feature_extractor.dinov3.pretrained_model_name - dinov3_feature_extractor = DinoV3FeatureExtractor(dinov3_model_name) - dynaclr_config = config.feature_extractor.dynaclr - dynaclr_ckpt_path = str(dynaclr_config.checkpoint) - dynaclr_encoder_cfg = resolve_dynaclr_encoder_cfg(config) - dynaclr_feature_extractor = DynaCLRFeatureExtractor( - checkpoint=dynaclr_config.checkpoint, - encoder_config=dynaclr_encoder_cfg, - ) - celldino_cfg = config.feature_extractor.celldino - if celldino_cfg.weights_path is not None: - celldino_weights_path = str(celldino_cfg.weights_path) - celldino_feature_extractor = CellDinoFeatureExtractor( - weights_path=celldino_weights_path, - img_size=int(celldino_cfg.img_size), - patch_size=int(celldino_cfg.patch_size), - ) + models = load_eval_models(config) - cache_ctx = init_cache_context( - config, - side="gt", - dinov3_model_name=dinov3_model_name, - dynaclr_ckpt_path=dynaclr_ckpt_path, - dynaclr_encoder_cfg=dynaclr_encoder_cfg, - celldino_weights_path=celldino_weights_path, - ) - pred_cache_ctx = init_cache_context( - config, - side="pred", - dinov3_model_name=dinov3_model_name, - dynaclr_ckpt_path=dynaclr_ckpt_path, - dynaclr_encoder_cfg=dynaclr_encoder_cfg, - celldino_weights_path=celldino_weights_path, - ) + cache_ctx, pred_cache_ctx = init_cache_contexts(config, models) _WORKER_STATE.update( { "initialized": True, - "seg_model": seg_model, - "dinov3": dinov3_feature_extractor, - "dynaclr": dynaclr_feature_extractor, - "celldino": celldino_feature_extractor, + "seg_model": models.seg_model, + "dinov3": models.dinov3, + "dynaclr": models.dynaclr, + "celldino": models.celldino, "cache_ctx": cache_ctx, "pred_cache_ctx": pred_cache_ctx, } @@ -673,8 +625,22 @@ def _worker_run_fov(config: DictConfig, pos_name: str, cuda_empty_every_n: int) return result -def evaluate_predictions(config: DictConfig): - """Evaluate predictions on all test images.""" +def evaluate_predictions(config: DictConfig, *, models: EvalModels | None = None): + """Evaluate predictions on all test images. + + Parameters + ---------- + config : DictConfig + Resolved eval config. + models : EvalModels | None, optional + Pre-loaded segmenter + feature extractors. When provided, the + inline ``load_eval_models(config)`` call is skipped — used by the + grouped multi-condition driver to amortize model loads across + conditions sharing the same target/extractors. Default ``None`` + preserves the historical single-condition behavior. Note: under + ``runtime.executor=process``, workers still load their own model + copies; this kwarg saves only the parent-side load. + """ from dynacell.evaluation.runtime import ( apply_thread_budget, dump_timings_csv, @@ -683,8 +649,6 @@ def evaluate_predictions(config: DictConfig): reset_timings, resolve_runtime, ) - from dynacell.evaluation.segmentation import prepare_segmentation_model - from dynacell.evaluation.utils import CellDinoFeatureExtractor, DinoV3FeatureExtractor, DynaCLRFeatureExtractor # Phase 1 runtime resolution: lock in executor + thread caps before any # heavy work. fov_workers may be provisional when "auto"; re-resolved in @@ -709,53 +673,17 @@ def evaluate_predictions(config: DictConfig): save_dir = Path(config.save.save_dir) save_dir.mkdir(parents=True, exist_ok=True) - seg_model = prepare_segmentation_model(config) - - dinov3_model_name = None - dynaclr_ckpt_path = None - dynaclr_encoder_cfg = None - celldino_weights_path = None - dinov3_feature_extractor = None - dynaclr_feature_extractor = None - celldino_feature_extractor = None + if config.compute_feature_metrics and io_config.cell_segmentation_path is None: + raise ValueError("io.cell_segmentation_path is required when compute_feature_metrics=true") - if config.compute_feature_metrics: - if io_config.cell_segmentation_path is None: - raise ValueError("io.cell_segmentation_path is required when compute_feature_metrics=true") - dinov3_model_name = config.feature_extractor.dinov3.pretrained_model_name - dinov3_feature_extractor = DinoV3FeatureExtractor(dinov3_model_name) - dynaclr_config = config.feature_extractor.dynaclr - dynaclr_ckpt_path = str(dynaclr_config.checkpoint) - dynaclr_encoder_cfg = resolve_dynaclr_encoder_cfg(config) - dynaclr_feature_extractor = DynaCLRFeatureExtractor( - checkpoint=dynaclr_config.checkpoint, - encoder_config=dynaclr_encoder_cfg, - ) - celldino_cfg = config.feature_extractor.celldino - if celldino_cfg.weights_path is not None: - celldino_weights_path = str(celldino_cfg.weights_path) - celldino_feature_extractor = CellDinoFeatureExtractor( - weights_path=celldino_weights_path, - img_size=int(celldino_cfg.img_size), - patch_size=int(celldino_cfg.patch_size), - ) + if models is None: + models = load_eval_models(config) + seg_model = models.seg_model + dinov3_feature_extractor = models.dinov3 + dynaclr_feature_extractor = models.dynaclr + celldino_feature_extractor = models.celldino - cache_ctx = init_cache_context( - config, - side="gt", - dinov3_model_name=dinov3_model_name, - dynaclr_ckpt_path=dynaclr_ckpt_path, - dynaclr_encoder_cfg=dynaclr_encoder_cfg, - celldino_weights_path=celldino_weights_path, - ) - pred_cache_ctx = init_cache_context( - config, - side="pred", - dinov3_model_name=dinov3_model_name, - dynaclr_ckpt_path=dynaclr_ckpt_path, - dynaclr_encoder_cfg=dynaclr_encoder_cfg, - celldino_weights_path=celldino_weights_path, - ) + cache_ctx, pred_cache_ctx = init_cache_contexts(config, models) seg_path = Path(io_config.cell_segmentation_path) if io_config.cell_segmentation_path is not None else None @@ -1156,22 +1084,200 @@ def _final_metrics_cache_valid(config: DictConfig) -> bool: return pixel_ok and mask_ok and feature_ok +def _load_cached_final_metrics(config: DictConfig) -> tuple[list, list, list]: + """Reload the (pixel, mask, feature) NPY caches saved by ``save_metrics``.""" + save_dir = Path(config.save.save_dir) + pixel = np.load(save_dir / config.save.pixel_metrics_filename, allow_pickle=True).tolist() + mask = np.load(save_dir / config.save.mask_metrics_filename, allow_pickle=True).tolist() + if config.compute_feature_metrics: + feature = np.load(save_dir / config.save.feature_metrics_filename, allow_pickle=True).tolist() + else: + feature = [] + return pixel, mask, feature + + +_MODEL_LOADING_FIELDS: tuple[str, ...] = ( + "target_name", + "feature_extractor", + "compute_feature_metrics", + "use_gpu", + # require_complete_cache flips ``prepare_segmentation_model`` between + # "load real SuperModel" and "return None" — letting a condition + # override it would mean the shared seg_model bundle is wrong for that + # condition (None when cache-miss expects a real model, or loaded but + # never used). Treat as a grouped invariant. + "io.require_complete_cache", +) + + +def _snapshot_field(cfg: DictConfig, cfg_field: str): + """Resolve a model-loading field to a comparable plain Python value.""" + node = OmegaConf.select(cfg, cfg_field, default=None) + if OmegaConf.is_config(node): + return OmegaConf.to_container(node, resolve=False) + return node + + +def _merge_condition(base: DictConfig, overrides: DictConfig | dict) -> DictConfig: + """Return a fresh DictConfig with ``overrides`` deep-merged into ``base``. + + Hydra composition always returns struct-mode configs. ``OmegaConf.merge`` + propagates struct mode from its first argument, so merging an overlay + that carries fields outside the base's schema (notably the per-condition + ``name`` label) raises ``ConfigKeyError``, and deleting the + ``conditions`` / ``name`` keys raises ``ConfigTypeError`` even when the + merge succeeds. The to_container round-trip below escapes struct mode + while still producing a config that's independent of ``base``. + """ + base_copy = OmegaConf.create(OmegaConf.to_container(base, resolve=False)) + merged = OmegaConf.merge(base_copy, OmegaConf.create(overrides)) + for key in ("conditions", "name"): + if key in merged: + del merged[key] + return merged # type: ignore[return-value] + + +def _check_grouped_field_invariants(base_snapshot: dict[str, object], merged: DictConfig, condition_name: str) -> None: + """Raise if a per-condition merged config disagrees with the baseline on model-loading fields. + + ``base_snapshot`` is computed once per grouped run from the + conditions-stripped base, so condition 0 is validated symmetrically + with all later conditions (a condition 0 overlay that sneaks in e.g. + ``target_name`` would be rejected, not silently adopted as the new + "base"). + """ + for cfg_field in _MODEL_LOADING_FIELDS: + merged_val = _snapshot_field(merged, cfg_field) + if base_snapshot[cfg_field] != merged_val: + raise ValueError( + f"Condition {condition_name!r}: overrides changed model-loading field " + f"{cfg_field!r}. Move it to the base config or run this condition separately." + ) + + +def evaluate_predictions_grouped(config: DictConfig) -> list[tuple[str, tuple]]: + """Run ``evaluate_predictions`` over a list of conditions sharing one model load. + + Reads ``config.conditions`` (list of per-condition overrides). For each + condition, merges its overrides into a copy of the base config and + calls :func:`evaluate_predictions` with the shared ``EvalModels``. + Models are loaded lazily on the first condition that misses its + ``_final_metrics_cache_valid`` short-circuit, so restarts where every + condition is cache-hit pay no cold-start. + + Conditions may freely override ``io.*``, ``save.*``, ``runtime.*``, + ``limit_positions``, and ``force_recompute.*``. They must NOT change + ``target_name``, ``feature_extractor.*``, ``compute_feature_metrics``, + or ``use_gpu`` — those gate which models get loaded. + + Parameters + ---------- + config : DictConfig + Eval config with an extra top-level ``conditions: [...]`` list. + Each entry is a dict-like overlay applied to the base. + + Returns + ------- + list[tuple[str, tuple]] + ``[(condition_name, (pixel_rows, mask_rows, feature_rows)), ...]`` + in input order. ``condition_name`` is taken from the entry's + ``name`` field, falling back to its index as a string. + + Notes + ----- + Under ``runtime.executor=process``, workers still load their own + model copies per condition (the pool is rebuilt inside each + ``evaluate_predictions`` call). Only the parent-side load is shared. + Use ``executor=serial`` to maximize the amortization benefit. + """ + conditions = OmegaConf.select(config, "conditions", default=None) + if not conditions: + raise ValueError("evaluate_predictions_grouped requires a non-empty top-level 'conditions' list") + + executor = OmegaConf.select(config, "runtime.executor", default="serial") + require_complete = bool(OmegaConf.select(config, "io.require_complete_cache", default=False)) + n_conditions = len(conditions) + if executor == "process" and n_conditions > 1: + if require_complete: + # Cache-only path: workers skip prepare_segmentation_model (returns + # None) and don't instantiate extractors. Pool re-spawn is the only + # overhead — mild informational note, not a warning. + print( + "[grouped] note: runtime.executor=process with require_complete_cache=true — " + f"workers re-init per condition but no models actually load. " + f"{n_conditions} pool spawns total; consider executor=serial to skip " + "spawn overhead entirely." + ) + else: + # Worst-case: each condition's worker pool independently loads + # SuperModel + DINOv3 + DynaCLR + CELL-DINO. Total wasted + # cold-start ≈ load_time × N_workers × N_conditions. + print( + "\n" + "[grouped] !!! WARNING: runtime.executor=process + " + f"{n_conditions} conditions + require_complete_cache=false !!!\n" + " Each condition's worker pool independently loads SuperModel + " + "DINOv3 + DynaCLR + CELL-DINO.\n" + " Expected waste: ~30-90 s × N_workers × " + f"{n_conditions} conditions of redundant cold-start.\n" + " Fix: set runtime.executor=serial to share the parent's pre-loaded " + "models across conditions.\n" + " Use process mode only when you need per-FOV parallelism within a " + "single condition.\n" + ) + + # Canonical baseline for invariant checks: the input config with the + # ``conditions`` list dropped. Snapshot once; conditions are validated + # against this, including condition 0. Round-trip through to_container + # to escape Hydra's struct-mode flag — ``del`` on a struct DictConfig + # raises ``ConfigTypeError``. + models_base = OmegaConf.create(OmegaConf.to_container(config, resolve=False)) + if "conditions" in models_base: + del models_base["conditions"] + base_snapshot = {field: _snapshot_field(models_base, field) for field in _MODEL_LOADING_FIELDS} + + models: EvalModels | None = None + + def get_models() -> EvalModels: + nonlocal models + if models is None: + print(f"[grouped] loading shared models for {len(conditions)} conditions ...") + models = load_eval_models(models_base) + return models + + results: list[tuple[str, tuple]] = [] + for idx, cond in enumerate(conditions): + name = ( + cond.get("name", str(idx)) if isinstance(cond, dict) else OmegaConf.select(cond, "name", default=str(idx)) + ) + merged = _merge_condition(config, cond) + apply_dataset_ref(merged) + _check_grouped_field_invariants(base_snapshot, merged, name) + print(f"[grouped] ({idx + 1}/{len(conditions)}) evaluating {name!r} → {merged.save.save_dir}") + + if _final_metrics_cache_valid(merged): + print(f"[grouped] ({idx + 1}/{len(conditions)}) {name!r}: reusing cached final metrics") + pixel_metrics, mask_metrics, feature_metrics = _load_cached_final_metrics(merged) + else: + pixel_metrics, mask_metrics, feature_metrics = evaluate_predictions(merged, models=get_models()) + save_metrics( + merged, + pixel_metrics=pixel_metrics, + mask_metrics=mask_metrics, + feature_metrics=feature_metrics, + ) + results.append((name, (pixel_metrics, mask_metrics, feature_metrics))) + + return results + + @hydra.main(version_base="1.2", config_path="_configs", config_name="eval") def evaluate_model(config: DictConfig): """Evaluate model on test images.""" apply_dataset_ref(config) - save_dir = Path(config.save.save_dir) - pixel_metrics_path = save_dir / config.save.pixel_metrics_filename - mask_metrics_path = save_dir / config.save.mask_metrics_filename - feature_metrics_path = save_dir / config.save.feature_metrics_filename if _final_metrics_cache_valid(config): print("Found existing metrics.") - pixel_metrics = np.load(pixel_metrics_path, allow_pickle=True).tolist() - mask_metrics = np.load(mask_metrics_path, allow_pickle=True).tolist() - if config.compute_feature_metrics: - feature_metrics = np.load(feature_metrics_path, allow_pickle=True).tolist() - else: - feature_metrics = [] + pixel_metrics, mask_metrics, feature_metrics = _load_cached_final_metrics(config) else: pixel_metrics, mask_metrics, feature_metrics = evaluate_predictions(config) save_metrics( @@ -1183,5 +1289,11 @@ def evaluate_model(config: DictConfig): return pixel_metrics, mask_metrics, feature_metrics +@hydra.main(version_base="1.2", config_path="_configs", config_name="eval_grouped") +def evaluate_model_grouped(config: DictConfig): + """Run grouped multi-condition eval, amortizing model loads across conditions.""" + return evaluate_predictions_grouped(config) + + if __name__ == "__main__": evaluate_model() diff --git a/applications/dynacell/src/dynacell/evaluation/precompute_cli.py b/applications/dynacell/src/dynacell/evaluation/precompute_cli.py index cd1861bfc..6c9898337 100644 --- a/applications/dynacell/src/dynacell/evaluation/precompute_cli.py +++ b/applications/dynacell/src/dynacell/evaluation/precompute_cli.py @@ -21,21 +21,18 @@ from dynacell.evaluation._ref_hook import apply_dataset_ref from dynacell.evaluation.metrics import build_crops +from dynacell.evaluation.model_loader import LoadFlags, init_gt_cache_context, load_eval_models from dynacell.evaluation.pipeline_cache import ( DeepFeatureBatcher, flush_manifest, fov_cp_features, fov_masks, - init_cache_context, - resolve_dynaclr_encoder_cfg, ) def precompute_gt_artifacts(config: DictConfig) -> None: """Build every GT-side artifact toggled on in ``config.build``.""" from dynacell.evaluation.runtime import apply_thread_budget, resolve_runtime - from dynacell.evaluation.segmentation import prepare_segmentation_model - from dynacell.evaluation.utils import CellDinoFeatureExtractor, DinoV3FeatureExtractor, DynaCLRFeatureExtractor if config.io.gt_cache_dir is None: raise ValueError("io.gt_cache_dir is required for dynacell precompute-gt") @@ -62,46 +59,28 @@ def precompute_gt_artifacts(config: DictConfig) -> None: "build.cp / build.dinov3 / build.dynaclr / build.celldino is true" ) - seg_model = prepare_segmentation_model(config) if build.masks else None - - dinov3_model_name = None - dynaclr_ckpt_path = None - dynaclr_encoder_cfg = None - celldino_weights_path = None - dinov3_feature_extractor = None - dynaclr_feature_extractor = None - celldino_feature_extractor = None - - if build.dinov3: - dinov3_model_name = config.feature_extractor.dinov3.pretrained_model_name - dinov3_feature_extractor = DinoV3FeatureExtractor(dinov3_model_name) - if build.dynaclr: - dynaclr_config = config.feature_extractor.dynaclr - dynaclr_ckpt_path = str(dynaclr_config.checkpoint) - dynaclr_encoder_cfg = resolve_dynaclr_encoder_cfg(config) - dynaclr_feature_extractor = DynaCLRFeatureExtractor( - checkpoint=dynaclr_config.checkpoint, - encoder_config=dynaclr_encoder_cfg, - ) - if build.celldino: - celldino_cfg = config.feature_extractor.celldino - if celldino_cfg.weights_path is None: - raise ValueError("feature_extractor.celldino.weights_path is required when build.celldino=true") - celldino_weights_path = str(celldino_cfg.weights_path) - celldino_feature_extractor = CellDinoFeatureExtractor( - weights_path=celldino_weights_path, - img_size=int(celldino_cfg.img_size), - patch_size=int(celldino_cfg.patch_size), - ) + # Stricter celldino gate than evaluate_predictions: precompute-gt + # requires an explicit weights path when build.celldino=true (the + # eval path soft-skips a null weights_path; here it would silently + # produce no cache, so we surface it as an error). + if build.celldino and config.feature_extractor.celldino.weights_path is None: + raise ValueError("feature_extractor.celldino.weights_path is required when build.celldino=true") - cache_ctx = init_cache_context( + models = load_eval_models( config, - side="gt", - dinov3_model_name=dinov3_model_name, - dynaclr_ckpt_path=dynaclr_ckpt_path, - dynaclr_encoder_cfg=dynaclr_encoder_cfg, - celldino_weights_path=celldino_weights_path, + flags=LoadFlags( + masks=bool(build.masks), + dinov3=bool(build.dinov3), + dynaclr=bool(build.dynaclr), + celldino=bool(build.celldino), + ), ) + seg_model = models.seg_model + dinov3_feature_extractor = models.dinov3 + dynaclr_feature_extractor = models.dynaclr + celldino_feature_extractor = models.celldino + + cache_ctx = init_gt_cache_context(config, models) gt_path = Path(config.io.gt_path) seg_path = Path(config.io.cell_segmentation_path) if config.io.cell_segmentation_path is not None else None diff --git a/applications/dynacell/tests/__init__.py b/applications/dynacell/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/applications/dynacell/tests/_eval_fixtures.py b/applications/dynacell/tests/_eval_fixtures.py new file mode 100644 index 000000000..fe43e9f6a --- /dev/null +++ b/applications/dynacell/tests/_eval_fixtures.py @@ -0,0 +1,213 @@ +"""Shared cache-only fixture helpers for dynacell eval integration tests. + +Underscore prefix keeps pytest from collecting this file as tests. Imported +by ``test_evaluation_pipeline_parallel_cpu.py`` and +``test_evaluation_grouped.py`` (and any future eval test that wants the +same scaffolding) — see CLAUDE.md "no ``sys.path`` edits" rule. + +Builds tiny iohub HCS plates + pre-populated mask caches so the +cache-only eval path (``target_name=er`` + ``require_complete_cache=true`` ++ ``compute_feature_metrics=false``) runs end-to-end without loading any +segmenter / extractor / cubic — small enough to fit a single-CPU CI in +under ~30 s. +""" + +from __future__ import annotations + +import importlib +import sys +from datetime import UTC, datetime +from pathlib import Path + +import numpy as np +from iohub.ngff import open_ome_zarr +from omegaconf import OmegaConf + +N_POSITIONS = 3 +T = 2 +D = 4 +H = 32 +W = 32 + + +def live_pipeline_module(): + """Force-import a fresh ``dynacell.evaluation.pipeline`` module. + + Other tests (e.g. ``test_evaluation_pipeline.py``) stub + ``dynacell.evaluation.metrics`` via ``monkeypatch.setitem(sys.modules, + ...)``. After teardown ``sys.modules`` may still hold the stubbed + modules AND the pipeline module's bound ``compute_pixel_metrics`` + reference. Clear every ``dynacell.evaluation.*`` cache entry before + re-import so we get fresh modules with real implementations. + """ + for name in list(sys.modules): + if name.startswith("dynacell.evaluation"): + sys.modules.pop(name, None) + return importlib.import_module("dynacell.evaluation.pipeline") + + +def make_hcs_plate(path: Path, channel_name: str, seed: int) -> None: + """Create a tiny HCS plate with ``N_POSITIONS`` positions of shape ``(T, 1, D, H, W)``. + + Positions are named ``A/1/{i}`` for ``i in range(N_POSITIONS)``. Pixel + values are deterministic from ``seed`` so pred and GT differ + predictably. + """ + rng = np.random.default_rng(seed) + with open_ome_zarr(path, mode="w", layout="hcs", channel_names=[channel_name], version="0.5") as plate: + for i in range(N_POSITIONS): + pos = plate.create_position("A", "1", str(i)) + data = rng.uniform(0.0, 1.0, size=(T, 1, D, H, W)).astype(np.float32) + pos.create_image("0", data) + + +def make_mask_cache( + cache_dir: Path, + plate_path: Path, + channel_name: str, + side: str, + target_name: str = "er", +) -> None: + """Prebuild a mask cache (zarr + manifest) for ``N_POSITIONS`` positions. + + Mask values are deterministic per (side, position, t): GT masks fill + the upper half of the (H, W) plane; pred masks fill the upper + quadrant. This makes DICE / IoU deterministic and non-trivially + different between sides. ``side`` is ``"gt"`` or ``"pred"``. + """ + from dynacell.evaluation.cache import ( + CACHE_SCHEMA_VERSION, + cache_paths, + save_manifest, + write_mask, + ) + + paths = cache_paths(cache_dir) + mask_channel = "target_seg" if side == "gt" else "prediction_seg" + + pos_names = [f"A/1/{i}" for i in range(N_POSITIONS)] + for pos_name in pos_names: + masks = np.zeros((T, D, H, W), dtype=bool) + if side == "gt": + masks[:, :, : H // 2, :] = True + else: + masks[:, :, : H // 2, : W // 2] = True + write_mask(paths, target_name, pos_name, masks, channel_name=mask_channel) + + source_tag: dict[str, str] = {} if side == "gt" else {"source": "prediction"} + now = datetime.now(UTC).isoformat() + manifest = { + "cache_schema_version": CACHE_SCHEMA_VERSION, + "gt" if side == "gt" else "pred": { + "plate_path": str(plate_path), + "channel_name": channel_name, + }, + "artifacts": { + "organelle_masks": { + target_name: { + "target_name": target_name, + "path": f"organelle_masks/{target_name}.zarr", + "built_at": now, + "positions": pos_names, + **source_tag, + } + }, + }, + } + save_manifest(paths, manifest) + + +def build_eval_config( + pred_path: Path, + gt_path: Path, + gt_cache_dir: Path, + pred_cache_dir: Path, + save_dir: Path, + *, + executor: str, + fov_workers: int, +): + """Compose a minimal cache-only eval config.""" + return OmegaConf.create( + { + "target_name": "er", + "io": { + "pred_path": str(pred_path), + "gt_path": str(gt_path), + "pred_channel_name": "prediction", + "gt_channel_name": "target", + "cell_segmentation_path": None, + "gt_cache_dir": str(gt_cache_dir), + "pred_cache_dir": str(pred_cache_dir), + "require_complete_cache": True, + }, + "pixel_metrics": { + "spacing": [1.0, 1.0, 1.0], + "fsc": None, + "spectral_pcc": None, + }, + "feature_metrics": { + "patch_size": 64, + "deep_feature_batch_threshold": 256, + }, + "use_gpu": False, + "compute_microssim": False, + "compute_feature_metrics": False, + "limit_positions": None, + "force_recompute": { + "all": False, + "gt_masks": False, + "gt_cp": False, + "gt_dinov3": False, + "gt_dynaclr": False, + "gt_celldino": False, + "pred_masks": False, + "pred_cp": False, + "pred_dinov3": False, + "pred_dynaclr": False, + "pred_celldino": False, + "final_metrics": False, + }, + "save": { + "save_dir": str(save_dir), + "pixel_csv_filename": "pixel_metrics.csv", + "pixel_metrics_filename": "pixel_metrics.npy", + "mask_csv_filename": "mask_metrics.csv", + "mask_metrics_filename": "mask_metrics.npy", + "feature_csv_filename": "feature_metrics.csv", + "feature_metrics_filename": "feature_metrics.npy", + }, + "runtime": { + "fov_workers": fov_workers, + "threads_per_worker": "auto", + "executor": executor, + "cuda_empty_cache_every_n_timepoints": 0, + "gc_collect_every_n_fovs": 0, + }, + } + ) + + +def build_fixture(root: Path): + """Build matched pred + GT plates and mask caches under ``root``. + + Returns ``(pred_path, gt_path, gt_cache_dir, pred_cache_dir)``. + """ + pred_path = root / "pred.zarr" + gt_path = root / "gt.zarr" + gt_cache_dir = root / "gt_cache" + pred_cache_dir = root / "pred_cache" + make_hcs_plate(pred_path, "prediction", seed=42) + make_hcs_plate(gt_path, "target", seed=7) + make_mask_cache(gt_cache_dir, gt_path, "target", side="gt") + make_mask_cache(pred_cache_dir, pred_path, "prediction", side="pred") + return pred_path, gt_path, gt_cache_dir, pred_cache_dir + + +def read_position_arrays(plate_path: Path) -> dict[str, np.ndarray]: + """Return ``{pos_name: data}`` for every position in an HCS plate.""" + out: dict[str, np.ndarray] = {} + with open_ome_zarr(plate_path, mode="r") as plate: + for name, pos in plate.positions(): + out[name] = np.asarray(pos.data[:]) + return out diff --git a/applications/dynacell/tests/test_evaluation_grouped.py b/applications/dynacell/tests/test_evaluation_grouped.py new file mode 100644 index 000000000..92473a325 --- /dev/null +++ b/applications/dynacell/tests/test_evaluation_grouped.py @@ -0,0 +1,306 @@ +"""Parity test for the grouped multi-condition eval driver. + +Builds two independent fixtures (each its own pred + GT plates + mask +caches), runs them once via ``evaluate_predictions_grouped`` and once +via two back-to-back ``evaluate_predictions`` calls, and asserts the +outputs are byte-equal per condition. + +Cache-only design (same shape as +``test_evaluation_pipeline_parallel_cpu.py``): ``target_name=er`` + +``io.require_complete_cache=true`` + ``compute_feature_metrics=false`` +means no segmenter / extractors / cubic are loaded — the test +validates the **loop structure** of the grouped driver, not the +model-sharing speedup (which is hermetic-untestable). + +Fixture-building helpers live in ``_eval_fixtures.py``. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +from omegaconf import OmegaConf + +from ._eval_fixtures import ( + N_POSITIONS, + T, + build_eval_config, + build_fixture, + live_pipeline_module, + read_position_arrays, +) + + +def _build_grouped_config( + cond_a_root: Path, + cond_b_root: Path, + save_root: Path, +): + """Construct the grouped config with two condition overlays.""" + pred_a, gt_a, gt_cache_a, pred_cache_a = build_fixture(cond_a_root) + pred_b, gt_b, gt_cache_b, pred_cache_b = build_fixture(cond_b_root) + + base = build_eval_config( + pred_a, + gt_a, + gt_cache_a, + pred_cache_a, + save_root / "_unused", + executor="serial", + fov_workers=1, + ) + # Strip the placeholder io/save — conditions own those entirely. + base["io"]["pred_path"] = None + base["io"]["gt_path"] = None + base["io"]["gt_cache_dir"] = None + base["io"]["pred_cache_dir"] = None + base["save"]["save_dir"] = None + + base["conditions"] = [ + { + "name": "cond_a", + "io": { + "pred_path": str(pred_a), + "gt_path": str(gt_a), + "gt_cache_dir": str(gt_cache_a), + "pred_cache_dir": str(pred_cache_a), + }, + "save": {"save_dir": str(save_root / "cond_a_grouped")}, + }, + { + "name": "cond_b", + "io": { + "pred_path": str(pred_b), + "gt_path": str(gt_b), + "gt_cache_dir": str(gt_cache_b), + "pred_cache_dir": str(pred_cache_b), + }, + "save": {"save_dir": str(save_root / "cond_b_grouped")}, + }, + ] + return base + + +def test_grouped_matches_sequential_per_condition(tmp_path: Path): + """Grouped eval produces byte-equal outputs to sequential per-condition runs.""" + cond_a_root = tmp_path / "fixture_a" + cond_b_root = tmp_path / "fixture_b" + cond_a_root.mkdir() + cond_b_root.mkdir() + save_root = tmp_path / "saves" + save_root.mkdir() + + pred_a, gt_a, gt_cache_a, pred_cache_a = build_fixture(cond_a_root) + pred_b, gt_b, gt_cache_b, pred_cache_b = build_fixture(cond_b_root) + save_a_seq = save_root / "cond_a_seq" + save_b_seq = save_root / "cond_b_seq" + save_a_seq.mkdir() + save_b_seq.mkdir() + + pipeline = live_pipeline_module() + cfg_a_seq = build_eval_config(pred_a, gt_a, gt_cache_a, pred_cache_a, save_a_seq, executor="serial", fov_workers=1) + cfg_b_seq = build_eval_config(pred_b, gt_b, gt_cache_b, pred_cache_b, save_b_seq, executor="serial", fov_workers=1) + pixel_a_seq, mask_a_seq, _ = pipeline.evaluate_predictions(cfg_a_seq) + pixel_b_seq, mask_b_seq, _ = pipeline.evaluate_predictions(cfg_b_seq) + + grouped_cfg = _build_grouped_config(cond_a_root, cond_b_root, save_root) + pipeline = live_pipeline_module() + results = pipeline.evaluate_predictions_grouped(grouped_cfg) + + assert [name for name, _ in results] == ["cond_a", "cond_b"] + (_, (pixel_a_grp, mask_a_grp, _)), (_, (pixel_b_grp, mask_b_grp, _)) = results + + def _sort_key(row: dict): + return (row["FOV"], row["Timepoint"]) + + assert sorted(pixel_a_grp, key=_sort_key) == sorted(pixel_a_seq, key=_sort_key) + assert sorted(mask_a_grp, key=_sort_key) == sorted(mask_a_seq, key=_sort_key) + assert sorted(pixel_b_grp, key=_sort_key) == sorted(pixel_b_seq, key=_sort_key) + assert sorted(mask_b_grp, key=_sort_key) == sorted(mask_b_seq, key=_sort_key) + + for cond_name, save_seq, save_grp in [ + ("cond_a", save_a_seq, save_root / "cond_a_grouped"), + ("cond_b", save_b_seq, save_root / "cond_b_grouped"), + ]: + arrs_seq = read_position_arrays(save_seq / "segmentation_results.zarr") + arrs_grp = read_position_arrays(save_grp / "segmentation_results.zarr") + assert set(arrs_seq.keys()) == set(arrs_grp.keys()), f"{cond_name} position set mismatch" + for pos_name, arr_seq in arrs_seq.items(): + np.testing.assert_array_equal(arr_seq, arrs_grp[pos_name], err_msg=f"{cond_name}/{pos_name}") + + assert len(pixel_a_grp) == N_POSITIONS * T + assert len(pixel_b_grp) == N_POSITIONS * T + + +def test_grouped_rejects_model_loading_field_overrides(tmp_path: Path): + """Per-condition overrides on model-loading fields must raise — including condition 0.""" + cond_a_root = tmp_path / "fixture_a" + cond_b_root = tmp_path / "fixture_b" + cond_a_root.mkdir() + cond_b_root.mkdir() + save_root = tmp_path / "saves" + save_root.mkdir() + + grouped_cfg = _build_grouped_config(cond_a_root, cond_b_root, save_root) + # Sabotage condition B's target_name. + grouped_cfg["conditions"][1]["target_name"] = "membrane" + + pipeline = live_pipeline_module() + with pytest.raises(ValueError, match="model-loading field"): + pipeline.evaluate_predictions_grouped(grouped_cfg) + + +def test_grouped_rejects_condition0_model_field_override(tmp_path: Path): + """A model-loading override smuggled into condition 0 must raise, not silently re-baseline. + + Earlier wiring used ``conditions[0]``-merged config as the invariant + baseline, which would have made condition 0 implicitly trusted to + redefine ``target_name`` / ``feature_extractor.*`` etc. Guards + against regression of that asymmetry. + """ + cond_a_root = tmp_path / "fixture_a" + cond_b_root = tmp_path / "fixture_b" + cond_a_root.mkdir() + cond_b_root.mkdir() + save_root = tmp_path / "saves" + save_root.mkdir() + + grouped_cfg = _build_grouped_config(cond_a_root, cond_b_root, save_root) + grouped_cfg["conditions"][0]["target_name"] = "membrane" + + pipeline = live_pipeline_module() + with pytest.raises(ValueError, match="model-loading field"): + pipeline.evaluate_predictions_grouped(grouped_cfg) + + +def test_grouped_warns_loudly_on_process_plus_cache_miss(tmp_path: Path, capsys): + """``executor=process`` + multiple conditions + cache-miss prints a loud WARNING. + + This combination is the worst case for amortization: each + condition's worker pool independently re-loads SuperModel + the + three deep extractors, so cold-start cost multiplies by + ``N_workers × N_conditions``. The driver should warn the user + upfront and recommend ``executor=serial``. + """ + cond_a_root = tmp_path / "fixture_a" + cond_b_root = tmp_path / "fixture_b" + cond_a_root.mkdir() + cond_b_root.mkdir() + save_root = tmp_path / "saves" + save_root.mkdir() + + grouped_cfg = _build_grouped_config(cond_a_root, cond_b_root, save_root) + grouped_cfg["runtime"]["executor"] = "process" + grouped_cfg["runtime"]["fov_workers"] = 2 + grouped_cfg["io"]["require_complete_cache"] = False + + # The grouped driver prints the warning before iterating conditions. + # The eval will then fail/succeed depending on cache state; we don't + # care — we just want the warning text in stdout. + pipeline = live_pipeline_module() + try: + pipeline.evaluate_predictions_grouped(grouped_cfg) + except Exception: + pass + + out = capsys.readouterr().out + assert "WARNING" in out, f"expected loud warning, got: {out!r}" + assert "executor=process" in out + assert "require_complete_cache=false" in out + assert "executor=serial" in out + + +def test_grouped_mild_note_on_process_plus_cache_only(tmp_path: Path, capsys): + """``executor=process`` + ``require_complete_cache=true`` gets a mild note, not the loud WARNING. + + Under the cache-only path, workers skip ``prepare_segmentation_model`` + and don't instantiate extractors — pool re-spawn is the only + overhead, which is small. The driver should inform but not alarm. + """ + cond_a_root = tmp_path / "fixture_a" + cond_b_root = tmp_path / "fixture_b" + cond_a_root.mkdir() + cond_b_root.mkdir() + save_root = tmp_path / "saves" + save_root.mkdir() + + grouped_cfg = _build_grouped_config(cond_a_root, cond_b_root, save_root) + grouped_cfg["runtime"]["executor"] = "process" + grouped_cfg["runtime"]["fov_workers"] = 2 + # require_complete_cache stays True from the cache-only fixture. + + pipeline = live_pipeline_module() + pipeline.evaluate_predictions_grouped(grouped_cfg) + out = capsys.readouterr().out + assert "[grouped] note:" in out + assert "no models actually load" in out + assert "WARNING" not in out, f"unexpected loud warning under cache-only path: {out!r}" + + +def test_grouped_rejects_require_complete_cache_override(tmp_path: Path): + """``io.require_complete_cache`` is a grouped invariant — overrides must raise. + + The base config's ``io.require_complete_cache`` determines whether + ``load_eval_models`` instantiates a real ``SuperModel`` or returns + ``None``. Letting a condition flip this would mean the shared models + bundle disagrees with the condition's actual needs (None when + cache-miss expects a real segmenter, or loaded but never used). Guard + by treating it as a model-loading invariant. + """ + cond_a_root = tmp_path / "fixture_a" + cond_b_root = tmp_path / "fixture_b" + cond_a_root.mkdir() + cond_b_root.mkdir() + save_root = tmp_path / "saves" + save_root.mkdir() + + grouped_cfg = _build_grouped_config(cond_a_root, cond_b_root, save_root) + grouped_cfg["conditions"][1]["io"]["require_complete_cache"] = False + + pipeline = live_pipeline_module() + with pytest.raises(ValueError, match="require_complete_cache"): + pipeline.evaluate_predictions_grouped(grouped_cfg) + + +def test_grouped_works_on_struct_mode_config(tmp_path: Path): + """Grouped driver must work on Hydra-composed (struct-mode) configs. + + Real ``dynacell evaluate-grouped`` invocations go through + ``@hydra.main`` which always produces a struct-mode ``DictConfig``. + ``OmegaConf.merge`` propagates struct mode, so merging an overlay + carrying a ``name`` label (outside the schema) raises + ``ConfigKeyError``, and stripping ``conditions`` / ``name`` with + ``del`` raises ``ConfigTypeError``. The driver must escape struct + mode internally; this test guards against regression. + """ + cond_a_root = tmp_path / "fixture_a" + cond_b_root = tmp_path / "fixture_b" + cond_a_root.mkdir() + cond_b_root.mkdir() + save_root = tmp_path / "saves" + save_root.mkdir() + + grouped_cfg = _build_grouped_config(cond_a_root, cond_b_root, save_root) + OmegaConf.set_struct(grouped_cfg, True) + + pipeline = live_pipeline_module() + results = pipeline.evaluate_predictions_grouped(grouped_cfg) + assert [name for name, _ in results] == ["cond_a", "cond_b"] + + +def test_grouped_rejects_empty_conditions(tmp_path: Path): + """Missing or empty 'conditions' list must raise an informative error.""" + cond_a_root = tmp_path / "fixture_a" + cond_b_root = tmp_path / "fixture_b" + cond_a_root.mkdir() + cond_b_root.mkdir() + save_root = tmp_path / "saves" + save_root.mkdir() + grouped_cfg = _build_grouped_config(cond_a_root, cond_b_root, save_root) + grouped_cfg["conditions"] = [] + + pipeline = live_pipeline_module() + with pytest.raises(ValueError, match="non-empty"): + pipeline.evaluate_predictions_grouped(grouped_cfg) diff --git a/applications/dynacell/tests/test_evaluation_pipeline_parallel_cpu.py b/applications/dynacell/tests/test_evaluation_pipeline_parallel_cpu.py index 692ffa5ba..2386be79c 100644 --- a/applications/dynacell/tests/test_evaluation_pipeline_parallel_cpu.py +++ b/applications/dynacell/tests/test_evaluation_pipeline_parallel_cpu.py @@ -21,208 +21,25 @@ SuperModel, DinoV3, DynaCLR, CellDino, or cubic — exercises spawn boot, DictConfig pickling, ``FovResult`` transport, parent-side HCS plate write, and the dataset-level aggregation contract. + +Fixture-building helpers live in ``_eval_fixtures.py`` (shared with +``test_evaluation_grouped.py``). """ from __future__ import annotations -import importlib -import sys -from datetime import UTC, datetime from pathlib import Path import numpy as np -import pytest # noqa: F401 -- imported for future @pytest.mark uses -from iohub.ngff import open_ome_zarr -from omegaconf import OmegaConf - - -def _live_pipeline_module(): - """Force-import a fresh ``dynacell.evaluation.pipeline`` module. - - Other tests (``test_evaluation_pipeline.py``) stub - ``dynacell.evaluation.metrics`` etc. via ``monkeypatch.setitem(sys.modules, - ...)``. After teardown, ``sys.modules`` may still hold the stubbed modules - AND the pipeline module's bound ``compute_pixel_metrics`` reference. Clear - every ``dynacell.evaluation.*`` cache entry before re-import so we get - fresh modules with real implementations. - """ - for name in list(sys.modules): - if name.startswith("dynacell.evaluation"): - sys.modules.pop(name, None) - return importlib.import_module("dynacell.evaluation.pipeline") - - -# Fixture dimensions — kept tiny so the test runs in CI (single-CPU) under ~30s. -N_POSITIONS = 3 -T = 2 -D = 4 -H = 32 -W = 32 - - -def _make_hcs_plate(path: Path, channel_name: str, seed: int) -> None: - """Create a tiny HCS plate with ``N_POSITIONS`` positions of shape ``(T, 1, D, H, W)``. - - Positions are named ``A/1/{i}`` for ``i in range(N_POSITIONS)``. Pixel - values are deterministic from ``seed`` so pred and GT differ predictably. - """ - rng = np.random.default_rng(seed) - with open_ome_zarr(path, mode="w", layout="hcs", channel_names=[channel_name], version="0.5") as plate: - for i in range(N_POSITIONS): - pos = plate.create_position("A", "1", str(i)) - data = rng.uniform(0.0, 1.0, size=(T, 1, D, H, W)).astype(np.float32) - pos.create_image("0", data) - - -def _make_mask_cache( - cache_dir: Path, - plate_path: Path, - channel_name: str, - side: str, # "gt" or "pred" - target_name: str = "er", -) -> None: - """Prebuild a mask cache (zarr + manifest) for ``N_POSITIONS`` positions. - - The mask values themselves are deterministic per (side, position, t): - GT masks are TRUE in the upper half of the (H, W) plane; pred masks - are TRUE in the upper quadrant. This guarantees the DICE / IoU - metrics are deterministic and non-trivially different between sides. - """ - from dynacell.evaluation.cache import ( - CACHE_SCHEMA_VERSION, - cache_paths, - save_manifest, - write_mask, - ) - - paths = cache_paths(cache_dir) - mask_channel = "target_seg" if side == "gt" else "prediction_seg" - - pos_names = [f"A/1/{i}" for i in range(N_POSITIONS)] - for pos_name in pos_names: - masks = np.zeros((T, D, H, W), dtype=bool) - if side == "gt": - masks[:, :, : H // 2, :] = True # upper half - else: - masks[:, :, : H // 2, : W // 2] = True # upper quadrant - write_mask(paths, target_name, pos_name, masks, channel_name=mask_channel) - - # Manifest entries must satisfy _validate_artifact_params: the - # masks_section.get(target_name) must equal {"target_name": ..., **source_tag}. - source_tag: dict[str, str] = {} if side == "gt" else {"source": "prediction"} - now = datetime.now(UTC).isoformat() - manifest = { - "cache_schema_version": CACHE_SCHEMA_VERSION, - "gt" if side == "gt" else "pred": { - "plate_path": str(plate_path), - "channel_name": channel_name, - }, - "artifacts": { - "organelle_masks": { - target_name: { - "target_name": target_name, - "path": f"organelle_masks/{target_name}.zarr", - "built_at": now, - "positions": pos_names, - **source_tag, - } - }, - }, - } - save_manifest(paths, manifest) - - -def _build_eval_config( - pred_path: Path, - gt_path: Path, - gt_cache_dir: Path, - pred_cache_dir: Path, - save_dir: Path, - *, - executor: str, - fov_workers: int, -): - """Compose a minimal eval config for the cache-only path.""" - return OmegaConf.create( - { - "target_name": "er", - "io": { - "pred_path": str(pred_path), - "gt_path": str(gt_path), - "pred_channel_name": "prediction", - "gt_channel_name": "target", - "cell_segmentation_path": None, - "gt_cache_dir": str(gt_cache_dir), - "pred_cache_dir": str(pred_cache_dir), - "require_complete_cache": True, - }, - "pixel_metrics": { - "spacing": [1.0, 1.0, 1.0], - "fsc": None, - "spectral_pcc": None, - }, - "feature_metrics": { - "patch_size": 64, - "deep_feature_batch_threshold": 256, - }, - "use_gpu": False, - "compute_microssim": False, - "compute_feature_metrics": False, - "limit_positions": None, - "force_recompute": { - "all": False, - "gt_masks": False, - "gt_cp": False, - "gt_dinov3": False, - "gt_dynaclr": False, - "gt_celldino": False, - "pred_masks": False, - "pred_cp": False, - "pred_dinov3": False, - "pred_dynaclr": False, - "pred_celldino": False, - "final_metrics": False, - }, - "save": { - "save_dir": str(save_dir), - "pixel_csv_filename": "pixel_metrics.csv", - "pixel_metrics_filename": "pixel_metrics.npy", - "mask_csv_filename": "mask_metrics.csv", - "mask_metrics_filename": "mask_metrics.npy", - "feature_csv_filename": "feature_metrics.csv", - "feature_metrics_filename": "feature_metrics.npy", - }, - "runtime": { - "fov_workers": fov_workers, - "threads_per_worker": "auto", - "executor": executor, - "cuda_empty_cache_every_n_timepoints": 0, - "gc_collect_every_n_fovs": 0, - }, - } - ) - - -def _build_fixture(root: Path): - """Build pred + GT plates and matching mask caches under ``root``.""" - pred_path = root / "pred.zarr" - gt_path = root / "gt.zarr" - gt_cache_dir = root / "gt_cache" - pred_cache_dir = root / "pred_cache" - _make_hcs_plate(pred_path, "prediction", seed=42) - _make_hcs_plate(gt_path, "target", seed=7) - _make_mask_cache(gt_cache_dir, gt_path, "target", side="gt") - _make_mask_cache(pred_cache_dir, pred_path, "prediction", side="pred") - return pred_path, gt_path, gt_cache_dir, pred_cache_dir - -def _read_position_arrays(plate_path: Path) -> dict[str, np.ndarray]: - """Return ``{pos_name: data}`` for every position in an HCS plate.""" - out: dict[str, np.ndarray] = {} - with open_ome_zarr(plate_path, mode="r") as plate: - for name, pos in plate.positions(): - out[name] = np.asarray(pos.data[:]) - return out +from ._eval_fixtures import ( + N_POSITIONS, + T, + build_eval_config, + build_fixture, + live_pipeline_module, + read_position_arrays, +) def test_serial_vs_process_evaluate_predictions_parity(tmp_path: Path): @@ -234,17 +51,17 @@ def test_serial_vs_process_evaluate_predictions_parity(tmp_path: Path): """ fixture_root = tmp_path / "fixture" fixture_root.mkdir() - pred_path, gt_path, gt_cache_dir, pred_cache_dir = _build_fixture(fixture_root) + pred_path, gt_path, gt_cache_dir, pred_cache_dir = build_fixture(fixture_root) save_dir_serial = tmp_path / "serial" save_dir_process = tmp_path / "process" # Re-import inside the test body — test_lazy_init.py pops dynacell.* # modules from sys.modules, which would invalidate a module-level import. - pipeline = _live_pipeline_module() + pipeline = live_pipeline_module() evaluate_predictions = pipeline.evaluate_predictions - cfg_serial = _build_eval_config( + cfg_serial = build_eval_config( pred_path, gt_path, gt_cache_dir, @@ -253,7 +70,7 @@ def test_serial_vs_process_evaluate_predictions_parity(tmp_path: Path): executor="serial", fov_workers=1, ) - cfg_process = _build_eval_config( + cfg_process = build_eval_config( pred_path, gt_path, gt_cache_dir, @@ -269,17 +86,13 @@ def test_serial_vs_process_evaluate_predictions_parity(tmp_path: Path): pixel_serial, mask_serial, _ = evaluate_predictions(cfg_serial) pixel_process, mask_process, _ = evaluate_predictions(cfg_process) - # Pixel + mask rows: same set after sort. We sort by (FOV, Timepoint) - # because the process branch already sorts by pos_name and within-FOV - # iteration is deterministic, but we sort defensively. def _sort_key(row: dict): return (row["FOV"], row["Timepoint"]) assert sorted(pixel_serial, key=_sort_key) == sorted(pixel_process, key=_sort_key) assert sorted(mask_serial, key=_sort_key) == sorted(mask_process, key=_sort_key) - # Per-position seg plate equality (both channels). - arrs_serial = _read_position_arrays(save_dir_serial / "segmentation_results.zarr") - arrs_process = _read_position_arrays(save_dir_process / "segmentation_results.zarr") + arrs_serial = read_position_arrays(save_dir_serial / "segmentation_results.zarr") + arrs_process = read_position_arrays(save_dir_process / "segmentation_results.zarr") assert set(arrs_serial.keys()) == set(arrs_process.keys()) for pos_name, arr_a in arrs_serial.items(): arr_b = arrs_process[pos_name] @@ -296,11 +109,11 @@ def test_serial_path_runs_without_segmenter_loaded(tmp_path: Path): """ fixture_root = tmp_path / "fixture" fixture_root.mkdir() - pred_path, gt_path, gt_cache_dir, pred_cache_dir = _build_fixture(fixture_root) + pred_path, gt_path, gt_cache_dir, pred_cache_dir = build_fixture(fixture_root) save_dir = tmp_path / "out" save_dir.mkdir() - cfg = _build_eval_config( + cfg = build_eval_config( pred_path, gt_path, gt_cache_dir, @@ -309,10 +122,7 @@ def test_serial_path_runs_without_segmenter_loaded(tmp_path: Path): executor="serial", fov_workers=1, ) - pipeline = _live_pipeline_module() + pipeline = live_pipeline_module() pixel_rows, mask_rows, _ = pipeline.evaluate_predictions(cfg) assert len(pixel_rows) == N_POSITIONS * T assert len(mask_rows) == N_POSITIONS * T - # Confirms the no-segmenter contract: prepare_segmentation_model - # returned None and evaluate_predictions still produced valid outputs - # for every (FOV, t) pair. From 582917537da4dd10d3231d21851cc1515e54c473 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 15:53:38 -0700 Subject: [PATCH 17/20] fix(dynacell): keep SuperModel under require_complete_cache when pred cache off MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `prepare_segmentation_model` was returning None unconditionally whenever `io.require_complete_cache=true`, on the assumption that all masks would be served from cache. That assumption breaks for nucleus/membrane targets when `io.pred_cache_dir=None`: the per-T loop in `_process_one_fov` falls back to `segment(predict[t], seg_model=seg_model)`, which raises `ValueError` because `seg_model` is None. Narrow the early return to cases where no path can reach `segment()` without a cached mask — organelle targets (workflow-based, no `seg_model` needed) or nucleus/membrane runs that also have `io.pred_cache_dir` configured. Reported by Copilot on PR #426. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../src/dynacell/evaluation/segmentation.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/applications/dynacell/src/dynacell/evaluation/segmentation.py b/applications/dynacell/src/dynacell/evaluation/segmentation.py index 55205921f..bd9e90d5e 100644 --- a/applications/dynacell/src/dynacell/evaluation/segmentation.py +++ b/applications/dynacell/src/dynacell/evaluation/segmentation.py @@ -92,11 +92,16 @@ def prepare_segmentation_model(config): Returns ``None`` for organelles that use classical (non-DL) workflows. Respects ``config.use_gpu`` when deciding whether to move models to GPU. - Returns ``None`` unconditionally when ``io.require_complete_cache=true``: - in that mode the caller asserts that all needed masks live in the cache, - so loading SuperModel (which hits a quilt download for nucleus/membrane - on a cold node) would be wasted work. If a mask cache-miss happens, the - cache layer raises ``StaleCacheError`` before any ``segment()`` call. + Returns ``None`` when ``io.require_complete_cache=true`` AND no per-FOV + call path can invoke ``segment()`` without a cached mask. Concretely: + organelle targets (``er``/``nucleoli``/``lysosomes``/``mitochondria``) + delegate to aicssegmentation workflows that don't take ``seg_model``, + so skipping the SuperModel load is safe. For ``nucleus``/``membrane``, + the per-T loop in ``pipeline._process_one_fov`` falls back to + ``segment(predict[t], seg_model=seg_model)`` whenever the pred-side + mask cache is disabled (``io.pred_cache_dir=None``); in that case we + must still load SuperModel even under ``require_complete_cache=true`` + or the fallback raises ``ValueError`` mid-loop. """ require_complete = bool(getattr(config.io, "require_complete_cache", False)) if config.target_name not in [ @@ -109,7 +114,9 @@ def prepare_segmentation_model(config): ]: raise ValueError(f"Invalid target_name in config: {config.target_name!r}") if require_complete: - return None + pred_cache_dir = getattr(config.io, "pred_cache_dir", None) + if config.target_name not in ("nucleus", "membrane") or pred_cache_dir is not None: + return None if config.target_name in ["nucleus", "membrane"]: _require_segmenter_model_zoo() if config.target_name == "nucleus": From bc399eb2963fec60f6ac33f7973fdfbb84c697cb Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 15:53:48 -0700 Subject: [PATCH 18/20] docs(dynacell): align gpu_serialization_lock docstring with implementation The docstring claimed "no CUDA available in the worker" was one of the no-op conditions, but the implementation only checks `gate` and `_GPU_LOCK_PATH`. On a CPU-only node with `use_gpu=true` the lock would still serialize across workers. Rather than add a runtime CUDA probe (callers already pass `gate=use_gpu` to opt out for CPU-only regions), correct the docstring to describe what the code actually does. Reported by Copilot on PR #426. Co-Authored-By: Claude Opus 4.6 (1M context) --- applications/dynacell/src/dynacell/evaluation/runtime.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/applications/dynacell/src/dynacell/evaluation/runtime.py b/applications/dynacell/src/dynacell/evaluation/runtime.py index bfe74532d..7abeea0dd 100644 --- a/applications/dynacell/src/dynacell/evaluation/runtime.py +++ b/applications/dynacell/src/dynacell/evaluation/runtime.py @@ -440,9 +440,11 @@ def gpu_serialization_lock(gate: bool = True) -> Iterator[None]: usage tracks the ``use_gpu`` flag, so CPU-only runs don't serialize across workers needlessly. - Other no-op conditions: ``_GPU_LOCK_PATH`` unset (parent / serial - mode) or no CUDA available in the worker. Used to wrap model loads - and per-FOV GPU operations under ``executor=process``. + Other no-op condition: ``_GPU_LOCK_PATH`` unset (parent / serial + mode). Used to wrap model loads and per-FOV GPU operations under + ``executor=process``. CPU-only execution is opted into at the call + site by passing ``gate=False`` (typically ``gate=config.use_gpu``); + the lock itself does NOT probe ``torch.cuda.is_available()``. """ if not gate or _GPU_LOCK_PATH is None: yield From 7b8e1378ac9b48a31b6d1f8d59f913ffdbcc26f4 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Wed, 20 May 2026 16:18:38 -0700 Subject: [PATCH 19/20] =?UTF-8?q?docs(dynacell):=20rewrite=20eval=20README?= =?UTF-8?q?=20=E2=80=94=20runtime/grouped/cache-only=20+=20trim=20bloat?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The eval README had grown stale relative to the runtime/grouped/cache work landed on this branch and PR #427: no mention of `io.require_complete_cache`'s cache-only fast path, no `runtime.*` block, no `dynacell evaluate-grouped`, no CELL-DINO. This rewrite folds those in and restructures: FOV parallelism + grouped eval move out of "Artifact caches" (where they were mis-classified) into a new top-level "Scaling" section; a new "Quick start" sits above Configuration; HF-cache operational details go to a new "HPC notes" section at the end. Bloat trim along the way: components table compressed (related modules combined, trivial files pushed to a note), HF-cache + external-users sections shortened, the 11-row force_recompute table collapsed to a YAML block plus 5 grouped descriptions, redundant Hydra group-discovery paragraphs removed. No flags, footguns, or CLI options dropped — every config knob and env var from the old version is still documented. 466 → 265 lines. Public-repo hygiene: replaced the hard-coded `/hpc/projects/comp.micro/...` default for the team-shared HF cache with a description that points at the `_DEFAULT_SHARED_HF_CACHE` constant in `dynacell/__main__.py` and the `DYNACELL_SHARED_HF_CACHE` env-var override. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../src/dynacell/evaluation/README.md | 418 +++++++----------- 1 file changed, 150 insertions(+), 268 deletions(-) diff --git a/applications/dynacell/src/dynacell/evaluation/README.md b/applications/dynacell/src/dynacell/evaluation/README.md index a817836e4..5ef5e174e 100644 --- a/applications/dynacell/src/dynacell/evaluation/README.md +++ b/applications/dynacell/src/dynacell/evaluation/README.md @@ -6,183 +6,71 @@ End-to-end evaluation pipeline for virtual staining predictions against fluoresc | Module | Purpose | |---|---| -| `pipeline.py` | Hydra-driven orchestrator. Loads prediction/GT OME-Zarr plates, computes per-FOV per-timepoint metrics, saves CSVs + NPYs + plots. CLI entrypoint: `dynacell evaluate`. | -| `metrics.py` | Pixel metrics (PCC, SSIM, NRMSE, PSNR, FSC resolution, spectral PCC, MicroMS3IM), mask metrics (Dice, IoU, precision, recall, accuracy, TP/FP/FN/TN), feature metrics split into `*_target_*` / `*_pred_*` / `*_pairwise` so GT-side work can be cached separately from predictions. | -| `segmentation.py` | Organelle-specific classical-CV segmentation via `aicssegmentation` workflows (`nucleus`, `membrane`, `nucleoli`, `lysosomes`, `er`, `mitochondria`). Used for mask metrics. | -| `cache.py` | Artifact cache: on-disk layout, manifest I/O, read/write helpers, staleness check. Keyed by source plate/channel plus `cell_segmentation_path` where relevant. | -| `pipeline_cache.py` | Per-FOV load-or-compute wrappers for GT and prediction masks/features. Honor `force_recompute.*` flags and the `io.require_complete_cache` contract. | -| `precompute_cli.py` | Hydra entrypoint for `dynacell precompute-gt`. Iterates GT positions and fills the cache; no eval loop. | -| `utils.py` | `DinoV3FeatureExtractor`, `DynaCLRFeatureExtractor`, pairwise feature-similarity helpers, `plot_metrics()` bar/violin plots. | -| `io.py` | OME-Zarr / tiff readers and writers, prediction preprocessing transforms. | -| `torch_ssim.py` | GPU-friendly PyTorch SSIM. | -| `formatting.py` | Metric table formatting helpers. | -| `spectral_pcc/` | Bandlimited spectral PCC diagnostics and bead simulations. | -| `_configs/eval.yaml` | Hydra config for `dynacell evaluate`, with `???` MISSING markers for dataset-specific fields. | -| `_configs/precompute.yaml` | Hydra config for `dynacell precompute-gt`; inherits eval, requires `io.gt_cache_dir`. | +| `pipeline.py` | Hydra orchestrator. CLIs: `dynacell evaluate` (single-condition) and `dynacell evaluate-grouped` (one model load, N I/O conditions). | +| `metrics.py` | Pixel, mask, and feature metrics (CP regionprops + DINOv3 + DynaCLR + CELL-DINO), computed symmetrically for GT/predictions and combined pairwise. | +| `segmentation.py` | `aicssegmentation` workflows + SuperModel for `nucleus`/`membrane`. | +| `cache.py`, `pipeline_cache.py` | Artifact cache: on-disk layout, manifest, identity check, per-FOV load-or-compute wrappers, batched `precompute_deep_features`. | +| `model_loader.py` | Shared `load_eval_models(config, flags=...)` returning an `EvalModels` bundle. Used by both `evaluate` and `precompute-gt`. | +| `runtime.py` | BLAS/OMP thread caps, `ProcessPoolExecutor` worker initializer, `gpu_serialization_lock`, region timers. | +| `precompute_cli.py` | `dynacell precompute-gt` — fills the GT cache without running the eval loop. | +| `utils.py` | `DinoV3FeatureExtractor`, `DynaCLRFeatureExtractor`, `CellDinoFeatureExtractor`, plot helpers. | +| `_configs/*.yaml` | Hydra schemas: `eval.yaml`, `precompute.yaml`, `eval_grouped.yaml`. | + +Other files (`io.py`, `torch_ssim.py`, `formatting.py`, `spectral_pcc/`) house readers, GPU SSIM, and bead/PSF diagnostics. ## Inputs - `io.pred_path` — model predictions, HCS OME-Zarr (channel: `io.pred_channel_name`) -- `io.gt_path` — fluorescence ground truth, HCS OME-Zarr (channel: `io.gt_channel_name`) -- `io.cell_segmentation_path` — *optional* precomputed cell segmentation, HCS OME-Zarr. Required only when `compute_feature_metrics=true` or when building CP/DINOv3/DynaCLR cache entries. Position layout must match GT/pred 1:1. -- `io.gt_cache_dir` — *optional* directory for the GT artifact cache. `null` (default) disables caching; set to a writable path to opt in. Required for `dynacell precompute-gt` and for `io.require_complete_cache=true`. -- `io.pred_cache_dir` — *optional* directory for the prediction artifact cache. `null` (default) preserves the old behavior. Must be distinct from `io.gt_cache_dir`. +- `io.gt_path` — fluorescence ground truth (channel: `io.gt_channel_name`) +- `io.cell_segmentation_path` — *optional* precomputed cell segmentation HCS OME-Zarr. Required when `compute_feature_metrics=true` or when building CP/DINOv3/DynaCLR/CELL-DINO cache entries. Position layout must match GT/pred 1:1. +- `io.gt_cache_dir`, `io.pred_cache_dir` — *optional* artifact cache directories; must be distinct. See [Caches](#caches). -## Running an evaluation - -`dynacell evaluate` is a Hydra entrypoint. Override any field on the CLI with `key=value`. - -Paths and settings that belong to a (target, marker, dataset) combination -live in named Hydra config groups, so most invocations only need to select -the right group and point at the prediction / output paths. Groups come from -two sources: the packaged schema under `src/dynacell/evaluation/_configs/` -and — on a repo checkout — HPC-bound groups under -`configs/benchmarks/virtual_staining/_internal/` that `dynacell.__main__` -exposes through two injected `hydra.searchpath` roots. See the table below. - -### Config groups - -| Group | Options | What it sets | Source | -|---|---|---|---| -| `target` | `er_sec61b`, `mito_tomm20`, `membrane`, `nucleus` | `target_name`, `benchmark.dataset_ref.target`. | `configs/benchmarks/virtual_staining/_internal/shared/eval/target/` | -| `predict_set` | `ipsc_confocal` | `benchmark.dataset_ref.dataset`. | in-package (`_configs/predict_set/`) | -| `feature_extractor/dinov3` | `lvd1689m` | `feature_extractor.dinov3.pretrained_model_name`. | in-package (`_configs/feature_extractor/dinov3/`) | -| `feature_extractor/dynaclr` | `default` | `feature_extractor.dynaclr.checkpoint` and 8-field `encoder` dict. | `configs/benchmarks/virtual_staining/_internal/shared/eval/feature_extractor/dynaclr/` | -| `leaf` | `///eval__` (8 canonical leaves) | Composes all of the above for a canonical benchmark run; see "Benchmark eval leaves" below. | `configs/benchmarks/virtual_staining/_internal/leaf/` (symlink tree) | - -`io.*` fields (`gt_path`, `cell_segmentation_path`, `gt_channel_name`, -`pred_channel_name`, `gt_cache_dir`) and `pixel_metrics.spacing` are now -owned by the dataset manifest (`dynacell/data/manifests.py`) and -spliced into the composed config by a post-compose hook in -`_ref_hook.py`. `pred_channel_name` is derived as -`{target_channel}_prediction`. - -- **In-package** groups (`predict_set`, `feature_extractor/dinov3`, - `spectral_pcc/*`) ship in the wheel: schema and path-free reference - values only. -- **Repo-checkout** groups (`target`, `feature_extractor/dynaclr`, `leaf`) - all live under `configs/benchmarks/virtual_staining/_internal/` — the - hidden support tree that keeps the top level of `virtual_staining/` - biology-only. They contain HPC paths, our DynaCLR checkpoint, and - benchmark-instance values — useless to external users. - `dynacell.__main__` injects two `hydra.searchpath` roots - (`_internal/` for the `leaf/` tree and `_internal/shared/eval/` for - `target/` and `feature_extractor/dynaclr/`) when running from a repo - checkout. Wheel installs without the repo silently omit these, and - external users supply their own via `--config-dir`. -- **Hydra only discovers `.yaml` files for group resolution**, so eval - group files under `_internal/shared/eval/`, the canonical eval leaves - at `///eval__.yaml`, and the - `_internal/leaf/` symlinks all use `.yaml`. - Lightning-side train and predict leaves stay `.yml` (they compose - through `viscy_utils.compose`, which is extension-agnostic). - -Selecting a group on the CLI: `=