Support checkpointing in FileReader #4954

staniewzki · 2023-07-25T13:43:37Z

Category:

New feature (non-breaking change which adds functionality)

Description:

This PR adds support for checkpointing to FileReader.

In the prefetching thread, a state snapshot is produced every epoch by the Loader, and stored in a queue.

Currently, only pad_last_batch=true is supported, because otherwise, the epoch might change mid-batch, which complicates checkpointing.

Additional information:

Affected modules and functionalities:

Key points relevant for the review:

Note: most of the changes in file_label_loader.h diff come from fixing the indentation there.

Tests:

Checklist

Documentation

DALI team only

Requirements

Implements new requirements
Affects existing requirements
N/A

REQ IDs: N/A

JIRA TASK: DALI-3527

dali/operators/reader/loader/loader.h

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

dali/pipeline/operator/operator.h

dali/operators/reader/loader/file_label_loader.h

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

staniewzki · 2023-07-26T14:35:18Z

dali/operators/reader/loader/file_label_loader.h

+    vector<string> files;
+    vector<int> labels;
+
+    has_files_arg_ = spec.TryGetRepeatedArgument(files, "files");
+    has_labels_arg_ = spec.TryGetRepeatedArgument(labels, "labels");
+    has_file_list_arg_ = spec.TryGetArgument(file_list_, "file_list");
+    has_file_root_arg_ = spec.TryGetArgument(file_root_, "file_root");
+    bool has_file_filters_arg = spec.TryGetRepeatedArgument(filters_, "file_filters");
+
+    // TODO(ksztenderski): CocoLoader inherits after FileLabelLoader and it doesn't work with
+    // GetArgument.
+    spec.TryGetArgument(case_sensitive_filter_, "case_sensitive_filter");
+
+    DALI_ENFORCE(has_file_root_arg_ || has_files_arg_ || has_file_list_arg_,
+      "``file_root`` argument is required when not using ``files`` or ``file_list``.");
+
+    DALI_ENFORCE(has_files_arg_ + has_file_list_arg_ <= 1,
+      "File paths can be provided through ``files`` or ``file_list`` but not both.");
+
+    DALI_ENFORCE(has_files_arg_ || !has_labels_arg_,
+      "The argument ``labels`` is valid only when file paths "
+      "are provided as ``files`` argument.");
+
+    DALI_ENFORCE(!has_file_filters_arg || filters_.size() > 0,
+                  "``file_filters`` list cannot be empty.");
+
+    if (has_file_list_arg_) {
+      DALI_ENFORCE(!file_list_.empty(), "``file_list`` argument cannot be empty");
+      if (!has_file_root_arg_) {
+        auto idx = file_list_.rfind(filesystem::dir_sep);
+        if (idx != string::npos) {
+          file_root_ = file_list_.substr(0, idx);
        }
      }
+    }

