Create stub files for the dynamic mode #6089

rostan-t · 2025-11-18T17:41:06Z

Category: New feature (non-breaking change which adds functionality)

Description:

Currently stub (*.pyi) files are only generated for "fn" and the legacy "ops" API. This PR adds stubs generation for dynamic mode.

Additional information:

Affected modules and functionalities:

New __init__.pyi files get created during the build process.

Key points relevant for the review:

Does the IDE / language server complain when using the API normally?
Does the IDE / language server properly report incorrect types?
Does the IDE properly suggest code completion?
You can use typing.reveal_type to see what type the type checker infers. For example:

Tests:

Checklist

Documentation

DALI team only

Requirements

Implements new requirements
Affects existing requirements
N/A

REQ IDs: N/A

JIRA TASK: DALI-4439

dali/python/nvidia/dali/_typing.py

dali/python/nvidia/dali/ops/_signatures.py

dali/python/nvidia/dali/_typing.py

dali/python/nvidia/dali/ops/_signatures.py

greptile-apps · 2025-11-21T17:03:20Z

Greptile Overview

Greptile Summary

This PR extends DALI's stub file (.pyi) generation to support the dynamic mode API, which previously only had stubs for the fn and ops APIs.

Key changes:

Added comprehensive stub generation for dynamic mode operators and functions in _signatures.py with proper handling of overloaded signatures based on input types (Tensor, Batch, TensorLike)
Correctly splits keyword arguments between __init__ (non-tensor kwargs) and __call__ (tensor kwargs only) to match runtime behavior in _op_builder.py
Added DLPack protocol and TensorLike type alias to _typing.py for dynamic mode type annotations
Made __cuda_array_interface__ a property in dynamic Tensor class for protocol compatibility
Updated fn and ops modules to pass api parameter for proper signature generation
Added _generated attribute to dynamically created operator classes/functions for stub detection
Modified stub generator script to include dynamic mode in the build process

The implementation properly handles the complex type system for dynamic mode, including tensor/batch broadcasting semantics and proper type narrowing through overloads.

Confidence Score: 4/5

This PR is safe to merge with very low risk - changes are well-structured and previous critical issues have been addressed
The PR adds a significant new feature (dynamic mode stub generation) with careful attention to type correctness. All previously identified critical issues (builtins module handling, tensor/non-tensor kwarg split, API parameter passing) have been fixed in recent commits. The main signature generation file (_signatures.py) is complex but the logic correctly mirrors the runtime behavior. Minor complexity in the code is inherent to handling multiple overloads and API variants.
Pay close attention to dali/python/nvidia/dali/ops/_signatures.py due to its complexity, though no critical issues were found

Important Files Changed

File Analysis

Filename	Score	Overview
dali/python/nvidia/dali/ops/_signatures.py	4/5	Adds comprehensive stub generation for dynamic mode with proper tensor/non-tensor kwarg split, API parameter threading, and type annotations
dali/python/nvidia/dali/experimental/dynamic/_op_builder.py	5/5	Adds `_generated` attribute to dynamically built operator classes and functions for stub generation detection
internal_tools/python_stub_generator.py	5/5	Adds dynamic mode stub generation to the build process

Sequence Diagram

sequenceDiagram
    participant Build as Build Process
    participant StubGen as python_stub_generator.py
    participant Sigs as _signatures.py
    participant OpBuilder as _op_builder.py
    participant Backend as DALI Backend
    
    Build->>StubGen: Generate stubs for wheel
    StubGen->>Sigs: gen_all_signatures(path, "fn")
    StubGen->>Sigs: gen_all_signatures(path, "ops")
    StubGen->>Sigs: gen_all_signatures(path, "dynamic")
    
    Sigs->>Backend: _registry._all_registered_ops()
    Backend-->>Sigs: List of schema names
    
    loop For each operator schema
        Sigs->>Backend: _b.TryGetSchema(schema_name)
        Backend-->>Sigs: OpSchema
        
        alt Dynamic API
            Sigs->>Sigs: _gen_dynamic_signature(schema, name)
            Sigs->>Sigs: _gen_dynamic_cls_signature() or _gen_dynamic_fun_signature()
            Sigs->>Sigs: _gen_dynamic_call_signature(include_only_inputs=True)
            Note over Sigs: __call__ gets tensor kwargs only
            Sigs->>Sigs: _call_signature(include_kwarg_inputs=False)
            Note over Sigs: __init__ gets non-tensor kwargs only
        else Fn API
            Sigs->>Sigs: _gen_fn_signature(schema, name)
            Sigs->>Sigs: _call_signature(api="fn", include_inputs=True)
        else Ops API
            Sigs->>Sigs: _gen_ops_signature(schema, name)
            Sigs->>Sigs: _call_signature(api="ops") for __init__ and __call__
        end
        
        Sigs->>Sigs: Write signature to .pyi file
    end
    
    alt Dynamic mode operators
        OpBuilder->>OpBuilder: build_operator_class(schema)
        OpBuilder->>OpBuilder: Set _generated = True
        Note over OpBuilder: Runtime: __init__ skips tensor args,<br/>__call__ includes tensor args only
    end
    
    Sigs-->>StubGen: Stub files created
    StubGen-->>Build: Stubs ready for packaging

