dex-python-sdk 1.20.0__tar.gz → 1.23.0__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 (72) hide show
  1. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/PKG-INFO +28 -18
  2. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/README.md +27 -17
  3. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_pull_worker.py +111 -7
  4. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/async_worker.py +2 -0
  5. dex_python_sdk-1.23.0/dex/dexpb/dex_pb2.py +551 -0
  6. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/dexpb/dex_pb2.pyi +12 -6
  7. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/worker.py +2 -0
  8. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/worker_options.py +7 -3
  9. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/pyproject.toml +1 -1
  10. dex_python_sdk-1.20.0/dex/dexpb/dex_pb2.py +0 -551
  11. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/LEGACY_NOTICES.md +0 -0
  12. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/LICENSE +0 -0
  13. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/__init__.py +0 -0
  14. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_async_value_hydrator.py +0 -0
  15. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_async_worker_dispatcher.py +0 -0
  16. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_async_worker_service.py +0 -0
  17. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_credential_transport.py +0 -0
  18. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_grpc_errors.py +0 -0
  19. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_invocation_context.py +0 -0
  20. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_native.pyi +0 -0
  21. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_server_protocol.py +0 -0
  22. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_utils.py +0 -0
  23. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_value_hydrator.py +0 -0
  24. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_value_mapper.py +0 -0
  25. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_worker_dispatcher.py +0 -0
  26. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/_worker_service.py +0 -0
  27. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/async_client.py +0 -0
  28. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/attribute.py +0 -0
  29. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/attribute_match.py +0 -0
  30. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/blob_cache.py +0 -0
  31. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/channel.py +0 -0
  32. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/client.py +0 -0
  33. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/client_options.py +0 -0
  34. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/codec.py +0 -0
  35. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/condition.py +0 -0
  36. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/context.py +0 -0
  37. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/credential.py +0 -0
  38. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/dexpb/__init__.py +0 -0
  39. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/dexpb/dex_pb2_grpc.py +0 -0
  40. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/flow.py +0 -0
  41. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/flow_config.py +0 -0
  42. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/flow_info.py +0 -0
  43. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/flow_options.py +0 -0
  44. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/flow_result.py +0 -0
  45. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/pull_registration.py +0 -0
  46. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/py.typed +0 -0
  47. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/rpc_invoke_options.py +0 -0
  48. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/runtime_errors.py +0 -0
  49. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/step.py +0 -0
  50. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/step_execution.py +0 -0
  51. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/stream.py +0 -0
  52. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/subflow.py +0 -0
  53. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/timer.py +0 -0
  54. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/wait.py +0 -0
  55. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/dex/wait_options.py +0 -0
  56. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/Cargo.lock +0 -0
  57. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/Cargo.toml +0 -0
  58. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache/Cargo.toml +0 -0
  59. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache/LICENSE +0 -0
  60. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache/README.md +0 -0
  61. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache/src/config.rs +0 -0
  62. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache/src/entry.rs +0 -0
  63. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache/src/error.rs +0 -0
  64. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache/src/format.rs +0 -0
  65. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache/src/lib.rs +0 -0
  66. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache/src/policy.rs +0 -0
  67. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache/src/store.rs +0 -0
  68. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache/tests/blob_cache_integration.rs +0 -0
  69. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache-python/Cargo.toml +0 -0
  70. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache-python/LICENSE +0 -0
  71. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/sdk-rust/crates/dex-blob-cache-python/README.md +0 -0
  72. {dex_python_sdk-1.20.0 → dex_python_sdk-1.23.0}/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.20.0
3
+ Version: 1.23.0
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
@@ -135,13 +135,18 @@ The Worker polls its sticky queue and every registered Flow type queue. Flow
135
135
  type names must be unique within a Dex server. `max_concurrent_tasks` defaults
136
136
  to the CPU count, and `idle_streams_per_queue` defaults to one. At capacity,
137
137
  the Worker rejects an assignment so Dex can reassign it. Each stream runs one
