From f4b72280fb885ecb0aad02aa06a78378a68703f9 Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 01:35:05 -0700 Subject: [PATCH 01/16] refactor(zarr): SaveZarrJob keeps its writers in a per-acquisition registry The open writers and the HCS plate/well bookkeeping were class-level, which meant per-acquisition only because each acquisition had its own save subprocess. ZarrWriterRegistry makes that explicit and injectable, so an acquisition saving in the main process gets its own. The subprocess keeps using SaveZarrJob.default_registry; clear_writers() and finalize_all_writers() delegate to it. Co-Authored-By: Claude Fable 5.1 --- software/control/core/job_processing.py | 123 ++++++++++-------- .../tests/control/core/test_zarr_writer.py | 58 ++++++++- 2 files changed, 124 insertions(+), 57 deletions(-) diff --git a/software/control/core/job_processing.py b/software/control/core/job_processing.py index ec27df65a..4e00697dd 100644 --- a/software/control/core/job_processing.py +++ b/software/control/core/job_processing.py @@ -490,6 +490,55 @@ class ZarrWriteResult: region_idx: int = 0 +@dataclass +class ZarrWriterRegistry: + """The open ZarrWriters of one acquisition, keyed by writer_key, plus the HCS metadata already written. + + 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. + """ + + writers: Dict[str, "ZarrWriter"] = field(default_factory=dict) + plates_written: Set[str] = field(default_factory=set) + wells_written: Set[str] = field(default_factory=set) + + _log: ClassVar = squid.logging.get_logger("ZarrWriterRegistry") + + 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 + + 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(self.writers.values()): + if writer.is_initialized and not writer.is_finalized: + try: + writer.abort() + except Exception as e: + self._log.warning(f"Error aborting writer during clear: {e}") + finally: + self._clear() + + def _clear(self) -> None: + self.writers.clear() + self.plates_written.clear() + self.wells_written.clear() + + @dataclass class SaveZarrJob(Job): """Job for saving images to Zarr v3 format using TensorStore. @@ -500,67 +549,31 @@ class SaveZarrJob(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) - # 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"]] = {} + default_registry: ClassVar["ZarrWriterRegistry"] = ZarrWriterRegistry() - # 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() + @property + def _registry(self) -> "ZarrWriterRegistry": + return self.registry if self.registry is not None else SaveZarrJob.default_registry @classmethod def clear_writers(cls) -> None: - """Clear all zarr writers, aborting any that are still active. - - Call at start of new acquisition to ensure clean state. - Uses try-finally to guarantee dictionaries are cleared even if abort fails. - """ - try: - for writer in list(cls._zarr_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}") - finally: - # Always clear dictionaries, even if abort loop fails - cls._zarr_writers.clear() - cls._hcs_plate_written.clear() - cls._hcs_wells_written.clear() + """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 all active zarr writers. - - Call at end of acquisition to ensure all data is written. - - 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 + """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,20 +585,20 @@ 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: @@ -684,7 +697,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 +741,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 +752,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 diff --git a/software/tests/control/core/test_zarr_writer.py b/software/tests/control/core/test_zarr_writer.py index c4e79e9ff..76784a283 100644 --- a/software/tests/control/core/test_zarr_writer.py +++ b/software/tests/control/core/test_zarr_writer.py @@ -1705,8 +1705,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 +2082,57 @@ 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 == {} From 7392656174fd6f0a55b2e3718bfa1b365637824f Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 01:37:34 -0700 Subject: [PATCH 02/16] feat(zarr): non-blocking writes: ZarrWriter.submit_frame() and SaveZarrJob.submit() Both return the tensorstore write future instead of waiting on it; write_frame() and run() are now submit then wait, unchanged in behavior. Needed by the in-process runner, which must not block the camera callback on a write. Co-Authored-By: Claude Fable 5.1 --- software/control/core/job_processing.py | 35 +++++--- software/control/core/zarr_writer.py | 46 +++++----- .../tests/control/core/test_zarr_writer.py | 84 +++++++++++++++++++ 3 files changed, 133 insertions(+), 32 deletions(-) diff --git a/software/control/core/job_processing.py b/software/control/core/job_processing.py index 4e00697dd..69be9e9a6 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 @@ -602,6 +602,18 @@ def _write_hcs_metadata_if_needed(self, region_id: str, fov: int) -> None: 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. " @@ -667,13 +679,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 @@ -760,13 +772,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! diff --git a/software/control/core/zarr_writer.py b/software/control/core/zarr_writer.py index 483c7a66e..3a871d9e9 100644 --- a/software/control/core/zarr_writer.py +++ b/software/control/core/zarr_writer.py @@ -651,17 +651,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,19 +691,29 @@ 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}") + 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 write to complete (blocking) - # TensorStore futures have a .result() method that blocks until complete - future.result() - if config.ndim == 5: + 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. + + This submits the write via :meth:`submit_frame` and waits on the resulting future before + returning, so the data is visible to other processes reading the same zarr store once this + returns. Callers that write many frames back-to-back should prefer :meth:`submit_frame` + with a bounded number of futures in flight. + + Args: + image: 2D image array (Y, X) + t: Time point index + c: Channel index + z: Z-slice index + fov: FOV index (required for 6D datasets, ignored for 5D) + """ + self.submit_frame(image, t, c, z, fov).result() + if self._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}") diff --git a/software/tests/control/core/test_zarr_writer.py b/software/tests/control/core/test_zarr_writer.py index 76784a283..030910585 100644 --- a/software/tests/control/core/test_zarr_writer.py +++ b/software/tests/control/core/test_zarr_writer.py @@ -2136,3 +2136,87 @@ def test_abort_all_seals_open_stores_incomplete(self, tmp_path): 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) From d687740e97009c9a333bb2f7ca30f3ba56fc1829 Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 01:39:44 -0700 Subject: [PATCH 03/16] feat(zarr): InProcessZarrRunner saves Zarr in the acquisition process JobRunner's interface, SaveZarrJob.submit() underneath: dispatch submits the write and returns in about a millisecond; a drain thread waits for each write, keeps the backpressure counters and reports JobResults on the output queue. Not wired to the workers yet. Co-Authored-By: Claude Fable 5.1 --- .../control/core/in_process_zarr_runner.py | 160 +++++++++++++++ .../core/test_in_process_zarr_runner.py | 189 ++++++++++++++++++ 2 files changed, 349 insertions(+) create mode 100644 software/control/core/in_process_zarr_runner.py create mode 100644 software/tests/control/core/test_in_process_zarr_runner.py 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..1dbdd3796 --- /dev/null +++ b/software/control/core/in_process_zarr_runner.py @@ -0,0 +1,160 @@ +"""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 Record + Z-Stack recording +writer does), 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 queue +import threading +from typing import Any, Optional, Tuple + +import squid.logging +from control.core.backpressure import BackpressureValues +from control.core.job_processing import AcquisitionInfo, JobResult, SaveZarrJob, ZarrWriterInfo, ZarrWriterRegistry + +_SENTINEL = object() + + +class InProcessZarrRunner: + runs_in_process = True + + 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, future, image_bytes, result), in submission order + self._output: "queue.Queue" = queue.Queue() + self._pending = 0 + self._pending_lock = threading.Lock() + self._thread = threading.Thread(target=self._drain, name="InProcessZarrRunner", daemon=True) + self._started = False + self._finished = False + self._killed = 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: + 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 + self._count_dispatched(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") + self._submitted.put_nowait((job, _FailedFuture(e), image_bytes, None)) + return True + self._submitted.put_nowait((job, 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: + """Stop tracking the outstanding writes; shutdown() then seals the stores as aborted.""" + self._killed = True + + def shutdown(self, timeout_s: float = 1.0, aborted: bool = False) -> None: + """Wait for the outstanding writes, then seal every store. + + Complete when the acquisition ran to its end; `acquisition_complete: False, aborted: True` + when it was aborted (by the user or by an error) or after kill(). Captured frames are + written either way. + """ + if not self._started or self._finished: + return + self._finished = True + self._submitted.put_nowait(_SENTINEL) + self._thread.join(timeout=max(timeout_s, 1.0)) + if self._thread.is_alive(): + self._log.warning("drain thread still waiting on a write after the shutdown timeout; sealing anyway") + if self._killed or aborted: + self.registry.abort_all() + elif not self.registry.finalize_all(): + self._log.error("ZARR FINALIZATION INCOMPLETE - Some data may not be saved correctly") + + # ---- the drain thread + + def _drain(self) -> None: + while True: + item = self._submitted.get() + if item is _SENTINEL: + return + job, future, image_bytes, result = item + try: + if future is not None: + future.result() + self._output.put_nowait(JobResult(job_id=job.job_id, result=result, exception=None)) + except Exception as e: + self._log.exception(f"Job {job.job_id} failed! Returning exception result.") + self._output.put_nowait(JobResult(job_id=job.job_id, result=None, exception=e)) + finally: + self._count_completed(image_bytes) + + def _count_dispatched(self, image_bytes: int) -> None: + with self._pending_lock: + self._pending += 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 + + def _count_completed(self, image_bytes: int) -> None: + with self._pending_lock: + self._pending -= 1 + 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) + with self._bp_pending_bytes.get_lock(): + self._bp_pending_bytes.value = max(0, self._bp_pending_bytes.value - image_bytes) + if self._bp_capacity_event is not None: + self._bp_capacity_event.set() + + +class _FailedFuture: + """A write that failed before it was submitted, reported through the same path as a failed write.""" + + def __init__(self, error: Exception): + self._error = error + + def result(self) -> Any: + raise self._error 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..97989b05d --- /dev/null +++ b/software/tests/control/core/test_in_process_zarr_runner.py @@ -0,0 +1,189 @@ +"""InProcessZarrRunner: SaveZarrJob in this process, with JobRunner's interface and counters.""" + +import json +import os +import queue +import time + +import numpy as np +import pytest + +import squid.abc +from control.core.backpressure import create_backpressure_values +from control.core.job_processing import CaptureInfo, JobImage, JobResult, SaveZarrJob, ZarrWriteResult, ZarrWriterInfo +from control.models import AcquisitionChannel, CameraSettings, IlluminationSettings + +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=CaptureInfo( + position=squid.abc.Pos(x_mm=0.0, y_mm=0.0, z_mm=0.0, theta_rad=None), + z_index=z, + capture_time=time.time(), + configuration=AcquisitionChannel( + name="BF", + illumination_settings=IlluminationSettings(illumination_channel="LED", intensity=50.0), + camera_settings=CameraSettings(exposure_time_ms=10.0, gain_mode=1.0), + ), + save_directory="/tmp/unused", + file_id=f"f_{z}", + region_id="A1", + fov=0, + configuration_idx=0, + time_point=0, + ), + capture_image=JobImage(image_array=np.full((8, 8), value, dtype=np.uint16)), + ) + + +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 _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) + 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 + + class _FailingFuture: + def result(self): + raise RuntimeError("disk full") + + monkeypatch.setattr(SaveZarrJob, "submit", lambda self: (_FailingFuture(), 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() + runner.dispatch(_job(0)) + (result,) = _results(runner, 1) + assert isinstance(result.exception, RuntimeError) and result.result is None + 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) + 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() From 037937ded98f7fce9854e366fdd42ca940442f5b Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 01:41:05 -0700 Subject: [PATCH 04/16] fix(zarr): InProcessZarrRunner refuses a dispatch after shutdown A frame dispatched after the stores are sealed would open a writer nothing finalizes and leave the backpressure counters raised. dispatch() now returns False, which the worker already treats as a reason to abort. Co-Authored-By: Claude Fable 5.1 --- .../control/core/in_process_zarr_runner.py | 4 ++++ .../core/test_in_process_zarr_runner.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/software/control/core/in_process_zarr_runner.py b/software/control/core/in_process_zarr_runner.py index 1dbdd3796..1f86ae3eb 100644 --- a/software/control/core/in_process_zarr_runner.py +++ b/software/control/core/in_process_zarr_runner.py @@ -65,6 +65,10 @@ def set_zarr_writer_info(self, zarr_writer_info: ZarrWriterInfo) -> None: # ---- the hot path: called from the camera frame callback def dispatch(self, job: SaveZarrJob) -> bool: + """Submit the write and return; False once shutdown() has sealed the stores (the worker aborts).""" + 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 diff --git a/software/tests/control/core/test_in_process_zarr_runner.py b/software/tests/control/core/test_in_process_zarr_runner.py index 97989b05d..8dd5cf95d 100644 --- a/software/tests/control/core/test_in_process_zarr_runner.py +++ b/software/tests/control/core/test_in_process_zarr_runner.py @@ -187,3 +187,21 @@ def test_shutdown_is_idempotent_and_before_start_is_safe(tmp_path): 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) From c802dd161ce09c38e6188ad1352ac005a647c173 Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 01:44:20 -0700 Subject: [PATCH 05/16] fix(zarr): InProcessZarrRunner dispatch and shutdown exclude each other; settle counters in tests A dispatch racing shutdown could land behind the sentinel and never drain. The tests also read the backpressure counters right after the result was queued, before the drain thread's finally had run. Co-Authored-By: Claude Fable 5.1 --- .../control/core/in_process_zarr_runner.py | 52 ++++++++++--------- .../core/test_in_process_zarr_runner.py | 19 +++++++ 2 files changed, 47 insertions(+), 24 deletions(-) diff --git a/software/control/core/in_process_zarr_runner.py b/software/control/core/in_process_zarr_runner.py index 1f86ae3eb..da90d4fd3 100644 --- a/software/control/core/in_process_zarr_runner.py +++ b/software/control/core/in_process_zarr_runner.py @@ -11,7 +11,7 @@ import queue import threading -from typing import Any, Optional, Tuple +from typing import Any, Optional import squid.logging from control.core.backpressure import BackpressureValues @@ -36,6 +36,8 @@ def __init__( self._output: "queue.Queue" = queue.Queue() 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 @@ -65,25 +67,26 @@ def set_zarr_writer_info(self, zarr_writer_info: ZarrWriterInfo) -> None: # ---- the hot path: called from the camera frame callback def dispatch(self, job: SaveZarrJob) -> bool: - """Submit the write and return; False once shutdown() has sealed the stores (the worker aborts).""" - 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 - self._count_dispatched(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") - self._submitted.put_nowait((job, _FailedFuture(e), image_bytes, None)) + """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 + self._count_dispatched(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") + self._submitted.put_nowait((job, _FailedFuture(e), image_bytes, None)) + return True + self._submitted.put_nowait((job, future, image_bytes, result)) return True - self._submitted.put_nowait((job, future, image_bytes, result)) - return True def output_queue(self) -> "queue.Queue": return self._output @@ -93,7 +96,7 @@ def has_pending(self) -> bool: return self._pending > 0 def kill(self) -> None: - """Stop tracking the outstanding writes; shutdown() then seals the stores as aborted.""" + """Mark the run aborted; shutdown() then seals the stores as aborted.""" self._killed = True def shutdown(self, timeout_s: float = 1.0, aborted: bool = False) -> None: @@ -103,10 +106,11 @@ def shutdown(self, timeout_s: float = 1.0, aborted: bool = False) -> None: when it was aborted (by the user or by an error) or after kill(). Captured frames are written either way. """ - if not self._started or self._finished: - return - self._finished = True - self._submitted.put_nowait(_SENTINEL) + with self._state_lock: + if not self._started or self._finished: + return + self._finished = True + self._submitted.put_nowait(_SENTINEL) self._thread.join(timeout=max(timeout_s, 1.0)) if self._thread.is_alive(): self._log.warning("drain thread still waiting on a write after the shutdown timeout; sealing anyway") diff --git a/software/tests/control/core/test_in_process_zarr_runner.py b/software/tests/control/core/test_in_process_zarr_runner.py index 8dd5cf95d..1154b628e 100644 --- a/software/tests/control/core/test_in_process_zarr_runner.py +++ b/software/tests/control/core/test_in_process_zarr_runner.py @@ -52,6 +52,21 @@ def _results(runner, n, timeout=10.0): 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") @@ -90,6 +105,7 @@ def test_counters_follow_jobrunner_contract(tmp_path): 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() @@ -122,9 +138,11 @@ def result(self): 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() @@ -141,6 +159,7 @@ def boom(self): 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() From 596624520948f0197535d266019ccd296ae3ff19 Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 01:49:25 -0700 Subject: [PATCH 06/16] feat(acquisition): save Zarr v3 in the acquisition process SaveZarrJob never gets a forked subprocess any more: MultiPointWorker gives it an InProcessZarrRunner. The TIFF formats keep the subprocess. Also seals the store at the end when multiprocessing is off, which the old in-process path never did, and seals an aborted run's stores as aborted. Co-Authored-By: Claude Fable 5.1 --- software/control/core/multi_point_worker.py | 44 ++++++++++- .../control/test_MultiPointController.py | 77 +++++++++++++++++++ 2 files changed, 117 insertions(+), 4 deletions(-) diff --git a/software/control/core/multi_point_worker.py b/software/control/core/multi_point_worker.py index dfcaa8f84..b22b729f9 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). @@ -698,9 +723,18 @@ 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() + aborted = bool(self.abort_requested_fn()) for job_class, job_runner in active_runners: + if getattr(job_runner, "runs_in_process", False): + # 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: + job_runner.shutdown(remaining_time, aborted=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() @@ -1418,6 +1452,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/tests/control/test_MultiPointController.py b/software/tests/control/test_MultiPointController.py index 1005b2378..f96c5d3d4 100644 --- a/software/tests/control/test_MultiPointController.py +++ b/software/tests/control/test_MultiPointController.py @@ -723,3 +723,80 @@ 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"): + attrs = json.load(open(p)).get("attributes", {}).get("_squid") + if attrs is not None: + out.append(attrs) + return out + + +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.""" + pytest.importorskip("tensorstore") + from control.core.in_process_zarr_runner import InProcessZarrRunner + + 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) + 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.""" + pytest.importorskip("tensorstore") + + control._def.MERGE_CHANNELS = False + monkeypatch.setattr(control._def, "FILE_SAVING_OPTION", control._def.FileSavingOption.ZARR_V3) + monkeypatch.setattr(control._def.Acquisition, "USE_MULTIPROCESSING", False) + 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) + 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.""" + pytest.importorskip("tensorstore") + import time + + 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()) + mpc.set_NZ(5) + add_some_coordinates(mpc) + select_some_configs(mpc, scope.objective_store.current_objective) + 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) From 0d9d2ae04a19bef76e11ed1942892b16669eefb6 Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 01:54:12 -0700 Subject: [PATCH 07/16] fix(acquisition): a Zarr v3 run keeps the warm save subprocess for a later TIFF run With Zarr saved in-process, handing the pre-warmed subprocess to a Zarr run only released it and forked a replacement. Nothing is forked for Zarr now, so a process that has used tensorstore (the NDViewer, or a recording) can no longer lose a later Zarr acquisition to the tensorstore fork abort. Co-Authored-By: Claude Fable 5.1 --- .../control/core/multi_point_controller.py | 16 ++++--- .../control/test_MultiPointController.py | 43 +++++++++++++++++++ 2 files changed, 52 insertions(+), 7 deletions(-) diff --git a/software/control/core/multi_point_controller.py b/software/control/core/multi_point_controller.py index 286aac9e8..753413328 100644 --- a/software/control/core/multi_point_controller.py +++ b/software/control/core/multi_point_controller.py @@ -23,6 +23,7 @@ from control.core.live_controller import LiveController from control.microscope import Microscope from control.core.multi_point_worker import MultiPointWorker +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 +309,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 diff --git a/software/tests/control/test_MultiPointController.py b/software/tests/control/test_MultiPointController.py index f96c5d3d4..ae7109dae 100644 --- a/software/tests/control/test_MultiPointController.py +++ b/software/tests/control/test_MultiPointController.py @@ -800,3 +800,46 @@ def test_abort_mid_zarr_run_seals_store_aborted(monkeypatch): 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.""" + pytest.importorskip("tensorstore") + from control.core import job_processing + + 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()) + 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)) + add_some_coordinates(mpc) + select_some_configs(mpc, scope.objective_store.current_objective) + 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.""" + 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) + 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}" From 739cba5d079e48cb4eae6cdf8daf0cf8526ef779 Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 02:10:11 -0700 Subject: [PATCH 08/16] fix(acquisition): a save failure ends the acquisition as an error A dead save subprocess, or a failed Zarr write surfacing at the end, used to be waited out and the run finished as if nothing had happened, with images never written. The worker now aborts with an error in both cases. Co-Authored-By: Claude Fable 5.1 --- software/control/core/multi_point_worker.py | 15 ++++- .../control/test_MultiPointController.py | 60 +++++++++++++++++++ 2 files changed, 73 insertions(+), 2 deletions(-) diff --git a/software/control/core/multi_point_worker.py b/software/control/core/multi_point_worker.py index b22b729f9..5b85146f3 100644 --- a/software/control/core/multi_point_worker.py +++ b/software/control/core/multi_point_worker.py @@ -708,8 +708,12 @@ def time_left(): # Drain results before shutdown self._summarize_runner_outputs(drain_all=True) + if self._acquisition_error_count > 0 and self._abort_on_failed_job: + self._log.error(f"{self._acquisition_error_count} save job(s) failed; the acquisition ends as an error") + self._abort_due_to_error() - # 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 @@ -725,7 +729,8 @@ def shutdown_runner(job_runner, timeout): self._log.info("Shutting down job runners...") remaining_time = time_left() - aborted = bool(self.abort_requested_fn()) + # 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 getattr(job_runner, "runs_in_process", False): # Its shutdown seals the stores (aborted ones as aborted); do it here so the @@ -897,6 +902,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() diff --git a/software/tests/control/test_MultiPointController.py b/software/tests/control/test_MultiPointController.py index ae7109dae..daf34d94b 100644 --- a/software/tests/control/test_MultiPointController.py +++ b/software/tests/control/test_MultiPointController.py @@ -843,3 +843,63 @@ def test_three_zarr_runs_in_one_session_all_sealed(monkeypatch): 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.""" + pytest.importorskip("tensorstore") + import logging + from control.core.job_processing import SaveZarrJob + + class _FailingFuture: + def result(self): + raise RuntimeError("disk full") + + real_submit = SaveZarrJob.submit + + def failing_submit(self): + future, result = real_submit(self) # open the store as usual, so there is a zarr.json to inspect + if future is not None: + future.result() + return _FailingFuture(), result + + monkeypatch.setattr(SaveZarrJob, "submit", failing_submit) + 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) + 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) From f9960f1e5f3bb9d7bc182d6e3c0c4b3603f37f73 Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 02:13:14 -0700 Subject: [PATCH 09/16] test(acquisition): a write failure that only the final drain sees ends the run as an error Co-Authored-By: Claude Fable 5.1 --- .../control/test_MultiPointController.py | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/software/tests/control/test_MultiPointController.py b/software/tests/control/test_MultiPointController.py index daf34d94b..ecaa96658 100644 --- a/software/tests/control/test_MultiPointController.py +++ b/software/tests/control/test_MultiPointController.py @@ -903,3 +903,51 @@ def test_dead_tiff_subprocess_aborts_acquisition(caplog): 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_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). + """ + pytest.importorskip("tensorstore") + import logging + from control.core.job_processing import SaveZarrJob + from control.core.multi_point_worker import MultiPointWorker, SummarizeResult + + class _FailingFuture: + def result(self): + raise RuntimeError("disk full") + + real_submit = SaveZarrJob.submit + + def failing_submit(self): + future, result = real_submit(self) # open the store as usual, so there is a zarr.json to inspect + if future is not None: + future.result() + return _FailingFuture(), result + + 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) + + monkeypatch.setattr(SaveZarrJob, "submit", failing_submit) + monkeypatch.setattr(MultiPointWorker, "_summarize_runner_outputs", final_drain_only) + 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) + with caplog.at_level(logging.ERROR): + mpc.run_acquisition() + assert tt.finished_event.wait(120) + assert mpc.multiPointWorker._abort_cause == "error" + assert any("save job(s) failed; the acquisition ends as an error" 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) From 8e92128b110b1df7c0ce4ec67d1bb45245b7ec86 Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 02:16:02 -0700 Subject: [PATCH 10/16] docs(zarr): Zarr v3 is written in the acquisition process Co-Authored-By: Claude Fable 5.1 --- software/docs/zarr-v3-format.md | 1 + 1 file changed, 1 insertion(+) 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) From 985cea51414bef2e77a20c2e2bf489b9e89b777e Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 02:31:52 -0700 Subject: [PATCH 11/16] fix(zarr): ZarrWriter needs no event loop, so a persistent camera thread can open stores run after run initialize() runs on the camera frame callback thread. It created and set an asyncio loop for that thread, and finalize()/abort() on the worker thread closed it, leaving the closed loop as the thread's current loop. Real cameras keep one read thread for their lifetime, so every store of the next Zarr acquisition failed with "Event loop is closed". Open the store and wait on pending writes with tensorstore futures' .result() instead; drop the loop bookkeeping. Spec, codecs, metadata and written bytes are unchanged. Co-Authored-By: Claude Fable 5.1 --- software/control/core/zarr_writer.py | 195 +++++++----------- .../core/test_in_process_zarr_runner.py | 40 ++++ .../tests/control/core/test_zarr_writer.py | 60 ++++++ 3 files changed, 173 insertions(+), 122 deletions(-) diff --git a/software/control/core/zarr_writer.py b/software/control/core/zarr_writer.py index 3a871d9e9..468b1f23a 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): @@ -356,118 +355,91 @@ def __init__(self, config: ZarrAcquisitionConfig): 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). @@ -733,13 +705,8 @@ def wait_for_pending(self, timeout_s: Optional[float] = None) -> int: count = len(self._pending_futures) log.debug(f"Waiting for {count} pending writes...") - # 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()) + for future in self._pending_futures: + future.result() self._pending_futures.clear() log.debug(f"Completed {count} pending writes") @@ -778,7 +745,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: @@ -811,23 +777,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/tests/control/core/test_in_process_zarr_runner.py b/software/tests/control/core/test_in_process_zarr_runner.py index 1154b628e..7d4dca589 100644 --- a/software/tests/control/core/test_in_process_zarr_runner.py +++ b/software/tests/control/core/test_in_process_zarr_runner.py @@ -224,3 +224,43 @@ def test_dispatch_after_shutdown_is_refused(tmp_path): 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) diff --git a/software/tests/control/core/test_zarr_writer.py b/software/tests/control/core/test_zarr_writer.py index 030910585..7ce2a8c25 100644 --- a/software/tests/control/core/test_zarr_writer.py +++ b/software/tests/control/core/test_zarr_writer.py @@ -2220,3 +2220,63 @@ def test_submit_in_simulated_io_mode_returns_no_future(self, tmp_path, monkeypat 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 From 6d47bb495490c478a76d084e212dce3ebfa9073b Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 02:32:33 -0700 Subject: [PATCH 12/16] fix(acquisition): close() shuts down an in-process Zarr runner instead of terminating it The abnormal-shutdown loop in MultiPointController.close() called terminate()/join()/kill() on every live runner. InProcessZarrRunner has no terminate()/join(), so close() logged an AttributeError and the open Zarr stores were never sealed. Shut it down with aborted=True instead. Co-Authored-By: Claude Fable 5.1 --- .../control/core/multi_point_controller.py | 7 +++++ .../control/test_MultiPointController.py | 26 +++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/software/control/core/multi_point_controller.py b/software/control/core/multi_point_controller.py index 753413328..178283de9 100644 --- a/software/control/core/multi_point_controller.py +++ b/software/control/core/multi_point_controller.py @@ -1234,6 +1234,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 getattr(job_runner, "runs_in_process", False): + # 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/tests/control/test_MultiPointController.py b/software/tests/control/test_MultiPointController.py index ecaa96658..6d53c60b5 100644 --- a/software/tests/control/test_MultiPointController.py +++ b/software/tests/control/test_MultiPointController.py @@ -905,6 +905,32 @@ def test_dead_tiff_subprocess_aborts_acquisition(caplog): 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.runs_in_process = True + 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. From 7749e2f15987f6554c619a625ca2029542b76ffb Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 02:34:27 -0700 Subject: [PATCH 13/16] fix(acquisition): abandoned save jobs end the run as an error When _finish_jobs timed out with jobs still pending it logged them as abandoned and killed the runner, but the run was still reported as completed, while the in-process Zarr runner sealed those stores as aborted. Missing data is an error: abort the run once after the loop. Co-Authored-By: Claude Fable 5.1 --- software/control/core/multi_point_worker.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/software/control/core/multi_point_worker.py b/software/control/core/multi_point_worker.py index 5b85146f3..b85aa3578 100644 --- a/software/control/core/multi_point_worker.py +++ b/software/control/core/multi_point_worker.py @@ -699,12 +699,17 @@ 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 self._summarize_runner_outputs(drain_all=True) From 20573a2b4a8eae9908d37bffd02a422a4db8918b Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 02:34:47 -0700 Subject: [PATCH 14/16] docs(zarr): correct comments and messages for in-process Zarr saving - in_process_zarr_runner: the process's tensorstore user today is the NDViewer (a recording writer will be). - SaveZarrJob: it writes through its registry, with zarr_writer_info injected by whichever runner dispatches it; the missing-info error names submit()/run(). - multi_point_controller: the pre-warm comments apply to TIFF formats; a Zarr v3 run gets (None, None). - test: open zarr.json with a with block. Co-Authored-By: Claude Fable 5.1 --- software/control/core/in_process_zarr_runner.py | 4 ++-- software/control/core/job_processing.py | 11 ++++++----- software/control/core/multi_point_controller.py | 8 +++++--- software/tests/control/test_MultiPointController.py | 3 ++- 4 files changed, 15 insertions(+), 11 deletions(-) diff --git a/software/control/core/in_process_zarr_runner.py b/software/control/core/in_process_zarr_runner.py index da90d4fd3..771b88de4 100644 --- a/software/control/core/in_process_zarr_runner.py +++ b/software/control/core/in_process_zarr_runner.py @@ -2,8 +2,8 @@ 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 Record + Z-Stack recording -writer does), because tensorstore aborts such a child on its first use. This runner presents +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. diff --git a/software/control/core/job_processing.py b/software/control/core/job_processing.py index 69be9e9a6..b6779cc0c 100644 --- a/software/control/core/job_processing.py +++ b/software/control/core/job_processing.py @@ -543,8 +543,9 @@ def _clear(self) -> None: class SaveZarrJob(Job): """Job for saving images to Zarr v3 format using TensorStore. - 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. + 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") @@ -616,9 +617,9 @@ def submit(self) -> Tuple[Optional[Any], ZarrWriteResult]: """ 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.submit()/run() requires zarr_writer_info but it is None. " + "The runner that dispatches it injects zarr_writer_info. " + "If running directly, set job.zarr_writer_info before calling submit() or run()." ) from control.core.io_simulation import is_simulation_enabled, simulated_zarr_write diff --git a/software/control/core/multi_point_controller.py b/software/control/core/multi_point_controller.py index 178283de9..bd7dee895 100644 --- a/software/control/core/multi_point_controller.py +++ b/software/control/core/multi_point_controller.py @@ -978,7 +978,8 @@ def finish_fn(): 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) + # (TIFF formats: starts a new one warming for next acquisition; a Zarr v3 run gets + # (None, None) and leaves the warm subprocess in place) prewarmed_runner, prewarmed_bp_values = self.get_prewarmed_job_runner() # Worker creation can fail - ensure runner is cleaned up on error @@ -1002,8 +1003,9 @@ def finish_fn(): ) 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. + # Note: for TIFF formats get_prewarmed_job_runner() already started a NEW pre-warmed + # runner, so we're cleaning up the one that was handed off to us. A Zarr v3 run got + # (None, None) and left the warm subprocess in place, so there is nothing to clean up. self._cleanup_prewarmed_runner( prewarmed_runner, context="after worker creation failure", diff --git a/software/tests/control/test_MultiPointController.py b/software/tests/control/test_MultiPointController.py index 6d53c60b5..593913f89 100644 --- a/software/tests/control/test_MultiPointController.py +++ b/software/tests/control/test_MultiPointController.py @@ -732,7 +732,8 @@ def _zarr_jsons_with_squid(mpc): out = [] for p in (Path(mpc.base_path) / mpc.experiment_ID).rglob("zarr.json"): - attrs = json.load(open(p)).get("attributes", {}).get("_squid") + with open(p) as f: + attrs = json.load(f).get("attributes", {}).get("_squid") if attrs is not None: out.append(attrs) return out From 8cc27fa14675817c2cd2ab08c7fb8c6bd5b3fb30 Mon Sep 17 00:00:00 2001 From: You Yan Date: Wed, 7 Oct 2026 02:47:25 -0700 Subject: [PATCH 15/16] test(acquisition): mock job runners in the close() tests declare themselves subprocesses A bare MagicMock answers runs_in_process truthily, so close() took the new in-process branch and never called terminate(). Co-Authored-By: Claude Fable 5.1 --- software/tests/control/core/test_backpressure.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) 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)] From 79f507568c98f2c088a934e1ea9b8e30ca2b5a70 Mon Sep 17 00:00:00 2001 From: You Yan Date: Fri, 9 Oct 2026 06:08:07 -0700 Subject: [PATCH 16/16] refactor(zarr): simplify the in-process Zarr saving branch - InProcessZarrRunner.shutdown() never stamps a store complete over a write it did not see finish: if the drain thread is still blocked after the bounded join, the stores are sealed aborted and shutdown() returns False, which the worker turns into an error end (Copilot review). - One failure policy: _summarize_job_result requests the abort when a save job fails and abort_on_failed_job is set; the per-FOV and end-of-run copies go. - The backpressure counter contract lives in backpressure.note_job_dispatched / note_job_completed, used by JobRunner and the in-process runner alike. - The drain thread holds the job id, not the job (and its frame), while a write is in flight; a failed submit is a concurrent.futures.Future instead of a bespoke class; isinstance(InProcessZarrRunner) replaces the duck-typed runs_in_process flag. - ZarrWriter loses the never-populated pending-futures API (wait_for_pending, pending_write_count) and write_frame is one line over submit_frame. - Tests: one Zarr v3 controller fixture and one failing-submit helper replace six copies of the same setup; runner tests reuse make_test_capture_info and gain the blocked-write shutdown case. Co-Authored-By: Claude Fable 5.1 --- software/control/core/backpressure.py | 24 ++++ .../control/core/in_process_zarr_runner.py | 86 +++++-------- software/control/core/job_processing.py | 45 +++---- .../control/core/multi_point_controller.py | 12 +- software/control/core/multi_point_worker.py | 17 ++- software/control/core/zarr_writer.py | 52 +------- .../core/test_in_process_zarr_runner.py | 70 ++++++---- .../tests/control/core/test_zarr_writer.py | 8 -- .../control/test_MultiPointController.py | 121 ++++++------------ software/tools/benchmark_save_formats.py | 10 +- 10 files changed, 170 insertions(+), 275 deletions(-) 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 index 771b88de4..897e95d2d 100644 --- a/software/control/core/in_process_zarr_runner.py +++ b/software/control/core/in_process_zarr_runner.py @@ -9,20 +9,19 @@ backpressure counters, and puts the JobResult on the output queue the worker polls. """ +import concurrent.futures import queue import threading -from typing import Any, Optional +from typing import Optional import squid.logging -from control.core.backpressure import BackpressureValues +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: - runs_in_process = True - def __init__( self, zarr_writer_info: Optional[ZarrWriterInfo] = None, @@ -32,8 +31,9 @@ def __init__( 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, future, image_bytes, result), in submission order + 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 @@ -41,7 +41,7 @@ def __init__( self._thread = threading.Thread(target=self._drain, name="InProcessZarrRunner", daemon=True) self._started = False self._finished = False - self._killed = False + self._aborted = False # ---- lifecycle, as JobRunner has it @@ -77,15 +77,18 @@ def dispatch(self, job: SaveZarrJob) -> bool: 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 - self._count_dispatched(image_bytes) + 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") - self._submitted.put_nowait((job, _FailedFuture(e), image_bytes, None)) - return True - self._submitted.put_nowait((job, future, image_bytes, result)) + 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": @@ -97,27 +100,30 @@ def has_pending(self) -> bool: def kill(self) -> None: """Mark the run aborted; shutdown() then seals the stores as aborted.""" - self._killed = True + self._aborted = True - def shutdown(self, timeout_s: float = 1.0, aborted: bool = False) -> None: - """Wait for the outstanding writes, then seal every store. + 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) or after kill(). Captured frames are - written either way. + 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 + return True self._finished = True self._submitted.put_nowait(_SENTINEL) self._thread.join(timeout=max(timeout_s, 1.0)) - if self._thread.is_alive(): - self._log.warning("drain thread still waiting on a write after the shutdown timeout; sealing anyway") - if self._killed or aborted: + 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 @@ -126,43 +132,15 @@ def _drain(self) -> None: item = self._submitted.get() if item is _SENTINEL: return - job, future, image_bytes, result = item + job_id, future, image_bytes, result = item try: if future is not None: future.result() - self._output.put_nowait(JobResult(job_id=job.job_id, result=result, exception=None)) + self._output.put_nowait(JobResult(job_id=job_id, result=result, exception=None)) except Exception as e: - self._log.exception(f"Job {job.job_id} failed! Returning exception result.") - self._output.put_nowait(JobResult(job_id=job.job_id, result=None, exception=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: - self._count_completed(image_bytes) - - def _count_dispatched(self, image_bytes: int) -> None: - with self._pending_lock: - self._pending += 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 - - def _count_completed(self, image_bytes: int) -> None: - with self._pending_lock: - self._pending -= 1 - 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) - with self._bp_pending_bytes.get_lock(): - self._bp_pending_bytes.value = max(0, self._bp_pending_bytes.value - image_bytes) - if self._bp_capacity_event is not None: - self._bp_capacity_event.set() - - -class _FailedFuture: - """A write that failed before it was submitted, reported through the same path as a failed write.""" - - def __init__(self, error: Exception): - self._error = error - - def result(self) -> Any: - raise self._error + 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 b6779cc0c..1d04a816e 100644 --- a/software/control/core/job_processing.py +++ b/software/control/core/job_processing.py @@ -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 ( @@ -531,12 +532,9 @@ def abort_all(self) -> None: except Exception as e: self._log.warning(f"Error aborting writer during clear: {e}") finally: - self._clear() - - def _clear(self) -> None: - self.writers.clear() - self.plates_written.clear() - self.wells_written.clear() + self.writers.clear() + self.plates_written.clear() + self.wells_written.clear() @dataclass @@ -617,9 +615,8 @@ def submit(self) -> Tuple[Optional[Any], ZarrWriteResult]: """ if self.zarr_writer_info is None: raise ValueError( - "SaveZarrJob.submit()/run() requires zarr_writer_info but it is None. " - "The runner that dispatches it injects zarr_writer_info. " - "If running directly, set job.zarr_writer_info before calling submit() or 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 @@ -874,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) @@ -1051,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 bd7dee895..8e998b7ed 100644 --- a/software/control/core/multi_point_controller.py +++ b/software/control/core/multi_point_controller.py @@ -23,6 +23,7 @@ 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 @@ -977,9 +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 - # (TIFF formats: starts a new one warming for next acquisition; a Zarr v3 run gets - # (None, None) and leaves the warm subprocess in place) + # 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 @@ -1002,10 +1001,7 @@ def finish_fn(): run_state_writer=self._run_state_writer, ) except Exception: - # Clean up pre-warmed runner if worker creation failed. - # Note: for TIFF formats get_prewarmed_job_runner() already started a NEW pre-warmed - # runner, so we're cleaning up the one that was handed off to us. A Zarr v3 run got - # (None, None) and left the warm subprocess in place, so there is nothing to clean up. + # 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", @@ -1236,7 +1232,7 @@ 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 getattr(job_runner, "runs_in_process", False): + 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)" diff --git a/software/control/core/multi_point_worker.py b/software/control/core/multi_point_worker.py index b85aa3578..8d7a19e08 100644 --- a/software/control/core/multi_point_worker.py +++ b/software/control/core/multi_point_worker.py @@ -711,11 +711,8 @@ def time_left(): # 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) - if self._acquisition_error_count > 0 and self._abort_on_failed_job: - self._log.error(f"{self._acquisition_error_count} save job(s) failed; the acquisition ends as an error") - self._abort_due_to_error() # 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. @@ -737,11 +734,12 @@ def shutdown_runner(job_runner, timeout): # 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 getattr(job_runner, "runs_in_process", False): + 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: - job_runner.shutdown(remaining_time, aborted=aborted) + 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 @@ -935,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: @@ -1120,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) diff --git a/software/control/core/zarr_writer.py b/software/control/core/zarr_writer.py index 468b1f23a..469ac415b 100644 --- a/software/control/core/zarr_writer.py +++ b/software/control/core/zarr_writer.py @@ -352,7 +352,6 @@ def __init__(self, config: ZarrAcquisitionConfig): """ self._config = config self._dataset = None - self._pending_futures: List[Any] = [] self._initialized = False self._finalized = False @@ -670,52 +669,8 @@ def submit_frame(self, image: np.ndarray, t: int, c: int, z: int, fov: Optional[ return self._dataset[fov, t, c, z, :, :].write(image) 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. - - This submits the write via :meth:`submit_frame` and waits on the resulting future before - returning, so the data is visible to other processes reading the same zarr store once this - returns. Callers that write many frames back-to-back should prefer :meth:`submit_frame` - with a bounded number of futures in flight. - - Args: - image: 2D image array (Y, X) - t: Time point index - c: Channel index - z: Z-slice index - fov: FOV index (required for 6D datasets, ignored for 5D) - """ + """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() - if self._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...") - - for future in self._pending_futures: - future.result() - - 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 finalize(self) -> None: """Finalize the dataset (blocking).""" @@ -725,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: @@ -756,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() diff --git a/software/tests/control/core/test_in_process_zarr_runner.py b/software/tests/control/core/test_in_process_zarr_runner.py index 7d4dca589..23f2f89c4 100644 --- a/software/tests/control/core/test_in_process_zarr_runner.py +++ b/software/tests/control/core/test_in_process_zarr_runner.py @@ -1,17 +1,18 @@ """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 -import squid.abc from control.core.backpressure import create_backpressure_values -from control.core.job_processing import CaptureInfo, JobImage, JobResult, SaveZarrJob, ZarrWriteResult, ZarrWriterInfo -from control.models import AcquisitionChannel, CameraSettings, IlluminationSettings +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") @@ -22,26 +23,17 @@ def _info(tmp_path, t_size=1, z_size=2): def _job(z, value=1): return SaveZarrJob( - capture_info=CaptureInfo( - position=squid.abc.Pos(x_mm=0.0, y_mm=0.0, z_mm=0.0, theta_rad=None), - z_index=z, - capture_time=time.time(), - configuration=AcquisitionChannel( - name="BF", - illumination_settings=IlluminationSettings(illumination_channel="LED", intensity=50.0), - camera_settings=CameraSettings(exposure_time_ms=10.0, gain_mode=1.0), - ), - save_directory="/tmp/unused", - file_id=f"f_{z}", - region_id="A1", - fov=0, - configuration_idx=0, - time_point=0, - ), + 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: @@ -130,11 +122,11 @@ def test_dispatch_injects_zarr_writer_info_and_registry(tmp_path): def test_failed_write_reaches_output_queue_as_exception(tmp_path, monkeypatch): from control.core.in_process_zarr_runner import InProcessZarrRunner - class _FailingFuture: - def result(self): - raise RuntimeError("disk full") - - monkeypatch.setattr(SaveZarrJob, "submit", lambda self: (_FailingFuture(), ZarrWriteResult(0, 0, 0, "BF", 0))) + 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() @@ -264,3 +256,33 @@ def _camera_thread(): 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 7ce2a8c25..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 diff --git a/software/tests/control/test_MultiPointController.py b/software/tests/control/test_MultiPointController.py index 593913f89..e926cdc7c 100644 --- a/software/tests/control/test_MultiPointController.py +++ b/software/tests/control/test_MultiPointController.py @@ -739,11 +739,9 @@ def _zarr_jsons_with_squid(mpc): return out -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.""" +def _zarr_v3_controller(monkeypatch): + """A simulated microscope and controller set up for a Zarr v3 multipoint run, with its tracker.""" pytest.importorskip("tensorstore") - from control.core.in_process_zarr_runner import InProcessZarrRunner - 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) @@ -751,6 +749,32 @@ def test_zarr_v3_multipoint_saves_in_process_and_seals(monkeypatch): 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) @@ -760,16 +784,8 @@ def test_zarr_v3_multipoint_saves_in_process_and_seals(monkeypatch): def test_zarr_is_sealed_with_multiprocessing_off(monkeypatch): """The in-process fallback used to leave the store unsealed; it must not any more.""" - pytest.importorskip("tensorstore") - - control._def.MERGE_CHANNELS = False - monkeypatch.setattr(control._def, "FILE_SAVING_OPTION", control._def.FileSavingOption.ZARR_V3) monkeypatch.setattr(control._def.Acquisition, "USE_MULTIPROCESSING", False) - 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) + mpc, tt = _zarr_v3_controller(monkeypatch) mpc.run_acquisition() assert tt.finished_event.wait(120) attrs = _zarr_jsons_with_squid(mpc) @@ -778,17 +794,10 @@ def test_zarr_is_sealed_with_multiprocessing_off(monkeypatch): 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.""" - pytest.importorskip("tensorstore") import time - 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()) + mpc, tt = _zarr_v3_controller(monkeypatch) mpc.set_NZ(5) - add_some_coordinates(mpc) - select_some_configs(mpc, scope.objective_store.current_objective) mpc.run_acquisition() assert tt.started_event.wait(30) deadline = time.monotonic() + 30 @@ -805,20 +814,13 @@ def test_abort_mid_zarr_run_seals_store_aborted(monkeypatch): 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.""" - pytest.importorskip("tensorstore") from control.core import job_processing - 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()) + 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)) - add_some_coordinates(mpc) - select_some_configs(mpc, scope.objective_store.current_objective) mpc.run_acquisition() assert tt.finished_event.wait(120) assert started == [], "a Zarr v3 run must not start a save subprocess" @@ -827,15 +829,7 @@ def test_zarr_v3_run_leaves_the_warm_subprocess_for_tiff(monkeypatch): 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.""" - 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) + mpc, tt = _zarr_v3_controller(monkeypatch) for i in range(3): tt.started_event.clear() tt.finished_event.clear() @@ -848,30 +842,10 @@ def test_three_zarr_runs_in_one_session_all_sealed(monkeypatch): 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.""" - pytest.importorskip("tensorstore") import logging - from control.core.job_processing import SaveZarrJob - - class _FailingFuture: - def result(self): - raise RuntimeError("disk full") - - real_submit = SaveZarrJob.submit - def failing_submit(self): - future, result = real_submit(self) # open the store as usual, so there is a zarr.json to inspect - if future is not None: - future.result() - return _FailingFuture(), result - - monkeypatch.setattr(SaveZarrJob, "submit", failing_submit) - 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) + 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) @@ -921,7 +895,6 @@ def test_close_shuts_down_in_process_runner_instead_of_terminating_it(): mpc._memory_monitor = None mpc.thread = None runner = MagicMock(spec=InProcessZarrRunner) - runner.runs_in_process = True runner.is_alive.return_value = True runner.terminate = MagicMock() # not part of InProcessZarrRunner; present only to prove it is not called worker = MagicMock() @@ -938,23 +911,9 @@ def test_zarr_write_error_surfacing_at_the_end_ends_the_run_as_error(monkeypatch 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). """ - pytest.importorskip("tensorstore") import logging - from control.core.job_processing import SaveZarrJob from control.core.multi_point_worker import MultiPointWorker, SummarizeResult - class _FailingFuture: - def result(self): - raise RuntimeError("disk full") - - real_submit = SaveZarrJob.submit - - def failing_submit(self): - future, result = real_submit(self) # open the store as usual, so there is a zarr.json to inspect - if future is not None: - future.result() - return _FailingFuture(), result - real_summarize = MultiPointWorker._summarize_runner_outputs def final_drain_only(self, drain_all=False): @@ -962,19 +921,13 @@ def final_drain_only(self, drain_all=False): return SummarizeResult(none_failed=True, had_results=False) return real_summarize(self, drain_all=drain_all) - monkeypatch.setattr(SaveZarrJob, "submit", failing_submit) + mpc, tt = _zarr_v3_controller(monkeypatch) + _make_zarr_writes_fail(monkeypatch) monkeypatch.setattr(MultiPointWorker, "_summarize_runner_outputs", final_drain_only) - 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) with caplog.at_level(logging.ERROR): mpc.run_acquisition() assert tt.finished_event.wait(120) assert mpc.multiPointWorker._abort_cause == "error" - assert any("save job(s) failed; the acquisition ends as an error" in r.getMessage() for r in caplog.records) + 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