基于Ray的音频分块并行读取与顺序保持方案咨询

需要基于Ray实现大规模音频文件并行推理，核心需求如下：

• 逐块读取音频文件，避免一次性加载全文件至内存
• 每块完成推理后逐块写入输出文件
• 保证单个音频文件的分块读写顺序严格一致
• 确保写入时的块完整性，避免输出文件损坏

已实现支持全文件加载的AudioDatasource与AudioDataSink，但分块场景下的顺序保持及写操作完整性问题尚未解决，现有代码如下：

class AudioDataSink(RowBasedFileDatasink):
    def __init__(self, path: str, file_format: str, model_sample_rate: int, **file_datasink_kwargs):
        super().__init__(path, file_format=file_format, **file_datasink_kwargs)
        self.model_sr = model_sample_rate
        self.file_format = file_format

    def write_row_to_file(self, row: dict[str, Any], file: "pyarrow.NativeFile"):
        audio = row["audio"]
        if self.model_sr != row["original_sr"]:
            audio = soxr.resample(audio, self.model_sr, row["original_sr"])
        sf.write(file, audio, self.model_sr)


class AudioDatasource(FileBasedDatasource):
    _WRITE_FILE_PER_ROW = True
    _NUM_THREADS_PER_TASK = 8

    def __init__(
        self,
        paths: str | list[str],
        model_sample_rate: int,
        **file_based_datasource_kwargs,
    ):
        super().__init__(
            paths,
            file_extensions=["wav", "mp3", "flac", "ogg", "aiff"],
            **file_based_datasource_kwargs,
        )
        self.model_sr = model_sample_rate

    def _read_stream(self, f: "pyarrow.NativeFile", path: str, **reader_args) -> Iterator[Block]:
        try:
            audio, sr = sf.read(f, dtype="float32", always_2d=True)
        except RuntimeError as e:
            raise ValueError(f"Error reading {f}: {e}")

        if sr != self.model_sr:
            audio = soxr.resample(audio, sr, self.model_sr)
        audio = np.transpose(audio)
        builder = DelegatingBlockBuilder()
        item = {"audio": audio, "original_sr": sr}
        builder.add(item)
        block = builder.build()
        yield block

以下是针对分块读写的改造方案，重点解决顺序保持和文件完整性问题：

1. 分块读取的AudioDatasource改造 #

新增分块时长参数，流式读取音频文件的每个块，并为每个块标记文件唯一标识、块序号和总块数，为后续有序写入提供依据。

import numpy as np
import soundfile as sf
import soxr
from ray.data.datasource import FileBasedDatasource, DelegatingBlockBuilder
from ray.data.block import Block
from typing import Iterator, Any, Union

class AudioDatasource(FileBasedDatasource):
    _WRITE_FILE_PER_ROW = True
    _NUM_THREADS_PER_TASK = 8

    def __init__(
        self,
        paths: Union[str, list[str]],
        model_sample_rate: int,
        block_duration: float = 10.0,  # 分块时长（秒），可按需调整
        **file_based_datasource_kwargs,
    ):
        super().__init__(
            paths,
            file_extensions=["wav", "mp3", "flac", "ogg", "aiff"],
            **file_based_datasource_kwargs,
        )
        self.model_sr = model_sample_rate
        self.block_samples = int(block_duration * model_sample_rate)  # 每个块的采样点数

    def _read_stream(self, f: "pyarrow.NativeFile", path: str, **reader_args) -> Iterator[Block]:
        try:
            # 获取文件基础信息
            info = sf.info(f)
            total_frames = info.frames
            original_sr = info.samplerate

            # 计算重采样后的总采样数
            if original_sr != self.model_sr:
                resampled_total = int(total_frames * self.model_sr / original_sr)
            else:
                resampled_total = total_frames

            block_idx = 0
            # 循环分块读取
            while True:
                # 计算当前块在原文件中的起始帧位置
                if original_sr != self.model_sr:
                    start_frame = int(block_idx * self.block_samples * original_sr / self.model_sr)
                else:
                    start_frame = block_idx * self.block_samples

                # 读取当前块
                audio, sr = sf.read(f, dtype="float32", always_2d=True, start=start_frame, frames=self.block_samples)
                if len(audio) == 0:
                    break

                # 重采样到模型采样率
                if sr != self.model_sr:
                    audio = soxr.resample(audio, sr, self.model_sr)
                audio = np.transpose(audio)

                builder = DelegatingBlockBuilder()
                # 携带排序所需的元数据
                item = {
                    "audio": audio,
                    "original_sr": original_sr,
                    "file_id": path,
                    "block_idx": block_idx,
                    "total_blocks": (resampled_total + self.block_samples - 1) // self.block_samples
                }
                builder.add(item)
                yield builder.build()
                block_idx += 1
        except RuntimeError as e:
            raise ValueError(f"读取文件 {path} 出错: {e}")

