From dcdcbda376231880ea02886976294fbac8b149f8 Mon Sep 17 00:00:00 2001 From: ramyaguru Date: Tue, 30 Jun 2026 14:20:19 -0400 Subject: [PATCH 01/33] PR0: Unwire (do not delete) PublishToCloudOp PublishToCloudOp consolidated per-batch temp HDF files into a final file, but SinkAndPublishOp.write_scan_file now does the single end-of-scan write directly to the publish folder (commit 2d089e3). The temp files it consumed are no longer produced, so it no-ops on every trigger. Rather than delete it (it may be useful for a future external/temp-file workflow), this just removes it from the live pipeline. That still simplifies the control path for the upcoming flush/header work: ControlOp no longer needs an output port or a processing_end branch. - Keep PublishToCloudOp class in publish.py, documented as not-currently-wired - pipeline.py: drop its construction, import, and both control flows - Stop SinkAndPublishOp emitting processing_end (+ drop its output port) - ControlOp: drop output port and processing_end forwarding Co-Authored-By: Claude Opus 4.8 --- pipeline/control.py | 15 +++++---------- pipeline/pipeline.py | 15 ++------------- pipeline/publish.py | 20 ++++++++++---------- 3 files changed, 17 insertions(+), 33 deletions(-) diff --git a/pipeline/control.py b/pipeline/control.py index 36aa720..976ded9 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -12,8 +12,8 @@ class ControlOp(Operator): """ Control operator for managing pipeline flow. - Handles control messages like flush and processing_end, - coordinating state across multiple operators. + Handles the flush control message, coordinating flush state across + multiple operators. """ def __init__(self, fragment, *args, @@ -35,17 +35,16 @@ def __init__(self, fragment, *args, def setup(self, spec: OperatorSpec): spec.input("input").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128) - spec.output("output") def compute(self, op_input, op_output, context): """Handle control messages.""" msg = op_input.receive("input") - + if msg == "flush": # Flush all flushable operators for op in self.flushable_ops: op.flush() - + # Publish flush message through the backend if available if self.publish_backend is not None: import numpy as np @@ -53,11 +52,7 @@ def compute(self, op_input, op_output, context): # Also signal ptycho consumers; harmless when ptycho is disabled # (no subscriber listens on this subject). self.publish_backend.publish("ptycho_flush", np.array([1])) - - elif msg == "processing_end": - # Forward processing_end signal - op_output.emit("processing_end", "output") - + else: self.logger.info(f"Received unknown message: {msg}") diff --git a/pipeline/pipeline.py b/pipeline/pipeline.py index 7a2316a..b8568eb 100644 --- a/pipeline/pipeline.py +++ b/pipeline/pipeline.py @@ -28,7 +28,7 @@ GatherOp ) from processing import MaskingOp -from publish import SinkAndPublishOp, PublishToCloudOp +from publish import SinkAndPublishOp from control import ControlOp # Try to import NATS (optional for testing) @@ -133,15 +133,6 @@ def compose(self): **self.kwargs('sink_and_publish_op'), name="sink_and_publish_op") - # ===== Cloud Publishing Operator ===== - publish_folder = self.kwargs('sink_and_publish_op')['publish_folder'] - temp_folder = self.kwargs('sink_and_publish_op')['temp_folder'] - - publish_to_cloud_op = PublishToCloudOp(self, - publish_folder=publish_folder, - temp_folder=temp_folder, - name="publish_to_cloud_op") - # ===== Control Operator ===== flushable_ops = [gather_op, position_src, sink_and_publish_op] @@ -205,10 +196,8 @@ def compose(self): self.add_flow(gather_op, ptycho_accum, {("output", "input")}) self.add_flow(ptycho_recon, ptycho_publish, {("output", "input")}) - # Control path: flush and completion signals + # Control path: flush signal (start message → unconditional idempotent flush) self.add_flow(img_src, control_op, {("flush", "input")}) - self.add_flow(sink_and_publish_op, control_op, {("processing_end", "input")}) - self.add_flow(control_op, publish_to_cloud_op, {("output", "trigger")}) def main(): diff --git a/pipeline/publish.py b/pipeline/publish.py index ca64e09..a09df58 100644 --- a/pipeline/publish.py +++ b/pipeline/publish.py @@ -116,7 +116,6 @@ def __init__(self, fragment, *args, def setup(self, spec: OperatorSpec): spec.input("input").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128).condition(ConditionType.NONE) - spec.output("processing_end").condition(ConditionType.NONE) def write_scan_file(self, series_id): """Write the buffered scan to a single HDF5 file. @@ -210,8 +209,6 @@ def compute(self, op_input, op_output, context): if self.publish_folder is not None and series_id is not None: self.write_scan_file(series_id) - op_output.emit("processing_end", "processing_end") - _n = self.processed_frame_count _b = self.processed_batch_count _elapsed = time.time() - series_start_time if series_start_time > 0 else 0 @@ -228,17 +225,20 @@ class PublishToCloudOp(Operator): NOTE: STXM saving now happens in SinkAndPublishOp.write_scan_file (a single end-of-scan write), so the per-batch temp files this op consolidated are no - longer produced. It is retained for any external/temp-file workflow and - no-ops gracefully when no temp file is present. + longer produced. It is retained for a possible future external/temp-file + workflow and no-ops gracefully when no temp file is present. + + NOT CURRENTLY WIRED into the pipeline (see pipeline.py). To re-enable, feed a + completion trigger into its "trigger" input (e.g. from SinkAndPublishOp). """ - + def __init__(self, fragment, publish_folder: str = None, temp_folder: str = None, *args, **kwargs): """ Initialize cloud publishing operator. - + Args: fragment: Holoscan fragment publish_folder: Final destination folder @@ -256,14 +256,14 @@ def compute(self, op_input, op_output, context): """Consolidate and publish dataset on trigger using metadata.""" # Receive trigger - metadata is automatically merged trigger = op_input.receive("trigger") - + if trigger == "processing_end": if self.publish_folder is None or self.temp_folder is None: return # Get the series ID from metadata that flowed from upstream series_id = self.metadata.get("series_id") - + if series_id is None: self.logger.warning("No series_id found in metadata, cannot publish") return @@ -299,7 +299,7 @@ def compute(self, op_input, op_output, context): f.attrs[key] = value self.logger.info(f"Published concatenated data to {publish_file}") - + # Remove temp file os.remove(temp_file) From 87515de81dfaeb4bd68b12d878c20b16375a9f76 Mon Sep 17 00:00:00 2001 From: ramyaguru Date: Tue, 30 Jun 2026 14:31:33 -0400 Subject: [PATCH 02/33] =?UTF-8?q?PR1:=20Flush=20plumbing=20=E2=80=94=20com?= =?UTF-8?q?pletion=20signal=20+=20lock-safe,=20idempotent=20flush?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Foundational plumbing for the header/tomography work. No behavior change for existing single-projection scans (flush still happens at the next start). - PtychoReconstructionOp: emit a one-shot "recon_complete" on a new "complete" output port when the final iteration is reached (reuses the existing is_last predicate, S9); guard + reset via _completed in flush(). - PtychoAccumulatorOp.flush: zero the GPU buffers INSIDE ptycho_state["lock"] so recon can't read half-zeroed buffers (R-3); no-op when nothing has been accumulated since the last flush (R-2 _dirty guard) so the unconditional start-flush is free when clean. - ControlOp: log "recon_complete" only — NO flush here, so a completed single-projection scan keeps its result until the next start/header. The flush-on-completion / per-projection semantics arrive with PR3. - pipeline.py: wire ptycho_recon "complete" -> control_op. Co-Authored-By: Claude Opus 4.8 --- pipeline/control.py | 8 ++++++ pipeline/pipeline.py | 2 ++ pipeline/ptychography_ops.py | 47 ++++++++++++++++++++++++++++++------ 3 files changed, 49 insertions(+), 8 deletions(-) diff --git a/pipeline/control.py b/pipeline/control.py index 976ded9..7e280a4 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -53,6 +53,14 @@ def compute(self, op_input, op_output, context): # (no subscriber listens on this subject). self.publish_backend.publish("ptycho_flush", np.array([1])) + elif msg == "recon_complete": + # Plumbing for PR2 (header preempt) / PR3 (tomography boundary). + # No flush here on purpose: a completed single-projection scan must + # keep its result until the next start/header. The flush-on- + # completion / per-projection semantics are added with the + # tomography work. + self.logger.info("Reconstruction complete signal received") + else: self.logger.info(f"Received unknown message: {msg}") diff --git a/pipeline/pipeline.py b/pipeline/pipeline.py index b8568eb..acd5cb3 100644 --- a/pipeline/pipeline.py +++ b/pipeline/pipeline.py @@ -195,6 +195,8 @@ def compose(self): if self.ptychography_enabled: self.add_flow(gather_op, ptycho_accum, {("output", "input")}) self.add_flow(ptycho_recon, ptycho_publish, {("output", "input")}) + # Completion signal → control (logged in PR1; drives flush in PR3) + self.add_flow(ptycho_recon, control_op, {("complete", "input")}) # Control path: flush signal (start message → unconditional idempotent flush) self.add_flow(img_src, control_op, {("flush", "input")}) diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 12f40ba..6751ad3 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -51,6 +51,9 @@ class PtychoAccumulatorOp(Operator): def __init__(self, fragment, *args, ptycho_state, **kwargs): self.ptycho_state = ptycho_state self.lock = ptycho_state["lock"] + # Tracks whether any frames have been accumulated since the last flush, + # so a redundant flush (e.g. the unconditional start-flush, R-2) is free. + self._dirty = False self.logger = logging.getLogger(kwargs.get("name", "PtychoAccumulatorOp")) super().__init__(fragment, *args, **kwargs) @@ -60,16 +63,25 @@ def setup(self, spec: OperatorSpec): ).condition(ConditionType.NONE) def flush(self): - """Reset fill level and zero out GPU buffers for a new series.""" + """Reset fill level and zero out GPU buffers for a new series. + + No-ops when nothing has been accumulated since the last flush (R-2), so + the unconditional start-flush is free when buffers are already clean. + Buffer zeroing happens under the lock so the reconstruction op can never + observe half-zeroed buffers (R-3). + """ with self.lock: + if not self._dirty: + return self.ptycho_state["filled_until"] = 0 - self.ptycho_state["raw_gpu"][:] = 0 - self.ptycho_state["positions_full"][:] = 0 - self.ptycho_state["tilts_full"][:] = 0 - # Clear auto-centre so the new scan re-derives its own scan centre - # from the first batch rather than reusing the previous scan's. - self.ptycho_state["scan_center_py"] = None - self.ptycho_state["scan_center_px"] = None + self.ptycho_state["raw_gpu"][:] = 0 + self.ptycho_state["positions_full"][:] = 0 + self.ptycho_state["tilts_full"][:] = 0 + # Clear auto-centre so the new scan re-derives its own scan centre + # from the first batch rather than reusing the previous scan's. + self.ptycho_state["scan_center_py"] = None + self.ptycho_state["scan_center_px"] = None + self._dirty = False self.logger.info("Flushed ptychography accumulator buffers") def compute(self, op_input, op_output, context): @@ -155,6 +167,7 @@ def compute(self, op_input, op_output, context): # Atomically update fill counter with self.lock: self.ptycho_state["filled_until"] = new_end + self._dirty = True # Summary when buffer is full if new_end >= self.ptycho_state["no_frames"]: @@ -292,6 +305,9 @@ def __init__( self.all_data_arrived = False self.post_stream_count = 0 self.initialized_gpu = False + # Emitted exactly once per scan when the final iteration is reached; + # reset on flush so the next scan/projection can signal again. + self._completed = False # Pristine reconstruction state, snapshotted on first GPU init and # used to reset the object (and optionally the probe) on flush. self._obj_initial = None @@ -304,6 +320,8 @@ def __init__( def setup(self, spec: OperatorSpec): spec.output("output").condition(ConditionType.NONE) + # Completion signal to ControlOp (plumbing for PR2/PR3). + spec.output("complete").condition(ConditionType.NONE) def flush(self): """Reset reconstruction state for a new scan. @@ -317,6 +335,7 @@ def flush(self): self.current_iteration = 0 self.all_data_arrived = False self.post_stream_count = 0 + self._completed = False if self.initialized_gpu: pty_model = self.ptycho_state["pty_model"] pty_model.obj.array_global[:] = self._obj_initial @@ -509,6 +528,18 @@ def compute(self, op_input, op_output, context): } op_output.emit(out, "output") + # Completion signal — emitted once when the scan's final iteration is + # reached (reuses the existing `is_last` predicate, S9). Plumbing for + # PR2 (header preempt) / PR3 (tomography boundary); ControlOp currently + # only logs it, so a completed single-projection scan keeps its result + # until the next start/header. + if is_last and not self._completed: + self._completed = True + op_output.emit("recon_complete", "complete") + self.logger.info( + "Reconstruction complete at iteration %d", self.current_iteration + ) + t_end = time.perf_counter() self.logger.info( "ITER_TIMING iter=%d n_filled=%d valid=%d total_ms=%.1f " From 6a4f9276990a4c2dc8d421f77fa542e2d6a5152d Mon Sep 17 00:00:00 2001 From: ramyaguru Date: Tue, 30 Jun 2026 23:14:26 -0400 Subject: [PATCH 03/33] =?UTF-8?q?PR1:=20Flush=20hardening=20=E2=80=94=20ra?= =?UTF-8?q?ce-safe=20deferred=20flush,=20flush-on-completion,=20STXM=20saf?= =?UTF-8?q?ety?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Builds on the flush plumbing so completion drives the flush (Task 3) and every flush is race-safe — a prerequisite for PR2's mid-stream header preemption. - ControlOp: recon_complete now flushes (result is already saved+published before it fires) and marks _flushed; the scan-start flush becomes a safety net that skips when already flushed (no double-flush) and only runs if the previous scan didn't complete. Factored into _do_flush(). - Deferred flush for GatherOp, PtychoAccumulatorOp, PtychoReconstructionOp: flush() sets a flag; the reset happens at the top of the next compute(), so it never mutates caches/GPU buffers/the PIE object concurrently with compute (fixes the GatherOp boolean-index race; safe for mid-stream preemption). recon flush also zeros the shared filled_until so it can't re-process the finished scan before the accumulator's own deferred flush runs. - recon_complete is gated on all_data_arrived, so it only fires when the scan is genuinely complete (not when a low total_iterations exhausts mid-stream). - SinkAndPublishOp: write_scan_file now saves-and-clears (self-contained), and flush() defensively writes any unwritten buffer before clearing so a scan's STXM file is never lost. Completion check uses >= (tolerates overshoot). Co-Authored-By: Claude Opus 4.8 --- pipeline/control.py | 52 +++++++++------- pipeline/data_io.py | 30 +++++++--- pipeline/ptychography_ops.py | 112 ++++++++++++++++++++++------------- pipeline/publish.py | 27 ++++++++- 4 files changed, 149 insertions(+), 72 deletions(-) diff --git a/pipeline/control.py b/pipeline/control.py index 7e280a4..64dcb65 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -32,34 +32,46 @@ def __init__(self, fragment, *args, self.logger = logging.getLogger(kwargs.get("name", "ControlOp")) self.flushable_ops = flushable_ops self.publish_backend = publish_backend - + # True once a completion (recon_complete) flush has run and no new scan + # has started since. Lets the scan-start flush skip when the buffers are + # already clean, so we don't double-flush (Task 3 flush-check-at-start). + self._flushed = False + def setup(self, spec: OperatorSpec): spec.input("input").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128) + def _do_flush(self): + """Flush all flushable operators and broadcast the flush signals.""" + for op in self.flushable_ops: + op.flush() + if self.publish_backend is not None: + import numpy as np + self.publish_backend.publish("stxm_flush", np.array([1])) # Simple signal + # Also signal ptycho consumers; harmless when ptycho is disabled + # (no subscriber listens on this subject). + self.publish_backend.publish("ptycho_flush", np.array([1])) + def compute(self, op_input, op_output, context): """Handle control messages.""" msg = op_input.receive("input") - if msg == "flush": - # Flush all flushable operators - for op in self.flushable_ops: - op.flush() - - # Publish flush message through the backend if available - if self.publish_backend is not None: - import numpy as np - self.publish_backend.publish("stxm_flush", np.array([1])) # Simple signal - # Also signal ptycho consumers; harmless when ptycho is disabled - # (no subscriber listens on this subject). - self.publish_backend.publish("ptycho_flush", np.array([1])) + if msg == "recon_complete": + # The recon finished its final iteration and has ALREADY saved + # (after_iteration -> pty_out) and published the result before emitting + # this, so flushing now is safe (Task 3: flush after the last iteration). + self.logger.info("Reconstruction complete — flushing for next scan") + self._do_flush() + self._flushed = True - elif msg == "recon_complete": - # Plumbing for PR2 (header preempt) / PR3 (tomography boundary). - # No flush here on purpose: a completed single-projection scan must - # keep its result until the next start/header. The flush-on- - # completion / per-projection semantics are added with the - # tomography work. - self.logger.info("Reconstruction complete signal received") + elif msg == "flush": + # Scan-start safety flush: only flush if the buffers aren't already + # clean from a completion flush. If the previous scan completed, this + # no-ops (no double flush); if it was interrupted, this cleans up. + if self._flushed: + self.logger.info("Start-flush skipped — already flushed on completion") + self._flushed = False + else: + self._do_flush() else: self.logger.info(f"Received unknown message: {msg}") diff --git a/pipeline/data_io.py b/pipeline/data_io.py index 73052d9..c0c0233 100644 --- a/pipeline/data_io.py +++ b/pipeline/data_io.py @@ -203,7 +203,7 @@ def compute(self, op_input, op_output, context): # Handle data message datasets = msg["datasets"] #print(datasets) - + # Extract position data x = np.array(datasets["/pi_x"]["data"]) #FMC_IN.VAL1.Mean y = np.array(datasets["/FMC_IN.VAL2.Mean"]["data"]) @@ -472,6 +472,12 @@ def __init__(self, fragment, *args, batch_size: int = 1, **kwargs): self.position_ids = np.zeros((0,), dtype=int) self.count = 0 self.batch_size = int(batch_size) + # Deferred flush: flush() sets this flag and the actual cache clear + # happens at the top of the next compute(), so it never mutates the + # caches while compute() is mid-synchronise. This avoids the boolean- + # index race (data_io.py "size of axis is 0 but ... 64") when a flush + # arrives mid-stream — e.g. PR2's header preemption. + self._flush_requested = False self.logger = logging.getLogger(kwargs.get("name", "GatherOp")) super().__init__(fragment, *args, **kwargs) @@ -482,23 +488,29 @@ def setup(self, spec: OperatorSpec): spec.output("output") def flush(self): - """Reset all cached data on flush.""" + """Request a cache reset. Deferred to the top of the next compute() so it + never clears the caches while compute() is mid-synchronise (thread-safe).""" + self._flush_requested = True + + def _perform_flush(self): + """Actually clear the caches — only ever called from compute().""" self.images = None self.image_ids = np.zeros((0,), dtype=int) self.positions = np.zeros((0, 4)) self.position_ids = np.zeros((0,), dtype=int) self.count = 0 self.logger.info( - f"[FLUSH VERIFY] GatherOp cleared: " - f"images={None if self.images is None else self.images.shape}, " - f"image_ids={self.image_ids.size}, " - f"positions={self.positions.shape}, " - f"position_ids={self.position_ids.size}, " - f"count={self.count}" + "[FLUSH VERIFY] GatherOp cleared: images=None, image_ids=0, " + "positions=(0, 4), position_ids=0, count=0" ) - def compute(self, op_input, op_output, context): + def compute(self, op_input, op_output, context): """Gather and synchronize image and position data.""" + # Perform any requested flush here — single-threaded w.r.t. the caches. + if self._flush_requested: + self._flush_requested = False + self._perform_flush() + # Receive image data images_dict = op_input.receive("images") if images_dict is not None: diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 6751ad3..9d88e6b 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -54,6 +54,11 @@ def __init__(self, fragment, *args, ptycho_state, **kwargs): # Tracks whether any frames have been accumulated since the last flush, # so a redundant flush (e.g. the unconditional start-flush, R-2) is free. self._dirty = False + # Deferred flush (see GatherOp): flush() sets this flag; the actual reset + # happens at the top of the next compute(), so it never zeros the GPU + # buffers while compute() is mid-write — race-safe even for a mid-stream + # flush (PR2 header preemption). + self._flush_requested = False self.logger = logging.getLogger(kwargs.get("name", "PtychoAccumulatorOp")) super().__init__(fragment, *args, **kwargs) @@ -63,28 +68,33 @@ def setup(self, spec: OperatorSpec): ).condition(ConditionType.NONE) def flush(self): - """Reset fill level and zero out GPU buffers for a new series. - - No-ops when nothing has been accumulated since the last flush (R-2), so - the unconditional start-flush is free when buffers are already clean. - Buffer zeroing happens under the lock so the reconstruction op can never - observe half-zeroed buffers (R-3). - """ + """Request a reset; performed at the top of the next compute() (deferred, + so it never races the buffer writes in compute).""" + self._flush_requested = True + + def _perform_flush(self): + """Reset fill level and zero the GPU buffers. Only ever called from + compute(), so it is single-threaded w.r.t. the buffer writes. No-ops when + nothing has been accumulated since the last flush (free redundant flush).""" + self._flush_requested = False + if not self._dirty: + return with self.lock: - if not self._dirty: - return self.ptycho_state["filled_until"] = 0 - self.ptycho_state["raw_gpu"][:] = 0 - self.ptycho_state["positions_full"][:] = 0 - self.ptycho_state["tilts_full"][:] = 0 - # Clear auto-centre so the new scan re-derives its own scan centre - # from the first batch rather than reusing the previous scan's. - self.ptycho_state["scan_center_py"] = None - self.ptycho_state["scan_center_px"] = None - self._dirty = False + self.ptycho_state["raw_gpu"][:] = 0 + self.ptycho_state["positions_full"][:] = 0 + self.ptycho_state["tilts_full"][:] = 0 + # Clear auto-centre so the new scan re-derives its own scan centre + # from the first batch rather than reusing the previous scan's. + self.ptycho_state["scan_center_py"] = None + self.ptycho_state["scan_center_px"] = None + self._dirty = False self.logger.info("Flushed ptychography accumulator buffers") def compute(self, op_input, op_output, context): + # Perform any requested flush here — single-threaded w.r.t. the buffers. + if self._flush_requested: + self._perform_flush() data = op_input.receive("input") if data is None: return @@ -308,6 +318,10 @@ def __init__( # Emitted exactly once per scan when the final iteration is reached; # reset on flush so the next scan/projection can signal again. self._completed = False + # Deferred flush (see GatherOp/accumulator): flush() sets this flag; the + # object/counter reset happens at the top of the next compute(), so it + # never races the PIE object update — race-safe for a mid-stream flush. + self._flush_requested = False # Pristine reconstruction state, snapshotted on first GPU init and # used to reset the object (and optionally the probe) on flush. self._obj_initial = None @@ -324,30 +338,42 @@ def setup(self, spec: OperatorSpec): spec.output("complete").condition(ConditionType.NONE) def flush(self): - """Reset reconstruction state for a new scan. + """Request a reset; performed at the top of the next compute() (deferred, + so the object/counter reset never races the PIE update in compute).""" + self._flush_requested = True + + def _perform_flush(self): + """Reset reconstruction state for a new scan. Only ever called from + compute(), so the object reset is single-threaded w.r.t. the PIE update. Resets iteration counters and the object to its initial guess. By default the probe (and its flux) are CARRIED OVER from the previous scan as a warm start, since consecutive scans usually share illumination. Set ``reset_probe=True`` to fully reset the probe too. """ + self._flush_requested = False + self.current_iteration = 0 + self.all_data_arrived = False + self.post_stream_count = 0 + self._completed = False + # Clear the shared fill counter too, so this op immediately sees "no data" + # and won't re-process the just-finished scan before the accumulator's own + # (deferred) flush zeros the buffers. Both ops resetting it to 0 is + # consistent; the accumulator's flush runs before it writes new frames. with self.lock: - self.current_iteration = 0 - self.all_data_arrived = False - self.post_stream_count = 0 - self._completed = False - if self.initialized_gpu: - pty_model = self.ptycho_state["pty_model"] - pty_model.obj.array_global[:] = self._obj_initial - pty_model.obj.array_global_old[:] = self._obj_initial - if self.reset_probe: - # Full reset: restore initial probe + flux so iteration 0 - # recomputes flux and re-normalises the probe. - pty_model.probe.array_states[:] = self._probe_initial - pty_model.source.flux = self._flux_initial - # else: leave the previous scan's probe and flux untouched. - # flux stays >= 0, so the iter-0 re-normalisation branch in - # compute() does not fire and the carried probe is preserved. + self.ptycho_state["filled_until"] = 0 + if self.initialized_gpu: + pty_model = self.ptycho_state["pty_model"] + pty_model.obj.array_global[:] = self._obj_initial + pty_model.obj.array_global_old[:] = self._obj_initial + if self.reset_probe: + # Full reset: restore initial probe + flux so iteration 0 + # recomputes flux and re-normalises the probe. + pty_model.probe.array_states[:] = self._probe_initial + pty_model.source.flux = self._flux_initial + # else: leave the previous scan's probe and flux untouched. + # flux stays >= 0, so the iter-0 re-normalisation branch in + # compute() does not fire and the carried probe is preserved. if self.reset_probe: self.logger.info( @@ -362,6 +388,10 @@ def flush(self): ) def compute(self, op_input, op_output, context): + # Perform any requested flush here — single-threaded w.r.t. the PIE update. + if self._flush_requested: + self._perform_flush() + # Snapshot fill level with self.lock: n_filled = self.ptycho_state["filled_until"] @@ -528,12 +558,14 @@ def compute(self, op_input, op_output, context): } op_output.emit(out, "output") - # Completion signal — emitted once when the scan's final iteration is - # reached (reuses the existing `is_last` predicate, S9). Plumbing for - # PR2 (header preempt) / PR3 (tomography boundary); ControlOp currently - # only logs it, so a completed single-projection scan keeps its result - # until the next start/header. - if is_last and not self._completed: + # Completion signal — emitted once, only when the scan is GENUINELY + # complete: all frames have arrived AND the final (post-stream) iteration + # is done. Gating on all_data_arrived is essential — without it, a low + # total_iterations exhausts mid-stream and fires "complete" while frames + # are still arriving, which would flush GatherOp mid-compute (race) and + # reset the object before the full scan is reconstructed. ControlOp + # flushes on this signal (Task 3: flush after the last iteration). + if is_last and self.all_data_arrived and not self._completed: self._completed = True op_output.emit("recon_complete", "complete") self.logger.info( diff --git a/pipeline/publish.py b/pipeline/publish.py index a09df58..44402ac 100644 --- a/pipeline/publish.py +++ b/pipeline/publish.py @@ -104,6 +104,9 @@ def __init__(self, fragment, *args, # Avoids per-batch HDF5 open/close on the compute hot path; the file # is written once at scan end (processing_end). self.scan_buffer = [] + # Save-state tracking so flush never discards an unwritten scan buffer. + self._written = False + self._series_id = None self.publish_folder = publish_folder self.publish_tensors = publish_tensors if publish_tensors is not None else [] @@ -134,12 +137,26 @@ def write_scan_file(self, series_id): with h5py.File(filepath, 'w') as f: f.create_dataset('stxm', data=data) self.logger.info(f"Wrote {data.shape[0]} frames to {filepath}") + # Self-contained: mark saved and clear the buffer so a later flush has + # nothing to discard. + self._written = True + self.scan_buffer = [] def flush(self): - """Reset counters and discard any buffered scan data on flush.""" + """Reset counters. If the buffer still holds data that was never written + (a flush arrived before the end-of-scan write), write it out first so a + scan's STXM file is never lost.""" + if (self.scan_buffer and not self._written + and self.publish_folder is not None and self._series_id is not None): + self.logger.warning( + "Flush with %d unwritten STXM batch(es) — writing before clearing", + len(self.scan_buffer), + ) + self.write_scan_file(self._series_id) self.processed_frame_count = 0 self.processed_batch_count = 0 self.scan_buffer = [] + self._written = False def compute(self, op_input, op_output, context): """Receive, publish, and save processed data using metadata.""" @@ -199,12 +216,16 @@ def compute(self, op_input, op_output, context): if self.publish_folder is not None: if len(arrays_to_publish) > 0: self.scan_buffer.append(np.concatenate(arrays_to_publish, axis=1)) + self._written = False # new unsaved data + self._series_id = series_id # remembered for a defensive flush-save self.processed_batch_count += 1 self.processed_frame_count += tensor.shape[0] - # Check if processing is complete using metadata from upstream - if self.processed_frame_count == series_frame_count: + # Check if processing is complete using metadata from upstream. Use >= + # (not ==) so a batch that overshoots the exact count (e.g. frame count + # not a multiple of batch_size) still triggers the end-of-scan save. + if series_frame_count > 0 and self.processed_frame_count >= series_frame_count: # Write the whole scan once, now that no more batches are arriving if self.publish_folder is not None and series_id is not None: self.write_scan_file(series_id) From a5beeda5cd6b1703693001b2112e9ed89dfcb427 Mon Sep 17 00:00:00 2001 From: ramyaguru Date: Thu, 2 Jul 2026 14:02:23 -0400 Subject: [PATCH 04/33] PR2: Live header operator + dynamic scan geometry (config cleanup) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On-the-fly scan-geometry reconfiguration from a live JSON header, with mid-reconstruction preemption, on top of PR1's race-safe deferred flushes. - header_io.py (NEW): HeaderRxOp — dedicated ZMQ SUB, validates the header {npoints_h,npoints_v,step_size_h,step_size_v,num_projections}, stages pending_geometry + requests preemption, updates shared scan_state, emits a "header" token to ControlOp. Rejects over-capacity grids before staging. - ptychography_setup.py: split init_ptycho_state into load_ptycho_model (grid-independent) + configure_scan_geometry (grid-dependent object sizing, re-points views, needs_gpu_reinit, rejects no_frames > capacity) + init_ptycho_state (load + allocate GPU buffers ONCE at max capacity [R-6] + one default configure). New ptycho_state keys: H, W, capacity, scan_state, preempt_requested/quiesced Events, pending_geometry, needs_gpu_reinit. - ptychography_ops.py: R-4 quiescence handshake at the top of recon.compute — on preempt, save the in-flight partial + emit recon_complete, set quiesced, then apply staged geometry while quiesced and re-init GPU for the new object. R-6: full-buffer summary sliced to [:no_frames]. - control.py: new "header" case -> flush for the reconfigured scan. - pipeline.py: always-present scan_state (S11); build HeaderRxOp from the header_src config; wire header_src -> control_op; pass scan_state to init_ptycho_state. - config_test/prod.yaml: remove npoints_*/step_size_*; add header_src block + max_npoints_h/v (buffer capacity) + default_step_size_h/v. Validated in-container (A400): launch, back-to-back regression, header-before-data + full recon on reconfigured geometry, mid-recon preemption with partial save, post-preempt recovery, oversized-grid reject. Assisted-By: Claude Opus 4.8 (1M context) --- pipeline/config_prod.yaml | 15 ++ pipeline/config_test.yaml | 21 ++- pipeline/control.py | 11 ++ pipeline/header_io.py | 154 +++++++++++++++++++ pipeline/pipeline.py | 33 +++- pipeline/ptychography_ops.py | 114 +++++++++++++- pipeline/ptychography_setup.py | 266 +++++++++++++++++++++++++-------- 7 files changed, 542 insertions(+), 72 deletions(-) create mode 100644 pipeline/header_io.py diff --git a/pipeline/config_prod.yaml b/pipeline/config_prod.yaml index d225900..f59c0d2 100644 --- a/pipeline/config_prod.yaml +++ b/pipeline/config_prod.yaml @@ -18,6 +18,13 @@ position_src: zmq_endpoint: "tcp://172.23.82.204:6666" receive_timeout_ms: 1000 +header_src: + # Dedicated ZMQ SUB socket for live scan-geometry headers (PR2). A JSON + # object with npoints_h/npoints_v/step_size_h/step_size_v/num_projections + # reconfigures scan geometry on the fly and preempts an in-flight recon. + zmq_endpoint: "tcp://172.23.82.204:6667" # production endpoint (placeholder — adjust) + receive_timeout_ms: 100 + masking_op: center_x: 257 center_y: 515 @@ -38,6 +45,14 @@ ptychography: ptyrex_config: "/workdir/test_data/pty_config_15keV_streamTest_6mm.json" scan_ID: [1, 1, 1] ID: [1, 1, 1] + # Scan grid now comes from the live header (header_src). max_npoints_* set the + # GPU buffer capacity (allocated once, never realloced — R-6); default_step_* + # is the startup geometry before any header arrives. (Legacy no_frames/R/ + # scan_range below are unused — dead keys, S3, left for a separate cleanup.) + max_npoints_h: 32 # capacity = 32*32 = 1024 frames (matches prior no_frames: 1024) + max_npoints_v: 32 # bump these (with GPU RAM in mind) if larger scans are needed + default_step_size_h: 0.25 + default_step_size_v: 0.25 no_frames: 1024 total_iterations: 20 post_stream_iterations: 1 diff --git a/pipeline/config_test.yaml b/pipeline/config_test.yaml index b516910..bdb338e 100644 --- a/pipeline/config_test.yaml +++ b/pipeline/config_test.yaml @@ -20,6 +20,14 @@ position_src: zmq_endpoint: "tcp://172.23.82.204:6666" # production endpoint receive_timeout_ms: 1000 # Increased timeout for testing +header_src: + # Dedicated ZMQ SUB socket for live scan-geometry headers (PR2). A JSON + # object with npoints_h/npoints_v/step_size_h/step_size_v/num_projections + # reconfigures scan geometry on the fly and preempts an in-flight recon. + #zmq_endpoint: "tcp://172.23.82.77:5557" # Local simulator + zmq_endpoint: "tcp://172.23.82.204:6667" # production endpoint (placeholder) + receive_timeout_ms: 100 # short so the blocking recv doesn't hold a worker thread + masking_op: center_x: 524 # Adjusted for 192x192 images (center) center_y: 287 # Adjusted for 192x192 images (center) @@ -40,10 +48,15 @@ ptychography: ptyrex_config: "/workdir/PtyREX/config_409907.json" scan_ID: [1, 1, 1] ID: [1, 1, 1] - npoints_h: 100 # horizontal scan points - npoints_v: 100 # vertical scan points - step_size_h: 0.25 # horizontal step size (microns) - step_size_v: 0.25 # vertical step size (microns) + # Scan grid now comes from the live header (header_src). These set the GPU + # buffer capacity (allocated once, never realloced — R-6) and the default + # geometry used at startup before any header arrives. + max_npoints_h: 100 # buffer capacity + startup default grid (horizontal) + max_npoints_v: 100 # buffer capacity + startup default grid (vertical) + # capacity = 100*100 = 10000 frames (matches the prior committed 100x100 grid; + # a header requesting more frames is rejected — bump these with GPU RAM in mind) + default_step_size_h: 0.25 # startup step size (microns), until a header arrives + default_step_size_v: 0.25 # startup step size (microns), until a header arrives total_iterations: 25 post_stream_iterations: 1 housekeeping_interval: 1 diff --git a/pipeline/control.py b/pipeline/control.py index 64dcb65..908dd86 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -63,6 +63,17 @@ def compute(self, op_input, op_output, context): self._do_flush() self._flushed = True + elif msg == "header": + # A live header reconfigures the scan for a new dataset. Flush so the + # STXM path saves+clears its current buffer before reconfiguration + # (SinkAndPublishOp.flush writes any unwritten scan). This works even + # when ptychography is disabled; when enabled, the recon's own + # recon_complete (on quiesce) also flushes — harmless, flush is + # idempotent. Mark _flushed so the following start-flush skips. + self.logger.info("Header received — flushing for reconfigured scan") + self._do_flush() + self._flushed = True + elif msg == "flush": # Scan-start safety flush: only flush if the buffers aren't already # clean from a completion flush. If the previous scan completed, this diff --git a/pipeline/header_io.py b/pipeline/header_io.py new file mode 100644 index 0000000..faf9d66 --- /dev/null +++ b/pipeline/header_io.py @@ -0,0 +1,154 @@ +""" +Header Input for the Holoscan Ptycho Pipeline + +Defines HeaderRxOp: a dedicated ZMQ SUB operator that listens for live +scan-geometry headers and reconfigures the pipeline on the fly (PR2). + +A header is a JSON object, e.g.:: + + {"npoints_h": 100, "npoints_v": 100, + "step_size_h": 0.25, "step_size_v": 0.25, + "num_projections": 1} + +On a valid header the operator: + 1. updates the always-present shared ``scan_state`` (projection count), + 2. stages the new geometry in ``ptycho_state`` and requests preemption of an + in-flight reconstruction (R-4 handshake; the recon op applies the geometry + once it has finished the current iteration, saved, and quiesced), and + 3. emits a ``"header"`` token to ControlOp so the STXM path flushes for the + new dataset (works even when ptychography is disabled). + +Malformed headers are logged and ignored without disturbing an in-flight scan. +""" + +import logging + +import zmq + +from holoscan.core import Operator, OperatorSpec, ConditionType + + +class HeaderRxOp(Operator): + """Receive live scan-geometry headers over a dedicated ZMQ SUB socket.""" + + def __init__( + self, + fragment, + *args, + zmq_endpoint: str = None, + receive_timeout_ms: int = 100, + scan_state: dict = None, + ptycho_state: dict = None, + **kwargs, + ): + """ + Args: + fragment: Holoscan fragment + zmq_endpoint: ZMQ endpoint to connect to (e.g. "tcp://host:5557") + receive_timeout_ms: recv timeout in ms (short so the blocking recv + does not hold a worker thread for long) + scan_state: always-present shared holder for projection/frame counts + ptycho_state: ptycho state (None when ptychography is disabled) + """ + self.logger = logging.getLogger(kwargs.get("name", "HeaderRxOp")) + logging.basicConfig(level=logging.INFO) + + self.endpoint = zmq_endpoint + context = zmq.Context() + self.socket = context.socket(zmq.SUB) + self.socket.setsockopt_string(zmq.SUBSCRIBE, "") + self.socket.setsockopt(zmq.RCVTIMEO, receive_timeout_ms) + + try: + self.socket.connect(self.endpoint) + except zmq.error.ZMQError: + self.logger.error("Failed to connect header socket to %s", self.endpoint) + + self.scan_state = scan_state + self.ptycho_state = ptycho_state + + super().__init__(fragment, *args, **kwargs) + + def setup(self, spec: OperatorSpec): + # Token to ControlOp; NONE condition so this source runs freely. + spec.output("header").condition(ConditionType.NONE) + + def _validate(self, msg): + """Validate a header dict; return the parsed tuple or None if malformed.""" + if not isinstance(msg, dict): + self.logger.warning("Header is not a JSON object: %r", msg) + return None + try: + npoints_h = int(msg["npoints_h"]) + npoints_v = int(msg["npoints_v"]) + step_size_h = float(msg["step_size_h"]) + step_size_v = float(msg["step_size_v"]) + num_projections = int(msg.get("num_projections", 1)) + except (KeyError, TypeError, ValueError) as exc: + self.logger.warning("Malformed header %r: %s", msg, exc) + return None + if ( + npoints_h <= 0 or npoints_v <= 0 + or step_size_h <= 0 or step_size_v <= 0 + or num_projections < 1 + ): + self.logger.warning("Header has non-positive values: %r", msg) + return None + return npoints_h, npoints_v, step_size_h, step_size_v, num_projections + + def compute(self, op_input, op_output, context): + try: + msg = self.socket.recv_json() + except zmq.error.Again: + return # recv timed out — nothing to do this tick + except Exception as exc: # noqa: BLE001 - don't let a bad message kill the op + self.logger.warning("Header receive error: %s", exc) + return + + parsed = self._validate(msg) + if parsed is None: + return # malformed — ignore, leave any in-flight scan untouched + + npoints_h, npoints_v, step_size_h, step_size_v, num_projections = parsed + self.logger.info( + "Received header: %d x %d points, step %.4g x %.4g µm, " + "num_projections=%d", + npoints_h, npoints_v, step_size_h, step_size_v, num_projections, + ) + + # Reject a grid that exceeds the pre-allocated GPU capacity BEFORE staging + # it (buffers are never realloced, R-6). Rejecting here keeps the bad + # header off the recon's apply path, which would otherwise raise inside + # compute() and take down the pipeline. + if self.ptycho_state is not None: + capacity = self.ptycho_state.get("capacity") + if capacity is not None and npoints_h * npoints_v > capacity: + self.logger.warning( + "Header grid %dx%d = %d frames exceeds capacity %d — rejected " + "(increase max_npoints_h/max_npoints_v). In-flight scan " + "untouched.", + npoints_h, npoints_v, npoints_h * npoints_v, capacity, + ) + return + + # 1. Update the always-present shared holder (S11). + if self.scan_state is not None: + self.scan_state["num_projections"] = num_projections + self.scan_state["current_projection"] = 0 + + # 2. Ptycho path: stage geometry + request preemption (R-4). The recon op + # applies configure_scan_geometry once it has quiesced, so no buffer + # view is swapped under a live PIE iteration. + if self.ptycho_state is not None: + with self.ptycho_state["lock"]: + self.ptycho_state["pending_geometry"] = { + "npoints_h": npoints_h, + "npoints_v": npoints_v, + "step_size_h": step_size_h, + "step_size_v": step_size_v, + } + self.ptycho_state["preempt_requested"].set() + + # 3. Notify ControlOp so the STXM path flushes for the new dataset + # (works even when ptychography is disabled). + op_output.emit("header", "header") diff --git a/pipeline/pipeline.py b/pipeline/pipeline.py index acd5cb3..c756ecb 100644 --- a/pipeline/pipeline.py +++ b/pipeline/pipeline.py @@ -30,6 +30,7 @@ from processing import MaskingOp from publish import SinkAndPublishOp from control import ControlOp +from header_io import HeaderRxOp # Try to import NATS (optional for testing) try: @@ -58,6 +59,14 @@ def __init__(self, *args, **kwargs): self.num_decompress_ops = 4 self.ptychography_enabled = False self.ptycho_state = None + # Always-present shared holder (S11) for projection/frame counts, written + # by the header op and read by both the STXM and ptycho paths. Populated + # in main(); the frame count is filled in by configure_scan_geometry. + self.scan_state = { + "no_frames": 0, + "num_projections": 1, + "current_projection": 0, + } super().__init__(*args, **kwargs) self.enable_metadata(True) @@ -178,6 +187,20 @@ def compose(self): publish_backend=publish_backend, name="control_op") + # ===== Header Source (live scan geometry, optional) ===== + # A dedicated ZMQ SUB socket for JSON scan-geometry headers. Reconfigures + # geometry on the fly and preempts an in-flight recon (R-4 handshake). + header_src = None + header_cfg = self.kwargs('header_src') + if header_cfg: + header_src = HeaderRxOp(self, + name="header_src", + scan_state=self.scan_state, + ptycho_state=self.ptycho_state, + **header_cfg) + else: + logger.warning("No header_src config — live scan headers disabled") + # ===== Connect Operators ===== # I/O: Image reception and decompression -> gather for i in range(self.num_decompress_ops): @@ -201,6 +224,12 @@ def compose(self): # Control path: flush signal (start message → unconditional idempotent flush) self.add_flow(img_src, control_op, {("flush", "input")}) + # Header path: live geometry header → control (flush for new dataset). + # The ptycho geometry reconfigure is driven separately via the R-4 + # handshake in ptycho_state (header op sets preempt_requested). + if header_src is not None: + self.add_flow(header_src, control_op, {("header", "input")}) + def main(): """Main entry point for STXM pipeline.""" @@ -236,7 +265,9 @@ def main(): from ptychography_setup import init_ptycho_state logger.info("Initialising ptychography state…") - app.ptycho_state = init_ptycho_state(ptycho_cfg) + # Pass the shared scan_state so configure_scan_geometry can mirror the + # frame count for the STXM path and header op (S11). + app.ptycho_state = init_ptycho_state(ptycho_cfg, app.scan_state) app.ptychography_enabled = True worker_threads = max(worker_threads, 8) logger.info("Ptychography enabled (worker_threads=%d)", worker_threads) diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 9d88e6b..04521e2 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -179,10 +179,12 @@ def compute(self, op_input, op_output, context): self.ptycho_state["filled_until"] = new_end self._dirty = True - # Summary when buffer is full + # Summary when buffer is full. Buffers are allocated at capacity (R-6), + # so slice to the logical no_frames rather than the full buffer extent. if new_end >= self.ptycho_state["no_frames"]: - all_py = cp.asnumpy(self.ptycho_state["positions_full"][0, 0, :]) - all_px = cp.asnumpy(self.ptycho_state["positions_full"][0, 1, :]) + no_frames = self.ptycho_state["no_frames"] + all_py = cp.asnumpy(self.ptycho_state["positions_full"][0, 0, :no_frames]) + all_px = cp.asnumpy(self.ptycho_state["positions_full"][0, 1, :no_frames]) pty_model = self.ptycho_state["pty_model"] obj_h = int(pty_model.obj.sz_glo[-2]) obj_w = int(pty_model.obj.sz_glo[-1]) @@ -387,11 +389,117 @@ def _perform_flush(self): "ptychography.reset_probe: true in the config." ) + # ------------------------------------------------------------------ + # Header preemption handshake (R-4) + + def _save_and_signal_complete(self, op_output): + """Persist the in-flight partial result, then emit recon_complete. + + Called at the top of compute() when a header preemption is requested, so + the current reconstruction is saved (finish-current-iteration → save) + before the geometry is reconfigured. No-ops the save when the recon has + not produced anything yet (idle / pre-first-iteration). + """ + if self.initialized_gpu and self.current_iteration > 0: + pty_data = self.ptycho_state["pty_data"] + pty_model = self.ptycho_state["pty_model"] + pty_params = self.ptycho_state["pty_params"] + # Bring the object/probe to host and run the end-of-iteration save + # (writes pty_out). We are reconfiguring next, so we do not push the + # arrays back to device (no to_device). + from_device(pty_model, pty_params) + setup.after_iteration(pty_data, pty_model, pty_params, pty_plot=None) + obj_2d = np.squeeze(cp.asnumpy(pty_model.obj.array_global)) + probe_2d = np.squeeze(cp.asnumpy(pty_model.probe.array_states)) + out = { + "object_phase": np.angle(obj_2d).astype(np.float32), + "object_amp": np.abs(obj_2d).astype(np.float32), + "probe_phase": np.angle(probe_2d).astype(np.float32), + "probe_amp": np.abs(probe_2d).astype(np.float32), + "iteration": self.current_iteration, + } + op_output.emit(out, "output") + self.logger.info( + "Saved partial result before preemption (iter %d)", + self.current_iteration, + ) + # Signal completion → ControlOp flushes all ops for the next dataset. + if not self._completed: + self._completed = True + op_output.emit("recon_complete", "complete") + + def _apply_pending_geometry(self): + """Apply the header's staged geometry while quiesced, then clear the + handshake. Safe because no PIE iteration is in flight (Phase 2).""" + from ptychography_setup import configure_scan_geometry + + with self.lock: + pending = self.ptycho_state.get("pending_geometry") + self.ptycho_state["pending_geometry"] = None + if pending is not None: + try: + configure_scan_geometry(self.ptycho_state, **pending) + self.logger.info("Applied new scan geometry from header") + except Exception: + # Never let a bad reconfigure take down the pipeline; keep the + # previous geometry and resume. (HeaderRxOp already rejects + # over-capacity grids; this guards anything unexpected.) + self.logger.exception( + "Failed to apply new scan geometry — keeping previous geometry" + ) + self.ptycho_state["needs_gpu_reinit"] = False + # A full geometry reconfigure subsumes any pending flush (the object is + # rebuilt from scratch), so drop a stale deferred flush that would + # otherwise reference the previous geometry's initial-object snapshot. + self._flush_requested = False + # Clear the handshake — next compute re-inits GPU for the new object. + self.ptycho_state["quiesced"].clear() + self.ptycho_state["preempt_requested"].clear() + + def _reset_for_new_geometry(self): + """Re-init reconstruction state after a geometry change. + + configure_scan_geometry rebuilt the object arrays on the host, so the + one-time GPU transfer and the pristine-object snapshot must be redone. + """ + self.ptycho_state["needs_gpu_reinit"] = False + self.initialized_gpu = False + self._obj_initial = None + self._probe_initial = None + self._flux_initial = None + self.current_iteration = 0 + self.all_data_arrived = False + self.post_stream_count = 0 + self._completed = False + self.logger.info("Recon reset for new scan geometry") + def compute(self, op_input, op_output, context): + # Header preemption handshake (R-4) takes priority over any flush so the + # in-flight partial is saved with the object still intact. + if self.ptycho_state["preempt_requested"].is_set(): + if not self.ptycho_state["quiesced"].is_set(): + # Phase 1: the current iteration is already finished (compute is + # atomic). Save the partial result + signal completion, then + # quiesce so the geometry can be re-pointed without racing a + # live PIE view. + self._save_and_signal_complete(op_output) + self.ptycho_state["quiesced"].set() + self.logger.info("Recon quiesced for header preemption") + else: + # Phase 2: still preempted and quiesced — apply the staged + # geometry now (nothing is touching the buffers) and clear the + # handshake. Next compute re-inits GPU for the new object. + self._apply_pending_geometry() + return + # Perform any requested flush here — single-threaded w.r.t. the PIE update. if self._flush_requested: self._perform_flush() + # Geometry changed while quiesced — re-init GPU state for the new object. + if self.ptycho_state.get("needs_gpu_reinit"): + self._reset_for_new_geometry() + # Snapshot fill level with self.lock: n_filled = self.ptycho_state["filled_until"] diff --git a/pipeline/ptychography_setup.py b/pipeline/ptychography_setup.py index 77e6947..2c9ceef 100644 --- a/pipeline/ptychography_setup.py +++ b/pipeline/ptychography_setup.py @@ -1,11 +1,15 @@ """ Ptychography State Initialization -Builds the shared ptycho_state dict at application launch by: -1. Loading PtyREX model from JSON config -2. Computing scan extent from npoints and step_size (matching PtyREX streaming) -3. Running one-time PtyREX setup (mirrors pre_process_reconstruct_stream) -4. Pre-allocating GPU buffers +Builds the shared ptycho_state dict at application launch, split into a +grid-INDEPENDENT one-time model load and a grid-DEPENDENT geometry +configuration so scan geometry can be reconfigured on the fly from a live +header (PR2) without reallocating GPU buffers (R-6): + +1. ``load_ptycho_model`` — load PtyREX model, jsplitter, detector pre-load +2. GPU buffers allocated once at the configured MAX capacity (never realloced) +3. ``configure_scan_geometry`` — grid-dependent object sizing + view re-pointing +4. ``init_ptycho_state`` — load + allocate-at-max + one default configure """ import threading @@ -54,12 +58,24 @@ def update(self, *args, **kwargs): pass -def init_ptycho_state(ptycho_cfg: dict) -> dict: - """Build ptycho_state from PtyREX JSON config + pipeline YAML overrides. +def _ensure_ptyrex_on_path(): + """Make the vendored PtyREX package importable (mirrors the original + lazy sys.path insertion so the host can import this module without PtyREX).""" + import importlib.util, os, sys + + _spec = importlib.util.find_spec("ptyrex") + if _spec and _spec.origin: + _ptyrex_root = os.path.dirname(os.path.dirname(_spec.origin)) + if _ptyrex_root not in sys.path: + sys.path.insert(0, _ptyrex_root) + - The only scan parameters required are npoints_h, npoints_v, step_size_h, - and step_size_v — the scan extent in pixels (R) and object size are - derived automatically, matching PtyREX's streaming workflow. +def load_ptycho_model(ptycho_cfg: dict): + """Grid-INDEPENDENT one-time load. + + Loads the PtyREX model from JSON, installs the single-rank jsplitter stub, + and runs the detector pre-load (crop geometry). Nothing here depends on the + scan grid, so it runs exactly once at startup. Parameters ---------- @@ -68,23 +84,11 @@ def init_ptycho_state(ptycho_cfg: dict) -> dict: Returns ------- - dict - Shared state containing PtyREX model objects and pre-allocated - GPU buffers. + (pty_data, pty_model, pty_params, H, W) + PtyREX objects plus the cropped detector frame size (H, W). """ - import importlib.util, os, sys - import cupy as cp - - _spec = importlib.util.find_spec("ptyrex") - if _spec and _spec.origin: - _ptyrex_root = os.path.dirname(os.path.dirname(_spec.origin)) - if _ptyrex_root not in sys.path: - sys.path.insert(0, _ptyrex_root) - + _ensure_ptyrex_on_path() from ptyrex.core.io import json_read - from ptyrex.reconstruct.core import setup - from ptyrex.reconstruct.iterator.process_pty_model import generate_grow_scan_params - from ptyrex.reconstruct.utils import numpy as utils_np # ── 1. Load PtyREX model from JSON config ────────────────────────── ptyrex_config_path = ptycho_cfg["ptyrex_config"] @@ -92,25 +96,6 @@ def init_ptycho_state(ptycho_cfg: dict) -> dict: ID = ptycho_cfg.get("ID", [1, 1, 1]) pty_data, pty_model, pty_params = json_read.load(ptyrex_config_path, scan_ID, ID) - # ── 2. Compute streaming parameters from npoints / step_size ─────── - npoints_h = ptycho_cfg["npoints_h"] - npoints_v = ptycho_cfg["npoints_v"] - step_size_h = ptycho_cfg["step_size_h"] - step_size_v = ptycho_cfg["step_size_v"] - - no_frames = npoints_h * npoints_v - # Scan extent in microns with 20% padding (same formula as PtyREX streaming) - N = [ - ((npoints_v - 1) * step_size_v) * 1.2, - ((npoints_h - 1) * step_size_h) * 1.2, - ] - logger.info( - "Scan: %d x %d points, step %.3f x %.3f µm → " - "N = [%.2f, %.2f] µm, %d frames", - npoints_h, npoints_v, step_size_h, step_size_v, - N[1], N[0], no_frames, - ) - pty_params.total_iterations = ptycho_cfg["total_iterations"] # Ensure string attributes expected by PtyREX save/config routines @@ -118,13 +103,76 @@ def init_ptycho_state(ptycho_cfg: dict) -> dict: pty_data.ID = str(pty_data.ID[0]) if isinstance(pty_data.ID, list) else str(pty_data.ID) pty_data.scan_ID = str(pty_data.scan_ID[0]) if isinstance(pty_data.scan_ID, list) else str(pty_data.scan_ID) - # ── 3. Dummy jsplitter for single-rank pipeline ──────────────────── + # ── Dummy jsplitter for single-rank pipeline ─────────────────────── pty_params.jsplitter = DummyJSplitter() - # ── 4. Initialise scan arrays (mirrors pre_process_reconstruct_stream) ─ + # ── Detector pre-load (crop geometry — grid-independent) ─────────── pty_data.pre_load(pty_model.detector, pty_params) H = int(pty_data.crop_bottom - pty_data.crop_top) W = int(pty_data.crop_right - pty_data.crop_left) + + return pty_data, pty_model, pty_params, H, W + + +def configure_scan_geometry( + ptycho_state: dict, + npoints_h: int, + npoints_v: int, + step_size_h: float, + step_size_v: float, +): + """Grid-DEPENDENT (re)configuration of scan geometry. + + Recomputes ``no_frames`` and the scan extent ``N``, re-runs the PtyREX + object-sizing setup, refreshes the detector pixel mask / dp transforms, and + re-points the scan arrays at the pre-allocated GPU buffers. **No GPU + reallocation** (R-6): a grid whose ``no_frames`` exceeds the configured + capacity is rejected with ``ValueError``. + + Safe to call at startup (from ``init_ptycho_state``) and again on a live + header, but ONLY while the reconstruction is quiesced (R-4 handshake) — it + re-points views the recon reads during a PIE iteration. + + Sets ``ptycho_state["needs_gpu_reinit"] = True`` so the reconstruction op + re-runs its one-time GPU transfer and re-snapshots the pristine object for + the new geometry. + """ + _ensure_ptyrex_on_path() + from ptyrex.reconstruct.core import setup + from ptyrex.reconstruct.iterator.process_pty_model import generate_grow_scan_params + from ptyrex.reconstruct.utils import numpy as utils_np + import cupy as cp + + pty_data = ptycho_state["pty_data"] + pty_model = ptycho_state["pty_model"] + pty_params = ptycho_state["pty_params"] + H = ptycho_state["H"] + W = ptycho_state["W"] + capacity = ptycho_state["capacity"] + + # ── Compute streaming parameters from npoints / step_size ────────── + no_frames = int(npoints_h) * int(npoints_v) + if no_frames > capacity: + raise ValueError( + f"Requested grid {npoints_h}x{npoints_v} = {no_frames} frames exceeds " + f"the pre-allocated capacity of {capacity} frames " + f"(increase max_npoints_h/max_npoints_v in the config). " + f"Buffers are never reallocated at runtime (R-6)." + ) + + # Scan extent in microns with 20% padding (same formula as PtyREX streaming) + N = [ + ((npoints_v - 1) * step_size_v) * 1.2, + ((npoints_h - 1) * step_size_h) * 1.2, + ] + logger.info( + "Configuring scan: %d x %d points, step %.3f x %.3f µm → " + "N = [%.2f, %.2f] µm, %d frames (capacity %d)", + npoints_h, npoints_v, step_size_h, step_size_v, + N[1], N[0], no_frames, capacity, + ) + + # ── Initialise scan arrays (mirrors pre_process_reconstruct_stream) ─ pty_data.raw = np.zeros((no_frames, H, W), dtype=np.uint32) pty_model.scan.positions = np.ones([no_frames, 2], np.float32) @@ -144,13 +192,15 @@ def init_ptycho_state(ptycho_cfg: dict) -> dict: pty_model.scan.sz = [pty_model.scan.positions.shape[0], 1] pty_data.reg_ind = pty_model.scan.reg_ind - # ── 5. PtyREX one-time setup — pass N in microns (not R in pixels) ─ + # ── PtyREX one-time setup — pass N in microns (not R in pixels) ──── pty_plot = DummyPtyPlot() pty_data, pty_model, pty_params, pty_plot = setup.before_reconstruction_stream( pty_data, pty_model, pty_params, pty_plot, N ) - # ── 6. Pixel mask + dp transforms (mirrors post_process_stream) ── + # ── Pixel mask + dp transforms (mirrors post_process_stream) ─────── + # Re-run every reconfigure so ordering matches the validated single-scan + # path (before_reconstruction_stream then get_pixel_mask); cheap. df, ff, dp = pty_data.get_pixel_mask(pty_model.detector, pty_params) pty_data.dp = dp[ pty_data.crop_top : pty_data.crop_bottom, @@ -164,10 +214,9 @@ def init_ptycho_state(ptycho_cfg: dict) -> dict: pty_model.detector.mask[pty_data.dp > 0] = 0 pty_model.detector.mask_inv[pty_data.dp > 0] = 1 - # Now apply the same dp transforms that post_process_stream does on - # its first iteration so that dp matches the preprocessed data layout - # and uses the inverted convention expected by the flux computation - # (dp == 1 → good pixel). + # Apply the same dp transforms that post_process_stream does on its first + # iteration so dp matches the preprocessed data layout and uses the inverted + # convention expected by the flux computation (dp == 1 → good pixel). det = pty_model.detector if det.orientation == "01": pty_data.dp = pty_data.dp[:, ::-1] @@ -187,26 +236,91 @@ def init_ptycho_state(ptycho_cfg: dict) -> dict: pty_model.obj.array_global_old[:] = pty_model.obj.array_global[:] generate_grow_scan_params(pty_params) - # ── 7. Pre-allocate GPU buffers ──────────────────────────────────── - raw_gpu = cp.zeros((no_frames, H, W), dtype=cp.uint32) - positions_full = cp.zeros((1, 2, no_frames), dtype=cp.float32) - tilts_full = cp.zeros((1, 2, no_frames), dtype=cp.float32) - + # ── Re-point scan arrays at the pre-allocated GPU buffers ────────── + # The buffers are allocated once at capacity in init_ptycho_state and never + # realloced; the accumulator/recon index them via ptycho_state directly and + # bound their reads to [:n_filled] (n_filled <= no_frames <= capacity). + positions_full = ptycho_state["positions_full"] + tilts_full = ptycho_state["tilts_full"] pty_model.scan.positions = positions_full pty_model.scan.tilts = tilts_full pty_model.scan.original = cp.zeros_like(positions_full) pty_model.scan.previous = cp.zeros_like(positions_full) + # ── Update shared state ──────────────────────────────────────────── + with ptycho_state["lock"]: + ptycho_state["no_frames"] = no_frames + ptycho_state["filled_until"] = 0 + ptycho_state["N"] = N + # Clear auto-centre so the new geometry re-derives its own scan centre. + ptycho_state["scan_center_py"] = None + ptycho_state["scan_center_px"] = None + # Signal the reconstruction op to re-init GPU state for the new object size. + ptycho_state["needs_gpu_reinit"] = True + + # Mirror the frame count into the always-present scan_state holder (S11) so + # the STXM path and header op can read it even when ptycho is disabled. + scan_state = ptycho_state.get("scan_state") + if scan_state is not None: + scan_state["no_frames"] = no_frames + scan_state["npoints_h"] = int(npoints_h) + scan_state["npoints_v"] = int(npoints_v) + scan_state["step_size_h"] = float(step_size_h) + scan_state["step_size_v"] = float(step_size_v) + logger.info( - "ptycho_state initialized: %d frames, image size %dx%d, " - "object size %s, %d total iterations", + "Scan geometry configured: %d frames, image size %dx%d, object size %s", no_frames, H, W, tuple(int(x) for x in pty_model.obj.sz_glo), - pty_params.total_iterations, ) - # ── 8. Assemble ptycho_state ─────────────────────────────────────── - return { + +def init_ptycho_state(ptycho_cfg: dict, scan_state: dict = None) -> dict: + """Build ptycho_state: load the model, allocate GPU buffers at the configured + MAX capacity, then configure a default scan geometry so the pipeline can run + before any header arrives. + + Parameters + ---------- + ptycho_cfg : dict + The ``ptychography`` section of the pipeline YAML config. Must provide + ``max_npoints_h``/``max_npoints_v`` (buffer capacity + default grid) and + ``default_step_size_h``/``default_step_size_v`` (startup step sizes). + scan_state : dict, optional + The always-present shared holder for projection/frame counts (S11). Its + ``no_frames`` is populated by the default configure below. + + Returns + ------- + dict + Shared state containing PtyREX model objects, pre-allocated GPU buffers + (at max capacity), geometry, and the preemption handshake primitives. + """ + import cupy as cp + + # ── 1. Grid-independent model load ───────────────────────────────── + pty_data, pty_model, pty_params, H, W = load_ptycho_model(ptycho_cfg) + + # ── 2. Capacity + default startup grid ───────────────────────────── + max_npoints_h = int(ptycho_cfg["max_npoints_h"]) + max_npoints_v = int(ptycho_cfg["max_npoints_v"]) + capacity = max_npoints_h * max_npoints_v + default_step_h = float(ptycho_cfg["default_step_size_h"]) + default_step_v = float(ptycho_cfg["default_step_size_v"]) + + # ── 3. Pre-allocate GPU buffers ONCE at max capacity (R-6) ───────── + raw_gpu = cp.zeros((capacity, H, W), dtype=cp.uint32) + positions_full = cp.zeros((1, 2, capacity), dtype=cp.float32) + tilts_full = cp.zeros((1, 2, capacity), dtype=cp.float32) + + logger.info( + "ptycho_state buffers allocated at capacity: %d frames, " + "image size %dx%d, %d total iterations", + capacity, H, W, pty_params.total_iterations, + ) + + # ── 4. Assemble ptycho_state (geometry filled in by configure) ───── + ptycho_state = { "pty_data": pty_data, "pty_model": pty_model, "pty_params": pty_params, @@ -214,9 +328,33 @@ def init_ptycho_state(ptycho_cfg: dict) -> dict: "positions_full": positions_full, "tilts_full": tilts_full, "filled_until": 0, - "no_frames": no_frames, + "no_frames": 0, # set by configure_scan_geometry + "H": H, + "W": W, + "capacity": capacity, "scan_center_py": None, "scan_center_px": None, - "N": N, + "N": None, # set by configure_scan_geometry "lock": threading.Lock(), + "scan_state": scan_state, + # Preemption handshake (R-4): header stages pending_geometry + sets + # preempt_requested; recon saves the partial, sets quiesced, then applies + # the geometry while quiesced and clears the flags. + "preempt_requested": threading.Event(), + "quiesced": threading.Event(), + "pending_geometry": None, + "needs_gpu_reinit": False, } + + # ── 5. Configure the default (startup) scan geometry ─────────────── + configure_scan_geometry( + ptycho_state, + npoints_h=max_npoints_h, + npoints_v=max_npoints_v, + step_size_h=default_step_h, + step_size_v=default_step_v, + ) + # First configure just did startup init; no reconfigure has happened yet. + ptycho_state["needs_gpu_reinit"] = False + + return ptycho_state From 6ec7d59f7dd68b8cfabcab24ad28db933bb6ddb7 Mon Sep 17 00:00:00 2001 From: ramyaguru Date: Thu, 2 Jul 2026 16:12:58 -0400 Subject: [PATCH 05/33] =?UTF-8?q?PR3:=20Tomography=20=E2=80=94=20multi-pro?= =?UTF-8?q?jection=20accumulation,=20per-projection=20recon=20+=20STXM=20o?= =?UTF-8?q?utput?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One arm/start wraps all projections; num_projections × no_frames stream continuously and the projection boundary is segmented by frame count (series_id shared across projections). Each projection early-stops once its frames are in and the in-flight PIE iteration finishes, then advances; single scans still run to total_iterations. - GatherOp: drain cached frames on series_finished (R-1). - PtychoAccumulatorOp: projection-boundary split with carry, backpressure once filled >= no_frames, advance_projection() resets fill level while preserving carry + centre. - PtychoReconstructionOp: tomography is_last = all_data_arrived; per-projection save ({series}_proj{NN}_recon.h5); self-advances on observing filled drop below no_frames (avoids double-complete); emits projection_complete (non-final) vs recon_complete (final). - ControlOp: projection_complete → scoped advance (accumulator + recon only, NOT GatherOp — its next-projection cache must survive). - SinkAndPublishOp: per-projection STXM save ({series}_proj{NN}.h5) reading scan_state; header_io sets scan_state no_frames; pipeline wires scan_state. - ITER_TIMING per-iteration timing diagnostic retained (sub-ms overhead). - Fix off-by-one in the dummy_img_index test path (series_frame_count - 1 → series_frame_count) exposed while testing; test-only path (default False). Container-tested on the A400 with a 2-projection header over a 2048-frame daqsim stream: both projections reconstruct and save (STXM + recon), correct 1024-frame counts, and the non-final/final completion signals distinguished. See DECISIONS_header_tomo.md / TESTING_header_tomo.md (incl. open item O1 on per-projection post-stream refinement iterations, deferred to PR4). Assisted-By: Claude Opus 4.8 (1M context) --- pipeline/DECISIONS_header_tomo.md | 34 ++++++ pipeline/TESTING_header_tomo.md | 62 ++++++++++ pipeline/control.py | 28 ++++- pipeline/data_io.py | 22 +++- pipeline/header_io.py | 5 +- pipeline/pipeline.py | 7 ++ pipeline/ptychography_ops.py | 196 +++++++++++++++++++++++++++--- pipeline/publish.py | 83 ++++++++++--- 8 files changed, 398 insertions(+), 39 deletions(-) create mode 100644 pipeline/DECISIONS_header_tomo.md create mode 100644 pipeline/TESTING_header_tomo.md diff --git a/pipeline/DECISIONS_header_tomo.md b/pipeline/DECISIONS_header_tomo.md new file mode 100644 index 0000000..48b0b94 --- /dev/null +++ b/pipeline/DECISIONS_header_tomo.md @@ -0,0 +1,34 @@ +# Decision log — header / dynamic-geometry / tomography feature + +Design decisions made while implementing the `feat/scan-header-tomo` branch (PR0→PR3), +recorded for review. Dates are when the decision was made. + +## Git / workflow +| # | Decision | Rationale | +|---|----------|-----------| +| 1 | Keep the beamline sim-test tweaks OUT of the feature commits — held in a reversible patch `test_data/sim_code_tweaks.patch` (position mapping `/FMC_IN.VAL1`, `PTYCHO_CENTER` override), applied for testing, reverted before committing. | The tweaks are local test scaffolding (localhost endpoints, fixed scan centre) that must not ship in production code. | +| 2 | Single feature branch `feat/scan-header-tomo` for the whole effort (renamed from `feat/pr1-flush-plumbing`), not stacked per-PR branches. | Simpler to manage/review as one branch. | +| 3 | Dropped the pre-existing `736faf4 "Commit before merge"` via rebase; `Dockerfile_bwell` kept locally (untracked), not pushed. | That commit mixed a Blackwell Dockerfile with debug prints later removed; not part of this feature. | +| 4 | Commit sign-off trailer is `Assisted-By: Claude Opus 4.8 (1M context)`, not `Co-Authored-By:`. | Reflects Claude's assistant role. Applies to all future commits. | + +## Architecture (from the plan `dls-holoscan-header-tomo-plan.md`) +| # | Decision | Rationale | +|---|----------|-----------| +| 5 | Header transport = a **dedicated ZMQ SUB socket** (`header_src`), separate from images/positions. | Keeps the geometry channel independent; always listening. | +| 6 | GPU buffers allocated **once at max capacity** (`max_npoints_h/v`), never realloced at runtime (R-6). A header requesting more frames than capacity is **rejected**. | Runtime cupy realloc under live op references risks fragmentation/leaks and a swap-under-recon race. | +| 7 | Header preemption uses a **quiescence handshake** (R-4): header stages geometry + sets `preempt_requested`; recon finishes the in-flight iteration → saves the partial → signals complete → quiesces → applies the new geometry → re-inits GPU. | The recon holds buffer *views* across a PIE iteration; geometry can only change safely while it's idle. | +| 8 | Flush model = flush-on-completion + a skip-if-clean safety flush at scan start (PR1). | Idempotent; removes stale-buffer bugs without double-flushing. | + +## Tomography (PR3) +| # | Decision | Rationale | +|---|----------|-----------| +| 9 | One `arm`/`start` wraps **all projections**; `num_projections × no_frames` frames stream continuously between a single start/end. Projection boundary is segmented **by frame count**. | Matches the acquisition model; `series_id` (Dectris series counter, one per arm) is shared across all projections. | +| 10 | Each projection **early-stops** as soon as its `no_frames` are accumulated + the in-flight iteration finishes (does NOT run full `total_iterations` per projection). Single-projection scans still run to `total_iterations`. | Throughput: keep up with a continuous multi-projection stream. | +| 11 | **Projection-boundary advance is scoped** (2026-07-02): flush ONLY the accumulator (reset GPU buffer) + recon (reset object/iters) + advance the STXM sink and `current_projection`. **Do NOT flush GatherOp** — its cached next-projection frames must survive. Only the FINAL projection triggers a full scan-end flush (incl. GatherOp). | A full flush at every boundary would drop the next projection's frames piling up in GatherOp's cache (R-5). | +| 12 | Per-projection output files named `{series_id}_proj{NN}.h5` for both STXM and ptycho recon (2026-07-02). Theta omitted from the name for now (≈0 in current test data; easy to add). | All projections share one `series_id`, so the projection index is mandatory to avoid overwrite (M2). | +| 13 | Single-buffer first with a bounded GatherOp cache; add double-buffering only if a load test shows the boundary backlog overflows (PR4). | Avoid paying 2× memory + complexity before evidence it's needed. | + +## Open questions (not blocking — revisit in PR4) +| # | Question | Context | +|---|----------|---------| +| O1 | Should tomography projections get a fixed number of **post-stream refinement iterations** before advancing, instead of stopping the instant all frames arrive? | For tomography (`num_projections > 1`), `is_last = all_data_arrived` (`ptychography_ops.py`), so a projection ends as soon as its `no_frames` are in — in the container test proj-0 got 7 PIE iterations and proj-1 only 4 (its frames were already cached by GatherOp during proj-0, so `all_data_arrived` tripped almost immediately). The single-scan branch grants `post_stream_iterations`; the tomography branch ignores them entirely. On a fast stream this means per-projection reconstructions are coarse. Lever: let each projection run N post-stream iterations after its frames arrive but before advancing, trading throughput for per-projection recon quality. | diff --git a/pipeline/TESTING_header_tomo.md b/pipeline/TESTING_header_tomo.md new file mode 100644 index 0000000..7a93b1c --- /dev/null +++ b/pipeline/TESTING_header_tomo.md @@ -0,0 +1,62 @@ +# Testing log — header / dynamic-geometry / tomography feature + +Branch `feat/scan-header-tomo` (PR0 → PR3). Tests run in the `ptycho-holoscan:ptyrex` +container on the **A400** (`CUDA_VISIBLE_DEVICES=1`), driven by the daqsim simulator +(`daqsim:latest`, Dectris SIMPLON emulator) streaming scan **409907** (32×32 = 1024 +frames, 515×515 uint32). Sim tweaks applied via `test_data/sim_code_tweaks.patch`; +geometry headers sent with `test_data/send_header.py`; scans triggered with +`dectris-hackathon/daqsim/trigger.py --stream both --nimages 1024`. + +## PR0 — Unwire PublishToCloudOp +- Regression: STXM + ptycho scans produce identical output with the op unwired. Pipeline + composes and runs. (Validated during the PR0/PR1 session.) + +## PR1 — Flush plumbing + hardening +- Single ptycho scan reconstructs to `total_iterations` and idles; `recon_complete` logged. +- Back-to-back second scan: start-flush safety net fires, clean second reconstruction, two + output HDFs. No-double-flush confirmed (`Start-flush skipped — already flushed on completion`). + +## PR2 — Live header operator + dynamic geometry — ALL PASS (2026-07-02) +Container run, logs captured to `test_data/run34.log`. + +| # | Test | Result | +|---|------|--------| +| 1 | Launch with `npoints_*` removed from config | ✅ `buffers allocated at capacity: 1024 frames`, `Scan geometry configured: 1024 frames … object size (…,523,523)` | +| 2 | Baseline scan + back-to-back regression | ✅ both reconstruct to iter 25 and flush; no-double-flush holds | +| 3 | Header before data → full scan on reconfigured geometry | ✅ full handshake (`Received header`→`Header received`→`Recon quiesced`→`Applied new scan geometry`→`Recon reset`); scan then reconstructs to completion, **all 1024 positions valid**. 2nd `before_reconstruction_stream` stable. | +| 4 | Header mid-reconstruction (preemption) | ✅ `Saved partial result before preemption (iter 24)` emitted **before** flush/reconfigure, then quiesce → reconfigure → reset | +| 4b | Recovery after preemption | ✅ `GPU initialisation complete` → fresh scan `Reconstruction complete at iteration 25` | +| 5 | Oversized grid (64×64 = 4096 > capacity 1024) | ✅ `Header grid 64x64 … exceeds capacity 1024 — rejected`; rejected before staging, pipeline stays alive | + +**Caveat:** reconfigure tested only to the *same* geometry (32×32 / 1.5 µm, object stayed +523×523). The reconfigure + GPU-reinit code path is fully exercised; a change to a +*different* object size with matching stream data is unverified (needs a 2nd dataset). + +## PR3 — Tomography / multi-projection — PASS (2026-07-02) +Container run (A400, `config_sim.yaml`), 2-projection header sent, 2048-frame stream +(`trigger.py --stream both --nimages 2048`, sequence_id 29). Logs in `test_data/run_pr3c.log`. + +| # | Test | Result | +|---|------|--------| +| 1 | Header sets tomography mode | ✅ `Received header … num_projections=2` → handshake → geometry applied | +| 2 | Projection 0 completes at frame boundary | ✅ `Wrote projection 0 (1024 frames) → 29_proj00.h5`, `Saved projection recon → 29_proj00_recon.h5`, `Projection 0/2 complete … signal=projection_complete` (early-stopped at iter 7) | +| 3 | Scoped advance (accum + recon reset, GatherOp cache preserved) | ✅ `Projection complete — advanced to projection 1`; accumulator advanced with carry preserved, recon object reset / probe carried, GatherOp **not** flushed | +| 4 | Projection 1 completes | ✅ `Wrote projection 1 (1024 frames) → 29_proj01.h5`, `Saved projection recon → 29_proj01_recon.h5`, `Projection 1/2 complete … signal=recon_complete` (final; early-stopped at iter 4) | +| 5 | Final flush | ✅ `Reconstruction complete — flushing for next scan` | + +All four files on disk (`series_id=29`, both projections × STXM + recon), 1024-frame counts +exact, and the non-final vs final completion signals correctly distinguished +(`projection_complete` for proj-0, `recon_complete` for proj-1). No stall. + +**Test-tweak note:** the run used `dummy_img_index=True` (continuous synthetic image IDs) +to work around a daqsim looped-replay artifact where image IDs reset per loop while +position IDs continue. Testing exposed an off-by-one in that path +(`data_io.py`: `series_frame_count - 1` produced IDs `-1..N-2`, leaving one frame +unmatched per scan) — fixed to `series_frame_count`. `dummy_img_index` reverted to +`False` before commit; a real single-series tomography acquisition (monotonic IDs) +would not hit the artifact. + +**Open item (see DECISIONS O1):** each projection early-stops the moment its frames +arrive, so proj-1 got fewer PIE iterations than proj-0 (cached frames tripped +`all_data_arrived` immediately). Per-projection post-stream refinement iterations +deferred to PR4. diff --git a/pipeline/control.py b/pipeline/control.py index 908dd86..db7f61b 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -19,19 +19,29 @@ class ControlOp(Operator): def __init__(self, fragment, *args, flushable_ops: list[Operator] = None, publish_backend = None, + ptycho_accum = None, + ptycho_recon = None, + scan_state: dict = None, **kwargs): """ Initialize control operator. - + Args: fragment: Holoscan fragment flushable_ops: List of operators that can be flushed publish_backend: Backend instance for publishing flush messages + ptycho_accum: PtychoAccumulatorOp (for the scoped projection advance) + ptycho_recon: PtychoReconstructionOp (for the scoped projection advance) + scan_state: shared holder whose current_projection is advanced at a + tomography projection boundary (PR3) """ super().__init__(fragment, *args, **kwargs) self.logger = logging.getLogger(kwargs.get("name", "ControlOp")) self.flushable_ops = flushable_ops self.publish_backend = publish_backend + self.ptycho_accum = ptycho_accum + self.ptycho_recon = ptycho_recon + self.scan_state = scan_state # True once a completion (recon_complete) flush has run and no new scan # has started since. Lets the scan-start flush skip when the buffers are # already clean, so we don't double-flush (Task 3 flush-check-at-start). @@ -63,6 +73,22 @@ def compute(self, op_input, op_output, context): self._do_flush() self._flushed = True + elif msg == "projection_complete": + # PR3 tomography per-projection boundary: SCOPED advance. Reset the + # accumulator's fill level (carry preserved) and bump current_projection. + # The recon self-advances once it observes filled_until drop (so it can't + # re-complete the same projection). Do NOT flush GatherOp — its cached + # next-projection frames must survive (decision #11). The STXM sink + # segments itself by frame count (S2), so it isn't touched here. + if self.ptycho_accum is not None: + self.ptycho_accum.advance_projection() + if self.scan_state is not None: + self.scan_state["current_projection"] += 1 + self.logger.info( + "Projection complete — advanced to projection %d", + self.scan_state["current_projection"], + ) + elif msg == "header": # A live header reconfigures the scan for a new dataset. Flush so the # STXM path saves+clears its current buffer before reconfiguration diff --git a/pipeline/data_io.py b/pipeline/data_io.py index c0c0233..f0c46cd 100644 --- a/pipeline/data_io.py +++ b/pipeline/data_io.py @@ -376,7 +376,7 @@ def compute(self, op_input, op_output, context): self.batch[self.current_index] = data if self.dummy_img_index: - self.batch_ids[self.current_index] = self.series_frame_count - 1 + self.batch_ids[self.current_index] = self.series_frame_count else: # self.logger.info(f"Received image with id: {data_id}") self.batch_ids[self.current_index] = data_id @@ -472,6 +472,9 @@ def __init__(self, fragment, *args, batch_size: int = 1, **kwargs): self.position_ids = np.zeros((0,), dtype=int) self.count = 0 self.batch_size = int(batch_size) + # R-1 (PR3): latched once the series-end metadata is seen, so the final + # partial batch (< batch_size) is drained instead of stranded. Reset on flush. + self._series_finished = False # Deferred flush: flush() sets this flag and the actual cache clear # happens at the top of the next compute(), so it never mutates the # caches while compute() is mid-synchronise. This avoids the boolean- @@ -499,6 +502,7 @@ def _perform_flush(self): self.positions = np.zeros((0, 4)) self.position_ids = np.zeros((0,), dtype=int) self.count = 0 + self._series_finished = False self.logger.info( "[FLUSH VERIFY] GatherOp cleared: images=None, image_ids=0, " "positions=(0, 4), position_ids=0, count=0" @@ -535,12 +539,24 @@ def compute(self, op_input, op_output, context): self.positions = np.concatenate([self.positions, positions]) self.position_ids = np.concatenate([self.position_ids, position_ids]) + # R-1 (PR3): once the series has ended, drain the remaining matched IDs + # even if fewer than batch_size, so a final partial batch (no_frames not a + # multiple of batch_size) reaches the accumulator instead of stranding and + # idling the pipeline forever. series_finished flows from the image source + # via metadata on the final batch; latch it so the drain persists. + try: + if self.metadata is not None and self.metadata.get("series_finished", False): + self._series_finished = True + except Exception: + pass + # Find common IDs between images and positions if self.images is not None and self.image_ids.size > 0 and self.position_ids.size > 0: common_ids = np.intersect1d(self.image_ids, self.position_ids).astype(int) - + + drain = self._series_finished and int(common_ids.size) > 0 #if common_ids.size > 0: - if int(common_ids.size) >= self.batch_size: + if int(common_ids.size) >= self.batch_size or drain: # Create vectorized masks for efficient filtering mask_positions = np.isin(self.position_ids, common_ids) mask_images = np.isin(self.image_ids, common_ids) diff --git a/pipeline/header_io.py b/pipeline/header_io.py index faf9d66..c51109c 100644 --- a/pipeline/header_io.py +++ b/pipeline/header_io.py @@ -131,10 +131,13 @@ def compute(self, op_input, op_output, context): ) return - # 1. Update the always-present shared holder (S11). + # 1. Update the always-present shared holder (S11). Set no_frames here too + # (not only via configure_scan_geometry) so the STXM sink can segment + # per projection even when ptychography is disabled. if self.scan_state is not None: self.scan_state["num_projections"] = num_projections self.scan_state["current_projection"] = 0 + self.scan_state["no_frames"] = npoints_h * npoints_v # 2. Ptycho path: stage geometry + request preemption (R-4). The recon op # applies configure_scan_geometry once it has quiesced, so no buffer diff --git a/pipeline/pipeline.py b/pipeline/pipeline.py index c756ecb..1bf3a41 100644 --- a/pipeline/pipeline.py +++ b/pipeline/pipeline.py @@ -139,11 +139,14 @@ def compose(self): sink_and_publish_op = SinkAndPublishOp(self, tensor2subject=tensor2subject, publish_backend=publish_backend, + scan_state=self.scan_state, **self.kwargs('sink_and_publish_op'), name="sink_and_publish_op") # ===== Control Operator ===== flushable_ops = [gather_op, position_src, sink_and_publish_op] + ptycho_accum = None # set below when ptychography is enabled + ptycho_recon = None # ===== Ptychography Branch (conditional) ===== if self.ptychography_enabled: @@ -170,6 +173,7 @@ def compose(self): housekeeping_interval=ptycho_cfg["housekeeping_interval"], publish_interval=ptycho_cfg["publish_interval"], reset_probe=ptycho_cfg.get("reset_probe", False), + publish_folder=sink_config.get("publish_folder"), name="ptycho_reconstruction", ) @@ -185,6 +189,9 @@ def compose(self): control_op = ControlOp(self, flushable_ops=flushable_ops, publish_backend=publish_backend, + ptycho_accum=ptycho_accum, + ptycho_recon=ptycho_recon, + scan_state=self.scan_state, name="control_op") # ===== Header Source (live scan geometry, optional) ===== diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 04521e2..76c80c1 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -59,6 +59,13 @@ def __init__(self, fragment, *args, ptycho_state, **kwargs): # buffers while compute() is mid-write — race-safe even for a mid-stream # flush (PR2 header preemption). self._flush_requested = False + # PR3 tomography: frames that straddle a projection boundary are split — + # the head fills the current projection, the tail is carried here and + # written into the next projection once the boundary advance (flush) has + # reset filled_until to 0. None when there is no pending carry. + self._carry = None + # Deferred per-projection advance (PR3), distinct from a full flush. + self._advance_requested = False self.logger = logging.getLogger(kwargs.get("name", "PtychoAccumulatorOp")) super().__init__(fragment, *args, **kwargs) @@ -68,15 +75,23 @@ def setup(self, spec: OperatorSpec): ).condition(ConditionType.NONE) def flush(self): - """Request a reset; performed at the top of the next compute() (deferred, - so it never races the buffer writes in compute).""" + """Request a FULL reset (scan start/end/preempt); performed at the top of + the next compute() (deferred, so it never races the buffer writes).""" self._flush_requested = True + def advance_projection(self): + """Request a per-projection advance (PR3): reset the fill level for the + next projection but PRESERVE the straddle carry and scan centre. Deferred + to the top of the next compute().""" + self._advance_requested = True + def _perform_flush(self): """Reset fill level and zero the GPU buffers. Only ever called from compute(), so it is single-threaded w.r.t. the buffer writes. No-ops when nothing has been accumulated since the last flush (free redundant flush).""" self._flush_requested = False + # A full flush is a scan boundary — any carried straddle-tail is stale. + self._carry = None if not self._dirty: return with self.lock: @@ -91,10 +106,42 @@ def _perform_flush(self): self._dirty = False self.logger.info("Flushed ptychography accumulator buffers") + def _perform_advance(self): + """Per-projection advance (PR3): reset filled_until to 0 for the next + projection while KEEPING the carry (written next compute) and the scan + centre (all projections share the same physical scan region). Buffers are + not zeroed — the next projection overwrites [0:no_frames] as it fills.""" + self._advance_requested = False + with self.lock: + self.ptycho_state["filled_until"] = 0 + self._dirty = False + self.logger.info("Accumulator advanced to next projection (carry preserved)") + def compute(self, op_input, op_output, context): # Perform any requested flush here — single-threaded w.r.t. the buffers. if self._flush_requested: self._perform_flush() + # Per-projection advance (PR3) — after a full flush takes precedence. + if self._advance_requested: + self._perform_advance() + + scan_state = self.ptycho_state.get("scan_state") or {} + num_projections = int(scan_state.get("num_projections", 1)) + no_frames = self.ptycho_state["no_frames"] + + # PR3: write a carried straddle-tail into the freshly-advanced projection + # first (the boundary flush has reset filled_until to 0). + if self._carry is not None and self.ptycho_state["filled_until"] == 0: + carry = self._carry + self._carry = None + self._accumulate(carry["images"], carry["positions"]) + + # PR3 backpressure: if this projection is full and awaiting the boundary + # advance (recon finalize + scoped flush), stop consuming input so the + # next projection's frames queue upstream instead of being dropped. + if num_projections > 1 and self.ptycho_state["filled_until"] >= no_frames: + return + data = op_input.receive("input") if data is None: return @@ -104,9 +151,32 @@ def compute(self, op_input, op_output, context): batch_size = images.shape[0] filled = self.ptycho_state["filled_until"] - if filled + batch_size > self.ptycho_state["no_frames"]: - return # buffer full, drop batch + if filled + batch_size <= no_frames: + # Batch fits within the current projection. + self._accumulate(images, positions) + elif num_projections > 1: + # PR3: batch straddles the projection boundary — write the head to + # fill this projection, carry the tail to the next one. + head = no_frames - filled + self._accumulate(images[:head], positions[:head]) + self._carry = {"images": images[head:], "positions": positions[head:]} + self.logger.info( + "Projection boundary: wrote %d frames, carrying %d to next projection", + head, batch_size - head, + ) + else: + return # single projection, buffer full → drop batch + def _accumulate(self, images, positions): + """Preprocess a batch and write it into the pre-allocated buffers. + + Extracted so the projection-boundary dispatch in compute() can call it + for a whole batch, a straddle-head, or a carried straddle-tail. + """ + batch_size = images.shape[0] + if batch_size == 0: + return + filled = self.ptycho_state["filled_until"] pty_data = self.ptycho_state["pty_data"] # H2D + crop @@ -304,10 +374,13 @@ def __init__( housekeeping_interval=10, publish_interval=5, reset_probe=False, + publish_folder=None, **kwargs, ): self.ptycho_state = ptycho_state self.lock = ptycho_state["lock"] + # PR3: folder for per-projection reconstruction HDF5 files (tomography). + self.publish_folder = publish_folder self.total_iterations = int(total_iterations) self.post_stream_iterations = int(post_stream_iterations) self.housekeeping_interval = int(housekeeping_interval) @@ -340,10 +413,30 @@ def setup(self, spec: OperatorSpec): spec.output("complete").condition(ConditionType.NONE) def flush(self): - """Request a reset; performed at the top of the next compute() (deferred, - so the object/counter reset never races the PIE update in compute).""" + """Request a FULL reset (scan end/header/start); performed at the top of + the next compute() (deferred, so it never races the PIE update).""" self._flush_requested = True + def _perform_advance(self): + """Per-projection reset (PR3): fresh object + iteration counters for the + next projection, probe carried over (warm start). Does NOT touch + filled_until — the accumulator owns the fill level across the boundary. + Driven by the recon itself once it observes filled_until drop below + no_frames (i.e. the accumulator has advanced), so it never re-completes + the same projection.""" + self.current_iteration = 0 + self.all_data_arrived = False + self.post_stream_count = 0 + self._completed = False + if self.initialized_gpu: + pty_model = self.ptycho_state["pty_model"] + pty_model.obj.array_global[:] = self._obj_initial + pty_model.obj.array_global_old[:] = self._obj_initial + if self.reset_probe: + pty_model.probe.array_states[:] = self._probe_initial + pty_model.source.flux = self._flux_initial + self.logger.info("Recon advanced to next projection (object reset, probe carried)") + def _perform_flush(self): """Reset reconstruction state for a new scan. Only ever called from compute(), so the object reset is single-threaded w.r.t. the PIE update. @@ -500,6 +593,22 @@ def compute(self, op_input, op_output, context): if self.ptycho_state.get("needs_gpu_reinit"): self._reset_for_new_geometry() + scan_state = self.ptycho_state.get("scan_state") or {} + num_projections = int(scan_state.get("num_projections", 1)) + no_frames = self.ptycho_state["no_frames"] + + # PR3 tomography: after completing a projection, idle until the accumulator + # has advanced (filled_until dropped below no_frames). Self-advancing on + # that observation — rather than being pushed an advance — guarantees we + # never re-complete the same projection while its buffer is still full. + if self._completed and num_projections > 1: + with self.lock: + n_now = self.ptycho_state["filled_until"] + if n_now < no_frames: + self._perform_advance() # fresh object/counters for the next projection + else: + return # still holding the finished projection + # Snapshot fill level with self.lock: n_filled = self.ptycho_state["filled_until"] @@ -507,11 +616,10 @@ def compute(self, op_input, op_output, context): if n_filled == 0: return - # ITER_TIMING instrumentation (diagnostic, uncommitted) + # ITER_TIMING instrumentation (per-iteration timing diagnostic; INFO level, + # ~sub-ms against a ~450ms PIE iteration) t_start = time.perf_counter() - no_frames = self.ptycho_state["no_frames"] - # Detect when all data has arrived if n_filled >= no_frames and not self.all_data_arrived: self.all_data_arrived = True @@ -639,11 +747,17 @@ def compute(self, op_input, op_output, context): pty_model.scan.positions = self.ptycho_state["positions_full"] pty_model.scan.tilts = self.ptycho_state["tilts_full"] - # Housekeeping (every N iterations or on last) - is_last = self.current_iteration >= self.total_iterations - 1 and ( - not self.all_data_arrived - or self.post_stream_count >= self.post_stream_iterations - ) + # Housekeeping (every N iterations or on last). For tomography + # (num_projections > 1) a projection completes as soon as all its frames + # are in — finish this in-flight iteration, then advance (plan PR3). For a + # single projection, run to total_iterations + post_stream (as before). + if num_projections > 1: + is_last = self.all_data_arrived + else: + is_last = self.current_iteration >= self.total_iterations - 1 and ( + not self.all_data_arrived + or self.post_stream_count >= self.post_stream_iterations + ) if ( self.current_iteration % self.housekeeping_interval == 0 or is_last @@ -675,10 +789,26 @@ def compute(self, op_input, op_output, context): # flushes on this signal (Task 3: flush after the last iteration). if is_last and self.all_data_arrived and not self._completed: self._completed = True - op_output.emit("recon_complete", "complete") - self.logger.info( - "Reconstruction complete at iteration %d", self.current_iteration - ) + if num_projections > 1: + # Tomography: save this projection, then advance (non-final) or + # end the scan (final projection). ControlOp does the scoped + # advance on "projection_complete" (accum+recon only, not gather). + self._save_projection_file() + current_proj = int(scan_state.get("current_projection", 0)) + if current_proj >= num_projections - 1: + signal = "recon_complete" # last projection → full scan end + else: + signal = "projection_complete" + op_output.emit(signal, "complete") + self.logger.info( + "Projection %d/%d complete at iteration %d (signal=%s)", + current_proj, num_projections, self.current_iteration, signal, + ) + else: + op_output.emit("recon_complete", "complete") + self.logger.info( + "Reconstruction complete at iteration %d", self.current_iteration + ) t_end = time.perf_counter() self.logger.info( @@ -704,6 +834,36 @@ def compute(self, op_input, op_output, context): # ------------------------------------------------------------------ + def _save_projection_file(self): + """PR3: write this projection's reconstruction to its own HDF5 file, + named with the shared series_id + projection index (M2 — all projections + of a tomography scan share one series_id, so the index is mandatory).""" + if self.publish_folder is None: + return + scan_state = self.ptycho_state.get("scan_state") or {} + series_id = scan_state.get("series_id", "unknown") + proj = int(scan_state.get("current_projection", 0)) + pty_model = self.ptycho_state["pty_model"] + obj_2d = np.squeeze(cp.asnumpy(pty_model.obj.array_global)) + probe_2d = np.squeeze(cp.asnumpy(pty_model.probe.array_states)) + import h5py + try: + os.makedirs(self.publish_folder, exist_ok=True) + path = os.path.join( + self.publish_folder, f"{series_id}_proj{proj:02d}_recon.h5" + ) + with h5py.File(path, "w") as f: + f.create_dataset("object_phase", data=np.angle(obj_2d).astype(np.float32)) + f.create_dataset("object_amp", data=np.abs(obj_2d).astype(np.float32)) + f.create_dataset("probe_phase", data=np.angle(probe_2d).astype(np.float32)) + f.create_dataset("probe_amp", data=np.abs(probe_2d).astype(np.float32)) + f.attrs["projection"] = proj + f.attrs["series_id"] = str(series_id) + f.attrs["iteration"] = int(self.current_iteration) + self.logger.info("Saved projection recon → %s", path) + except Exception: + self.logger.exception("Failed to save projection recon file") + def _init_gpu(self): """One-time transfer of static model data to GPU.""" pty_data = self.ptycho_state["pty_data"] diff --git a/pipeline/publish.py b/pipeline/publish.py index 44402ac..73b1f47 100644 --- a/pipeline/publish.py +++ b/pipeline/publish.py @@ -80,6 +80,7 @@ def __init__(self, fragment, *args, publish_backend: PublishBackend = None, backend: str = "nats", backend_endpoint: str = None, + scan_state: dict = None, **kwargs): """ Initialize sink and publish operator. @@ -100,13 +101,19 @@ def __init__(self, fragment, *args, self.processed_frame_count = 0 self.processed_batch_count = 0 - # In-memory accumulator of per-batch arrays for the current scan. + # In-memory accumulator of per-batch arrays for the current scan/projection. # Avoids per-batch HDF5 open/close on the compute hot path; the file - # is written once at scan end (processing_end). + # is written once at scan end (or per projection for tomography). self.scan_buffer = [] # Save-state tracking so flush never discards an unwritten scan buffer. self._written = False self._series_id = None + # PR3 tomography: shared holder (num_projections, no_frames) + a LOCAL + # per-projection index/counter so the STXM path segments itself (S2), + # independent of the ptycho path's current_projection. + self.scan_state = scan_state + self._projection = 0 + self._proj_frame_count = 0 self.publish_folder = publish_folder self.publish_tensors = publish_tensors if publish_tensors is not None else [] @@ -142,6 +149,23 @@ def write_scan_file(self, series_id): self._written = True self.scan_buffer = [] + def _write_projection_file(self, series_id, projection): + """PR3: write the current projection's buffered STXM to its own file, + named with the shared series_id + projection index (M2), then clear the + buffer for the next projection.""" + if self.publish_folder is None or not self.scan_buffer: + return + os.makedirs(self.publish_folder, exist_ok=True) + data = np.concatenate(self.scan_buffer, axis=0) + filepath = os.path.join(self.publish_folder, f"{series_id}_proj{projection:02d}.h5") + with h5py.File(filepath, 'w') as f: + f.create_dataset('stxm', data=data) + f.attrs['projection'] = projection + f.attrs['series_id'] = str(series_id) + self.logger.info(f"Wrote projection {projection} ({data.shape[0]} frames) to {filepath}") + self._written = True + self.scan_buffer = [] + def flush(self): """Reset counters. If the buffer still holds data that was never written (a flush arrived before the end-of-scan write), write it out first so a @@ -152,11 +176,18 @@ def flush(self): "Flush with %d unwritten STXM batch(es) — writing before clearing", len(self.scan_buffer), ) - self.write_scan_file(self._series_id) + # Name with the projection index if we're mid-tomography. + scan_state = self.scan_state or {} + if int(scan_state.get("num_projections", 1)) > 1: + self._write_projection_file(self._series_id, self._projection) + else: + self.write_scan_file(self._series_id) self.processed_frame_count = 0 self.processed_batch_count = 0 self.scan_buffer = [] self._written = False + self._projection = 0 + self._proj_frame_count = 0 def compute(self, op_input, op_output, context): """Receive, publish, and save processed data using metadata.""" @@ -222,19 +253,39 @@ def compute(self, op_input, op_output, context): self.processed_batch_count += 1 self.processed_frame_count += tensor.shape[0] - # Check if processing is complete using metadata from upstream. Use >= - # (not ==) so a batch that overshoots the exact count (e.g. frame count - # not a multiple of batch_size) still triggers the end-of-scan save. - if series_frame_count > 0 and self.processed_frame_count >= series_frame_count: - # Write the whole scan once, now that no more batches are arriving - if self.publish_folder is not None and series_id is not None: - self.write_scan_file(series_id) - - _n = self.processed_frame_count - _b = self.processed_batch_count - _elapsed = time.time() - series_start_time if series_start_time > 0 else 0 - _rate = _n/_elapsed if _elapsed > 0 else 0 - self.logger.info(f"{_n} processed in {_elapsed:.1f}s. speed: {_rate:.1f} Hz (in {_b} batches)") + # Share series_id so the ptycho per-projection recon files can be named + # with the same identifier (PR3). + if self.scan_state is not None and series_id is not None: + self.scan_state["series_id"] = series_id + + scan_state = self.scan_state or {} + num_projections = int(scan_state.get("num_projections", 1)) + proj_no_frames = int(scan_state.get("no_frames", 0)) + + if num_projections > 1 and proj_no_frames > 0: + # Tomography (S2): save one STXM file per projection, segmented by + # frame count — self-contained, no ControlOp round-trip. Uses >= with + # a count carry (I1); exact frame boundaries require no_frames to be a + # multiple of batch_size (a few overshoot frames otherwise land in the + # current projection's file). + self._proj_frame_count += tensor.shape[0] + if self._proj_frame_count >= proj_no_frames: + if self.publish_folder is not None and series_id is not None: + self._write_projection_file(series_id, self._projection) + self._proj_frame_count -= proj_no_frames # carry overshoot count + self._projection += 1 + else: + # Single scan: write the whole series once (existing behaviour). Use + # >= (not ==) so a batch overshooting the exact count still triggers. + if series_frame_count > 0 and self.processed_frame_count >= series_frame_count: + if self.publish_folder is not None and series_id is not None: + self.write_scan_file(series_id) + + _n = self.processed_frame_count + _b = self.processed_batch_count + _elapsed = time.time() - series_start_time if series_start_time > 0 else 0 + _rate = _n/_elapsed if _elapsed > 0 else 0 + self.logger.info(f"{_n} processed in {_elapsed:.1f}s. speed: {_rate:.1f} Hz (in {_b} batches)") class PublishToCloudOp(Operator): From 761fd35478159221ba89d3c126ece8c89a394766 Mon Sep 17 00:00:00 2001 From: ramyaguru Date: Thu, 2 Jul 2026 17:19:12 -0400 Subject: [PATCH 06/33] =?UTF-8?q?PR4:=20Double-buffering=20=E2=80=94=20pro?= =?UTF-8?q?jection=20ping-pong=20to=20decouple=20accumulation=20from=20fin?= =?UTF-8?q?alize?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Motivated by moving to a much faster detector: single-buffer backpressures the source across the ~one-iteration finalize window (measured ~490 ms), which throttles a free-running detector and risks dropping positions (PUB/SUB, no flow control). Double-buffering keeps the accumulator draining while the recon finalizes. Two GPU buffer sets (ping-pong). While the recon finalizes projection N on the read buffer, the accumulator fills N+1 into the write buffer. On fill it flips to the free buffer (else clean backpressure fallback when the recon is lapped by a full projection — data is never clobbered). Single-projection scans always use buffer 0 (no flip). - ptychography_setup.py: allocate raw_gpu/positions_full/tilts_full as 2-element lists (once, at max capacity — R-6); add ping-pong state (filled_until[], write_idx, read_idx, buf_free[], num_buffers) reset in configure_scan_geometry. - PtychoAccumulatorOp: _try_flip() on fill (else backpressure + a one-shot "Recon lagging" warning); all writes indexed by write_idx; _perform_flush zeros both buffers + resets the ping-pong idempotently. - PtychoReconstructionOp: reads indexed by read_idx; _flip_read() releases the finished buffer, OWNS current_projection (no ControlOp round-trip → no save-vs-index race), advances to the next buffer, resets the object (probe carried); final projection emits recon_complete and idles. - control.py: dropped the now-dead projection_complete branch; ControlOp keeps recon_complete / header / flush. Container-validated on the A400 (see DESIGN_pr4_double_buffer.md): regression (exact 1024/projection, same iteration budget), overlap (accumulator fills the next projection during a stretched finalize without backpressure), and lapping fallback (clean backpressure + recovery when the recon is lapped). No drops, no deadlock. Open item O1 (per-projection post-stream iterations) is now more visible — pre-filled projections finalize with few refinement iterations. Assisted-By: Claude Opus 4.8 (1M context) --- pipeline/DESIGN_pr4_double_buffer.md | 173 +++++++++++++++++++ pipeline/control.py | 22 +-- pipeline/ptychography_ops.py | 244 ++++++++++++++++----------- pipeline/ptychography_setup.py | 57 +++++-- 4 files changed, 371 insertions(+), 125 deletions(-) create mode 100644 pipeline/DESIGN_pr4_double_buffer.md diff --git a/pipeline/DESIGN_pr4_double_buffer.md b/pipeline/DESIGN_pr4_double_buffer.md new file mode 100644 index 0000000..451f5fd --- /dev/null +++ b/pipeline/DESIGN_pr4_double_buffer.md @@ -0,0 +1,173 @@ +# Design — PR4 double-buffering (projection ping-pong) + +Branch `feat/tomo-pr4-loadtest`. Motivated by moving to a **much faster detector**; +the single-buffer design measured sufficient at 3 kHz on the sim but with shrinking +margin as frame rate / frame size grow. + +## Why (measured, 2026-07-02) + +Load test on daqsim (scan409907, 2 projections, A400). Instrumentation tagged +`PR4-MEASURE` (`data_io.py` GatherOp cache HWM; `ptychography_ops.py` finalize-window +timing + `PR4_ACCUM_QUEUE_CAP` / `PR4_FINALIZE_DELAY_MS` knobs). + +- **Run A (baseline, 500 fps, queue=128):** finalize window **~490 ms** (≈ one PIE + iteration + HDF save). GatherOp cache HWM **256 frames**. Both projections exact. +- **Run B (queue=8, ~2 s injected finalize):** GatherOp cache HWM **336 frames**, + both projections exact, **no drops, no deadlock abort** at 4× the 500 ms timeout. + +**Conclusions:** +1. The current single-buffer design degrades **gracefully** — backpressure is lossless, + and the `stop_on_deadlock` (500 ms) tripwire does **not** fire during a long finalize + because the recon thread is busy (counts as progress). So PR4 does **not** need to + touch the deadlock timeout. +2. Zero loss on the sim is a **sim artifact**: daqsim streams images (PUSH) and positions + from one loop, so pipeline backpressure on the PUSH socket stalls the whole sim, + incidentally pausing positions. **A real detector + PandA are independent free-running + sources** — backpressure during the finalize window can't throttle them, so frames/ + positions land in a full socket and drop (positions especially: PUB/SUB, no flow + control, no CONFLATE, ~1000 RCVHWM). +3. So the single-buffer failure mode against a real fast detector is **source throttling / + silent position loss during the ~1-iteration finalize window**, not a crash. Backlog + ≈ `acq_rate × finalize_window`; the finalize window grows with object/frame size. + +Double-buffering removes the finalize-window backpressure: the accumulator never stops +draining, so a free-running detector is never throttled and positions never back up. + +## How single-buffer works today (baseline) + +One shared GPU buffer set in `ptycho_state`: `raw_gpu (capacity,H,W)`, `positions_full`, +`tilts_full`, one `filled_until` counter. The PtyREX model (`pty_model.scan.positions`, +`pty_data.raw_expanded`) holds **views** into these buffers. + +- `PtychoAccumulatorOp` writes batches into `raw_gpu[filled:new_end]`, advances + `filled_until`. On a projection boundary it splits the straddling batch (head fills the + projection, tail → `self._carry`). +- Once `filled_until >= no_frames` (tomography), the accumulator **backpressures**: + `compute` returns before `receive()` (`ptychography_ops.py:147`). Incoming batches wait + in the `gather → accum` `DOUBLE_BUFFER` queue (capacity 128 ≈ 8k frames). +- `PtychoReconstructionOp` runs the final iteration on the full buffer, saves the + per-projection HDF (~490 ms), emits `projection_complete`. +- `ControlOp` → `accum.advance_projection()` → next accum tick resets `filled_until=0`, + writes the carry, resumes draining. Recon observes `filled < no_frames`, resets object + (probe carried) for the next projection. + +The gap: for the whole finalize window the accumulator is **not draining**, so the source +is backpressured. That is what double-buffering eliminates. + +## Double-buffer design (2-buffer ping-pong) + +### State (`ptycho_state`, allocated once at `capacity` per R-6 — now ×2) +``` +raw_gpu: [bufA, bufB] # two (capacity,H,W) arrays +positions_full: [posA, posB] +tilts_full: [tiltA, tiltB] +filled_until: [fillA, fillB] # per-buffer fill level (under lock) +write_idx: 0 # buffer the accumulator writes (accum owns) +read_idx: 0 # buffer the recon reads (recon owns) +buffer_ready: [Event, Event] # buffer i full & handed to recon +buffer_free: [Event, Event] # buffer i free for the accumulator (both set at init) +``` +GPU cost: 2× `raw_gpu` (~64 MB each at 1024×128×128×f32 → ~128 MB total). Fine on the +A400; scales 2× with frame size — watch on the faster detector. + +### Accumulator (`PtychoAccumulatorOp`) +- Write into `raw_gpu[write_idx]` until `filled_until[write_idx] == no_frames`. +- **On full (tomography): flip instead of backpressure.** + 1. `buffer_ready[write_idx].set()` (hand this projection to the recon). + 2. If `buffer_free[1-write_idx]` is set → flip `write_idx`, `filled_until[write_idx]=0`, + clear `buffer_free[write_idx]`, write the carried straddle-tail, keep draining. + 3. Else (recon still on the other buffer — we've lapped it) → **backpressure `return` + as today.** This is the graceful fallback for sustained slowness: double-buffer buys + exactly one projection of runway, not infinite throughput. +- Single projection (`num_projections == 1`): unchanged — always buffer 0, no flip. + +### Reconstruction (`PtychoReconstructionOp`) +- Read views from `raw_gpu[read_idx]` / `positions_full[read_idx]` (index the existing + view re-point that already runs each `compute`). +- Run iterations; when `filled_until[read_idx] == no_frames` finalize + save the + projection HDF (as today). +- **After finalizing:** `buffer_free[read_idx].set()`, clear `buffer_ready[read_idx]`, + flip `read_idx`, reset object (probe carried) for the next projection. +- Per-projection object reset is the existing advance logic, keyed off the flip. + +### Coordination +`write_idx` is touched only by the accumulator, `read_idx` only by the recon; they observe +each other through `buffer_ready` / `buffer_free` events and per-buffer `filled_until` +(under the existing `ptycho_state["lock"]`). Ping-pong invariant: the accumulator never +writes a buffer whose `buffer_free` is clear (recon still reading it) — enforced by step +2/3 above. + +### Flush / preempt / advance +- **Full flush (scan end / R-2 start safety / PR2 header preempt):** reset BOTH buffers + (`filled_until=[0,0]`, zero both `raw_gpu`), `write_idx=read_idx=0`, `buffer_free` both + set, `buffer_ready` both clear, drop the carry. Fold into `_perform_flush` + + `_reset_for_new_geometry`. +- **Per-projection advance** is now implicit in the flip — the explicit + `advance_projection()` / `projection_complete → control → advance` round-trip can be + simplified or kept as the buffer-free signal. Decision point below. + +## Open design decisions (resolve before implementing) +- **D1 — keep or retire the `projection_complete → ControlOp → advance` round-trip?** + With the flip owning the advance, ControlOp's role shrinks to flush-on-final only. Keep + it for the final-projection flush; the per-projection advance becomes buffer-local. +- **D2 — 2 buffers vs an N-ring.** Start with 2 (one projection of runway, matches the + plan). Parameterize `num_buffers` if a load test later shows one projection isn't enough + headroom for the faster detector. +- **D3 — event objects vs simple int flags under the lock.** Events are clean but add + threading objects to `ptycho_state`; two ints (`ready_mask`, `free_mask`) under the + existing lock may be simpler and match the existing deferred-flag style. +- **D4 — does the STXM path need the same treatment?** `SinkAndPublishOp` saves per + projection but doesn't hold a GPU buffer across a long finalize; it likely doesn't need + double-buffering, but confirm it isn't coupled to the ptycho backpressure via the shared + `gather` output. + +## Testing plan +- **Regression:** the PR3 2-projection tomo test still produces exact 1024-frame + projections + all 4 files (`{series}_proj00/01` × STXM + recon). +- **Ping-pong proof:** instrument the flip (`write_idx`/`read_idx` transitions) and confirm + the accumulator keeps draining (no backpressure `return`) through a boundary while the + recon finalizes — i.e. GatherOp cache HWM stays flat and the accum input queue does not + back up during the finalize window. +- **Lapping fallback:** inject a long finalize (`PR4_FINALIZE_DELAY_MS`) longer than one + projection's accumulation so the recon is lapped; confirm the accumulator falls back to + clean backpressure (no data loss) rather than overwriting a buffer the recon is reading. +- **Preempt + flush:** header preemption mid-tomo resets both buffers correctly. + +## Resolved decisions (as implemented) +- **D1** — kept ControlOp for `recon_complete`/`header`/`flush` only. The per-projection + `projection_complete` round-trip is **gone**: the accumulator self-flips write buffers, + and the **recon owns `current_projection`** + the read-buffer flip (`_flip_read`), which + removes the save-vs-bump filename race. `current_projection` still resets to 0 on a new + header (`header_io.py`). +- **D2** — exactly 2 buffers (`NUM_BUFFERS = 2`), one projection of runway. `num_buffers` + is parameterised so an N-ring is a one-constant change if the faster detector needs more. +- **D3** — int-under-lock coordination (`write_idx`/`read_idx`/`filled_until[]`/`buf_free[]` + under `ptycho_state["lock"]`); no new Event objects. +- **D4** — STXM path untouched (own `_projection` counter, no iterative recon, no buffer + held across a finalize). + +## Validation results (container, A400, 2026-07-02) +All three design tests pass (`test_data/run_pr4c/d/e.log`): +- **Test 1 — regression** (2 proj, no stress): both projections reconstruct to exactly + 1024, all 4 files, same iteration budget as single-buffer (proj0=7, proj1=5). Ping-pong + logs confirm the flips (`Accumulator flipped to buffer 1` → `flipped read buffer to 1`). +- **Test 2 — overlap** (2 proj, 1.5 s injected finalize): the accumulator flipped to buf1 + and **fully accumulated proj-1 during proj-0's 2 s finalize** — never backpressured, + GatherOp cache flat. This is the win vs single-buffer Run B (which stalled and relied on + the sim's PUSH throttle). +- **Test 3 — lapping fallback** (3 proj, 3 s injected finalize): when the recon was lapped + by a full projection, the accumulator logged `Recon lagging … backpressuring upstream` + and **recovered** once the recon released the buffer — clean backpressure, no clobbering, + all 3 projections at 1024, no drops/deadlock. + +**O1 is now more urgent (observed):** when the recon lags acquisition, later projections +arrive **pre-filled** and finalize with near-zero refinement iterations (Test 2 proj-1 got +1 iteration vs 5 when filled concurrently). Double-buffering makes this visible; the fix is +per-projection post-stream iterations (deferred O1). On the faster detector this will matter. + +## Note +All `PR4-MEASURE` instrumentation + the sim tweaks (position mapping, `PTYCHO_CENTER`, +`dummy_img_index=True`) were working-tree-only test scaffolding, reverted before this commit +(the one kept production log is the `Recon lagging …` backpressure warning). The +`num_buffers`/ping-pong state lives in `ptycho_state`; `config_sim.yaml` + `Dockerfile_bwell` +remain untracked local test scaffolding. diff --git a/pipeline/control.py b/pipeline/control.py index db7f61b..e20a024 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -73,21 +73,13 @@ def compute(self, op_input, op_output, context): self._do_flush() self._flushed = True - elif msg == "projection_complete": - # PR3 tomography per-projection boundary: SCOPED advance. Reset the - # accumulator's fill level (carry preserved) and bump current_projection. - # The recon self-advances once it observes filled_until drop (so it can't - # re-complete the same projection). Do NOT flush GatherOp — its cached - # next-projection frames must survive (decision #11). The STXM sink - # segments itself by frame count (S2), so it isn't touched here. - if self.ptycho_accum is not None: - self.ptycho_accum.advance_projection() - if self.scan_state is not None: - self.scan_state["current_projection"] += 1 - self.logger.info( - "Projection complete — advanced to projection %d", - self.scan_state["current_projection"], - ) + # PR4: tomography projection boundaries no longer round-trip through + # ControlOp. With double-buffering the accumulator flips write buffers + # itself and the recon owns current_projection + the read-buffer flip + # (avoiding a save-vs-bump race), so there is no "projection_complete" + # signal any more — the recon emits "recon_complete" only on the FINAL + # projection, handled above. ControlOp is kept for recon_complete / header + # / flush. elif msg == "header": # A live header reconfigures the scan for a new dataset. Flush so the diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 76c80c1..f19bd2e 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -64,8 +64,9 @@ def __init__(self, fragment, *args, ptycho_state, **kwargs): # written into the next projection once the boundary advance (flush) has # reset filled_until to 0. None when there is no pending carry. self._carry = None - # Deferred per-projection advance (PR3), distinct from a full flush. - self._advance_requested = False + # One-shot flag so the "recon lagging" backpressure warning fires once per + # stall, not every 10 ms tick. + self._lapped = False self.logger = logging.getLogger(kwargs.get("name", "PtychoAccumulatorOp")) super().__init__(fragment, *args, **kwargs) @@ -79,68 +80,92 @@ def flush(self): the next compute() (deferred, so it never races the buffer writes).""" self._flush_requested = True - def advance_projection(self): - """Request a per-projection advance (PR3): reset the fill level for the - next projection but PRESERVE the straddle carry and scan centre. Deferred - to the top of the next compute().""" - self._advance_requested = True + def _reset_pingpong(self): + """Reset the PR4 double-buffer ping-pong to its canonical scan-start state + under the lock. Idempotent, so the accumulator's and recon's deferred + flushes can both call it in any order without disagreeing.""" + nbuf = self.ptycho_state["num_buffers"] + with self.lock: + self.ptycho_state["filled_until"] = [0] * nbuf + self.ptycho_state["write_idx"] = 0 + self.ptycho_state["read_idx"] = 0 + self.ptycho_state["buf_free"] = [i != 0 for i in range(nbuf)] def _perform_flush(self): - """Reset fill level and zero the GPU buffers. Only ever called from - compute(), so it is single-threaded w.r.t. the buffer writes. No-ops when - nothing has been accumulated since the last flush (free redundant flush).""" + """Reset fill levels + ping-pong and zero the GPU buffers. Only ever called + from compute(), so it is single-threaded w.r.t. the buffer writes. No-ops + when nothing has been accumulated since the last flush (free redundant + flush).""" self._flush_requested = False # A full flush is a scan boundary — any carried straddle-tail is stale. self._carry = None if not self._dirty: return - with self.lock: - self.ptycho_state["filled_until"] = 0 - self.ptycho_state["raw_gpu"][:] = 0 - self.ptycho_state["positions_full"][:] = 0 - self.ptycho_state["tilts_full"][:] = 0 + self._reset_pingpong() + for b in range(self.ptycho_state["num_buffers"]): + self.ptycho_state["raw_gpu"][b][:] = 0 + self.ptycho_state["positions_full"][b][:] = 0 + self.ptycho_state["tilts_full"][b][:] = 0 # Clear auto-centre so the new scan re-derives its own scan centre # from the first batch rather than reusing the previous scan's. self.ptycho_state["scan_center_py"] = None self.ptycho_state["scan_center_px"] = None self._dirty = False - self.logger.info("Flushed ptychography accumulator buffers") - - def _perform_advance(self): - """Per-projection advance (PR3): reset filled_until to 0 for the next - projection while KEEPING the carry (written next compute) and the scan - centre (all projections share the same physical scan region). Buffers are - not zeroed — the next projection overwrites [0:no_frames] as it fills.""" - self._advance_requested = False + self.logger.info("Flushed ptychography accumulator buffers (both)") + + def _try_flip(self): + """PR4: move the write cursor to the next buffer for the next projection, + iff that buffer is free (the recon has released it). Returns False when the + other buffer is still owned by the recon (we've lapped it by a full + projection) so the caller backpressures instead of overwriting live data.""" + nbuf = self.ptycho_state["num_buffers"] with self.lock: - self.ptycho_state["filled_until"] = 0 - self._dirty = False - self.logger.info("Accumulator advanced to next projection (carry preserved)") + other = (self.ptycho_state["write_idx"] + 1) % nbuf + if not self.ptycho_state["buf_free"][other]: + return False + self.ptycho_state["buf_free"][other] = False + self.ptycho_state["write_idx"] = other + self.ptycho_state["filled_until"][other] = 0 + self.logger.info("Accumulator flipped to buffer %d for next projection", other) + return True def compute(self, op_input, op_output, context): # Perform any requested flush here — single-threaded w.r.t. the buffers. if self._flush_requested: self._perform_flush() - # Per-projection advance (PR3) — after a full flush takes precedence. - if self._advance_requested: - self._perform_advance() scan_state = self.ptycho_state.get("scan_state") or {} num_projections = int(scan_state.get("num_projections", 1)) no_frames = self.ptycho_state["no_frames"] - - # PR3: write a carried straddle-tail into the freshly-advanced projection - # first (the boundary flush has reset filled_until to 0). - if self._carry is not None and self.ptycho_state["filled_until"] == 0: + w = self.ptycho_state["write_idx"] + + # PR4: the current write buffer is full. For tomography, flip to the next + # buffer so we keep draining projection N+1 while the recon finalizes N on + # the read buffer. If the other buffer isn't free yet (recon hasn't + # released it — we've lapped it by a full projection), fall back to + # backpressure so frames queue upstream rather than being dropped. Single + # projection never flips. + if num_projections > 1 and self.ptycho_state["filled_until"][w] >= no_frames: + if not self._try_flip(): + if not self._lapped: # warn once per stall, not every tick + self._lapped = True + self.logger.warning( + "Recon lagging: write buffer %d full but the recon still " + "holds the other buffer — backpressuring upstream. Both " + "double-buffer slots are in use; the detector is outrunning " + "reconstruction.", w, + ) + return # lapped → clean backpressure fallback + self._lapped = False + w = self.ptycho_state["write_idx"] + + # Write a carried straddle-tail into the (freshly-flipped, empty) write + # buffer first. + if self._carry is not None and self.ptycho_state["filled_until"][w] == 0: carry = self._carry self._carry = None self._accumulate(carry["images"], carry["positions"]) - - # PR3 backpressure: if this projection is full and awaiting the boundary - # advance (recon finalize + scoped flush), stop consuming input so the - # next projection's frames queue upstream instead of being dropped. - if num_projections > 1 and self.ptycho_state["filled_until"] >= no_frames: - return + w = self.ptycho_state["write_idx"] data = op_input.receive("input") if data is None: @@ -150,13 +175,14 @@ def compute(self, op_input, op_output, context): positions = np.asarray(data["positions"]) # (N, 4) [x, y, z, theta] batch_size = images.shape[0] - filled = self.ptycho_state["filled_until"] + filled = self.ptycho_state["filled_until"][w] if filled + batch_size <= no_frames: # Batch fits within the current projection. self._accumulate(images, positions) elif num_projections > 1: - # PR3: batch straddles the projection boundary — write the head to - # fill this projection, carry the tail to the next one. + # Batch straddles the projection boundary — write the head to fill this + # projection; carry the tail. Next compute sees the buffer full, flips, + # and writes the carry into the new buffer. head = no_frames - filled self._accumulate(images[:head], positions[:head]) self._carry = {"images": images[head:], "positions": positions[head:]} @@ -176,7 +202,8 @@ def _accumulate(self, images, positions): batch_size = images.shape[0] if batch_size == 0: return - filled = self.ptycho_state["filled_until"] + w = self.ptycho_state["write_idx"] + filled = self.ptycho_state["filled_until"][w] pty_data = self.ptycho_state["pty_data"] # H2D + crop @@ -222,11 +249,11 @@ def _accumulate(self, images, positions): positions_txyz = positions[:, [3, 0, 1, 2]] pos_y, pos_x = self._transform_positions(positions_txyz) - # Write into pre-allocated buffers + # Write into pre-allocated buffers (the current write buffer, PR4) new_end = filled + batch_size - self.ptycho_state["raw_gpu"][filled:new_end] = images_gpu - self.ptycho_state["positions_full"][0, 0, filled:new_end] = cp.asarray(pos_y) - self.ptycho_state["positions_full"][0, 1, filled:new_end] = cp.asarray(pos_x) + self.ptycho_state["raw_gpu"][w][filled:new_end] = images_gpu + self.ptycho_state["positions_full"][w][0, 0, filled:new_end] = cp.asarray(pos_y) + self.ptycho_state["positions_full"][w][0, 1, filled:new_end] = cp.asarray(pos_x) # Diagnostics on first batch if filled == 0: @@ -244,17 +271,17 @@ def _accumulate(self, images, positions): pos_y.min(), pos_y.max(), pos_x.min(), pos_x.max(), ) - # Atomically update fill counter + # Atomically update fill counter (for the current write buffer) with self.lock: - self.ptycho_state["filled_until"] = new_end + self.ptycho_state["filled_until"][w] = new_end self._dirty = True # Summary when buffer is full. Buffers are allocated at capacity (R-6), # so slice to the logical no_frames rather than the full buffer extent. if new_end >= self.ptycho_state["no_frames"]: no_frames = self.ptycho_state["no_frames"] - all_py = cp.asnumpy(self.ptycho_state["positions_full"][0, 0, :no_frames]) - all_px = cp.asnumpy(self.ptycho_state["positions_full"][0, 1, :no_frames]) + all_py = cp.asnumpy(self.ptycho_state["positions_full"][w][0, 0, :no_frames]) + all_px = cp.asnumpy(self.ptycho_state["positions_full"][w][0, 1, :no_frames]) pty_model = self.ptycho_state["pty_model"] obj_h = int(pty_model.obj.sz_glo[-2]) obj_w = int(pty_model.obj.sz_glo[-1]) @@ -418,12 +445,11 @@ def flush(self): self._flush_requested = True def _perform_advance(self): - """Per-projection reset (PR3): fresh object + iteration counters for the - next projection, probe carried over (warm start). Does NOT touch - filled_until — the accumulator owns the fill level across the boundary. - Driven by the recon itself once it observes filled_until drop below - no_frames (i.e. the accumulator has advanced), so it never re-completes - the same projection.""" + """Per-projection reset: fresh object + iteration counters for the next + projection, probe carried over (warm start). Does NOT touch the buffers — + the read buffer was just released and the recon now reads the other one. + Called from _flip_read (PR4) when the recon moves to the next projection's + buffer, so it never re-completes the projection it just finished.""" self.current_iteration = 0 self.all_data_arrived = False self.post_stream_count = 0 @@ -437,6 +463,21 @@ def _perform_advance(self): pty_model.source.flux = self._flux_initial self.logger.info("Recon advanced to next projection (object reset, probe carried)") + def _flip_read(self, scan_state): + """PR4: finished the current read buffer — release it back to the + accumulator, advance current_projection (the recon owns it, so the file + save can't race a ControlOp bump), move the read cursor to the next + buffer, and reset the object/counters for the next projection.""" + nbuf = self.ptycho_state["num_buffers"] + with self.lock: + r = self.ptycho_state["read_idx"] + self.ptycho_state["buf_free"][r] = True # accumulator may reuse it + self.ptycho_state["read_idx"] = (r + 1) % nbuf + scan_state["current_projection"] = int( + scan_state.get("current_projection", 0) + ) + 1 + self._perform_advance() # fresh object/counters (probe carried) + def _perform_flush(self): """Reset reconstruction state for a new scan. Only ever called from compute(), so the object reset is single-threaded w.r.t. the PIE update. @@ -451,12 +492,16 @@ def _perform_flush(self): self.all_data_arrived = False self.post_stream_count = 0 self._completed = False - # Clear the shared fill counter too, so this op immediately sees "no data" - # and won't re-process the just-finished scan before the accumulator's own - # (deferred) flush zeros the buffers. Both ops resetting it to 0 is - # consistent; the accumulator's flush runs before it writes new frames. + # Reset the shared ping-pong (PR4) too, so this op immediately sees "no + # data" and won't re-process the just-finished scan before the + # accumulator's own (deferred) flush zeros the buffers. Idempotent with the + # accumulator's identical reset — order-independent. + nbuf = self.ptycho_state["num_buffers"] with self.lock: - self.ptycho_state["filled_until"] = 0 + self.ptycho_state["filled_until"] = [0] * nbuf + self.ptycho_state["write_idx"] = 0 + self.ptycho_state["read_idx"] = 0 + self.ptycho_state["buf_free"] = [i != 0 for i in range(nbuf)] if self.initialized_gpu: pty_model = self.ptycho_state["pty_model"] pty_model.obj.array_global[:] = self._obj_initial @@ -597,21 +642,17 @@ def compute(self, op_input, op_output, context): num_projections = int(scan_state.get("num_projections", 1)) no_frames = self.ptycho_state["no_frames"] - # PR3 tomography: after completing a projection, idle until the accumulator - # has advanced (filled_until dropped below no_frames). Self-advancing on - # that observation — rather than being pushed an advance — guarantees we - # never re-complete the same projection while its buffer is still full. - if self._completed and num_projections > 1: - with self.lock: - n_now = self.ptycho_state["filled_until"] - if n_now < no_frames: - self._perform_advance() # fresh object/counters for the next projection - else: - return # still holding the finished projection + # PR4: after the FINAL projection completes we idle until the scan-end + # flush (ControlOp flushes on recon_complete). Non-final projections reset + # _completed in _flip_read on the same tick, so this only catches the + # final-idle case — the recon never re-processes a finished projection. + if self._completed: + return - # Snapshot fill level + # Snapshot fill level of the buffer we're reading (PR4 double-buffer). with self.lock: - n_filled = self.ptycho_state["filled_until"] + r = self.ptycho_state["read_idx"] + n_filled = self.ptycho_state["filled_until"][r] if n_filled == 0: return @@ -626,7 +667,7 @@ def compute(self, op_input, op_output, context): self.post_stream_count = 0 pty_model = self.ptycho_state["pty_model"] pty_model.scan.original = cp.copy( - self.ptycho_state["positions_full"] + self.ptycho_state["positions_full"][r] ) self.logger.info("All %d frames arrived", no_frames) @@ -650,19 +691,19 @@ def compute(self, op_input, op_output, context): pty_model = self.ptycho_state["pty_model"] pty_params = self.ptycho_state["pty_params"] - pty_model.scan.positions = self.ptycho_state["positions_full"][ + pty_model.scan.positions = self.ptycho_state["positions_full"][r][ :, :, :n_filled ] - pty_model.scan.tilts = self.ptycho_state["tilts_full"][ + pty_model.scan.tilts = self.ptycho_state["tilts_full"][r][ :, :, :n_filled ] - pty_data.raw_expanded = self.ptycho_state["raw_gpu"][:n_filled][ + pty_data.raw_expanded = self.ptycho_state["raw_gpu"][r][:n_filled][ cp.newaxis, :, :, : ] # Flux normalization — compute once on first iteration if self.current_iteration == 0 and pty_model.source.flux < 0: - raw_cpu = cp.asnumpy(self.ptycho_state["raw_gpu"][:n_filled]) + raw_cpu = cp.asnumpy(self.ptycho_state["raw_gpu"][r][:n_filled]) dp = pty_data.dp pty_model.source.flux = float(np.sum( np.sum(raw_cpu, 0)[dp == 1] @@ -717,8 +758,8 @@ def compute(self, op_input, op_output, context): "No valid positions (0/%d in object bounds), skipping iteration", n_filled, ) - pty_model.scan.positions = self.ptycho_state["positions_full"] - pty_model.scan.tilts = self.ptycho_state["tilts_full"] + pty_model.scan.positions = self.ptycho_state["positions_full"][r] + pty_model.scan.tilts = self.ptycho_state["tilts_full"][r] return pty_params.frame_IDs = cp.asnumpy(valid_ids) @@ -743,9 +784,9 @@ def compute(self, op_input, op_output, context): combine_subsets_stream(pty_model, pty_params, recon_data) t_pie = time.perf_counter() - # Restore full buffer references for next accumulator writes - pty_model.scan.positions = self.ptycho_state["positions_full"] - pty_model.scan.tilts = self.ptycho_state["tilts_full"] + # Restore full (read-buffer) references after the sliced PIE view + pty_model.scan.positions = self.ptycho_state["positions_full"][r] + pty_model.scan.tilts = self.ptycho_state["tilts_full"][r] # Housekeeping (every N iterations or on last). For tomography # (num_projections > 1) a projection completes as soon as all its frames @@ -790,20 +831,33 @@ def compute(self, op_input, op_output, context): if is_last and self.all_data_arrived and not self._completed: self._completed = True if num_projections > 1: - # Tomography: save this projection, then advance (non-final) or - # end the scan (final projection). ControlOp does the scoped - # advance on "projection_complete" (accum+recon only, not gather). + # Tomography: save this projection using the CURRENT index, then + # either end the scan (final) or flip to the next projection's + # buffer. PR4: the recon owns current_projection and the read-buffer + # flip itself — no ControlOp round-trip — so it can never save the + # next projection under a stale index (the old race). self._save_projection_file() current_proj = int(scan_state.get("current_projection", 0)) if current_proj >= num_projections - 1: - signal = "recon_complete" # last projection → full scan end + # Final projection → full scan end. ControlOp flushes on this. + op_output.emit("recon_complete", "complete") + self.logger.info( + "Projection %d/%d complete (final) at iteration %d", + current_proj, num_projections, self.current_iteration, + ) + # _completed stays True → idle until the scan-end flush. else: - signal = "projection_complete" - op_output.emit(signal, "complete") - self.logger.info( - "Projection %d/%d complete at iteration %d (signal=%s)", - current_proj, num_projections, self.current_iteration, signal, - ) + # Release the finished read buffer to the accumulator and move + # to the next projection's buffer (which the accumulator has + # been filling meanwhile). Resets _completed → resume next tick. + iter_done = self.current_iteration # _flip_read resets it to 0 + self._flip_read(scan_state) + self.logger.info( + "Projection %d/%d complete at iteration %d — flipped read " + "buffer to %d for next projection", + current_proj, num_projections, iter_done, + self.ptycho_state["read_idx"], + ) else: op_output.emit("recon_complete", "complete") self.logger.info( diff --git a/pipeline/ptychography_setup.py b/pipeline/ptychography_setup.py index 2c9ceef..8be34f7 100644 --- a/pipeline/ptychography_setup.py +++ b/pipeline/ptychography_setup.py @@ -240,17 +240,26 @@ def configure_scan_geometry( # The buffers are allocated once at capacity in init_ptycho_state and never # realloced; the accumulator/recon index them via ptycho_state directly and # bound their reads to [:n_filled] (n_filled <= no_frames <= capacity). + # PR4: point the model at buffer 0 to start; the recon re-points these views + # to positions_full[read_idx] / raw_gpu[read_idx] each compute, so this is + # just the initial binding + the shape source for scan.original/previous. positions_full = ptycho_state["positions_full"] tilts_full = ptycho_state["tilts_full"] - pty_model.scan.positions = positions_full - pty_model.scan.tilts = tilts_full - pty_model.scan.original = cp.zeros_like(positions_full) - pty_model.scan.previous = cp.zeros_like(positions_full) + pty_model.scan.positions = positions_full[0] + pty_model.scan.tilts = tilts_full[0] + pty_model.scan.original = cp.zeros_like(positions_full[0]) + pty_model.scan.previous = cp.zeros_like(positions_full[0]) # ── Update shared state ──────────────────────────────────────────── + # A (re)configure is a clean scan boundary: reset the ping-pong to buffer 0, + # both fill levels empty, buffer 0 claimed for the first projection. with ptycho_state["lock"]: ptycho_state["no_frames"] = no_frames - ptycho_state["filled_until"] = 0 + nbuf = ptycho_state["num_buffers"] + ptycho_state["filled_until"] = [0] * nbuf + ptycho_state["write_idx"] = 0 + ptycho_state["read_idx"] = 0 + ptycho_state["buf_free"] = [i != 0 for i in range(nbuf)] ptycho_state["N"] = N # Clear auto-centre so the new geometry re-derives its own scan centre. ptycho_state["scan_center_py"] = None @@ -309,14 +318,21 @@ def init_ptycho_state(ptycho_cfg: dict, scan_state: dict = None) -> dict: default_step_v = float(ptycho_cfg["default_step_size_v"]) # ── 3. Pre-allocate GPU buffers ONCE at max capacity (R-6) ───────── - raw_gpu = cp.zeros((capacity, H, W), dtype=cp.uint32) - positions_full = cp.zeros((1, 2, capacity), dtype=cp.float32) - tilts_full = cp.zeros((1, 2, capacity), dtype=cp.float32) + # PR4 double-buffering: TWO buffer sets (ping-pong). While the recon + # finalizes projection N on the read buffer, the accumulator fills + # projection N+1 into the write buffer — so the accumulator never stops + # draining and a free-running detector is never backpressured across the + # ~one-iteration finalize window. Single-projection scans always use + # buffer 0 (no flip). Cost: 2× raw_gpu (~tens of MB). + NUM_BUFFERS = 2 + raw_gpu = [cp.zeros((capacity, H, W), dtype=cp.uint32) for _ in range(NUM_BUFFERS)] + positions_full = [cp.zeros((1, 2, capacity), dtype=cp.float32) for _ in range(NUM_BUFFERS)] + tilts_full = [cp.zeros((1, 2, capacity), dtype=cp.float32) for _ in range(NUM_BUFFERS)] logger.info( - "ptycho_state buffers allocated at capacity: %d frames, " - "image size %dx%d, %d total iterations", - capacity, H, W, pty_params.total_iterations, + "ptycho_state buffers allocated at capacity: %d frames × %d buffers " + "(double-buffered), image size %dx%d, %d total iterations", + capacity, NUM_BUFFERS, H, W, pty_params.total_iterations, ) # ── 4. Assemble ptycho_state (geometry filled in by configure) ───── @@ -324,10 +340,21 @@ def init_ptycho_state(ptycho_cfg: dict, scan_state: dict = None) -> dict: "pty_data": pty_data, "pty_model": pty_model, "pty_params": pty_params, - "raw_gpu": raw_gpu, - "positions_full": positions_full, - "tilts_full": tilts_full, - "filled_until": 0, + "raw_gpu": raw_gpu, # list[NUM_BUFFERS] of (capacity,H,W) + "positions_full": positions_full, # list[NUM_BUFFERS] + "tilts_full": tilts_full, # list[NUM_BUFFERS] + "num_buffers": NUM_BUFFERS, + # PR4 ping-pong state (all touched under "lock"): + # filled_until[i] — fill level of buffer i + # write_idx — buffer the accumulator writes (accumulator owns) + # read_idx — buffer the recon reads (recon owns) + # buf_free[i] — buffer i holds no data the recon still needs, so the + # accumulator may claim it for a new projection. Init: + # buffer 0 is claimed for the first projection. + "filled_until": [0] * NUM_BUFFERS, + "write_idx": 0, + "read_idx": 0, + "buf_free": [i != 0 for i in range(NUM_BUFFERS)], "no_frames": 0, # set by configure_scan_geometry "H": H, "W": W, From c57511ba5c556c64b79f7faefb8d2044f7a5673e Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Fri, 31 Jul 2026 11:23:20 +0100 Subject: [PATCH 07/33] Fixing the race between flush and ptycho reconstruction which was causing the recon to fail when a new header was sent mid-reconstruction. Plan can be found in fix_header_race.md. Start with PR 1: Control flush split (refactor, behavior-preserving) --- fix_header_race.md | 306 +++++++++++++++++++++++++++++++++++++++++++ pipeline/control.py | 31 +++-- pipeline/pipeline.py | 10 +- 3 files changed, 334 insertions(+), 13 deletions(-) create mode 100644 fix_header_race.md diff --git a/fix_header_race.md b/fix_header_race.md new file mode 100644 index 0000000..2bf4e0e --- /dev/null +++ b/fix_header_race.md @@ -0,0 +1,306 @@ +# Fix Header Race: Aligned STXM + Ptycho Transition Plan (Revised) + +## Goal +Prevent the header/preemption race that can crash ptychography while keeping STXM and ptycho branches aligned across scan transitions. + +## Decisions Locked In +- Header-preemption completion and end-of-scan completion are both treated as full completion. +- Branch alignment is strict: both branches resume from the same first post-transition batch. +- Data arriving during transition should be preserved where possible. +- Overflow during blocked transition is fail-fast. +- Overflow fail-fast behavior: ControlOp publishes a final error signal first, then raises. +- max_blocked_frames is config-driven (default policy value: 10 x batch_size). + +## Why The Old Draft Was Not Sufficient +- It did not explicitly separate ptycho flush request from ptycho flush execution. +- It released the transition barrier too early. +- It did not define overflow behavior and signaling. +- It did not define a deterministic transition phase model. + +## Transition State Model + +### Shared state (thread-safe) +- **File:** pipeline/pipeline.py +- **Where:** StxmApp shared state initialization +- **Add:** + - transition_blocked_event (threading.Event) + - transition_phase: idle | waiting_quiesce | waiting_flush_exec + - ptycho_accum_flushed (bool) + - ptycho_recon_flushed (bool) + - max_blocked_frames (config value) + - transition_error (optional text for diagnostics) + +**Reason** +- Multiple scheduler worker threads can read/write transition state concurrently. +- Event plus explicit phase prevents ambiguous behavior. + +## Authoritative Transition Sequence + +1. **Header is received** in pipeline/header_io.py + - Validate header. + - Stage pending_geometry. + - Set preempt_requested. + - Set transition_blocked_event. + - Set transition_phase = waiting_quiesce. + - Clear ptycho flush ack flags. + - Emit header token to ControlOp. + +2. **ControlOp receives header** in pipeline/control.py + - Execute STXM-only flush. + - Do not request ptycho flush here. + - Keep transition blocked. + +3. **Ptycho recon sees preempt_requested** in pipeline/ptychography_ops.py + - Save partial result if needed. + - Emit recon_complete. + - Enter quiesced preemption path. + +4. **ControlOp receives recon_complete while phase=waiting_quiesce** + - Request ptycho flush now (request only): + - call ptycho_accum.flush() + - call ptycho_recon.flush() + - Set transition_phase = waiting_flush_exec. + +5. **Actual ptycho flush execution occurs later in operator compute** + - Accumulator: _perform_flush executes at top of next compute if requested. + - Recon: _perform_flush executes at top of next compute if requested. + - Each sets its corresponding ack flag when _perform_flush has actually run. + +6. **Barrier release condition** + - Only release transition when BOTH are true: + - ptycho_accum_flushed + - ptycho_recon_flushed + - Then: + - clear transition_blocked_event + - set transition_phase = idle + - reset ack flags for next transition + +This is the key timing rule: +- Ptycho flush is requested in ControlOp on recon_complete during waiting_quiesce. +- Ptycho flush is executed inside ptycho operators in _perform_flush during their next compute tick. +- Barrier is released only after both executions are acknowledged. + +## Concrete Code Changes + +### 1) Shared transition primitives +- **File:** pipeline/pipeline.py +- **Where:** StxmApp.__init__ and compose wiring +- **Change:** + - Add thread-safe transition fields to shared state. + - Pass shared state into GatherOp and ControlOp. + +### 2) Gather alignment gate and overflow accounting +- **File:** pipeline/data_io.py +- **Where:** GatherOp.__init__, setup, compute +- **Change:** + - Store shared transition state reference. + - Before emit, block when transition_blocked_event is set. + - Keep matched data cached while blocked. + - Track blocked cached frame count. + - Load max_blocked_frames from config/state. + +### 3) Gather fail-fast path +- **File:** pipeline/data_io.py and pipeline/control.py +- **Where:** GatherOp.compute overflow branch + ControlOp input handling +- **Change:** + - When blocked cache exceeds max_blocked_frames: + - set transition_error in shared state + - emit control message, e.g. transition_overflow + - ControlOp handles transition_overflow by: + - publishing final error signal + - raising RuntimeError + +### 4) Header starts transition cleanly +- **File:** pipeline/header_io.py +- **Where:** after pending geometry staging and preempt_requested +- **Change:** + - set transition_blocked_event + - set transition_phase=waiting_quiesce + - clear prior ack flags + - keep current scan_state geometry update behavior + +### 5) Split ControlOp flush ownership +- **File:** pipeline/control.py +- **Where:** constructor + helpers +- **Change:** + - Replace single flushable_ops with stxm_flush_ops and ptycho_flush_ops. + - Add helpers: + - do_stxm_flush() + - request_ptycho_flush() + - do_full_flush() (for non-transition full completion path) + +### 6) Control header branch +- **File:** pipeline/control.py +- **Where:** msg == header +- **Change:** + - STXM-only flush. + - Never request ptycho flush in this branch. + +### 7) Control recon_complete branch (phase-aware) +- **File:** pipeline/control.py +- **Where:** msg == recon_complete +- **Change:** + - If transition_phase == waiting_quiesce: + - request ptycho flush + - set waiting_flush_exec + - do not release barrier + - Else: + - existing full-completion behavior + +### 8) Control flush/start branch hardening +- **File:** pipeline/control.py +- **Where:** msg == flush +- **Change:** + - If transition blocked, do not request ptycho flush from this path. + - Use STXM-only or no-op policy to avoid race reintroduction. + +### 9) Ptycho flush ack on execution +- **File:** pipeline/ptychography_ops.py +- **Where:** + - PtychoAccumulatorOp._perform_flush + - PtychoReconstructionOp._perform_flush +- **Change:** + - Set ack flags when each _perform_flush has actually executed. + +### 10) Barrier release at safe point only +- **File:** pipeline/ptychography_ops.py or pipeline/control.py (single owner chosen) +- **Where:** after both ack flags observed true +- **Change:** + - Release transition barrier only when both ptycho flush executions are confirmed. + - Do not release solely at end of _apply_pending_geometry. + +### 11) Optional extra preemption guard +- **File:** pipeline/ptychography_ops.py +- **Where:** immediately before reconstruction_data/combine launch +- **Change:** + - Add second preempt check to reduce chance of one extra iteration starting. + +## Config Changes +- **File:** pipeline/config_test.yaml and pipeline/config_prod.yaml +- **Add under scheduler or a new transition section:** + - max_blocked_frames: integer +- **Default policy suggestion:** + - max_blocked_frames = 10 x image_src.batch_size (computed if unset) + +## Logging and Observability +- **File:** pipeline/header_io.py + - Log transition start, phase, and header id/shape. +- **File:** pipeline/control.py + - Log phase transitions. + - Log ptycho flush request moment. + - Log final error signal publication before raise. +- **File:** pipeline/ptychography_ops.py + - Log each actual _perform_flush execution and ack set. + - Log barrier release with both ack flags. +- **File:** pipeline/data_io.py + - Log blocked-cache growth and threshold crossing. + +## Expected Outcome +- No ptycho mid-iteration invalid-state reset from header/start control paths. +- Deterministic and explicit timing for ptycho flush request vs execution. +- Strict STXM/ptycho alignment through transition barrier. +- Bounded blocked buffering with explicit fail-fast and error signaling. + +## Notes And Remaining Non-goals +- Multiple rapid header bursts are still out of scope for this pass. +- Best-effort data preservation is targeted, but overflow path intentionally stops the run. + +## PR-sized Implementation Plan + +### PR 1: Control flush split (refactor, behavior-preserving) +- **Purpose:** Separate STXM vs ptycho flush ownership in ControlOp with minimal behavior change. +- **Includes:** + - Split ControlOp operator groups into `stxm_flush_ops` and `ptycho_flush_ops`. + - Add helper methods for STXM-only flush, ptycho-only flush request, and full flush. + - Update compose wiring to pass split groups. +- **Files:** + - pipeline/control.py + - pipeline/pipeline.py +- **Verification:** + - Existing single-scan behavior remains unchanged. + - Existing flush topics still publish as before. + +### PR 2: Transition state primitives + header phase start +- **Purpose:** Introduce thread-safe transition state and start transition on header. +- **Includes:** + - Add transition fields in shared state (`transition_blocked_event`, `transition_phase`, ack flags, `transition_error`). + - Add `max_blocked_frames` state value from config/default policy. + - Header path sets `waiting_quiesce` and blocks transition on valid header when ptycho is enabled. +- **Files:** + - pipeline/pipeline.py + - pipeline/header_io.py +- **Verification:** + - Header during active recon sets blocked event and phase to `waiting_quiesce`. + - No data-path behavior change yet. + +### PR 3: Phase-aware ControlOp handling +- **Purpose:** Make header and recon_complete handling deterministic by phase. +- **Includes:** + - Header branch performs STXM-only flush and never requests ptycho flush. + - recon_complete branch: + - if `waiting_quiesce`: request ptycho flush and move to `waiting_flush_exec` + - else: retain existing full-completion behavior + - flush/start branch hardening while transition is blocked. +- **Files:** + - pipeline/control.py +- **Verification:** + - Header no longer triggers immediate ptycho flush request. + - recon_complete in preemption path triggers ptycho flush request exactly once. + +### PR 4: Ptycho flush execution ack + safe barrier release +- **Purpose:** Release transition only after ptycho flush has actually executed. +- **Includes:** + - Set `ptycho_accum_flushed` in accumulator `_perform_flush`. + - Set `ptycho_recon_flushed` in recon `_perform_flush`. + - Release blocked event only when both ack flags are true. + - Reset phase to `idle` and clear acks for next transition. +- **Files:** + - pipeline/ptychography_ops.py + - pipeline/control.py (if release ownership is centralized) +- **Verification:** + - Logs clearly show: request -> execution ack -> barrier release. + - No early release before both ptycho flush executions. + +### PR 5: Gather alignment gate + fail-fast overflow signaling +- **Purpose:** Enforce strict branch alignment and bounded blocked buffering. +- **Includes:** + - Pass shared state into GatherOp. + - Block Gather emit while transition is blocked; keep matched data cached. + - Track blocked cached frame count. + - Overflow path (`> max_blocked_frames`): + - Gather emits `transition_overflow` control message and sets transition error context. + - ControlOp publishes final error signal first, then raises RuntimeError. +- **Files:** + - pipeline/data_io.py + - pipeline/control.py + - pipeline/pipeline.py + - pipeline/config_test.yaml + - pipeline/config_prod.yaml +- **Verification:** + - Normal transition preserves and resumes aligned batches. + - Forced overflow path publishes final error signal before raise. + +### PR 6: Hardening and observability +- **Purpose:** Improve resilience and diagnosability. +- **Includes:** + - Optional second preemption check just before PIE iteration launch. + - Log phase transitions, flush requests, flush execution acks, and barrier release. + - Tighten comments to document state machine semantics. +- **Files:** + - pipeline/ptychography_ops.py + - pipeline/control.py + - pipeline/header_io.py + - pipeline/data_io.py +- **Verification:** + - Repeated header-during-recon runs show deterministic ordering. + - Logs are sufficient to reconstruct transition timeline end-to-end. + +## Suggested Merge Order +1. PR 1 +2. PR 2 +3. PR 3 +4. PR 4 +5. PR 5 +6. PR 6 + +This order minimizes risk by landing structure first, then phase logic, then barrier correctness, then buffered alignment/fail-fast, then hardening. diff --git a/pipeline/control.py b/pipeline/control.py index e20a024..3dac8a1 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -17,7 +17,8 @@ class ControlOp(Operator): """ def __init__(self, fragment, *args, - flushable_ops: list[Operator] = None, + stxm_flush_ops: list[Operator] = None, + ptycho_flush_ops: list[Operator] = None, publish_backend = None, ptycho_accum = None, ptycho_recon = None, @@ -28,7 +29,8 @@ def __init__(self, fragment, *args, Args: fragment: Holoscan fragment - flushable_ops: List of operators that can be flushed + stxm_flush_ops: STXM-side operators that can be flushed + ptycho_flush_ops: Ptycho-side operators that can be flushed publish_backend: Backend instance for publishing flush messages ptycho_accum: PtychoAccumulatorOp (for the scoped projection advance) ptycho_recon: PtychoReconstructionOp (for the scoped projection advance) @@ -37,7 +39,8 @@ def __init__(self, fragment, *args, """ super().__init__(fragment, *args, **kwargs) self.logger = logging.getLogger(kwargs.get("name", "ControlOp")) - self.flushable_ops = flushable_ops + self.stxm_flush_ops = stxm_flush_ops or [] + self.ptycho_flush_ops = ptycho_flush_ops or [] self.publish_backend = publish_backend self.ptycho_accum = ptycho_accum self.ptycho_recon = ptycho_recon @@ -50,10 +53,20 @@ def __init__(self, fragment, *args, def setup(self, spec: OperatorSpec): spec.input("input").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128) - def _do_flush(self): - """Flush all flushable operators and broadcast the flush signals.""" - for op in self.flushable_ops: + def _do_stxm_flush(self): + """Flush STXM-side operators only.""" + for op in self.stxm_flush_ops: op.flush() + + def _do_ptycho_flush(self): + """Flush ptycho-side operators only.""" + for op in self.ptycho_flush_ops: + op.flush() + + def _do_full_flush(self): + """Flush STXM + ptycho operators and broadcast flush signals.""" + self._do_stxm_flush() + self._do_ptycho_flush() if self.publish_backend is not None: import numpy as np self.publish_backend.publish("stxm_flush", np.array([1])) # Simple signal @@ -70,7 +83,7 @@ def compute(self, op_input, op_output, context): # (after_iteration -> pty_out) and published the result before emitting # this, so flushing now is safe (Task 3: flush after the last iteration). self.logger.info("Reconstruction complete — flushing for next scan") - self._do_flush() + self._do_full_flush() self._flushed = True # PR4: tomography projection boundaries no longer round-trip through @@ -89,7 +102,7 @@ def compute(self, op_input, op_output, context): # recon_complete (on quiesce) also flushes — harmless, flush is # idempotent. Mark _flushed so the following start-flush skips. self.logger.info("Header received — flushing for reconfigured scan") - self._do_flush() + self._do_full_flush() self._flushed = True elif msg == "flush": @@ -100,7 +113,7 @@ def compute(self, op_input, op_output, context): self.logger.info("Start-flush skipped — already flushed on completion") self._flushed = False else: - self._do_flush() + self._do_full_flush() else: self.logger.info(f"Received unknown message: {msg}") diff --git a/pipeline/pipeline.py b/pipeline/pipeline.py index 1bf3a41..5c6b92f 100644 --- a/pipeline/pipeline.py +++ b/pipeline/pipeline.py @@ -144,7 +144,8 @@ def compose(self): name="sink_and_publish_op") # ===== Control Operator ===== - flushable_ops = [gather_op, position_src, sink_and_publish_op] + stxm_flush_ops = [gather_op, position_src, sink_and_publish_op] + ptycho_flush_ops = [] ptycho_accum = None # set below when ptychography is enabled ptycho_recon = None @@ -183,11 +184,12 @@ def compose(self): name="ptycho_publish", ) - flushable_ops.append(ptycho_accum) - flushable_ops.append(ptycho_recon) + ptycho_flush_ops.append(ptycho_accum) + ptycho_flush_ops.append(ptycho_recon) control_op = ControlOp(self, - flushable_ops=flushable_ops, + stxm_flush_ops=stxm_flush_ops, + ptycho_flush_ops=ptycho_flush_ops, publish_backend=publish_backend, ptycho_accum=ptycho_accum, ptycho_recon=ptycho_recon, From c50dfa28d6355d3dcf5507338dfd61524f67f57d Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Fri, 31 Jul 2026 11:34:34 +0100 Subject: [PATCH 08/33] Tested PR 1 and it all seems to work. Continuing with fixing the header race. Now PR 2: Transition state primitives + header phase start --- pipeline/header_io.py | 7 +++++++ pipeline/pipeline.py | 16 +++++++++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/pipeline/header_io.py b/pipeline/header_io.py index c51109c..50a28a3 100644 --- a/pipeline/header_io.py +++ b/pipeline/header_io.py @@ -151,6 +151,13 @@ def compute(self, op_input, op_output, context): "step_size_v": step_size_v, } self.ptycho_state["preempt_requested"].set() + if self.scan_state is not None: + self.scan_state["transition_blocked_event"].set() + self.scan_state["transition_phase"] = "waiting_quiesce" + self.scan_state["ptycho_accum_flushed"] = False + self.scan_state["ptycho_recon_flushed"] = False + self.scan_state["transition_error"] = None + self.logger.info("Header received — transition blocked, waiting for recon quiesce") # 3. Notify ControlOp so the STXM path flushes for the new dataset # (works even when ptychography is disabled). diff --git a/pipeline/pipeline.py b/pipeline/pipeline.py index 5c6b92f..7b39d6d 100644 --- a/pipeline/pipeline.py +++ b/pipeline/pipeline.py @@ -14,6 +14,7 @@ """ import logging +import threading from argparse import ArgumentParser from holoscan.core import Application @@ -66,6 +67,12 @@ def __init__(self, *args, **kwargs): "no_frames": 0, "num_projections": 1, "current_projection": 0, + "transition_blocked_event": threading.Event(), + "transition_phase": "idle", + "ptycho_accum_flushed": False, + "ptycho_recon_flushed": False, + "transition_error": None, + "max_blocked_frames": None, } super().__init__(*args, **kwargs) self.enable_metadata(True) @@ -260,16 +267,23 @@ def main(): # Load config to make kwargs available app.config(args.config) + image_src_config = app.kwargs('image_src') + # Get scheduler parameters from config via kwargs scheduler_config = app.kwargs('scheduler') num_decompress_ops = scheduler_config.get('num_decompress_ops', 4) worker_threads = scheduler_config.get('worker_threads', 6) + ptycho_cfg = app.kwargs("ptychography") + default_blocked_frames = int(image_src_config.get("batch_size", 100)) * 10 + app.scan_state["max_blocked_frames"] = int( + ptycho_cfg.get("max_blocked_frames", default_blocked_frames) + ) if ptycho_cfg else default_blocked_frames + # Set num_decompress_ops - will be used in compose() when run() is called app.num_decompress_ops = num_decompress_ops # Ptychography setup (before compose) - ptycho_cfg = app.kwargs("ptychography") if ptycho_cfg and ptycho_cfg.get("enabled", False): from ptychography_setup import init_ptycho_state From 874343105d786416c55ec0c05e37be7bfdf9bb82 Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Fri, 31 Jul 2026 11:45:12 +0100 Subject: [PATCH 09/33] PR 1 and PR 2 work. Able to send a header mid-reconstruction and holoscan reacts well. Now PR 3: Phase-aware ControlOp handling. Header no longer triggers immediate ptycho flush request. recon_complete in preemption path triggers ptycho flush request exactly once. --- pipeline/control.py | 35 ++++++++++++++++++++++++++++++----- 1 file changed, 30 insertions(+), 5 deletions(-) diff --git a/pipeline/control.py b/pipeline/control.py index 3dac8a1..ff14567 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -63,6 +63,10 @@ def _do_ptycho_flush(self): for op in self.ptycho_flush_ops: op.flush() + def _request_ptycho_flush(self): + """Request ptycho-side flush without doing any STXM work.""" + self._do_ptycho_flush() + def _do_full_flush(self): """Flush STXM + ptycho operators and broadcast flush signals.""" self._do_stxm_flush() @@ -77,11 +81,27 @@ def _do_full_flush(self): def compute(self, op_input, op_output, context): """Handle control messages.""" msg = op_input.receive("input") + transition_blocked = False + transition_phase = "idle" + if self.scan_state is not None: + transition_blocked = bool(self.scan_state.get("transition_blocked_event")) + if transition_blocked: + transition_blocked = self.scan_state["transition_blocked_event"].is_set() + transition_phase = self.scan_state.get("transition_phase", "idle") if msg == "recon_complete": # The recon finished its final iteration and has ALREADY saved # (after_iteration -> pty_out) and published the result before emitting # this, so flushing now is safe (Task 3: flush after the last iteration). + if transition_blocked and transition_phase == "waiting_quiesce": + self.logger.info( + "Recon quiesced for header transition — requesting ptycho flush" + ) + self._request_ptycho_flush() + if self.scan_state is not None: + self.scan_state["transition_phase"] = "waiting_flush_exec" + return + self.logger.info("Reconstruction complete — flushing for next scan") self._do_full_flush() self._flushed = True @@ -97,18 +117,23 @@ def compute(self, op_input, op_output, context): elif msg == "header": # A live header reconfigures the scan for a new dataset. Flush so the # STXM path saves+clears its current buffer before reconfiguration - # (SinkAndPublishOp.flush writes any unwritten scan). This works even - # when ptychography is disabled; when enabled, the recon's own - # recon_complete (on quiesce) also flushes — harmless, flush is - # idempotent. Mark _flushed so the following start-flush skips. + # (SinkAndPublishOp.flush writes any unwritten scan). Ptycho flush is + # deferred until the recon quiesces and emits recon_complete. self.logger.info("Header received — flushing for reconfigured scan") - self._do_full_flush() + self._do_stxm_flush() self._flushed = True elif msg == "flush": # Scan-start safety flush: only flush if the buffers aren't already # clean from a completion flush. If the previous scan completed, this # no-ops (no double flush); if it was interrupted, this cleans up. + if transition_blocked: + self.logger.info( + "Start-flush deferred for blocked transition — STXM-only flush now, ptycho waits for recon quiesce" + ) + self._do_stxm_flush() + self._flushed = True + return if self._flushed: self.logger.info("Start-flush skipped — already flushed on completion") self._flushed = False From f18f9668f35355f838a2e48237cb05af995fbb70 Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Fri, 31 Jul 2026 11:56:19 +0100 Subject: [PATCH 10/33] PR 3 was not sending stxm_flush to visualiser anymore, so the STXM visualisation wasn't flushing, it was just accumulating and looking strange. So we added a couple of lines in pipeline/control.py lines 124-126 which send a stxm_flush message to the visualiser --- pipeline/control.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pipeline/control.py b/pipeline/control.py index ff14567..2a77e87 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -121,6 +121,9 @@ def compute(self, op_input, op_output, context): # deferred until the recon quiesces and emits recon_complete. self.logger.info("Header received — flushing for reconfigured scan") self._do_stxm_flush() + if self.publish_backend is not None: + import numpy as np + self.publish_backend.publish("stxm_flush", np.array([1])) self._flushed = True elif msg == "flush": From 80f0225b1fab3366c61c25bf4a3ea8674f368a92 Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Fri, 31 Jul 2026 12:05:21 +0100 Subject: [PATCH 11/33] PR 1, 2 and 3 now work after fixing stxm flushing for visualiser. Now PR 4: Ptycho flush execution ack + safe barrier release. Release transition only after ptycho flush has actually executed. --- pipeline/ptychography_ops.py | 37 ++++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index f19bd2e..562b0f3 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -40,6 +40,35 @@ ) +def _ack_flush_and_maybe_release(scan_state, ack_key, logger): + """Mark one ptycho flush-execution ack and release the transition barrier + only when both ptycho acks are present in waiting_flush_exec.""" + if not scan_state: + return + + blocked_event = scan_state.get("transition_blocked_event") + if blocked_event is None or not blocked_event.is_set(): + return + + if scan_state.get("transition_phase", "idle") != "waiting_flush_exec": + return + + scan_state[ack_key] = True + logger.info("Transition flush ack set: %s=True", ack_key) + + if ( + scan_state.get("ptycho_accum_flushed", False) + and scan_state.get("ptycho_recon_flushed", False) + ): + blocked_event.clear() + scan_state["transition_phase"] = "idle" + scan_state["ptycho_accum_flushed"] = False + scan_state["ptycho_recon_flushed"] = False + logger.info( + "Transition barrier released after both ptycho flush executions" + ) + + class PtychoAccumulatorOp(Operator): """Fast batch accumulator for ptychography. @@ -97,9 +126,13 @@ def _perform_flush(self): when nothing has been accumulated since the last flush (free redundant flush).""" self._flush_requested = False + scan_state = self.ptycho_state.get("scan_state") or {} # A full flush is a scan boundary — any carried straddle-tail is stale. self._carry = None if not self._dirty: + _ack_flush_and_maybe_release( + scan_state, "ptycho_accum_flushed", self.logger + ) return self._reset_pingpong() for b in range(self.ptycho_state["num_buffers"]): @@ -112,6 +145,7 @@ def _perform_flush(self): self.ptycho_state["scan_center_px"] = None self._dirty = False self.logger.info("Flushed ptychography accumulator buffers (both)") + _ack_flush_and_maybe_release(scan_state, "ptycho_accum_flushed", self.logger) def _try_flip(self): """PR4: move the write cursor to the next buffer for the next projection, @@ -488,6 +522,7 @@ def _perform_flush(self): illumination. Set ``reset_probe=True`` to fully reset the probe too. """ self._flush_requested = False + scan_state = self.ptycho_state.get("scan_state") or {} self.current_iteration = 0 self.all_data_arrived = False self.post_stream_count = 0 @@ -527,6 +562,8 @@ def _perform_flush(self): "ptychography.reset_probe: true in the config." ) + _ack_flush_and_maybe_release(scan_state, "ptycho_recon_flushed", self.logger) + # ------------------------------------------------------------------ # Header preemption handshake (R-4) From 1eae67f65e7ffaeb6d73281e00b184f36212f7ae Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Fri, 31 Jul 2026 12:13:52 +0100 Subject: [PATCH 12/33] PR 4 works. Moving onto PR 5: Gather alignment gate + fail-fast overflow signaling; Enforce strict branch alignment and bounded blocked buffering. This implements the overflow fail error (ie determine the max number of buffered frames while ptycho state quiesces and flushes). --- pipeline/config_prod.yaml | 1 + pipeline/config_test.yaml | 1 + pipeline/control.py | 10 ++++++ pipeline/data_io.py | 64 ++++++++++++++++++++++++++++++++++++++- pipeline/pipeline.py | 2 ++ 5 files changed, 77 insertions(+), 1 deletion(-) diff --git a/pipeline/config_prod.yaml b/pipeline/config_prod.yaml index f59c0d2..d931389 100644 --- a/pipeline/config_prod.yaml +++ b/pipeline/config_prod.yaml @@ -58,6 +58,7 @@ ptychography: post_stream_iterations: 1 housekeeping_interval: 2 publish_interval: 1 + max_blocked_frames: 640 reset_probe: false # false = carry previous scan's probe forward (warm start); true = full probe reset each scan scan_range: [[0.0, 0.0], [0.0, 0.0]] # [[x1, y1], [x2, y2]] motor coords (µm) R: [351.0, 364.0] # scan range in pixels; must match actual scan extent diff --git a/pipeline/config_test.yaml b/pipeline/config_test.yaml index bdb338e..df99ded 100644 --- a/pipeline/config_test.yaml +++ b/pipeline/config_test.yaml @@ -61,4 +61,5 @@ ptychography: post_stream_iterations: 1 housekeeping_interval: 1 publish_interval: 1 + max_blocked_frames: 1000 reset_probe: true # false = carry previous scan's probe forward (warm start); true = full probe reset each scan diff --git a/pipeline/control.py b/pipeline/control.py index 2a77e87..cd029a4 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -143,6 +143,16 @@ def compute(self, op_input, op_output, context): else: self._do_full_flush() + elif msg == "transition_overflow": + err = "Transition overflow reported by GatherOp" + if self.scan_state is not None: + err = self.scan_state.get("transition_error") or err + self.logger.error(err) + if self.publish_backend is not None: + import numpy as np + self.publish_backend.publish("transition_error", np.array([1])) + raise RuntimeError(err) + else: self.logger.info(f"Received unknown message: {msg}") diff --git a/pipeline/data_io.py b/pipeline/data_io.py index f0c46cd..421a113 100644 --- a/pipeline/data_io.py +++ b/pipeline/data_io.py @@ -459,7 +459,14 @@ class GatherOp(Operator): gather -> masking_op -> publish """ - def __init__(self, fragment, *args, batch_size: int = 1, **kwargs): + def __init__( + self, + fragment, + *args, + batch_size: int = 1, + scan_state: dict = None, + **kwargs, + ): """ Initialize gather operator. @@ -472,6 +479,8 @@ def __init__(self, fragment, *args, batch_size: int = 1, **kwargs): self.position_ids = np.zeros((0,), dtype=int) self.count = 0 self.batch_size = int(batch_size) + self.scan_state = scan_state + self._default_max_blocked_frames = int(self.batch_size) * 10 # R-1 (PR3): latched once the series-end metadata is seen, so the final # partial batch (< batch_size) is drained instead of stranded. Reset on flush. self._series_finished = False @@ -481,6 +490,10 @@ def __init__(self, fragment, *args, batch_size: int = 1, **kwargs): # index race (data_io.py "size of axis is 0 but ... 64") when a flush # arrives mid-stream — e.g. PR2's header preemption. self._flush_requested = False + # PR5: one-shot overflow signal while blocked (avoid repeated control spam) + # plus a small counter for transition observability. + self._transition_overflow_reported = False + self._last_blocked_common = -1 self.logger = logging.getLogger(kwargs.get("name", "GatherOp")) super().__init__(fragment, *args, **kwargs) @@ -489,6 +502,24 @@ def setup(self, spec: OperatorSpec): spec.input("images").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128).condition(ConditionType.NONE) spec.input("positions").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128).condition(ConditionType.NONE) spec.output("output") + spec.output("control").condition(ConditionType.NONE) + + def _transition_blocked(self): + if self.scan_state is None: + return False + blocked_event = self.scan_state.get("transition_blocked_event") + return bool(blocked_event is not None and blocked_event.is_set()) + + def _max_blocked_frames(self): + if self.scan_state is None: + return self._default_max_blocked_frames + value = self.scan_state.get("max_blocked_frames") + if value is None: + return self._default_max_blocked_frames + try: + return int(value) + except (TypeError, ValueError): + return self._default_max_blocked_frames def flush(self): """Request a cache reset. Deferred to the top of the next compute() so it @@ -503,6 +534,8 @@ def _perform_flush(self): self.position_ids = np.zeros((0,), dtype=int) self.count = 0 self._series_finished = False + self._transition_overflow_reported = False + self._last_blocked_common = -1 self.logger.info( "[FLUSH VERIFY] GatherOp cleared: images=None, image_ids=0, " "positions=(0, 4), position_ids=0, count=0" @@ -557,6 +590,35 @@ def compute(self, op_input, op_output, context): drain = self._series_finished and int(common_ids.size) > 0 #if common_ids.size > 0: if int(common_ids.size) >= self.batch_size or drain: + if self._transition_blocked(): + blocked_common = int(common_ids.size) + if blocked_common != self._last_blocked_common: + self._last_blocked_common = blocked_common + self.logger.info( + "Transition blocked: caching %d matched frames in GatherOp", + blocked_common, + ) + + max_blocked_frames = self._max_blocked_frames() + if blocked_common > max_blocked_frames: + err = ( + "Gather blocked-cache overflow: matched=" + f"{blocked_common} exceeds max_blocked_frames=" + f"{max_blocked_frames}" + ) + if self.scan_state is not None: + self.scan_state["transition_error"] = err + if not self._transition_overflow_reported: + self._transition_overflow_reported = True + self.logger.error(err) + op_output.emit("transition_overflow", "control") + return + return + + # Transition resumed / not blocked: clear one-shot overflow latch. + self._transition_overflow_reported = False + self._last_blocked_common = -1 + # Create vectorized masks for efficient filtering mask_positions = np.isin(self.position_ids, common_ids) mask_images = np.isin(self.image_ids, common_ids) diff --git a/pipeline/pipeline.py b/pipeline/pipeline.py index 7b39d6d..c1fac79 100644 --- a/pipeline/pipeline.py +++ b/pipeline/pipeline.py @@ -113,6 +113,7 @@ def compose(self): gather_op = GatherOp(self, PeriodicCondition(self, int(0.01 * 1e9)), batch_size=self.kwargs('image_src')['batch_size'], + scan_state=self.scan_state, name="gather_op") # ===== Masking Operator (Processing - computes intensities) ===== @@ -239,6 +240,7 @@ def compose(self): # Control path: flush signal (start message → unconditional idempotent flush) self.add_flow(img_src, control_op, {("flush", "input")}) + self.add_flow(gather_op, control_op, {("control", "input")}) # Header path: live geometry header → control (flush for new dataset). # The ptycho geometry reconfigure is driven separately via the R-4 From 013ef6e27a7047f777b8fe0a826ef71fffd12555 Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Fri, 31 Jul 2026 12:19:13 +0100 Subject: [PATCH 13/33] Undoing the overflow fail behaviour, as it was conflicting with the data accumulation. --- pipeline/config_prod.yaml | 1 - pipeline/config_test.yaml | 1 - pipeline/control.py | 10 ---------- pipeline/data_io.py | 36 ++---------------------------------- pipeline/pipeline.py | 1 - 5 files changed, 2 insertions(+), 47 deletions(-) diff --git a/pipeline/config_prod.yaml b/pipeline/config_prod.yaml index d931389..f59c0d2 100644 --- a/pipeline/config_prod.yaml +++ b/pipeline/config_prod.yaml @@ -58,7 +58,6 @@ ptychography: post_stream_iterations: 1 housekeeping_interval: 2 publish_interval: 1 - max_blocked_frames: 640 reset_probe: false # false = carry previous scan's probe forward (warm start); true = full probe reset each scan scan_range: [[0.0, 0.0], [0.0, 0.0]] # [[x1, y1], [x2, y2]] motor coords (µm) R: [351.0, 364.0] # scan range in pixels; must match actual scan extent diff --git a/pipeline/config_test.yaml b/pipeline/config_test.yaml index df99ded..bdb338e 100644 --- a/pipeline/config_test.yaml +++ b/pipeline/config_test.yaml @@ -61,5 +61,4 @@ ptychography: post_stream_iterations: 1 housekeeping_interval: 1 publish_interval: 1 - max_blocked_frames: 1000 reset_probe: true # false = carry previous scan's probe forward (warm start); true = full probe reset each scan diff --git a/pipeline/control.py b/pipeline/control.py index cd029a4..2a77e87 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -143,16 +143,6 @@ def compute(self, op_input, op_output, context): else: self._do_full_flush() - elif msg == "transition_overflow": - err = "Transition overflow reported by GatherOp" - if self.scan_state is not None: - err = self.scan_state.get("transition_error") or err - self.logger.error(err) - if self.publish_backend is not None: - import numpy as np - self.publish_backend.publish("transition_error", np.array([1])) - raise RuntimeError(err) - else: self.logger.info(f"Received unknown message: {msg}") diff --git a/pipeline/data_io.py b/pipeline/data_io.py index 421a113..350ed01 100644 --- a/pipeline/data_io.py +++ b/pipeline/data_io.py @@ -480,7 +480,6 @@ def __init__( self.count = 0 self.batch_size = int(batch_size) self.scan_state = scan_state - self._default_max_blocked_frames = int(self.batch_size) * 10 # R-1 (PR3): latched once the series-end metadata is seen, so the final # partial batch (< batch_size) is drained instead of stranded. Reset on flush. self._series_finished = False @@ -490,9 +489,7 @@ def __init__( # index race (data_io.py "size of axis is 0 but ... 64") when a flush # arrives mid-stream — e.g. PR2's header preemption. self._flush_requested = False - # PR5: one-shot overflow signal while blocked (avoid repeated control spam) - # plus a small counter for transition observability. - self._transition_overflow_reported = False + # PR5: track blocked matched-frame growth for observability. self._last_blocked_common = -1 self.logger = logging.getLogger(kwargs.get("name", "GatherOp")) @@ -502,7 +499,6 @@ def setup(self, spec: OperatorSpec): spec.input("images").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128).condition(ConditionType.NONE) spec.input("positions").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128).condition(ConditionType.NONE) spec.output("output") - spec.output("control").condition(ConditionType.NONE) def _transition_blocked(self): if self.scan_state is None: @@ -510,17 +506,6 @@ def _transition_blocked(self): blocked_event = self.scan_state.get("transition_blocked_event") return bool(blocked_event is not None and blocked_event.is_set()) - def _max_blocked_frames(self): - if self.scan_state is None: - return self._default_max_blocked_frames - value = self.scan_state.get("max_blocked_frames") - if value is None: - return self._default_max_blocked_frames - try: - return int(value) - except (TypeError, ValueError): - return self._default_max_blocked_frames - def flush(self): """Request a cache reset. Deferred to the top of the next compute() so it never clears the caches while compute() is mid-synchronise (thread-safe).""" @@ -534,7 +519,6 @@ def _perform_flush(self): self.position_ids = np.zeros((0,), dtype=int) self.count = 0 self._series_finished = False - self._transition_overflow_reported = False self._last_blocked_common = -1 self.logger.info( "[FLUSH VERIFY] GatherOp cleared: images=None, image_ids=0, " @@ -598,25 +582,9 @@ def compute(self, op_input, op_output, context): "Transition blocked: caching %d matched frames in GatherOp", blocked_common, ) - - max_blocked_frames = self._max_blocked_frames() - if blocked_common > max_blocked_frames: - err = ( - "Gather blocked-cache overflow: matched=" - f"{blocked_common} exceeds max_blocked_frames=" - f"{max_blocked_frames}" - ) - if self.scan_state is not None: - self.scan_state["transition_error"] = err - if not self._transition_overflow_reported: - self._transition_overflow_reported = True - self.logger.error(err) - op_output.emit("transition_overflow", "control") - return return - # Transition resumed / not blocked: clear one-shot overflow latch. - self._transition_overflow_reported = False + # Transition resumed / not blocked. self._last_blocked_common = -1 # Create vectorized masks for efficient filtering diff --git a/pipeline/pipeline.py b/pipeline/pipeline.py index c1fac79..76f8d92 100644 --- a/pipeline/pipeline.py +++ b/pipeline/pipeline.py @@ -240,7 +240,6 @@ def compose(self): # Control path: flush signal (start message → unconditional idempotent flush) self.add_flow(img_src, control_op, {("flush", "input")}) - self.add_flow(gather_op, control_op, {("control", "input")}) # Header path: live geometry header → control (flush for new dataset). # The ptycho geometry reconfigure is driven separately via the R-4 From f9317cabc80b5c9b80166ba6b466377423814de2 Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Fri, 31 Jul 2026 12:43:15 +0100 Subject: [PATCH 14/33] Revert "Undoing the overflow fail behaviour, as it was conflicting with the data accumulation." This reverts commit 013ef6e27a7047f777b8fe0a826ef71fffd12555. All PR 5 needs to be un-done and roll back to PR 4, but using git revert rather than git reset for future clarity. --- pipeline/config_prod.yaml | 1 + pipeline/config_test.yaml | 1 + pipeline/control.py | 10 ++++++++++ pipeline/data_io.py | 36 ++++++++++++++++++++++++++++++++++-- pipeline/pipeline.py | 1 + 5 files changed, 47 insertions(+), 2 deletions(-) diff --git a/pipeline/config_prod.yaml b/pipeline/config_prod.yaml index f59c0d2..d931389 100644 --- a/pipeline/config_prod.yaml +++ b/pipeline/config_prod.yaml @@ -58,6 +58,7 @@ ptychography: post_stream_iterations: 1 housekeeping_interval: 2 publish_interval: 1 + max_blocked_frames: 640 reset_probe: false # false = carry previous scan's probe forward (warm start); true = full probe reset each scan scan_range: [[0.0, 0.0], [0.0, 0.0]] # [[x1, y1], [x2, y2]] motor coords (µm) R: [351.0, 364.0] # scan range in pixels; must match actual scan extent diff --git a/pipeline/config_test.yaml b/pipeline/config_test.yaml index bdb338e..df99ded 100644 --- a/pipeline/config_test.yaml +++ b/pipeline/config_test.yaml @@ -61,4 +61,5 @@ ptychography: post_stream_iterations: 1 housekeeping_interval: 1 publish_interval: 1 + max_blocked_frames: 1000 reset_probe: true # false = carry previous scan's probe forward (warm start); true = full probe reset each scan diff --git a/pipeline/control.py b/pipeline/control.py index 2a77e87..cd029a4 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -143,6 +143,16 @@ def compute(self, op_input, op_output, context): else: self._do_full_flush() + elif msg == "transition_overflow": + err = "Transition overflow reported by GatherOp" + if self.scan_state is not None: + err = self.scan_state.get("transition_error") or err + self.logger.error(err) + if self.publish_backend is not None: + import numpy as np + self.publish_backend.publish("transition_error", np.array([1])) + raise RuntimeError(err) + else: self.logger.info(f"Received unknown message: {msg}") diff --git a/pipeline/data_io.py b/pipeline/data_io.py index 350ed01..421a113 100644 --- a/pipeline/data_io.py +++ b/pipeline/data_io.py @@ -480,6 +480,7 @@ def __init__( self.count = 0 self.batch_size = int(batch_size) self.scan_state = scan_state + self._default_max_blocked_frames = int(self.batch_size) * 10 # R-1 (PR3): latched once the series-end metadata is seen, so the final # partial batch (< batch_size) is drained instead of stranded. Reset on flush. self._series_finished = False @@ -489,7 +490,9 @@ def __init__( # index race (data_io.py "size of axis is 0 but ... 64") when a flush # arrives mid-stream — e.g. PR2's header preemption. self._flush_requested = False - # PR5: track blocked matched-frame growth for observability. + # PR5: one-shot overflow signal while blocked (avoid repeated control spam) + # plus a small counter for transition observability. + self._transition_overflow_reported = False self._last_blocked_common = -1 self.logger = logging.getLogger(kwargs.get("name", "GatherOp")) @@ -499,6 +502,7 @@ def setup(self, spec: OperatorSpec): spec.input("images").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128).condition(ConditionType.NONE) spec.input("positions").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128).condition(ConditionType.NONE) spec.output("output") + spec.output("control").condition(ConditionType.NONE) def _transition_blocked(self): if self.scan_state is None: @@ -506,6 +510,17 @@ def _transition_blocked(self): blocked_event = self.scan_state.get("transition_blocked_event") return bool(blocked_event is not None and blocked_event.is_set()) + def _max_blocked_frames(self): + if self.scan_state is None: + return self._default_max_blocked_frames + value = self.scan_state.get("max_blocked_frames") + if value is None: + return self._default_max_blocked_frames + try: + return int(value) + except (TypeError, ValueError): + return self._default_max_blocked_frames + def flush(self): """Request a cache reset. Deferred to the top of the next compute() so it never clears the caches while compute() is mid-synchronise (thread-safe).""" @@ -519,6 +534,7 @@ def _perform_flush(self): self.position_ids = np.zeros((0,), dtype=int) self.count = 0 self._series_finished = False + self._transition_overflow_reported = False self._last_blocked_common = -1 self.logger.info( "[FLUSH VERIFY] GatherOp cleared: images=None, image_ids=0, " @@ -582,9 +598,25 @@ def compute(self, op_input, op_output, context): "Transition blocked: caching %d matched frames in GatherOp", blocked_common, ) + + max_blocked_frames = self._max_blocked_frames() + if blocked_common > max_blocked_frames: + err = ( + "Gather blocked-cache overflow: matched=" + f"{blocked_common} exceeds max_blocked_frames=" + f"{max_blocked_frames}" + ) + if self.scan_state is not None: + self.scan_state["transition_error"] = err + if not self._transition_overflow_reported: + self._transition_overflow_reported = True + self.logger.error(err) + op_output.emit("transition_overflow", "control") + return return - # Transition resumed / not blocked. + # Transition resumed / not blocked: clear one-shot overflow latch. + self._transition_overflow_reported = False self._last_blocked_common = -1 # Create vectorized masks for efficient filtering diff --git a/pipeline/pipeline.py b/pipeline/pipeline.py index 76f8d92..c1fac79 100644 --- a/pipeline/pipeline.py +++ b/pipeline/pipeline.py @@ -240,6 +240,7 @@ def compose(self): # Control path: flush signal (start message → unconditional idempotent flush) self.add_flow(img_src, control_op, {("flush", "input")}) + self.add_flow(gather_op, control_op, {("control", "input")}) # Header path: live geometry header → control (flush for new dataset). # The ptycho geometry reconfigure is driven separately via the R-4 From f739ca75265b307e0ed5a167ebabc9ae0070251d Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Fri, 31 Jul 2026 12:44:29 +0100 Subject: [PATCH 15/33] Revert "PR 4 works. Moving onto PR 5: Gather alignment gate + fail-fast overflow signaling; Enforce strict branch alignment and bounded blocked buffering. This implements the overflow fail error (ie determine the max number of buffered frames while ptycho state quiesces and flushes)." This reverts commit 1eae67f65e7ffaeb6d73281e00b184f36212f7ae. Fully reverting PR 5. --- pipeline/config_prod.yaml | 1 - pipeline/config_test.yaml | 1 - pipeline/control.py | 10 ------ pipeline/data_io.py | 64 +-------------------------------------- pipeline/pipeline.py | 2 -- 5 files changed, 1 insertion(+), 77 deletions(-) diff --git a/pipeline/config_prod.yaml b/pipeline/config_prod.yaml index d931389..f59c0d2 100644 --- a/pipeline/config_prod.yaml +++ b/pipeline/config_prod.yaml @@ -58,7 +58,6 @@ ptychography: post_stream_iterations: 1 housekeeping_interval: 2 publish_interval: 1 - max_blocked_frames: 640 reset_probe: false # false = carry previous scan's probe forward (warm start); true = full probe reset each scan scan_range: [[0.0, 0.0], [0.0, 0.0]] # [[x1, y1], [x2, y2]] motor coords (µm) R: [351.0, 364.0] # scan range in pixels; must match actual scan extent diff --git a/pipeline/config_test.yaml b/pipeline/config_test.yaml index df99ded..bdb338e 100644 --- a/pipeline/config_test.yaml +++ b/pipeline/config_test.yaml @@ -61,5 +61,4 @@ ptychography: post_stream_iterations: 1 housekeeping_interval: 1 publish_interval: 1 - max_blocked_frames: 1000 reset_probe: true # false = carry previous scan's probe forward (warm start); true = full probe reset each scan diff --git a/pipeline/control.py b/pipeline/control.py index cd029a4..2a77e87 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -143,16 +143,6 @@ def compute(self, op_input, op_output, context): else: self._do_full_flush() - elif msg == "transition_overflow": - err = "Transition overflow reported by GatherOp" - if self.scan_state is not None: - err = self.scan_state.get("transition_error") or err - self.logger.error(err) - if self.publish_backend is not None: - import numpy as np - self.publish_backend.publish("transition_error", np.array([1])) - raise RuntimeError(err) - else: self.logger.info(f"Received unknown message: {msg}") diff --git a/pipeline/data_io.py b/pipeline/data_io.py index 421a113..f0c46cd 100644 --- a/pipeline/data_io.py +++ b/pipeline/data_io.py @@ -459,14 +459,7 @@ class GatherOp(Operator): gather -> masking_op -> publish """ - def __init__( - self, - fragment, - *args, - batch_size: int = 1, - scan_state: dict = None, - **kwargs, - ): + def __init__(self, fragment, *args, batch_size: int = 1, **kwargs): """ Initialize gather operator. @@ -479,8 +472,6 @@ def __init__( self.position_ids = np.zeros((0,), dtype=int) self.count = 0 self.batch_size = int(batch_size) - self.scan_state = scan_state - self._default_max_blocked_frames = int(self.batch_size) * 10 # R-1 (PR3): latched once the series-end metadata is seen, so the final # partial batch (< batch_size) is drained instead of stranded. Reset on flush. self._series_finished = False @@ -490,10 +481,6 @@ def __init__( # index race (data_io.py "size of axis is 0 but ... 64") when a flush # arrives mid-stream — e.g. PR2's header preemption. self._flush_requested = False - # PR5: one-shot overflow signal while blocked (avoid repeated control spam) - # plus a small counter for transition observability. - self._transition_overflow_reported = False - self._last_blocked_common = -1 self.logger = logging.getLogger(kwargs.get("name", "GatherOp")) super().__init__(fragment, *args, **kwargs) @@ -502,24 +489,6 @@ def setup(self, spec: OperatorSpec): spec.input("images").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128).condition(ConditionType.NONE) spec.input("positions").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128).condition(ConditionType.NONE) spec.output("output") - spec.output("control").condition(ConditionType.NONE) - - def _transition_blocked(self): - if self.scan_state is None: - return False - blocked_event = self.scan_state.get("transition_blocked_event") - return bool(blocked_event is not None and blocked_event.is_set()) - - def _max_blocked_frames(self): - if self.scan_state is None: - return self._default_max_blocked_frames - value = self.scan_state.get("max_blocked_frames") - if value is None: - return self._default_max_blocked_frames - try: - return int(value) - except (TypeError, ValueError): - return self._default_max_blocked_frames def flush(self): """Request a cache reset. Deferred to the top of the next compute() so it @@ -534,8 +503,6 @@ def _perform_flush(self): self.position_ids = np.zeros((0,), dtype=int) self.count = 0 self._series_finished = False - self._transition_overflow_reported = False - self._last_blocked_common = -1 self.logger.info( "[FLUSH VERIFY] GatherOp cleared: images=None, image_ids=0, " "positions=(0, 4), position_ids=0, count=0" @@ -590,35 +557,6 @@ def compute(self, op_input, op_output, context): drain = self._series_finished and int(common_ids.size) > 0 #if common_ids.size > 0: if int(common_ids.size) >= self.batch_size or drain: - if self._transition_blocked(): - blocked_common = int(common_ids.size) - if blocked_common != self._last_blocked_common: - self._last_blocked_common = blocked_common - self.logger.info( - "Transition blocked: caching %d matched frames in GatherOp", - blocked_common, - ) - - max_blocked_frames = self._max_blocked_frames() - if blocked_common > max_blocked_frames: - err = ( - "Gather blocked-cache overflow: matched=" - f"{blocked_common} exceeds max_blocked_frames=" - f"{max_blocked_frames}" - ) - if self.scan_state is not None: - self.scan_state["transition_error"] = err - if not self._transition_overflow_reported: - self._transition_overflow_reported = True - self.logger.error(err) - op_output.emit("transition_overflow", "control") - return - return - - # Transition resumed / not blocked: clear one-shot overflow latch. - self._transition_overflow_reported = False - self._last_blocked_common = -1 - # Create vectorized masks for efficient filtering mask_positions = np.isin(self.position_ids, common_ids) mask_images = np.isin(self.image_ids, common_ids) diff --git a/pipeline/pipeline.py b/pipeline/pipeline.py index c1fac79..7b39d6d 100644 --- a/pipeline/pipeline.py +++ b/pipeline/pipeline.py @@ -113,7 +113,6 @@ def compose(self): gather_op = GatherOp(self, PeriodicCondition(self, int(0.01 * 1e9)), batch_size=self.kwargs('image_src')['batch_size'], - scan_state=self.scan_state, name="gather_op") # ===== Masking Operator (Processing - computes intensities) ===== @@ -240,7 +239,6 @@ def compose(self): # Control path: flush signal (start message → unconditional idempotent flush) self.add_flow(img_src, control_op, {("flush", "input")}) - self.add_flow(gather_op, control_op, {("control", "input")}) # Header path: live geometry header → control (flush for new dataset). # The ptycho geometry reconfigure is driven separately via the R-4 From e9d99e95918f49fcd50aaa348a30d9a6a74e2fad Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Fri, 31 Jul 2026 12:53:37 +0100 Subject: [PATCH 16/33] Rolled back to PR 4. Just added a few logs and comments. --- pipeline/control.py | 7 +++++++ pipeline/data_io.py | 37 +++++++++++++++++++++++++++++++++++- pipeline/header_io.py | 16 ++++++++++++++++ pipeline/ptychography_ops.py | 10 ++++++++-- 4 files changed, 67 insertions(+), 3 deletions(-) diff --git a/pipeline/control.py b/pipeline/control.py index 2a77e87..6097dd3 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -88,6 +88,12 @@ def compute(self, op_input, op_output, context): if transition_blocked: transition_blocked = self.scan_state["transition_blocked_event"].is_set() transition_phase = self.scan_state.get("transition_phase", "idle") + self.logger.info( + "Control message=%s (transition_blocked=%s phase=%s)", + msg, + transition_blocked, + transition_phase, + ) if msg == "recon_complete": # The recon finished its final iteration and has ALREADY saved @@ -100,6 +106,7 @@ def compute(self, op_input, op_output, context): self._request_ptycho_flush() if self.scan_state is not None: self.scan_state["transition_phase"] = "waiting_flush_exec" + self.logger.info("Transition state -> waiting_flush_exec") return self.logger.info("Reconstruction complete — flushing for next scan") diff --git a/pipeline/data_io.py b/pipeline/data_io.py index f0c46cd..1f3b75e 100644 --- a/pipeline/data_io.py +++ b/pipeline/data_io.py @@ -459,7 +459,14 @@ class GatherOp(Operator): gather -> masking_op -> publish """ - def __init__(self, fragment, *args, batch_size: int = 1, **kwargs): + def __init__( + self, + fragment, + *args, + batch_size: int = 1, + scan_state: dict = None, + **kwargs, + ): """ Initialize gather operator. @@ -472,6 +479,7 @@ def __init__(self, fragment, *args, batch_size: int = 1, **kwargs): self.position_ids = np.zeros((0,), dtype=int) self.count = 0 self.batch_size = int(batch_size) + self.scan_state = scan_state # R-1 (PR3): latched once the series-end metadata is seen, so the final # partial batch (< batch_size) is drained instead of stranded. Reset on flush. self._series_finished = False @@ -481,6 +489,8 @@ def __init__(self, fragment, *args, batch_size: int = 1, **kwargs): # index race (data_io.py "size of axis is 0 but ... 64") when a flush # arrives mid-stream — e.g. PR2's header preemption. self._flush_requested = False + # PR5: track blocked matched-frame growth for observability. + self._last_blocked_common = -1 self.logger = logging.getLogger(kwargs.get("name", "GatherOp")) super().__init__(fragment, *args, **kwargs) @@ -490,6 +500,12 @@ def setup(self, spec: OperatorSpec): spec.input("positions").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128).condition(ConditionType.NONE) spec.output("output") + def _transition_blocked(self): + if self.scan_state is None: + return False + blocked_event = self.scan_state.get("transition_blocked_event") + return bool(blocked_event is not None and blocked_event.is_set()) + def flush(self): """Request a cache reset. Deferred to the top of the next compute() so it never clears the caches while compute() is mid-synchronise (thread-safe).""" @@ -503,6 +519,7 @@ def _perform_flush(self): self.position_ids = np.zeros((0,), dtype=int) self.count = 0 self._series_finished = False + self._last_blocked_common = -1 self.logger.info( "[FLUSH VERIFY] GatherOp cleared: images=None, image_ids=0, " "positions=(0, 4), position_ids=0, count=0" @@ -557,6 +574,24 @@ def compute(self, op_input, op_output, context): drain = self._series_finished and int(common_ids.size) > 0 #if common_ids.size > 0: if int(common_ids.size) >= self.batch_size or drain: + blocked_common = int(common_ids.size) + if self._transition_blocked(): + if blocked_common != self._last_blocked_common: + self._last_blocked_common = blocked_common + self.logger.info( + "Transition blocked: caching %d matched frames in GatherOp", + blocked_common, + ) + return + + # Transition resumed / not blocked. + if self._last_blocked_common >= 0: + self.logger.info( + "Transition unblocked: resuming GatherOp emit with %d cached matched frames", + blocked_common, + ) + self._last_blocked_common = -1 + # Create vectorized masks for efficient filtering mask_positions = np.isin(self.position_ids, common_ids) mask_images = np.isin(self.image_ids, common_ids) diff --git a/pipeline/header_io.py b/pipeline/header_io.py index 50a28a3..72fd8a6 100644 --- a/pipeline/header_io.py +++ b/pipeline/header_io.py @@ -135,9 +135,21 @@ def compute(self, op_input, op_output, context): # (not only via configure_scan_geometry) so the STXM sink can segment # per projection even when ptychography is disabled. if self.scan_state is not None: + prev_phase = self.scan_state.get("transition_phase", "idle") + prev_blocked = False + blocked_event = self.scan_state.get("transition_blocked_event") + if blocked_event is not None: + prev_blocked = blocked_event.is_set() self.scan_state["num_projections"] = num_projections self.scan_state["current_projection"] = 0 self.scan_state["no_frames"] = npoints_h * npoints_v + self.logger.info( + "Header accepted: no_frames=%d num_projections=%d (prev_phase=%s, prev_blocked=%s)", + self.scan_state["no_frames"], + self.scan_state["num_projections"], + prev_phase, + prev_blocked, + ) # 2. Ptycho path: stage geometry + request preemption (R-4). The recon op # applies configure_scan_geometry once it has quiesced, so no buffer @@ -157,8 +169,12 @@ def compute(self, op_input, op_output, context): self.scan_state["ptycho_accum_flushed"] = False self.scan_state["ptycho_recon_flushed"] = False self.scan_state["transition_error"] = None + self.logger.info( + "Transition state -> waiting_quiesce (blocked=True, acks reset)" + ) self.logger.info("Header received — transition blocked, waiting for recon quiesce") # 3. Notify ControlOp so the STXM path flushes for the new dataset # (works even when ptychography is disabled). + self.logger.info("Emitting header control token") op_output.emit("header", "header") diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 562b0f3..c2dfd02 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -54,7 +54,13 @@ def _ack_flush_and_maybe_release(scan_state, ack_key, logger): return scan_state[ack_key] = True - logger.info("Transition flush ack set: %s=True", ack_key) + logger.info( + "Transition flush ack set: %s=True (phase=%s accum=%s recon=%s)", + ack_key, + scan_state.get("transition_phase", "idle"), + scan_state.get("ptycho_accum_flushed", False), + scan_state.get("ptycho_recon_flushed", False), + ) if ( scan_state.get("ptycho_accum_flushed", False) @@ -65,7 +71,7 @@ def _ack_flush_and_maybe_release(scan_state, ack_key, logger): scan_state["ptycho_accum_flushed"] = False scan_state["ptycho_recon_flushed"] = False logger.info( - "Transition barrier released after both ptycho flush executions" + "Transition barrier released after both ptycho flush executions (phase=idle, blocked=False)" ) From 3c337aac01c5711e749893439a64aaabac94edd5 Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Fri, 31 Jul 2026 14:22:44 +0100 Subject: [PATCH 17/33] Pipeline published stxm_flush message for visualiser so that it knows to clear/flush visualisation when a new projection comes in. --- pipeline/publish.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/pipeline/publish.py b/pipeline/publish.py index 73b1f47..5c5856e 100644 --- a/pipeline/publish.py +++ b/pipeline/publish.py @@ -188,7 +188,14 @@ def flush(self): self._written = False self._projection = 0 self._proj_frame_count = 0 - + + def _publish_stxm_flush(self): + """Notify downstream consumers that the current STXM projection is done.""" + if self.backend is None: + return + import numpy as np + self.backend.publish("stxm_flush", np.array([1])) + def compute(self, op_input, op_output, context): """Receive, publish, and save processed data using metadata.""" # Initialize backend on first call @@ -272,6 +279,10 @@ def compute(self, op_input, op_output, context): if self._proj_frame_count >= proj_no_frames: if self.publish_folder is not None and series_id is not None: self._write_projection_file(series_id, self._projection) + + # Notify the visualizer / downstream consumers that the + # previous projection is complete and should be cleared. + self._publish_stxm_flush() self._proj_frame_count -= proj_no_frames # carry overshoot count self._projection += 1 else: From 94ab59b971d78c0b139c3af30b1530ab637d7029 Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Fri, 31 Jul 2026 15:42:18 +0100 Subject: [PATCH 18/33] Setting scan_center_py to None in ptychography_ops for each projection when buffer flips; this will hopefully help with out-of-bound problems at syn_probe_intensity_kernel --- pipeline/ptychography_ops.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index c2dfd02..5d1f577 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -166,6 +166,10 @@ def _try_flip(self): self.ptycho_state["buf_free"][other] = False self.ptycho_state["write_idx"] = other self.ptycho_state["filled_until"][other] = 0 + # Per-projection auto-centering: the first batch written into the + # new buffer must derive a fresh center for that projection. + self.ptycho_state["scan_center_py"] = None + self.ptycho_state["scan_center_px"] = None self.logger.info("Accumulator flipped to buffer %d for next projection", other) return True @@ -397,9 +401,12 @@ def _transform_positions(self, positions_txyz): halfview = self.ptycho_state["N"][0]/1.2/2 * 1e-6 * pty_model.scan.scale[0] self.ptycho_state["scan_center_py"] = float(cp.mean(py)) + halfview self.ptycho_state["scan_center_px"] = float(cp.mean(px)) + scan_state = self.ptycho_state.get("scan_state") or {} + projection = int(scan_state.get("current_projection", 0)) self.logger.info( - "Auto-centring scan: center_py=%.6e m, center_px=%.6e m " + "Auto-centring projection %d: center_py=%.6e m, center_px=%.6e m " "(theta=%.2f°)", + projection, self.ptycho_state["scan_center_py"], self.ptycho_state["scan_center_px"], theta, From 256085729ab77bee4b580737a92ac4c1e653ec6a Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Fri, 31 Jul 2026 15:52:08 +0100 Subject: [PATCH 19/33] Very small change: changing vis_stxm visualisation of ptycho obj phase and probe phase to grayscal rather than twilight. --- pipeline/vis/vis_stxm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pipeline/vis/vis_stxm.py b/pipeline/vis/vis_stxm.py index 41ddbde..390b979 100644 --- a/pipeline/vis/vis_stxm.py +++ b/pipeline/vis/vis_stxm.py @@ -225,9 +225,9 @@ def build_combined_figure(): placeholder = np.zeros((64, 64)) * np.nan ptycho_axes_info = [ - (ax_obj_phase, "twilight", "object_phase"), + (ax_obj_phase, "gray", "object_phase"), (ax_obj_amp, "gray", "object_amp"), - (ax_prb_phase, "twilight", "probe_phase"), + (ax_prb_phase, "gray", "probe_phase"), (ax_prb_amp, "gray", "probe_amp"), ] ptycho_ims = {} From feea7c33af9497197d0ebfd731a6ba3442e83b23 Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Mon, 3 Aug 2026 16:56:15 +0100 Subject: [PATCH 20/33] Change flux calculation so that probe is always normalised to flux, and flux is only calculated if json file value is < 0 --- pipeline/ptychography_ops.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 5d1f577..27415c0 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -752,13 +752,14 @@ def compute(self, op_input, op_output, context): ] # Flux normalization — compute once on first iteration - if self.current_iteration == 0 and pty_model.source.flux < 0: + if self.current_iteration == 0: raw_cpu = cp.asnumpy(self.ptycho_state["raw_gpu"][r][:n_filled]) dp = pty_data.dp - pty_model.source.flux = float(np.sum( - np.sum(raw_cpu, 0)[dp == 1] - ) / raw_cpu.shape[0]) - self.logger.info("Computed flux = %.2f from %d frames", pty_model.source.flux, n_filled) + if pty_model.source.flux < 0: + pty_model.source.flux = float(np.sum( + np.sum(raw_cpu, 0)[dp == 1] + ) / raw_cpu.shape[0]) + self.logger.info("Computed flux = %.2f from %d frames", pty_model.source.flux, n_filled) for trial_idx in range(pty_model.scan.tris_n): pty_model.probe.array_states[:, :, :, :, trial_idx, :, :] = setPower( pty_model.probe.array_states[:, :, :, :, trial_idx, :, :], From 73ebe5a5fa1d503f7e91c17eb56d367250582285 Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Mon, 3 Aug 2026 18:06:02 +0100 Subject: [PATCH 21/33] Add projection_advanced flag to ptycho_ops so that if we flip to the next projection current iteration doesn't do +1. This means that all projections will start from iteration 0. This makes sure they'll calculate flux and do probe normalisation. --- pipeline/ptychography_ops.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 27415c0..101004b 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -903,6 +903,7 @@ def compute(self, op_input, op_output, context): # been filling meanwhile). Resets _completed → resume next tick. iter_done = self.current_iteration # _flip_read resets it to 0 self._flip_read(scan_state) + projection_advanced = True self.logger.info( "Projection %d/%d complete at iteration %d — flipped read " "buffer to %d for next projection", @@ -935,7 +936,9 @@ def compute(self, op_input, op_output, context): n_filled, no_frames, ) - self.current_iteration += 1 + + if not projection_advanced: + self.current_iteration += 1 # ------------------------------------------------------------------ From a908209c4c623a5c95fa0ac01919370f6bdf103c Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Mon, 3 Aug 2026 18:23:38 +0100 Subject: [PATCH 22/33] Add projection_advancecd = False flag before completion block. --- pipeline/ptychography_ops.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 101004b..70cea50 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -879,6 +879,7 @@ def compute(self, op_input, op_output, context): # are still arriving, which would flush GatherOp mid-compute (race) and # reset the object before the full scan is reconstructed. ControlOp # flushes on this signal (Task 3: flush after the last iteration). + projection_advanced = False if is_last and self.all_data_arrived and not self._completed: self._completed = True if num_projections > 1: From 5c3ae99b66f3e72b7a610d5195b91b465d93eec8 Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Tue, 11 Aug 2026 10:42:55 +0100 Subject: [PATCH 23/33] A few changes: added reduced visualiser (probe amp + object phase only). Also increased padding zone (this is important with fast scans selun/ck3m because they may overshoot). Stopped printing start message from detector (too long). --- pipeline/data_io.py | 2 +- pipeline/ptychography_ops.py | 2 +- pipeline/ptychography_setup.py | 4 ++-- pipeline/vis/vis_stxm.py | 40 +++++++++++++++++++++++++++++++++- 4 files changed, 43 insertions(+), 5 deletions(-) diff --git a/pipeline/data_io.py b/pipeline/data_io.py index 1f3b75e..373eaff 100644 --- a/pipeline/data_io.py +++ b/pipeline/data_io.py @@ -45,7 +45,7 @@ def receive_cbor_message(zmq_message) -> tuple[str, cbor2.CBORTag, int, dict]: if msg_type == "image": compressed_image, image_id, msg_content = msg["data"]["threshold_1"], msg["image_id"], None elif msg_type == "start": - print(f"{msg_type} message content: {msg}") + print(f"Received {msg_type} message.") # content: {msg} compressed_image, image_id, msg_content = None, None, msg elif msg_type == "end": print(f"{msg_type} message content: {msg}") diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 70cea50..1ff44f5 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -398,7 +398,7 @@ def _transform_positions(self, positions_txyz): # Auto-centre: capture scan centre from first batch if self.ptycho_state["scan_center_py"] is None: - halfview = self.ptycho_state["N"][0]/1.2/2 * 1e-6 * pty_model.scan.scale[0] + halfview = self.ptycho_state["N"][0]/2/2 * 1e-6 * pty_model.scan.scale[0] self.ptycho_state["scan_center_py"] = float(cp.mean(py)) + halfview self.ptycho_state["scan_center_px"] = float(cp.mean(px)) scan_state = self.ptycho_state.get("scan_state") or {} diff --git a/pipeline/ptychography_setup.py b/pipeline/ptychography_setup.py index 8be34f7..ff3eef3 100644 --- a/pipeline/ptychography_setup.py +++ b/pipeline/ptychography_setup.py @@ -162,8 +162,8 @@ def configure_scan_geometry( # Scan extent in microns with 20% padding (same formula as PtyREX streaming) N = [ - ((npoints_v - 1) * step_size_v) * 1.2, - ((npoints_h - 1) * step_size_h) * 1.2, + ((npoints_v - 1) * step_size_v) * 2, + ((npoints_h - 1) * step_size_h) * 2, ] logger.info( "Configuring scan: %d x %d points, step %.3f x %.3f µm → " diff --git a/pipeline/vis/vis_stxm.py b/pipeline/vis/vis_stxm.py index 390b979..27c9f6a 100644 --- a/pipeline/vis/vis_stxm.py +++ b/pipeline/vis/vis_stxm.py @@ -233,6 +233,39 @@ def build_combined_figure(): ptycho_ims = {} for ax, cmap, key in ptycho_axes_info: im = ax.imshow(placeholder, cmap=cmap, interpolation='nearest', aspect='equal') + ax.invert_xaxis() + fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + ptycho_ims[key] = im + + return fig, ax_stxm_outer, ax_stxm_inner, ptycho_ims + +def build_combined_figure_reduced(): + """2-row x 2-col layout: STXM top, ptycho object phase + probe modulus bottom.""" + plt.style.use('dark_background') + matplotlib.rcParams.update({'font.size': 8}) + + fig = plt.figure(figsize=(10, 10)) + gs = fig.add_gridspec(2, 2, hspace=0.35, wspace=0.3) + + ax_stxm_outer = fig.add_subplot(gs[0, 0]) + ax_stxm_inner = fig.add_subplot(gs[0, 1]) + ax_obj_phase = fig.add_subplot(gs[1, 0]) + ax_prb_amp = fig.add_subplot(gs[1, 1]) + + ax_stxm_outer.set_title("STXM Outer") + ax_stxm_inner.set_title("STXM Inner") + ax_obj_phase.set_title("Object Phase") + ax_prb_amp.set_title("Probe Amplitude") + + placeholder = np.zeros((64, 64)) * np.nan + ptycho_axes_info = [ + (ax_obj_phase, "gray", "object_phase"), + (ax_prb_amp, "gray", "probe_amp"), + ] + ptycho_ims = {} + for ax, cmap, key in ptycho_axes_info: + im = ax.imshow(placeholder, cmap=cmap, interpolation='nearest', aspect='equal') + ax.invert_xaxis() fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) ptycho_ims[key] = im @@ -316,7 +349,12 @@ def animate_combined(i): threading.Thread(target=receive_stxm_data, args=(sub_backend,), daemon=True).start() threading.Thread(target=receive_ptycho_data, args=(sub_backend,), daemon=True).start() - fig, ax_stxm_outer, ax_stxm_inner, ptycho_ims = build_combined_figure() + reduced_flag = True + + if reduced_flag: + fig, ax_stxm_outer, ax_stxm_inner, ptycho_ims = build_combined_figure_reduced() + else: + fig, ax_stxm_outer, ax_stxm_inner, ptycho_ims = build_combined_figure() if all(v is not None for v in (args.xmin, args.xmax, args.ymin, args.ymax)): ax_stxm_outer.set_xlim(-args.xmax, -args.xmin) From fc548a550179c8ffffa2c096cd9ab18c0c7a9089 Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Tue, 11 Aug 2026 19:37:56 +0100 Subject: [PATCH 24/33] Change the auto-centering so that only takes into account half of the batch - this requires the batch size to be at least 2x the nX, but it's useful for cases where there is a lot of overshoot in x. Also I've included a sign for the calculation of y centre, because it changes in a tomography projection to projection. --- pipeline/ptychography_ops.py | 11 +++++++---- pipeline/ptychography_setup.py | 2 +- pipeline/vis/vis_stxm.py | 3 ++- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 1ff44f5..4688282 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -398,11 +398,14 @@ def _transform_positions(self, positions_txyz): # Auto-centre: capture scan centre from first batch if self.ptycho_state["scan_center_py"] is None: - halfview = self.ptycho_state["N"][0]/2/2 * 1e-6 * pty_model.scan.scale[0] - self.ptycho_state["scan_center_py"] = float(cp.mean(py)) + halfview - self.ptycho_state["scan_center_px"] = float(cp.mean(px)) - scan_state = self.ptycho_state.get("scan_state") or {} projection = int(scan_state.get("current_projection", 0)) + halfview = self.ptycho_state["N"][0]/2/2 * 1e-6 * pty_model.scan.scale[0] + sign = 1 if projection % 2 == 0 else -1 + batch_size_here = py.shape[0] + self.ptycho_state["scan_center_py"] = float(cp.mean(py)) + sign * halfview + self.ptycho_state["scan_center_px"] = float(cp.mean(px[(batch_size_here//2):])) + scan_state = self.ptycho_state.get("scan_state") or {} + self.logger.info( "Auto-centring projection %d: center_py=%.6e m, center_px=%.6e m " "(theta=%.2f°)", diff --git a/pipeline/ptychography_setup.py b/pipeline/ptychography_setup.py index ff3eef3..6dd6e37 100644 --- a/pipeline/ptychography_setup.py +++ b/pipeline/ptychography_setup.py @@ -160,7 +160,7 @@ def configure_scan_geometry( f"Buffers are never reallocated at runtime (R-6)." ) - # Scan extent in microns with 20% padding (same formula as PtyREX streaming) + # Scan extent in microns with 20% padding (same formula as PtyREX streaming) ## increased it on 11/09/26 to 100% padding for cases with position overshoots N = [ ((npoints_v - 1) * step_size_v) * 2, ((npoints_h - 1) * step_size_h) * 2, diff --git a/pipeline/vis/vis_stxm.py b/pipeline/vis/vis_stxm.py index 27c9f6a..19f2e70 100644 --- a/pipeline/vis/vis_stxm.py +++ b/pipeline/vis/vis_stxm.py @@ -266,7 +266,8 @@ def build_combined_figure_reduced(): for ax, cmap, key in ptycho_axes_info: im = ax.imshow(placeholder, cmap=cmap, interpolation='nearest', aspect='equal') ax.invert_xaxis() - fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + #fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + fig.colorbar(im, ax=ax) #, fraction=0.046, pad=0.04) ptycho_ims[key] = im return fig, ax_stxm_outer, ax_stxm_inner, ptycho_ims From 90ab64679449338069fe807477a609e64d688e3d Mon Sep 17 00:00:00 2001 From: Oriol Roche i Morgo Date: Tue, 11 Aug 2026 19:40:57 +0100 Subject: [PATCH 25/33] move scan_state definition up in ptychography_ops, transform_positions --- pipeline/ptychography_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 4688282..e9621ae 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -398,13 +398,13 @@ def _transform_positions(self, positions_txyz): # Auto-centre: capture scan centre from first batch if self.ptycho_state["scan_center_py"] is None: + scan_state = self.ptycho_state.get("scan_state") or {} projection = int(scan_state.get("current_projection", 0)) halfview = self.ptycho_state["N"][0]/2/2 * 1e-6 * pty_model.scan.scale[0] sign = 1 if projection % 2 == 0 else -1 batch_size_here = py.shape[0] self.ptycho_state["scan_center_py"] = float(cp.mean(py)) + sign * halfview self.ptycho_state["scan_center_px"] = float(cp.mean(px[(batch_size_here//2):])) - scan_state = self.ptycho_state.get("scan_state") or {} self.logger.info( "Auto-centring projection %d: center_py=%.6e m, center_px=%.6e m " From 247c674c1c8c3aac02571fa9740b053ba6ab115c Mon Sep 17 00:00:00 2001 From: UriRoche Date: Tue, 21 Jul 2026 11:54:26 +0200 Subject: [PATCH 26/33] Testing pipeline --- pipeline/config_test.yaml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/pipeline/config_test.yaml b/pipeline/config_test.yaml index bdb338e..6dd78e0 100644 --- a/pipeline/config_test.yaml +++ b/pipeline/config_test.yaml @@ -5,8 +5,8 @@ scheduler: worker_threads: 4 image_src: - zmq_endpoint: "tcp://172.23.82.48:31001" # Production endpoint - #zmq_endpoint: "tcp://172.23.82.77:5555" # Local simulator + #zmq_endpoint: "tcp://172.23.82.48:31001" # Production endpoint + zmq_endpoint: "tcp://172.23.82.77:5555" # Local simulator receive_timeout_ms: 1000 # Increased timeout for testing batch_size: 100 # Larger batch for testing @@ -16,16 +16,16 @@ decompress_op: data_dtype: uint16 position_src: - # zmq_endpoint: "tcp://172.23.82.77:5556" # Local simulator - zmq_endpoint: "tcp://172.23.82.204:6666" # production endpoint + zmq_endpoint: "tcp://172.23.82.77:5556" # Local simulator + #zmq_endpoint: "tcp://172.23.82.204:6666" # production endpoint receive_timeout_ms: 1000 # Increased timeout for testing header_src: # Dedicated ZMQ SUB socket for live scan-geometry headers (PR2). A JSON # object with npoints_h/npoints_v/step_size_h/step_size_v/num_projections # reconfigures scan geometry on the fly and preempts an in-flight recon. - #zmq_endpoint: "tcp://172.23.82.77:5557" # Local simulator - zmq_endpoint: "tcp://172.23.82.204:6667" # production endpoint (placeholder) + zmq_endpoint: "tcp://172.23.82.77:5558" # Local simulator + #zmq_endpoint: "tcp://172.23.82.204:6667" # production endpoint (placeholder) receive_timeout_ms: 100 # short so the blocking recv doesn't hold a worker thread masking_op: From 4632a3a21e6fd80455424a63c0ed48a01ab7e6d2 Mon Sep 17 00:00:00 2001 From: UriRoche Date: Fri, 31 Jul 2026 12:25:08 +0200 Subject: [PATCH 27/33] config changes for testing --- pipeline/config_test.yaml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/pipeline/config_test.yaml b/pipeline/config_test.yaml index 6dd78e0..d1d50d4 100644 --- a/pipeline/config_test.yaml +++ b/pipeline/config_test.yaml @@ -8,7 +8,7 @@ image_src: #zmq_endpoint: "tcp://172.23.82.48:31001" # Production endpoint zmq_endpoint: "tcp://172.23.82.77:5555" # Local simulator receive_timeout_ms: 1000 # Increased timeout for testing - batch_size: 100 # Larger batch for testing + batch_size: 200 # Larger batch for testing decompress_op: # data_size: [192, 192] # Match selun test data (192x192 images) @@ -51,13 +51,13 @@ ptychography: # Scan grid now comes from the live header (header_src). These set the GPU # buffer capacity (allocated once, never realloced — R-6) and the default # geometry used at startup before any header arrives. - max_npoints_h: 100 # buffer capacity + startup default grid (horizontal) + max_npoints_h: 200 # buffer capacity + startup default grid (horizontal) max_npoints_v: 100 # buffer capacity + startup default grid (vertical) # capacity = 100*100 = 10000 frames (matches the prior committed 100x100 grid; # a header requesting more frames is rejected — bump these with GPU RAM in mind) - default_step_size_h: 0.25 # startup step size (microns), until a header arrives - default_step_size_v: 0.25 # startup step size (microns), until a header arrives - total_iterations: 25 + default_step_size_h: 0.2 # startup step size (microns), until a header arrives + default_step_size_v: 0.2 # startup step size (microns), until a header arrives + total_iterations: 35 post_stream_iterations: 1 housekeeping_interval: 1 publish_interval: 1 From 08bca3fe41f23a3b1e6fe3fc921c837ecf2cf4dc Mon Sep 17 00:00:00 2001 From: UriRoche Date: Mon, 3 Aug 2026 17:58:03 +0200 Subject: [PATCH 28/33] Testing larger batch size --- pipeline/config_test.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pipeline/config_test.yaml b/pipeline/config_test.yaml index d1d50d4..e8d54b5 100644 --- a/pipeline/config_test.yaml +++ b/pipeline/config_test.yaml @@ -8,7 +8,7 @@ image_src: #zmq_endpoint: "tcp://172.23.82.48:31001" # Production endpoint zmq_endpoint: "tcp://172.23.82.77:5555" # Local simulator receive_timeout_ms: 1000 # Increased timeout for testing - batch_size: 200 # Larger batch for testing + batch_size: 600 # Larger batch for testing decompress_op: # data_size: [192, 192] # Match selun test data (192x192 images) @@ -61,4 +61,4 @@ ptychography: post_stream_iterations: 1 housekeeping_interval: 1 publish_interval: 1 - reset_probe: true # false = carry previous scan's probe forward (warm start); true = full probe reset each scan + reset_probe: true # false = carry previous scan's probe forward (warm start); true = full probe reset each scan From 7fd9e3b98b23cf8e210747ae1bdb8a567ade6380 Mon Sep 17 00:00:00 2001 From: UriRoche Date: Tue, 11 Aug 2026 11:45:40 +0200 Subject: [PATCH 29/33] Added a -2 flag for flux for when we're using a probe to start the scan. In that case, we don't need to re-calculate the flux. Also changed config file for testing at the beamline. --- pipeline/config_test.yaml | 22 +++++++++++----------- pipeline/ptychography_ops.py | 27 ++++++++++++++++----------- 2 files changed, 27 insertions(+), 22 deletions(-) diff --git a/pipeline/config_test.yaml b/pipeline/config_test.yaml index e8d54b5..4ed8000 100644 --- a/pipeline/config_test.yaml +++ b/pipeline/config_test.yaml @@ -5,15 +5,15 @@ scheduler: worker_threads: 4 image_src: - #zmq_endpoint: "tcp://172.23.82.48:31001" # Production endpoint + #zmq_endpoint: "tcp://172.23.82.131:31001" # Production endpoint zmq_endpoint: "tcp://172.23.82.77:5555" # Local simulator receive_timeout_ms: 1000 # Increased timeout for testing - batch_size: 600 # Larger batch for testing + batch_size: 320 # Larger batch for testing decompress_op: - # data_size: [192, 192] # Match selun test data (192x192 images) - data_size: [514, 1030] # Match selun test data (192x192 images) - data_dtype: uint16 + data_size: [190, 190] # Match selun test data (192x192 images) + # data_size: [514, 1030] # Match selun test data (192x192 images) + data_dtype: uint32 position_src: zmq_endpoint: "tcp://172.23.82.77:5556" # Local simulator @@ -25,12 +25,12 @@ header_src: # object with npoints_h/npoints_v/step_size_h/step_size_v/num_projections # reconfigures scan geometry on the fly and preempts an in-flight recon. zmq_endpoint: "tcp://172.23.82.77:5558" # Local simulator - #zmq_endpoint: "tcp://172.23.82.204:6667" # production endpoint (placeholder) + #zmq_endpoint: "tcp://172.23.82.32:6667" # production endpoint (placeholder) receive_timeout_ms: 100 # short so the blocking recv doesn't hold a worker thread masking_op: - center_x: 524 # Adjusted for 192x192 images (center) - center_y: 287 # Adjusted for 192x192 images (center) + center_x: 95 # Adjusted for 192x192 images (center) + center_y: 95 # Adjusted for 192x192 images (center) radius: 15 # Adjusted proportionally sink_and_publish_op: @@ -51,12 +51,12 @@ ptychography: # Scan grid now comes from the live header (header_src). These set the GPU # buffer capacity (allocated once, never realloced — R-6) and the default # geometry used at startup before any header arrives. - max_npoints_h: 200 # buffer capacity + startup default grid (horizontal) - max_npoints_v: 100 # buffer capacity + startup default grid (vertical) + max_npoints_h: 150 # buffer capacity + startup default grid (horizontal) + max_npoints_v: 150 # buffer capacity + startup default grid (vertical) # capacity = 100*100 = 10000 frames (matches the prior committed 100x100 grid; # a header requesting more frames is rejected — bump these with GPU RAM in mind) default_step_size_h: 0.2 # startup step size (microns), until a header arrives - default_step_size_v: 0.2 # startup step size (microns), until a header arrives + default_step_size_v: 0.15 # startup step size (microns), until a header arrives total_iterations: 35 post_stream_iterations: 1 housekeeping_interval: 1 diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index e9621ae..142e9d1 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -758,17 +758,22 @@ def compute(self, op_input, op_output, context): if self.current_iteration == 0: raw_cpu = cp.asnumpy(self.ptycho_state["raw_gpu"][r][:n_filled]) dp = pty_data.dp - if pty_model.source.flux < 0: - pty_model.source.flux = float(np.sum( - np.sum(raw_cpu, 0)[dp == 1] - ) / raw_cpu.shape[0]) - self.logger.info("Computed flux = %.2f from %d frames", pty_model.source.flux, n_filled) - for trial_idx in range(pty_model.scan.tris_n): - pty_model.probe.array_states[:, :, :, :, trial_idx, :, :] = setPower( - pty_model.probe.array_states[:, :, :, :, trial_idx, :, :], - pty_model.source.flux, - ) - self.logger.info("Probe power normalized to flux") + + if not pty_model.source.flux == -2: + + if pty_model.source.flux < 0: + pty_model.source.flux = float(np.sum( + np.sum(raw_cpu, 0)[dp == 1] + ) / raw_cpu.shape[0]) + self.logger.info("Computed flux = %.2f from %d frames", pty_model.source.flux, n_filled) + + + for trial_idx in range(pty_model.scan.tris_n): + pty_model.probe.array_states[:, :, :, :, trial_idx, :, :] = setPower( + pty_model.probe.array_states[:, :, :, :, trial_idx, :, :], + pty_model.source.flux, + ) + self.logger.info("Probe power normalized to flux") pty_params.current_iteration = cp.asarray(min( self.current_iteration, self.total_iterations - 1 From e10adf689adde7a5ac07762950074e152a3bfa5e Mon Sep 17 00:00:00 2001 From: UriRoche Date: Tue, 11 Aug 2026 20:38:40 +0200 Subject: [PATCH 30/33] making tests --- pipeline/config_test.yaml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/pipeline/config_test.yaml b/pipeline/config_test.yaml index 4ed8000..9c4e8d6 100644 --- a/pipeline/config_test.yaml +++ b/pipeline/config_test.yaml @@ -5,8 +5,8 @@ scheduler: worker_threads: 4 image_src: - #zmq_endpoint: "tcp://172.23.82.131:31001" # Production endpoint - zmq_endpoint: "tcp://172.23.82.77:5555" # Local simulator + zmq_endpoint: "tcp://172.23.82.131:31001" # Production endpoint + #zmq_endpoint: "tcp://172.23.82.77:5555" # Local simulator receive_timeout_ms: 1000 # Increased timeout for testing batch_size: 320 # Larger batch for testing @@ -16,16 +16,16 @@ decompress_op: data_dtype: uint32 position_src: - zmq_endpoint: "tcp://172.23.82.77:5556" # Local simulator - #zmq_endpoint: "tcp://172.23.82.204:6666" # production endpoint + #zmq_endpoint: "tcp://172.23.82.77:5556" # Local simulator + zmq_endpoint: "tcp://172.23.82.204:6666" # production endpoint receive_timeout_ms: 1000 # Increased timeout for testing header_src: # Dedicated ZMQ SUB socket for live scan-geometry headers (PR2). A JSON # object with npoints_h/npoints_v/step_size_h/step_size_v/num_projections # reconfigures scan geometry on the fly and preempts an in-flight recon. - zmq_endpoint: "tcp://172.23.82.77:5558" # Local simulator - #zmq_endpoint: "tcp://172.23.82.32:6667" # production endpoint (placeholder) + #zmq_endpoint: "tcp://172.23.82.77:5558" # Local simulator + zmq_endpoint: "tcp://172.23.82.32:6667" # production endpoint (placeholder) receive_timeout_ms: 100 # short so the blocking recv doesn't hold a worker thread masking_op: From 18c308e2e5848e6bbed9a4a319174ab3ee153426 Mon Sep 17 00:00:00 2001 From: UriRoche Date: Thu, 13 Aug 2026 13:21:40 +0200 Subject: [PATCH 31/33] Some offline tests with Ramya and Denis + adding selun_mask_unbinned mask file. --- pipeline/config_test.yaml | 26 +++++++++++++------------- selun_mask_unbinned.h5 | Bin 0 -> 38532 bytes 2 files changed, 13 insertions(+), 13 deletions(-) create mode 100644 selun_mask_unbinned.h5 diff --git a/pipeline/config_test.yaml b/pipeline/config_test.yaml index 9c4e8d6..573f018 100644 --- a/pipeline/config_test.yaml +++ b/pipeline/config_test.yaml @@ -5,32 +5,32 @@ scheduler: worker_threads: 4 image_src: - zmq_endpoint: "tcp://172.23.82.131:31001" # Production endpoint - #zmq_endpoint: "tcp://172.23.82.77:5555" # Local simulator + #zmq_endpoint: "tcp://172.23.82.131:31001" # Production endpoint + zmq_endpoint: "tcp://172.23.82.77:5555" # Local simulator receive_timeout_ms: 1000 # Increased timeout for testing - batch_size: 320 # Larger batch for testing + batch_size: 200 # Larger batch for testing decompress_op: - data_size: [190, 190] # Match selun test data (192x192 images) - # data_size: [514, 1030] # Match selun test data (192x192 images) - data_dtype: uint32 + #data_size: [190, 190] # Match selun test data (192x192 images) + data_size: [514, 1030] # Match selun test data (192x192 images) + data_dtype: uint16 position_src: - #zmq_endpoint: "tcp://172.23.82.77:5556" # Local simulator - zmq_endpoint: "tcp://172.23.82.204:6666" # production endpoint + zmq_endpoint: "tcp://172.23.82.77:5556" # Local simulator + #zmq_endpoint: "tcp://172.23.82.204:6666" # production endpoint receive_timeout_ms: 1000 # Increased timeout for testing header_src: # Dedicated ZMQ SUB socket for live scan-geometry headers (PR2). A JSON # object with npoints_h/npoints_v/step_size_h/step_size_v/num_projections # reconfigures scan geometry on the fly and preempts an in-flight recon. - #zmq_endpoint: "tcp://172.23.82.77:5558" # Local simulator - zmq_endpoint: "tcp://172.23.82.32:6667" # production endpoint (placeholder) + zmq_endpoint: "tcp://172.23.82.77:5558" # Local simulator + #zmq_endpoint: "tcp://172.23.82.32:6667" # production endpoint (placeholder) receive_timeout_ms: 100 # short so the blocking recv doesn't hold a worker thread masking_op: - center_x: 95 # Adjusted for 192x192 images (center) - center_y: 95 # Adjusted for 192x192 images (center) + center_x: 520 # Adjusted for 192x192 images (center) + center_y: 273 # Adjusted for 192x192 images (center) radius: 15 # Adjusted proportionally sink_and_publish_op: @@ -45,7 +45,7 @@ sink_and_publish_op: ptychography: enabled: true - ptyrex_config: "/workdir/PtyREX/config_409907.json" + ptyrex_config: "/workdir/PtyREX/config_410183.json" scan_ID: [1, 1, 1] ID: [1, 1, 1] # Scan grid now comes from the live header (header_src). These set the GPU diff --git a/selun_mask_unbinned.h5 b/selun_mask_unbinned.h5 new file mode 100644 index 0000000000000000000000000000000000000000..684c51875c5bf033c321c1eec943944c85e89f8e GIT binary patch literal 38532 zcmeI)u}Z^090%Y_+F~pSx`~rxAHh{!WwmuE}M_b*!Zirjcgm&1THOuZZZ;vnOOgS#suX3Zxb&~6A zof=h$pAdHQ`JFuee7}pc!5~a~mhnr?w}p>RYBlAA`mY3rgTvGKz?b={OG@8!&I-o$ z=wTGwC6B{$Vad8Ezr$7<@FwT-xBMRD;L7ik3$f<6Ez0jC=chd0-RS-{w*M->=f@}g znBkJU{hZtH@#k_}=Ar%+s6r>LNrtOJZ`Prz(}}m-eNk;q-+QrV{M|}-q!sB*S0*0U zbsDJ2Z^~9R-b_QbpHnUSacQiy&gZJPvC`@IasmVh5FkK+009C72oNAZfB*pk1PBly zK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF z5FkK+009C72oNAZfB*pk1PBlyK!Cs+38a6kUZV+{Dggon2oNAZfB*pk1PBlyK!5-N s0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0)Hs*1=E;IApigX literal 0 HcmV?d00001 From f2e0f289aa116e7dda7e340e10c11688a575b4b9 Mon Sep 17 00:00:00 2001 From: UriRoche Date: Sat, 29 Aug 2026 13:44:47 +0200 Subject: [PATCH 32/33] Changed ptychography_ops to avoid shadows from previous projections appearing in the next projection. The main solution was initialising the pty_model.obj.array_global_kernel with every new projection. Also, simplified the scan_center calculation to be 0 for y and mean(px) for x -- this is not a long term solution but it'll do for now. --- pipeline/ptychography_ops.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 142e9d1..b345374 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -166,6 +166,9 @@ def _try_flip(self): self.ptycho_state["buf_free"][other] = False self.ptycho_state["write_idx"] = other self.ptycho_state["filled_until"][other] = 0 + self.ptycho_state["raw_gpu"][other][:] = 0 + self.ptycho_state["positions_full"][other][:] = 0 + self.ptycho_state["tilts_full"][other][:] = 0 # Per-projection auto-centering: the first batch written into the # new buffer must derive a fresh center for that projection. self.ptycho_state["scan_center_py"] = None @@ -401,10 +404,15 @@ def _transform_positions(self, positions_txyz): scan_state = self.ptycho_state.get("scan_state") or {} projection = int(scan_state.get("current_projection", 0)) halfview = self.ptycho_state["N"][0]/2/2 * 1e-6 * pty_model.scan.scale[0] - sign = 1 if projection % 2 == 0 else -1 + + pyi = cp.mean(py[0:5]) + pyf = cp.mean(py[-5:-1]) + #sign = 1 if projection % 2 == 0 else -1 + sign = 1 if pyi < pyf else -1 + batch_size_here = py.shape[0] - self.ptycho_state["scan_center_py"] = float(cp.mean(py)) + sign * halfview - self.ptycho_state["scan_center_px"] = float(cp.mean(px[(batch_size_here//2):])) + self.ptycho_state["scan_center_py"] = -0.0 #float(cp.mean(py)) + sign * halfview + self.ptycho_state["scan_center_px"] = float(cp.mean(px)) #float(cp.mean(px[(batch_size_here//2):])) self.logger.info( "Auto-centring projection %d: center_py=%.6e m, center_px=%.6e m " @@ -508,6 +516,8 @@ def _perform_advance(self): pty_model = self.ptycho_state["pty_model"] pty_model.obj.array_global[:] = self._obj_initial pty_model.obj.array_global_old[:] = self._obj_initial + pty_model.obj.array_global_kernel[:] = self._obj_kernel_initial + pty_model.obj.array_global_kernel_old[:] = self._obj_kernel_initial if self.reset_probe: pty_model.probe.array_states[:] = self._probe_initial pty_model.source.flux = self._flux_initial @@ -522,6 +532,10 @@ def _flip_read(self, scan_state): with self.lock: r = self.ptycho_state["read_idx"] self.ptycho_state["buf_free"][r] = True # accumulator may reuse it + self.ptycho_state["raw_gpu"][r][:] = 0 + self.ptycho_state["positions_full"][r][:] = 0 + self.ptycho_state["tilts_full"][r][:] = 0 + self.ptycho_state["filled_until"][r] = 0 self.ptycho_state["read_idx"] = (r + 1) % nbuf scan_state["current_projection"] = int( scan_state.get("current_projection", 0) @@ -998,6 +1012,7 @@ def _init_gpu(self): # can reset the object back to its initial guess on flush. The probe # snapshot is only used when reset_probe is enabled. self._obj_initial = pty_model.obj.array_global.copy() + self._obj_kernel_initial = pty_model.obj.array_global_kernel.copy() self._probe_initial = pty_model.probe.array_states.copy() self._flux_initial = pty_model.source.flux # may be < 0 (auto) From 2f3842428808b3023b72bfbc199e9f93fc759965 Mon Sep 17 00:00:00 2001 From: UriRoche Date: Sat, 29 Aug 2026 13:45:17 +0200 Subject: [PATCH 33/33] Testing from holoscan and simulated stream. Changing config_test.yaml. --- pipeline/config_test.yaml | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/pipeline/config_test.yaml b/pipeline/config_test.yaml index 573f018..ee525b1 100644 --- a/pipeline/config_test.yaml +++ b/pipeline/config_test.yaml @@ -8,12 +8,12 @@ image_src: #zmq_endpoint: "tcp://172.23.82.131:31001" # Production endpoint zmq_endpoint: "tcp://172.23.82.77:5555" # Local simulator receive_timeout_ms: 1000 # Increased timeout for testing - batch_size: 200 # Larger batch for testing + batch_size: 180 # Larger batch for testing decompress_op: - #data_size: [190, 190] # Match selun test data (192x192 images) - data_size: [514, 1030] # Match selun test data (192x192 images) - data_dtype: uint16 + data_size: [190, 190] # Match selun test data (192x192 images) + #data_size: [514, 1030] # Match selun test data (192x192 images) + data_dtype: uint32 position_src: zmq_endpoint: "tcp://172.23.82.77:5556" # Local simulator @@ -29,9 +29,9 @@ header_src: receive_timeout_ms: 100 # short so the blocking recv doesn't hold a worker thread masking_op: - center_x: 520 # Adjusted for 192x192 images (center) - center_y: 273 # Adjusted for 192x192 images (center) - radius: 15 # Adjusted proportionally + center_x: 94 # Adjusted for 192x192 images (center) + center_y: 94 # Adjusted for 192x192 images (center) + radius: 20 # Adjusted proportionally sink_and_publish_op: publish_tensors: ["positions", "positions_ids", "inner", "outer", "intensity_ids"] @@ -45,20 +45,20 @@ sink_and_publish_op: ptychography: enabled: true - ptyrex_config: "/workdir/PtyREX/config_410183.json" + ptyrex_config: "/workdir/PtyREX/config_414287.json" scan_ID: [1, 1, 1] ID: [1, 1, 1] # Scan grid now comes from the live header (header_src). These set the GPU # buffer capacity (allocated once, never realloced — R-6) and the default # geometry used at startup before any header arrives. - max_npoints_h: 150 # buffer capacity + startup default grid (horizontal) - max_npoints_v: 150 # buffer capacity + startup default grid (vertical) + max_npoints_h: 200 # buffer capacity + startup default grid (horizontal) + max_npoints_v: 100 # buffer capacity + startup default grid (vertical) # capacity = 100*100 = 10000 frames (matches the prior committed 100x100 grid; # a header requesting more frames is rejected — bump these with GPU RAM in mind) default_step_size_h: 0.2 # startup step size (microns), until a header arrives - default_step_size_v: 0.15 # startup step size (microns), until a header arrives + default_step_size_v: 0.2 # startup step size (microns), until a header arrives total_iterations: 35 post_stream_iterations: 1 housekeeping_interval: 1 publish_interval: 1 - reset_probe: true # false = carry previous scan's probe forward (warm start); true = full probe reset each scan + reset_probe: false # false = carry previous scan's probe forward (warm start); true = full probe reset each scan