138
- task at a time; Step heartbeat and stream writes use the same stream. After tasks finish, the Worker
139
- gracefully closes surplus idle streams while preserving readiness. After load drops,
138
+ task at a time; Step heartbeat and stream writes use the same stream. After
139
+ tasks finish, the Worker gracefully closes surplus idle streams while preserving
140
+ readiness. The SDK sends an upstream Pong every 20 seconds while a task runs and
141
+ stops before its final frame. This interval is internal and needs no Worker
142
+ option. After load drops,
140
143
  the steady-state bound is `max_concurrent_tasks + (1 + number of Flow types) *
141
- idle_streams_per_queue` streams. The transient peak during a burst remains bounded
142
- by the matching caps. Size these options for `matching.maxConcurrentTasksCap` per
143
- app Worker ID per matching owner (256 by default) and `matching.maxStreams` per
144
- Matching instance (100,000 by default).
144
+ idle_streams_per_queue` streams.
145
+ Allow up to one more stream per queue for 30 seconds after reserve declines,
146
+ or up to `1 + number of Flow types` additional streams. The transient peak
147
+ during a burst remains bounded by the matching caps. Size these options for
148
+ `matching.maxConcurrentTasksCap` per app Worker ID per matching owner (256 by
149
+ default) and `matching.maxStreams` per Matching instance (100,000 by default).
145
150
 
146
151
  The Worker rejects an unregistered Step or RPC by name. Dex retries that task
147
152
  on another pull Worker without consuming a Step attempt, including when the
@@ -206,11 +211,11 @@ limit is at least two, otherwise zero. With the CPU-count default, Steps get one
206
211
  fewer slot than before. Set it explicitly to zero to opt out. With one slot and
207
212
  no explicit reserve, Worker start logs one warning: RPCs may wait behind
208
213
  long-running Steps; use a limit of at least two or set the reserve explicitly.
209
- Flows must stay responsive; Dex cannot prioritize synchronous RPCs like workflow
210
- engines prioritize reads. The value must be nonnegative and less than
211
- `max_concurrent_tasks` after its CPU-count default is resolved; invalid values
212
- fail construction. RPC and Step handlers count together: WaitFor and Execute
213
- tasks, including local attempts, are declined with `NO_CAPACITY` when the busy
214
+ Flows must stay responsive because RPCs share Worker task capacity with Steps.
215
+ The value must be nonnegative and less than `max_concurrent_tasks` after its
216
+ CPU-count default is resolved; invalid values fail construction. RPC and Step
217
+ handlers count together. WaitFor and Execute tasks, including local attempts,
218
+ are declined with `NO_CAPACITY` when the busy
214
219
  count reaches the task limit minus the reserve. API-direct and transactional
215
220
  InvokeWorkerRPC tasks can use the full limit.
216
221
 
@@ -222,15 +227,20 @@ A decline lets Dex try another Worker without consuming an attempt. If all
222
227
  Workers remain at their Step limit, a Step can wait until the dispatch wait
223
228
  expires, then receive `NO_WORKER_AVAILABLE`, consume an attempt, and follow its
224
229
  retry policy, as when all Workers are full. Local attempts wait the local
225
- dispatch wait, then fall back according to their policy. Each queue still keeps
226
- its configured idle streams; a reserve decline does not close a stream or change
227
- registration readiness.
230
+ dispatch wait, then fall back according to their policy. A queue that declines
231
+ a Step for reserved RPC capacity keeps one additional idle stream until 30
232
+ seconds after its last decline. Other queues retain their configured count.
233
+ The Worker does not open the extra stream while all task slots are busy. Set
234
+ `idle_streams_per_queue` to at least two for heavy Step backlogs. A reserve decline
235
+ does not close a stream or change registration readiness.
228
236
 
229
237
  Each reserve decline records the Flow type through the dex._pull_worker logger
230
238
  at DEBUG level. Export the cumulative `pull_reserved_rpc_declines` Worker value
231
- to distinguish reserve pressure from full-capacity rejections. The count
232
- excludes full-capacity rejections and unknown names, never resets, and remains
233
- available after shutdown. Listening Workers return zero.
239
+ to distinguish reserve pressure from full-capacity rejections. The count excludes
240
+ full-capacity rejections and unknown names. It counts decline events, including
241
+ repeated declines of the same Step dispatch while Dex retries it. Read its rate
242
+ to assess reserve pressure. It never resets and remains available after shutdown.
243
+ Listening Workers return zero.
234
244
 
235
245
  ## Flow state I/O
236
246
 
@@ -119,13 +119,18 @@ The Worker polls its sticky queue and every registered Flow type queue. Flow
119
119
  type names must be unique within a Dex server. `max_concurrent_tasks` defaults
120
120
  to the CPU count, and `idle_streams_per_queue` defaults to one. At capacity,
121
121
  the Worker rejects an assignment so Dex can reassign it. Each stream runs one
122
- task at a time; Step heartbeat and stream writes use the same stream. After tasks finish, the Worker
123
- gracefully closes surplus idle streams while preserving readiness. After load drops,
122
+ task at a time; Step heartbeat and stream writes use the same stream. After
123
+ tasks finish, the Worker gracefully closes surplus idle streams while preserving
124
+ readiness. The SDK sends an upstream Pong every 20 seconds while a task runs and
125
+ stops before its final frame. This interval is internal and needs no Worker
126
+ option. After load drops,
124
127
  the steady-state bound is `max_concurrent_tasks + (1 + number of Flow types) *
