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

[jantonguirao]

Category:

New feature (S3 support for webdataset and RecordIO index files, plus integration tests and
documentation)

Description:

DALI's s3:// support had no test coverage. This adds integration tests that start a mock S3 server
in the background, seed a bucket, and assert that the readers return byte-identical results to the
same data read from local disk.

The mock server is moto, started and torn down by the test module itself, so no external service or
CI-side fixture is needed.

While covering fn.readers.webdataset, it became clear that index_paths only ever worked for local
files: ParseIndexFile opened it through a raw std::ifstream, unlike paths (the tar shards),
which already goes through FileStream::Open and therefore already supports s3://. This PR adds
that support: index_paths can now point at S3 (or any other FileStream-backed scheme) alongside
the shards it describes. The index is streamed through FileStreamBuf (the same utility
IndexedFileLoader already uses for local index files) rather than read into memory in full, since
index_paths is a public, caller-controlled argument and the index can be arbitrarily large. Wiring
this up also surfaced and fixed a latent bug in FileStreamBuf itself: the read exactly at
end-of-file could issue an out-of-range request on a remote stream (a local file just returns 0
bytes; S3 raises on a byte range starting at or past the object's end) - previously unreachable since
nothing had used FileStreamBuf against a remote FileStream before.

Review raised the same question for the other index-file readers: fn.readers.tfrecord was already
fine (IndexedFileLoader::ReadIndexFile has gone through FileStream::Open/FileStreamBuf since
#5515), but fn.readers.mxnet's RecordIOLoader::ReadIndexFile had the identical local-only
std::ifstream gap webdataset_loader.cc just lost - its shards (paths_) already used
FileStream::Open, only its index didn't. Fixed the same way, in scope here, with matching S3
coverage.

Two implementation details in the mock-server helper are load-bearing and are commented in the source
so they are not "simplified" away later:

• The server runs in a subprocess, not in-process. Pipeline::Build is bound without
  py::call_guard<py::gil_scoped_release>, and S3 object listing happens inside Build(), so an
  in-process server thread is GIL-starved and the request eventually times out in libcurl
  (curlCode: 28). Every moto example shows the in-process form, hence the comment.
• The endpoint is always an IP literal. DALI does not set useVirtualAddressing, so the AWS SDK
  only selects path-style addressing when the endpoint host is an IP. With a host name the bucket is
  prepended to the host instead, which requires a matching DNS entry (and MINIO_DOMAIN, for a real
  MinIO).

Additional information:

Affected modules and functionalities:

• dali/operators/reader/loader/webdataset_loader.cc/.h: index_paths now opens through
  FileStream::Open and is parsed through a streaming FileStreamBuf, instead of a local-only
  std::ifstream read fully into memory.
• dali/operators/reader/loader/recordio_loader.h: same fix for fn.readers.mxnet's index file.
• dali/util/file.h: fixes a latent end-of-file bug in FileStreamBuf that only a remote
  FileStream can trigger (see above).
• dali/test/python/s3_test_utils.py (new): mock server lifecycle, env setup, bucket helpers and
  the require_mock_server/require_s3_support preconditions. Mock-only: nothing here points at a
  real endpoint, so credentials/region/bucket are fixed rather than configurable.
• dali/test/python/reader/test_s3.py (new): the tests.
• qa/TL0_python-self-test-readers-decoders/test_nofw.sh: adds moto, flask, flask-cors and
  boto3 to pip_packages. Plain moto does not pull in flask; that lives in the moto[server]
  extra, which also drags in cfn-lint, docker and more.
• qa/TL1_python-self-test_conda/test_nofw.sh: adds the same packages, so the conda suite gets real
  S3 coverage instead of skipping the module.
• qa/TL0_python-self-test_tegra/test_body.sh: excludes test_s3.py explicitly - unlike conda,
  aarch64 wheel availability makes installing moto/boto3 there impractical.
• qa/TL0_python-self-test-readers-decoders/test_body.sh: the sanitizer-skip removed from
  test_body.sh for this module was actually redundant - SKIP_TESTS="test_numpy.py test_s3.py"
  here already excludes it under DALI_ENABLE_SANITIZERS.
• docs/env_vars.rst: documents AWS_ENDPOINT_URL, which the tests rely on and which was
  previously undocumented.

Key points relevant for the review:

• No DALI source changes were needed for the file-reader coverage - S3ClientManager already reads
  AWS_ENDPOINT_URL. The webdataset index_paths change is the one real source change in this PR.
• require_mock_server/require_s3_support raise instead of skipping when boto3/moto are missing or
  BUILD_AWSSDK is off, mirroring the equivalent decision already made on the GCS PR: every suite
  that runs dali/test/python/reader is expected to have these, so a missing dependency there is an
  environment bug, not something that should silently skip the coverage.
• It is skipped under sanitizers because qa/leak.sup has no suppressions for the AWS SDK, libcurl
  or OpenSSL globals, and the mock server subprocess would inherit LD_PRELOAD.
• Signing is deliberately not covered. moto does not verify SigV4, so the key-with-a-space case
  pins URL path construction rather than canonical-request encoding. Signing fidelity would need a
  separate run against MinIO or real S3; not attempted here.
• Cost is about 16 s, CPU only (device_id=None), on a suite that already runs for several minutes.
  Roughly 9 s of that is the listing-pagination case, which seeds 1100 objects; the multi-chunk
  webdataset-index test is marked slow and adds a few more seconds.

Tests:

• Existing tests apply
• New tests added
  • Python tests
  • GTests
  • Benchmark
  • Other
• N/A

New module reader/test_s3.py, covering:

Test	What it exercises
test_file_reader_file_root	ListObjectsV2 with a prefix, label assignment, file_filters over keys
test_file_reader_files_arg	HeadObject + ranged GetObject, a key containing a space, a 0-byte object
test_file_reader_listing_pagination	continuation tokens across the 1000-key page limit (1100 keys)
test_webdataset_index_inferred	index inferred from the tar, i.e. many small ranged GETs
test_webdataset_local_index	S3 shard with a local index file
test_webdataset_remote_index	S3 shard with the index file also on S3
test_webdataset_remote_index_multi_chunk	a ~137 KiB remote index, forcing multiple FileStreamBuf reads (including the end-of-file boundary case)
test_mxnet_reader_local_index	S3 RecordIO shard with a local index file
test_mxnet_reader_remote_index	S3 RecordIO shard with the index file also on S3
test_file_reader_missing_object	error surface for a missing key
test_file_reader_missing_bucket	error surface for a missing bucket

Readers that do not route through FileStream::Open (FITS, LMDB-backed caffe/caffe2, nemo_asr,
sequence, video) are out of scope. fn.readers.numpy over s3:// is also left out of this PR - it
hits a separate bug, fixed in DALI-4887; it can be added to this matrix once that lands.

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-4010

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; the latest change correctly fixes the MXNet test pipeline’s multi-output structure, and no outstanding findings remain.

Summary

The PR adds mock-S3 integration coverage and enables WebDataset and RecordIO index files to be streamed through FileStream from S3-compatible storage.

• Adds S3-backed reader tests and mock-server lifecycle utilities.
• Streams remote WebDataset and MXNet index files without loading them fully into memory.
• Corrects FileStreamBuf handling of partial reads and end-of-file boundaries.
• Moves the process-wide S3 client manager implementation into one shared-library definition.
• Documents AWS_ENDPOINT_URL and updates relevant QA dependencies and exclusions.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  P[Reader paths and index_paths] --> FS[FileStream::Open]
  FS -->|Local path| L[Local FileStream]
  FS -->|s3:// URI| S[S3 FileStream]
  L --> B[FileStreamBuf]
  S --> B
  B --> I[Streaming index parser]
  I --> W[WebDataset loader]
  I --> M[RecordIO loader]

Loading

_{Reviews (18) · Last reviewed commit: "Fix mxnet_pipe: wrap fn.readers.mxnet's ..."}

[Copilot]

🟡 Changes recommended

Unresolved critical unsupported-build handling and multiple moderate external-S3 lifecycle and configuration issues remain.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds S3 reader integration tests using a subprocess-based Moto server, with optional external S3-compatible endpoints.

Changes:

• Adds S3 lifecycle, seeding, environment, and capability helpers.
• Tests file readers, WebDataset readers, pagination, and error handling.
• Updates dependencies, sanitizer exclusions, and AWS_ENDPOINT_URL documentation.

File summaries

File	Changes and review comments
qa/TL0_python-self-test-readers-decoders/test_nofw.sh	Adds Moto, Flask, Flask-CORS, and boto3 dependencies.
qa/TL0_python-self-test-readers-decoders/test_body.sh	Skips the S3 module under sanitizers.
docs/env_vars.rst	Documents AWS_ENDPOINT_URL.
dali/test/python/s3_test_utils.py	Provides S3 server and setup helpers. Critical (2 votes): the unsupported-build probe does not exercise FileStream::Open, so BUILD_AWSSDK=OFF builds fail instead of skipping. Moderate: external runs do not reliably handle missing boto3 (1), overwrite standard AWS credentials (2), can hang before the startup timeout (1), force us-east-1 (1), and do not support external-region bucket creation (1).
dali/test/python/reader/test_s3.py	Adds file-reader and WebDataset tests, pagination, and missing-resource checks. Moderate: external runs leave seeded objects behind (2) and rely on a potentially pre-existing missing-bucket name (1).

Review details

Suppressed comments (4)

dali/test/python/reader/test_s3.py:191

• The negative test relies on dali-no-such-bucket never existing, but that is not guaranteed when DALI_TEST_S3_ENDPOINT points at a persistent MinIO or real S3 service (the name could already be present or owned by the test account). In that case listing succeeds or returns a different error and the NoSuchBucket assertion fails. Use a per-run unique bucket name or make the missing-bucket case configurable/skip it for external services.

def test_file_reader_missing_bucket():
    with assert_raises(RuntimeError, glob="*NoSuchBucket*"):
        file_pipe(file_root="s3://dali-no-such-bucket/data").build()

dali/test/python/s3_test_utils.py:193

• When DALI_TEST_S3_ENDPOINT is set, this returns before checking boto3, but setUpModule still calls s3_client() and that function unconditionally imports boto3 at line 166. Consequently an external MinIO/real-S3 run without boto3 fails with an uncaught import error instead of being skipped; check boto3 unconditionally and only make moto.server conditional on the absence of the external endpoint.

    if os.environ.get("DALI_TEST_S3_ENDPOINT"):
        return

dali/test/python/s3_test_utils.py:94

• This blocking readline() is executed before _wait_until_serving() applies startup_timeout_s. If moto import or server.start() hangs before printing a port, the whole reader suite hangs indefinitely rather than failing after the configured 60 seconds; read/poll the pipe with the deadline (or add a watchdog) before waiting for the startup line.

        line = self._proc.stdout.readline()  # the child reports the port it bound

dali/test/python/s3_test_utils.py:29

• REGION is forced to us-east-1 and then exported over both standard region variables, so the advertised external real-S3 mode cannot use a bucket in another region: the unconditional create_bucket call will require a location constraint, and DALI requests will be signed with the wrong region. Allow the external test region to be configured and create the bucket with the matching LocationConstraint when needed.

REGION = "us-east-1"

• Files reviewed: 5/5 changed files
• Comments generated: 3
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[jantonguirao]

!build

[dali-automaton]

CI MESSAGE: [67423372]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [67423372]: BUILD FAILED

[JanuszL]

I wonder if we can extend package installation logic to catch uninstall-no-record-file and then apply --ignore-installed` to pip cmd.

[jantonguirao]

Good idea, and pip itself now recommends exactly that — in pip 26.2 the hint attached to this error changed from --force-reinstall --no-deps to --ignore-installed --no-deps. I prototyped it and it does work, but it has enough sharp edges that I would rather send it on its own than fold it into a test PR.

What I found, on a venv with a blinker 1.7.0 whose RECORD I deleted to emulate the apt layout:

1. The string to match is not stable across pip versions. pip 25.1 prints error: uninstall-no-record-file; pip 22.0 prints ERROR: Cannot uninstall blinker 1.7.0, RECORD file not found. with no reference token at all; and there is the sibling uninstall-distutils-installed-package. So the match has to be a three-way alternation, and all of it goes to stderr, so the capture needs 2>&1.
2. --ignore-installed is not scoped to the offending package — it clears the resolver's entire view of installed distributions, so nothing is uninstalled and everything is overwritten in place. After a successful retry, site-packages holds both blinker-1.7.0.dist-info and blinker-1.9.0.dist-info, importlib.metadata.version("blinker") reports the stale 1.7.0, pip list agrees, and files dropped between the two versions survive.
3. It does not heal the environment. A later pip install "blinker>=1.9" in the same container fails with the same error, because the RECORD-less dist-info is still there. So the fallback rescues each command individually and leaves the underlying state broken. (pip's older --force-reinstall --no-deps suggestion does not work either — same error.)
4. The shell side has two traps. With set -e in force, the obvious { cmd1 || cmd2; echo $? > f; } | tee log silently loses the status: errexit kills the pipeline subshell on the failing AND-OR list before the echo runs. And $PIPESTATUS is only meaningful with set -o pipefail — a bare pipeline || ret=$? would read tee's 0 and turn a failed install into a pass. The version I have wraps the AND-OR list in an if, passes the status through a temp file, keeps ${install_cmd} unquoted so globs and multi-package commands behave as today, and returns pip's real exit code so the call sites are unaffected.

Since install_pip_pkg is on the path of all 51 suites that source qa/test_template.sh, plus the CUDA wheel installs, the numpy<2 downgrade in TL0_python-self-test-core and the TF install in TL1_tensorflow-dali_test, I will send it as its own PR and relax this pin to flask<4 in that same commit, so that PR's CI run is itself the proof the new path fires. The upper bound stays because moto is pinned at 5.2.3 and moto.server sits on flask + flask-cors.

For the record, the version facts in the comment check out — flask 3.0.3 requires blinker>=1.6.2, flask 3.1.0 requires blinker>=1.9 — and this pin is exactly what the last CI failure was: that pipeline ran the unpinned flask revision and died in pip install flask with uninstall-no-record-file before a single test executed.

[jantonguirao]

Opened as #6485.

One adjustment from what I said above: it doesn't relax the flask pin in the same commit after all. That pin lives on this PR's branch, which hasn't merged, so qa/TL0_python-self-test-readers-decoders/test_nofw.sh on main has no flask entry yet to relax - there's nothing there for #6485 to touch. Once both land, relaxing flask==3.0.3 to flask<4 here is a one-line follow-up, and it'll double as the first real CI exercise of the new retry path.

[JanuszL]

I think you can merge #6485 (already reviewed) and adjust the code here.

[jantonguirao]

#6485 is approved, checks green, and mergeable. Once it lands on main I'll relax flask==3.0.3 to flask<4 here in the same PR so this branch's CI actually exercises the new retry path.

[jantonguirao]

Pushed 27f41c3, which addresses the review. Summary for anyone reading the diff rather than the threads:

Change	Raised by
Bound the mock-server startup read with a watchdog — a child alive but stuck before printing its port hung the whole suite, since the 60 s deadline was only armed afterwards	Greptile P2
Send CreateBucketConfiguration outside us-east-1 (and only outside it) — the old call broke for anyone with AWS_DEFAULT_REGION exported, not just for real S3	Greptile P1
Never delete the bucket; teardown removes the run's own prefix only	Greptile P1
Leave credentials to the SDK chain when DALI_TEST_S3_ENDPOINT is set, instead of exporting dummy keys over a profile / SSO session / instance role	Copilot
Generate the webdataset index inside the try, so a failure there no longer leaks every object already uploaded	found while verifying the above
Raise on the per-key errors DeleteObjects returns in its response rather than raising	found while verifying the above
Derive the bucket used by the error paths from the per-run uuid — a hardcoded dali-no-such-bucket can be owned by someone else on real S3, which answers AccessDenied, not NoSuchBucket	Copilot (suppressed comment)
Require boto3 in both modes and moto only in the mock one	Copilot (suppressed comment)

Two threads are left open on purpose: the single-vs-double backtick one, where the new heading matches all 22 headings on that page and changing only it would look worse, and the install_pip_pkg / --ignore-installed one, which I would rather land as its own PR — the reasoning is in the thread.

Verified locally before pushing, against moto 5.2.3 (the pinned version), with black and flake8 clean:

• default mock path: 7/7 pass
• AWS_DEFAULT_REGION=eu-west-1: 7/7 pass (the previous create_bucket call fails there with IllegalLocationConstraintException)
• external-endpoint path against a standing server pre-seeded with a bucket and a foreign object: 7/7 pass, bucket and foreign object untouched, the run's prefix gone afterwards
• startup watchdog with the child patched to time.sleep(600) and a 3 s timeout: raises after 3.0 s where it previously blocked indefinitely

One thing worth flagging: the earlier CI run never executed these tests. It died in pip install flask — that revision had flask unpinned, so it pulled 3.1.3, which wants blinker>=1.9, and the preinstalled apt blinker cannot be uninstalled. The pin now in the branch avoids it, but this next run is the first time the module itself will have run in CI.

[jantonguirao]

!build

[dali-automaton]

CI MESSAGE: [67947457]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [67947457]: BUILD FAILED

[jantonguirao]

!build

[dali-automaton]

CI MESSAGE: [68014495]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [68014495]: BUILD PASSED

[dali-automaton]

CI MESSAGE: [68235175]: BUILD STARTED

[jantonguirao]

Added dali/util/file_test.cc (wired into CMake, same convention as uri_test.cc) with an adversarial gtest for FileStreamBuf, mocking two real, divergent backends:

• ShortReadStream models ODirectFileStream (raw read() syscall) — can legitimately return fewer bytes than requested without that meaning EOF.
• RemoteLikeStream models S3FileStream — never returns 0 for EOF; a byte range starting at or past the object end raises instead.

Before pushing, I compiled and ran this exact test file, unmodified, against the header as it stood at each of the last 3 commits (standalone harness, not the full dali_test binary — same header file, just not going through the whole DALI CMake build):

Version	GenuinePartialReadsDoNotTruncate	RemoteStreamNeverReadsPastEnd	RemoteStreamExactMultipleOfBuffer
before 9249de1 (no eof_/remaining_ tracking)	✅ pass	❌ throws	❌ throws
9249de1 ("short read == EOF")	❌ truncates to 3/200 bytes	✅ pass	❌ still throws
9389d49 (current, remaining_-tracked retry)	✅ pass	✅ pass	✅ pass

The middle row is the interesting one: the first fix didn't even fully solve what it targeted — an exact-buffer-size-multiple payload never triggers a "short" read, so eof_ never got set, and it still throws on the boundary. Only the remaining_-based version passes all three. Pushed as ff31c53.

[dali-automaton]

CI MESSAGE: [68238308]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [68240550]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [68240550]: BUILD PASSED

[dali-automaton]

CI MESSAGE: [68348405]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [68348405]: BUILD FAILED

[dali-automaton]

CI MESSAGE: [68364637]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [68365175]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [68365175]: BUILD FAILED

[dali-automaton]

CI MESSAGE: [68392609]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [68392609]: BUILD PASSED