Add PyTorch DataLoader Evaluator plugin #6112

JanuszL · 2025-12-05T10:39:16Z

Introduces a lightweight diagnostic tool for identifying data loading
bottlenecks in PyTorch training pipelines.
This change adds Loader Evaluator inside pytorch DALI plugin, a
jupyter notebook tutorial, and a documentation page with tests
LoaderEvaluator class wraps PyTorch DataLoader with performance monitoring
Two operation modes: 'log' (normal iteration with metrics) and 'replay'
(cached batches for ideal performance simulation)
PerformanceMetrics class for detailed performance tracking and bottleneck
analysis
In-memory batch caching for replay mode to simulate ideal data loading
Comprehensive test suite and documentation with example notebook
The tool helps users compare real vs. ideal data loading performance and
identify optimization opportunities.

Authored-by: Albert Wolant [email protected]

Category:

New feature (non-breaking change which adds functionality)

Description:

Introduces a lightweight diagnostic tool for identifying data loading
bottlenecks in PyTorch training pipelines.
This change adds Loader Evaluator inside pytorch DALI plugin, a
jupyter notebook tutorial, and a documentation page with tests
LoaderEvaluator class wraps PyTorch DataLoader with performance monitoring
Two operation modes: 'log' (normal iteration with metrics) and 'replay'
(cached batches for ideal performance simulation)
PerformanceMetrics class for detailed performance tracking and bottleneck
analysis
In-memory batch caching for replay mode to simulate ideal data loading
Comprehensive test suite and documentation with example notebook
The tool helps users compare real vs. ideal data loading performance and
identify optimization opportunities.

Additional information:

Affected modules and functionalities:

new module in Pytorch plugin
new example
new test for it
new documentation page describing the overall idea

Key points relevant for the review:

overall idea and flow

Tests:

Checklist

Documentation

DALI team only

Requirements

Implements new requirements
Affects existing requirements
N/A

REQ IDs: N/A

JIRA TASK: DALI-4299

review-notebook-app · 2025-12-05T10:39:22Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

dali/test/python/test_pytorch_loader_evaluator.py

dali-automaton · 2025-12-05T10:45:29Z

CI MESSAGE: [39670512]: BUILD STARTED

JanuszL · 2025-12-05T10:47:30Z

!build

dali-automaton · 2025-12-05T10:50:27Z

CI MESSAGE: [39670654]: BUILD STARTED

greptile-apps · 2025-12-05T10:55:55Z

Greptile Overview

Greptile Summary

This PR adds a new LoaderEvaluator diagnostic tool to the PyTorch plugin for identifying data loading bottlenecks in training pipelines.

Key Changes:

LoaderEvaluator class wraps PyTorch DataLoader with two modes: "log" (normal iteration with metrics) and "replay" (cached batches for ideal performance simulation)
Comprehensive test suite covering initialization, modes, metrics collection, and edge cases
Documentation includes RST pages and a Jupyter notebook tutorial demonstrating the bottleneck detection workflow

Implementation Notes:

The tool caches batches during construction in replay mode, then replays them to simulate zero-overhead data loading
Performance metrics include batch times, total time, and cache status
Thread-safe caching using threading.Lock

Minor Issue Found:

Unused variable batch_count in _replay_mode_iter() method

Confidence Score: 4/5

This PR is safe to merge - adds a new self-contained diagnostic tool with comprehensive tests and no changes to existing functionality.
Score of 4 reflects a well-structured new feature with good test coverage and documentation. Only minor issue is an unused variable that doesn't affect functionality.
Minor cleanup needed in dali/python/nvidia/dali/plugin/pytorch/loader_evaluator/loader.py (unused variable).

Important Files Changed

File Analysis

Filename	Score	Overview
dali/python/nvidia/dali/plugin/pytorch/loader_evaluator/loader.py	4/5	Core LoaderEvaluator class with caching and metrics collection. Contains one unused variable `batch_count`.
dali/test/python/test_pytorch_loader_evaluator.py	5/5	Comprehensive test suite covering log/replay modes, edge cases, and metrics collection.
docs/examples/frameworks/pytorch/loader_evaluator/pytorch_data_loader_evaluator.ipynb	5/5	Well-documented Jupyter notebook demonstrating bottleneck detection workflow.
docs/plugins/pytorch_data_loader_evaluator.rst	5/5	Detailed documentation explaining the tool's purpose, usage, and comparison with other profiling tools.

Sequence Diagram

sequenceDiagram
    participant User as Training Code
    participant LE as LoaderEvaluator
    participant DL as PyTorch DataLoader
    participant Cache as Batch Cache

    alt Log Mode
        User->>LE: __iter__()
        LE->>DL: iter(dataloader)
        loop For each batch
            LE->>DL: next()
            DL-->>LE: batch
            Note over LE: Record batch_time
            LE-->>User: yield batch
        end
        Note over LE: Set end_time
    else Replay Mode (Construction)
        User->>LE: __init__(mode="replay")
        LE->>DL: iter(dataloader)
        loop Cache batches
            DL-->>LE: batch
            LE->>Cache: append(batch)
        end
        Note over LE: cache_ready = True
    end

    alt Replay Mode (Iteration)
        User->>LE: __iter__()
        loop For original_length iterations
            LE->>Cache: get batch[i % cache_size]
            Cache-->>LE: cached batch
            Note over LE: Record batch_time
            LE-->>User: yield batch
        end
    end

