dex-python-sdk 1.12.0__tar.gz → 1.14.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (69) hide show
  1. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/PKG-INFO +34 -7
  2. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/README.md +33 -6
  3. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/__init__.py +8 -0
  4. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_pull_worker.py +83 -20
  5. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/async_worker.py +71 -1
  6. dex_python_sdk-1.14.1/dex/pull_registration.py +181 -0
  7. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/worker.py +71 -2
  8. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/worker_options.py +6 -0
  9. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/pyproject.toml +1 -1
  10. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/LEGACY_NOTICES.md +0 -0
  11. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/LICENSE +0 -0
  12. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_async_value_hydrator.py +0 -0
  13. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_async_worker_dispatcher.py +0 -0
  14. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_async_worker_service.py +0 -0
  15. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_credential_transport.py +0 -0
  16. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_grpc_errors.py +0 -0
  17. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_invocation_context.py +0 -0
  18. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_native.pyi +0 -0
  19. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_server_protocol.py +0 -0
  20. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_utils.py +0 -0
  21. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_value_hydrator.py +0 -0
  22. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_value_mapper.py +0 -0
  23. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_worker_dispatcher.py +0 -0
  24. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/_worker_service.py +0 -0
  25. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/async_client.py +0 -0
  26. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/attribute.py +0 -0
  27. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/attribute_match.py +0 -0
  28. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/blob_cache.py +0 -0
  29. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/channel.py +0 -0
  30. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/client.py +0 -0
  31. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/client_options.py +0 -0
  32. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/codec.py +0 -0
  33. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/condition.py +0 -0
  34. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/context.py +0 -0
  35. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/credential.py +0 -0
  36. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/dexpb/__init__.py +0 -0
  37. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/dexpb/dex_pb2.py +0 -0
  38. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/dexpb/dex_pb2.pyi +0 -0
  39. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/dexpb/dex_pb2_grpc.py +0 -0
  40. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/flow.py +0 -0
  41. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/flow_config.py +0 -0
  42. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/flow_info.py +0 -0
  43. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/flow_options.py +0 -0
  44. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/flow_result.py +0 -0
  45. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/py.typed +0 -0
  46. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/rpc_invoke_options.py +0 -0
  47. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/runtime_errors.py +0 -0
  48. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/step.py +0 -0
  49. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/step_execution.py +0 -0
  50. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/stream.py +0 -0
  51. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/subflow.py +0 -0
  52. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/timer.py +0 -0
  53. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/wait.py +0 -0
  54. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/dex/wait_options.py +0 -0
  55. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/Cargo.lock +0 -0
  56. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/Cargo.toml +0 -0
  57. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/crates/dex-blob-cache/Cargo.toml +0 -0
  58. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/crates/dex-blob-cache/LICENSE +0 -0
  59. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/crates/dex-blob-cache/src/config.rs +0 -0
  60. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/crates/dex-blob-cache/src/entry.rs +0 -0
  61. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/crates/dex-blob-cache/src/error.rs +0 -0
  62. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/crates/dex-blob-cache/src/format.rs +0 -0
  63. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/crates/dex-blob-cache/src/lib.rs +0 -0
  64. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/crates/dex-blob-cache/src/policy.rs +0 -0
  65. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/crates/dex-blob-cache/src/store.rs +0 -0
  66. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/crates/dex-blob-cache/tests/blob_cache_integration.rs +0 -0
  67. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/crates/dex-blob-cache-python/Cargo.toml +0 -0
  68. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/crates/dex-blob-cache-python/LICENSE +0 -0
  69. {dex_python_sdk-1.12.0 → dex_python_sdk-1.14.1}/sdk-rust/crates/dex-blob-cache-python/src/lib.rs +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dex-python-sdk
3
- Version: 1.12.0
3
+ Version: 1.14.1
4
4
  Requires-Dist: grpcio>=1.83.0
5
5
  Requires-Dist: grpcio-status>=1.83.0
6
6
  Requires-Dist: protobuf>=7.35.1
@@ -94,11 +94,13 @@ The Worker polls its sticky queue and every registered Flow type queue. Flow
94
94
  type names must be unique within a Dex server. `max_concurrent_tasks` defaults
95
95
  to the CPU count, and `idle_streams_per_queue` defaults to one. At capacity,
96
96
  the Worker rejects an assignment so Dex can reassign it. Each stream runs one
97
- task at a time; Step heartbeat and stream writes use the same stream. A Worker
98
- may hold up to `(1 + number of Flow types) * (idle_streams_per_queue +
99
- max_concurrent_tasks)` streams because busy streams are replaced with idle
100
- streams and retained after tasks finish. Size these options for the server's
101
- per-Worker stream cap (256 by default).
97
+ task at a time; Step heartbeat and stream writes use the same stream. After tasks finish, the Worker
98
+ gracefully closes surplus idle streams while preserving readiness. After load drops,
99
+ the steady-state bound is `max_concurrent_tasks + (1 + number of Flow types) *
100
+ idle_streams_per_queue` streams. The transient peak during a burst remains bounded
101
+ by the matching caps. Size these options for `matching.maxConcurrentTasksCap` per
102
+ app Worker ID per matching owner (256 by default) and `matching.maxStreams` per
103
+ Matching instance (100,000 by default).
102
104
 
