Skip to content

feat(python): payload codecs with global chain and per-task registry#365

Merged
pratyush618 merged 5 commits into
masterfrom
feat/python-codecs
Jul 6, 2026
Merged

feat(python): payload codecs with global chain and per-task registry#365
pratyush618 merged 5 commits into
masterfrom
feat/python-codecs

Conversation

@pratyush618

@pratyush618 pratyush618 commented Jul 5, 2026

Copy link
Copy Markdown
Collaborator

Summary

Adds a payload-codec layer to the Python SDK with wire formats that are part of the cross-SDK contract, so codec-framed payloads interoperate across SDKs.

  • New taskito/codecs.py: PayloadCodec protocol, GzipCodec (zip-bomb-guarded decode, 64 MiB default cap), AesGcmCodec ([12B nonce][ciphertext||16B tag], encryption extra), HmacCodec ([32B mac][body], constant-time compare), and CodecSerializer (encode in list order, decode in reverse).
  • Queue(codec=...) global chain wraps the queue serializer, covering payloads and results; Queue(codecs={...}) + @queue.task(codecs=[...]) applies named codecs per task (payload only — results stay on the queue serializer).
  • New CryptoError(SerializationError) for decrypt/verify failures.
  • Dedup keys hash the pre-codec bytes so idempotent=True stays deterministic under per-task encryption codecs.

Result-path fix

Worker paths previously wrote results as raw cloudpickle regardless of the configured serializer, while JobResult.result() read them back with the queue serializer — silently broken for any non-default serializer, and incompatible with a codec chain. All four result-write sites (thread-pool worker and native-async executor in Rust, async_support executor, prefork child) now route through a new Queue._serialize_result(), and the saga orchestrator's result read is aligned to the queue-level serializer.

Backward compatibility: results persisted before this change are raw cloudpickle; the default SmartSerializer reads untagged legacy bytes, so existing queues upgrade cleanly. Non-default serializers could never read results correctly before, so this is strictly an improvement. Jobs persisted before enabling a codec chain cannot be decoded through it — drain the queue first.

Tests

  • 28 codec tests (tests/core/test_codecs.py): unit round-trips, tamper/short/cap rejection, chain reversibility, worker e2e for global chain and per-task codecs, idempotency-under-encryption regression.
  • 3 result-path regressions (tests/core/test_result_serialization.py) plus updated native-async executor tests.
  • Full suite: 1110 passed, 4 skipped; Rust workspace tests green.

@coderabbitai

coderabbitai Bot commented Jul 5, 2026

Copy link
Copy Markdown

Review Change Stack

Warning

Review limit reached

You’ve reached a temporary PR review limit under our Fair Usage Limits Policy.

Your recent review volume is higher than typical usage, so adaptive limits are currently applied.

Next review available in: 30 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 622d42a5-171f-4c7a-9735-760229b01d44

📥 Commits

Reviewing files that changed from the base of the PR and between cc7a01c and 40e577a.

📒 Files selected for processing (2)
  • sdks/python/taskito/codecs.py
  • sdks/python/tests/core/test_codecs.py
📝 Walkthrough

Walkthrough

Introduces a composable payload codec framework (Gzip, AES-GCM, HMAC codecs plus a CodecSerializer wrapper and CryptoError) integrated into Queue's construction, enqueue/dedup ordering, per-task codec registration, workflow/fan-out payload encoding, and result serialization across native async, async executor, and prefork worker paths, with accompanying tests.

Changes

Payload codec framework and queue integration

Layer / File(s) Summary
Codec classes, exceptions, and public exports
sdks/python/taskito/codecs.py, sdks/python/taskito/exceptions.py, sdks/python/taskito/__init__.py
Adds PayloadCodec protocol, GzipCodec, AesGcmCodec, HmacCodec, CodecSerializer, and CryptoError, re-exported from the package.
Queue construction, payload encoding, and enqueue/dedup ordering
sdks/python/taskito/app.py
Adds codec/codecs parameters to Queue.__init__, new _apply_task_codecs/_encode_payload helpers, reorders enqueue/enqueue_many dedup-key hashing to occur on pre-codec bytes, applies codecs before size validation, and updates batching and result serialization paths.
Per-task codec registration and predicate re-enqueue
sdks/python/taskito/mixins/decorators.py, sdks/python/taskito/mixins/predicates.py
Extends task() decorator with a codecs parameter, wraps per-task serializers with CodecSerializer, and routes periodic and defer re-enqueue payloads through _encode_payload.
Result serialization via queue hooks
crates/taskito-python/src/native_async/task_executor.rs, crates/taskito-python/src/py_worker.rs, sdks/python/taskito/async_support/executor.py, sdks/python/taskito/prefork/child.py, sdks/python/taskito/workflows/saga/orchestrator.py
Routes result serialization through queue._serialize_result (falling back to cloudpickle) and updates saga orchestrator to use the queue-level serializer for forward payload deserialization.
Workflow and fan-out codec-encoded payloads
sdks/python/taskito/workflows/builder.py, sdks/python/taskito/workflows/tracker/fan_out.py
Uses queue._encode_payload/_apply_task_codecs for workflow node, fan-out, and fan-in payloads.
Codec and serialization test suites
sdks/python/tests/core/test_codecs.py, sdks/python/tests/core/test_result_serialization.py, sdks/python/tests/worker/test_native_async.py
Adds codec unit/integration tests and updates async executor tests to configure and assert _serialize_result usage.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant Queue
  participant CodecChain
  participant Storage
  Caller->>Queue: enqueue(task_name, args, kwargs)
  Queue->>Queue: serialize (args, kwargs)
  Queue->>Queue: compute dedup key from pre-codec bytes
  Queue->>CodecChain: _apply_task_codecs(task_name, payload)
  CodecChain-->>Queue: encoded payload
  Queue->>Queue: validate payload size (post-codec)
  Queue->>Storage: persist job with encoded payload
