Add ExternalSource support to transparent pipelining

[rostan-t]

Category:

New feature (non-breaking change which adds functionality)

Description:

#6395 adds ndd.ExternalSource. This PR makes it usable with transparent pipelining.

It's usable in two ways:

1. Through .compiled(batch_size) to iterate on it.
2. By calling it directly in a compiled loop to add it as a co-source.

In the second case, we need to make sure that every source seen during tracing is called (exactly once) during each compiled iteration as the executor pulls data anyway, which could create misaligned states.

This PR generalizes the concept of source in _compile.py and decouples CompileContext from readers. The SupportsCompile protocol is implemented by both Reader and ExternalSource allowing both to be used as sources through duck typing.

Additional information:

Affected modules and functionalities:

Transparent pipelining.

Key points relevant for the review:

This is a follow-up for #6395.

Tests:

• Existing tests apply
• New tests added
  • Python tests
  • 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-4745

[greptile-apps]

Greptile Summary

This PR extends transparent pipelining to support ExternalSource as both a compiled loop root (via .compiled(batch_size)) and as a co-source feeder (called inside another source's compiled loop). The implementation generalises the existing CompileContext/CompiledEpochIterator machinery behind a SupportsCompile protocol so that both Reader and ExternalSource participate without special-casing.

• _compile.py: CompileSource becomes a dataclass tracking multiple sources; CompileContext moves from a single source to a sources list with underrun detection (_root_stopped/_require_consumed); CompiledEpochIterator becomes an ABC with concrete _ReaderEpochIterator and _ExternalSourceEpochIterator subclasses.
• _external_source.py: Adds _Role enum, compiled() method, _trace_pull/_compiled_call dispatch, and _source_callback used by DALI's fn.external_source node wired into the compiled pipeline.
• _ops.py / _op_builder.py: Reader implements the SupportsCompile protocol; _original_tensor_args is removed in favour of _raw_tensor_args fed through the new _wire_arg helper.
• Tests: 25+ new tests covering cycle semantics, multi-output, broadcast, GPU placement, break/reuse, role-conflict errors, and feeder underrun/over-read detection.

Confidence Score: 5/5

Safe to merge; the refactoring is well-structured and the new ExternalSource compile paths are guarded with thorough error detection.

The multi-source compile context, role-locking, underrun detection, and teardown paths all look correct. The tracing/compiled dispatch for both Reader and ExternalSource is clean. Tests cover cycle semantics, multi-output, broadcast, GPU, break/reuse, and all error scenarios. One minor style nit (zip without strict) is the only issue found.

No files require special attention.

Important Files Changed

Filename	Overview
dali/python/nvidia/dali/experimental/dynamic/_compile.py	Refactors CompileContext to support multiple sources; introduces SupportsCompile protocol, CompiledEpochIterator ABC, and splits into _ReaderEpochIterator/_ExternalSourceEpochIterator. Offset assignment, pipeline wiring, and stop/underrun detection all look correct.
dali/python/nvidia/dali/experimental/dynamic/_external_source.py	Adds _Role enum and compiled-mode support (_trace_pull, _compiled_call, _source_callback, _wire_pipeline). Role transitions, source locking, and teardown paths are handled correctly.
dali/python/nvidia/dali/experimental/dynamic/_ops.py	Extracts _wire_arg, adds _teardown_compile/_wire_pipeline/_shape_result/_transfer_into/_make_epoch_iterator to Reader; removes _original_tensor_args in favour of _raw_tensor_args. One style nit: zip without strict in _shape_result.
dali/python/nvidia/dali/experimental/dynamic/_op_builder.py	Removes _original_tensor_args tracking; minor import reordering. No logic changes.
dali/test/python/experimental_mode/test_compile.py	Adds 25+ new tests covering ES as root and feeder with correctness assertions and error-path coverage.
dali/test/python/experimental_mode/test_external_source.py	Adds test_device_resolved_at_call (multi-GPU) with correct guard.
dali/test/python/experimental_mode/ndd_utils.py	Adds TypeGuard annotation to _is_compiled; minor import reordering.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant User as User loop
    participant ES as ExternalSource root
    participant Feeder as ExternalSource feeder
    participant Iter as EpochIterator
    participant Ctx as CompileContext
    participant Pipe as DALI Pipeline

    User->>ES: "compiled(batch_size=4)"
    ES->>Iter: make_iterator(batch_size)
    Iter->>Ctx: new CompileContext(batch_size)
    Note over Iter,Ctx: TRACING epoch
    Iter->>ES: _trace_pull(ctx)
    ES->>Ctx: add_source sources[0]
    ES->>Ctx: _mark_read sources[0]
    Iter->>User: first_batch via _emit_step
    User->>Feeder: __call__()
    Feeder->>Ctx: add_source sources[1]
    Feeder->>Ctx: _mark_read sources[1]
    Feeder->>User: feeder_batch
    User->>Iter: resume
    Iter->>Ctx: _require_consumed checks sources[1:]
    Iter->>Ctx: build_pipeline
    Ctx->>Pipe: wire fn.external_source and build
    Note over Iter,Pipe: State = COMPILED
    loop Each compiled iteration
        Iter->>Ctx: run_pipeline
        Ctx->>Pipe: pipeline.run
        Pipe->>ES: _source_callback prefetch
        Pipe->>Feeder: _source_callback prefetch
        Pipe->>Ctx: pipeline_outputs
        Iter->>User: batch via _emit_step
        User->>Feeder: __call__() compiled_call
        Feeder->>Ctx: _mark_read sources[1]
        User->>Iter: resume
        Iter->>Ctx: _require_consumed
    end
    Note over ES: StopIteration marks _root_stopped
    Iter->>Pipe: pipeline.reset for next epoch

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant User as User loop
    participant ES as ExternalSource root
    participant Feeder as ExternalSource feeder
    participant Iter as EpochIterator
    participant Ctx as CompileContext
    participant Pipe as DALI Pipeline

    User->>ES: "compiled(batch_size=4)"
    ES->>Iter: make_iterator(batch_size)
    Iter->>Ctx: new CompileContext(batch_size)
    Note over Iter,Ctx: TRACING epoch
    Iter->>ES: _trace_pull(ctx)
    ES->>Ctx: add_source sources[0]
    ES->>Ctx: _mark_read sources[0]
    Iter->>User: first_batch via _emit_step
    User->>Feeder: __call__()
    Feeder->>Ctx: add_source sources[1]
    Feeder->>Ctx: _mark_read sources[1]
    Feeder->>User: feeder_batch
    User->>Iter: resume
    Iter->>Ctx: _require_consumed checks sources[1:]
    Iter->>Ctx: build_pipeline
    Ctx->>Pipe: wire fn.external_source and build
    Note over Iter,Pipe: State = COMPILED
    loop Each compiled iteration
        Iter->>Ctx: run_pipeline
        Ctx->>Pipe: pipeline.run
        Pipe->>ES: _source_callback prefetch
        Pipe->>Feeder: _source_callback prefetch
        Pipe->>Ctx: pipeline_outputs
        Iter->>User: batch via _emit_step
        User->>Feeder: __call__() compiled_call
        Feeder->>Ctx: _mark_read sources[1]
        User->>Iter: resume
        Iter->>Ctx: _require_consumed
    end
    Note over ES: StopIteration marks _root_stopped
    Iter->>Pipe: pipeline.reset for next epoch

Loading

_{Reviews (10): Last reviewed commit: "Refactor compile loop stop handling and ..." | Re-trigger Greptile}

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [55677891]: BUILD STARTED

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [55678954]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [55678954]: BUILD PASSED

[banasraf]

This changes when device="gpu" is resolved. Previously we stored the string and as_tensor/as_batch resolved it when the source was called, inside the active EvalContext. Now _as_device(device) resolves it in the constructor and captures whichever GPU is current at that time.

For example, if the source is created on GPU 0 and later compiled with EvalContext(device_id=1), its data will still be created on GPU 0 while the pipeline runs on GPU 1.

Could we keep the device unresolved until execution, or resolve it when the compiled context becomes active? It would also be useful to test creating the source outside an explicit EvalContext and running it inside one.

[rostan-t]

Fixed by be28803

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [56746334]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [56746334]: BUILD FAILED

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [57248017]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [57248017]: BUILD FAILED

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [57391158]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [57391158]: BUILD PASSED