dex-python-sdk 1.11.0__tar.gz → 1.14.0__tar.gz

{dex_python_sdk-1.11.0 → dex_python_sdk-1.14.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dex-python-sdk
-Version: 1.11.0
+Version: 1.14.0
 Requires-Dist: grpcio>=1.83.0
 Requires-Dist: grpcio-status>=1.83.0
 Requires-Dist: protobuf>=7.35.1
@@ -100,6 +100,14 @@ max_concurrent_tasks)` streams because busy streams are replaced with idle
 streams and retained after tasks finish. Size these options for the server's
 per-Worker stream cap (256 by default).
 
+The Worker rejects an unregistered Step or RPC by name. Dex retries that task
+on another pull Worker without consuming a Step attempt, including when the
+Flow was sticky to the rejecting Worker. If no capable Worker accepts it within
+the server's dispatch wait (10 seconds by default, or the local Step attempt),
+the Step follows its failure and retry policy; an RPC caller receives the named
+error. This reassignment requires a Dex Server that supports it. Against older
+Servers, the Worker returns the not-registered error as before.
+
 The generated `app_worker_id` uses the first nonempty `DEX_WORKER_INSTANCE`,
 `POD_NAME`, `HOSTNAME`, or system hostname, plus eight random base32 characters.
 An override must uniquely identify one Worker process and change after restart.
@@ -114,7 +122,7 @@ an active task to finish, then waits up to `stop_drain_timeout` for the server
 to close the stream after the final frame. An idle stream reconnects
 immediately. An older server without pull support produces one explicit error
 log and is retried after one minute. `start()` reports startup
-synchronization errors, while later stream errors are logged and retried.
+synchronization errors, while later stream errors are logged and retried by default.
 `stop()` drains for at most
 `stop_drain_timeout`, which defaults to 30 seconds. Afterward it requests
 cancellation and waits for handlers to return. Synchronous handlers should
@@ -122,6 +130,31 @@ check `Context.is_cancellation_requested()` during long operations. Streams
 that sent a final frame wait for server EOF until the remaining Stop drain
 deadline before the channel closes.
 
+Call `worker.wait_pull_registered(timeout=30)` after starting a synchronous
+Worker on another thread. For AsyncWorker, start it in another task and use
+`await worker.wait_pull_registered(timeout=30)`. The timeout is in seconds;
+`None` waits indefinitely. Each waits for every Flow type queue and the app
+Worker ID queue to register at least once. Timeout raises `TimeoutError`;
+async cancellation cancels only that wait. Startup failures reach waiters.
+A permanent rejection before readiness raises `PullRegistrationError`, whose
+`queue` and `error` identify the rejected queue and original transport error.
+Listening mode and stop before readiness raise `RuntimeError`.
+
+Use `worker.pull_registration_state()` for subsequent readiness probes. Its
+immutable `queues` expose `queue`, `has_registered_stream`, and
+`last_permanent_error`. After initial registration, `is_ready` stays true across
+transient failures, reconnects, and stop. It becomes false only when a queue
+has a permanent error and no registered stream. Re-registration clears the
+error. The startup wait keeps succeeding after initial registration, including
+after stop. Multiple callers may wait independently before or after start.
+
+`PullWorkerOptions(stop_on_permanent_stream_rejection=True)` makes the first
+permanent rejection leaving a queue without a registered stream drain within
+`stop_drain_timeout`, then raise `PullRegistrationError` from `start()`. This
+applies before and after readiness. Another registered stream for that queue
+prevents shutdown. The default is False and retries indefinitely. Always call
+`stop()` or `close()` afterward; both remain safe after `start()` returns.
+
 ## Flow state I/O
 
 Applications read and write Flow state through typed RPCs. Client and AsyncClient

{dex_python_sdk-1.11.0 → dex_python_sdk-1.14.0}/README.md

@@ -84,6 +84,14 @@ max_concurrent_tasks)` streams because busy streams are replaced with idle
 streams and retained after tasks finish. Size these options for the server's
 per-Worker stream cap (256 by default).
 