103
105
  The Worker rejects an unregistered Step or RPC by name. Dex retries that task
104
106
  on another pull Worker without consuming a Step attempt, including when the
@@ -122,7 +124,7 @@ an active task to finish, then waits up to `stop_drain_timeout` for the server
122
124
  to close the stream after the final frame. An idle stream reconnects
123
125
  immediately. An older server without pull support produces one explicit error
124
126
  log and is retried after one minute. `start()` reports startup
125
- synchronization errors, while later stream errors are logged and retried.
127
+ synchronization errors, while later stream errors are logged and retried by default.
126
128
  `stop()` drains for at most
127
129
  `stop_drain_timeout`, which defaults to 30 seconds. Afterward it requests
128
130
  cancellation and waits for handlers to return. Synchronous handlers should
@@ -130,6 +132,31 @@ check `Context.is_cancellation_requested()` during long operations. Streams
130
132
  that sent a final frame wait for server EOF until the remaining Stop drain
131
133
  deadline before the channel closes.
132
134
 
135
+ Call `worker.wait_pull_registered(timeout=30)` after starting a synchronous
136
+ Worker on another thread. For AsyncWorker, start it in another task and use
137
+ `await worker.wait_pull_registered(timeout=30)`. The timeout is in seconds;
138
+ `None` waits indefinitely. Each waits for every Flow type queue and the app
139
+ Worker ID queue to register at least once. Timeout raises `TimeoutError`;
140
+ async cancellation cancels only that wait. Startup failures reach waiters.
141
+ A permanent rejection before readiness raises `PullRegistrationError`, whose
142
+ `queue` and `error` identify the rejected queue and original transport error.
143
+ Listening mode and stop before readiness raise `RuntimeError`.
144
+
145
+ Use `worker.pull_registration_state()` for subsequent readiness probes. Its
146
+ immutable `queues` expose `queue`, `has_registered_stream`, and
147
+ `last_permanent_error`. After initial registration, `is_ready` stays true across
148
+ transient failures, reconnects, and stop. It becomes false only when a queue
149
+ has a permanent error and no registered stream. Re-registration clears the
150
+ error. The startup wait keeps succeeding after initial registration, including
151
+ after stop. Multiple callers may wait independently before or after start.
152
+
153
+ `PullWorkerOptions(stop_on_permanent_stream_rejection=True)` makes the first
154
+ permanent rejection leaving a queue without a registered stream drain within
155
+ `stop_drain_timeout`, then raise `PullRegistrationError` from `start()`. This
156
+ applies before and after readiness. Another registered stream for that queue
157
+ prevents shutdown. The default is False and retries indefinitely. Always call
158
+ `stop()` or `close()` afterward; both remain safe after `start()` returns.
159
+
133
160
  ## Flow state I/O
134
161
 
135
162
  Applications read and write Flow state through typed RPCs. Client and AsyncClient
@@ -78,11 +78,13 @@ The Worker polls its sticky queue and every registered Flow type queue. Flow
78
78
  type names must be unique within a Dex server. `max_concurrent_tasks` defaults
79
79
  to the CPU count, and `idle_streams_per_queue` defaults to one. At capacity,
80
80
  the Worker rejects an assignment so Dex can reassign it. Each stream runs one
81
- task at a time; Step heartbeat and stream writes use the same stream. A Worker
82
- may hold up to `(1 + number of Flow types) * (idle_streams_per_queue +
83
- max_concurrent_tasks)` streams because busy streams are replaced with idle
84
- streams and retained after tasks finish. Size these options for the server's
85
- per-Worker stream cap (256 by default).
81
+ task at a time; Step heartbeat and stream writes use the same stream. After tasks finish, the Worker
82
+ gracefully closes surplus idle streams while preserving readiness. After load drops,
83
+ the steady-state bound is `max_concurrent_tasks + (1 + number of Flow types) *
84
+ idle_streams_per_queue` streams. The transient peak during a burst remains bounded
85
+ by the matching caps. Size these options for `matching.maxConcurrentTasksCap` per
86
+ app Worker ID per matching owner (256 by default) and `matching.maxStreams` per
87
+ Matching instance (100,000 by default).
86
88
 
87
89
  The Worker rejects an unregistered Step or RPC by name. Dex retries that task
88
90
  on another pull Worker without consuming a Step attempt, including when the
@@ -106,7 +108,7 @@ an active task to finish, then waits up to `stop_drain_timeout` for the server
106
108
  to close the stream after the final frame. An idle stream reconnects
107
109
  immediately. An older server without pull support produces one explicit error
108
110
  log and is retried after one minute. `start()` reports startup
