[GDS] wait_buffers can prematurely recycle in-flight buffers and corrupt NumPy reader output

Version

2.1.0, official nvidia-dali-cuda130==2.1.0 wheel.

The affected code is in gds_mem.cc at v2.1.0.
Source inspection also found the same file content on main at
3e66711c1605cba7de43fecfc3d41882d07bd4fc.
The runtime reproduction below was performed against the official 2.1.0 wheel;
main was checked at source level only.

Describe the bug.

GDSStagingEngine::wait_buffers() can return staging buffers to the free pool
while a newer batch of GPU copies still needs their contents. A subsequent
cuFileRead() into a buffer obtained through the normal engine API can then
overwrite the source of a pending copy, silently mixing data from different
files in the output.

We investigated this after intermittent cross-file corruption with
fn.readers.numpy(device="gpu"). Saved outputs included bytes from other input
files, including NPY header bytes inside an array payload. Reading on CPU and
then copying to GPU did not exhibit the same corruption in the observed runs.
The controlled staging-engine reproduction below isolates a concrete race;
it does not rerun the full training workload.

Expected: A staging buffer becomes available only after all submitted GPU
copies that read from it have completed. Waiting for an older recording of a
reusable event must not authorize reclaiming a newer batch.

Actual: With the original implementation, the reproduction exposes 32
still-pending buffers in ready_buffers_ and leases another pending buffer to a
reader. A later file read overwrites that buffer. One complete 2 MiB output
block, expected to contain the integer 66, instead contains 999 from another
file. CUDA and cuFile calls complete without reporting an error.

Interleaving:

Normal API calls exhaust the default 64-buffer pool: 33 buffers are in
scheduled_ for batch G1 and 31 are in unscheduled_. The submission
condition is unscheduled_.size() > 32.
Thread A enters wait_buffers(), releases the mutex, and synchronizes
ready_ for G1. G1 completes. A has not yet reacquired the mutex.
Thread B enters, safely reclaims G1's 33 buffers, and performs two more normal
read/copy operations. The 31 queued copies plus these two trigger batch G2.
B updates scheduled_ and re-records ready_ for G2.
G2 is still pending when A reacquires the mutex.
A unconditionally moves the current scheduled_ into ready_buffers_
and clears it, although A only waited for G1. These are now G2's buffers.
A reader obtains one of these buffers through get_staging_buffer() and
calls cuFileRead() for another file. When G2 executes, it copies the
overwritten contents into the earlier output.

The event re-recording is deliberately scheduled after A's synchronize call
returns, so this does not depend on ambiguous timing inside
cudaEventSynchronize() itself.

Minimum reproducible example

The self-contained script below generates synthetic files and builds a C++/CUDA
test against the installed DALI library. It requires no private dataset,
PyTorch, or distributed training setup. Requirements are Linux x86_64, a
CUDA-13-compatible GPU, the official 2.1.0 wheel with headers, g++, and a CUDA
toolkit with cuFile development headers/libraries. Use an output filesystem
that supports O_DIRECT and your local cuFile configuration.

Save the code in the collapsed section as repro.py, then run:

python -m pip install nvidia-dali-cuda130==2.1.0
export CUDA_HOME=/path/to/cuda-13.0
CUDA_VISIBLE_DEVICES=0 python repro.py --work-dir ./gds-race-output

If DALI is outside the selected Python environment, add
--dali-dir /path/to/site-packages/nvidia/dali.
The default GPU architecture is native; --arch sm_120 can override it.

The default mode runs the original and event-recheck variants consecutively.
In original mode, a zero exit code means the test successfully reproduced
corruption. In event-recheck mode, a zero exit code requires zero corruption.

Instrumentation and scope:

get_staging_buffer, allocation, submission, return_ready, and GPU copies
use the real DALI shared library. File reads use real cuFileRead calls.
The executable interposes only wait_buffers, using its v2.1.0 source body
with scheduling hooks before/after the real cudaEventSynchronize, or that
same body with the proposed change below. Installed libraries are unchanged.
A CPU barrier pauses A after synchronization but before relocking. GPU gates
keep G1/G2 pending until explicitly released. Normal public API calls create
the 33/31 pool state; the test does not insert entries into private queues.
-fno-access-control permits inspecting invariants and checking mutex state.
The pools are warmed before gating to avoid allocation synchronization
interfering with the intended interleaving. A watchdog releases stalled
gates, but any run in which it fires fails the test.
This verifies the engine race under controlled scheduling, including actual
data corruption and the proposed fix. It is not an end-to-end NumPy pipeline
or training test, and does not measure natural occurrence probability.

