nemo_curator.stages.audio.alm.pretrain.finalize
nemo_curator.stages.audio.alm.pretrain.finalize
Driver-side prepare/finalize helpers and shard mergers.
Pipeline writers (manifest, metrics, audio tar) all emit one shard file
per replica. These helpers run on the driver around pipeline.run()
to clean up stale shards, merge fresh shards into the user-visible
output files, and reconcile manifest rows against the final tar.
Module Contents
Functions
API
Merge per-replica audio tar shards into output_path.
Reads every <output_path>.shard-*.tar written by the extractor
workers, copies their members into a single fresh tar at
output_path in lexicographic member-name order (matches
Energon expectations for indexed tar datasets), and removes the
shards. Re-write via Python tarfile instead of byte-level
concatenation so the merger is portable, handles padding/header
boundaries correctly, and tolerates shards left without trailing
zero-blocks by workers that were ray.kill-ed before
teardown().
Reconcile the merged metrics summary against the post-reconcile manifest.
Worker shards are written before _reconcile_manifest_with_tar can
prune manifest rows whose audio is missing or unreadable, so the
initial _merge_metrics_shards summary overcounts on the output
side by exactly the number of rows the reconcile pass removed. This
helper:
- Increments
dropped.missing_audioanddropped.corrupted_audiowith the reconcile pass’s drop counts (so reconcile drops are attributable separately from worker-side filters likeempty/overlap). - Rebuilds the output-side counters —
num_output_snippets,output_total_segments,output_total_duration_sec,snippet_duration_histogram_30s, and eachper_original[*].out_*field — by walking the now-authoritative (post-reconcile) manifest. Input-side and worker-side dropped counters are left untouched.
No-op when both reconcile counts are zero (the merged metrics already match the manifest) or when the metrics file doesn’t exist (dry run / reconcile-in-isolation tests).
Drop manifest rows whose audio path field isn’t a valid, readable tar member with a positive duration.
Returns (dropped_missing, dropped_unreadable) so the caller can
surface the counts (e.g. patch them into the merged metrics summary).
A no-op when the tar file doesn’t exist (dry-run, or all tar shards
were empty). See finalize_audio_pretrain_outputs for why this
is needed even when the pipeline reports success.
The tar itself is left as-is. Removing dropped members would
require rewriting the whole archive; downstream consumers iterate
the manifest and look up by audio_filepath_key, so orphan members
are harmless.
Merge per-worker shards into the final manifest, metrics JSON, and audio tar.
Call once on the driver, AFTER pipeline.run() returns
successfully. Reads all manifest + metrics + tar shards written by
the writer / aggregator / extractor stages, concatenates / combines
them, writes the final user-facing files at the user-provided paths,
and removes the shards.
After the audio tar is built, reconciles the manifest against the tar. A manifest row is dropped if either:
- its
audio_filepath_keyvalue is not a member of the tar, or - the corresponding tar member’s audio header is unreadable or reports zero frames / zero sample rate.
Check (1) guards against the Xenna failure mode where a worker is
ray.kill-ed between writing a JSONL line for a snippet and
flushing the snippet’s audio bytes to its tar shard. Check (2)
guards against truncated members (worker killed mid-payload-write
with a header but no body) and corrupted writes that would surface
to downstream consumers (WebDataset / Energon) as a runtime decode
error instead of a clean filter.
Reconcile drops are surfaced in the merged metrics under
dropped.missing_audio (check 1) and dropped.corrupted_audio
(check 2), so the user can attribute post-pipeline integrity drops
separately from worker-side filters (empty, overlap, …).
Delete any pre-existing shards from prior runs.
Call this once on the driver, BEFORE pipeline.run(). Multi-worker
backends would race on cleanup if we did it inside a stage’s
setup(), so we keep cleanup driver-only.