2. 有序分块写入的AudioDataSink改造 #

使用Ray Actor管理单个文件的写入状态，确保同一文件的块按序号顺序写入，同时独占文件句柄避免并发写入冲突，保证文件完整性。

import ray
import soundfile as sf
import soxr
import os
from ray.data.datasource import RowBasedFileDatasink
from typing import Any, Union

# 单个文件的有序写入Actor
@ray.remote
class AudioFileWriter:
    def __init__(self, output_path: str, model_sr: int, file_format: str, total_blocks: int, channels: int = 1):
        self.output_path = output_path
        self.model_sr = model_sr
        self.file_format = file_format
        self.total_blocks = total_blocks
        self.written_blocks = set()
        self.channels = channels
        self.file_handle = self._init_file()

    def _init_file(self):
        # 初始化流式写入句柄，适配不同音频格式
        return sf.SoundFile(
            self.output_path,
            mode='w',
            samplerate=self.model_sr,
            channels=self.channels,
            format=self.file_format
        )

    def write_block(self, audio_block: np.ndarray, block_idx: int):
        # 避免重复写入
        if block_idx in self.written_blocks:
            return
        # 检查块顺序，确保按序号写入
        expected_idx = len(self.written_blocks)
        if block_idx != expected_idx:
            raise ValueError(f"文件 {self.output_path} 期望写入块 {expected_idx}，收到块 {block_idx}")
        # 写入音频块（转置回soundfile要求的格式）
        self.file_handle.write(audio_block.T)
        self.written_blocks.add(block_idx)
        # 所有块写入完成后关闭文件
        if len(self.written_blocks) == self.total_blocks:
            self.file_handle.close()

class AudioDataSink(RowBasedFileDatasink):
    def __init__(self, path: str, file_format: str, model_sample_rate: int, **file_datasink_kwargs):
        super().__init__(path, file_format=file_format, **file_datasink_kwargs)
        self.model_sr = model_sample_rate
        self.file_format = file_format
        self.writers = {}  # 缓存每个文件的写入Actor

    def write_row_to_file(self, row: dict[str, Any], file: "pyarrow.NativeFile"):
        file_id = row["file_id"]
        output_filename = os.path.basename(file_id)
        output_path = os.path.join(self.path, output_filename)
        block_idx = row["block_idx"]
        total_blocks = row["total_blocks"]
        audio = row["audio"]

        # 按需重采样
        if self.model_sr != row["original_sr"]:
            audio = soxr.resample(audio, self.model_sr, row["original_sr"])

        # 获取或创建当前文件的写入Actor
        if output_path not in self.writers:
            # 自动识别音频通道数
            channels = audio.shape[0]
            self.writers[output_path] = AudioFileWriter.remote(
                output_path, self.model_sr, self.file_format, total_blocks, channels
            )
        writer = self.writers[output_path]

        # 调用Actor写入块，确保顺序执行
        ray.get(writer.write_block.remote(audio, block_idx))

关键说明 #

• 顺序保证：每个块携带block_idx序号，Actor内部检查当前块是否为预期的下一个块，强制按顺序写入，避免Ray任务并行导致的乱序。
• 文件完整性：通过Actor独占文件句柄，使用SoundFile的流式写入API，避免多线程/进程并发写入冲突；所有块写入完成后统一关闭文件，防止文件损坏。
• 内存优化：分块读取每次仅加载一个块的音频数据，内存占用与块大小正相关，适合处理GB级超大音频文件。
• 并行效率：Ray原生支持多文件并行处理，不同文件的写入Actor独立运行，不互相干扰，最大化利用集群资源。

内容的提问来源于stack exchange，提问作者Luke