From a422ff6026f43aa601b587c90ee9f6d52148ea9e Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Sun, 4 Oct 2026 03:18:33 -0400 Subject: [PATCH] Wait for admitted progress callbacks before returning terminal executions Closes #18 Refs OpenHCSDev/openhcs#573 and OpenHCSDev/openhcs#394 --- src/zmqruntime/execution/client.py | 71 ++++++-- src/zmqruntime/execution/progress_stream.py | 53 ++++-- src/zmqruntime/execution/server.py | 26 ++- src/zmqruntime/execution/wait_policy.py | 2 + src/zmqruntime/messages.py | 13 ++ tests/test_execution.py | 178 +++++++++++++++++++- tests/test_progress_stream.py | 3 + 7 files changed, 314 insertions(+), 32 deletions(-) diff --git a/src/zmqruntime/execution/client.py b/src/zmqruntime/execution/client.py index 6fafe6b..204d0e2 100644 --- a/src/zmqruntime/execution/client.py +++ b/src/zmqruntime/execution/client.py @@ -7,6 +7,7 @@ import threading import uuid from abc import ABC, abstractmethod +from dataclasses import replace from typing import Generic, TypeVar import zmq @@ -18,6 +19,7 @@ ) from zmqruntime.execution.responses import ( ExecutionSubmissionResponse, + ExecutionWaitResult, WireRequest, WireResponse, ) @@ -132,7 +134,20 @@ def wait_for_completion( known_server_process_is_alive=self.known_server_process_is_alive, owned_server_process_exit=self.owned_server_process_exit, ) - return waiter.wait(execution_id, policy) + result = waiter.wait(execution_id, policy) + record = ExecutionWaitResult.from_wire(result).execution + if ( + record is not None + and record.progress_sequence + and self._progress_registration is not None + and self._progress_stream is not None + ): + self._progress_stream.wait_for_delivery( + execution_id, + lambda: self._delivered_progress_through(execution_id, record.progress_sequence), + timeout=policy.status_timeout_ms / 1000, + ) + return result def execute( self, @@ -255,8 +270,8 @@ def disconnect(self): raise listener_error def _ensure_progress_subscription(self, *, timeout_ms: int = 5000) -> None: - self._start_progress_listener() if self._progress_registration is not None: + self._start_progress_listener() return response = self._send_control_request( { @@ -269,9 +284,15 @@ def _ensure_progress_subscription(self, *, timeout_ms: int = 5000) -> None: raise RuntimeError(f"Progress registration response missing status: {response}") if response[MessageFields.STATUS] != ResponseType.OK.value: raise RuntimeError(f"Progress registration failed: {response}") - self._progress_registration = CallbackSubscription( - self._unregister_progress, - ) + with self._progress_lock: + for execution_id, payload in response[MessageFields.PROGRESS_OBSERVATIONS].items(): + baseline = ExecutionProgressObservation.from_wire(payload) + previous = self._progress_by_execution_id.get(execution_id) + if previous is not None: + baseline = replace(baseline, delivery_error=previous.delivery_error) + self._progress_by_execution_id[execution_id] = baseline + self._progress_registration = CallbackSubscription(self._unregister_progress) + self._start_progress_listener() def _unregister_progress(self) -> bool: if not self.is_connected(): @@ -292,16 +313,44 @@ def enable_progress_stream(self) -> None: def _record_progress(self, data: dict) -> None: execution_id = data[MessageFields.EXECUTION_ID] + observation = ExecutionProgressObservation.from_wire( + { + "sequence": data[MessageFields.PROGRESS_SEQUENCE], + "event": data, + } + ) with self._progress_lock: current = self._progress_by_execution_id.get(execution_id) - observation = ( - ExecutionProgressObservation.first(data) - if current is None - else current.followed_by(data) - ) + if current is not None and observation.sequence <= current.sequence: + return # Events admitted before this subscription's baseline. + expected = 1 if current is None else current.sequence + 1 + delivery_error = None if current is None else current.delivery_error + if observation.sequence != expected: + delivery_error = ( + f"Progress for {execution_id} skipped sequence {expected}; " + f"received {observation.sequence}" + ) + observation = replace(observation, delivery_error=delivery_error) self._progress_by_execution_id[execution_id] = observation + if delivery_error is not None: + raise RuntimeError(delivery_error) if self.progress_callback is not None: - self.progress_callback(data) + try: + self.progress_callback(data) + except Exception as error: + with self._progress_lock: + self._progress_by_execution_id[execution_id] = replace( + observation, delivery_error=f"Progress callback failed: {error}" + ) + raise + + def _delivered_progress_through(self, execution_id: str, sequence: int) -> bool: + observation = self.progress_observation(execution_id) + if observation is None: + return False + if observation.delivery_error is not None: + raise RuntimeError(observation.delivery_error) + return observation.sequence >= sequence def progress_observation( self, diff --git a/src/zmqruntime/execution/progress_stream.py b/src/zmqruntime/execution/progress_stream.py index 7262f1d..5e8c125 100644 --- a/src/zmqruntime/execution/progress_stream.py +++ b/src/zmqruntime/execution/progress_stream.py @@ -5,7 +5,6 @@ import json import logging import threading -import time from collections.abc import Callable, Mapping from dataclasses import dataclass from types import MappingProxyType @@ -16,7 +15,7 @@ WireResponse, WireValue, # noqa: F401 -- resolves WireResponse's recursive annotation in this module ) -from zmqruntime.messages import validate_progress_payload +from zmqruntime.messages import MessageFields, validate_progress_payload logger = logging.getLogger(__name__) @@ -43,6 +42,7 @@ class ExecutionProgressObservation: sequence: int event: WireResponse + delivery_error: str | None = None def __post_init__(self) -> None: if self.sequence < 1: @@ -53,23 +53,13 @@ def __post_init__(self) -> None: _freeze_wire_value(self.event), ) - @classmethod - def first(cls, event: WireResponse) -> ExecutionProgressObservation: - """Create the first retained observation for one execution.""" - - return cls(sequence=1, event=event) - - def followed_by(self, event: WireResponse) -> ExecutionProgressObservation: - """Return the next immutable observation for the same execution.""" - - return type(self)(sequence=self.sequence + 1, event=event) - def as_wire(self) -> dict: """Return a detached JSON-compatible projection.""" return { type(self).sequence.__name__: self.sequence, type(self).event.__name__: _thaw_wire_value(self.event), + type(self).delivery_error.__name__: self.delivery_error, } @classmethod @@ -82,7 +72,9 @@ def from_wire(cls, payload: WireResponse) -> ExecutionProgressObservation: raise TypeError("Execution progress sequence must be an integer") if not isinstance(event, Mapping): raise TypeError("Execution progress event must be a mapping") - return cls(sequence=sequence, event=event) + return cls( + sequence=sequence, event=event, delivery_error=payload.get(cls.delivery_error.__name__) + ) class ProgressStreamSubscriber: @@ -98,6 +90,8 @@ def __init__( self._thread: threading.Thread | None = None self._stop_event = threading.Event() self._lifecycle_lock = threading.RLock() + self._delivery_condition = threading.Condition() + self._dispatching_execution_id: str | None = None def start(self) -> None: with self._lifecycle_lock: @@ -123,6 +117,23 @@ def stop(self, timeout: float = 2.0) -> None: if self._thread is thread: self._thread = None + def wait_for_delivery( + self, + execution_id: str, + delivered: Callable[[], bool], + timeout: float, + ) -> None: + """Wait for this execution's callback dispatch, not merely its socket read.""" + with self._delivery_condition: + complete = self._delivery_condition.wait_for( + lambda: self._dispatching_execution_id != execution_id and delivered(), + timeout=timeout, + ) + if not complete: + raise TimeoutError( + f"Progress delivery for {execution_id} did not complete within {timeout:.3f}s" + ) + def _listen_loop(self) -> None: logger.info("Progress listener loop started") message_count = 0 @@ -130,12 +141,13 @@ def _listen_loop(self) -> None: while not self._stop_event.is_set(): socket = self._socket_provider() if socket is None: - time.sleep(0.1) + self._stop_event.wait(0.1) + continue + if not socket.poll(timeout=50, flags=zmq.POLLIN): continue try: message = socket.recv_string(zmq.NOBLOCK) except zmq.Again: - time.sleep(0.05) continue if self._dispatch_message(message): @@ -155,7 +167,14 @@ def _dispatch_message(self, message: str) -> bool: try: data = json.loads(message) validate_progress_payload(data) - self._callback(data) + with self._delivery_condition: + self._dispatching_execution_id = data.get(MessageFields.EXECUTION_ID) + try: + self._callback(data) + finally: + with self._delivery_condition: + self._dispatching_execution_id = None + self._delivery_condition.notify_all() except Exception as error: logger.exception("Progress message dispatch failed: %s", error) return False diff --git a/src/zmqruntime/execution/server.py b/src/zmqruntime/execution/server.py index 7962375..6bf3c91 100644 --- a/src/zmqruntime/execution/server.py +++ b/src/zmqruntime/execution/server.py @@ -16,6 +16,7 @@ from zmqruntime.config import ZMQConfig from zmqruntime.execution.lifecycle import InMemoryExecutionLifecycleEngine +from zmqruntime.execution.progress_stream import ExecutionProgressObservation from zmqruntime.messages import ( CancelRequest, ControlMessageType, @@ -409,10 +410,20 @@ def _handle_register_progress(self, msg): request, error = self._validate_and_parse(msg, ProgressRegistrationRequest) if error: return error - self._progress_subscribers.add(request.client_id) + with self._progress_publish_lock: + self._progress_subscribers.add(request.client_id) + observations = { + execution_id: ExecutionProgressObservation( + sequence=record.progress_sequence, + event=record.progress_event, + ).as_wire() + for execution_id, record in self.active_executions.items() + if record.progress_event is not None + } return { MessageFields.STATUS: ResponseType.OK.value, MessageFields.MESSAGE: "Progress subscriber registered", + MessageFields.PROGRESS_OBSERVATIONS: observations, MessageFields.CLIENT_ID: request.client_id, MessageFields.PROGRESS_SUBSCRIBERS: len(self._progress_subscribers), } @@ -433,7 +444,18 @@ def send_progress_update(self, progress_update: dict) -> None: from zmqruntime.messages import validate_progress_payload validate_progress_payload(progress_update) - self.progress_queue.put(progress_update) + execution_id = progress_update[MessageFields.EXECUTION_ID] + with self._progress_publish_lock: + record = self.active_executions[execution_id] + observation = ExecutionProgressObservation( + sequence=record.progress_sequence + 1, + event=progress_update, + ) + event = observation.as_wire()["event"] + record.progress_sequence = observation.sequence + event[MessageFields.PROGRESS_SEQUENCE] = observation.sequence + record.progress_event = event + self.progress_queue.put(event) def _get_worker_info(self): try: diff --git a/src/zmqruntime/execution/wait_policy.py b/src/zmqruntime/execution/wait_policy.py index e07ce8d..d0eb9b5 100644 --- a/src/zmqruntime/execution/wait_policy.py +++ b/src/zmqruntime/execution/wait_policy.py @@ -75,6 +75,7 @@ def wait(self, execution_id: str, policy: WaitPolicy) -> dict[str, WireValue]: result = { MessageFields.STATUS: ExecutionStatus.FAILED.value, MessageFields.EXECUTION_ID: execution_id, + MessageFields.EXECUTION: status_response[MessageFields.EXECUTION], } if record.error is not None: result[MessageFields.MESSAGE] = record.error @@ -83,6 +84,7 @@ def wait(self, execution_id: str, policy: WaitPolicy) -> dict[str, WireValue]: return { MessageFields.STATUS: ExecutionStatus.CANCELLED.value, MessageFields.EXECUTION_ID: execution_id, + MessageFields.EXECUTION: status_response[MessageFields.EXECUTION], MessageFields.MESSAGE: "Execution was cancelled", } diff --git a/src/zmqruntime/messages.py b/src/zmqruntime/messages.py index 25a1d22..ec1b6ef 100644 --- a/src/zmqruntime/messages.py +++ b/src/zmqruntime/messages.py @@ -302,6 +302,8 @@ class MessageFields: METADATA = "metadata" CLIENT_ID = "client_id" PROGRESS_SUBSCRIBERS = "progress_subscribers" + PROGRESS_SEQUENCE = "progress_sequence" + PROGRESS_OBSERVATIONS = "progress_observations" PROCESS_IDENTITY = "process_identity" CONTROL_CAPABILITIES = "control_capabilities" APPLICATION = "application" @@ -796,6 +798,14 @@ class ExecutionRecord: results_summary: Optional[Dict[str, Any]] = None compile_only: bool = False metadata: Dict[str, Any] = field(default_factory=dict) + progress_sequence: int = 0 + progress_event: dict | None = field(default=None, repr=False) + + def __post_init__(self) -> None: + if isinstance(self.progress_sequence, bool) or not isinstance(self.progress_sequence, int): + raise TypeError("Execution progress sequence must be an integer") + if self.progress_sequence < 0: + raise ValueError("Execution progress sequence cannot be negative") def set_extra(self, key: str, value: Any) -> None: self.metadata[key] = value @@ -835,6 +845,7 @@ def _to_transport_value(value: Any) -> Any: MessageFields.END_TIME: _to_transport_value(self.end_time), MessageFields.ERROR: _to_transport_value(self.error), MessageFields.COMPILE_ONLY: self.compile_only, + MessageFields.PROGRESS_SEQUENCE: self.progress_sequence, } if self.traceback is not None: result[MessageFields.TRACEBACK] = _to_transport_value(self.traceback) @@ -855,6 +866,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "ExecutionRecord": MessageFields.TRACEBACK, MessageFields.RESULTS_SUMMARY, MessageFields.COMPILE_ONLY, + MessageFields.PROGRESS_SEQUENCE, } metadata = {k: v for k, v in data.items() if k not in known} return cls( @@ -869,6 +881,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "ExecutionRecord": results_summary=data.get(MessageFields.RESULTS_SUMMARY), compile_only=bool(data.get(MessageFields.COMPILE_ONLY, False)), metadata=metadata, + progress_sequence=data.get(MessageFields.PROGRESS_SEQUENCE, 0), ) diff --git a/tests/test_execution.py b/tests/test_execution.py index 3495430..4cd2b7d 100644 --- a/tests/test_execution.py +++ b/tests/test_execution.py @@ -567,7 +567,10 @@ def connect(self, timeout: float = 10.0): def _send_control_request(self, request, timeout_ms=5000): if request[MessageFields.TYPE] == ControlMessageType.REGISTER_PROGRESS.value: - return {MessageFields.STATUS: ResponseType.OK.value} + return { + MessageFields.STATUS: ResponseType.OK.value, + MessageFields.PROGRESS_OBSERVATIONS: {}, + } return request @@ -619,7 +622,10 @@ def _start_progress_listener(self): def _send_control_request(self, request, timeout_ms=5000): self.sent_requests.append((request, timeout_ms)) if request.get(MessageFields.TYPE) == ControlMessageType.REGISTER_PROGRESS.value: - return {MessageFields.STATUS: ResponseType.OK.value} + return { + MessageFields.STATUS: ResponseType.OK.value, + MessageFields.PROGRESS_OBSERVATIONS: {}, + } if request.get(MessageFields.TYPE) == ControlMessageType.UNREGISTER_PROGRESS.value: return {MessageFields.STATUS: ResponseType.OK.value} return request @@ -683,6 +689,9 @@ def test_execution_client_retains_immutable_progress_observations_per_execution( total=2, ).to_dict() + first[MessageFields.PROGRESS_SEQUENCE] = 1 + second[MessageFields.PROGRESS_SEQUENCE] = 2 + other[MessageFields.PROGRESS_SEQUENCE] = 1 client._record_progress(first) client._record_progress(second) client._record_progress(other) @@ -1566,3 +1575,168 @@ def __setattr__(self, name, value): assert published[0]["execution"]["start_time"] == 2.0 assert published[0]["execution"]["end_time"] == 5.0 assert published[0]["execution"]["status"] == ExecutionStatus.COMPLETE.value + + +@pytest.mark.parametrize( + "terminal_status", [ExecutionStatus.COMPLETE, ExecutionStatus.FAILED, ExecutionStatus.CANCELLED] +) +def test_terminal_status_waits_for_actual_progress_callback_tail(terminal_status): + from zmqruntime.execution.progress_stream import ProgressStreamSubscriber + from zmqruntime.messages import ExecutionRecord + from zmqruntime.subscription import CallbackSubscription + + server = DummyExecutionServer(port=5555) + record = ExecutionRecord("tail-1", "plate-1", None, terminal_status.value) + server.active_executions[record.execution_id] = record + for number in (1, 2): + server.send_progress_update( + TaskProgress( + task_id=record.execution_id, + phase="axis_completed" if number == 2 else "step_completed", + status="success", + percent=number * 50.0, + timestamp=float(number), + completed=number, + total=2, + ).to_dict() + ) + events = [server.progress_queue.get_nowait(), server.progress_queue.get_nowait()] + client = DummyExecutionClient() + callback_started = threading.Event() + release_callback = threading.Event() + status_observed = threading.Event() + finished = threading.Event() + received = [] + + def callback(event): + received.append(event["phase"]) + if event[MessageFields.PROGRESS_SEQUENCE] == 1: + callback_started.set() + assert release_callback.wait(1) + + client.progress_callback = callback + client._progress_stream = ProgressStreamSubscriber(lambda: None, client._record_progress) + client._progress_registration = CallbackSubscription(lambda: True) + + def poll(_execution_id, *, timeout_ms): + status_observed.set() + return { + MessageFields.STATUS: ResponseType.OK.value, + MessageFields.EXECUTION: record.to_dict(), + } + + client.poll_status = poll + results = [] + + def wait(): + results.append(client.wait_for_completion(record.execution_id, poll_interval=0)) + finished.set() + + dispatcher = threading.Thread( + target=lambda: [client._progress_stream._dispatch_message(json.dumps(e)) for e in events] + ) + waiter = threading.Thread(target=wait) + dispatcher.start() + assert callback_started.wait(1) + waiter.start() + assert status_observed.wait(1) + assert not finished.wait(0.02) + release_callback.set() + dispatcher.join(1) + waiter.join(1) + assert not dispatcher.is_alive() and not waiter.is_alive() + assert finished.is_set() + assert received == ["step_completed", "axis_completed"] + assert results[0][MessageFields.EXECUTION][MessageFields.PROGRESS_SEQUENCE] == 2 + + +def test_progress_registration_baseline_and_loss_have_explicit_delivery_laws(): + from zmqruntime.execution.progress_stream import ProgressStreamSubscriber + from zmqruntime.messages import ExecutionRecord + + server = DummyExecutionServer(port=5555) + record = ExecutionRecord("late-1", "plate-1", None, ExecutionStatus.COMPLETE.value) + server.active_executions[record.execution_id] = record + payload = TaskProgress( + task_id=record.execution_id, + phase="running", + status="running", + percent=50.0, + timestamp=1.0, + completed=1, + total=2, + ).to_dict() + server.send_progress_update(payload) + baseline = server._handle_register_progress( + { + MessageFields.TYPE: ControlMessageType.REGISTER_PROGRESS.value, + MessageFields.CLIENT_ID: "late-client", + } + ) + client = DummyExecutionClient() + client._send_control_request = lambda _request, timeout_ms=5000: baseline + client._start_progress_listener = lambda: None + client._ensure_progress_subscription() + assert client.progress_observation(record.execution_id).sequence == 1 + received = [] + client.progress_callback = received.append + client._record_progress(server.progress_queue.get_nowait()) + assert received == [] # A late subscriber does not fabricate historical callbacks. + server.send_progress_update(payload) + second = server.progress_queue.get_nowait() + client._record_progress(second) + assert len(received) == 1 + + client._progress_stream = ProgressStreamSubscriber(lambda: None, client._record_progress) + client.poll_status = lambda _execution_id, timeout_ms: { + MessageFields.STATUS: ResponseType.OK.value, + MessageFields.EXECUTION: record.to_dict(), + } + client.wait_for_completion(record.execution_id, poll_interval=0) + server.send_progress_update(payload) + with pytest.raises(TimeoutError, match="Progress delivery"): + client.wait_for_completion(record.execution_id, poll_interval=0, status_timeout_ms=10) + missing = server.progress_queue.get_nowait() + missing[MessageFields.PROGRESS_SEQUENCE] += 1 + assert not client._progress_stream._dispatch_message(json.dumps(missing)) + with pytest.raises(RuntimeError, match="skipped sequence 3"): + client.wait_for_completion(record.execution_id, poll_interval=0) + client._progress_registration = None + client._ensure_progress_subscription() + with pytest.raises(RuntimeError, match="skipped sequence 3"): + client.wait_for_completion(record.execution_id, poll_interval=0) + + +def test_terminal_progress_callback_failure_is_not_completed_delivery(): + from zmqruntime.execution.progress_stream import ProgressStreamSubscriber + from zmqruntime.messages import ExecutionRecord + from zmqruntime.subscription import CallbackSubscription + + client = DummyExecutionClient() + client._progress_registration = CallbackSubscription(lambda: True) + client._progress_stream = ProgressStreamSubscriber(lambda: None, client._record_progress) + + def fail(_event): + raise ValueError("observer rejected tail") + + client.progress_callback = fail + event = TaskProgress( + task_id="callback-1", + phase="axis_completed", + status="success", + percent=100.0, + timestamp=1.0, + completed=1, + total=1, + ).to_dict() + event[MessageFields.PROGRESS_SEQUENCE] = 1 + assert not client._progress_stream._dispatch_message(json.dumps(event)) + record = ExecutionRecord( + "callback-1", "plate-1", None, ExecutionStatus.COMPLETE.value, progress_sequence=1 + ) + client.poll_status = lambda _execution_id, timeout_ms: { + MessageFields.STATUS: ResponseType.OK.value, + MessageFields.EXECUTION: record.to_dict(), + } + with pytest.raises(RuntimeError, match="observer rejected tail"): + client.wait_for_completion(record.execution_id, poll_interval=0) diff --git a/tests/test_progress_stream.py b/tests/test_progress_stream.py index 8f3005a..78abbda 100644 --- a/tests/test_progress_stream.py +++ b/tests/test_progress_stream.py @@ -56,6 +56,9 @@ class OneMessageSocket: def __init__(self): self.sent = False + def poll(self, timeout, flags): + return not self.sent + def recv_string(self, _flags): if self.sent: import zmq