Add S3 storage integration tests, and S3 support for webdataset index files

dali/operators/reader/loader/recordio_loader.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2017-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -49,12 +49,20 @@ class RecordIOLoader : public IndexedFileLoader {
     DALI_ENFORCE(index_paths.size() == 1,
         "RecordIOReader supports only a single index file");
     const std::string& path = index_paths[0];
-    std::ifstream index_file(path);
-    DALI_ENFORCE(index_file.good(),
+    // FileStream::Open dispatches on the URI scheme (e.g. s3://), so the index file can live
+    // next to the shards it describes, remote or local, same as indexed_file_loader.h.
+    auto index_file = FileStream::Open(path);
+    auto index_file_cleanup = AtScopeExit([&index_file] {
+      if (index_file)
+        index_file->Close();
+    });
+    FileStreamBuf<> stream_buf(index_file.get());
+    std::istream fin(&stream_buf);
+    DALI_ENFORCE(fin.good(),
         "Could not open RecordIO index file. Provided path: \"" + path + "\"");
     std::vector<size_t> temp;
     size_t index, offset, prev_offset = -1;
-    while (index_file >> index >> offset) {
+    while (fin >> index >> offset) {
       temp.push_back(offset);
     }
     DALI_ENFORCE(!temp.empty(),
@@ -80,7 +88,6 @@ class RecordIOLoader : public IndexedFileLoader {
       indices_.emplace_back(temp.back() - file_offsets[file_offset_index],
                             size, file_offset_index);
     }
-    index_file.close();
   }
 
   void ReadSample(IndexedFileLoaderSample& sample) override {

dali/operators/reader/loader/webdataset_loader.cc

@@ -1,4 +1,4 @@
-// Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -20,6 +20,7 @@
 #include <memory>
 #include <numeric>
 #include <random>
+#include <sstream>
 #include <tuple>
 #include <utility>
 #include "dali/core/common.h"
@@ -60,7 +61,7 @@ inline MissingExtBehavior ParseMissingExtBehavior(std::string missing_component_
 
 inline void ParseSampleDesc(std::vector<SampleDesc>& samples_container,
                             std::vector<ComponentDesc>& components_container,
-                            std::ifstream& index_file, const std::string& index_path, int64_t line,
+                            std::istream& index_file, const std::string& index_path, int64_t line,
                             int index_version) {
   // Preparing the SampleDesc
   samples_container.emplace_back();
@@ -118,7 +119,13 @@ inline int ParseIndexVersion(const string& version_str) {
 inline void ParseIndexFile(std::vector<SampleDesc>& samples_container,
                            std::vector<ComponentDesc>& components_container,
                            const std::string& index_path) {
-  std::ifstream index_file(index_path);
+  // FileStream::Open dispatches on the URI scheme (e.g. s3://), so index files can live next to
+  // the shards they describe, remote or local. index_paths is a public, caller-controlled
+  // argument and the index can be arbitrarily large, so it's streamed through FileStreamBuf
+  // (also used by IndexedFileLoader) rather than read into memory in full.
+  auto file = FileStream::Open(index_path);
+  FileStreamBuf<64 * 1024> file_buf(file.get());
+  std::istream index_file(&file_buf);
 
   // Index Checking
   std::string global_meta;

dali/operators/reader/loader/webdataset_loader.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -15,7 +15,6 @@
 #ifndef DALI_OPERATORS_READER_LOADER_WEBDATASET_LOADER_H_
 #define DALI_OPERATORS_READER_LOADER_WEBDATASET_LOADER_H_
 
-#include <fstream>
 #include <memory>
 #include <set>
 #include <string>

dali/test/python/reader/test_s3.py

@@ -0,0 +1,294 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import io
+import os
+import struct
+import tarfile
+import tempfile
+from concurrent.futures import ThreadPoolExecutor
+
+import nvidia.dali.fn as fn
+from nvidia.dali import pipeline_def
+
+import s3_test_utils as s3
+from nose_utils import assert_raises, attr
+from test_utils import compare_pipelines
+
+batch_size = 4
+num_threads = 4
+
+g_server = None
+g_tmpdir = None
+g_root = None
+g_files = None
+g_tar = None
+g_index = None
+g_rec = None
+g_rec_index = None
+g_client = None
+
+DATA_PREFIX = f"{s3.PREFIX}/data"
+WDS_PREFIX = f"{s3.PREFIX}/wds"
+MANY_PREFIX = f"{s3.PREFIX}/many"
+REC_PREFIX = f"{s3.PREFIX}/rec"
+
+
+def _make_local_dataset(root):
+    files = []
+    for cls in ("class_a", "class_b"):
+        os.makedirs(os.path.join(root, cls))
+        for i in range(4):
+            name = os.path.join(cls, f"{i} x.dat")  # space on purpose: NVIDIA/DALI#5521
+            with open(os.path.join(root, name), "w") as f:
+                f.write(f"Contents of {name}.\n")
+            files.append(name)
+    empty = os.path.join("class_a", "empty.dat")  # 0-byte -> HeadObject branch
+    open(os.path.join(root, empty), "w").close()
+    files.append(empty)
+    return sorted(files)
+
+
+def _make_local_tar(path, num_samples=8):
+    with tarfile.open(path, "w") as tar:
+        for i in range(num_samples):
+            for ext, payload in (("txt", b"sample-%d" % i), ("cls", b"%d" % i)):
+                info = tarfile.TarInfo(f"{i:04}.{ext}")
+                info.size = len(payload)
+                tar.addfile(info, io.BytesIO(payload))
+
+
+def _make_local_recordio(path, index_path, num_samples=8):
+    """A minimal MXNet RecordIO shard: magic (0xced7230a) + length/cflag word (cflag=0, no
+    continuation) + a 24-byte ImageRecordIOHeader (flag=0 -> a single float label, no extra
+    labels) + the sample payload. The index is plain text, "<ignored> <byte offset>" per line -
+    RecordIOLoader::ReadIndexFile only reads the second column (dali/operators/reader/loader/
+    recordio_loader.h)."""
+    with open(path, "wb") as rec, open(index_path, "w") as idx:
+        offset = 0
+        for i in range(num_samples):
+            payload = struct.pack("<IfQQ", 0, float(i), 0, 0) + b"sample-%d" % i
+            record = struct.pack("<II", 0xCED7230A, len(payload)) + payload
+            rec.write(record)
+            idx.write(f"{i} {offset}\n")
+            offset += len(record)
+
+
+def _seed_many_objects(client, count=1100):
+    def put(i):
+        client.put_object(Bucket=s3.BUCKET, Key=f"{MANY_PREFIX}/c/{i:04}.dat", Body=b"x")
+
+    with ThreadPoolExecutor(16) as ex:
+        list(ex.map(put, range(count)))
+
+
+def setUpModule():
+    global g_server, g_tmpdir, g_root, g_files, g_tar, g_index, g_rec, g_rec_index, g_client
+    s3.require_mock_server()
+
+    g_tmpdir = tempfile.TemporaryDirectory()
+    g_root = os.path.join(g_tmpdir.name, "data")
+    os.makedirs(g_root)
+    g_files = _make_local_dataset(g_root)
+    g_tar = os.path.join(g_tmpdir.name, "shard0.tar")
+    _make_local_tar(g_tar)
+    g_rec = os.path.join(g_tmpdir.name, "train.rec")
+    g_rec_index = os.path.join(g_tmpdir.name, "train.idx")
+    _make_local_recordio(g_rec, g_rec_index)
+
+    g_server = s3.make_server()
+    try:
+        endpoint = g_server.start()
+        s3.export_s3_env(endpoint)
+        s3.require_s3_support()
+        g_client = s3.s3_client(endpoint)
+        s3.create_bucket(g_client, s3.BUCKET)
+        s3.upload_dir(g_client, s3.BUCKET, g_root, DATA_PREFIX)
+        g_client.upload_file(g_tar, s3.BUCKET, f"{WDS_PREFIX}/shard0.tar")
+        g_client.upload_file(g_rec, s3.BUCKET, f"{REC_PREFIX}/train.rec")
+        g_client.upload_file(g_rec_index, s3.BUCKET, f"{REC_PREFIX}/train.idx")
+
+        # everything past the first upload has to stay inside this try: unittest skips
+        # tearDownModule once setUpModule raised, so the call below is the only thing that
+        # removes what was uploaded
+        import webdataset_base as base
+
+        g_index = base.generate_temp_index_file(g_tar)
+        g_client.upload_file(g_index.name, s3.BUCKET, f"{WDS_PREFIX}/shard0.index")
+    except Exception:
+        tearDownModule()
+        raise
+
+
+def tearDownModule():
+    global g_server, g_tmpdir, g_index, g_client
+    try:
+        # The mock server, and everything uploaded to it, is thrown away wholesale.
+        g_client = None
+        if g_server is not None:
+            g_server.stop()
+            g_server = None
+    finally:
+        if g_index is not None:
+            g_index.close()
+            g_index = None
+        if g_tmpdir is not None:
+            g_tmpdir.cleanup()
+            g_tmpdir = None
+
+
+@pipeline_def(batch_size=batch_size, num_threads=num_threads, device_id=None)
+def file_pipe(file_root=None, files=None, dont_use_mmap=False):
+    return tuple(
+        fn.readers.file(
+            file_root=file_root,
+            files=files,
+            name="Reader",
+            dont_use_mmap=dont_use_mmap,
+            file_filters=["*.dat"],
+        )
+    )
+
+
+@pipeline_def(batch_size=batch_size, num_threads=num_threads, device_id=None)
+def wds_pipe(paths, index_paths=None, dont_use_mmap=False):
+    return tuple(
+        fn.readers.webdataset(
+            paths=paths, index_paths=index_paths, ext=["txt", "cls"], dont_use_mmap=dont_use_mmap
+        )
+    )
+
+
+@pipeline_def(batch_size=batch_size, num_threads=num_threads, device_id=None)
+def mxnet_pipe(path, index_path, dont_use_mmap=False):
+    return tuple(fn.readers.mxnet(path=path, index_path=index_path, dont_use_mmap=dont_use_mmap))
+
+
+def test_file_reader_file_root():
+    compare_pipelines(
+        file_pipe(file_root=f"s3://{s3.BUCKET}/{DATA_PREFIX}"),
+        file_pipe(file_root=g_root, dont_use_mmap=True),  # mmap cannot map the 0-byte file
+        batch_size,
+        3,
+    )
+
+
+def test_file_reader_files_arg():
+    compare_pipelines(
+        file_pipe(files=[f"s3://{s3.BUCKET}/{DATA_PREFIX}/{f}" for f in g_files]),
+        file_pipe(files=[os.path.join(g_root, f) for f in g_files], dont_use_mmap=True),
+        batch_size,
+        2,
+    )
+
+
+@attr("slow")
+def test_file_reader_listing_pagination():
+    # seeded here rather than in setUpModule: it costs ~7 s and nothing else needs it
+    _seed_many_objects(s3.s3_client(g_server.endpoint_url))
+    pipe = file_pipe(file_root=f"s3://{s3.BUCKET}/{MANY_PREFIX}")
+    pipe.build()
+    assert pipe.reader_meta("Reader")["epoch_size"] == 1100, pipe.reader_meta("Reader")
+
+
+def test_webdataset_index_inferred():
+    compare_pipelines(
+        wds_pipe(paths=f"s3://{s3.BUCKET}/{WDS_PREFIX}/shard0.tar"),
+        wds_pipe(paths=g_tar, dont_use_mmap=True),
+        batch_size,
+        2,
+    )
+
+
+def test_webdataset_local_index():
+    compare_pipelines(
+        wds_pipe(paths=f"s3://{s3.BUCKET}/{WDS_PREFIX}/shard0.tar", index_paths=[g_index.name]),
+        wds_pipe(paths=g_tar, index_paths=[g_index.name], dont_use_mmap=True),
+        batch_size,
+        2,
+    )
+
+
+def test_webdataset_remote_index():
+    compare_pipelines(
+        wds_pipe(
+            paths=f"s3://{s3.BUCKET}/{WDS_PREFIX}/shard0.tar",
+            index_paths=[f"s3://{s3.BUCKET}/{WDS_PREFIX}/shard0.index"],
+        ),
+        wds_pipe(paths=g_tar, index_paths=[g_index.name], dont_use_mmap=True),
+        batch_size,
+        2,
+    )
+
+
+@attr("slow")
+def test_webdataset_remote_index_multi_chunk():
+    # 3000 samples puts the index well past FileStreamBuf's 64 KiB read chunk
+    # (webdataset_loader.cc), so building this pipeline exercises multiple Read() round trips
+    # against S3 - including the one at the exact end of the index - not just the single-chunk
+    # case every other webdataset test here happens to hit.
+    import webdataset_base as base
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        big_tar = os.path.join(tmpdir, "big.tar")
+        _make_local_tar(big_tar, num_samples=3000)
+        big_index = base.generate_temp_index_file(big_tar)
+        try:
+            prefix = f"{WDS_PREFIX}/big"
+            g_client.upload_file(big_tar, s3.BUCKET, f"{prefix}/shard0.tar")
+            g_client.upload_file(big_index.name, s3.BUCKET, f"{prefix}/shard0.index")
+            compare_pipelines(
+                wds_pipe(
+                    paths=f"s3://{s3.BUCKET}/{prefix}/shard0.tar",
+                    index_paths=[f"s3://{s3.BUCKET}/{prefix}/shard0.index"],
+                ),
+                wds_pipe(paths=big_tar, index_paths=[big_index.name], dont_use_mmap=True),
+                batch_size,
+                2,
+            )
+        finally:
+            big_index.close()
+
+
+def test_mxnet_reader_local_index():
+    compare_pipelines(
+        mxnet_pipe(path=f"s3://{s3.BUCKET}/{REC_PREFIX}/train.rec", index_path=g_rec_index),
+        mxnet_pipe(path=g_rec, index_path=g_rec_index, dont_use_mmap=True),
+        batch_size,
+        2,
+    )
+
+
+def test_mxnet_reader_remote_index():
+    compare_pipelines(
+        mxnet_pipe(
+            path=f"s3://{s3.BUCKET}/{REC_PREFIX}/train.rec",
+            index_path=f"s3://{s3.BUCKET}/{REC_PREFIX}/train.idx",
+        ),
+        mxnet_pipe(path=g_rec, index_path=g_rec_index, dont_use_mmap=True),
+        batch_size,
+        2,
+    )
+
+
+def test_file_reader_missing_object():
+    with assert_raises(RuntimeError, glob="*S3 Object not found. bucket=*object=*"):
+        pipe = file_pipe(files=[f"s3://{s3.BUCKET}/{DATA_PREFIX}/no-such-object.dat"])
+        pipe.build()
+        pipe.run()
+
+
+def test_file_reader_missing_bucket():
+    with assert_raises(RuntimeError, glob="*NoSuchBucket*"):
+        file_pipe(file_root=f"s3://{s3.MISSING_BUCKET}/data").build()