greptile-apps

Additional Comments (3)

dali/python/nvidia/dali/plugin/pytorch/loader_evaluator/metrics.py, line 137-145 (link)

logic: json.dump() will fail with TypeError because get_summary() returns numpy types (np.float64 from np.mean(), np.std(), etc.) which are not JSON serializable.

Consider converting numpy types to native Python types before serialization, e.g.:

def save_metrics(self, filename: str):
    """Save metrics to a JSON file."""
    metrics = self.get_summary()
    bottlenecks = self.identify_bottlenecks()

    # Convert numpy types to native Python types for JSON serialization
    def convert_numpy(obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, (np.floating, np.integer)):
            return obj.item()
        elif isinstance(obj, dict):
            return {k: convert_numpy(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [convert_numpy(i) for i in obj]
        return obj

    output = convert_numpy({"metrics": metrics, "bottlenecks": bottlenecks})

    with open(filename, "w") as f:
        json.dump(output, f, indent=2)

    print(f"Performance metrics saved to: {filename}")

dali/python/nvidia/dali/plugin/pytorch/loader_evaluator/metrics.py, line 100-102 (link)

style: Potential division by zero if all batch_times are 0 (edge case).
dali/python/nvidia/dali/plugin/pytorch/loader_evaluator/__init__.py, line 1 (link)

style: Missing Apache 2.0 license header that other files in this repository have. Same applies to loader.py and metrics.py.

_{9 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

- Introduces a lightweight diagnostic tool for identifying data loading bottlenecks in PyTorch training pipelines. - This change adds Loader Evaluator inside pytorch DALI plugin, a jupyter notebook tutorial, and a documentation page with tests - LoaderEvaluator class wraps PyTorch DataLoader with performance monitoring Two operation modes: 'log' (normal iteration with metrics) and 'replay' (cached batches for ideal performance simulation) - PerformanceMetrics class for detailed performance tracking and bottleneck analysis - In-memory batch caching for replay mode to simulate ideal data loading - Comprehensive test suite and documentation with example notebook - The tool helps users compare real vs. ideal data loading performance and identify optimization opportunities. Authored-by: Albert Wolant <[email protected]> Signed-off-by: Janusz Lisiecki <[email protected]>

JanuszL · 2025-12-05T12:11:10Z

@greptileai please review again.

greptile-apps

Additional Comments (1)

dali/python/nvidia/dali/plugin/pytorch/loader_evaluator/__init__.py, line 16 (link)

syntax: Missing metrics.py module. This import will raise ModuleNotFoundError: No module named 'nvidia.dali.plugin.pytorch.loader_evaluator.metrics' at runtime.

Either create the missing metrics.py file with the PerformanceMetrics class, or remove this import and the corresponding __all__ entry.

_{8 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

Signed-off-by: Janusz Lisiecki <[email protected]>

JanuszL · 2025-12-05T12:46:05Z

!build

dali-automaton · 2025-12-05T12:50:28Z

CI MESSAGE: [39675479]: BUILD STARTED

JanuszL · 2025-12-05T12:53:58Z

@greptileai please review again.

greptile-apps

Additional Comments (2)

dali/python/nvidia/dali/plugin/pytorch/loader_evaluator/loader.py, line 85 (link)

style: Type hint assumes specific batch structure but DataLoader can return any type. Should this support generic batch types instead of assuming Tuple[TorchTensor, TorchTensor]?

_{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!}
dali/python/nvidia/dali/plugin/pytorch/loader_evaluator/loader.py, line 146-163 (link)

style: Double break logic could be simplified with a single loop condition

_{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!}

_{8 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

JanuszL · 2025-12-05T13:11:37Z

@greptileai please review again.

Signed-off-by: Janusz Lisiecki <[email protected]>

greptile-apps

Additional Comments (1)

dali/python/nvidia/dali/plugin/pytorch/loader_evaluator/loader.py, line 143 (link)

style: Unused variable batch_count - initialized but never used.

_{8 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

JanuszL · 2025-12-05T13:19:39Z

!build

dali-automaton · 2025-12-05T13:25:21Z

CI MESSAGE: [39676718]: BUILD STARTED

qa/TL0_python-self-test-core/test_body.sh

docs/plugins/pytorch_data_loader_evaluator.rst

mdabek-nvidia · 2025-12-05T14:51:10Z

docs/examples/frameworks/pytorch/loader_evaluator/pytorch_data_loader_evaluator.ipynb

@@ -0,0 +1,451 @@
+{


Line #10. import sys
I think that sys is not used anywhere.

Reply via ReviewNB

mdabek-nvidia · 2025-12-05T14:51:10Z

docs/examples/frameworks/pytorch/loader_evaluator/pytorch_data_loader_evaluator.ipynb

@@ -0,0 +1,451 @@
+{


Line #1. # Wrap your dataloader with LoaderEvaluator (this is the only change!)
Nitpick: there is also a change below dataloader became monitored_dataloader

Reply via ReviewNB

Do you think it should be:

dataloader = LoaderEvaluator( dataloader, mode="replay", num_cached_batches=len(dataloader) // 10 )

instead?

Yes, this would be clear, that we are just "updating" the dataloader.

Signed-off-by: Janusz Lisiecki <[email protected]>

dali-automaton · 2025-12-05T17:33:13Z

CI MESSAGE: [39670654]: BUILD FAILED

dali-automaton · 2025-12-05T17:50:19Z

CI MESSAGE: [39676718]: BUILD FAILED

JanuszL · 2025-12-05T18:53:11Z

!build

dali-automaton · 2025-12-05T18:55:46Z

CI MESSAGE: [39693133]: BUILD STARTED

dali-automaton · 2025-12-05T21:14:49Z

CI MESSAGE: [39693133]: BUILD FAILED

JanuszL marked this pull request as draft December 5, 2025 10:39

github-advanced-security bot found potential problems Dec 5, 2025

View reviewed changes

dali/test/python/test_pytorch_loader_evaluator.py Fixed Show fixed Hide fixed

JanuszL force-pushed the loader_Evaluator branch 2 times, most recently from 6766528 to a008b7c Compare December 5, 2025 10:45

greptile-apps bot reviewed Dec 5, 2025

View reviewed changes

mdabek-nvidia self-assigned this Dec 5, 2025

JanuszL force-pushed the loader_Evaluator branch from a008b7c to 7d47417 Compare December 5, 2025 11:12

greptile-apps bot reviewed Dec 5, 2025

View reviewed changes

JanuszL force-pushed the loader_Evaluator branch 3 times, most recently from c414911 to b7d6958 Compare December 5, 2025 12:43

Fix

41896c6

Signed-off-by: Janusz Lisiecki <[email protected]>

JanuszL force-pushed the loader_Evaluator branch from b7d6958 to 41896c6 Compare December 5, 2025 12:45

greptile-apps bot reviewed Dec 5, 2025

View reviewed changes

More fixes

f61d369

Signed-off-by: Janusz Lisiecki <[email protected]>

JanuszL force-pushed the loader_Evaluator branch from ff53b06 to f61d369 Compare December 5, 2025 13:13

greptile-apps bot reviewed Dec 5, 2025

View reviewed changes

mdabek-nvidia reviewed Dec 5, 2025

View reviewed changes

qa/TL0_python-self-test-core/test_body.sh Outdated Show resolved Hide resolved

docs/plugins/pytorch_data_loader_evaluator.rst Show resolved Hide resolved

mdabek-nvidia reviewed Dec 5, 2025

View reviewed changes

JanuszL added 2 commits December 5, 2025 16:02

more fixes

d306b1b

Signed-off-by: Janusz Lisiecki <[email protected]>

Fix

1c41dac

Signed-off-by: Janusz Lisiecki <[email protected]>

Add PyTorch DataLoader Evaluator plugin #6112

Are you sure you want to change the base?

Add PyTorch DataLoader Evaluator plugin #6112

Uh oh!

Conversation

JanuszL commented Dec 5, 2025

Category:

Description:

Additional information:

Affected modules and functionalities:

Key points relevant for the review:

Tests:

Checklist

Documentation

DALI team only

Requirements

Uh oh!

review-notebook-app bot commented Dec 5, 2025

Uh oh!

Uh oh!

dali-automaton commented Dec 5, 2025

Uh oh!

JanuszL commented Dec 5, 2025

Uh oh!

dali-automaton commented Dec 5, 2025

Uh oh!

greptile-apps bot commented Dec 5, 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 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Additional Comments (3)

Uh oh!

JanuszL commented Dec 5, 2025

Uh oh!

greptile-apps bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Additional Comments (1)

Uh oh!

JanuszL commented Dec 5, 2025

Uh oh!

dali-automaton commented Dec 5, 2025

Uh oh!

JanuszL commented Dec 5, 2025

Uh oh!

greptile-apps bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Additional Comments (2)

Uh oh!

JanuszL commented Dec 5, 2025

Uh oh!

greptile-apps bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Additional Comments (1)

Uh oh!

JanuszL commented Dec 5, 2025

Uh oh!

dali-automaton commented Dec 5, 2025

Uh oh!

Uh oh!

Uh oh!

mdabek-nvidia Dec 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JanuszL Dec 5, 2025

Choose a reason for hiding this comment

Uh oh!

mdabek-nvidia Dec 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

greptile-apps bot commented Dec 5, 2025 •

edited

Loading

greptile-apps bot left a comment •

edited

Loading

greptile-apps bot left a comment •

edited

Loading

greptile-apps bot left a comment •

edited

Loading

greptile-apps bot left a comment •

edited

Loading

mdabek-nvidia Dec 5, 2025 •

edited

Loading

mdabek-nvidia Dec 5, 2025 •

edited

Loading