The exact single-file reproducer was packaged, extracted into a separate
directory, compiled, and run against the official wheel before this report.

Complete self-contained reproducer: repro.py

#!/usr/bin/env python3
"""Self-contained controlled DALI 2.1.0 GDS staging-engine race repro.

Run: python repro.py --help
The embedded NVIDIA method retains its original Apache-2.0 notice.
"""

from pathlib import Path
import runpy
import tempfile

FILES = {
'gpu_engine_test.cpp': r'''// Real DALI shared-library engine, real CUDA copies, and real cuFileRead calls.
// Only wait_buffers is interposed, using the original or event-recheck body
// plus explicit test-only scheduling hooks.
#include <cuda_runtime_api.h>
#include <cufile.h>
#include <dlfcn.h>
#include <fcntl.h>
#include <unistd.h>
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <iostream>
#include <memory>
#include <mutex>
#include <optional>
#include <stdexcept>
#include <string>
#include <thread>
#include <unordered_set>
#include <vector>
#include "gpu_test_hooks.hpp"

#ifndef TEST_DALI_VERSION
#define TEST_DALI_VERSION "unknown"
#endif
#ifndef TEST_RECHECK
#define TEST_RECHECK 0
#endif
#ifndef TEST_MODE
#if TEST_RECHECK
#define TEST_MODE "event-recheck"
#else
#define TEST_MODE "original"
#endif
#endif

extern "C" cudaError_t launch_gpu_gate(cudaStream_t, uint32_t *, int);
using Engine = dali::gds::GDSStagingEngine;
using Buffer = dali::gds::GDSStagingBuffer;

static void require(bool value, const std::string &message) {
  if (!value) throw std::runtime_error(message);
}
static void check_cuda(cudaError_t status, const char *where) {
  if (status != cudaSuccess) throw std::runtime_error(std::string(where) + ": " + cudaGetErrorString(status));
}
#define GPU(call) check_cuda((call), #call)

struct Control {
  std::mutex mutex;
  std::condition_variable cv;
  bool before = false, after = false, release = false;
  std::exception_ptr worker_error;
  std::atomic<int> hook_calls{0};
  std::atomic<bool> finished{false}, watchdog_fired{false};
  volatile uint32_t *host_flags = nullptr;
};
static Control control;
static thread_local bool pause_waiter_a = false;
static std::FILE *trace_file = nullptr;
static std::mutex trace_mutex;

static void trace(const char *phase, const void *buffer = nullptr, int file_id = -1,
                  size_t offset = 0, size_t bytes = 0) {
  std::lock_guard<std::mutex> lock(trace_mutex);
  auto ns = std::chrono::steady_clock::now().time_since_epoch().count();
  auto thread = std::hash<std::thread::id>{}(std::this_thread::get_id());
  if (trace_file) {
    std::fprintf(trace_file, "%lld,%zu,%s,%p,%d,%zu,%zu\n", static_cast<long long>(ns),
                 thread, phase, buffer, file_id, offset, bytes);
    std::fflush(trace_file);
  }
}

void BeforeRealEventSynchronize() {
  ++control.hook_calls;
  trace("wait_before_real_cuda_sync");
  if (pause_waiter_a) {
    std::lock_guard<std::mutex> lock(control.mutex);
    control.before = true;
    control.cv.notify_all();
  }
}
void AfterRealEventSynchronize() {
  trace("wait_after_real_cuda_sync");
  if (!pause_waiter_a) return;
  std::unique_lock<std::mutex> lock(control.mutex);
  control.after = true;
  control.cv.notify_all();
  control.cv.wait(lock, [] { return control.release; });
}

static void release_gpu_gates() {
  if (control.host_flags) {
    __sync_synchronize();
    control.host_flags[0] = 1;
    control.host_flags[1] = 1;
    __sync_synchronize();
  }
}

struct Source {
  int fd = -1;
  CUfileHandle_t handle = nullptr;
  explicit Source(const std::string &path) {
    fd = open(path.c_str(), O_RDONLY | O_DIRECT);
    require(fd >= 0, "open O_DIRECT failed: " + path);
    CUfileDescr_t descriptor{};
    descriptor.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD;
    descriptor.handle.fd = fd;
    auto status = cuFileHandleRegister(&handle, &descriptor);
    require(status.err == CU_FILE_SUCCESS, "cuFileHandleRegister failed: " + std::to_string(status.err));
  }
  ~Source() {
    if (handle) cuFileHandleDeregister(handle);
    if (fd >= 0) close(fd);
  }
  void read(void *gpu, size_t bytes, size_t offset, int file_id) {
    trace("cufile_read_begin", gpu, file_id, offset, bytes);
    size_t done = 0;
    while (done < bytes) {
      auto got = cuFileRead(handle, gpu, bytes - done, offset + done, done);
      require(got > 0, "cuFileRead failed/short EOF: " + std::to_string(got));
      done += got;
    }
    trace("cufile_read_end", gpu, file_id, offset, bytes);
  }
};

struct CopyRecord { void *source; int expected; };

int main(int argc, char **argv) {
  require(argc == 3, "usage: test input_directory output_prefix");
  std::string directory = argv[1], prefix = argv[2];
  trace_file = std::fopen((prefix + ".trace.csv").c_str(), "w");
  require(trace_file, "Cannot create trace");
  std::fprintf(trace_file, "steady_clock_ticks,thread,phase,gpu_buffer,file_id,file_offset,bytes\n");
  std::unique_ptr<Engine> engine;
  std::optional<Buffer> held;
  std::thread waiter, watchdog;
  cudaStream_t stream = nullptr;
  uint64_t *output = nullptr;
  uint32_t *mapped_host = nullptr, *mapped_device = nullptr;
  std::unique_ptr<Source> source_a, source_b, source_c;
  bool cufile_opened = false;
  int exit_code = 1;
  try {
    GPU(cudaSetDevice(0));
    cudaDeviceProp properties{};
    GPU(cudaGetDeviceProperties(&properties, 0));
    int driver = 0, runtime = 0, cufile = 0;
    GPU(cudaDriverGetVersion(&driver));
    GPU(cudaRuntimeGetVersion(&runtime));
    auto opened = cuFileDriverOpen();
    require(opened.err == CU_FILE_SUCCESS, "cuFileDriverOpen failed: " + std::to_string(opened.err));
    cufile_opened = true;
    cuFileGetVersion(&cufile);
    // Some cuFile releases close stdin. Keep file descriptors 0-2 out of the file handle pool.
    if (fcntl(STDIN_FILENO, F_GETFL) < 0) {
      int null_fd = open("/dev/null", O_RDONLY);
      if (null_fd != STDIN_FILENO) { dup2(null_fd, STDIN_FILENO); close(null_fd); }
    }
    size_t chunk = dali::gds::GetGDSChunkSize();
    size_t words = chunk / sizeof(uint64_t);
    source_a = std::make_unique<Source>(directory + "/source_a.bin");
    source_b = std::make_unique<Source>(directory + "/source_b.bin");
    source_c = std::make_unique<Source>(directory + "/source_c.bin");
    GPU(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
    GPU(cudaMalloc(reinterpret_cast<void **>(&output), 128 * chunk));
    GPU(cudaHostAlloc(reinterpret_cast<void **>(&mapped_host), 2 * sizeof(uint32_t), cudaHostAllocMapped));
    mapped_host[0] = mapped_host[1] = 0;
    GPU(cudaHostGetDevicePointer(reinterpret_cast<void **>(&mapped_device), mapped_host, 0));
    control.host_flags = mapped_host;
    watchdog = std::thread([] {
      for (int i = 0; i < 300 && !control.finished; ++i)
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
      if (!control.finished) {
        control.watchdog_fired = true;
        release_gpu_gates();
        std::lock_guard<std::mutex> lock(control.mutex);
        control.release = true;
        control.cv.notify_all();
      }
    });
    engine = std::make_unique<Engine>(0);
    engine->set_stream(stream);
    std::vector<CopyRecord> records;
    auto read_and_copy = [&](Engine &target, int index, Source &file, int source_id, size_t offset) {
      auto buffer = target.get_staging_buffer();
      try {
        file.read(buffer.at(0), chunk, offset, source_id);
        if (&target == engine.get()) records.push_back({buffer.at(0), index + 1});
        trace("copy_to_client", buffer.at(0), source_id, offset, chunk);
        target.copy_to_client(output + index * words, chunk, std::move(buffer));
      } catch (...) {
        target.return_unused(std::move(buffer));
        throw;
      }
    };

    // Warm genuine GDS and scratchpad pools BEFORE closing a GPU gate. This
    // avoids cudaMalloc synchronization becoming an artificial blocking cause.
    {
      Engine warm(0);
      warm.set_stream(stream);
      for (int i = 0; i < 64; ++i) read_and_copy(warm, i, *source_a, 0, i * chunk);
      warm.commit();
      GPU(cudaStreamSynchronize(stream));
    }
    std::cerr << "REAL_GPU_WARMUP_DONE " << TEST_DALI_VERSION << std::endl;
    GPU(launch_gpu_gate(stream, mapped_device, 0));
    for (int i = 0; i < 64; ++i) read_and_copy(*engine, i, *source_a, 0, i * chunk);
    require(engine->blocks_.size() == 64 && engine->ready_buffers_.empty() &&
            engine->scheduled_.size() == 33 && engine->unscheduled_.size() == 31,
            "Cannot create expected 64-buffer legal state");
    require(cudaEventQuery(engine->ready_) == cudaErrorNotReady, "Generation 1 is not gated");
    trace("legal_full_pool_g1_scheduled33_unscheduled31");
    waiter = std::thread([&] {
      try {
        GPU(cudaSetDevice(0));
        pause_waiter_a = true;
        held.emplace(engine->get_staging_buffer());
      } catch (...) {
        std::lock_guard<std::mutex> lock(control.mutex);
        control.worker_error = std::current_exception();
        control.cv.notify_all();
      }
    });
    {
      std::unique_lock<std::mutex> lock(control.mutex);
      require(control.cv.wait_for(lock, std::chrono::seconds(10), [] { return control.before || control.worker_error; }),
              "DALI wait_buffers interposition did not execute");
      if (control.worker_error) std::rethrow_exception(control.worker_error);
    }
    mapped_host[0] = 1;
    __sync_synchronize();
    {
      std::unique_lock<std::mutex> lock(control.mutex);
      require(control.cv.wait_for(lock, std::chrono::seconds(10), [] { return control.after || control.worker_error; }),
              "Real CUDA event G1 did not synchronize");
      if (control.worker_error) std::rethrow_exception(control.worker_error);
    }
    require(cudaEventQuery(engine->ready_) == cudaSuccess, "G1 is not complete at post-sync barrier");
    GPU(launch_gpu_gate(stream, mapped_device, 1));
    bool b_enters = engine->lock_.try_lock();
    if (b_enters) engine->lock_.unlock();
    require(b_enters, "The mutex should be unlocked while A is paused");
    read_and_copy(*engine, 64, *source_b, 1, 0);
    read_and_copy(*engine, 65, *source_b, 1, chunk);
    {
      std::lock_guard<std::mutex> lock(control.mutex);
      control.release = true;
    }
    control.cv.notify_all();
    waiter.join();
    if (control.worker_error) std::rethrow_exception(control.worker_error);
    require(records.size() == 66, "Unexpected copy record count");
    std::unordered_set<void *> pending;
    for (int i = 33; i < 66; ++i) pending.insert(records[i].source);
    int unsafe_ready = 0;
    for (void *p : engine->ready_buffers_) unsafe_ready += pending.count(p);
    bool initial_lease_pending = pending.count(held->at(0));
    int ready_count = engine->ready_buffers_.size();
    int scheduled_count = engine->scheduled_.size();
    require(cudaEventQuery(engine->ready_) == cudaErrorNotReady, "G2 must still be copying/gated");
    if (TEST_RECHECK) {
      require(unsafe_ready == 0 && !initial_lease_pending, "Fixed variant still exposes pending buffers");
    }
    else require(b_enters && unsafe_ready > 0, "Original variant did not prematurely recycle buffers");
    int extra_reads = 0;
    if (!TEST_RECHECK) {
      while (!pending.count(held->at(0))) {
        source_c->read(held->at(0), chunk, 0, 2);
        engine->copy_to_client(output + (66 + extra_reads) * words, chunk, std::move(*held));
        held.reset();
        require(++extra_reads <= 31, "Unexpected ready-buffer traversal");
        held.emplace(engine->get_staging_buffer());
      }
    }
    // This is actual file-to-GPU IO, not a host-side model or cudaMemset stand-in.
    source_c->read(held->at(0), chunk, 0, 2);
    trace("different_file_overwrote_leased_gpu_buffer", held->at(0), 2, 0, chunk);
    mapped_host[1] = 1;
    __sync_synchronize();
    GPU(cudaStreamSynchronize(stream));
    std::vector<uint64_t> host(66 * words);
    GPU(cudaMemcpy(host.data(), output, host.size() * sizeof(uint64_t), cudaMemcpyDeviceToHost));
    uint64_t bad_words = 0;
    int bad_samples = 0, first_bad = -1;
    uint64_t first_value = 0;
    for (int i = 0; i < 66; ++i) {
      bool bad = false;
      for (size_t j = 0; j < words; ++j) {
        auto value = host[i * words + j];
        if (value != static_cast<uint64_t>(i + 1)) {
          ++bad_words;
          bad = true;
          if (first_bad < 0) { first_bad = i; first_value = value; }
        }
      }
      if (bad) {
        ++bad_samples;
        std::ofstream dump(prefix + ".bad_" + std::to_string(i) + ".bin", std::ios::binary);
        dump.write(reinterpret_cast<const char *>(host.data() + i * words), chunk);
      }
    }
    require(!control.watchdog_fired, "Watchdog released the GPU gates; result invalid");
    require(TEST_RECHECK ? bad_words == 0 : bad_words > 0, "Unexpected GPU data comparison result");
    Dl_info dali_location{}, cufile_location{};
    dladdr(dlsym(RTLD_DEFAULT, "_ZN4dali3gds16GDSStagingEngine18get_staging_bufferEPv"), &dali_location);
    dladdr(reinterpret_cast<void *>(&cuFileRead), &cufile_location);
    std::cout << "{\"dali_version\":\"" << TEST_DALI_VERSION << "\",\"fixed\":" << (TEST_RECHECK ? "true" : "false")
              << ",\"fix_kind\":\"" << TEST_MODE << "\""
              << ",\"actual_cuda\":true,\"actual_cufile_read\":true,\"gpu\":\"" << properties.name
              << "\",\"driver_version\":" << driver << ",\"cuda_runtime\":" << runtime << ",\"cufile_version\":" << cufile
              << ",\"chunk_bytes\":" << chunk << ",\"buffers\":64,\"commit_trigger\":33,\"reader_workers\":2"
              << ",\"wait_hook_calls\":" << control.hook_calls
              << ",\"ready_after_a\":" << ready_count << ",\"scheduled_after_a\":" << scheduled_count
              << ",\"premature_ready_buffers\":" << unsafe_ready
              << ",\"first_lease_still_pending\":" << (initial_lease_pending ? "true" : "false")
              << ",\"extra_reads\":" << extra_reads << ",\"bad_samples\":" << bad_samples
              << ",\"bad_uint64_values\":" << bad_words << ",\"first_bad_sample\":" << first_bad
              << ",\"first_wrong_value\":" << first_value << ",\"overwrite_file_value\":999"
              << ",\"get_buffer_library\":\"" << (dali_location.dli_fname ? dali_location.dli_fname : "unknown")
              << "\",\"cufile_library\":\"" << (cufile_location.dli_fname ? cufile_location.dli_fname : "unknown")
              << "\",\"watchdog_fired\":false}" << std::endl;
    exit_code = 0;
  } catch (const std::exception &error) {
    std::cerr << "REAL_GPU_TEST_ERROR " << error.what() << std::endl;
  }
  release_gpu_gates();
  {
    std::lock_guard<std::mutex> lock(control.mutex);
    control.release = true;
  }
  control.cv.notify_all();
  if (waiter.joinable()) waiter.join();
  cudaDeviceSynchronize();
  if (held && engine) { engine->return_unused(std::move(*held)); held.reset(); }
  if (engine && stream) {
    try { engine->commit(); cudaStreamSynchronize(stream); } catch (...) {}
  }
  engine.reset();
  source_c.reset(); source_b.reset(); source_a.reset();
  if (cufile_opened) cuFileDriverClose();
  if (stream) cudaStreamDestroy(stream);
  if (output) cudaFree(output);
  control.host_flags = nullptr;
  if (mapped_host) cudaFreeHost(mapped_host);
  control.finished = true;
  if (watchdog.joinable()) watchdog.join();
  if (trace_file) std::fclose(trace_file);
  return exit_code;
}
''',
'gpu_gate.cu': r'''#include <cuda_runtime.h>
#include <cstdint>

// A test-only GPU gate. The host/watchdog can always release mapped host flags
// without waiting for any CUDA stream or holding a DALI mutex.
__global__ void wait_for_host_flag(volatile uint32_t *flags, int index) {
  while (flags[index] == 0) __nanosleep(1000);
  __threadfence_system();
}

extern "C" cudaError_t launch_gpu_gate(cudaStream_t stream, uint32_t *flags, int index) {
  wait_for_host_flag<<<1, 1, 0, stream>>>(flags, index);
  return cudaGetLastError();
}
''',
'gpu_test_hooks.hpp': r'''#pragma once
#include <cuda_runtime_api.h>
#include "dali/core/cuda_error.h"
#include "dali/operators/reader/gds_mem.h"
void BeforeRealEventSynchronize();
void AfterRealEventSynchronize();
''',
'original_wait.cpp': r'''// Copyright (c) 2022-2023, 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.

// Original wait_buffers from DALI v2.1.0 (fd3f55f88b5f41f05e061d86843ddfd3da88ad83).
#include "gpu_test_hooks.hpp"
namespace dali { namespace gds {
void GDSStagingEngine::wait_buffers(std::unique_lock<std::mutex> &lock) {
  if (scheduled_.empty() && !unscheduled_.empty())
    commit_no_lock();
  lock.unlock();
  CUDA_CALL(cudaEventSynchronize(ready_));
  lock.lock();
  for (void *ptr : scheduled_)
    ready_buffers_.insert(ptr);
  scheduled_.clear();
}

}}
''',
'run_repro.py': r'''"""Build and run a controlled real-CUDA/real-cuFile GDS staging-engine repro.

Requires Linux, a compatible DALI wheel with C++ headers, a CUDA toolkit with
cuFile headers/libraries, g++, and storage readable through cuFileRead.
The installed DALI libraries are not modified.
"""

import argparse
import importlib.util
import json
import os
from pathlib import Path
import shutil
import subprocess


HERE = Path(__file__).resolve().parent


def run(command, work, label, env=None, timeout=180):
    with (work / (label + '.stdout')).open('w') as out:
        with (work / (label + '.stderr')).open('w') as err:
            result = subprocess.run(
                command, cwd=work, env=env, stdout=out, stderr=err,
                timeout=timeout,
            )
    if result.returncode:
        raise RuntimeError(
            label + ' failed:\n'
            + (work / (label + '.stderr')).read_text(errors='replace')[-9000:]
        )
    return (work / (label + '.stdout')).read_text()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--dali-dir', type=Path, help='Directory containing libdali_operators.so and include/')
    parser.add_argument('--cuda-home', type=Path, help='CUDA toolkit directory, including cuFile development files')
    parser.add_argument('--arch', default='native', help='nvcc GPU architecture; default: native')
    parser.add_argument('--mode', choices=['original', 'event-recheck', 'both'], default='both')
    parser.add_argument('--work-dir', type=Path, default=Path.cwd() / 'gds-event-race-output')
    args = parser.parse_args()

    dali = args.dali_dir
    if dali is None:
        spec = importlib.util.find_spec('nvidia.dali')
        if spec is None:
            parser.error('Install DALI or specify --dali-dir.')
        dali = Path(next(iter(spec.submodule_search_locations)))
    dali = dali.resolve()

    cuda = args.cuda_home
    if cuda is None:
        configured = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
        nvcc = shutil.which('nvcc')
        cuda = Path(configured) if configured else (Path(nvcc).resolve().parent.parent if nvcc else None)
    if cuda is None:
        parser.error('Set CUDA_HOME, add nvcc to PATH, or specify --cuda-home.')
    cuda = cuda.resolve()
    for path in [dali / 'libdali_operators.so', dali / 'include/dali/operators/reader/gds_mem.h',
                 cuda / 'bin/nvcc', cuda / 'include/cufile.h']:
        if not path.is_file():
            parser.error('Required file not found: ' + str(path))

    work = args.work_dir.resolve()
    work.mkdir(parents=True, exist_ok=True)
    data = work / 'input'
    data.mkdir(exist_ok=True)
    chunk = 2 << 20
    for name, values in [('source_a.bin', range(1, 65)),
                         ('source_b.bin', [65, 66]), ('source_c.bin', [999])]:
        with (data / name).open('wb') as handle:
            for value in values:
                handle.write(value.to_bytes(8, 'little') * (chunk // 8))

    run([str(cuda / 'bin/nvcc'), '-std=c++17', '-arch=' + args.arch,
         '-c', str(HERE / 'gpu_gate.cu'), '-o', str(work / 'gpu_gate.o'),
         '-Xcompiler', '-fPIC'], work, 'build_gate')

    original = (HERE / 'original_wait.cpp').read_text()
    recycle = '  for (void *ptr : scheduled_)\n    ready_buffers_.insert(ptr);\n  scheduled_.clear();'
    assert original.count(recycle) == 1
    modes = ['original', 'event-recheck'] if args.mode == 'both' else [args.mode]
    for mode in modes:
        recheck = mode == 'event-recheck'
        source = original
        if recheck:
            source = source.replace(
                recycle,
                '  // Another waiter may have re-recorded ready_ during the unlocked wait.\n'
                '  // Query its current state under the mutex before reclaiming buffers.\n'
                '  return_ready();',
            )
        source = source.replace(
            '  CUDA_CALL(cudaEventSynchronize(ready_));',
            '  BeforeRealEventSynchronize();\n'
            '  CUDA_CALL(cudaEventSynchronize(ready_));\n'
            '  AfterRealEventSynchronize();',
        )
        method = work / (mode + '_wait.cpp')
        method.write_text(source)
        executable = work / mode
        command = [
            'g++', '-std=c++17', '-O0', '-g', '-pthread', '-fno-access-control',
            '-DTEST_RECHECK=' + str(int(recheck)),
            '-DTEST_MODE="' + mode + '"', '-DTEST_DALI_VERSION="official_2_1"',
            '-I' + str(HERE), '-I' + str(dali / 'include'), '-I' + str(cuda / 'include'),
            str(HERE / 'gpu_engine_test.cpp'), str(method), str(work / 'gpu_gate.o'),
            '-L' + str(dali), '-L' + str(cuda / 'lib64'), '-Wl,--export-dynamic',
            '-Wl,-rpath,' + str(dali), '-Wl,-rpath,' + str(dali / '.libs'),
            '-Wl,-rpath,' + str(cuda / 'lib64'),
            '-ldali_operators', '-ldali', '-ldali_kernels', '-ldali_core',
            '-lcudart', '-lcufile', '-ldl', '-o', str(executable),
        ]
        run(command, work, 'build_' + mode)
        env = os.environ.copy()
        env.setdefault('CUDA_VISIBLE_DEVICES', '0')
        env['DALI_GDS_CHUNK_SIZE'] = '2M'
        env['LD_LIBRARY_PATH'] = ':'.join(
            [str(cuda / 'lib64'), str(dali), str(dali / '.libs')]
            + ([env['LD_LIBRARY_PATH']] if env.get('LD_LIBRARY_PATH') else [])
        )
        print('Running', mode, flush=True)
        text = run([str(executable), str(data), str(work / mode)],
                   work, mode, env=env, timeout=90)
        result = json.loads(text.strip().splitlines()[-1])
        (work / (mode + '.json')).write_text(json.dumps(result, indent=2) + '\n')
        print(json.dumps(result, indent=2), flush=True)
    print('PASS: original mode reproduced corruption; event-recheck mode requires zero corruption.'
          if args.mode == 'both' else 'PASS: expected assertions for ' + args.mode, flush=True)


if __name__ == '__main__':
    main()
''',
}