greptile-apps

_{7 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

dali/python/nvidia/dali/ops/_signatures.py

greptile-apps

_{7 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

dali/python/nvidia/dali/ops/_signatures.py

greptile-apps

_{7 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

dali/python/nvidia/dali/ops/_signatures.py

greptile-apps

_{7 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

dali/python/nvidia/dali/ops/_signatures.py

rostan-t · 2025-11-24T09:41:56Z

!build

dali-automaton · 2025-11-24T09:48:17Z

CI MESSAGE: [39065210]: BUILD STARTED

dali-automaton · 2025-11-24T09:49:29Z

CI MESSAGE: [39065210]: BUILD FAILED

dali-automaton · 2025-11-24T09:52:46Z

CI MESSAGE: [39065347]: BUILD STARTED

dali-automaton · 2025-11-24T10:48:30Z

CI MESSAGE: [39065347]: BUILD FAILED

dali-automaton · 2025-11-24T11:49:11Z

CI MESSAGE: [39070781]: BUILD STARTED

greptile-apps

_{7 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

dali/python/nvidia/dali/ops/_signatures.py

dali-automaton · 2025-11-24T12:11:42Z

CI MESSAGE: [39070781]: BUILD FAILED

dali-automaton · 2025-11-24T14:43:45Z

CI MESSAGE: [39076892]: BUILD STARTED

dali-automaton · 2025-11-24T15:09:35Z

CI MESSAGE: [39076892]: BUILD FAILED

dali-automaton · 2025-11-24T15:46:57Z

CI MESSAGE: [39079519]: BUILD STARTED

greptile-apps

_{8 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

dali-automaton · 2025-11-24T16:38:38Z

CI MESSAGE: [39079519]: BUILD FAILED

rostan-t · 2025-11-25T09:16:39Z

!build

dali/python/nvidia/dali/ops/_signatures.py

…n definitions Signed-off-by: Rostan Tabet <[email protected]>

Signed-off-by: Rostan Tabet <[email protected]>

The stubs generation script imports dali.exerimental.dynamic, which requires NVTX and makefun to be importable. Signed-off-by: Rostan Tabet <[email protected]>

Signed-off-by: Rostan Tabet <[email protected]>

greptile-apps

_{9 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

rostan-t · 2025-11-27T00:27:00Z

!build

dali-automaton · 2025-11-27T00:30:28Z

CI MESSAGE: [39226828]: BUILD STARTED

dali-automaton · 2025-11-27T01:45:35Z

CI MESSAGE: [39226828]: BUILD PASSED

github-advanced-security bot found potential problems Nov 18, 2025

View reviewed changes

rostan-t force-pushed the ndd_pyi branch from 5dfbde4 to e0e2173 Compare November 18, 2025 17:53

github-advanced-security bot found potential problems Nov 18, 2025

View reviewed changes

dali/python/nvidia/dali/_typing.py Dismissed Show dismissed Hide dismissed

github-advanced-security bot found potential problems Nov 20, 2025

View reviewed changes

dali/python/nvidia/dali/_typing.py Dismissed Show dismissed Hide dismissed

github-advanced-security bot found potential problems Nov 21, 2025

View reviewed changes

dali/python/nvidia/dali/ops/_signatures.py Fixed Show fixed Hide fixed

rostan-t marked this pull request as ready for review November 21, 2025 16:59

greptile-apps bot reviewed Nov 21, 2025

View reviewed changes

dali/python/nvidia/dali/ops/_signatures.py Outdated Show resolved Hide resolved

dali/python/nvidia/dali/ops/_signatures.py Outdated Show resolved Hide resolved

github-advanced-security bot found potential problems Nov 24, 2025

View reviewed changes

dali/python/nvidia/dali/ops/_signatures.py Dismissed Show dismissed Hide dismissed

greptile-apps bot reviewed Nov 24, 2025

View reviewed changes

dali/python/nvidia/dali/ops/_signatures.py Outdated Show resolved Hide resolved

greptile-apps bot reviewed Nov 24, 2025

View reviewed changes

dali/python/nvidia/dali/ops/_signatures.py Outdated Show resolved Hide resolved

dali/python/nvidia/dali/ops/_signatures.py Outdated Show resolved Hide resolved

greptile-apps bot reviewed Nov 24, 2025

View reviewed changes

dali/python/nvidia/dali/ops/_signatures.py Show resolved Hide resolved

dali-automaton assigned szalpal and jantonguirao Nov 24, 2025

jantonguirao approved these changes Nov 24, 2025

View reviewed changes

greptile-apps bot reviewed Nov 24, 2025

View reviewed changes

dali/python/nvidia/dali/ops/_signatures.py Outdated Show resolved Hide resolved

greptile-apps bot reviewed Nov 24, 2025

View reviewed changes

github-advanced-security bot found potential problems Nov 25, 2025

View reviewed changes

dali/python/nvidia/dali/ops/_signatures.py Fixed Show fixed Hide fixed

rostan-t added 23 commits November 26, 2025 18:07

Use the dynamic mode doc, omit excluded args and re-export pure Pytho…

b8dfba1

…n definitions Signed-off-by: Rostan Tabet <[email protected]>

Improve the annotation of the device argument

261fe77

Signed-off-by: Rostan Tabet <[email protected]>

Support tensor-like in signatures

a50ccd3

Signed-off-by: Rostan Tabet <[email protected]>

Start generalization to support class ops

fbdb884

Signed-off-by: Rostan Tabet <[email protected]>

Add stubs for class operators

30477cc

Signed-off-by: Rostan Tabet <[email protected]>

Only add next_epoch to readers, not _samples and _batches

420bfd9

Signed-off-by: Rostan Tabet <[email protected]>

Resolve ambiguity when symbols conflict

b01e022

Signed-off-by: Rostan Tabet <[email protected]>

Skip deprecated operators and fix the builtins conflict resolution

82f708e

Signed-off-by: Rostan Tabet <[email protected]>

Fix typo in _api_to_module

f7d61db

Signed-off-by: Rostan Tabet <[email protected]>

Use builtins instead of __builtins__

7d59b5b

Signed-off-by: Rostan Tabet <[email protected]>

Fix next_epoch return type

1989335

Signed-off-by: Rostan Tabet <[email protected]>

Fix Literal syntax for Python 3.10

e8779b6

Signed-off-by: Rostan Tabet <[email protected]>

Use newline instead of \n

ad06315

Signed-off-by: Rostan Tabet <[email protected]>

pip-install NVTX and makefun earlier.

e39dc24

The stubs generation script imports dali.exerimental.dynamic, which requires NVTX and makefun to be importable. Signed-off-by: Rostan Tabet <[email protected]>

Make sure that nvidia_dali_path is a Path

3f66fd4

Signed-off-by: Rostan Tabet <[email protected]>

Do not skip deprecated operators

04fac09

Signed-off-by: Rostan Tabet <[email protected]>

Fix ndd.ops stubs

58976e8

Signed-off-by: Rostan Tabet <[email protected]>

Allow kwargs to accept batches

7f7046f

Signed-off-by: Rostan Tabet <[email protected]>

Add ndd.DType to signatures

7ab2db6

Signed-off-by: Rostan Tabet <[email protected]>

Reexport types in the dynamic mode header

153b8ed

Signed-off-by: Rostan Tabet <[email protected]>

Fix Reader.next_epoch

f1f91eb

Signed-off-by: Rostan Tabet <[email protected]>

Fix long lines

41952cb

Signed-off-by: Rostan Tabet <[email protected]>

Fix handling of include_kwarg_inputs

1e929e7

Signed-off-by: Rostan Tabet <[email protected]>

rostan-t force-pushed the ndd_pyi branch from 287f4be to 1e929e7 Compare November 27, 2025 00:22

greptile-apps bot reviewed Nov 27, 2025

View reviewed changes

rostan-t merged commit 8cb9b6d into NVIDIA:main Nov 27, 2025
7 checks passed

rostan-t deleted the ndd_pyi branch November 27, 2025 09:45

Create stub files for the dynamic mode #6089

Create stub files for the dynamic mode #6089

Uh oh!

Conversation

rostan-t commented Nov 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Category: New feature (non-breaking change which adds functionality)

Description:

Additional information:

Affected modules and functionalities:

Key points relevant for the review:

Tests:

Checklist

Documentation

DALI team only

Requirements

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

greptile-apps bot commented Nov 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Overview

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

rostan-t commented Nov 24, 2025

Uh oh!

dali-automaton commented Nov 24, 2025

Uh oh!

dali-automaton commented Nov 24, 2025

Uh oh!

dali-automaton commented Nov 24, 2025

Uh oh!

dali-automaton commented Nov 24, 2025

Uh oh!

dali-automaton commented Nov 24, 2025

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dali-automaton commented Nov 24, 2025

Uh oh!

dali-automaton commented Nov 24, 2025

Uh oh!

dali-automaton commented Nov 24, 2025

Uh oh!

dali-automaton commented Nov 24, 2025

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

dali-automaton commented Nov 24, 2025

rostan-t commented Nov 18, 2025 •

edited

Loading

greptile-apps bot commented Nov 21, 2025 •

edited

Loading