Loading audio files from archives

🐛 Describe the bug

I've been playing around with torchdata as a replacement for the webdataset library. My main use-case is reading data from network-attached file systems (such as ceph), which implies streaming from e.g. .tar files, which is something webdataset is designed for.

In the following code I have the following relative file system:
data.zip

├── file
│   ├── 19-198-0000.flac
│   └── 19-198-0000.wav
├── tar
│   ├── flac.tar
│   └── wav.tar
└── zip
    ├── flac.zip
    └── wav.zip

Where each .zip or .tar archive contains respectively the 19-198-0000.flac or 19-198-0000.wav file taken from the LibriSpeech dataset.

From my reading of the documentation, this seams the easiest way to read from the archive:

import torchaudio.backend.sox_io_backend as tab

from torchdata.datapipes.iter import (
    FileLister,
    FileOpener,
    TarArchiveLoader,
    ZipArchiveLoader,
    Mapper,
)


def audio_stream_to_tensor(element):
    path, stream = element

    audio_tensor, sample_rate = tab.load(stream)

    return audio_tensor

dp = FileLister(".", masks=["wav.tar"], recursive=True)
dp = FileOpener(dp, mode="b")
dp = TarArchiveLoader(dp, mode="r")
dp = Mapper(dp, audio_stream_to_tensor)

for x in dp:
    print(x) # tensor([[0.0044, 0.0033, 0.0031,  ..., 0.0047, 0.0060, 0.0060]])

This works :)! However, it fails when we try to read the flac.tar

dp = FileLister(".", masks=["flac.tar"], recursive=True)
dp = FileOpener(dp, mode="b")
dp = TarArchiveLoader(dp, mode="r")
dp = Mapper(dp, audio_stream_to_tensor)

for x in dp:
    print(x)

formats: can't open input file `': FLAC ERROR whilst decoding metadata
Traceback (most recent call last):
  File "/home/nik/phd/repo/librispeech/playground/example.py", line 35, in <module>
    for x in dp:
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torch/utils/data/datapipes/_typing.py", line 514, in wrap_generator
    response = gen.send(None)
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torch/utils/data/datapipes/iter/callable.py", line 116, in __iter__
    yield self._apply_fn(data)
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torch/utils/data/datapipes/iter/callable.py", line 81, in _apply_fn
    return self.fn(data)
  File "/home/nik/phd/repo/librispeech/playground/example.py", line 15, in audio_stream_to_tensor
    audio_tensor, sample_rate = tab.load(stream)
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torchaudio/backend/sox_io_backend.py", line 220, in load
    return _fallback_load_fileobj(filepath, frame_offset, num_frames, normalize, channels_first, format)
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torchaudio/io/_compat.py", line 109, in load_audio_fileobj
    s = torchaudio._torchaudio_ffmpeg.StreamReaderFileObj(src, format, None, 4096)
RuntimeError: Failed to open the input "StreamWrapper<<ExFileObject name='./tar/flac.tar'>>" (Invalid data found when processing input).
This exception is thrown by __iter__ of MapperIterDataPipe(datapipe=TarArchiveLoaderIterDataPipe, fn=audio_stream_to_tensor, input_col=None, output_col=None)

Similarly for ZipArchiveLoader, reading from wav.zip works, while flac.zip returns a similar error:

formats: can't open input file `': FLAC ERROR whilst decoding metadata
Traceback (most recent call last):
  File "/home/nik/phd/repo/librispeech/playground/example.py", line 35, in <module>
    for x in dp:
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torch/utils/data/datapipes/_typing.py", line 514, in wrap_generator
    response = gen.send(None)
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torch/utils/data/datapipes/iter/callable.py", line 116, in __iter__
    yield self._apply_fn(data)
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torch/utils/data/datapipes/iter/callable.py", line 81, in _apply_fn
    return self.fn(data)
  File "/home/nik/phd/repo/librispeech/playground/example.py", line 15, in audio_stream_to_tensor
    audio_tensor, sample_rate = tab.load(stream)
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torchaudio/backend/sox_io_backend.py", line 220, in load
    return _fallback_load_fileobj(filepath, frame_offset, num_frames, normalize, channels_first, format)
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torchaudio/io/_compat.py", line 109, in load_audio_fileobj
    s = torchaudio._torchaudio_ffmpeg.StreamReaderFileObj(src, format, None, 4096)
RuntimeError: Failed to open the input "StreamWrapper<<zipfile.ZipExtFile name='19-198-0000.flac' mode='r' compress_type=deflate>>" (Invalid data found when processing input).
This exception is thrown by __iter__ of MapperIterDataPipe(datapipe=ZipArchiveLoaderIterDataPipe, fn=audio_stream_to_tensor, input_col=None, output_col=None)

Moreover, adding torchaudio.info to the map function also leads to the same issue for .wav files:

def audio_stream_to_tensor_and_meta(element):
    path, stream = element

    meta = tab.info(stream)
    print(meta)
    audio_tensor, sample_rate = tab.load(stream)

    return audio_tensor, meta

dp = FileLister(".", masks=["wav.tar"], recursive=True)
dp = FileOpener(dp, mode="b")
dp = TarArchiveLoader(dp, mode="r")
dp = Mapper(dp, audio_stream_to_tensor_and_meta)

for x in dp:
    print(x)

AudioMetaData(sample_rate=16000, num_frames=31440, num_channels=1, bits_per_sample=16, encoding=PCM_S)
formats: can't determine type of file `'
Traceback (most recent call last):
  File "/home/nik/phd/repo/librispeech/playground/example.py", line 36, in <module>
    for x in dp:
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torch/utils/data/datapipes/_typing.py", line 514, in wrap_generator
    response = gen.send(None)
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torch/utils/data/datapipes/iter/callable.py", line 116, in __iter__
    yield self._apply_fn(data)
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torch/utils/data/datapipes/iter/callable.py", line 81, in _apply_fn
    return self.fn(data)
  File "/home/nik/phd/repo/librispeech/playground/example.py", line 25, in audio_stream_to_tensor_and_meta
    audio_tensor, sample_rate = tab.load(stream)
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torchaudio/backend/sox_io_backend.py", line 220, in load
    return _fallback_load_fileobj(filepath, frame_offset, num_frames, normalize, channels_first, format)
  File "/home/nik/phd/repo/librispeech/.venv/lib/python3.10/site-packages/torchaudio/io/_compat.py", line 109, in load_audio_fileobj
    s = torchaudio._torchaudio_ffmpeg.StreamReaderFileObj(src, format, None, 4096)
RuntimeError: Failed to open the input "StreamWrapper<<ExFileObject name='./tar/wav.tar'>>" (Invalid data found when processing input).
This exception is thrown by __iter__ of MapperIterDataPipe(datapipe=TarArchiveLoaderIterDataPipe, fn=audio_stream_to_tensor_and_meta, input_col=None, output_col=None)

So I assume that the issues stem from the fact that the stream provided by torchdata is not seekable, or at least the buffer is not large enough?

Versions

PyTorch version: 1.12.1+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.10.4 (main, Apr 20 2022, 11:26:44) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-46-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.5.119
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3070
Nvidia driver version: 495.29.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.2
[pip3] torch==1.12.1
[pip3] torchaudio==0.12.1
[pip3] torchdata==0.4.1
[conda] Could not collect

Start with TarArchiveLoader and ZipArchiveLoader, then trace the torchaudio backend load and info calls used by the Mapper functions. Reproduce the examples for WAV and FLAC streams, including the metadata call; done means archive-contained audio can be decoded and inspected through these pipelines.