nemo_curator.stages.audio.alm.pretrain.finalize

View as Markdown

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

NameDescription
_build_final_summary-
_collect_reconciled_output_stats-
_merge_manifest_shards-
_merge_metrics_shards-
_merge_tar_shardsMerge per-replica audio tar shards into output_path.
_patch_metrics_post_reconcileReconcile the merged metrics summary against the post-reconcile manifest.
_reconcile_manifest_with_tarDrop manifest rows whose audio path field isn’t a valid, readable
finalize_audio_pretrain_outputsMerge per-worker shards into the final manifest, metrics JSON, and audio tar.
prepare_audio_pretrain_outputsDelete any pre-existing shards from prior runs.

API

nemo_curator.stages.audio.alm.pretrain.finalize._build_final_summary(
per_original: dict[str, dict[str, typing.Any]],
durations: list[float],
filtered_examples: list[str] | None = None
) -> dict[str, typing.Any]
nemo_curator.stages.audio.alm.pretrain.finalize._collect_reconciled_output_stats(
manifest_path: str
) -> tuple[dict[str, dict[str, typing.Any]], list[float]]
nemo_curator.stages.audio.alm.pretrain.finalize._merge_manifest_shards(
output_path: str
) -> None
nemo_curator.stages.audio.alm.pretrain.finalize._merge_metrics_shards(
metrics_path: str
) -> None
nemo_curator.stages.audio.alm.pretrain.finalize._merge_tar_shards(
output_path: str
) -> None

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().

nemo_curator.stages.audio.alm.pretrain.finalize._patch_metrics_post_reconcile(
metrics_path: str,
manifest_path: str,
dropped_missing: int,
dropped_unreadable: int
) -> None

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:

  1. Increments dropped.missing_audio and dropped.corrupted_audio with the reconcile pass’s drop counts (so reconcile drops are attributable separately from worker-side filters like empty / overlap).
  2. Rebuilds the output-side counters — num_output_snippets, output_total_segments, output_total_duration_sec, snippet_duration_histogram_30s, and each per_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).

nemo_curator.stages.audio.alm.pretrain.finalize._reconcile_manifest_with_tar(
manifest_path: str,
tar_path: str,
audio_filepath_key: str = 'audio_filepath'
) -> tuple[int, int]

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.

nemo_curator.stages.audio.alm.pretrain.finalize.finalize_audio_pretrain_outputs(
output_manifest_path: str,
metrics_path: str,
output_audio_tar_path: str,
audio_filepath_key: str = 'audio_filepath'
) -> None

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:

  1. its audio_filepath_key value is not a member of the tar, or
  2. 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, …).

nemo_curator.stages.audio.alm.pretrain.finalize.prepare_audio_pretrain_outputs(
output_manifest_path: str,
metrics_path: str,
output_audio_tar_path: str
) -> None

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.