Fix S3 root-level object discovery, document remote file_root, add S3 video tests

dali/operators/reader/loader/discover_files_s3.cc

@@ -51,7 +51,8 @@ std::vector<FileLabelEntry> s3_discover_files(const std::string &file_root,
       [&](const std::string &object_key, size_t object_size) {
         auto p = std::filesystem::relative(object_key, parent_object_key);
         auto path_elems = count_elems(p);
-        assert(path_elems >= 2);
+        if (path_elems < 2)
+          return;  // objects directly under the root belong to no class
         if (path_elems > 2)
           return;  // we only look at one subdir level
         const auto& subdir = p.begin()->native();

dali/operators/video/frames_decoder_base.cc

@@ -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.
@@ -18,6 +18,8 @@
 #include "dali/core/error_handling.h"
 #include "dali/core/small_vector.h"
 #include "dali/operators/video/video_utils.h"
+#include "dali/util/file.h"
+#include "dali/util/uri.h"
 
 namespace dali {
 
@@ -70,8 +72,63 @@ int64_t seek_memory_video_file(void *data_ptr, int64_t new_position, int origin)
   return memory_video_file->Seek(new_position, origin);
 }
 
+// The callbacks below are called from libavformat (C code) - exceptions must not escape them.
+
+int read_input_stream(void *data_ptr, uint8_t *av_io_buffer, int av_io_buffer_size) {
+  InputStream *stream = static_cast<InputStream *>(data_ptr);
+  try {
+    // Some streams (e.g. S3) fail when asked to read past the end, so we check it here.
+    ssize_t left = stream->SSize() - stream->TellRead();
+    if (left <= 0)
+      return AVERROR_EOF;
+    size_t n = stream->Read(av_io_buffer, std::min<ssize_t>(left, av_io_buffer_size));
+    return n > 0 ? static_cast<int>(n) : AVERROR_EOF;
+  } catch (const std::exception &e) {
+    DALI_WARN(make_string("Error while reading video data: ", e.what()));
+    return AVERROR(EIO);
+  }
+}
+
+int64_t seek_input_stream(void *data_ptr, int64_t new_position, int origin) {
+  InputStream *stream = static_cast<InputStream *>(data_ptr);
+  try {
+    origin &= ~AVSEEK_FORCE;
+    switch (origin) {
+      case AVSEEK_SIZE:
+        return stream->Size();
+      case SEEK_SET:
+      case SEEK_CUR:
+      case SEEK_END:
+        stream->SeekRead(new_position, origin);
+        return stream->TellRead();
+      default:
+        return AVERROR(EINVAL);
+    }
+  } catch (const std::out_of_range &) {
+    return AVERROR(EINVAL);
+  } catch (const std::exception &e) {
+    DALI_WARN(make_string("Error while seeking in video data: ", e.what()));
+    return AVERROR(EIO);
+  }
+}
+
 }  // namespace detail
 
+namespace {
+
+/**
+ * @brief Checks whether the path should be opened with FileStream::Open instead of libavformat.
+ *
+ * libavformat is built only with the "file" protocol, so any other scheme (e.g. s3://, gs://)
+ * is handled by DALI's FileStream.
+ */
+bool IsRemotePath(const std::string &path) {
+  auto uri = URI::Parse(path, URI::ParseOpts::AllowNonEscaped);
+  return uri.valid() && uri.scheme() != "file";
+}
+
+}  // namespace
+
 int FramesDecoderBase::OpenFile(const std::string& filename) {
   LOG_LINE << "Opening file " << filename << std::endl;
   ctx_.reset(avformat_alloc_context());
@@ -86,21 +143,43 @@ int FramesDecoderBase::OpenFile(const std::string& filename) {
 
 int FramesDecoderBase::OpenMemoryFile(MemoryVideoFile &memory_video_file) {
   LOG_LINE << "Opening memory file" << std::endl;
+  static constexpr int DEFAULT_AV_BUFFER_SIZE = (1 << 15);
+  return OpenCustomIO(&memory_video_file,
+                      detail::read_memory_video_file,
+                      detail::seek_memory_video_file,
+                      DEFAULT_AV_BUFFER_SIZE);
+}
+
+int FramesDecoderBase::OpenInputStream(InputStream &stream) {
+  LOG_LINE << "Opening input stream" << std::endl;
+  // libavformat assumes that the stream starts at the current position
+  stream.SeekRead(0, SEEK_SET);
+  // Each read may result in a separate request to a remote storage, so we use a larger buffer.
+  static constexpr int INPUT_STREAM_AV_BUFFER_SIZE = (1 << 20);
+  return OpenCustomIO(&stream,
+                      detail::read_input_stream,
+                      detail::seek_input_stream,
+                      INPUT_STREAM_AV_BUFFER_SIZE);
+}
+
+int FramesDecoderBase::OpenCustomIO(void *opaque,
+                                    int (*read_packet)(void *opaque, uint8_t *buf, int buf_size),
+                                    int64_t (*seek)(void *opaque, int64_t offset, int whence),
+                                    int buffer_size) {
   ctx_.reset(avformat_alloc_context());
   DALI_ENFORCE(ctx_, "Could not alloc avformat context");
 
-  static constexpr int DEFAULT_AV_BUFFER_SIZE = (1 << 15);
-  uint8_t* buffer = static_cast<uint8_t*>(av_malloc(DEFAULT_AV_BUFFER_SIZE));
+  uint8_t* buffer = static_cast<uint8_t*>(av_malloc(buffer_size));
   DALI_ENFORCE(buffer, "Could not alloc avio context buffer");
 
   auto avio_ctx = avio_alloc_context(
     buffer,
-    DEFAULT_AV_BUFFER_SIZE,
+    buffer_size,
     0,
-    &memory_video_file,
-    detail::read_memory_video_file,
+    opaque,
+    read_packet,
     nullptr,
-    detail::seek_memory_video_file);
+    seek);
 
   if (!avio_ctx) {
     av_freep(&buffer);
@@ -218,7 +297,13 @@ FramesDecoderBase::FramesDecoderBase(const std::string &filename, DALIImageType
   DALI_ENFORCE(image_type == DALI_YCbCr || image_type == DALI_RGB,
                make_string("Invalid image type: ", image_type));
   image_type_ = image_type;
-  int ret = OpenFile(filename);
+  int ret = -1;
+  if (IsRemotePath(filename)) {
+    input_stream_ = FileStream::Open(filename);
+    ret = OpenInputStream(*input_stream_);
+  } else {
+    ret = OpenFile(filename);
+  }
   if (ret < 0) {
     DALI_WARN(make_string("Failed to open video file \"", Filename(), "\", due to ",
                           av_error_string(ret)));
@@ -255,6 +340,30 @@ FramesDecoderBase::FramesDecoderBase(const char *memory_file, size_t memory_file
   next_frame_idx_ = 0;
 }
 
+FramesDecoderBase::FramesDecoderBase(std::unique_ptr<InputStream> stream,
+                                     std::string_view source_info, DALIImageType image_type) {
+  av_log_set_level(AV_LOG_ERROR);
+  filename_ = source_info;
+  DALI_ENFORCE(stream, "The input stream must not be null");
+  DALI_ENFORCE(image_type == DALI_YCbCr || image_type == DALI_RGB,
+               make_string("Invalid image type: ", image_type));
+  image_type_ = image_type;
+  input_stream_ = std::move(stream);
+  int ret = OpenInputStream(*input_stream_);
+  if (ret < 0) {
+    DALI_WARN(make_string("Failed to open video file \"", Filename(), "\" due to: ",
+                          av_error_string(ret)));
+    return;
+  }
+
+  packet_.reset(av_packet_alloc());
+  DALI_ENFORCE(packet_, "Could not allocate av packet");
+
+  is_valid_ = true;
+  can_seek_ = true;
+  next_frame_idx_ = 0;
+}
+
 void FramesDecoderBase::ParseNumFrames() {
   num_frames_ = 0;
   while (true) {
@@ -480,6 +589,11 @@ void FramesDecoderBase::Reset() {
     DALI_ENFORCE(ret >= 0,
                  make_string("Could not open video file from memory buffer due to: ",
                              av_error_string(ret)));
+  } else if (input_stream_) {
+    ret = OpenInputStream(*input_stream_);
+    DALI_ENFORCE(ret >= 0,
+                 make_string("Could not open video file \"", Filename(),
+                    "\" due to: ", av_error_string(ret)));
   } else {
     ret = OpenFile(Filename());
     DALI_ENFORCE(ret >= 0,

dali/operators/video/frames_decoder_base.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2021-2022, 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.
@@ -30,6 +30,7 @@ extern "C" {
 #include "dali/core/boundary.h"
 #include "dali/core/common.h"
 #include "dali/core/span.h"
+#include "dali/core/stream.h"
 #include "dali/core/tensor_shape.h"
 #include "dali/core/unique_handle.h"
 #include "dali/pipeline/data/types.h"
@@ -177,7 +178,11 @@ class DLL_PUBLIC FramesDecoderBase {
   /**
    * @brief Initialize the decoder from a file.
    *
-   * @param filename Path to a video file.
+   * Local paths (with no scheme or with the ``file://`` scheme) are opened directly by
+   * libavformat. URLs with any other scheme (e.g. ``s3://``, ``gs://``) are opened with
+   * FileStream::Open and read through the InputStream interface.
+   *
+   * @param filename Path or URL of a video file.
    * @param image_type Image type of the video.
    */
   explicit FramesDecoderBase(const std::string &filename, DALIImageType image_type = DALI_RGB);
@@ -192,6 +197,17 @@ class DLL_PUBLIC FramesDecoderBase {
   explicit FramesDecoderBase(const char *memory_file, size_t memory_file_size,
                              std::string_view source_info = {}, DALIImageType image_type = DALI_RGB);
 
+  /**
+   * @brief Initialize the decoder from an input stream.
+   *
+   * @param stream Stream with video file data; the decoder takes ownership of it.
+   *               The stream must cover the entire video file, including the header.
+   * @param source_info Source information for the video file.
+   */
+  explicit FramesDecoderBase(std::unique_ptr<InputStream> stream,
+                             std::string_view source_info = {},
+                             DALIImageType image_type = DALI_RGB);
+
   /**
    * @brief Number of frames in the video. It returns 0, if this information is unavailable.
    */
@@ -387,7 +403,7 @@ class DLL_PUBLIC FramesDecoderBase {
   FramesDecoderBase& operator=(FramesDecoderBase&&) = default;
 
   std::string Filename() {
-    return filename_.size() ? filename_ : "memory file";
+    return filename_.size() ? filename_ : input_stream_ ? "input stream" : "memory file";
   }
 
   bool IsValid() {
@@ -419,6 +435,7 @@ class DLL_PUBLIC FramesDecoderBase {
 
   int OpenFile(const std::string& filename);
   int OpenMemoryFile(MemoryVideoFile& memory_video_file);
+  int OpenInputStream(InputStream& stream);
 
   FrameIndex index_;
 
@@ -441,11 +458,20 @@ class DLL_PUBLIC FramesDecoderBase {
 
   std::string GetAllStreamInfo() const;
 
+  /**
+   * @brief Opens the format context with custom I/O, reading through the provided callbacks.
+   */
+  int OpenCustomIO(void *opaque,
+                   int (*read_packet)(void *opaque, uint8_t *buf, int buf_size),
+                   int64_t (*seek)(void *opaque, int64_t offset, int whence),
+                   int buffer_size);
+
   int channels_ = 3;
   bool is_vfr_ = false;
 
   std::string filename_ = {};
   std::unique_ptr<MemoryVideoFile> memory_video_file_;
+  std::unique_ptr<InputStream> input_stream_;
 
   int num_frames_ = -1;
   bool can_seek_ = true;  // at first, we assume that the video is seekable

dali/operators/video/frames_decoder_cpu.cc

@@ -34,6 +34,12 @@ FramesDecoderCpu::FramesDecoderCpu(const char *memory_file, size_t memory_file_s
   is_valid_ = is_valid_ && SelectVideoStream();
 }
 
+FramesDecoderCpu::FramesDecoderCpu(std::unique_ptr<InputStream> stream,
+                                   std::string_view source_info, DALIImageType image_type)
+  : FramesDecoderBase(std::move(stream), source_info, image_type) {
+  is_valid_ = is_valid_ && SelectVideoStream();
+}
+
 void FramesDecoderCpu::CopyFrame(uint8_t *dst, const uint8_t *src) {
   std::memcpy(dst, src, FrameSize());
 }

dali/operators/video/frames_decoder_cpu.h

@@ -45,6 +45,19 @@ class DLL_PUBLIC FramesDecoderCpu : public FramesDecoderBase {
    */
   FramesDecoderCpu(const char *memory_file, size_t memory_file_size, std::string_view = {}, DALIImageType image_type = DALI_RGB);
 
+  /**
+   * @brief Construct a new FramesDecoder object.
+   *
+   * @param stream Input stream with video file data; the decoder takes ownership of it.
+   * @param source_info Source info of the video file.
+   * @param image_type Image type of the video.
+   *
+   * @note This constructor assumes that the stream covers the entire video file,
+   * including the header.
+   */
+  explicit FramesDecoderCpu(std::unique_ptr<InputStream> stream, std::string_view source_info = {},
+                            DALIImageType image_type = DALI_RGB);
+
   FramesDecoderCpu(FramesDecoderCpu&&) = default;
 
   bool ReadNextFrame(uint8_t *data) override;

dali/operators/video/frames_decoder_gpu.cc

@@ -443,6 +443,18 @@ FramesDecoderGpu::FramesDecoderGpu(const char *memory_file, size_t memory_file_s
   }
 }
 
+FramesDecoderGpu::FramesDecoderGpu(std::unique_ptr<InputStream> input_stream,
+                                   std::string_view source_info, cudaStream_t stream,
+                                   DALIImageType image_type)
+    : FramesDecoderBase(std::move(input_stream), source_info, image_type),
+      frame_buffer_(num_decode_surfaces_),
+      stream_(stream) {
+  is_valid_ = is_valid_ && SelectVideoStream();
+  if (is_valid_) {
+    InitGpuParser();
+  }
+}
+
 int FramesDecoderGpu::ProcessPictureDecode(CUVIDPICPARAMS *picture_params) {
   // Sending empty packet will call this callback.
   // If we want to flush the decoder, we do not need to do anything here

dali/operators/video/frames_decoder_gpu.h

@@ -151,6 +151,20 @@ class DLL_PUBLIC FramesDecoderGpu : public FramesDecoderBase {
                    std::string_view source_info = {}, cudaStream_t stream = 0,
                    DALIImageType image_type = DALI_RGB);
 
+  /**
+   * @brief Construct a new FramesDecoder object.
+   *
+   * @param input_stream Input stream with video file data; the decoder takes ownership of it.
+   * @param source_info Source info of the video file.
+   * @param stream CUDA stream to use for decoding.
+   * @param image_type Image type of the video.
+   * @note This constructor assumes that the input stream covers the entire video file,
+   * including the header.
+   */
+  explicit FramesDecoderGpu(std::unique_ptr<InputStream> input_stream,
+                            std::string_view source_info = {}, cudaStream_t stream = 0,
+                            DALIImageType image_type = DALI_RGB);
+
   bool ReadNextFrame(uint8_t *data) override;
 
   void SeekFrame(int frame_id) override;