+The Worker rejects an unregistered Step or RPC by name. Dex retries that task
+on another pull Worker without consuming a Step attempt, including when the
+Flow was sticky to the rejecting Worker. If no capable Worker accepts it within
+the server's dispatch wait (10 seconds by default, or the local Step attempt),
+the Step follows its failure and retry policy; an RPC caller receives the named
+error. This reassignment requires a Dex Server that supports it. Against older
+Servers, the Worker returns the not-registered error as before.
+
 The generated `app_worker_id` uses the first nonempty `DEX_WORKER_INSTANCE`,
 `POD_NAME`, `HOSTNAME`, or system hostname, plus eight random base32 characters.
 An override must uniquely identify one Worker process and change after restart.
@@ -98,7 +106,7 @@ an active task to finish, then waits up to `stop_drain_timeout` for the server
 to close the stream after the final frame. An idle stream reconnects
 immediately. An older server without pull support produces one explicit error
 log and is retried after one minute. `start()` reports startup
-synchronization errors, while later stream errors are logged and retried.
+synchronization errors, while later stream errors are logged and retried by default.
 `stop()` drains for at most
 `stop_drain_timeout`, which defaults to 30 seconds. Afterward it requests
 cancellation and waits for handlers to return. Synchronous handlers should
@@ -106,6 +114,31 @@ check `Context.is_cancellation_requested()` during long operations. Streams
 that sent a final frame wait for server EOF until the remaining Stop drain
 deadline before the channel closes.
 
+Call `worker.wait_pull_registered(timeout=30)` after starting a synchronous
+Worker on another thread. For AsyncWorker, start it in another task and use
+`await worker.wait_pull_registered(timeout=30)`. The timeout is in seconds;
+`None` waits indefinitely. Each waits for every Flow type queue and the app
+Worker ID queue to register at least once. Timeout raises `TimeoutError`;
+async cancellation cancels only that wait. Startup failures reach waiters.
+A permanent rejection before readiness raises `PullRegistrationError`, whose
+`queue` and `error` identify the rejected queue and original transport error.
+Listening mode and stop before readiness raise `RuntimeError`.
+
+Use `worker.pull_registration_state()` for subsequent readiness probes. Its
+immutable `queues` expose `queue`, `has_registered_stream`, and
+`last_permanent_error`. After initial registration, `is_ready` stays true across
+transient failures, reconnects, and stop. It becomes false only when a queue
+has a permanent error and no registered stream. Re-registration clears the
+error. The startup wait keeps succeeding after initial registration, including
+after stop. Multiple callers may wait independently before or after start.
+
+`PullWorkerOptions(stop_on_permanent_stream_rejection=True)` makes the first
+permanent rejection leaving a queue without a registered stream drain within
+`stop_drain_timeout`, then raise `PullRegistrationError` from `start()`. This
+applies before and after readiness. Another registered stream for that queue
+prevents shutdown. The default is False and retries indefinitely. Always call
+`stop()` or `close()` afterward; both remain safe after `start()` returns.
+
 ## Flow state I/O
 
 Applications read and write Flow state through typed RPCs. Client and AsyncClient

{dex_python_sdk-1.11.0 → dex_python_sdk-1.14.0}/dex/__init__.py

@@ -113,6 +113,11 @@ from dex.timer import Timer
 from dex.wait import Wait
 from dex.wait_options import WaitForAttributeOptions, WaitForStepCompletionOptions
 from dex.worker import Worker