Loading
sequenceDiagram
  participant Executor
  participant QueueRef
  participant Cloudpickle
  Executor->>QueueRef: check _queue_ref
  alt queue present
    Executor->>QueueRef: _serialize_result(task_name, result)
    QueueRef-->>Executor: serialized bytes
  else no queue
    Executor->>Cloudpickle: dumps(result)
    Cloudpickle-->>Executor: serialized bytes
  end
Loading

Possibly related PRs

  • ByteVeda/taskito#351: Both PRs introduce a payload codec concept with a codec-chain wrapper around a serializer, using matching PayloadCodec/CodecSerializer abstractions.
  • ByteVeda/taskito#353: Both PRs implement per-task payload codec chaining/selection with named codecs applied through enqueue and reversed on the worker.

Suggested labels: python, tests

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly summarizes the main change: adding Python payload codecs with a global chain and per-task registry.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch feat/python-codecs

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
sdks/python/taskito/app.py (1)

343-365: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win

_encode_payload() should enforce the payload cap itself

Several callers bypass the external _check_payload_size() step (sdks/python/taskito/mixins/predicates.py::_reenqueue_after_defer, sdks/python/taskito/workflows/builder.py::_compile, and sdks/python/taskito/workflows/tracker/fan_out.py), so oversized task payloads can still reach the backend. Moving the size check into _encode_payload() would make the limit consistent and let _dispatch_batched_payload() drop its extra guard.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@sdks/python/taskito/app.py` around lines 343 - 365, `_encode_payload()` in
`app.py` should enforce `max_payload_bytes` directly so oversized payloads are
rejected even when callers skip `_check_payload_size()`. Add the size check
after serialization/codecs in `_encode_payload()` using the existing
`_check_payload_size()` helper and task name, and then remove the redundant
guard from `_dispatch_batched_payload()` so all paths go through the same limit
enforcement.
🧹 Nitpick comments (3)
crates/taskito-python/src/native_async/task_executor.rs (1)

152-165: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Duplicated result-serialization block with py_worker.rs.

This exact block (None-check, _queue_ref lookup, _serialize_result/cloudpickle.dumps fallback, .extract()) is duplicated verbatim in crates/taskito-python/src/py_worker.rs (lines 210-223). Consider extracting a shared helper (e.g. serialize_task_result(context_mod, job, result) -> PyResult<Option<Vec<u8>>>) in a common module so both worker paths stay in sync as this codec-aware logic evolves.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/taskito-python/src/native_async/task_executor.rs` around lines 152 -
165, The result-serialization logic in task_executor.rs is duplicated in
py_worker.rs, so extract the None-check, _queue_ref lookup,
_serialize_result/cloudpickle.dumps fallback, and byte extraction into a shared
helper such as serialize_task_result(context_mod, job, result) returning
PyResult<Option<Vec<u8>>>. Update both TaskExecutor and the corresponding
py_worker path to call the shared helper so the codec-aware behavior stays
consistent in one place.
crates/taskito-python/src/py_worker.rs (1)

210-223: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Same duplicated serialization block as native_async/task_executor.rs.