109
- synchronization errors, while later stream errors are logged and retried.
111
+ synchronization errors, while later stream errors are logged and retried by default.
110
112
  `stop()` drains for at most
111
113
  `stop_drain_timeout`, which defaults to 30 seconds. Afterward it requests
112
114
  cancellation and waits for handlers to return. Synchronous handlers should
@@ -114,6 +116,31 @@ check `Context.is_cancellation_requested()` during long operations. Streams
114
116
  that sent a final frame wait for server EOF until the remaining Stop drain
115
117
  deadline before the channel closes.
116
118
 
119
+ Call `worker.wait_pull_registered(timeout=30)` after starting a synchronous
120
+ Worker on another thread. For AsyncWorker, start it in another task and use
121
+ `await worker.wait_pull_registered(timeout=30)`. The timeout is in seconds;
122
+ `None` waits indefinitely. Each waits for every Flow type queue and the app
123
+ Worker ID queue to register at least once. Timeout raises `TimeoutError`;
124
+ async cancellation cancels only that wait. Startup failures reach waiters.
125
+ A permanent rejection before readiness raises `PullRegistrationError`, whose
126
+ `queue` and `error` identify the rejected queue and original transport error.
127
+ Listening mode and stop before readiness raise `RuntimeError`.
128
+
129
+ Use `worker.pull_registration_state()` for subsequent readiness probes. Its
130
+ immutable `queues` expose `queue`, `has_registered_stream`, and
131
+ `last_permanent_error`. After initial registration, `is_ready` stays true across
132
+ transient failures, reconnects, and stop. It becomes false only when a queue
133
+ has a permanent error and no registered stream. Re-registration clears the
134
+ error. The startup wait keeps succeeding after initial registration, including
135
+ after stop. Multiple callers may wait independently before or after start.
136
+
137
+ `PullWorkerOptions(stop_on_permanent_stream_rejection=True)` makes the first
138
+ permanent rejection leaving a queue without a registered stream drain within
139
+ `stop_drain_timeout`, then raise `PullRegistrationError` from `start()`. This
140
+ applies before and after readiness. Another registered stream for that queue
141
+ prevents shutdown. The default is False and retries indefinitely. Always call
142
+ `stop()` or `close()` afterward; both remain safe after `start()` returns.
143
+
117
144
  ## Flow state I/O
118
145
 
119
146
  Applications read and write Flow state through typed RPCs. Client and AsyncClient
@@ -113,6 +113,11 @@ from dex.timer import Timer
113
113
  from dex.wait import Wait
114
114
  from dex.wait_options import WaitForAttributeOptions, WaitForStepCompletionOptions
115
115
  from dex.worker import Worker
116
+ from dex.pull_registration import (
117
+ PullRegistrationError,
118
+ PullRegistrationState,
119
+ PullQueueRegistrationState,
120
+ )
116
121
  from dex.worker_options import PullWorkerOptions, WorkerOptions, WorkerTarget
117
122
 