+from dex.pull_registration import (
+    PullRegistrationError,
+    PullRegistrationState,
+    PullQueueRegistrationState,
+)
 from dex.worker_options import PullWorkerOptions, WorkerOptions, WorkerTarget
 
 __all__ = [
@@ -214,6 +219,9 @@ __all__ = [
     "WorkerOptions",
     "WorkerTarget",
     "PullWorkerOptions",
+    "PullRegistrationError",
+    "PullRegistrationState",
+    "PullQueueRegistrationState",
     "WaitForAttributeOptions",
     "WaitForStepCompletionOptions",
     "dead_end",

{dex_python_sdk-1.11.0 → dex_python_sdk-1.14.0}/dex/_pull_worker.py

@@ -28,6 +28,7 @@ from dex._server_protocol import sdk_version
 from dex.dexpb import dex_pb2 as pb
 from dex.dexpb import dex_pb2_grpc
 from dex.flow import Registry
+from dex.pull_registration import PullRegistrationError, _PullRegistration
 from dex.worker_options import PullWorkerOptions
 
 _LOGGER = logging.getLogger(__name__)
@@ -82,9 +83,16 @@ class PullWorkerRuntime:
         app_worker_id: str | None = None,
         credential_runner: AsyncCredentialRunner | None = None,
     ) -> None:
+        self._registry = registry
         self.app_worker_id = app_worker_id or resolve_app_worker_id(registry, options)
         self.flow_types = sorted(registry._registered_flows)
         self.queues = (self.app_worker_id, *self.flow_types)
+        self.registration = _PullRegistration(self.queues)
+        self._stop_on_permanent_stream_rejection = (
+            options.stop_on_permanent_stream_rejection
+        )
+        self._permanent_error: PullRegistrationError | None = None
+        self._rejection_stop_task: asyncio.Task[None] | None = None
         self.max_concurrent_tasks = options.max_concurrent_tasks or os.cpu_count() or 1
         self.idle_streams_per_queue = options.idle_streams_per_queue or 1
         self.stop_drain_timeout = options.stop_drain_timeout.total_seconds() or 30
@@ -103,6 +111,26 @@ class PullWorkerRuntime:
         self._busy_tasks = 0
         self._unsupported_logged = False
 
+    def unknown_task_reason(
+        self, task: pb.PollWorkerTask
+    ) -> pb.PollWorkerTaskRejectReason | None:
+        kind = task.WhichOneof("request")
+        if kind == "rpc":
+            flow = self._registry._registered_flows.get(task.rpc.flow_type)
+            if flow is None or task.rpc.rpc_name not in flow.rpcs:
+                return pb.POLL_WORKER_TASK_REJECT_REASON_UNKNOWN_RPC
+        elif kind in ("wait_for", "execute"):
+            request = task.wait_for if kind == "wait_for" else task.execute
+            flow = self._registry._registered_flows.get(request.flow_type)
+            if flow is None:
+                return pb.POLL_WORKER_TASK_REJECT_REASON_UNKNOWN_STEP
+            if kind == "execute" and request.step_type == "sys:timeout_handler":
+                if not flow.has_timeout_handler:
+                    return pb.POLL_WORKER_TASK_REJECT_REASON_UNKNOWN_STEP
+            elif request.step_type not in flow.steps:
+                return pb.POLL_WORKER_TASK_REJECT_REASON_UNKNOWN_STEP
+        return None
+
     async def run(self) -> None:
         self._loop = asyncio.get_running_loop()
         self._loop_thread_id = threading.get_ident()
@@ -110,9 +138,10 @@ class PullWorkerRuntime:
         self._stopping = asyncio.Event()
         self._stop_completed = asyncio.Event()
         self._started.set()
-        channel = outbound_worker_task_channel(self._server_address)
-        stub = dex_pb2_grpc.WorkerTaskServiceStub(channel)  # type: ignore[no-untyped-call]
+        channel: grpc.aio.Channel | None = None
         try:
+            channel = outbound_worker_task_channel(self._server_address)
+            stub = dex_pb2_grpc.WorkerTaskServiceStub(channel)  # type: ignore[no-untyped-call]
             self._wake.set()
             while not self._stopping.is_set():
                 await self._wake.wait()
@@ -123,9 +152,14 @@ class PullWorkerRuntime:
                 *(stream.run_task for stream in tuple(self._streams)),
                 return_exceptions=True,
             )
+            if self._rejection_stop_task is not None:
+                await self._rejection_stop_task
+            if self._permanent_error is not None:
+                raise self._permanent_error
         finally:
             try:
-                await channel.close(grace=None)
+                if channel is not None:
+                    await channel.close(grace=None)
             finally:
                 self._loop_thread_id = None
 
@@ -171,6 +205,7 @@ class PullWorkerRuntime:
         self.signal()
 
     async def stop(self) -> None:
+        self.registration.record_stopping()
         if self._stopping is None:
             return
         if self._stopping.is_set():
@@ -242,6 +277,7 @@ class PullStream:
         self.handler_task: asyncio.Task[None] | None = None
         self.task_id = ""
         self.is_busy = False
+        self._is_registered = False
         self.is_draining = False
         self.is_half_closing = False
         self._active = threading.Event()
@@ -249,6 +285,7 @@ class PullStream:
         self._send_lock = asyncio.Lock()
         self._is_reconnecting = False
         self._deadline_handle: asyncio.TimerHandle | None = None
+        self._supports_unknown_name_rejection = False
 
     def on_done(self, completed: asyncio.Task[None]) -> None:
         self.runtime._stream_done(self, completed)
@@ -267,7 +304,24 @@ class PullStream:
                 _LOGGER.exception(
                     "Worker task stream failed", extra={"queue": self.queue}
                 )
-            if self.runtime.is_stopping:
+            finally:
+                code = error.code() if error is not None else None
+                rejection = self.runtime.registration.record_stream_result(
+                    self.queue,
+                    self._is_registered,
+                    error if code in _PERMANENT_CODES else None,
+                )
+                self._is_registered = False
+                if (
+                    rejection is not None
+                    and self.runtime._stop_on_permanent_stream_rejection
+                    and self.runtime._permanent_error is None
+                ):
+                    self.runtime._permanent_error = rejection
+                    self.runtime._rejection_stop_task = asyncio.create_task(
+                        self.runtime.stop()
+                    )
+            if self.runtime.is_stopping or self.runtime._permanent_error is not None:
                 break
             if self._is_reconnecting:
                 self._is_reconnecting = False
@@ -301,6 +355,7 @@ class PullStream:
     async def _poll(self) -> bool:
         if self.runtime.is_stopping:
             return False
+        self._supports_unknown_name_rejection = False
         if self.runtime._credential_runner is None:
             self._call = self.stub.PollWorkerTask()
         else:
@@ -364,7 +419,12 @@ class PullStream:
                         return False
                     if kind != "registered":
                         raise RuntimeError("Worker task stream did not register")
+                    self._supports_unknown_name_rejection = (
+                        frame.registered.supports_unknown_name_rejection
+                    )
                     was_registered = True
+                    self._is_registered = True
+                    self.runtime.registration.record_registered_stream(self.queue)
                     continue
                 if kind == "task":
                     await self._start_task(frame.task)
@@ -443,6 +503,14 @@ class PullStream:
                     pass
 
     async def _start_task(self, task: pb.PollWorkerTask) -> None:
+        unknown_reason = (
+            self.runtime.unknown_task_reason(task)
+            if self._supports_unknown_name_rejection
+            else None
+        )
+        if unknown_reason is not None:
+            await self._reject_task(task, unknown_reason)
+            return
         if self.is_busy or self.is_draining or not self.runtime.accept_task():
             await self._reject_task(task)
             return
@@ -458,12 +526,16 @@ class PullStream:
             )
         self.runtime.signal()
 
-    async def _reject_task(self, task: pb.PollWorkerTask) -> None:
+    async def _reject_task(
+        self,
+        task: pb.PollWorkerTask,
+        reason: pb.PollWorkerTaskRejectReason = pb.POLL_WORKER_TASK_REJECT_REASON_NO_CAPACITY,
+    ) -> None:
         await self._send(
             pb.PollWorkerTaskRequest(
                 reject=pb.PollWorkerTaskReject(
                     task_id=task.task_id,
-                    reason=pb.POLL_WORKER_TASK_REJECT_REASON_NO_CAPACITY,
+                    reason=reason,
                 )
             )
         )