if __name__ == "__main__":
    with tempfile.TemporaryDirectory(prefix="dali-gds-repro-") as directory:
        root = Path(directory)
        for name, content in FILES.items():
            (root / name).write_text(content)
        runpy.run_path(str(root / "run_repro.py"), run_name="__main__")

Relevant log output

Both cases use 2 reader threads, one GPU copy stream, 64 buffers, a 33-copy
submission trigger, and 2,097,152-byte chunks.

Metric	Original	Event recheck
`wait_hook_calls`	2	2
`ready_after_a`	63	30
`scheduled_after_a`	0	33
`premature_ready_buffers`	32	0
`first_lease_still_pending`	true	false
`bad_samples`	1	0
`bad_uint64_values`	262144	0
`watchdog_fired`	false	false

The original run's output block at zero-based index 65 should contain 66 but
contains 999 throughout. The saved corrupt block is byte-for-byte equal to
the entire synthetic source_c.bin file (2 MiB). Their SHA-256 is:

52965276086f2f6c9028e8592114413493703c92f168204a572ec7befc62ac3b

The reproducer saves per-case JSON, stdout/stderr, a CSV trace containing
thread IDs, buffer addresses and file offsets, and corrupt blocks as
.bad_*.bin for the original run.

Other/Misc.

Reproduction environment:

Ubuntu 22.04.5 LTS, x86_64; kernel 5.15.0-127-generic; glibc 2.35.
One NVIDIA RTX 6000D, compute capability 12.0, selected for the test.
NVIDIA driver 580.95.05.
CUDA toolkit/runtime 13.0; nvcc 13.0.88.
cuFile library 1.15.1; cuFileGetVersion() returns 1150.
Python 3.12.13 and GCC/G++ 11.4.0.
Input files on NFS; opened with O_DIRECT and read using cuFileRead.
The test does not assert whether cuFile uses direct GDS or compatibility mode.
DALI_GDS_CHUNK_SIZE=2M is explicitly set by the runner.

Suggested fix: After reacquiring the mutex, query the event's current state
using the existing return_ready() helper instead of unconditionally
reclaiming scheduled_:

 void GDSStagingEngine::wait_buffers(std::unique_lock<std::mutex> &lock) {
   if (scheduled_.empty() && !unscheduled_.empty())
     commit_no_lock();
   lock.unlock();
   CUDA_CALL(cudaEventSynchronize(ready_));
   lock.lock();
-  for (void *ptr : scheduled_)
-    ready_buffers_.insert(ptr);
-  scheduled_.clear();
+  // Another waiter may have re-recorded ready_ during the unlocked wait.
+  // Query its current state under the mutex before reclaiming buffers.
+  return_ready();
 }

return_ready() already checks cudaEventQuery(ready_), leaves pending
buffers alone on cudaErrorNotReady, and propagates other CUDA errors. The
mutex keeps the event recording and scheduled-buffer set consistent during
the query and reclamation. The existing outer loop in get_staging_buffer()
can use already-free buffers or wait again if needed.

This preserves the unlocked wait. The fixed reproduction still allows B to
enter and submit G2 while A is paused, and reports no premature buffers or
corrupt output. The runtime comparison interposes the changed method; a full
DALI rebuild and end-to-end training validation are outside this test.

Could you confirm the reclamation invariant and whether this change, together
with a regression test for the interleaving above, is appropriate?

Check for duplicates

I searched the repository's issues and pull requests for wait_buffers,
GDSStagingEngine, and GDS corruption, and found no duplicate describing
this race. The GDS-related setup/import failure in Help needed with GDS support #4718 is a different
symptom and failure path.

@hhkl1989

@hhkl1989 thank you for your report. #6624 should fix that according to your suggestion.

Navigation Menu

Repository navigation

[GDS] wait_buffers can prematurely recycle in-flight buffers and corrupt NumPy reader output #6623

Version

Describe the bug.

Minimum reproducible example

Relevant log output

Other/Misc.

Check for duplicates

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[GDS] wait_buffers can prematurely recycle in-flight buffers and corrupt NumPy reader output #6623

Description

Version

Describe the bug.

Minimum reproducible example

Relevant log output

Other/Misc.

Check for duplicates

Activity

JanuszL commented on Nov 6, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions