diff --git a/software/control/core/backpressure.py b/software/control/core/backpressure.py index 039382ba5..84129aea4 100644 --- a/software/control/core/backpressure.py +++ b/software/control/core/backpressure.py @@ -32,6 +32,8 @@ "BackpressureController", "BackpressureStats", "BackpressureValues", + "note_job_dispatched", + "note_job_completed", "create_backpressure_values", ] @@ -53,6 +55,28 @@ class BackpressureStats: is_throttled: bool +def note_job_dispatched(pending_jobs, pending_bytes, image_bytes: int) -> None: + """Raise the shared counters for one dispatched job. No-op when the runner was built without them.""" + if pending_jobs is None: + return + with pending_jobs.get_lock(): + pending_jobs.value += 1 + with pending_bytes.get_lock(): + pending_bytes.value += image_bytes + + +def note_job_completed(pending_jobs, pending_bytes, capacity_event, image_bytes: int) -> None: + """Lower the shared counters for one finished job (success or failure) and wake a throttled producer.""" + if pending_jobs is None: + return + with pending_jobs.get_lock(): + pending_jobs.value = max(0, pending_jobs.value - 1) + with pending_bytes.get_lock(): + pending_bytes.value = max(0, pending_bytes.value - image_bytes) + if capacity_event is not None: + capacity_event.set() + + def create_backpressure_values() -> BackpressureValues: """Create multiprocessing primitives for cross-process backpressure tracking. diff --git a/software/control/core/in_process_zarr_runner.py b/software/control/core/in_process_zarr_runner.py new file mode 100644 index 000000000..897e95d2d --- /dev/null +++ b/software/control/core/in_process_zarr_runner.py @@ -0,0 +1,146 @@ +"""Zarr v3 saving in the acquisition process. + +tensorstore compresses and writes on its own threads and hands back a future, so a save +subprocess buys nothing for Zarr and costs a copy of every image through a pipe. It also +cannot be fork()ed once this process has used tensorstore (the NDViewer does; a recording +writer will), because tensorstore aborts such a child on its first use. This runner presents +the part of JobRunner's interface the workers use and runs SaveZarrJob here instead: dispatch +submits the write and returns; a drain thread waits for each write in order, keeps the +backpressure counters, and puts the JobResult on the output queue the worker polls. +""" + +import concurrent.futures +import queue +import threading +from typing import Optional + +import squid.logging +from control.core.backpressure import BackpressureValues, note_job_completed, note_job_dispatched +from control.core.job_processing import AcquisitionInfo, JobResult, SaveZarrJob, ZarrWriterInfo, ZarrWriterRegistry + +_SENTINEL = object() + + +class InProcessZarrRunner: + def __init__( + self, + zarr_writer_info: Optional[ZarrWriterInfo] = None, + bp_values: Optional[BackpressureValues] = None, + ): + self._log = squid.logging.get_logger(self.__class__.__name__) + self._zarr_writer_info = zarr_writer_info + self._bp_pending_jobs, self._bp_pending_bytes, self._bp_capacity_event = bp_values or (None, None, None) + self.registry = ZarrWriterRegistry() + self._submitted: "queue.Queue" = queue.Queue() # (job_id, future, image_bytes, result), in submission order + self._output: "queue.Queue" = queue.Queue() + # Counted up before the submit, as JobRunner does, so has_pending() never says False mid-flight. + self._pending = 0 + self._pending_lock = threading.Lock() + # dispatch and shutdown exclude each other, so no job lands behind the drain thread's sentinel + self._state_lock = threading.Lock() + self._thread = threading.Thread(target=self._drain, name="InProcessZarrRunner", daemon=True) + self._started = False + self._finished = False + self._aborted = False + + # ---- lifecycle, as JobRunner has it + + def start(self) -> None: + self._thread.start() + self._started = True + + def is_ready(self) -> bool: + return self._started and not self._finished + + def wait_ready(self, timeout_s: float = 5.0) -> bool: + return self.is_ready() + + def is_alive(self) -> bool: + return self._started and self._thread.is_alive() + + def set_acquisition_info(self, acquisition_info: Optional[AcquisitionInfo]) -> None: + """Accepted for interface parity; SaveZarrJob does not use it.""" + + def set_zarr_writer_info(self, zarr_writer_info: ZarrWriterInfo) -> None: + self._zarr_writer_info = zarr_writer_info + + # ---- the hot path: called from the camera frame callback + + def dispatch(self, job: SaveZarrJob) -> bool: + """Submit the write and return True; return False, leaving the job untouched, once shutdown() has run.""" + with self._state_lock: + if self._finished: + self._log.warning(f"Job {job.job_id} dispatched after shutdown; refusing it") + return False + if self._zarr_writer_info is None: + raise ValueError("Cannot dispatch SaveZarrJob: InProcessZarrRunner has no zarr_writer_info") + job.zarr_writer_info = self._zarr_writer_info + job.registry = self.registry + image_bytes = job.capture_image.image_array.nbytes if job.capture_image.image_array is not None else 0 + with self._pending_lock: + self._pending += 1 + note_job_dispatched(self._bp_pending_jobs, self._bp_pending_bytes, image_bytes) + try: + future, result = job.submit() + except Exception as e: + # Same outcome as a failed write: the worker sees the exception on the output queue. + self._log.exception(f"Job {job.job_id} could not be submitted") + future, result = concurrent.futures.Future(), None + future.set_exception(e) + # Only the id goes to the drain thread: tensorstore has copied the frame, so nothing keeps it alive. + self._submitted.put_nowait((job.job_id, future, image_bytes, result)) + return True + + def output_queue(self) -> "queue.Queue": + return self._output + + def has_pending(self) -> bool: + with self._pending_lock: + return self._pending > 0 + + def kill(self) -> None: + """Mark the run aborted; shutdown() then seals the stores as aborted.""" + self._aborted = True + + def shutdown(self, timeout_s: float = 1.0, aborted: bool = False) -> bool: + """Wait for the outstanding writes, then seal every store. Returns False if a write was still in flight. + + Complete when the acquisition ran to its end; `acquisition_complete: False, aborted: True` + when it was aborted (by the user or by an error), after kill(), or when the writes did not + finish within the timeout (a store is never stamped complete over a write still in flight). + Captured frames are written either way. + """ + with self._state_lock: + if not self._started or self._finished: + return True + self._finished = True + self._submitted.put_nowait(_SENTINEL) + self._thread.join(timeout=max(timeout_s, 1.0)) + drained = not self._thread.is_alive() + if not drained: + self._log.error("a Zarr write is still in flight after the shutdown timeout; sealing the stores as aborted") + if self._aborted or aborted or not drained: + self.registry.abort_all() + elif not self.registry.finalize_all(): + self._log.error("ZARR FINALIZATION INCOMPLETE - Some data may not be saved correctly") + return drained + + # ---- the drain thread + + def _drain(self) -> None: + while True: + item = self._submitted.get() + if item is _SENTINEL: + return + job_id, future, image_bytes, result = item + try: + if future is not None: + future.result() + self._output.put_nowait(JobResult(job_id=job_id, result=result, exception=None)) + except Exception as e: + self._log.exception(f"Job {job_id} failed! Returning exception result.") + self._output.put_nowait(JobResult(job_id=job_id, result=None, exception=e)) + finally: + with self._pending_lock: + self._pending -= 1 + note_job_completed(self._bp_pending_jobs, self._bp_pending_bytes, self._bp_capacity_event, image_bytes) diff --git a/software/control/core/job_processing.py b/software/control/core/job_processing.py index ec27df65a..1d04a816e 100644 --- a/software/control/core/job_processing.py +++ b/software/control/core/job_processing.py @@ -6,7 +6,7 @@ import json from datetime import datetime from contextlib import contextmanager -from typing import ClassVar, Dict, Generic, List, Optional, Set, Tuple, TypeVar, Union +from typing import Any, ClassVar, Dict, Generic, List, Optional, Set, Tuple, TypeVar, Union from uuid import uuid4 from dataclasses import dataclass, field @@ -20,6 +20,7 @@ from control._def import ZProjectionMode, DownsamplingMethod import squid.abc import squid.logging +from control.core.backpressure import note_job_completed, note_job_dispatched from control.models import AcquisitionChannel from control.core import utils_ome_tiff_writer as ome_tiff_writer from control.core.memory_profiler import ( @@ -491,76 +492,87 @@ class ZarrWriteResult: @dataclass -class SaveZarrJob(Job): - """Job for saving images to Zarr v3 format using TensorStore. +class ZarrWriterRegistry: + """The open ZarrWriters of one acquisition, keyed by writer_key, plus the HCS metadata already written. - Uses a process-local ZarrWriter that is initialized lazily on first write. - The zarr_writer_info field is injected by JobRunner.dispatch() before the job runs. + Each acquisition gets its own registry (the in-process runner owns one); the save subprocess + uses SaveZarrJob.default_registry, which lives as long as the subprocess does. """ - _log: ClassVar = squid.logging.get_logger("SaveZarrJob") - zarr_writer_info: Optional[ZarrWriterInfo] = field(default=None) - - # Class-level writer storage keyed by output_path. - # SAFETY: JobRunner runs in a multiprocessing.Process (not threads), so each - # worker process has its own independent copy of this class variable. - # WARNING: This dict is NOT thread-safe. DO NOT use SaveZarrJob with threading - # (e.g., ThreadPoolExecutor) - it will cause race conditions and data corruption. - _zarr_writers: ClassVar[Dict[str, "ZarrWriter"]] = {} + writers: Dict[str, "ZarrWriter"] = field(default_factory=dict) + plates_written: Set[str] = field(default_factory=set) + wells_written: Set[str] = field(default_factory=set) - # Track HCS metadata that has been written (plate path -> True, well path -> True) - _hcs_plate_written: ClassVar[Set[str]] = set() - _hcs_wells_written: ClassVar[Set[str]] = set() + _log: ClassVar = squid.logging.get_logger("ZarrWriterRegistry") - @classmethod - def clear_writers(cls) -> None: - """Clear all zarr writers, aborting any that are still active. + def finalize_all(self) -> bool: + """Finalize every open writer. Returns False if any failed. Empties the writers (HCS bookkeeping is kept).""" + failed = [] + for path, writer in list(self.writers.items()): + if writer.is_initialized and not writer.is_finalized: + try: + writer.finalize() + self._log.info(f"Finalized zarr writer: {path}") + except Exception as e: + self._log.error(f"Error finalizing writer {path}: {e}") + failed.append(path) + self.writers.clear() + if failed: + self._log.error(f"Failed to finalize {len(failed)} zarr writers: {failed}") + return False + return True - Call at start of new acquisition to ensure clean state. - Uses try-finally to guarantee dictionaries are cleared even if abort fails. - """ + def abort_all(self) -> None: + """Seal every open writer as aborted. Empties the registry even if a writer fails to abort.""" try: - for writer in list(cls._zarr_writers.values()): + for writer in list(self.writers.values()): if writer.is_initialized and not writer.is_finalized: try: writer.abort() except Exception as e: - cls._log.warning(f"Error aborting writer during clear: {e}") + self._log.warning(f"Error aborting writer during clear: {e}") finally: - # Always clear dictionaries, even if abort loop fails - cls._zarr_writers.clear() - cls._hcs_plate_written.clear() - cls._hcs_wells_written.clear() + self.writers.clear() + self.plates_written.clear() + self.wells_written.clear() - @classmethod - def finalize_all_writers(cls) -> bool: - """Finalize all active zarr writers. - Call at end of acquisition to ensure all data is written. +@dataclass +class SaveZarrJob(Job): + """Job for saving images to Zarr v3 format using TensorStore. - Returns: - True if all writers finalized successfully, False if any failed. - """ - failed_paths = [] - for path, writer in list(cls._zarr_writers.items()): - if writer.is_initialized and not writer.is_finalized: - try: - writer.finalize() - cls._log.info(f"Finalized zarr writer: {path}") - except Exception as e: - cls._log.error(f"Error finalizing writer {path}: {e}") - failed_paths.append(path) - cls._zarr_writers.clear() - if failed_paths: - cls._log.error(f"Failed to finalize {len(failed_paths)} zarr writers: {failed_paths}") - return False - return True + Writes through its registry's ZarrWriter for the store (an InProcessZarrRunner's own + registry, or default_registry in the save subprocess), initialized lazily on first write. + The zarr_writer_info field is injected by whichever runner dispatches the job. + """ + + _log: ClassVar = squid.logging.get_logger("SaveZarrJob") + zarr_writer_info: Optional[ZarrWriterInfo] = field(default=None) + # The acquisition's writers. None = default_registry (the save subprocess, which has one acquisition + # per process). The in-process runner injects its own, as JobRunner.dispatch injects zarr_writer_info. + registry: Optional["ZarrWriterRegistry"] = field(default=None) + + default_registry: ClassVar["ZarrWriterRegistry"] = ZarrWriterRegistry() + + @property + def _registry(self) -> "ZarrWriterRegistry": + return self.registry if self.registry is not None else SaveZarrJob.default_registry + + @classmethod + def clear_writers(cls) -> None: + """Abort and forget the default registry's writers (start of an acquisition in the subprocess).""" + cls.default_registry.abort_all() + + @classmethod + def finalize_all_writers(cls) -> bool: + """Finalize the default registry's writers (end of the subprocess). Returns False if any failed.""" + return cls.default_registry.finalize_all() def _write_hcs_metadata_if_needed(self, region_id: str, fov: int) -> None: """Write HCS plate and well metadata if not already written. Called when a new writer is initialized for an HCS acquisition. - Uses class-level sets to track which plate/well metadata has been written. + Uses the registry's sets to track which plate/well metadata has been written. Args: region_id: Well ID (e.g., "A1", "B12") @@ -572,28 +584,39 @@ def _write_hcs_metadata_if_needed(self, region_id: str, fov: int) -> None: # Write plate metadata (once per acquisition) plate_path = info.get_plate_path() - if plate_path not in self._hcs_plate_written: + if plate_path not in self._registry.plates_written: rows, cols, wells = info.get_hcs_structure() write_plate_metadata(plate_path, rows, cols, wells, plate_name="plate") - self._hcs_plate_written.add(plate_path) + self._registry.plates_written.add(plate_path) self._log.info(f"Wrote HCS plate metadata: {len(wells)} wells") # Write well metadata (once per well) well_path = info.get_well_path(region_id) - if well_path not in self._hcs_wells_written: + if well_path not in self._registry.wells_written: # Get FOV count for this well fov_count = info.get_fov_count(region_id) fields = list(range(fov_count)) write_well_metadata(well_path, fields) - self._hcs_wells_written.add(well_path) + self._registry.wells_written.add(well_path) self._log.debug(f"Wrote HCS well metadata for {region_id}: {fov_count} fields") def run(self) -> ZarrWriteResult: + """Write the frame and wait for it: what the save subprocess does, one job at a time.""" + future, result = self.submit() + if future is not None: + future.result() + return result + + def submit(self) -> Tuple[Optional[Any], ZarrWriteResult]: + """Submit the frame write without waiting for it. + + Returns the tensorstore future (None in simulated-I/O mode, where the write already + happened) and the ZarrWriteResult to report once the future completes. + """ if self.zarr_writer_info is None: raise ValueError( - "SaveZarrJob.run() requires zarr_writer_info but it is None. " - "This job must be dispatched via JobRunner.dispatch(), which injects zarr_writer_info. " - "If running directly, set job.zarr_writer_info before calling run()." + "SaveZarrJob requires zarr_writer_info but it is None; the runner that dispatches the job " + "injects it (set job.zarr_writer_info when running one directly)." ) from control.core.io_simulation import is_simulation_enabled, simulated_zarr_write @@ -654,13 +677,13 @@ def run(self) -> ZarrWriteResult: f"SaveZarrJob {self.job_id}: simulated write of {bytes_written} bytes " f"to {output_path} (image shape={image.shape})" ) - return result + return None, result - self._save_zarr(image, info, output_path) - return result + future = self._submit_write(image, info, output_path) + return future, result - def _save_zarr(self, image: np.ndarray, info: CaptureInfo, output_path: str) -> None: - """Write image to zarr dataset using TensorStore. + def _submit_write(self, image: np.ndarray, info: CaptureInfo, output_path: str) -> Any: + """Submit the image write to the zarr dataset using TensorStore; returns the write future. Args: image: Image array to write @@ -684,7 +707,7 @@ def _save_zarr(self, image: np.ndarray, info: CaptureInfo, output_path: str) -> else: writer_key = output_path # Unique per FOV - if writer_key not in self._zarr_writers: + if writer_key not in self._registry.writers: if is_hcs or not use_6d_fov: # 5D shape: (T, C, Z, Y, X) - one writer per FOV shape = ( @@ -728,7 +751,7 @@ def _save_zarr(self, image: np.ndarray, info: CaptureInfo, output_path: str) -> except Exception as e: self._log.error(f"Failed to initialize zarr writer for {output_path}: {e}") raise - self._zarr_writers[writer_key] = writer + self._registry.writers[writer_key] = writer if is_hcs: mode_str = "HCS 5D" # Write HCS plate and well metadata @@ -739,7 +762,7 @@ def _save_zarr(self, image: np.ndarray, info: CaptureInfo, output_path: str) -> mode_str = "non-HCS 5D per-FOV" self._log.info(f"Initialized zarr writer ({mode_str}): {output_path}") - writer = self._zarr_writers[writer_key] + writer = self._registry.writers[writer_key] # Write frame t = info.time_point or 0 @@ -747,13 +770,12 @@ def _save_zarr(self, image: np.ndarray, info: CaptureInfo, output_path: str) -> z = info.z_index if is_hcs or not use_6d_fov: - # 5D write - writer.write_frame(image, t=t, c=c, z=z) - self._log.debug(f"Wrote frame t={t}, c={c}, z={z} to {output_path}") + future = writer.submit_frame(image, t=t, c=c, z=z) # 5D + self._log.debug(f"Submitted frame t={t}, c={c}, z={z} to {output_path}") else: - # 6D write with FOV index - writer.write_frame(image, t=t, c=c, z=z, fov=fov) - self._log.debug(f"Wrote frame t={t}, c={c}, z={z}, fov={fov} to {output_path}") + future = writer.submit_frame(image, t=t, c=c, z=z, fov=fov) # 6D with FOV index + self._log.debug(f"Submitted frame t={t}, c={c}, z={z}, fov={fov} to {output_path}") + return future # These are debugging jobs - they should not be used in normal usage! @@ -849,11 +871,7 @@ def dispatch(self, job: Job): # has_pending() to return False while job is still in flight. with self._pending_count.get_lock(): self._pending_count.value += 1 - if self._bp_pending_jobs is not None: - with self._bp_pending_jobs.get_lock(): - self._bp_pending_jobs.value += 1 - with self._bp_pending_bytes.get_lock(): - self._bp_pending_bytes.value += image_bytes + note_job_dispatched(self._bp_pending_jobs, self._bp_pending_bytes, image_bytes) try: self._input_queue.put_nowait(job) @@ -1026,21 +1044,15 @@ def run(self): with self._pending_count.get_lock(): self._pending_count.value -= 1 - # Backpressure tracking: decrement counters immediately when job completes. - # Backpressure tracks queue memory, not subprocess memory. - if self._bp_pending_jobs is not None: - with self._bp_pending_jobs.get_lock(): - self._bp_pending_jobs.value = max(0, self._bp_pending_jobs.value - 1) - - # Decrement image bytes - if job.capture_image and job.capture_image.image_array is not None: - image_bytes = job.capture_image.image_array.nbytes - with self._bp_pending_bytes.get_lock(): - self._bp_pending_bytes.value = max(0, self._bp_pending_bytes.value - image_bytes) - - # Signal capacity available for all job completions - if self._bp_capacity_event is not None: - self._bp_capacity_event.set() + # Backpressure tracks queue memory, not subprocess memory: release this job's share now. + image_bytes = ( + job.capture_image.image_array.nbytes + if job.capture_image and job.capture_image.image_array is not None + else 0 + ) + note_job_completed( + self._bp_pending_jobs, self._bp_pending_bytes, self._bp_capacity_event, image_bytes + ) # Finalize any zarr writers that are still open try: diff --git a/software/control/core/multi_point_controller.py b/software/control/core/multi_point_controller.py index 286aac9e8..8e998b7ed 100644 --- a/software/control/core/multi_point_controller.py +++ b/software/control/core/multi_point_controller.py @@ -23,6 +23,8 @@ from control.core.live_controller import LiveController from control.microscope import Microscope from control.core.multi_point_worker import MultiPointWorker +from control.core.in_process_zarr_runner import InProcessZarrRunner +from control.core.job_processing import SaveSettings from control.core.objective_store import ObjectiveStore from control.core.memory_profiler import MemoryMonitor, log_memory from control.microcontroller import Microcontroller @@ -308,20 +310,21 @@ def _cleanup_prewarmed_runner( self._log.error(f"Error shutting down pre-warmed runner {context}: {e}") def get_prewarmed_job_runner(self) -> Tuple[Optional["JobRunner"], Optional["BackpressureValues"]]: - """Get the pre-warmed job runner and its shared backpressure values. + """The pre-warmed save subprocess and its backpressure values for the acquisition about to start. - Returns: - Tuple of (runner, bp_values) where: - - runner: JobRunner instance or None if not available - - bp_values: BackpressureValues tuple or None - - The runner and values are cleared (so they're only used once). + Zarr v3 is saved in the acquisition process (InProcessZarrRunner), so a Zarr run gets + (None, None) and the warm subprocess stays ready for a later TIFF run. Any other format + consumes it (the references are cleared so it is only used once) and a fresh one starts + warming for the next acquisition. Usage: runner, bp_values = controller.get_prewarmed_job_runner() worker = MultiPointWorker(..., prewarmed_job_runner=runner, prewarmed_bp_values=bp_values) """ + if SaveSettings().file_saving_option == control._def.FileSavingOption.ZARR_V3: + return None, None + runner = self._prewarmed_job_runner bp_values = self._prewarmed_bp_values @@ -975,8 +978,7 @@ def finish_fn(): except Exception as e: self._log.warning(f"Failed to write acquisition watchdog start state: {e}") - # Get pre-warmed job runner and its shared backpressure values - # (starts a new one warming for next acquisition) + # The pre-warmed save subprocess and its backpressure values (None, None for a Zarr v3 run) prewarmed_runner, prewarmed_bp_values = self.get_prewarmed_job_runner() # Worker creation can fail - ensure runner is cleaned up on error @@ -999,9 +1001,7 @@ def finish_fn(): run_state_writer=self._run_state_writer, ) except Exception: - # Clean up pre-warmed runner if worker creation failed. - # Note: get_prewarmed_job_runner() already started a NEW pre-warmed runner, - # so we're cleaning up the one that was handed off to us. + # The runner handed to us is ours to clean up; a replacement is already warming (None for Zarr v3). self._cleanup_prewarmed_runner( prewarmed_runner, context="after worker creation failure", @@ -1232,6 +1232,13 @@ def close(self, timeout_s: float = 5.0) -> None: for job_class, job_runner in job_runners: try: if job_runner is not None and job_runner.is_alive(): + if isinstance(job_runner, InProcessZarrRunner): + # Not a process: shut it down so its open stores are sealed (as aborted). + self._log.warning( + f"Shutting down {job_class.__name__} in-process runner (abnormal shutdown)" + ) + job_runner.shutdown(timeout_s=self._PROCESS_TERMINATE_TIMEOUT_S, aborted=True) + continue self._log.warning(f"Terminating {job_class.__name__} job runner (abnormal shutdown)") job_runner.terminate() job_runner.join(timeout=self._PROCESS_TERMINATE_TIMEOUT_S) diff --git a/software/control/core/multi_point_worker.py b/software/control/core/multi_point_worker.py index dfcaa8f84..8d7a19e08 100644 --- a/software/control/core/multi_point_worker.py +++ b/software/control/core/multi_point_worker.py @@ -47,6 +47,7 @@ JobRunner, JobResult, ) +from control.core.in_process_zarr_runner import InProcessZarrRunner from control.core.mosaic_utils import ( calculate_overlap_pixels, parse_well_id, @@ -311,15 +312,26 @@ def __init__( mode_str = "per-FOV 5D (OME-NGFF compliant)" self._log.info(f"ZARR_V3 output: {mode_str}, base path: {self.experiment_path}") - # Use pre-warmed job runner if available, otherwise create new ones. - # IMPORTANT: Only use pre-warmed runner if BOTH runner AND backpressure values + # Zarr is saved in this process; the TIFF formats keep the save subprocess. + # Use the pre-warmed subprocess if available, otherwise create new ones. + # IMPORTANT: Only use the pre-warmed runner if BOTH runner AND backpressure values # are available. Using a runner without matching backpressure values would cause # the BackpressureController to track different counters than the JobRunner. can_use_prewarmed = prewarmed_job_runner is not None and prewarmed_bp_values is not None used_prewarmed = False for job_class in job_classes: job_runner = None - if Acquisition.USE_MULTIPROCESSING: + if job_class is SaveZarrJob: + job_runner = self._create_in_process_zarr_runner(zarr_writer_info) + if can_use_prewarmed and not used_prewarmed: + # Nothing here will use the subprocess the controller handed over. + self._log.info("Zarr v3 saves in-process; releasing the pre-warmed save subprocess") + try: + prewarmed_job_runner.shutdown(timeout_s=1.0) + except Exception as e: + self._log.error(f"Error shutting down the unused pre-warmed runner: {e}") + can_use_prewarmed = False + elif Acquisition.USE_MULTIPROCESSING: # Try to use pre-warmed runner for the first job class if can_use_prewarmed and not used_prewarmed: if prewarmed_job_runner.is_ready(): @@ -434,6 +446,19 @@ def _is_well_based_acquisition(self) -> bool: ) return True + def _create_in_process_zarr_runner(self, zarr_writer_info: Optional[ZarrWriterInfo]) -> InProcessZarrRunner: + """Zarr is saved in this process (see in_process_zarr_runner.py); the runner shares our backpressure counters.""" + runner = InProcessZarrRunner( + zarr_writer_info=zarr_writer_info, + bp_values=( + self._backpressure.pending_jobs_value, + self._backpressure.pending_bytes_value, + self._backpressure.capacity_event, + ), + ) + runner.start() + return runner + def _abort_due_to_error(self) -> None: """Abort the run due to an internal error (vs a user abort). @@ -674,17 +699,23 @@ def time_left(): time.sleep(0.1) else: # Timed out - kill any runners that still have pending jobs + abandoned = False for job_class, job_runner in active_runners: if job_runner.has_pending(): self._log.error( f"Timed out after {timeout_s} [s] waiting for jobs to finish. Pending jobs for {job_class.__name__} abandoned!!!" ) job_runner.kill() + abandoned = True + if abandoned: + # Data is missing, so the run ends as an error (and the Zarr stores are sealed as aborted). + self._abort_due_to_error() - # Drain results before shutdown + # Drain results before shutdown (a failed job seen here aborts the run too, via _summarize_job_result) self._summarize_runner_outputs(drain_all=True) - # Shut down all job runners in parallel (in background to avoid blocking on subprocess termination). + # Shut down the save subprocesses in parallel (in background to avoid blocking on subprocess termination); + # the in-process Zarr runner is shut down inline below, because its shutdown seals the stores. # Using daemon threads is safe here because: # 1. All jobs are complete and results are already drained # 2. The subprocess termination is best-effort cleanup only @@ -698,9 +729,20 @@ def shutdown_runner(job_runner, timeout): except Exception as e: log.error(f"Error shutting down job runner in background: {e}") - self._log.info("Shutting down job runners (non-blocking)...") + self._log.info("Shutting down job runners...") remaining_time = time_left() + # A user Stop, an error abort, or a crash out of run(): the store must not claim completeness. + aborted = bool(self.abort_requested_fn()) or self._run_state_fatal for job_class, job_runner in active_runners: + if isinstance(job_runner, InProcessZarrRunner): + # Its shutdown seals the stores (aborted ones as aborted); do it here so the + # acquisition is not reported finished before the metadata says so. + try: + if not job_runner.shutdown(remaining_time, aborted=aborted): + self._abort_due_to_error() # a write was still in flight: the stores were sealed aborted + except Exception as e: + log.error(f"Error shutting down the in-process Zarr runner: {e}") + continue t = threading.Thread(target=shutdown_runner, args=(job_runner, remaining_time), daemon=True) t.start() @@ -863,6 +905,12 @@ def _summarize_runner_outputs(self, drain_all: bool = False) -> SummarizeResult: if out_queue is None: # Queue was cleared during shutdown continue + if not drain_all and not job_runner.is_alive() and job_runner.has_pending(): + self._log.error( + f"{job_class.__name__} save subprocess exited (exit code {getattr(job_runner, 'exitcode', None)}) " + f"with jobs pending; aborting the acquisition" + ) + self._abort_due_to_error() while True: try: job_result: JobResult = out_queue.get_nowait() @@ -885,6 +933,9 @@ def _summarize_job_result(self, job_result: JobResult) -> bool: if job_result.exception is not None: self._log.error(f"Error while running job {job_result.job_id}: {job_result.exception}") self._acquisition_error_count += 1 + if self._abort_on_failed_job: + self._log.error("A save job failed, aborting acquisition because abort_on_failed_job=True") + self._abort_due_to_error() # Send Slack error notification if self._slack_notifier is not None: @@ -1070,9 +1121,7 @@ def run_coordinate_acquisition(self, current_path): with self._timing.get_timer("job result summaries"): result = self._summarize_runner_outputs() if not result.none_failed and self._abort_on_failed_job: - self._log.error("Some jobs failed, aborting acquisition because abort_on_failed_job=True") - self._abort_due_to_error() - return + return # _summarize_job_result has requested the abort with self._timing.get_timer("move_to_coordinate"): self.move_to_coordinate(coordinate_mm, region_id, fov) @@ -1418,6 +1467,8 @@ def _image_callback(self, camera_frame: CameraFrame): return else: try: + # In-process fallback (multiprocessing off) for the TIFF job classes; + # SaveZarrJob always has an InProcessZarrRunner. # NOTE(imo): We don't have any way of people using results, so for now just # grab and ignore it. result = job.run() diff --git a/software/control/core/zarr_writer.py b/software/control/core/zarr_writer.py index 483c7a66e..469ac415b 100644 --- a/software/control/core/zarr_writer.py +++ b/software/control/core/zarr_writer.py @@ -11,7 +11,6 @@ - Blosc compression with LZ4/Zstd codecs """ -import asyncio import json import os from dataclasses import dataclass, field @@ -341,8 +340,8 @@ def write_well_metadata( class ZarrWriter: """Zarr v3 writer for use in job processing. - Directly uses TensorStore without asyncio overhead for write operations. - Only uses asyncio for initialization and finalization where it's unavoidable. + Uses TensorStore futures directly and needs no event loop, so a store can be + initialized on one thread (e.g. the camera frame callback) and finalized on another. """ def __init__(self, config: ZarrAcquisitionConfig): @@ -353,121 +352,93 @@ def __init__(self, config: ZarrAcquisitionConfig): """ self._config = config self._dataset = None - self._pending_futures: List[Any] = [] self._initialized = False self._finalized = False - self._loop: Optional[asyncio.AbstractEventLoop] = None - self._owns_loop = False # True if we created the loop ourselves - - def _get_loop(self) -> asyncio.AbstractEventLoop: - """Get or create the event loop (only used for init/finalize).""" - if self._loop is None or self._loop.is_closed(): - try: - self._loop = asyncio.get_event_loop() - self._owns_loop = False # Using existing loop - except RuntimeError: - self._loop = asyncio.new_event_loop() - asyncio.set_event_loop(self._loop) - self._owns_loop = True # We created this loop - return self._loop def initialize(self) -> None: - """Initialize the zarr dataset (blocking). - - Uses asyncio only for TensorStore's async open operation. - """ + """Initialize the zarr dataset (blocking).""" if self._initialized: log.warning("Writer already initialized") return - try: - ts = _get_tensorstore() - config = self._config + ts = _get_tensorstore() + config = self._config - # Build TensorStore spec - os.makedirs(os.path.dirname(config.output_path), exist_ok=True) + # Build TensorStore spec + os.makedirs(os.path.dirname(config.output_path), exist_ok=True) - chunk_shape = _get_chunk_shape(config) - shard_shape = _get_shard_shape(config) - compression_codec = _get_compression_codec(config.compression) + chunk_shape = _get_chunk_shape(config) + shard_shape = _get_shard_shape(config) + compression_codec = _get_compression_codec(config.compression) - # Dimension names and transpose order depend on 5D vs 6D - if config.ndim == 5: - transpose_order = [4, 3, 2, 1, 0] # Reverse order for C-contiguous layout - else: - transpose_order = [5, 4, 3, 2, 1, 0] # Reverse order for C-contiguous layout + # Dimension names and transpose order depend on 5D vs 6D + if config.ndim == 5: + transpose_order = [4, 3, 2, 1, 0] # Reverse order for C-contiguous layout + else: + transpose_order = [5, 4, 3, 2, 1, 0] # Reverse order for C-contiguous layout + + # Determine if we need sharding (when chunk != shard) + use_sharding = chunk_shape != shard_shape - # Determine if we need sharding (when chunk != shard) - use_sharding = chunk_shape != shard_shape + # Build inner codec chain (with or without compression) + inner_codecs = [ + {"name": "transpose", "configuration": {"order": transpose_order}}, + {"name": "bytes", "configuration": {"endian": "little"}}, + ] + if compression_codec is not None: + inner_codecs.append(compression_codec) - # Build inner codec chain (with or without compression) - inner_codecs = [ - {"name": "transpose", "configuration": {"order": transpose_order}}, - {"name": "bytes", "configuration": {"endian": "little"}}, + if use_sharding: + codecs = [ + { + "name": "sharding_indexed", + "configuration": { + "chunk_shape": list(chunk_shape), + "codecs": inner_codecs, + "index_codecs": [ + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "crc32c"}, + ], + }, + } ] - if compression_codec is not None: - inner_codecs.append(compression_codec) + chunk_config = list(shard_shape) + else: + codecs = inner_codecs + chunk_config = list(chunk_shape) - if use_sharding: - codecs = [ - { - "name": "sharding_indexed", - "configuration": { - "chunk_shape": list(chunk_shape), - "codecs": inner_codecs, - "index_codecs": [ - {"name": "bytes", "configuration": {"endian": "little"}}, - {"name": "crc32c"}, - ], - }, - } - ] - chunk_config = list(shard_shape) - else: - codecs = inner_codecs - chunk_config = list(chunk_shape) - - spec = { - "driver": "zarr3", - "kvstore": {"driver": "file", "path": config.output_path}, - "metadata": { - "shape": list(config.shape), - "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": chunk_config}}, - "chunk_key_encoding": {"name": "default"}, - "data_type": _dtype_to_zarr(config.dtype), - "codecs": codecs, - "fill_value": 0, - }, - } + spec = { + "driver": "zarr3", + "kvstore": {"driver": "file", "path": config.output_path}, + "metadata": { + "shape": list(config.shape), + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": chunk_config}}, + "chunk_key_encoding": {"name": "default"}, + "data_type": _dtype_to_zarr(config.dtype), + "codecs": codecs, + "fill_value": 0, + }, + } - log.info( - f"Initializing Zarr v3 dataset: {config.output_path}, " - f"shape={config.shape}, chunks={chunk_shape}, shards={shard_shape}, " - f"compression={config.compression.value}" + log.info( + f"Initializing Zarr v3 dataset: {config.output_path}, " + f"shape={config.shape}, chunks={chunk_shape}, shards={shard_shape}, " + f"compression={config.compression.value}" + ) + + # Warn if overwriting existing data + if os.path.exists(config.output_path): + log.warning( + "Zarr v3 dataset path already exists and will be overwritten: %s", + config.output_path, ) - # Warn if overwriting existing data - if os.path.exists(config.output_path): - log.warning( - "Zarr v3 dataset path already exists and will be overwritten: %s", - config.output_path, - ) - - # Use asyncio only for the TensorStore open operation - async def _open(): - return await ts.open(spec, create=True, delete_existing=True) - - loop = self._get_loop() - self._dataset = loop.run_until_complete(_open()) - self._initialized = True - log.info("Zarr v3 dataset initialized successfully") - - # Write metadata synchronously (just file I/O) - self._write_zarr_metadata() - except Exception: - # Clean up event loop if we created it during failed initialization - self._cleanup_event_loop() - raise + self._dataset = ts.open(spec, create=True, delete_existing=True).result() + self._initialized = True + log.info("Zarr v3 dataset initialized successfully") + + # Write metadata synchronously (just file I/O) + self._write_zarr_metadata() def _is_ome_ngff_array_path(self) -> bool: """Check if output_path is an OME-NGFF array path (needs group-level metadata). @@ -651,17 +622,13 @@ def _write_zarr_metadata(self) -> None: log.error(f"Failed to write zarr metadata to {zarr_json_path}: {e}") raise RuntimeError(f"Failed to write zarr metadata: {e}") from e - def write_frame(self, image: np.ndarray, t: int, c: int, z: int, fov: Optional[int] = None) -> None: - """Write a single frame and block until the TensorStore write completes. + def submit_frame(self, image: np.ndarray, t: int, c: int, z: int, fov: Optional[int] = None): + """Submit a single frame write and return its TensorStore future WITHOUT waiting. - This method submits an asynchronous write via TensorStore's write API and - then waits on the resulting future (via future.result()) before returning. - This blocks the calling thread until the write has finished. The underlying - disk I/O is handled asynchronously by TensorStore; this method synchronously - waits for that async operation to complete. - - This ensures data is visible to other processes reading the same zarr store - before this method returns. + The returned object's ``.result()`` blocks until the write is committed (and raises if it + failed). TensorStore copies the source array as part of the write, so ``image`` may be + released once the future completes. Keeping several submissions in flight lets TensorStore + overlap encode, file creation, write and fsync of consecutive frames. Args: image: 2D image array (Y, X) @@ -695,54 +662,15 @@ def write_frame(self, image: np.ndarray, t: int, c: int, z: int, fov: Optional[i if image.dtype != config.dtype: image = image.astype(config.dtype) - # Write using TensorStore and wait for completion - # This ensures data is flushed before we notify the viewer if config.ndim == 5: - future = self._dataset[t, c, z, :, :].write(image) log.debug(f"Writing frame t={t}, c={c}, z={z}") - else: - future = self._dataset[fov, t, c, z, :, :].write(image) - log.debug(f"Writing frame fov={fov}, t={t}, c={c}, z={z}") - - # Wait for write to complete (blocking) - # TensorStore futures have a .result() method that blocks until complete - future.result() - if config.ndim == 5: - log.debug(f"Write complete for frame t={t}, c={c}, z={z}") - else: - log.debug(f"Write complete for frame fov={fov}, t={t}, c={c}, z={z}") - - def wait_for_pending(self, timeout_s: Optional[float] = None) -> int: - """Wait for pending writes (blocking). - - Args: - timeout_s: Optional timeout in seconds (not currently enforced) - - Returns: - Number of writes completed - """ - if not self._pending_futures: - return 0 - - count = len(self._pending_futures) - log.debug(f"Waiting for {count} pending writes...") + return self._dataset[t, c, z, :, :].write(image) + log.debug(f"Writing frame fov={fov}, t={t}, c={c}, z={z}") + return self._dataset[fov, t, c, z, :, :].write(image) - # Wait for all TensorStore futures in parallel using asyncio.gather - # This is more efficient than waiting sequentially - async def _wait_all(): - await asyncio.gather(*self._pending_futures) - - loop = self._get_loop() - loop.run_until_complete(_wait_all()) - - self._pending_futures.clear() - log.debug(f"Completed {count} pending writes") - return count - - @property - def pending_write_count(self) -> int: - """Number of writes currently pending.""" - return len(self._pending_futures) + def write_frame(self, image: np.ndarray, t: int, c: int, z: int, fov: Optional[int] = None) -> None: + """submit_frame() and wait for it, so the data is on disk for other readers once this returns.""" + self.submit_frame(image, t, c, z, fov).result() def finalize(self) -> None: """Finalize the dataset (blocking).""" @@ -752,9 +680,6 @@ def finalize(self) -> None: log.info("Finalizing Zarr v3 dataset...") - # Wait for all pending writes - self.wait_for_pending() - # Update metadata with completion status (in zarr.json attributes) zarr_json_path = self._get_metadata_zarr_json_path() try: @@ -772,7 +697,6 @@ def finalize(self) -> None: # Don't raise - data is already written, just log the metadata issue self._finalized = True - self._cleanup_event_loop() log.info(f"Zarr v3 dataset finalized: {self._config.output_path}") def abort(self) -> None: @@ -784,8 +708,6 @@ def abort(self) -> None: log.warning("Aborting Zarr writer...") try: - # Clear pending futures (don't wait for them) - self._pending_futures.clear() # Mark as incomplete in metadata (in zarr.json attributes) zarr_json_path = self._get_metadata_zarr_json_path() @@ -805,23 +727,8 @@ def abort(self) -> None: finally: # Always mark as finalized and cleanup resources self._finalized = True - self._cleanup_event_loop() log.warning(f"Zarr writer aborted: {self._config.output_path}") - def _cleanup_event_loop(self) -> None: - """Clean up the event loop to prevent resource leaks. - - Only closes the loop if we created it ourselves (via new_event_loop). - Loops obtained from get_event_loop() are shared and should not be closed. - """ - if self._loop is not None and self._owns_loop and not self._loop.is_closed(): - try: - self._loop.close() - except Exception as e: - log.warning(f"Error closing event loop: {e}") - self._loop = None - self._owns_loop = False - @property def is_initialized(self) -> bool: return self._initialized diff --git a/software/docs/zarr-v3-format.md b/software/docs/zarr-v3-format.md index cbed8daa4..e4c07584c 100644 --- a/software/docs/zarr-v3-format.md +++ b/software/docs/zarr-v3-format.md @@ -7,6 +7,7 @@ This document describes the Zarr v3 output format for Squid acquisitions. Squid supports saving acquisition data in Zarr v3 format with OME-NGFF 0.5 metadata. This format provides: - **High performance**: TensorStore backend with sharding for ~200 MB/s write speed +- **Written in the acquisition process**: Zarr v3 is saved by `ZarrWriter` through `InProcessZarrRunner`, not by the save subprocess the TIFF formats use. tensorstore does its I/O on its own threads, and it aborts any `fork()`ed child of a process that has used it, so a subprocess would gain nothing and must never be forked for Zarr. - **Compression options**: None, Fast (LZ4), Balanced (Zstd), Best (Zstd level 9) - **Streaming support**: Data can be read during acquisition - **OME-NGFF 0.5 metadata**: Standard metadata format (note: viewer support for Zarr v3 is still emerging) diff --git a/software/tests/control/core/test_backpressure.py b/software/tests/control/core/test_backpressure.py index bdedde3fa..aea57d5db 100644 --- a/software/tests/control/core/test_backpressure.py +++ b/software/tests/control/core/test_backpressure.py @@ -743,12 +743,14 @@ def test_close_handles_exception_in_abort(self): def test_close_terminates_live_job_runners(self): """close() terminates job runners that are still alive.""" from unittest.mock import MagicMock + from control.core.job_processing import JobRunner from control.core.multi_point_controller import MultiPointController controller = self._create_mock_controller() controller.acquisition_in_progress.return_value = False - mock_job_runner = MagicMock() + # spec=JobRunner: a save subprocess (a bare MagicMock would answer runs_in_process truthily) + mock_job_runner = MagicMock(spec=JobRunner) mock_job_runner.is_alive.side_effect = [True, False] controller.multiPointWorker = MagicMock() controller.multiPointWorker._job_runners = [(SlowJob, mock_job_runner)] @@ -761,12 +763,14 @@ def test_close_terminates_live_job_runners(self): def test_close_force_kills_stubborn_runners(self): """close() force kills job runners that don't respond to terminate.""" from unittest.mock import MagicMock + from control.core.job_processing import JobRunner from control.core.multi_point_controller import MultiPointController controller = self._create_mock_controller() controller.acquisition_in_progress.return_value = False - mock_job_runner = MagicMock() + # spec=JobRunner: a save subprocess (a bare MagicMock would answer runs_in_process truthily) + mock_job_runner = MagicMock(spec=JobRunner) mock_job_runner.is_alive.side_effect = [True, True, False] controller.multiPointWorker = MagicMock() controller.multiPointWorker._job_runners = [(SlowJob, mock_job_runner)] diff --git a/software/tests/control/core/test_in_process_zarr_runner.py b/software/tests/control/core/test_in_process_zarr_runner.py new file mode 100644 index 000000000..23f2f89c4 --- /dev/null +++ b/software/tests/control/core/test_in_process_zarr_runner.py @@ -0,0 +1,288 @@ +"""InProcessZarrRunner: SaveZarrJob in this process, with JobRunner's interface and counters.""" + +import concurrent.futures +import json +import os +import queue +import threading +import time + +import numpy as np +import pytest + +from control.core.backpressure import create_backpressure_values +from control.core.job_processing import JobImage, JobResult, SaveZarrJob, ZarrWriteResult, ZarrWriterInfo +from tests.control.core.test_zarr_writer import make_test_capture_info + +pytest.importorskip("tensorstore") + + +def _info(tmp_path, t_size=1, z_size=2): + return ZarrWriterInfo(base_path=str(tmp_path / "acq"), t_size=t_size, c_size=1, z_size=z_size) + + +def _job(z, value=1): + return SaveZarrJob( + capture_info=make_test_capture_info(z_index=z), + capture_image=JobImage(image_array=np.full((8, 8), value, dtype=np.uint16)), + ) + + +def _failed_future(exc): + future = concurrent.futures.Future() + future.set_exception(exc) + return future + + +def _results(runner, n, timeout=10.0): + out, deadline = [], time.monotonic() + timeout + while len(out) < n and time.monotonic() < deadline: + try: + out.append(runner.output_queue().get(timeout=0.1)) + except queue.Empty: + pass + return out + + +def _wait_settled(runner, bp_values=None, timeout=5.0): + # The drain thread queues the JobResult before its finally lowers the counters and sets capacity. + def settled(): + if runner.has_pending(): + return False + if bp_values is None: + return True + pending_jobs, pending_bytes, capacity = bp_values + return pending_jobs.value == 0 and pending_bytes.value == 0 and capacity.is_set() + + deadline = time.monotonic() + timeout + while not settled() and time.monotonic() < deadline: + time.sleep(0.01) + + +def _squid_attrs(info): + # The OME group metadata (with _squid) lives in the group above the "0" array (ZarrWriter._get_metadata_zarr_json_path). + group_json = os.path.join(os.path.dirname(info.get_output_path("A1", 0)), "zarr.json") + with open(group_json) as f: + return json.load(f)["attributes"]["_squid"] + + +def test_dispatch_writes_and_reports_a_result(tmp_path): + from control.core.in_process_zarr_runner import InProcessZarrRunner + + info = _info(tmp_path) + runner = InProcessZarrRunner(zarr_writer_info=info) + runner.start() + assert runner.is_ready() and runner.wait_ready(timeout_s=1.0) and runner.is_alive() + assert runner.dispatch(_job(0, value=7)) is True + (result,) = _results(runner, 1) + assert isinstance(result, JobResult) and result.exception is None + assert isinstance(result.result, ZarrWriteResult) and result.result.z_index == 0 + runner.shutdown() + assert _squid_attrs(info)["acquisition_complete"] is True + import tensorstore as ts + + written = ts.open( + {"driver": "zarr3", "kvstore": {"driver": "file", "path": info.get_output_path("A1", 0)}} + ).result() + assert int(written[0, 0, 0, 0, 0].read().result()) == 7 + + +def test_counters_follow_jobrunner_contract(tmp_path): + from control.core.in_process_zarr_runner import InProcessZarrRunner + + pending_jobs, pending_bytes, capacity = create_backpressure_values() + runner = InProcessZarrRunner(zarr_writer_info=_info(tmp_path), bp_values=(pending_jobs, pending_bytes, capacity)) + runner.start() + capacity.clear() + runner.dispatch(_job(0)) + runner.dispatch(_job(1)) + _results(runner, 2) + _wait_settled(runner, (pending_jobs, pending_bytes, capacity)) + assert not runner.has_pending() + assert pending_jobs.value == 0 and pending_bytes.value == 0 + assert capacity.is_set() + runner.shutdown() + + +def test_dispatch_injects_zarr_writer_info_and_registry(tmp_path): + from control.core.in_process_zarr_runner import InProcessZarrRunner + + runner = InProcessZarrRunner() + runner.set_zarr_writer_info(_info(tmp_path)) + runner.set_acquisition_info(None) # accepted and ignored, as the workers call it + runner.start() + job = _job(0) + runner.dispatch(job) + assert job.zarr_writer_info is not None and job.registry is runner.registry + assert SaveZarrJob.default_registry.writers == {} + _results(runner, 1) + runner.shutdown() + + +def test_failed_write_reaches_output_queue_as_exception(tmp_path, monkeypatch): + from control.core.in_process_zarr_runner import InProcessZarrRunner + + monkeypatch.setattr( + SaveZarrJob, + "submit", + lambda self: (_failed_future(RuntimeError("disk full")), ZarrWriteResult(0, 0, 0, "BF", 0)), + ) + pending_jobs, pending_bytes, capacity = create_backpressure_values() + runner = InProcessZarrRunner(zarr_writer_info=_info(tmp_path), bp_values=(pending_jobs, pending_bytes, capacity)) + runner.start() + capacity.clear() + runner.dispatch(_job(0)) + (result,) = _results(runner, 1) + assert isinstance(result.exception, RuntimeError) and result.result is None + _wait_settled(runner, (pending_jobs, pending_bytes, capacity)) + assert pending_jobs.value == 0 and pending_bytes.value == 0 + runner.shutdown() + + +def test_submit_failure_is_reported_not_raised(tmp_path, monkeypatch): + from control.core.in_process_zarr_runner import InProcessZarrRunner + + def boom(self): + raise OSError("cannot create store") + + monkeypatch.setattr(SaveZarrJob, "submit", boom) + runner = InProcessZarrRunner(zarr_writer_info=_info(tmp_path)) + runner.start() + assert runner.dispatch(_job(0)) is True + (result,) = _results(runner, 1) + assert isinstance(result.exception, OSError) + _wait_settled(runner) + assert not runner.has_pending() + runner.shutdown() + + +def test_kill_then_shutdown_aborts_open_stores(tmp_path): + from control.core.in_process_zarr_runner import InProcessZarrRunner + + info = _info(tmp_path) + runner = InProcessZarrRunner(zarr_writer_info=info) + runner.start() + runner.dispatch(_job(0)) + _results(runner, 1) + runner.kill() + runner.shutdown() + attrs = _squid_attrs(info) + assert attrs["acquisition_complete"] is False and attrs["aborted"] is True + + +def test_shutdown_aborted_waits_for_writes_then_seals_aborted(tmp_path): + """A user abort keeps the frames already captured and marks the store aborted, not complete.""" + from control.core.in_process_zarr_runner import InProcessZarrRunner + + info = _info(tmp_path) + runner = InProcessZarrRunner(zarr_writer_info=info) + runner.start() + runner.dispatch(_job(0, value=5)) + runner.shutdown(aborted=True) + attrs = _squid_attrs(info) + assert attrs["acquisition_complete"] is False and attrs["aborted"] is True + import tensorstore as ts + + written = ts.open( + {"driver": "zarr3", "kvstore": {"driver": "file", "path": info.get_output_path("A1", 0)}} + ).result() + assert int(written[0, 0, 0, 0, 0].read().result()) == 5 + + +def test_shutdown_is_idempotent_and_before_start_is_safe(tmp_path): + from control.core.in_process_zarr_runner import InProcessZarrRunner + + runner = InProcessZarrRunner(zarr_writer_info=_info(tmp_path)) + runner.shutdown() # never started: nothing to do + runner.start() + runner.shutdown() + runner.shutdown() + assert not runner.is_alive() + + +def test_dispatch_after_shutdown_is_refused(tmp_path): + """A frame after the stores are sealed must not reopen a writer or raise the counters.""" + from control.core.in_process_zarr_runner import InProcessZarrRunner + + info = _info(tmp_path) + pending_jobs, pending_bytes, capacity = create_backpressure_values() + runner = InProcessZarrRunner(zarr_writer_info=info, bp_values=(pending_jobs, pending_bytes, capacity)) + runner.start() + runner.shutdown() + job = _job(0) + assert runner.dispatch(job) is False + assert job.zarr_writer_info is None and job.registry is None + assert not runner.has_pending() + assert pending_jobs.value == 0 and pending_bytes.value == 0 + assert runner.registry.writers == {} + assert not os.path.exists(info.base_path) + + +def test_two_runs_dispatched_from_one_persistent_thread(tmp_path): + # A real camera delivers every acquisition's frames on one read thread that lives as long as the camera. + import threading + + from control.core.in_process_zarr_runner import InProcessZarrRunner + + requests: "queue.Queue" = queue.Queue() + + def _camera_thread(): + while True: + item = requests.get() + if item is None: + return + runner, job, done = item + try: + done.put(runner.dispatch(job)) + except BaseException as e: + done.put(e) + + thread = threading.Thread(target=_camera_thread, name="camera-callback", daemon=True) + thread.start() + try: + for run in (1, 2): + info = _info(tmp_path / f"run{run}") + runner = InProcessZarrRunner(zarr_writer_info=info) + runner.start() + for z in (0, 1): + done: "queue.Queue" = queue.Queue() + requests.put((runner, _job(z, value=run), done)) + assert done.get(timeout=10) is True + results = _results(runner, 2) + runner.shutdown() + assert len(results) == 2 + assert [r.exception for r in results] == [None, None] + assert _squid_attrs(info)["acquisition_complete"] is True + finally: + requests.put(None) + thread.join(timeout=10) + + +def test_shutdown_with_a_write_still_in_flight_seals_aborted_and_reports_it(tmp_path, monkeypatch): + """A bounded shutdown must never stamp a store complete over a write it did not see finish.""" + from control.core.in_process_zarr_runner import InProcessZarrRunner + + release = threading.Event() + real_submit = SaveZarrJob.submit + + def blocked_submit(self): + future, result = real_submit(self) # open the store as usual, so there is a zarr.json to inspect + future.result() + blocked = concurrent.futures.Future() + threading.Thread(target=lambda: (release.wait(), blocked.set_result(None)), daemon=True).start() + return blocked, result + + monkeypatch.setattr(SaveZarrJob, "submit", blocked_submit) + info = _info(tmp_path) + runner = InProcessZarrRunner(zarr_writer_info=info) + runner.start() + runner.dispatch(_job(0)) + try: + assert runner.shutdown(timeout_s=0.1) is False + attrs = _squid_attrs(info) + assert attrs["acquisition_complete"] is False and attrs["aborted"] is True + assert runner.has_pending() + finally: + release.set() + _wait_settled(runner) + assert not runner.has_pending() diff --git a/software/tests/control/core/test_zarr_writer.py b/software/tests/control/core/test_zarr_writer.py index c4e79e9ff..372e43ff9 100644 --- a/software/tests/control/core/test_zarr_writer.py +++ b/software/tests/control/core/test_zarr_writer.py @@ -940,10 +940,6 @@ def test_write_multiple_frames(self, temp_dir): test_image = np.ones((32, 32), dtype=np.uint16) * (t * 100 + c * 10 + z) writer.write_frame(test_image, t=t, c=c, z=z) - # Wait for all writes - completed = writer.wait_for_pending() - assert completed >= 0 - writer.finalize() assert writer.is_finalized @@ -967,7 +963,6 @@ def test_write_and_verify_data(self, temp_dir): # Write a known pattern test_image = np.arange(32 * 32, dtype=np.uint16).reshape((32, 32)) writer.write_frame(test_image, t=0, c=0, z=0) - writer.wait_for_pending() writer.finalize() # Read back and verify @@ -1141,7 +1136,6 @@ def test_6d_write_multiple_fovs(self, temp_dir): test_image = np.ones((32, 32), dtype=np.uint16) * (fov + 1) * 100 writer.write_frame(test_image, t=0, c=0, z=0, fov=fov) - writer.wait_for_pending() writer.finalize() assert writer.is_finalized @@ -1170,7 +1164,6 @@ def test_6d_write_and_verify_data(self, temp_dir): test_images.append(test_image) writer.write_frame(test_image, t=0, c=0, z=0, fov=fov) - writer.wait_for_pending() writer.finalize() # Read back and verify each FOV - 6D indexing: [fov, t, c, z, y, x] @@ -1430,7 +1423,6 @@ def test_dtype_auto_conversion_to_uint16(self, temp_dir, source_dtype): # Write image with different dtype image = np.ones((32, 32), dtype=source_dtype) * 100 writer.write_frame(image, t=0, c=0, z=0) - writer.wait_for_pending() writer.finalize() # Read back and verify dtype @@ -1705,8 +1697,8 @@ def tracking_info(msg, *args, **kwargs): # Well metadata should be written exactly once (even with 3 FOVs) # The tracking set should have exactly one entry for this well assert ( - len(SaveZarrJob._hcs_wells_written) == 1 - ), f"Well metadata should be tracked once, got {len(SaveZarrJob._hcs_wells_written)}" + len(SaveZarrJob.default_registry.wells_written) == 1 + ), f"Well metadata should be tracked once, got {len(SaveZarrJob.default_registry.wells_written)}" # Verify the actual file exists plate_path = os.path.join(temp_dir, "plate.ome.zarr") @@ -2082,3 +2074,201 @@ def test_save_zarr_job_uses_the_settings_it_carries(tmp_path, monkeypatch): written = json.load(f)["attributes"]["_squid"] assert written["chunk_mode"] == ZarrChunkMode.TILED_256.value assert written["compression"] == ZarrCompression.BEST.value + + +class TestZarrWriterRegistry: + def test_registry_is_per_instance(self, tmp_path): + """Two registries never see each other's writers or HCS bookkeeping.""" + from control.core.job_processing import ZarrWriterRegistry + + a, b = ZarrWriterRegistry(), ZarrWriterRegistry() + a.plates_written.add("/plate") + assert "/plate" not in b.plates_written + assert a.writers is not b.writers + + def test_job_writes_through_its_registry(self, tmp_path): + from control.core.job_processing import ZarrWriterRegistry + + registry = ZarrWriterRegistry() + info = ZarrWriterInfo(base_path=str(tmp_path / "acq"), t_size=1, c_size=1, z_size=1) + job = SaveZarrJob( + capture_info=make_test_capture_info(region_id="A1", fov=0), + capture_image=JobImage(image_array=np.full((16, 16), 3, dtype=np.uint16)), + zarr_writer_info=info, + registry=registry, + ) + job.run() + assert list(registry.writers) == [info.get_output_path("A1", 0)] + assert SaveZarrJob.default_registry.writers == {} + assert registry.finalize_all() is True + import tensorstore as ts + + written = ts.open( + {"driver": "zarr3", "kvstore": {"driver": "file", "path": info.get_output_path("A1", 0)}} + ).result() + assert np.array_equal(written[0, 0, 0].read().result(), np.full((16, 16), 3, dtype=np.uint16)) + + def test_abort_all_seals_open_stores_incomplete(self, tmp_path): + import json + from control.core.job_processing import ZarrWriterRegistry + + registry = ZarrWriterRegistry() + info = ZarrWriterInfo(base_path=str(tmp_path / "acq"), t_size=2, c_size=1, z_size=1) + job = SaveZarrJob( + capture_info=make_test_capture_info(region_id="A1", fov=0), + capture_image=JobImage(image_array=np.zeros((16, 16), dtype=np.uint16)), + zarr_writer_info=info, + registry=registry, + ) + job.run() + registry.abort_all() + # The OME group metadata lives in the group above the "0" array (see _get_metadata_zarr_json_path). + group_json = os.path.join(os.path.dirname(info.get_output_path("A1", 0)), "zarr.json") + with open(group_json) as f: + attrs = json.load(f)["attributes"]["_squid"] + assert attrs["acquisition_complete"] is False and attrs["aborted"] is True + assert registry.writers == {} + + +class TestZarrWriterSubmitFrame: + def _writer(self, tmp_path): + from control.core.zarr_writer import ZarrAcquisitionConfig, ZarrWriter + + config = ZarrAcquisitionConfig( + output_path=str(tmp_path / "store.zarr"), shape=(1, 1, 2, 16, 16), dtype=np.uint16, pixel_size_um=1.0 + ) + writer = ZarrWriter(config) + writer.initialize() + return writer + + def test_submit_frame_returns_the_write_future(self, tmp_path): + writer = self._writer(tmp_path) + future = writer.submit_frame(np.full((16, 16), 9, dtype=np.uint16), t=0, c=0, z=1) + assert future is not None + future.result() + writer.finalize() + import tensorstore as ts + + written = ts.open( + {"driver": "zarr3", "kvstore": {"driver": "file", "path": str(tmp_path / "store.zarr")}} + ).result() + assert int(written[0, 0, 1, 0, 0].read().result()) == 9 + + def test_submit_frame_validates_like_write_frame(self, tmp_path): + writer = self._writer(tmp_path) + with pytest.raises(ValueError): + writer.submit_frame(np.zeros((16, 16), dtype=np.uint16), t=0, c=0, z=2) + writer.finalize() + with pytest.raises(RuntimeError): + writer.submit_frame(np.zeros((16, 16), dtype=np.uint16), t=0, c=0, z=0) + + +class TestSaveZarrJobSubmit: + def _job(self, tmp_path, value): + from control.core.job_processing import ZarrWriterRegistry + + info = ZarrWriterInfo( + base_path=str(tmp_path / "acq"), t_size=1, c_size=1, z_size=2, channel_names=["BF"], pixel_size_um=1.0 + ) + return ( + SaveZarrJob( + capture_info=make_test_capture_info(region_id="A1", fov=0, z_index=value), + capture_image=JobImage(image_array=np.full((16, 16), value + 1, dtype=np.uint16)), + zarr_writer_info=info, + registry=ZarrWriterRegistry(), + ), + info, + ) + + def test_submit_returns_a_future_and_the_result_to_report(self, tmp_path): + job, info = self._job(tmp_path, 0) + future, result = job.submit() + assert isinstance(result, ZarrWriteResult) and result.z_index == 0 + future.result() # the write completes without run() ever waiting for it + import tensorstore as ts + + written = ts.open( + {"driver": "zarr3", "kvstore": {"driver": "file", "path": info.get_output_path("A1", 0)}} + ).result() + assert int(written[0, 0, 0, 0, 0].read().result()) == 1 + + def test_run_is_submit_then_wait(self, tmp_path, monkeypatch): + job, _ = self._job(tmp_path, 1) + waited = [] + + class _Future: + def result(self): + waited.append(True) + + expected = ZarrWriteResult(fov=0, time_point=0, z_index=1, channel_name="c", region_idx=0) + monkeypatch.setattr(job, "submit", lambda: (_Future(), expected)) + result = job.run() + assert waited == [True] and result.z_index == 1 + + def test_submit_in_simulated_io_mode_returns_no_future(self, tmp_path, monkeypatch): + import control._def + + monkeypatch.setattr(control._def, "SIMULATED_DISK_IO_ENABLED", True) + job, _ = self._job(tmp_path, 0) + future, result = job.submit() + assert future is None and isinstance(result, ZarrWriteResult) + + +class TestZarrWriterAcrossThreads: + """A camera's frame callback thread outlives an acquisition: the writer must not leave per-thread state behind.""" + + def test_initialize_on_persistent_thread_after_finalize_elsewhere(self, tmp_path): + import queue + import threading + + import tensorstore as ts + + from control.core.zarr_writer import ZarrAcquisitionConfig, ZarrWriter + + requests: "queue.Queue" = queue.Queue() + + def _camera_thread(): + while True: + item = requests.get() + if item is None: + return + work, reply = item + try: + reply.put((work(), None)) + except BaseException as e: # hand the failure to the main thread + reply.put((None, e)) + + thread = threading.Thread(target=_camera_thread, name="camera-callback", daemon=True) + thread.start() + + def _on_camera_thread(work): + reply: "queue.Queue" = queue.Queue() + requests.put((work, reply)) + value, error = reply.get(timeout=60) + if error is not None: + raise error + return value + + def _open_and_write(path, value): + config = ZarrAcquisitionConfig( + output_path=path, shape=(1, 1, 1, 16, 16), dtype=np.uint16, pixel_size_um=1.0 + ) + writer = ZarrWriter(config) + writer.initialize() + writer.write_frame(np.full((16, 16), value, dtype=np.uint16), t=0, c=0, z=0) + return writer + + try: + paths = [str(tmp_path / "run1.zarr"), str(tmp_path / "run2.zarr")] + for value, path in enumerate(paths, start=1): + writer = _on_camera_thread(lambda: _open_and_write(path, value)) + writer.finalize() # end of the run: finalized on another thread + finally: + requests.put(None) + thread.join(timeout=60) + + for value, path in enumerate(paths, start=1): + with open(os.path.join(path, "zarr.json")) as f: + assert json.load(f)["attributes"]["_squid"]["acquisition_complete"] is True + written = ts.open({"driver": "zarr3", "kvstore": {"driver": "file", "path": path}}).result() + assert int(written[0, 0, 0, 0, 0].read().result()) == value diff --git a/software/tests/control/test_MultiPointController.py b/software/tests/control/test_MultiPointController.py index 1005b2378..e926cdc7c 100644 --- a/software/tests/control/test_MultiPointController.py +++ b/software/tests/control/test_MultiPointController.py @@ -723,3 +723,211 @@ def test_run_acquisition_continues_when_stop_live_times_out_but_mcu_recovers(): mpc.thread.join(10) assert tt.image_count == mpc.get_acquisition_image_count() + + +def _zarr_jsons_with_squid(mpc): + """The _squid attributes of every OME group zarr.json of the run just finished (array zarr.json files carry none).""" + import json + from pathlib import Path + + out = [] + for p in (Path(mpc.base_path) / mpc.experiment_ID).rglob("zarr.json"): + with open(p) as f: + attrs = json.load(f).get("attributes", {}).get("_squid") + if attrs is not None: + out.append(attrs) + return out + + +def _zarr_v3_controller(monkeypatch): + """A simulated microscope and controller set up for a Zarr v3 multipoint run, with its tracker.""" + pytest.importorskip("tensorstore") + control._def.MERGE_CHANNELS = False + monkeypatch.setattr(control._def, "FILE_SAVING_OPTION", control._def.FileSavingOption.ZARR_V3) + scope = control.microscope.Microscope.build_from_global_config(True) + tt = TestAcquisitionTracker() + mpc = ts.get_test_multi_point_controller(microscope=scope, callbacks=tt.get_callbacks()) + add_some_coordinates(mpc) + select_some_configs(mpc, scope.objective_store.current_objective) + return mpc, tt + + +def _make_zarr_writes_fail(monkeypatch, message="disk full"): + """Every SaveZarrJob opens its store as usual (so there is a zarr.json to inspect) but its write fails.""" + import concurrent.futures + from control.core.job_processing import SaveZarrJob + + real_submit = SaveZarrJob.submit + + def failing_submit(self): + future, result = real_submit(self) + if future is not None: + future.result() + failed = concurrent.futures.Future() + failed.set_exception(RuntimeError(message)) + return failed, result + + monkeypatch.setattr(SaveZarrJob, "submit", failing_submit) + + +def test_zarr_v3_multipoint_saves_in_process_and_seals(monkeypatch): + """A Zarr v3 multipoint run saves through the in-process runner and ends with acquisition_complete=True.""" + from control.core.in_process_zarr_runner import InProcessZarrRunner + + mpc, tt = _zarr_v3_controller(monkeypatch) + mpc.run_acquisition() + assert tt.finished_event.wait(120) + assert isinstance(mpc.multiPointWorker._job_runners[0][1], InProcessZarrRunner) + attrs = _zarr_jsons_with_squid(mpc) + assert attrs and all(a["acquisition_complete"] is True for a in attrs) + + +def test_zarr_is_sealed_with_multiprocessing_off(monkeypatch): + """The in-process fallback used to leave the store unsealed; it must not any more.""" + monkeypatch.setattr(control._def.Acquisition, "USE_MULTIPROCESSING", False) + mpc, tt = _zarr_v3_controller(monkeypatch) + mpc.run_acquisition() + assert tt.finished_event.wait(120) + attrs = _zarr_jsons_with_squid(mpc) + assert attrs and all(a["acquisition_complete"] is True for a in attrs) + + +def test_abort_mid_zarr_run_seals_store_aborted(monkeypatch): + """Stop during a Zarr run: the run ends promptly and the stores say aborted, not complete.""" + import time + + mpc, tt = _zarr_v3_controller(monkeypatch) + mpc.set_NZ(5) + mpc.run_acquisition() + assert tt.started_event.wait(30) + deadline = time.monotonic() + 30 + while tt.image_count < 3 and time.monotonic() < deadline: + time.sleep(0.05) + assert tt.image_count >= 3 + t_abort = time.monotonic() + mpc.request_abort_aquisition() + assert tt.finished_event.wait(60) + assert time.monotonic() - t_abort < 30, "abort must not wait out the job timeouts" + attrs = _zarr_jsons_with_squid(mpc) + assert attrs and all(a["acquisition_complete"] is False and a.get("aborted") is True for a in attrs) + + +def test_zarr_v3_run_leaves_the_warm_subprocess_for_tiff(monkeypatch): + """A Zarr v3 acquisition neither consumes nor restarts the pre-warmed subprocess, and starts none of its own.""" + from control.core import job_processing + + mpc, tt = _zarr_v3_controller(monkeypatch) + warm_before = mpc._prewarmed_job_runner + assert warm_before is not None + started = [] # only subprocesses started from here on count; the controller's own pre-warm is above + monkeypatch.setattr(job_processing.JobRunner, "start", lambda self: started.append(self)) + mpc.run_acquisition() + assert tt.finished_event.wait(120) + assert started == [], "a Zarr v3 run must not start a save subprocess" + assert mpc._prewarmed_job_runner is warm_before + + +def test_three_zarr_runs_in_one_session_all_sealed(monkeypatch): + """Back-to-back Zarr runs in one process each get fresh writers and each store is sealed complete.""" + mpc, tt = _zarr_v3_controller(monkeypatch) + for i in range(3): + tt.started_event.clear() + tt.finished_event.clear() + mpc.start_new_experiment(f"zarr_session_{i}") + mpc.run_acquisition() + assert tt.finished_event.wait(120), f"run {i} did not finish" + attrs = _zarr_jsons_with_squid(mpc) + assert attrs and all(a["acquisition_complete"] is True for a in attrs), f"run {i}" + + +def test_zarr_write_error_ends_the_run_as_error(monkeypatch, caplog): + """A failed Zarr write must end the run as an error and seal the store incomplete, not finish quietly.""" + import logging + + mpc, tt = _zarr_v3_controller(monkeypatch) + _make_zarr_writes_fail(monkeypatch) + with caplog.at_level(logging.ERROR): + mpc.run_acquisition() + assert tt.finished_event.wait(120) + assert mpc.multiPointWorker._abort_cause == "error" + assert any("disk full" in r.getMessage() for r in caplog.records) + attrs = _zarr_jsons_with_squid(mpc) + assert attrs and all(a["acquisition_complete"] is False for a in attrs) + + +def test_dead_tiff_subprocess_aborts_acquisition(caplog): + """A save subprocess that dies with jobs pending must abort the run, not let it finish with data missing.""" + import logging + import queue + from unittest.mock import MagicMock + from control.core.job_processing import SaveImageJob + from control.core.multi_point_worker import MultiPointWorker + + worker = MultiPointWorker.__new__(MultiPointWorker) # only _summarize_runner_outputs' own fields are needed + worker._log = logging.getLogger("squid.MultiPointWorker") + worker._acquisition_error_count = 0 + worker._slack_notifier = None + runner = MagicMock() + runner.output_queue.return_value = MagicMock(get_nowait=MagicMock(side_effect=queue.Empty)) + runner.is_alive.return_value = False + runner.has_pending.return_value = True + runner.exitcode = -6 + worker._job_runners = [(SaveImageJob, runner)] + worker._abort_due_to_error = MagicMock() + with caplog.at_level(logging.ERROR): + worker._summarize_runner_outputs() + worker._abort_due_to_error.assert_called_once() + assert any("save subprocess exited" in r.getMessage() for r in caplog.records) + + +def test_close_shuts_down_in_process_runner_instead_of_terminating_it(): + """close() on an abnormal shutdown must seal an in-process runner's stores, not treat it as a process.""" + import logging + from unittest.mock import MagicMock + from control.core.in_process_zarr_runner import InProcessZarrRunner + from control.core.job_processing import SaveZarrJob + from control.core.multi_point_controller import MultiPointController + + mpc = MultiPointController.__new__(MultiPointController) # close() only needs the fields set below + mpc._log = logging.getLogger("squid.MultiPointController") + mpc._prewarmed_job_runner = None + mpc._prewarmed_bp_values = None + mpc._memory_monitor = None + mpc.thread = None + runner = MagicMock(spec=InProcessZarrRunner) + runner.is_alive.return_value = True + runner.terminate = MagicMock() # not part of InProcessZarrRunner; present only to prove it is not called + worker = MagicMock() + worker._job_runners = [(SaveZarrJob, runner)] + mpc.multiPointWorker = worker + mpc.close() + runner.shutdown.assert_called_once_with(timeout_s=MultiPointController._PROCESS_TERMINATE_TIMEOUT_S, aborted=True) + runner.terminate.assert_not_called() + + +def test_zarr_write_error_surfacing_at_the_end_ends_the_run_as_error(monkeypatch, caplog): + """A failed Zarr write that only _finish_jobs' final drain sees must end the run as an error, store incomplete. + + Why this test exists: the per-FOV poll already aborts on a failure it sees; this one checks the failure that only + the final drain sees, by holding every result back from the per-FOV poll (drain_all=False). + """ + import logging + from control.core.multi_point_worker import MultiPointWorker, SummarizeResult + + real_summarize = MultiPointWorker._summarize_runner_outputs + + def final_drain_only(self, drain_all=False): + if not drain_all: # the per-FOV poll in the acquisition loop sees nothing + return SummarizeResult(none_failed=True, had_results=False) + return real_summarize(self, drain_all=drain_all) + + mpc, tt = _zarr_v3_controller(monkeypatch) + _make_zarr_writes_fail(monkeypatch) + monkeypatch.setattr(MultiPointWorker, "_summarize_runner_outputs", final_drain_only) + with caplog.at_level(logging.ERROR): + mpc.run_acquisition() + assert tt.finished_event.wait(120) + assert mpc.multiPointWorker._abort_cause == "error" + assert any("A save job failed, aborting acquisition" in r.getMessage() for r in caplog.records) + attrs = _zarr_jsons_with_squid(mpc) + assert attrs and all(a["acquisition_complete"] is False for a in attrs) diff --git a/software/tools/benchmark_save_formats.py b/software/tools/benchmark_save_formats.py index 1066e0e7c..9372469b3 100644 --- a/software/tools/benchmark_save_formats.py +++ b/software/tools/benchmark_save_formats.py @@ -49,10 +49,7 @@ def benchmark_zarr_write(num_frames: int, image_shape: tuple, detailed: bool = F writer.write_frame(image, t=0, c=c, z=0) queue_time = time.perf_counter() - queue_start - # Wait for pending writes - wait_start = time.perf_counter() - writer.wait_for_pending() - wait_time = time.perf_counter() - wait_start + wait_time = 0.0 # write_frame() blocks until the write is committed finalize_start = time.perf_counter() writer.finalize() @@ -114,9 +111,7 @@ def benchmark_zarr_no_compression(num_frames: int, image_shape: tuple, detailed: writer.write_frame(image, t=0, c=c, z=0) queue_time = time.perf_counter() - queue_start - wait_start = time.perf_counter() - writer.wait_for_pending() - wait_time = time.perf_counter() - wait_start + wait_time = 0.0 # write_frame() blocks until the write is committed finalize_start = time.perf_counter() writer.finalize() @@ -186,7 +181,6 @@ async def _open(): start = time.perf_counter() for c in range(num_frames): writer.write_frame(image, t=0, c=c, z=0) - writer.wait_for_pending() elapsed = time.perf_counter() - start return elapsed