125
- idle_streams_per_queue` streams. The transient peak during a burst remains bounded
126
- by the matching caps. Size these options for `matching.maxConcurrentTasksCap` per
127
- app Worker ID per matching owner (256 by default) and `matching.maxStreams` per
128
- Matching instance (100,000 by default).
128
+ idle_streams_per_queue` streams.
129
+ Allow up to one more stream per queue for 30 seconds after reserve declines,
130
+ or up to `1 + number of Flow types` additional streams. The transient peak
131
+ during a burst remains bounded by the matching caps. Size these options for
132
+ `matching.maxConcurrentTasksCap` per app Worker ID per matching owner (256 by
133
+ default) and `matching.maxStreams` per Matching instance (100,000 by default).
129
134
 
130
135
  The Worker rejects an unregistered Step or RPC by name. Dex retries that task
131
136
  on another pull Worker without consuming a Step attempt, including when the
@@ -190,11 +195,11 @@ limit is at least two, otherwise zero. With the CPU-count default, Steps get one
190
195
  fewer slot than before. Set it explicitly to zero to opt out. With one slot and
191
196
  no explicit reserve, Worker start logs one warning: RPCs may wait behind
192
197
  long-running Steps; use a limit of at least two or set the reserve explicitly.
193
- Flows must stay responsive; Dex cannot prioritize synchronous RPCs like workflow
194
- engines prioritize reads. The value must be nonnegative and less than
195
- `max_concurrent_tasks` after its CPU-count default is resolved; invalid values
196
- fail construction. RPC and Step handlers count together: WaitFor and Execute
197
- tasks, including local attempts, are declined with `NO_CAPACITY` when the busy
198
+ Flows must stay responsive because RPCs share Worker task capacity with Steps.
199
+ The value must be nonnegative and less than `max_concurrent_tasks` after its
200
+ CPU-count default is resolved; invalid values fail construction. RPC and Step
201
+ handlers count together. WaitFor and Execute tasks, including local attempts,
202
+ are declined with `NO_CAPACITY` when the busy
198
203
  count reaches the task limit minus the reserve. API-direct and transactional
199
204
  InvokeWorkerRPC tasks can use the full limit.
200
205
 
@@ -206,15 +211,20 @@ A decline lets Dex try another Worker without consuming an attempt. If all
206
211
  Workers remain at their Step limit, a Step can wait until the dispatch wait
207
212
  expires, then receive `NO_WORKER_AVAILABLE`, consume an attempt, and follow its
208
213
  retry policy, as when all Workers are full. Local attempts wait the local
209
- dispatch wait, then fall back according to their policy. Each queue still keeps
210
- its configured idle streams; a reserve decline does not close a stream or change
211
- registration readiness.
214
+ dispatch wait, then fall back according to their policy. A queue that declines
215
+ a Step for reserved RPC capacity keeps one additional idle stream until 30
216
+ seconds after its last decline. Other queues retain their configured count.
217
+ The Worker does not open the extra stream while all task slots are busy. Set
218
+ `idle_streams_per_queue` to at least two for heavy Step backlogs. A reserve decline
219
+ does not close a stream or change registration readiness.
212
220
 
213
221
  Each reserve decline records the Flow type through the dex._pull_worker logger
214
222
  at DEBUG level. Export the cumulative `pull_reserved_rpc_declines` Worker value
