Remove pickle from parallel external source shared memory messages

[devin-ai-integration]

Category:

Bug fix (non-breaking change which fixes an issue) — security hardening

Description:

Parallel ExternalSource (PES) exchanged its internal task/result messages and sample meta-data through shared memory encoded with pickle. Anyone able to modify a shared-memory chunk (e.g. a compromised worker) could make the other side execute arbitrary code on pickle.loads.

This PR replaces pickle in the PES shared-memory protocol with a small, fixed-layout, little-endian struct-based format (_BinaryWriter / _BinaryReader in shared_batch.py):

ScheduledTask  = u8 tag=1 | q context_i | q scheduled_i | q epoch_start | q minibatch_i
                 | u8 kind: 0 -> 6q SampleRange(start, end, iteration, epoch, slice_start, slice_end)
                            1 -> u8 arg: 0 -> ()  | 1 -> q int  | 2 -> qq BatchInfo(iteration, epoch_idx)
CompletedTask  = u8 tag=2 | q worker_id | q context_i | q scheduled_i | q minibatch_i
                 | ? has_meta [QQ meta_offset, meta_size]
                 | u8 exc: 0 none | 1 StopIteration(str) | 2 RuntimeError(str)
                 | ? has_tb [str traceback]
samples meta   = Q count | sample*
sample         = u8 0 -> QQ offset, nbytes | str dtype.str | I ndim | ndim*Q shape
               | u8 1/2 -> Q n | sample*  (tuple / list, nesting limited to 64)

The reader validates every read against the buffer bounds, rejects unknown tags, trailing bytes, invalid SampleRanges, invalid / object dtypes and overly deep nesting, raising RuntimeError("Malformed shared memory message: ..."). Worker exceptions were already normalized to StopIteration/RuntimeError by CompletedTask.failed, so the format carries only those two types plus their message.

Additionally (nice-to-have from the ticket), ShmHandle::CreateHandle now uses memfd_create so the shared memory never has a name in the filesystem, even briefly; it falls back to the previous mkstemp + shm_open + shm_unlink path when memfd_create is unavailable (ENOSYS, or EPERM/EACCES/EINVAL from seccomp/old kernels). The fd is still passed to workers over Unix sockets, which works the same for memfds.

Callback passing reassessment: callbacks (and the optional py_callback_pickler output) are still pickled, but only in the parent -> worker direction, as multiprocessing.Process arguments at worker start-up (spawn/forkserver) — they never go through the shared-memory chunks and the parent never unpickles anything coming from workers. Executing the user's callback is arbitrary code execution by design, so this is left unchanged.

Additional information:

Affected modules and functionalities:

• dali/python/nvidia/dali/_multiproc/shared_batch.py: new binary codec (serialize_message / deserialize_message, serialize_samples_meta / deserialize_samples_meta); read_shm_message, write_shm_message and SharedBatchWriter use it. The previous restricted unpickler is removed.
• dali/python/nvidia/dali/_multiproc/messages.py, worker.py: docstrings only.
• dali/core/os/shared_mem.cc: memfd_create with shm_open fallback.

Key points relevant for the review:

• Batch-mode task arguments are limited to what ExternalSource actually produces: no argument, an int iteration, or BatchInfo.
• Integer fields are encoded as signed/unsigned 64-bit; sample shapes/offsets as unsigned 64-bit.
• The C++ change could only be syntax/lint-checked locally (no full build); CI will exercise it.

Tests:

• Existing tests apply
  • test_external_source_parallel_shared_batch.py, test_pool.py, test_external_source_parallel*.py
• New tests added
  • Python tests: test_external_source_parallel_shared_batch.py — round trips of ScheduledTask (sample range, sliced range, int / BatchInfo / empty batch args), CompletedTask (success/failure), nested tuple/list samples and more dtypes; rejection of pickle payloads, truncated/trailing data, bad tags, invalid dtypes, object dtypes, deep nesting.
  • GTests
  • Benchmark
  • Other
• N/A

Checklist

Documentation

• Existing documentation applies
• Documentation updated
  • Docstring
  • Doxygen
  • RST
  • Jupyter
  • Other
• N/A

DALI team only

Requirements

• Implements new requirements
• Affects existing requirements
• N/A

REQ IDs: N/A

JIRA TASK: DALI-4559

Link to Devin session: https://nvidia-cloud.devinenterprise.com/sessions/28dadd725c2143758024982108533e2c
Open in Devin Desktop: https://nvidia-cloud.devinenterprise.com/desktop/session/28dadd725c2143758024982108533e2c?variant=devin

[devin-ai-integration]

I'll fix CI failures and address comments from users with write access. I'll skip comments containing "(aside)".

• Disable automatic comment, CI, and merge conflict monitoring

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[devin-ai-integration]

Intentional: _read_batch_args returns the positional-argument tuple for the batch callback, which is either empty (callback takes no argument) or holds a single iteration / BatchInfo.

[greptile-apps]

Confidence Score: 3/5

[High risk] Replaces pickle serialization with binary format in inter-process messages.

The PR is not ready to merge until nesting failures are reported through the task result and the protocol tags satisfy the repository's enum requirement.

Findings

1.  Nested samples lose their error ▶
2.  Protocol tags use magic integers ▶

Summary

The PR replaces pickle in parallel ExternalSource shared-memory messages and sample metadata with a bounded binary format, and adds an anonymous memfd path with a fallback.

• The follow-up changes align writer and reader nesting limits and make nested test comparisons strict.
• Serialization errors from the new writer guard still need to reach the parent as meaningful task failures.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  C[Worker callback] --> D[Dispatcher thread]
  D --> W[SharedBatchWriter]
  W -->|Valid metadata| Q[CompletedTask queue]
  W -->|Nesting error| X[Dispatcher exits]
  X --> P[Parent receives worker-interruption error]

Loading

_{Reviews (2) · Last reviewed commit: "Apply sample nesting limit when writing ..."}

[greptile-apps]

Nested samples lose their error

When a parallel callback returns a sample nested more than 64 levels deep, this guard raises in the worker's dispatcher thread, outside the callback-error handler. The dispatcher sends no CompletedTask, so the parent reports a generic worker-interruption error instead of explaining the nesting problem. Catch serialization failures in the dispatcher and return them as failed tasks.

Knowledge Base Used: Python API and framework iterators

[greptile-apps]

Protocol tags use magic integers

The new protocol uses integer constants to dispatch message, task, argument, exception, and sample kinds. This violates the repository directive to use Enum or IntEnum for type and mode dispatch. Use typed tag groups so the protocol's cases remain explicit.

Rule Used: Use enums (or IntEnum/Enum in Python) for type/mode/format dispatch — not magic strings or magic integers. Self-documenting and compiler-checkable. (source)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!