ESPnet3 Parallel

espnet3/parallel/ is the shared provider/runner layer ESPnet3 uses to run one piece of work (one dataset item, one shard) in parallel across a local process, a local GPU, or a Dask-backed cluster (slurm, sge, kube, ...), without the stage code itself knowing which backend is active.

The pattern has two halves:

• a Provider (EnvironmentProvider) builds the runtime objects a job needs once — a dataset, a model, a device — and exposes them the same way whether execution ends up local or distributed.
• a Runner (BaseRunner) takes a provider, plans work into shards, dispatches forward() over each shard's items, and merges shard outputs into one final result.

EnvironmentProvider          BaseRunner
  build_env_local()  ---->   forward(idx, **env)
  build_worker_setup_fn()    open_writers / write_record / close_writers
                             merge(shard_dirs)

Stage code only ever constructs a provider, constructs a runner around it, and calls the runner — shard planning, resume/locking, and the local-vs-Dask dispatch all live inside BaseRunner.

Interactive guide

Components overview

ESPnet3 parallel processing is built around four components. Each has a clear responsibility. Colors are used consistently throughout this guide.

Driver-sideWorker-sideProvider / initforward()done / complete

BaseRunner

The orchestrator (Driver)

Splits data into shards, dispatches them to workers, and merges results when all workers finish. It never processes data itself.

__call__()_plan_shards()merge()

forward()

The work unit (Worker)

A pure function that processes one sample or batch. Must be a @staticmethod — no self — so Dask can serialize and send it to remote workers.

@staticmethodpickle-safe

EnvironmentProvider

The initialization blueprint

Defines how and where to build the dataset and model — locally on the driver, or once per worker process. Two methods, two different timings.

build_env_local()build_worker_setup_fn()

Cluster / Client

The execution environment

Set once with set_parallel(config). Changing env switches from running everything locally to submitting jobs to a SLURM cluster — no other code changes needed.

env: localenv: slurmenv: local_gpu

ℹ

You don't need to know Dask. ESPnet3 uses Dask internally, but you only interact with set_parallel(config). The BaseRunner handles everything else.

Build your parallel config

Select an execution environment and fill in the values. The YAML config and the generated SLURM script update in real time.

Execution environment

env execution modelocal — same process (dev / debug)local_gpu — multiple GPUs on one machineslurm — SLURM cluster

n_workers number of workers (SLURM: job count)

parallel_config.yaml● live

parallel:
  env: local
  n_workers: 4
  options: {}

ℹ

Other cluster types (PBS, SGE, LSF, HTCondor, ...): Change env to "pbs", "sge", "lsf", or "htcondor". The available options vary by scheduler — see the dask-jobqueue documentation.

Execution modes compared

A single value in parallel_config.env fundamentally changes where execution happens. Switch tabs to see the flow for each mode.

🖥 Your machine (Driver)

__call__(indices)

Plans shards, writes manifest.json

↓ _run_local()

_run_one_shard()

Processes shards sequentially (no parallelism)

↓

forward(idx, dataset, model)

Runs in the same process as the driver

↓ merge()

merge(shard_dirs)

Combines all shards into final output

📁 Filesystem

output_dir/

manifest.json
split.0/ → done
split.1/ → done
...

Local mode: No Dask, no workers. The driver processes all shards in sequence. Progress is shown with tqdm.

# local mode (config.yaml)
parallel:
  env: local
  n_workers: 1    # not used in local mode
  options: {}

EnvironmentProvider timing

The biggest source of confusion with EnvironmentProvider is that the two methods are called at completely different times and places.

build_env_local()called: on the driver, right now

1You call runner(indices)…

2_run_local() calls provider.build_env_local() directly

3The returned dict (dataset, model, ...) is used immediately

local mode flow

env = provider.build_env_local()
# dict is returned immediately
forward(idx, **env)

build_worker_setup_fn()called: on each worker at startup

1Driver calls provider.build_worker_setup_fn() → returns a function

2That function is wrapped in DictReturnWorkerPlugin and sent to workers

3When each worker starts, setup() is called — dataset/model are built on the worker

distributed mode flow

setup_fn = provider.build_worker_setup_fn()
# nothing is initialized yet
# setup_fn is a closure
plugin = DictReturnWorkerPlugin(setup_fn)
# when a worker starts:
env = setup_fn() # runs on the worker

⚠

Why return a function instead of the env directly?
SLURM workers run on different machines. If the driver initialized the model and tried to send it, the full model weights would need to be serialized and transferred over the network. Instead, only the lightweight initialization recipe (closure) is sent, and each worker builds its own copy locally — significantly reducing transfer overhead.

Automatic env injection: wrap_func_with_worker_env

When calling forward() on a worker, you don't pass dataset or model explicitly. wrap_func_with_worker_env inspects the function signature and matches argument names against the worker's env keys — injecting them automatically.

Argument names matching env keys are injected automatically

@staticmethod
def forward(idx, dataset, model, device, **env):
# dataset, model, device → injected from worker env
# idx → passed by client.map()

# Worker env (returned by setup_fn()):
env = { "dataset": ds, "model": md, "device": dev, ... }
↑ key names match forward's argument names → auto-injected

Implementation example: subclassing InferenceProvider

class EnvironmentProvider(ABC):
    """Abstract base — implement both methods."""

    @abstractmethod
    def build_env_local(self) -> dict:
        """Local execution: called on the driver immediately."""
        ...

    @abstractmethod
    def build_worker_setup_fn(self) -> Callable:
        """Distributed execution: the returned function is called on each worker."""
        ...

Shard lifecycle walkthrough

What happens internally between calling runner(range(200)) and getting the result back. Step through each stage.

① Indices received

Calling runner(range(200)) triggers BaseRunner.__call__(). All indices are converted to a list. If batch_size is set, they are grouped into batches here.

ℹ

Where is output_dir set? It is passed as a constructor argument to the Runner. All shard files, manifest.json, and final outputs are written under this directory.

runner = MyRunner(
    provider,
    output_dir="exp/decode",   # ← set here
    shard_subdir="train",       # shards go under output_dir/train/
    resume=True,               # skip already-done shards
)
# When using collect_stats(), output_dir is passed as a function
# argument and forwarded to the Runner automatically.

indices = range(200) → 200 indices

0

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

…

# First thing BaseRunner.__call__ does
indices = list(indices)
if self.batch_size is not None:
    indices = [
        list(indices[i : i + self.batch_size])
        for i in range(0, len(indices), self.batch_size)
    ]

Where this is used today

• infer() uses InferenceProvider / InferenceRunner — see Inference Provider.
• collect_stats uses CollectStatsInferenceProvider / CollectStatsRunner — see Data Preparation.
• The LibriTTS VITS recipe uses a custom provider/runner pair to extract x-vectors — see X-Vector Extraction for a full worked example.

Where to go next

Provider / Runner

The base contracts: env construction, forward, writer hooks, parallel config, resume/locking.

Inference Provider

The stage-facing provider/runner pair used by infer().

Data Preparation

Using the same pattern for dataset preparation and collect_stats.

X-Vector Extraction

A full worked example: a custom Provider and Runner for speaker embeddings.