215
- to distinguish reserve pressure from full-capacity rejections. The count
216
- excludes full-capacity rejections and unknown names, never resets, and remains
217
- available after shutdown. Listening Workers return zero.
223
+ to distinguish reserve pressure from full-capacity rejections. The count excludes
224
+ full-capacity rejections and unknown names. It counts decline events, including
225
+ repeated declines of the same Step dispatch while Dex retries it. Read its rate
226
+ to assess reserve pressure. It never resets and remains available after shutdown.
227
+ Listening Workers return zero.
218
228
 
219
229
  ## Flow state I/O
220
230
 
@@ -16,6 +16,7 @@ import random
16
16
  import re
17
17
  import socket
18
18
  import threading
19
+ import time
19
20
  from concurrent.futures import ThreadPoolExecutor
20
21
  from datetime import datetime, timezone
21
22
  from typing import Any
@@ -37,6 +38,7 @@ from dex.worker_options import PullWorkerOptions
37
38
 
38
39
  _LOGGER = logging.getLogger(__name__)
39
40
  _GRPC_AIO_EOF = getattr(grpc.aio, "EOF")
41
+ _BUSY_PULL_STREAM_KEEPALIVE_SECONDS = 20
40
42
  _PERMANENT_CODES = {
41
43
  grpc.StatusCode.UNAUTHENTICATED,
42
44
  grpc.StatusCode.PERMISSION_DENIED,
@@ -86,6 +88,7 @@ class PullWorkerRuntime:
86
88
  executor: ThreadPoolExecutor | None,
87
89
  app_worker_id: str | None = None,
88
90
  credential_runner: AsyncCredentialRunner | None = None,
91
+ reserve_pressure_linger: float = 30.0,
89
92
  ) -> None:
90
93
  self._registry = registry
91
94
  self.app_worker_id = app_worker_id or resolve_app_worker_id(registry, options)
@@ -107,6 +110,9 @@ class PullWorkerRuntime:
107
110
  options.reserved_rpc_tasks is None and self.max_concurrent_tasks == 1
108
111
  )
109
112
  self.reserved_rpc_declines = 0
113
+ self._last_reserve_decline: dict[str, float] = {}
114
+ self._reserve_pressure_timers: dict[str, asyncio.TimerHandle] = {}
115
+ self._reserve_pressure_linger = reserve_pressure_linger
110
116
  self.idle_streams_per_queue = options.idle_streams_per_queue or 1
111
117
  self.stop_drain_timeout = options.stop_drain_timeout.total_seconds() or 30
112
118
  self._dispatcher = dispatcher
@@ -196,12 +202,26 @@ class PullWorkerRuntime:
196
202
  and not stream.is_half_closing
197
203
  for stream in self._streams
198
204
  )
199
- for _ in range(self.idle_streams_per_queue - idle):
205
+ desired_idle = self.desired_idle_streams_per_queue(queue)
206
+ creation_target = (
207
+ self.idle_streams_per_queue
208
+ if self._busy_tasks == self.max_concurrent_tasks
209
+ else desired_idle
210
+ )
211
+ for _ in range(creation_target - idle):
200
212
  stream = PullStream(self, stub, queue)
201
213
  self._streams.add(stream)
202
214
  stream.run_task = asyncio.create_task(stream.run())
203
215
  stream.run_task.add_done_callback(stream.on_done)
204
216
 
217
+ def desired_idle_streams_per_queue(self, queue: str) -> int:
218
+ last_decline = self._last_reserve_decline.get(queue)
219
+ return self.idle_streams_per_queue + int(
220
+ self.reserved_rpc_tasks > 0
221
+ and last_decline is not None
222
+ and time.monotonic() - last_decline < self._reserve_pressure_linger
223
+ )
224
+
205
225
  def _stream_done(self, stream: PullStream, completed: asyncio.Task[None]) -> None:
206
226
  self._streams.discard(stream)
207
227
  if not completed.cancelled():
@@ -210,7 +230,9 @@ class PullWorkerRuntime:
210
230
  _LOGGER.error("Worker task stream failed", exc_info=failure)
211
231
  self.signal()
212
232
 
213
- def accept_task(self, task: pb.PollWorkerTask | None = None) -> bool:
233
+ def accept_task(
234
+ self, task: pb.PollWorkerTask | None = None, queue: str | None = None
235
+ ) -> bool:
214
236
  if self._stopping is None or self._stopping.is_set():
215
237
  return False
216
238
  if self._busy_tasks >= self.max_concurrent_tasks:
@@ -223,6 +245,18 @@ class PullWorkerRuntime:
223
245
  and self.unknown_task_reason(task) is None
224
246
  ):
225
247
  self.reserved_rpc_declines += 1
248
+ queue_name = queue or self.app_worker_id
249
+ self._last_reserve_decline[queue_name] = time.monotonic()
250
+ existing_timer = self._reserve_pressure_timers.pop(queue_name, None)
251
+ if existing_timer is not None:
252
+ existing_timer.cancel()
253
+ if self._loop is not None:
254
+ self._reserve_pressure_timers[queue_name] = self._loop.call_later(
255
+ self._reserve_pressure_linger,
256
+ self._expire_reserve_pressure,
257
+ queue_name,
258
+ )
259
+ self.signal()
226
260
  request = task.wait_for if task.HasField("wait_for") else task.execute
227
261
  _LOGGER.debug(
228
262
  "decline Step task for reserved RPC capacity: %s", request.flow_type
@@ -232,6 +266,21 @@ class PullWorkerRuntime:
232
266
  self.signal()
233
267
  return True
234
268
 
269
+ def _expire_reserve_pressure(self, queue: str) -> None:
270
+ last_decline = self._last_reserve_decline.get(queue)
271
+ if last_decline is None:
272
+ return
273
+ remaining = self._reserve_pressure_linger - (time.monotonic() - last_decline)
274
+ if remaining > 0:
275
+ assert self._loop is not None
276
+ self._reserve_pressure_timers[queue] = self._loop.call_later(
277
+ remaining, self._expire_reserve_pressure, queue
278
+ )
279
+ return
280
+ self._last_reserve_decline.pop(queue, None)
281
+ self._reserve_pressure_timers.pop(queue, None)
282
+ self.signal()
283
+
235
284
  def finish_task(self) -> None:
236
285
  self._busy_tasks -= 1
237
286
  self.signal()
@@ -250,6 +299,9 @@ class PullWorkerRuntime:
250
299
  asyncio.get_running_loop().time() + self.stop_drain_timeout
251
300
  )
252
301
  self._stopping.set()
302
+ for timer in self._reserve_pressure_timers.values():
303
+ timer.cancel()
304
+ self._reserve_pressure_timers.clear()
253
305
  try:
254
306
  for stream in tuple(self._streams):
255
307
  stream.drain()
@@ -317,6 +369,8 @@ class PullStream:
317
369
  self._active = threading.Event()
318
370
  self._call: Any = None
319
371
  self._send_lock = asyncio.Lock()
372
+ self._keepalive_task: asyncio.Task[None] | None = None
373
+ self._is_final_frame_sent = False
320
374
  self._is_reconnecting = False
321
375
  self._stream_metadata: Any = None
322
376
  self._deadline_handle: asyncio.TimerHandle | None = None
@@ -424,6 +478,7 @@ class PullStream:
424
478
  if self.runtime.is_stopping:
425
479
  return False
426
480
  self._supports_unknown_name_rejection = False
481
+ self._is_final_frame_sent = False
427
482
  self._stream_metadata = None
428
483
  if self.runtime._credential_runner is None:
429
484
  self._call = self.stub.PollWorkerTask()
@@ -471,9 +526,11 @@ class PullStream:
471
526
  and frame.WhichOneof("frame") == "task"
472
527
  ):
473
528
  await self._reject_task(frame.task)
474
- self.is_half_closing = True
475
- self.runtime.signal()
476
- await self._call.done_writing()
529
+ async with self._send_lock:
530
+ self._stop_keepalive()
531
+ self.is_half_closing = True
532
+ self.runtime.signal()
533
+ await self._call.done_writing()
477
534
  await self._await_stream_eof(reader)
478
535
  reader = None
479
536
  self._is_reconnecting = not self._is_retiring
@@ -515,6 +572,8 @@ class PullStream:
515
572
  self._is_reconnecting = True
516
573
  return was_registered
517
574
  finally:
575
+ async with self._send_lock:
576
+ self._stop_keepalive()
518
577
  retirement.cancel()
519
578
  await asyncio.gather(retirement, return_exceptions=True)
520
579
  self.is_half_closing = False
@@ -586,12 +645,18 @@ class PullStream:
586
645
  if unknown_reason is not None:
587
646
  await self._reject_task(task, unknown_reason)
588
647
  return
589
- if self.is_busy or self.is_draining or not self.runtime.accept_task(task):
648
+ if (
649
+ self.is_busy
650
+ or self.is_draining
651
+ or not self.runtime.accept_task(task, self.queue)
652
+ ):
590
653
  await self._reject_task(task)
591
654
  return
592
655
  self.is_busy = True
593
656
  self.task_id = task.task_id
594
657
  self._active.set()
658
+ self._is_final_frame_sent = False
659
+ self._keepalive_task = asyncio.create_task(self._send_busy_keepalives())
595
660
  self.handler_task = asyncio.create_task(self._execute(task))
596
661
  if task.HasField("deadline"):
597
662
  deadline = task.deadline.ToDatetime(tzinfo=timezone.utc)
@@ -716,14 +781,53 @@ class PullStream:
716
781
  async def _send_frame(self, frame: pb.PollWorkerTaskFrame) -> None:
717
782
  await self._send(pb.PollWorkerTaskRequest(task_frame=frame))
718
783
 
784
+ async def _send_busy_keepalives(self) -> None:
785
+ try:
786
+ while True:
787
+ await asyncio.sleep(_BUSY_PULL_STREAM_KEEPALIVE_SECONDS)
788
+ async with self._send_lock:
789
+ if self._is_final_frame_sent or self.is_half_closing:
790
+ return
791
+ await self._call.write(
792
+ pb.PollWorkerTaskRequest(pong=pb.PollWorkerTaskPong())
793
+ )
794
+ except asyncio.CancelledError:
795
+ return
796
+ except Exception:
797
+ _LOGGER.exception("send Worker keepalive Pong failed")
798
+
799
+ def _stop_keepalive(self) -> None:
800
+ if self._keepalive_task is not None:
801
+ self._keepalive_task.cancel()
802
+ self._keepalive_task = None
803
+
719
804
  async def _send(self, request: pb.PollWorkerTaskRequest) -> None:
720
805
  async with self._send_lock:
806
+ if request.HasField("task_frame"):
807
+ frame = request.task_frame
808
+ if (
809
+ frame.HasField("rpc")
810
+ or frame.HasField("worker_error")
811
+ or frame.HasField("wait_for")
812
+ and frame.wait_for.WhichOneof("output") == "result"
813
+ or frame.HasField("execute")
814
+ and frame.execute.WhichOneof("output") == "result"
815
+ ):
816
+ self._stop_keepalive()
817
+ self._is_final_frame_sent = True
818
+ if request.HasField("pong") and (
819
+ self._is_final_frame_sent or self.is_half_closing
820
+ ):
821
+ return
721
822
  await self._call.write(request)
722
823
 
723
824
  async def _finish_handler(self) -> None:
724
825
  if self.handler_task is None:
725
826
  return
726
827
  await asyncio.gather(self.handler_task, return_exceptions=True)
828
+ async with self._send_lock:
829
+ self._stop_keepalive()
830
+ self._is_final_frame_sent = False
727
831
  if self._deadline_handle is not None:
728
832
  self._deadline_handle.cancel()
729
833
  self._deadline_handle = None
@@ -744,7 +848,7 @@ class PullStream:
744
848
  and not candidate.is_draining
745
849
  for candidate in self.runtime._streams
746
850
  )
747
- if idle_replacements >= self.runtime.idle_streams_per_queue:
851
+ if idle_replacements >= self.runtime.desired_idle_streams_per_queue(self.queue):
748
852
  self.is_draining = True
749
853
  self._is_retiring = True
750
854
  self._retirement_requested.set()
@@ -225,6 +225,8 @@ class AsyncWorker:
225
225
  def pull_reserved_rpc_declines(self) -> int:
226
226
  """Return cumulative Step declines caused by reserved RPC capacity.
227
227
 
228
+ The counter counts decline events. Dex may decline the same Step dispatch
229
+ repeatedly while retrying it, so read the counter as a rate.
228
230
  Full-capacity rejections and unknown names are excluded. The count never
229
231
  resets during this Worker's lifetime and stays available after stop.
230
232
  Reading it alongside handlers is safe. Listening Workers return zero.