See the companion comment on crates/taskito-python/src/native_async/task_executor.rs (lines 152-165) — this block is byte-for-byte identical. Worth extracting to a shared helper to avoid drift between the two worker paths.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/taskito-python/src/py_worker.rs` around lines 210 - 223, The result
serialization logic in `py_worker.rs` is duplicated verbatim in the native async
worker path, so extract the shared behavior into a common helper and have both
`py_worker` and `native_async/task_executor` call it. Keep the helper
responsible for the `result.is_none()` handling, `_queue_ref` lookup,
`_serialize_result` vs `cloudpickle.dumps` fallback, and `bytes` extraction so
the two worker paths stay in sync.
sdks/python/taskito/mixins/decorators.py (1)

641-711: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

periodic() doesn't expose a codecs= param, so scheduled tasks can't opt into per-task codecs.

Both the plain-task branch (line 694: self.task(name=name, queue=queue)(fn)) and the workflow-launcher branch call self.task(...) without forwarding a codecs= list, so a periodic task can never get an entry in _task_codecs, even though _encode_payload is now used to build its payload. This mirrors the pre-existing serializer= gap, but is worth closing now that codec support is being wired through the queue.

♻️ Possible fix
 def periodic(
     self,
     cron: str,
     name: str | None = None,
     args: tuple = (),
     kwargs: dict | None = None,
     queue: str = "default",
     timezone: str | None = None,
+    codecs: list[str] | None = None,
 ) -> Callable[[Callable[..., Any]], TaskWrapper]:
     ...
-            wrapper = self.task(name=name, queue=queue)(fn)
+            wrapper = self.task(name=name, queue=queue, codecs=codecs)(fn)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@sdks/python/taskito/mixins/decorators.py` around lines 641 - 711, The
periodic decorator in `periodic()` is not forwarding per-task codec
configuration, so scheduled tasks can’t register their codecs even though
payloads are encoded through `_encode_payload`. Update both branches inside
`decorator()`—the normal task registration via `self.task(...)` and the workflow
launcher registration via `@self.task(...)`—to accept and pass through a
`codecs=` parameter. Make sure the codec list is propagated into the underlying
task registration so entries are created alongside `_periodic_configs` and the
generated launcher task stays consistent with `TaskWrapper` setup.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@sdks/python/taskito/app.py`:
- Around line 343-365: `_encode_payload()` in `app.py` should enforce
`max_payload_bytes` directly so oversized payloads are rejected even when
callers skip `_check_payload_size()`. Add the size check after
serialization/codecs in `_encode_payload()` using the existing
`_check_payload_size()` helper and task name, and then remove the redundant
guard from `_dispatch_batched_payload()` so all paths go through the same limit
enforcement.

---

Nitpick comments:
In `@crates/taskito-python/src/native_async/task_executor.rs`:
- Around line 152-165: The result-serialization logic in task_executor.rs is
duplicated in py_worker.rs, so extract the None-check, _queue_ref lookup,
_serialize_result/cloudpickle.dumps fallback, and byte extraction into a shared
helper such as serialize_task_result(context_mod, job, result) returning
PyResult<Option<Vec<u8>>>. Update both TaskExecutor and the corresponding
py_worker path to call the shared helper so the codec-aware behavior stays
consistent in one place.

In `@crates/taskito-python/src/py_worker.rs`:
- Around line 210-223: The result serialization logic in `py_worker.rs` is
duplicated verbatim in the native async worker path, so extract the shared
behavior into a common helper and have both `py_worker` and
`native_async/task_executor` call it. Keep the helper responsible for the
`result.is_none()` handling, `_queue_ref` lookup, `_serialize_result` vs
`cloudpickle.dumps` fallback, and `bytes` extraction so the two worker paths
stay in sync.

In `@sdks/python/taskito/mixins/decorators.py`:
- Around line 641-711: The periodic decorator in `periodic()` is not forwarding
per-task codec configuration, so scheduled tasks can’t register their codecs
even though payloads are encoded through `_encode_payload`. Update both branches
inside `decorator()`—the normal task registration via `self.task(...)` and the
workflow launcher registration via `@self.task(...)`—to accept and pass through
a `codecs=` parameter. Make sure the codec list is propagated into the
underlying task registration so entries are created alongside
`_periodic_configs` and the generated launcher task stays consistent with
`TaskWrapper` setup.

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: ad0e23e8-a503-4931-88dc-f84c17de9817

📥 Commits

Reviewing files that changed from the base of the PR and between 3ae7945 and cc7a01c.

📒 Files selected for processing (16)
  • crates/taskito-python/src/native_async/task_executor.rs
  • crates/taskito-python/src/py_worker.rs
  • sdks/python/taskito/__init__.py
  • sdks/python/taskito/app.py
  • sdks/python/taskito/async_support/executor.py
  • sdks/python/taskito/codecs.py
  • sdks/python/taskito/exceptions.py
  • sdks/python/taskito/mixins/decorators.py
  • sdks/python/taskito/mixins/predicates.py
  • sdks/python/taskito/prefork/child.py
  • sdks/python/taskito/workflows/builder.py
  • sdks/python/taskito/workflows/saga/orchestrator.py
  • sdks/python/taskito/workflows/tracker/fan_out.py
  • sdks/python/tests/core/test_codecs.py
  • sdks/python/tests/core/test_result_serialization.py
  • sdks/python/tests/worker/test_native_async.py

@pratyush618
pratyush618 merged commit 50c923b into master Jul 6, 2026
35 checks passed
@pratyush618
pratyush618 deleted the feat/python-codecs branch July 6, 2026 00:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant