torchaudio.io.StreamReader crashes on repeated seek / next

🐛 Describe the bug

If I repeatedly run the following script, it will fail before reaching "End" most of the time. About 75% of the time I run the script. Sometimes it crashes after the first seek, sometimes after the second, and sometimes it finishes completely.

I can overcome the issue by creating a new StreamReader instance every time I seek, but presumably this workaround will slow things down.

• This issue occurs with torch '2.1.2+cu121',
• It does NOT happen with torch '2.0.1+cu117'.

import requests
import os
import sys
import torchaudio
import code

print(torchaudio.__version__)

url = "https://assets.allsamplefiles.com/mp4/ns/60s/sample-file-quad-hd.mp4"
src = "D:\\sample-file-quad-hd.mp4"

if not os.path.isfile(src):
    print(f"Downloading video file {url}")
    response = requests.get(url)
    if response.status_code == 200:
        with open(src, "wb") as file:
            file.write(response.content)
    else:
        print("FAILED")
        sys.exit(0)
    print("Done downloading video file")

assert os.path.isfile(src)

stream_params = {
    "frames_per_chunk": 2,
    "buffer_chunk_size": -1,
    "stream_index": None,
    "decoder": "h264_cuvid",
    "decoder_option": {
        "deint": "2",
        "gpu": "0",
        "surfaces": "32",
    },
    "hw_accel": "cuda:0",
}

start_times = [0.0, 1.0677333333333334, 2.135466666666667]

stream_reader = torchaudio.io.StreamReader(src=src)
print("add_video_stream")
stream_reader.add_video_stream(**stream_params)

for i, start_time in enumerate(start_times):
    print(f"Before {i}")
    stream_reader.seek(start_time)
    chunk = next(stream_reader.stream())

print("End")

Note: the warning seems unrelated.

For example, here is some output:

(py311) C:\Users\Bank\Desktop>python torch_test.py
2.1.0+cu121
Downloading video file https://assets.allsamplefiles.com/mp4/ns/60s/sample-file-quad-hd.mp4
Done downloading video file
[W conversion.cpp:412] Warning: The output format NV12 is selected. This will be implicitly converted to YUV444P, in which all the color components Y, U, V have the same dimension. (function operator ())
Before 0
Before 1

Versions

Collecting environment information...
PyTorch version: 2.1.2+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Pro
GCC version: Could not collect
Clang version: Could not collect
CMake version: version 3.26.0-rc5
Libc version: N/A

Python version: 3.11.5 | packaged by Anaconda, Inc. | (main, Sep 11 2023, 13:26:23) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19045-SP0
Is CUDA available: True
CUDA runtime version: 11.7.64
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA GeForce RTX 2070 SUPER
GPU 1: NVIDIA GeForce RTX 2070 SUPER

Nvidia driver version: 537.42
cuDNN version: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\bin\cudnn_ops_train64_8.dll
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=3696
DeviceID=CPU0
Family=179
L2CacheSize=10240
L2CacheSpeed=
Manufacturer=GenuineIntel
MaxClockSpeed=3696
Name=Intel(R) Core(TM) i9-10900X CPU @ 3.70GHz
ProcessorType=3
Revision=21767

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.1
[pip3] torch==2.1.2+cu121
[pip3] torchaudio==2.1.2+cu121
[pip3] torchvision==0.16.2+cu121
[conda] cudatoolkit 11.1.1 hb074779_12 conda-forge
[conda] numpy 1.23.5 pypi_0 pypi

Start at the torchaudio.io.StreamReader entry point and run the supplied reproduction script with repeated seek and next calls. Compare the reported behavior between the listed torch versions and inspect the seek/stream interaction, including the CUDA decoder configuration. Done means repeated seeks complete without a crash and the regression is covered by an appropriate test.