118
123
  __all__ = [
@@ -214,6 +219,9 @@ __all__ = [
214
219
  "WorkerOptions",
215
220
  "WorkerTarget",
216
221
  "PullWorkerOptions",
222
+ "PullRegistrationError",
223
+ "PullRegistrationState",
224
+ "PullQueueRegistrationState",
217
225
  "WaitForAttributeOptions",
218
226
  "WaitForStepCompletionOptions",
219
227
  "dead_end",
@@ -28,6 +28,7 @@ from dex._server_protocol import sdk_version
28
28
  from dex.dexpb import dex_pb2 as pb
29
29
  from dex.dexpb import dex_pb2_grpc
30
30
  from dex.flow import Registry
31
+ from dex.pull_registration import PullRegistrationError, _PullRegistration
31
32
  from dex.worker_options import PullWorkerOptions
32
33
 
33
34
  _LOGGER = logging.getLogger(__name__)
@@ -86,6 +87,12 @@ class PullWorkerRuntime:
86
87
  self.app_worker_id = app_worker_id or resolve_app_worker_id(registry, options)
87
88
  self.flow_types = sorted(registry._registered_flows)
88
89
  self.queues = (self.app_worker_id, *self.flow_types)
90
+ self.registration = _PullRegistration(self.queues)
91
+ self._stop_on_permanent_stream_rejection = (
92
+ options.stop_on_permanent_stream_rejection
93
+ )
94
+ self._permanent_error: PullRegistrationError | None = None
95
+ self._rejection_stop_task: asyncio.Task[None] | None = None
89
96
  self.max_concurrent_tasks = options.max_concurrent_tasks or os.cpu_count() or 1
90
97
  self.idle_streams_per_queue = options.idle_streams_per_queue or 1
91
98
  self.stop_drain_timeout = options.stop_drain_timeout.total_seconds() or 30
@@ -131,9 +138,10 @@ class PullWorkerRuntime:
131
138
  self._stopping = asyncio.Event()
132
139
  self._stop_completed = asyncio.Event()
133
140
  self._started.set()
134
- channel = outbound_worker_task_channel(self._server_address)
135
- stub = dex_pb2_grpc.WorkerTaskServiceStub(channel) # type: ignore[no-untyped-call]
141
+ channel: grpc.aio.Channel | None = None
136
142
  try:
143
+ channel = outbound_worker_task_channel(self._server_address)
144
+ stub = dex_pb2_grpc.WorkerTaskServiceStub(channel) # type: ignore[no-untyped-call]
137
145
  self._wake.set()
138
146
  while not self._stopping.is_set():
139
147
  await self._wake.wait()
@@ -144,9 +152,14 @@ class PullWorkerRuntime:
144
152
  *(stream.run_task for stream in tuple(self._streams)),
145
153
  return_exceptions=True,
146
154
  )
155
+ if self._rejection_stop_task is not None:
156
+ await self._rejection_stop_task
157
+ if self._permanent_error is not None:
158
+ raise self._permanent_error
147
159
  finally:
148
160
  try:
149
- await channel.close(grace=None)
161
+ if channel is not None:
162
+ await channel.close(grace=None)
150
163
  finally:
151
164
  self._loop_thread_id = None
152
165
 
@@ -157,6 +170,8 @@ class PullWorkerRuntime:
157
170
  def _reconcile(self, stub: Any) -> None:
158
171
  if self._stopping is None or self._stopping.is_set():
159
172
  return
173
+ for stream in tuple(self._streams):
174
+ stream.try_retire()
160
175
  for queue in self.queues:
161
176
  idle = sum(
162
177
  stream.queue == queue
@@ -192,6 +207,7 @@ class PullWorkerRuntime:
192
207
  self.signal()
193
208
 
194
209
  async def stop(self) -> None:
210
+ self.registration.record_stopping()
195
211
  if self._stopping is None:
196
212
  return
197
213
  if self._stopping.is_set():
@@ -263,7 +279,10 @@ class PullStream:
263
279
  self.handler_task: asyncio.Task[None] | None = None
264
280
  self.task_id = ""
265
281
  self.is_busy = False
282
+ self._is_registered = False
266
283
  self.is_draining = False
284
+ self._is_retiring = False
285
+ self._retirement_requested = asyncio.Event()
267
286
  self.is_half_closing = False
268
287
  self._active = threading.Event()
269
288
  self._call: Any = None
@@ -289,7 +308,28 @@ class PullStream:
289
308
  _LOGGER.exception(
290
309
  "Worker task stream failed", extra={"queue": self.queue}
291
310
  )
292
- if self.runtime.is_stopping:
311
+ finally:
312
+ code = error.code() if error is not None else None
313
+ rejection = self.runtime.registration.record_stream_result(
314
+ self.queue,
315
+ self._is_registered,
316
+ error if code in _PERMANENT_CODES else None,
317
+ )
318
+ self._is_registered = False
319
+ if (
320
+ rejection is not None
321
+ and self.runtime._stop_on_permanent_stream_rejection
322
+ and self.runtime._permanent_error is None
323
+ ):
324
+ self.runtime._permanent_error = rejection
325
+ self.runtime._rejection_stop_task = asyncio.create_task(
326
+ self.runtime.stop()
327
+ )
328
+ if (
329
+ self.runtime.is_stopping
330
+ or self.runtime._permanent_error is not None
331
+ or self._is_retiring
332
+ ):
293
333
  break
294
334
  if self._is_reconnecting:
295
335
  self._is_reconnecting = False
@@ -346,11 +386,12 @@ class PullStream:
346
386
  )
347
387
  was_registered = False
348
388
  reader: asyncio.Task[Any] | None = None
389
+ retirement = asyncio.create_task(self._retirement_requested.wait())
349
390
  try:
350
391
  while not self.runtime.is_stopping or self.is_busy:
351
392
  if reader is None:
352
393
  reader = asyncio.create_task(self._call.read())
353
- waiting: set[asyncio.Task[Any]] = {reader}
394
+ waiting: set[asyncio.Task[Any]] = {reader, retirement}
354
395
  if self.handler_task is not None:
355
396
  waiting.add(self.handler_task)
356
397
  completed, _ = await asyncio.wait(
@@ -358,22 +399,23 @@ class PullStream:
358
399
  )
359
400
  if self.handler_task in completed:
360
401
  await self._finish_handler()
361
- if self.is_draining:
362
- if reader.done():
363
- frame = reader.result()
364
- reader = None
365
- if (
366
- frame is not _GRPC_AIO_EOF
367
- and frame.WhichOneof("frame") == "task"
368
- ):
369
- await self._reject_task(frame.task)
370
- self.is_half_closing = True
371
- self.runtime.signal()
372
- await self._call.done_writing()
373
- await self._await_stream_eof(reader)
402
+ self.try_retire()
403
+ if self.is_draining and not self.is_busy:
404
+ if reader.done():
405
+ frame = reader.result()
374
406
  reader = None
375
- self._is_reconnecting = True
376
- return was_registered
407
+ if (
408
+ frame is not _GRPC_AIO_EOF
409
+ and frame.WhichOneof("frame") == "task"
410
+ ):
411
+ await self._reject_task(frame.task)
412
+ self.is_half_closing = True
413
+ self.runtime.signal()
414
+ await self._call.done_writing()
415
+ await self._await_stream_eof(reader)
416
+ reader = None
417
+ self._is_reconnecting = not self._is_retiring
418
+ return was_registered
377
419
  if reader not in completed:
378
420
  continue
379
421
  frame = reader.result()
@@ -391,6 +433,9 @@ class PullStream:
391
433
  frame.registered.supports_unknown_name_rejection
392
434
  )
393
435
  was_registered = True
436
+ self._is_registered = True
437
+ self.runtime.registration.record_registered_stream(self.queue)
438
+ self.runtime.signal()
394
439
  continue
395
440
  if kind == "task":
396
441
  await self._start_task(frame.task)
@@ -408,6 +453,8 @@ class PullStream:
408
453
  self._is_reconnecting = True
409
454
  return was_registered
410
455
  finally:
456
+ retirement.cancel()
457
+ await asyncio.gather(retirement, return_exceptions=True)
411
458
  self.is_half_closing = False
412
459
  self.cancel_handler()
413
460
  if self.handler_task is not None:
@@ -624,6 +671,22 @@ class PullStream:
624
671
  self._active.clear()
625
672
  self.runtime.finish_task()
626
673
 
674
+ def try_retire(self) -> None:
675
+ if self.is_busy or self.is_draining or not self._is_registered:
676
+ return
677
+ idle_replacements = sum(
678
+ candidate is not self
679
+ and candidate.queue == self.queue
680
+ and candidate._is_registered
681
+ and not candidate.is_busy
682
+ and not candidate.is_draining
683
+ for candidate in self.runtime._streams
684
+ )
685
+ if idle_replacements >= self.runtime.idle_streams_per_queue:
686
+ self.is_draining = True
687
+ self._is_retiring = True
688
+ self._retirement_requested.set()
689
+
627
690
  def cancel_handler(self) -> None:
628
691
  self._active.clear()
629
692
  if self.runtime._executor is None and self.handler_task is not None:
@@ -30,6 +30,7 @@ from dex.blob_cache import BlobCache
30
30
  from dex.dexpb import dex_pb2 as pb
31
31
  from dex.dexpb import dex_pb2_grpc
32
32
  from dex.flow import Registry
33
+ from dex.pull_registration import PullRegistrationState
33
34
  from dex.worker_options import WorkerOptions, WorkerTarget
34
35
 
35
36
 
@@ -100,6 +101,7 @@ class AsyncWorker:
100
101
  self._server_protocol,
101
102
  )
102
103
  self._server = None
104
+ self._has_start_called = False
103
105
  self._pull = None
104
106
  if self.options.pull is not None:
105
107
  self._pull = PullWorkerRuntime(
@@ -167,11 +169,77 @@ class AsyncWorker:
167
169
 
168
170
  ``start`` may be awaited exactly once. It contacts FlowService before binding
169
171
  and completes only after ``stop`` terminates the listener or pull streams.
170
- Later pull stream failures are logged and retried.
172
+ Permanent pull stream rejections raise PullRegistrationError when
173
+ stop_on_permanent_stream_rejection is enabled; otherwise streams retry.
171
174
 
172
175
  Raises:
176
+ PullRegistrationError: A permanent stream rejection stops this Worker.
173
177
  RuntimeError: If lifecycle state, index synchronization, or binding fails.
174
178
  """
179
+ if self._has_start_called:
180
+ raise RuntimeError("Worker can only start once")
181
+ self._has_start_called = True
182
+ start_error: BaseException | None = None
183
+ try:
184
+ await self._start()
185
+ except BaseException as failure:
186
+ start_error = failure
187
+ raise
188
+ finally:
189
+ if self._pull is not None:
190
+ self._pull.registration.finish_start(start_error)
191
+ if self._state == "running":
192
+ await self.stop()
193
+
194
+ async def wait_pull_registered(self, timeout: float | None = None) -> None:
195
+ """Wait until every pull queue has registered at least once.
196
+
197
+ Start the Worker separately before waiting, or start it concurrently.
198
+ Successful registration latches, including after stop. Use
199
+ pull_registration_state for subsequent readiness probes. A permanent
200
+ rejection before readiness raises PullRegistrationError only when no
201
+ other stream is registered for that queue. Startup failures wake waiters.
202
+ Multiple callers may wait independently, including before start.
203
+
204
+ Args:
205
+ timeout: Maximum seconds to wait; None waits indefinitely.
206
+
207
+ Raises:
208
+ PullRegistrationError: A queue is permanently rejected before readiness.
209
+ TimeoutError: The timeout expires.
210
+ asyncio.CancelledError: The calling task is canceled; the Worker keeps running.
211
+ RuntimeError: This is a listening Worker, or it stops before readiness.
212
+ BaseException: The original startup failure, if start fails.
213
+
214
+ Examples:
215
+ >>> start_task = asyncio.create_task(worker.start())
216
+
217
+ >>> await worker.wait_pull_registered(timeout=30)
218
+ >>> is_ready = worker.pull_registration_state().is_ready
219
+ """
220
+ if self._pull is None:
221
+ raise RuntimeError("wait_pull_registered requires a pull Worker")
222
+ await self._pull.registration.wait_async(timeout)
223
+
224
+ def pull_registration_state(self) -> PullRegistrationState:
225
+ """Return an immutable snapshot of pull queue registrations.
226
+
227
+ Transient failures, reconnects, and stop preserve initial readiness.
228
+ A permanent rejection makes readiness false only while that queue has no
229
+ registered stream. Re-registration clears its error. Calls are safe
230
+ concurrently with start, stop, and other waiters. See wait_pull_registered
231
+ for an example.
232
+
233
+ Returns:
234
+ Current per-queue state, or an empty nonready state in listening mode.
235
+ """
236
+ return (
237
+ PullRegistrationState()
238
+ if self._pull is None
239
+ else self._pull.registration.snapshot()
240
+ )
241
+
242
+ async def _start(self) -> None:
175
243
  if self._state != "created":
176
244
  raise RuntimeError(f"AsyncWorker cannot start from state {self._state}")
177
245
  try:
@@ -247,6 +315,8 @@ class AsyncWorker:
247
315
  Half-closed pull streams wait for server EOF until the remaining drain
248
316
  deadline before the channel closes.
249
317
  """
318
+ if self._pull is not None:
319
+ self._pull.registration.record_stopping()
250
320
  if self._state in ("stopped", "closed"):
251
321
  return
252
322
  self._state = "stopping"
@@ -0,0 +1,181 @@
1
+ # Copyright (c) 2026 Super Durable, Inc.
2
+ #
3
+ # Licensed under the Sustainable Use License 1.0.
4
+ # You may not use this file except in compliance with the License.
5
+ # See the LICENSE file in the repository root.
6
+ #
7
+ # SPDX-License-Identifier: LicenseRef-Sustainable-Use-1.0
8
+
9
+ import asyncio
10
+ import threading
11
+ import time
12
+ from concurrent.futures import Future
13
+ from dataclasses import dataclass
14
+
15
+
16
+ class PullRegistrationError(RuntimeError):
17
+ """Identify a permanent rejection leaving a pull queue without a registered stream.
18
+
19
+ Raised by ``wait_pull_registered`` before initial readiness, or by ``start``
20
+ when ``stop_on_permanent_stream_rejection`` is enabled. Inspect ``queue``
21
+ and ``error`` for the rejected queue and original gRPC error. Permanent codes
22
+ are INVALID_ARGUMENT, RESOURCE_EXHAUSTED, UNIMPLEMENTED, PERMISSION_DENIED,
23
+ and UNAUTHENTICATED.
24
+
25
+ Attributes:
26
+ queue: Rejected Flow type or app Worker ID queue.
27
+ error: Original permanent transport error, also exposed as ``__cause__``.
28
+ """
29
+
30
+ def __init__(self, queue: str, error: BaseException) -> None:
31
+ """Create a queue-specific rejection.
32
+
33
+ Args:
34
+ queue: Rejected queue name.
35
+ error: Original permanent transport error.
36
+ """
37
+ super().__init__(f"pull Worker queue {queue!r} registration rejected: {error}")
38
+ self.queue = queue
39
+ self.error = error
40
+ self.__cause__ = error
41
+
42
+
43
+ @dataclass(frozen=True)
44
+ class PullQueueRegistrationState:
45
+ """Describe one queue in an immutable readiness snapshot.
46
+
47
+ Attributes:
48
+ queue: Flow type or app Worker ID queue.
49
+ has_registered_stream: Whether a stream is currently registered.
50
+ last_permanent_error: Last permanent transport error, or ``None``.
51
+ A new Registered acknowledgement clears it. Treat errors as immutable.
52
+ """
53
+
54
+ queue: str
55
+ has_registered_stream: bool
56
+ last_permanent_error: BaseException | None
57
+
58
+
59
+ @dataclass(frozen=True)
60
+ class PullRegistrationState:
61
+ """Describe pull registration readiness independently of the Worker lifecycle.
62
+
63
+ Obtain snapshots with ``worker.pull_registration_state()``. After every queue
64
+ has registered once, transient failures, reconnects, and stop keep readiness.
65
+ A permanent error with no registered stream makes ``is_ready`` false until
66
+ re-registration. A listening Worker returns ``False`` and an empty tuple.
67
+
68
+ Attributes:
69
+ is_ready: Whether all queues registered once and no queue is permanently
70
+ rejected without a registered stream.
71
+ queues: Immutable per-queue snapshots, ordered by app Worker ID then Flow type.
72
+
73
+ Examples:
74
+ >>> state = worker.pull_registration_state()
75
+ >>> is_ready = state.is_ready
76
+ """
77
+
78
+ is_ready: bool = False
79
+ queues: tuple[PullQueueRegistrationState, ...] = ()
80
+
81
+
82
+ class _PullRegistration:
83
+ def __init__(self, queues: tuple[str, ...]) -> None:
84
+ self._lock = threading.Lock()
85
+ self._queues = queues
86
+ self._counts = dict.fromkeys(queues, 0)
87
+ self._registered: set[str] = set()
88
+ self._errors: dict[str, BaseException] = {}
89
+ self._initial_error: PullRegistrationError | None = None
90
+ self._has_registered_all = False
91
+ self._is_stopping = False
92
+ self._has_start_returned = False
93
+ self._start_error: BaseException | None = None
94
+ self._changed: Future[None] = Future()
95
+
96
+ def record_registered_stream(self, queue: str) -> None:
97
+ with self._lock:
98
+ self._counts[queue] += 1
99
+ self._registered.add(queue)
100
+ self._errors.pop(queue, None)
101
+ self._has_registered_all = len(self._registered) == len(self._queues)
102
+ self._notify()
103
+
104
+ def record_stream_result(
105
+ self, queue: str, was_registered: bool, error: BaseException | None
106
+ ) -> PullRegistrationError | None:
107
+ with self._lock:
108
+ if was_registered:
109
+ self._counts[queue] -= 1
110
+ rejection = None
111
+ if error is not None:
112
+ self._errors[queue] = error
113
+ if self._counts[queue] == 0:
114
+ rejection = PullRegistrationError(queue, error)
115
+ if not self._has_registered_all and self._initial_error is None:
116
+ self._initial_error = rejection
117
+ self._notify()
118
+ return rejection
119
+
120
+ def record_stopping(self) -> None:
121
+ with self._lock:
122
+ self._is_stopping = True
123
+ self._notify()
124
+
125
+ def finish_start(self, error: BaseException | None) -> None:
126
+ with self._lock:
127
+ self._has_start_returned = True
128
+ self._start_error = error
129
+ self._notify()
130
+
131
+ def snapshot(self) -> PullRegistrationState:
132
+ with self._lock:
133
+ queues = tuple(
134
+ PullQueueRegistrationState(
135
+ queue, self._counts[queue] > 0, self._errors.get(queue)
136
+ )
137
+ for queue in self._queues
138
+ )
139
+ return PullRegistrationState(
140
+ self._has_registered_all
141
+ and not any(
142
+ not queue.has_registered_stream
143
+ and queue.last_permanent_error is not None
144
+ for queue in queues
145
+ ),
146
+ queues,
147
+ )
148
+
149
+ def wait(self, timeout: float | None) -> None:
150
+ deadline = None if timeout is None else time.monotonic() + timeout
151
+ while (changed := self._next_change()) is not None:
152
+ changed.result(
153
+ None if deadline is None else max(0, deadline - time.monotonic())
154
+ )
155
+
156
+ async def wait_async(self, timeout: float | None) -> None:
157
+ if self._next_change() is not None:
158
+ await asyncio.wait_for(self._wait_async(), timeout)
159
+
160
+ async def _wait_async(self) -> None:
161
+ while (changed := self._next_change()) is not None:
162
+ await asyncio.shield(asyncio.wrap_future(changed))
163
+
164
+ def _next_change(self) -> Future[None] | None:
165
+ with self._lock:
166
+ if self._has_registered_all:
167
+ return None
168
+ if self._initial_error is not None:
169
+ raise self._initial_error
170
+ if self._has_start_returned:
171
+ if self._start_error is not None:
172
+ raise self._start_error
173
+ raise RuntimeError("pull Worker stopped before registration")
174
+ if self._is_stopping:
175
+ raise RuntimeError("pull Worker stopped before registration")
176
+ return self._changed
177
+
178
+ def _notify(self) -> None:
179
+ changed = self._changed
180
+ self._changed = Future()
181
+ changed.set_result(None)
@@ -35,6 +35,7 @@ from dex.blob_cache import BlobCache
35
35
  from dex.dexpb import dex_pb2 as pb
36
36
  from dex.dexpb import dex_pb2_grpc
37
37
  from dex.flow import Registry
38
+ from dex.pull_registration import PullRegistrationState
38
39
  from dex.worker_options import WorkerOptions, WorkerTarget
39
40
 
40
41
 
@@ -117,6 +118,7 @@ class Worker:
117
118
  thread_name_prefix="dex-python-handler",
118
119
  )
119
120
  self._server = None
121
+ self._has_start_called = False
120
122
  self._pull = None
121
123
  if self.options.pull is not None:
122
124
  self._pull = PullWorkerRuntime(
@@ -188,12 +190,78 @@ class Worker:
188
190
 
189
191
  ``start`` may be called exactly once. It first contacts FlowService using the
190
192
  configured synchronization timeout, then starts a listener or outbound
191
- pull streams and waits until ``stop`` finishes. Later stream failures are
192
- logged and retried; they do not raise from ``start``.
193
+ pull streams and waits until ``stop`` finishes. Permanent stream rejections
194
+ raise PullRegistrationError when stop_on_permanent_stream_rejection is
195
+ enabled; otherwise streams retry.
193
196
 
194
197
  Raises:
198
+ PullRegistrationError: A permanent stream rejection stops this Worker.
195
199
  RuntimeError: If lifecycle state, index synchronization, or binding fails.
196
200
  """
201
+ with self._lock:
202
+ if self._has_start_called or self._state != "created":
203
+ raise RuntimeError(f"Worker cannot start from state {self._state}")
204
+ self._has_start_called = True
205
+ start_error: BaseException | None = None
206
+ try:
207
+ self._start()
208
+ except BaseException as failure:
209
+ start_error = failure
210
+ raise
211
+ finally:
212
+ if self._pull is not None:
213
+ self._pull.registration.finish_start(start_error)
214
+ if self._state == "running":
215
+ self.stop()
216
+
217
+ def wait_pull_registered(self, timeout: float | None = None) -> None:
218
+ """Wait until every pull queue has registered at least once.
219
+
220
+ Start the Worker separately before waiting, or start it concurrently.
221
+ Successful registration latches, including after stop. Use
222
+ pull_registration_state for subsequent readiness probes. A permanent
223
+ rejection before readiness raises PullRegistrationError only when no
224
+ other stream is registered for that queue. Startup failures wake waiters.
225
+ Multiple callers may wait independently, including before start.
226
+
227
+ Args:
228
+ timeout: Maximum seconds to wait; None waits indefinitely.
229
+
230
+ Raises:
231
+ PullRegistrationError: A queue is permanently rejected before readiness.
232
+ TimeoutError: The timeout expires.
233
+ RuntimeError: This is a listening Worker, or it stops before readiness.
234
+ BaseException: The original startup failure, if start fails.
235
+
236
+ Examples:
237
+ >>> start_thread = threading.Thread(target=worker.start)
238
+ >>> start_thread.start()
239
+ >>> worker.wait_pull_registered(timeout=30)
240
+ >>> is_ready = worker.pull_registration_state().is_ready
241
+ """
242
+ if self._pull is None:
243
+ raise RuntimeError("wait_pull_registered requires a pull Worker")
244
+ self._pull.registration.wait(timeout)
245
+
246
+ def pull_registration_state(self) -> PullRegistrationState:
247
+ """Return an immutable snapshot of pull queue registrations.
248
+
249
+ Transient failures, reconnects, and stop preserve initial readiness.
250
+ A permanent rejection makes readiness false only while that queue has no
251
+ registered stream. Re-registration clears its error. Calls are safe
252
+ concurrently with start, stop, and other waiters. See wait_pull_registered
253
+ for an example.
254
+
255
+ Returns:
256
+ Current per-queue state, or an empty nonready state in listening mode.
257
+ """
258
+ return (
259
+ PullRegistrationState()
260
+ if self._pull is None
261
+ else self._pull.registration.snapshot()
262
+ )
263
+
264
+ def _start(self) -> None:
197
265
  with self._lock:
198
266
  if self._state != "created":
199
267
  raise RuntimeError(f"Worker cannot start from state {self._state}")
@@ -291,6 +359,7 @@ class Worker:
291
359
  was_created = self._state == "created"
292
360
  self._state = "stopping"
293
361
  if self._pull is not None:
362
+ self._pull.registration.record_stopping()
294
363
  if not was_created:
295
364
  self._pull.stop_threadsafe()
296
365
  else:
@@ -77,6 +77,11 @@ class PullWorkerOptions:
77
77
  the final frame. During Stop, half-closed streams use the remaining
78
78
  drain deadline.
79
79
 
80
+ stop_on_permanent_stream_rejection: Whether the first permanent rejection
81
+ leaving a queue without a registered stream drains the Worker and
82
+ raises PullRegistrationError from start. Defaults to False, which
83
+ retries indefinitely. Applies before and after initial readiness.
84
+
80
85
  Raises:
81
86
  ValueError: If a numeric option is negative or app_worker_id is invalid.
82
87
 
@@ -91,6 +96,7 @@ class PullWorkerOptions:
91
96
  idle_streams_per_queue: int = 0
92
97
  app_worker_id: str = ""
93
98
  stop_drain_timeout: timedelta = timedelta(0)
99
+ stop_on_permanent_stream_rejection: bool = False
94
100
 
95
101
  def __post_init__(self) -> None:
96
102
  if self.max_concurrent_tasks < 0:
@@ -1,7 +1,7 @@
1
1
  [project]
2
2
  name = "dex-python-sdk"
3
3
  # Publish uses tag sdk-python/v*; CI stamps this for release builds.
4
- version = "1.12.0"
4
+ version = "1.14.1"
5
5
  description = "Python SDK for the Dex workflow engine"
6
6
  authors = [{name = "Super Durable"}]
7
7
  readme = "README.md"
File without changes