-      if (has_files_arg_) {
-        DALI_ENFORCE(files.size() > 0, "``files`` specified an empty list.");
-        if (has_labels_arg_) {
-          DALI_ENFORCE(files.size() == labels.size(), make_string("Provided ", labels.size(),
-            " labels for ", files.size(), " files."));
+    if (has_files_arg_) {
+      DALI_ENFORCE(files.size() > 0, "``files`` specified an empty list.");
+      if (has_labels_arg_) {
+        DALI_ENFORCE(files.size() == labels.size(), make_string("Provided ", labels.size(),
+          " labels for ", files.size(), " files."));

+        for (int i = 0, n = files.size(); i < n; i++)
+          image_label_pairs_.emplace_back(std::move(files[i]), labels[i]);
+      } else {
          for (int i = 0, n = files.size(); i < n; i++)
-            image_label_pairs_.emplace_back(std::move(files[i]), labels[i]);
-        } else {
-            for (int i = 0, n = files.size(); i < n; i++)
-              image_label_pairs_.emplace_back(std::move(files[i]), i);
-        }
+            image_label_pairs_.emplace_back(std::move(files[i]), i);
      }
+    }

-      /*
-      * Those options are mutually exclusive as `shuffle_after_epoch` will make every shard looks differently
-      * after each epoch so coexistence with `stick_to_shard` doesn't make any sense
-      * Still when `shuffle_after_epoch` we will set `stick_to_shard` internally in the FileLabelLoader so all
-      * DALI instances will do shuffling after each epoch
-      */
-      DALI_ENFORCE(!(shuffle_after_epoch_  && stick_to_shard_),
-                   "shuffle_after_epoch and stick_to_shard cannot be both true");
-      DALI_ENFORCE(!(shuffle_after_epoch_ && shuffle_),
-                   "shuffle_after_epoch and random_shuffle cannot be both true");
-      /*
-       * Imply `stick_to_shard` from  `shuffle_after_epoch`
-       */
-      if (shuffle_after_epoch_) {
-        stick_to_shard_ = true;
-      }
+    /*
+     * Those options are mutually exclusive as `shuffle_after_epoch` will make every shard looks differently
+     * after each epoch so coexistence with `stick_to_shard` doesn't make any sense
+     * Still when `shuffle_after_epoch` we will set `stick_to_shard` internally in the FileLabelLoader so all
+     * DALI instances will do shuffling after each epoch
+     */
+    DALI_ENFORCE(!(shuffle_after_epoch_  && stick_to_shard_),
+                  "shuffle_after_epoch and stick_to_shard cannot be both true");
+    DALI_ENFORCE(!(shuffle_after_epoch_ && shuffle_),
+                  "shuffle_after_epoch and random_shuffle cannot be both true");
+    /*
+     * Imply `stick_to_shard` from  `shuffle_after_epoch`
+     */
+    if (shuffle_after_epoch_) {
+      stick_to_shard_ = true;
+    }


Lengthy diff, but this is just an indentation fix.

dali/operators/reader/loader/file_label_loader.cc

stiepan · 2023-07-26T14:30:20Z

dali/operators/reader/loader/file_label_loader.h

+  std::vector<FileLabelLoaderState> state_queue_;
+  Index state_queue_front_;
+  Index state_queue_back_;
+  int checkpoint_epoch_;
+  std::mutex state_queue_mutex_;
+  using Loader<CPUBackend, ImageLabelWrapper>::checkpointing_;


For the record of what we talked offline:

Can we move it all to the Loader base? And make the FileLabelLoaderState just a LoaderCheckpoint? There does not seem to be anything file-reader specific in here.

We don't exactly know how other loaders will play out with the featrue, so I wouldn't push for that at that moment if it wasn't for the fact that feature flag and pushing new loader checkpoint already leaked to the base class. Now, the base loader has the dummy "push method" and no corresponding pop.

As a safegaurd for the readers we do not test for now you could add a flag to loader

template <typename Backend, typename LoadTarget, bool support_checkpointing=false> class Loader {...}

and somwhere inside to be used instead of plain checkpointing_

bool EnableCheckpointing() { // no need for virtual methods return support_checkpointing && checkpointing_; }

In the future, the enable_checkpointing could become the actual type of loder's checkpoint/state to be put in the queue.

dali/operators/reader/reader_op_test.cc

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

dali/operators/reader/reader_op.h

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

dali/operators/reader/reader_op_test.cc

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

staniewzki · 2023-07-31T08:59:37Z

!build

dali-automaton · 2023-07-31T09:05:05Z

CI MESSAGE: [9175984]: BUILD STARTED

dali-automaton · 2023-07-31T10:45:47Z

CI MESSAGE: [9175984]: BUILD PASSED

This PR adds support for checkpointing to FileReader. In the prefetching thread, a state snapshot is produced every epoch by the Loader, and stored in a queue. Currently, only pad_last_batch=true is supported, because otherwise, the epoch might change mid-batch, which complicates checkpointing. --------- Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

stiepan reviewed Jul 25, 2023

View reviewed changes

dali/operators/reader/loader/loader.h Show resolved Hide resolved

staniewzki force-pushed the file-reader-checkpointing branch from 99c56ca to 5395b50 Compare July 26, 2023 09:07

staniewzki added 5 commits July 26, 2023 11:08

Baseline functionality

036d74c

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

Modify tests

2b3ee24

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

Update comments

da013d2

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

Allow 'shuffle_after_epoch', add more tests

aed0be8

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

Adjust tests

ed5db55

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

staniewzki force-pushed the file-reader-checkpointing branch from 5395b50 to ed5db55 Compare July 26, 2023 09:09

stiepan reviewed Jul 26, 2023

View reviewed changes

dali/pipeline/operator/operator.h Show resolved Hide resolved

stiepan reviewed Jul 26, 2023

View reviewed changes

dali/operators/reader/loader/file_label_loader.h Outdated Show resolved Hide resolved

staniewzki added 2 commits July 26, 2023 14:36

Add more tests

fddf547

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

Adjust typo

b43d723

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

staniewzki marked this pull request as ready for review July 26, 2023 13:00

staniewzki commented Jul 26, 2023

View reviewed changes

stiepan reviewed Jul 26, 2023

View reviewed changes

stiepan self-assigned this Jul 26, 2023

stiepan reviewed Jul 26, 2023

View reviewed changes

dali/operators/reader/reader_op_test.cc Show resolved Hide resolved

Adjust todo message

d8eadd6

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

staniewzki force-pushed the file-reader-checkpointing branch from 5dd95b1 to d8eadd6 Compare July 27, 2023 08:00

staniewzki added 3 commits July 27, 2023 10:03

Remove added line

ad3e171

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

Fix linter issue

37aaf7c

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

Adjust tests

3f12834

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

jantonguirao assigned banasraf Jul 27, 2023

Mark the checkpointing support in templates

5f1a70e

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

stiepan reviewed Jul 27, 2023

View reviewed changes

dali/operators/reader/reader_op.h Outdated Show resolved Hide resolved

Fix documentation

313c5e0

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

stiepan reviewed Jul 27, 2023

View reviewed changes

dali/operators/reader/reader_op_test.cc Show resolved Hide resolved

stiepan reviewed Jul 27, 2023

View reviewed changes

dali/operators/reader/reader_op_test.cc Outdated Show resolved Hide resolved

staniewzki added 2 commits July 28, 2023 08:50

Fix sharding

8d087e1

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

Remove redundant info from snapshot

55f5631

Signed-off-by: Michał Staniewski <mstaniewski@nvidia.com>

stiepan approved these changes Jul 28, 2023

View reviewed changes

banasraf approved these changes Jul 31, 2023

View reviewed changes

staniewzki merged commit d52fd08 into NVIDIA:main Jul 31, 2023
5 checks passed

szkarpinski mentioned this pull request Oct 10, 2023

Add Fast-Forward to Loader base #5088

Merged

18 tasks

szkarpinski mentioned this pull request Nov 3, 2023

Add checkpointing support to Video Reader #5139

Merged

18 tasks

szkarpinski mentioned this pull request Nov 30, 2023

Support checkpointing in Numpy reader #5198

Merged

18 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support checkpointing in FileReader #4954

Support checkpointing in FileReader #4954

staniewzki commented Jul 25, 2023 •

edited

Loading

staniewzki Jul 26, 2023

stiepan Jul 26, 2023 •

edited

Loading

staniewzki commented Jul 31, 2023

dali-automaton commented Jul 31, 2023

dali-automaton commented Jul 31, 2023

Support checkpointing in FileReader #4954

Support checkpointing in FileReader #4954

Conversation

staniewzki commented Jul 25, 2023 • edited Loading

Category:

Description:

Additional information:

Affected modules and functionalities:

Key points relevant for the review:

Tests:

Checklist

Documentation

DALI team only

Requirements

staniewzki Jul 26, 2023

Choose a reason for hiding this comment

stiepan Jul 26, 2023 • edited Loading

Choose a reason for hiding this comment

staniewzki commented Jul 31, 2023

dali-automaton commented Jul 31, 2023

dali-automaton commented Jul 31, 2023

staniewzki commented Jul 25, 2023 •

edited

Loading

stiepan Jul 26, 2023 •

edited

Loading