{dex_python_sdk-1.11.0 → dex_python_sdk-1.14.0}/dex/async_worker.py

@@ -30,6 +30,7 @@ from dex.blob_cache import BlobCache
 from dex.dexpb import dex_pb2 as pb
 from dex.dexpb import dex_pb2_grpc
 from dex.flow import Registry
+from dex.pull_registration import PullRegistrationState
 from dex.worker_options import WorkerOptions, WorkerTarget
 
 
@@ -100,6 +101,7 @@ class AsyncWorker:
             self._server_protocol,
         )
         self._server = None
+        self._has_start_called = False
         self._pull = None
         if self.options.pull is not None:
             self._pull = PullWorkerRuntime(
@@ -167,11 +169,77 @@ class AsyncWorker:
 
         ``start`` may be awaited exactly once. It contacts FlowService before binding
         and completes only after ``stop`` terminates the listener or pull streams.
-        Later pull stream failures are logged and retried.
+        Permanent pull stream rejections raise PullRegistrationError when
+        stop_on_permanent_stream_rejection is enabled; otherwise streams retry.
 
         Raises:
+            PullRegistrationError: A permanent stream rejection stops this Worker.
             RuntimeError: If lifecycle state, index synchronization, or binding fails.
         """
+        if self._has_start_called:
+            raise RuntimeError("Worker can only start once")
+        self._has_start_called = True
+        start_error: BaseException | None = None
+        try:
+            await self._start()
+        except BaseException as failure:
+            start_error = failure
+            raise
+        finally:
+            if self._pull is not None:
+                self._pull.registration.finish_start(start_error)
+                if self._state == "running":
+                    await self.stop()
+
+    async def wait_pull_registered(self, timeout: float | None = None) -> None:
+        """Wait until every pull queue has registered at least once.
+
+        Start the Worker separately before waiting, or start it concurrently.
+        Successful registration latches, including after stop. Use
+        pull_registration_state for subsequent readiness probes. A permanent
+        rejection before readiness raises PullRegistrationError only when no
+        other stream is registered for that queue. Startup failures wake waiters.
+        Multiple callers may wait independently, including before start.
+
+        Args:
+            timeout: Maximum seconds to wait; None waits indefinitely.
+
+        Raises:
+            PullRegistrationError: A queue is permanently rejected before readiness.
+            TimeoutError: The timeout expires.
+            asyncio.CancelledError: The calling task is canceled; the Worker keeps running.
+            RuntimeError: This is a listening Worker, or it stops before readiness.
+            BaseException: The original startup failure, if start fails.
+
+        Examples:
+            >>> start_task = asyncio.create_task(worker.start())
+
+            >>> await worker.wait_pull_registered(timeout=30)
+            >>> is_ready = worker.pull_registration_state().is_ready
+        """
+        if self._pull is None:
+            raise RuntimeError("wait_pull_registered requires a pull Worker")
+        await self._pull.registration.wait_async(timeout)
+
+    def pull_registration_state(self) -> PullRegistrationState:
+        """Return an immutable snapshot of pull queue registrations.
+
+        Transient failures, reconnects, and stop preserve initial readiness.
+        A permanent rejection makes readiness false only while that queue has no
+        registered stream. Re-registration clears its error. Calls are safe
+        concurrently with start, stop, and other waiters. See wait_pull_registered
+        for an example.
+
+        Returns:
+            Current per-queue state, or an empty nonready state in listening mode.
+        """
+        return (
+            PullRegistrationState()
+            if self._pull is None
+            else self._pull.registration.snapshot()
+        )
+
+    async def _start(self) -> None:
         if self._state != "created":
             raise RuntimeError(f"AsyncWorker cannot start from state {self._state}")
         try:
@@ -247,6 +315,8 @@ class AsyncWorker:
         Half-closed pull streams wait for server EOF until the remaining drain
         deadline before the channel closes.
         """
+        if self._pull is not None:
+            self._pull.registration.record_stopping()
         if self._state in ("stopped", "closed"):
             return
         self._state = "stopping"