Restrict worker result deserialization

[JanuszL]

Category:

Bug fix

Description:

Restrict deserialization of worker-to-parent shared-memory results to the internal multiprocessing protocol object graph.

Addresses:
The multiprocessing subsystem uses pickle.loads() to deserialize data from shared memory regions written by worker processes. In shared_batch.py, both deserialize_sample_meta() (line 162) and read_shm_message() (line 296) call pickle.loads() on data read from shared memory buffers. If a worker process is compromised (e.g., through a malicious data source callback), it could write crafted pickle payloads to shared memory that would achieve arbitrary code execution in the parent process when deserialized.

Additional information:

Affected modules and functionalities:

• Parallel ExternalSource worker result messages and shared-batch metadata

Key points relevant for the review:

• Parent-to-worker task data remains trusted and supports generic user batch arguments.
• Worker-to-parent result data accepts only internal result/metadata classes, NumPy dtypes, RuntimeError, and StopIteration.
• Non-StopIteration worker exceptions are normalized before serialization.

Tests:

• Existing tests apply
• New tests added
  • Python tests
    • test_external_source_parallel_shared_batch.test_restricted_unpickler_rejects_unsafe_globals
      -test_external_source_parallel_shared_batch.test_restricted_unpickler_normalizes_stop_iteration_subclasses
  • 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: N/A

[greptile-apps]

Greptile Summary

The PR restricts worker-result deserialization to the internal multiprocessing protocol and normalizes worker exceptions into allowlisted built-in types.

• Adds a restricted unpickler for worker-to-parent messages and shared-batch metadata.
• Preserves trusted generic deserialization for parent-to-worker task arguments.
• Converts StopIteration subclasses to built-in StopIteration while retaining termination semantics.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains from the previously reported StopIteration handling issues.

The current normalization converts StopIteration subclasses into an allowlisted built-in StopIteration, and the parent still recognizes the deserialized exception as end-of-data.

Important Files Changed

Filename	Overview
dali/python/nvidia/dali/_multiproc/messages.py	Normalizes arbitrary worker exceptions to RuntimeError and StopIteration, fully addressing the previously reported subclass behavior.
dali/python/nvidia/dali/_multiproc/shared_batch.py	Introduces an allowlisted worker-result unpickler while retaining generic deserialization for trusted parent-originated tasks.
dali/test/python/test_external_source_parallel_shared_batch.py	Tests rejection of unsafe globals and verifies that StopIteration subclasses retain termination semantics after serialization.

Sequence Diagram

sequenceDiagram
    participant Parent
    participant Queue as Shared-memory queue
    participant Worker
    Parent->>Queue: "Task descriptor (worker_id = -1)"
    Queue->>Worker: Trusted task payload
    Worker->>Queue: CompletedTask and batch metadata
    Queue->>Parent: Worker result descriptor
    Parent->>Parent: Restricted protocol deserialization

Loading

_{Reviews (3): Last reviewed commit: "Preserve worker StopIteration semantics" | Re-trigger Greptile}

[JanuszL]

!build

[JanuszL]

!build

[dali-automaton]

CI MESSAGE: [61556838]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [61556838]: BUILD PASSED

[rostan-t]

Why trust the parent but not the children? Is it because it's really impractical (or impossible) to add validation for the types passed by the parent or is there something else?

[JanuszL]

The trust boundary is deliberately worker → parent: a compromised worker can write a forged shared-memory result that the parent would otherwise deserialize. Parent → worker messages originate in the controlling process and may contain arbitrary user batch arguments, so restricting them to the internal result protocol would break supported inputs.