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.

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