diff --git a/nemo_curator/stages/audio/_agent_ready.py b/nemo_curator/stages/audio/_agent_ready.py new file mode 100644 index 0000000000..efeeddc6be --- /dev/null +++ b/nemo_curator/stages/audio/_agent_ready.py @@ -0,0 +1,318 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field, is_dataclass +from enum import Enum +from typing import TYPE_CHECKING, Any, ClassVar, Literal + +if TYPE_CHECKING: + from collections.abc import Mapping + +AudioForm = Literal["file", "waveform"] +ProducedForm = Literal["tensor", "disk"] +Cardinality = Literal["1:1", "1:1 nested-list", "1:N fan-out", "N:1", "filter"] +Dispatch = Literal["process", "process_batch", "auto"] +# How a stage handles per-item failures at runtime. "unknown" means the stage +# has not declared a uniform policy (the default for most stages today). +ErrorPolicy = Literal["skip", "fail", "annotate", "unknown"] + +# Semantic role vocabulary. Because key *names* are agent-configurable +# (config-knobs-only standardization: every read/written key is a ``*_key`` +# constructor field that an agent may rename), two stages can only be chained +# reliably by matching the *role* a key plays, not its literal string. Roles are +# resolved from the invariant ``*_key`` field name (see ``_roles.py``), so they +# survive an agent renaming a key's value. ``"unknown"`` is the safe default and +# never blocks composition. +Role = Literal[ + "audio_filepath", + "waveform", + "sample_rate", + "duration", + "num_samples", + "segments", + "diar_segments", + "vad_segments", + "overlap_segments", + "timestamps", + "start", + "end", + "start_ms", + "end_ms", + "text", + "pred_text", + "reference_text", + "words", + "alignment", + "speaker_id", + "num_speakers", + "score", + "metrics", + "prediction", + "segment_num", + "original_file", + "item_id", + "windows", + "manifest_path", + "output_path", + "unknown", +] + + +@dataclass(frozen=True) +class ParamSpec: + """A single constructor parameter an agent can set on a stage. + + Usually derived automatically from the stage's dataclass fields (or + ``__init__`` signature) via + :func:`nemo_curator.stages.audio._agent_registry.stage_params`, but a stage + may also override/augment entries in ``StageContract.params``. + """ + + name: str + type: str = "Any" + default: Any = None + required: bool = False + choices: list[Any] | None = None # populated for Literal[...] params + description: str | None = None + role: Role | None = None # semantic role of the key this param configures + + +@dataclass(frozen=True) +class IOSpec: + """Task data keys and audio forms read or written by a stage.""" + + data_keys: list[str] = field(default_factory=list) + segment_data_keys: list[str] = field(default_factory=list) + accepts: list[AudioForm] = field(default_factory=list) + produces: list[ProducedForm] = field(default_factory=list) + + +@dataclass(frozen=True) +class Gates: + """Execution gates or side effects an agent should know before wrapping a stage.""" + + writes_to_disk: bool = False + requires_gpu: bool = False + requires_internet_first_run: bool = False + requires_ffmpeg: bool = False + lifecycle_side_effects: bool = False + runtime_secrets: list[str] = field(default_factory=list) + # Serializability contract for sinks (distinguishes the two JSON sinks): + # requires_serializable_input — this stage serializes task.data as-is (e.g. + # raw json.dumps) and will fail on a resident tensor/audio blob. + # sanitizes_output — this stage strips tensors/audio blobs, so anything + # downstream of it is serialization-safe. + requires_serializable_input: bool = False + sanitizes_output: bool = False + + +@dataclass(frozen=True) +class SizeEnvelope: + """Coarse size and memory hints for agent planning.""" + + max_input_sec: float | None = None + allowed_sample_rates: list[int] | None = None + channels: Literal["mono", "stereo", "any"] = "any" + memory_hint: str | None = None + + +@dataclass(frozen=True) +class StaticHints: + """Instance-independent hints a stage may declare for discovery/planning. + + Lets a stage expose ``cardinality_options``, ``gates``, ``dispatch``, + ``error_policy``, ``description`` and ``stage_id`` *without* being + instantiated (see :meth:`AgentReady.describe_static`). Declared on a class + via the ``AGENT_STATIC`` ClassVar; every field defaults so declaring it is + fully optional and additive. + """ + + cardinality_options: list[str] = field(default_factory=list) + gates: Gates = field(default_factory=Gates) + dispatch: Dispatch = "auto" + error_policy: ErrorPolicy = "unknown" + description: str | None = None + stage_id: str | None = None + + +@dataclass(frozen=True) +class StageContract: + """Read-only discovery contract for an agent-ready processing stage.""" + + reads: IOSpec = field(default_factory=IOSpec) + writes: IOSpec = field(default_factory=IOSpec) + reads_one_of: list[IOSpec] = field(default_factory=list) + metadata_reads: list[str] = field(default_factory=list) + metadata_writes: list[str] = field(default_factory=list) + cardinality: Cardinality = "1:1" + cardinality_options: list[str] = field(default_factory=list) + iteration_key: str | None = None + preserves_upstream_keys: bool = True + wrappable: bool = True + size_envelope: SizeEnvelope = field(default_factory=SizeEnvelope) + gates: Gates = field(default_factory=Gates) + # Agent-facing metadata (advisory; defaults keep older describe() calls valid). + stage_id: str | None = None # stable semantic id; defaults to the class name when None + description: str | None = None # one-line human summary for planners/UIs + params: list[ParamSpec] = field(default_factory=list) # usually auto-derived at discovery time + dispatch: Dispatch = "auto" # "auto" => infer from the stage at runtime + error_policy: ErrorPolicy = "unknown" + # Resolved-key-value -> semantic role. Populated at discovery time by + # ``_agent_registry.build_contract``; empty when a contract is built by hand. + key_roles: dict[str, Role] = field(default_factory=dict) + # True when the stage only implements ``process_batch`` (``process`` raises). + # Auto-derived at discovery time; agents must not call ``process`` on these. + batch_only: bool = False + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-safe dict of this contract (``json.dumps`` never raises).""" + return { + "reads": asdict(self.reads), + "writes": asdict(self.writes), + "reads_one_of": [asdict(spec) for spec in self.reads_one_of], + "metadata_reads": list(self.metadata_reads), + "metadata_writes": list(self.metadata_writes), + "cardinality": self.cardinality, + "cardinality_options": list(self.cardinality_options), + "iteration_key": self.iteration_key, + "preserves_upstream_keys": self.preserves_upstream_keys, + "wrappable": self.wrappable, + "size_envelope": asdict(self.size_envelope), + "gates": asdict(self.gates), + "stage_id": self.stage_id, + "description": self.description, + "params": [_paramspec_to_dict(p) for p in self.params], + "dispatch": self.dispatch, + "error_policy": self.error_policy, + "key_roles": dict(self.key_roles), + "batch_only": self.batch_only, + } + + +def _jsonable_default(value: Any) -> Any: # noqa: ANN401, PLR0911 (complexity accepted: one early return per JSON type family) + """Coerce an arbitrary param default to a JSON-serializable value. + + Non-serializable objects (tensors, models, callables) become a short + ``""`` sentinel rather than raising. + """ + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, (list, tuple, set)): + return [_jsonable_default(v) for v in value] + if isinstance(value, dict): + return {str(k): _jsonable_default(v) for k, v in value.items()} + if isinstance(value, Enum): + return value.value + if is_dataclass(value) and not isinstance(value, type): + try: + return {k: _jsonable_default(v) for k, v in asdict(value).items()} + except Exception: # noqa: BLE001 + return f"" + return f"" + + +def _paramspec_to_dict(p: ParamSpec) -> dict[str, Any]: + return { + "name": p.name, + "type": p.type, + "default": _jsonable_default(p.default), + "required": p.required, + "choices": None if p.choices is None else [_jsonable_default(c) for c in p.choices], + "description": p.description, + "role": p.role, + } + + +def _json_type(type_str: str | None) -> str | None: + """Map a rendered ParamSpec.type string to a JSON-Schema type (or None to omit).""" + if not type_str: + return None + t = type_str.replace(" ", "").replace("|None", "") + if t.startswith("Optional["): + t = t[len("Optional[") : -1] if t.endswith("]") else t + scalar = {"str": "string", "int": "integer", "float": "number", "bool": "boolean"} + if t in scalar: + return scalar[t] + lowered = t.lower() + if lowered.startswith(("list", "sequence", "tuple")): + return "array" + if lowered.startswith(("dict", "mapping")): + return "object" + return None + + +def to_json_schema(params: list[ParamSpec]) -> dict[str, Any]: + """Build a JSON-Schema ``object`` describing a stage's configurable params. + + Suitable as the argument schema for an agent's stage-configuration form. + """ + properties: dict[str, Any] = {} + required: list[str] = [] + for p in params: + schema: dict[str, Any] = {} + json_type = _json_type(p.type) + if json_type is not None: + schema["type"] = json_type + if p.choices is not None: + schema["enum"] = [_jsonable_default(c) for c in p.choices] + if p.default is not None: + schema["default"] = _jsonable_default(p.default) + if p.description: + schema["description"] = p.description + if p.role: + schema["x-role"] = p.role + properties[p.name] = schema + if p.required: + required.append(p.name) + out: dict[str, Any] = {"type": "object", "properties": properties} + if required: + out["required"] = required + return out + + +class AgentReady: + """Mixin for stages that expose a read-only agent discovery contract.""" + + # Opt-in, instance-independent discovery hints. Annotated as ClassVar so + # dataclass stages do NOT treat these as fields. All optional/additive. + AGENT_STATIC: ClassVar[StaticHints | None] = None + # Set True on stages whose ``process`` raises (only ``process_batch`` works). + BATCH_ONLY: ClassVar[bool] = False + # Rare per-field role overrides keyed by ``*_key`` field name. Consulted + # before the shared ``_roles.KEY_ROLES`` table. + KEY_ROLE_OVERRIDES: ClassVar[Mapping[str, Role]] = {} + + def describe(self) -> StageContract: + raise NotImplementedError + + @classmethod + def describe_static(cls) -> StageContract: + """Instance-free contract for discovery/planning. + + Uses class defaults + ``AGENT_STATIC`` and never runs ``__init__`` side + effects, so it is safe for stages with required constructor args. Prefer + the instance-level :meth:`describe` (or + ``_agent_registry.build_contract``) when resolved key *values* are + needed. + """ + from nemo_curator.stages.audio._agent_registry import static_contract + + return static_contract(cls) + + +# (resolve_contract was removed: dead code whose signature promised class +# acceptance the body rejected. Instances: use stage.describe() / build_contract; +# classes: use describe_static() / static_contract.) diff --git a/nemo_curator/stages/audio/_residency.py b/nemo_curator/stages/audio/_residency.py new file mode 100644 index 0000000000..135768b3d2 --- /dev/null +++ b/nemo_curator/stages/audio/_residency.py @@ -0,0 +1,157 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import contextlib +import os +import tempfile +from typing import TYPE_CHECKING, Any, Literal + +import soundfile as sf + +from nemo_curator.stages.audio.common import ensure_waveform_2d, load_audio_file + +if TYPE_CHECKING: + from collections.abc import Callable + +InputResidency = Literal["file", "waveform", "auto"] + + +def resolve_audio( # noqa: PLR0913 (complexity accepted: keyword-only residency/key knobs mirror the stage fields) + item: dict[str, Any], + *, + residency: InputResidency = "auto", + audio_filepath_key: str = "audio_filepath", + waveform_key: str = "waveform", + sample_rate_key: str = "sample_rate", + mono: bool = True, + loader: Callable[..., tuple[Any, int]] | None = None, +) -> tuple[Any, int] | None: + """Return ``(waveform_2d, sample_rate)`` from tensor keys or a file path. + + ``auto`` prefers an existing waveform, then falls back to file loading. + ``waveform`` never falls back to disk. ``file`` always loads from the + configured path key. + + ``loader`` overrides the file-loading callable (default + :func:`~nemo_curator.stages.audio.common.load_audio_file`); stages pass + their own module-level symbol so callers can patch it at the stage module. + """ + waveform = item.get(waveform_key) + sample_rate = item.get(sample_rate_key) + if residency != "file" and waveform is not None and sample_rate is not None: + return ensure_waveform_2d(waveform), int(sample_rate) + + if residency == "waveform": + return None + + path = item.get(audio_filepath_key) + if path: + expanded = os.path.expanduser(str(path)) + if os.path.exists(expanded): + return (loader or load_audio_file)(expanded, mono=mono) + return None + + +def _as_soundfile_array(waveform: Any) -> Any: # noqa: ANN401 + waveform = ensure_waveform_2d(waveform) + if hasattr(waveform, "detach"): + waveform = waveform.detach() + if hasattr(waveform, "cpu"): + waveform = waveform.cpu() + if hasattr(waveform, "numpy"): + waveform = waveform.numpy() + if getattr(waveform, "ndim", 0) == 2: # noqa: PLR2004 - 2 == a (channels, samples) 2-D array + channels, samples = waveform.shape + if channels == 1: + return waveform[0] + if channels < samples: + return waveform.T + return waveform + + +def resolve_audio_path( # noqa: PLR0913 (complexity accepted: keyword-only residency/key knobs mirror the stage fields) + item: dict[str, Any], + *, + residency: InputResidency = "auto", + audio_filepath_key: str = "audio_filepath", + waveform_key: str = "waveform", + sample_rate_key: str = "sample_rate", + temp_dir: str | None = None, + register_temp: list[str] | None = None, +) -> str | None: + """Return an audio path, writing a temp WAV when only a waveform exists. + + When a temp WAV is materialized from an in-memory waveform and + ``register_temp`` is provided, the temp path is appended to that list so the + caller can delete it after use (see :func:`cleanup_temp_files`). Without + ``register_temp`` the caller is responsible for cleanup itself. + """ + path = item.get(audio_filepath_key) + local_path: str | None = None + if residency != "waveform" and path: + local_path = os.path.expanduser(str(path)) + if os.path.exists(local_path): + return local_path + # Protocol-prefixed paths (file://, http(s)://, s3://, ...) were handled + # by the stages' own fsspec machinery before the residency layer existed; + # keep accepting them when the target exists remotely. + if "://" in str(path): + try: + from fsspec.core import url_to_fs + + fs, fspath = url_to_fs(str(path)) + if fs.exists(fspath): + return path + except Exception: # noqa: BLE001, S110 - unknown protocol/creds -> deliberate fall-through + pass + + if residency == "file": + # Pre-residency stages handed unverified paths straight to their own + # downstream machinery (ffmpeg/NeMo/fsspec) and let it report the + # failure; keep that contract instead of gating on os.path.exists. + return local_path + + waveform = item.get(waveform_key) + sample_rate = item.get(sample_rate_key) + if waveform is None or sample_rate is None: + return local_path + + fd, tmp = tempfile.mkstemp(suffix=".wav", dir=temp_dir) + os.close(fd) + sf.write(tmp, _as_soundfile_array(waveform), int(sample_rate)) + if register_temp is not None: + register_temp.append(tmp) + return tmp + + +def cleanup_temp_files(paths: list[str] | None) -> None: + """Best-effort removal of temp files created by :func:`resolve_audio_path`.""" + for path in paths or (): + with contextlib.suppress(OSError): + os.remove(path) + + +def produce_audio_filepath( + item: dict[str, Any], + new_path: str, + *, + key: str = "audio_filepath", + original_key: str = "original_audio_filepath", +) -> None: + """Update a canonical audio path while preserving the first prior value.""" + if key in item and original_key not in item: + item[original_key] = item[key] + item[key] = new_path diff --git a/nemo_curator/stages/audio/common.py b/nemo_curator/stages/audio/common.py index a27c1f54bd..4cb80af249 100644 --- a/nemo_curator/stages/audio/common.py +++ b/nemo_curator/stages/audio/common.py @@ -25,6 +25,7 @@ from loguru import logger from nemo_curator.backends.base import NodeInfo, WorkerMetadata +from nemo_curator.stages.audio._agent_ready import AgentReady, Gates, IOSpec, StageContract from nemo_curator.stages.base import CompositeStage, ProcessingStage from nemo_curator.stages.file_partitioning import FilePartitioningStage from nemo_curator.tasks import AudioTask, EmptyTask, FileGroupTask @@ -41,7 +42,7 @@ def get_audio_duration(audio_filepath: str) -> float: @dataclass -class GetAudioDurationStage(ProcessingStage[AudioTask, AudioTask]): +class GetAudioDurationStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """Compute audio duration from the file at *audio_filepath_key* and store the result under *duration_key*. @@ -65,6 +66,12 @@ def inputs(self) -> tuple[list[str], list[str]]: def outputs(self) -> tuple[list[str], list[str]]: return [], [self.duration_key] + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"]), + writes=IOSpec(data_keys=[self.duration_key]), + ) + def process(self, task: AudioTask) -> AudioTask: t0 = time.perf_counter() audio_filepath = task.data[self.audio_filepath_key] @@ -74,7 +81,7 @@ def process(self, task: AudioTask) -> AudioTask: return task -class PreserveByValueStage(ProcessingStage[AudioTask, AudioTask]): +class PreserveByValueStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """Filter entries by comparing *input_value_key* against *target_value*. Returns ``None`` from ``process()`` to drop entries that fail the @@ -87,6 +94,7 @@ class PreserveByValueStage(ProcessingStage[AudioTask, AudioTask]): """ name: str = "PreserveByValueStage" + BATCH_ONLY = True # process() raises; only process_batch is implemented (agent-discovery hint) def __init__( self, @@ -108,6 +116,13 @@ def inputs(self) -> tuple[list[str], list[str]]: def outputs(self) -> tuple[list[str], list[str]]: return [], [self.input_value_key] + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec(data_keys=[self.input_value_key]), + writes=IOSpec(data_keys=[self.input_value_key]), + cardinality="filter", + ) + def process(self, task: AudioTask) -> AudioTask | None: msg = "PreserveByValueStage only supports process_batch" raise NotImplementedError(msg) @@ -133,7 +148,7 @@ def process_batch(self, tasks: list[AudioTask]) -> list[AudioTask]: @dataclass -class ManifestReaderStage(ProcessingStage[FileGroupTask, AudioTask]): +class ManifestReaderStage(AgentReady, ProcessingStage[FileGroupTask, AudioTask]): """Read JSONL manifest files from a FileGroupTask and emit one AudioTask per line. Uses line-by-line streaming via fsspec (no Pandas) to keep memory at ~1x file size. @@ -177,9 +192,16 @@ def ray_stage_spec(self) -> dict[str, Any]: def num_workers(self) -> int | None: return 1 + def describe(self) -> StageContract: + return StageContract( + writes=IOSpec(data_keys=["audio_filepath"]), + cardinality="1:N fan-out", + gates=Gates(lifecycle_side_effects=True), + ) + @dataclass -class ManifestReader(CompositeStage[EmptyTask, AudioTask]): +class ManifestReader(AgentReady, CompositeStage[EmptyTask, AudioTask]): """Composite stage for reading JSONL manifests. Decomposes into: @@ -227,9 +249,12 @@ def get_description(self) -> str: parts.append(f"with target blocksize {self.blocksize}") return ", ".join(parts) + def describe(self) -> StageContract: + return StageContract(cardinality="1:N fan-out", wrappable=False) + @dataclass -class ManifestWriterStage(ProcessingStage[AudioTask, AudioTask]): +class ManifestWriterStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """Append a single AudioTask to a JSONL manifest file. The output file is truncated once in ``setup()`` (called on the driver) @@ -290,6 +315,18 @@ def process(self, task: AudioTask) -> AudioTask: def num_workers(self) -> int | None: return 1 + def describe(self) -> StageContract: + return StageContract( + gates=Gates( + writes_to_disk=True, + lifecycle_side_effects=True, + # Serializes task.data as-is via json.dumps; a resident tensor + # (e.g. a waveform) will crash it. Route through + # AudioToDocumentStage (which sanitizes) if one may be present. + requires_serializable_input=True, + ), + ) + def load_audio_file(audio_path: str, mono: bool = True) -> tuple[torch.Tensor, int]: """Load audio file and return waveform tensor (channels, samples) and sample rate.""" @@ -327,6 +364,13 @@ def resolve_waveform_from_item( item['audio_filepath'], resolves missing sample_rate from file header. Updates item in-place when loading from file. Returns None if resolution fails. + + .. note:: + The canonical resolver is :func:`nemo_curator.stages.audio._residency.resolve_audio`. + This helper is retained for its unique behavior — reading ``sample_rate`` from the + file header *without* reloading an already-present waveform, and writing the loaded + waveform/sample_rate back into ``item`` — which ``resolve_audio`` does not replicate. + Prefer ``resolve_audio`` in new code. """ waveform = item.get("waveform") sample_rate = item.get("sample_rate") diff --git a/nemo_curator/stages/audio/preprocessing/mono_conversion.py b/nemo_curator/stages/audio/preprocessing/mono_conversion.py index ed60661235..4aecb27982 100755 --- a/nemo_curator/stages/audio/preprocessing/mono_conversion.py +++ b/nemo_curator/stages/audio/preprocessing/mono_conversion.py @@ -27,19 +27,23 @@ """ import os +import tempfile from dataclasses import dataclass, field +import soundfile as sf import torch from loguru import logger -from nemo_curator.stages.audio.common import load_audio_file +from nemo_curator.stages.audio._agent_ready import AgentReady, Gates, IOSpec, StageContract +from nemo_curator.stages.audio._residency import produce_audio_filepath, resolve_audio +from nemo_curator.stages.audio.common import ensure_waveform_2d, load_audio_file from nemo_curator.stages.base import ProcessingStage from nemo_curator.stages.resources import Resources from nemo_curator.tasks import AudioTask @dataclass -class MonoConversionStage(ProcessingStage[AudioTask, AudioTask]): +class MonoConversionStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """ Audio mono conversion and sample rate verification stage. @@ -50,11 +54,43 @@ class MonoConversionStage(ProcessingStage[AudioTask, AudioTask]): output_sample_rate: Expected sample rate in Hz (default: 48000) audio_filepath_key: Key in data dict for audio file path strict_sample_rate: If True, reject audio with wrong sample rate + waveform_key: Key in data dict for the in-memory mono waveform tensor. + sample_rate_key: Key in data dict for the waveform sample rate. + is_mono_key: Key where the mono flag is written. + duration_key: Key where the audio duration in seconds is written. + num_samples_key: Key where the number of samples is written. + output_audio_filepath_key: Key where the written mono WAV path is stored + (write_to_disk=True only). + original_audio_filepath_key: Key preserving the pre-conversion path when + update_audio_filepath=True. + input_residency: Which input to use — "waveform" (in-memory only), "file" + (audio_filepath only), or "auto" (waveform first, file fallback; default). + keep_waveform_in_task: If True (default), store the mono waveform and sample + rate in task.data for downstream in-memory consumers. + write_to_disk: If True, write the converted mono audio to a WAV file. + write_to_disk without output_dir writes WAV files to the system temp dir + and nothing cleans them up; in multi-node runs point output_dir at a + shared filesystem. + update_audio_filepath: If True (with write_to_disk), repoint audio_filepath_key + at the written mono WAV and keep the old path under original_audio_filepath_key. + output_dir: Directory for the written WAV files (default: system temp dir). """ output_sample_rate: int = 48000 audio_filepath_key: str = "audio_filepath" + waveform_key: str = "waveform" + sample_rate_key: str = "sample_rate" + is_mono_key: str = "is_mono" + duration_key: str = "duration" + num_samples_key: str = "num_samples" + output_audio_filepath_key: str = "mono_audio_filepath" + original_audio_filepath_key: str = "original_audio_filepath" strict_sample_rate: bool = True + input_residency: str = "auto" + keep_waveform_in_task: bool = True + write_to_disk: bool = False + update_audio_filepath: bool = False + output_dir: str | None = None name: str = "MonoConversion" batch_size: int = 1 @@ -67,33 +103,91 @@ def inputs(self) -> tuple[list[str], list[str]]: return [], [] def outputs(self) -> tuple[list[str], list[str]]: - return [], ["waveform", "sample_rate", "is_mono", "duration", "num_samples"] - - def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: + outputs = [ + self.waveform_key, + self.sample_rate_key, + self.is_mono_key, + self.duration_key, + self.num_samples_key, + ] + if self.write_to_disk: + outputs.append(self.output_audio_filepath_key) + if self.update_audio_filepath: + outputs.append(self.audio_filepath_key) + return [], outputs + + def describe(self) -> StageContract: + produces = [] + if self.keep_waveform_in_task: + produces.append("tensor") + if self.write_to_disk: + produces.append("disk") + writes = [ + self.is_mono_key, + self.duration_key, + self.num_samples_key, + ] + if self.keep_waveform_in_task: + writes.extend([self.waveform_key, self.sample_rate_key]) + if self.write_to_disk: + writes.append(self.output_audio_filepath_key) + if self.update_audio_filepath: + writes.append(self.audio_filepath_key) + return StageContract( + reads=IOSpec(data_keys=[self.audio_filepath_key], accepts=["file", "waveform"]), + writes=IOSpec(data_keys=writes, produces=produces), + gates=Gates(writes_to_disk=self.write_to_disk), + ) + + def _write_audio(self, waveform: torch.Tensor, sample_rate: int, task: AudioTask) -> str: + output_dir = self.output_dir or tempfile.gettempdir() + os.makedirs(output_dir, exist_ok=True) + stem = os.path.splitext(os.path.basename(str(task.data.get(self.audio_filepath_key, "audio"))))[0] + fd, path = tempfile.mkstemp(prefix=f"{stem}_mono_", suffix=".wav", dir=output_dir) + os.close(fd) + audio = waveform.detach().cpu() + arr = audio[0].numpy() if audio.shape[0] == 1 else audio.T.numpy() + sf.write(path, arr, sample_rate) + return path + + def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: # noqa: C901 (complexity accepted: residency/sample-rate branch matrix; no refactor pre-PR) """ Convert audio to mono and verify sample rate. Mutates task.data in-place with waveform data. Returns task if successful, [] if doesn't meet requirements. """ - audio_filepath = task.data.get(self.audio_filepath_key) - - if not audio_filepath or not os.path.exists(audio_filepath): - logger.error(f"Audio file not found: {audio_filepath}") + try: + resolved = resolve_audio( + task.data, + residency=self.input_residency, # type: ignore[arg-type] + audio_filepath_key=self.audio_filepath_key, + waveform_key=self.waveform_key, + sample_rate_key=self.sample_rate_key, + mono=False, + loader=load_audio_file, # module-level symbol: patchable at this module, as pre-residency + ) + except (OSError, RuntimeError) as e: # corrupt/unreadable audio -> skip the row, don't crash the batch + logger.error(f"Failed to load audio for {task.data.get(self.audio_filepath_key)!r}: {e}") + return [] + if resolved is None: + logger.error(f"Audio input not found for key {self.audio_filepath_key!r}") return [] try: - waveform, sample_rate = load_audio_file(audio_filepath, mono=False) + waveform, sample_rate = resolved + waveform = ensure_waveform_2d(waveform) if sample_rate <= 0: - logger.error(f"Invalid sample rate ({sample_rate}) in {audio_filepath}") + logger.error(f"Invalid sample rate ({sample_rate}) in audio input") return [] num_channels = waveform.shape[0] if self.strict_sample_rate and sample_rate != self.output_sample_rate: + audio_source = task.data.get(self.audio_filepath_key, self.waveform_key) logger.warning( - f"Sample rate {sample_rate}Hz != expected {self.output_sample_rate}Hz: {audio_filepath}" + f"Sample rate {sample_rate}Hz != expected {self.output_sample_rate}Hz: {audio_source}" ) return [] @@ -103,14 +197,26 @@ def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: else: mono_waveform = waveform - task.data["waveform"] = mono_waveform - task.data["sample_rate"] = sample_rate - task.data["is_mono"] = True - task.data["duration"] = mono_waveform.shape[1] / sample_rate - task.data["num_samples"] = mono_waveform.shape[1] + if self.keep_waveform_in_task: + task.data[self.waveform_key] = mono_waveform + task.data[self.sample_rate_key] = sample_rate + task.data[self.is_mono_key] = True + task.data[self.duration_key] = mono_waveform.shape[1] / sample_rate + task.data[self.num_samples_key] = mono_waveform.shape[1] + + if self.write_to_disk: + path = self._write_audio(mono_waveform, sample_rate, task) + task.data[self.output_audio_filepath_key] = path + if self.update_audio_filepath: + produce_audio_filepath( + task.data, + path, + key=self.audio_filepath_key, + original_key=self.original_audio_filepath_key, + ) except (OSError, RuntimeError) as e: - logger.error(f"Error processing {audio_filepath}: {e}") + logger.error(f"Error processing audio input: {e}") return [] else: return task diff --git a/tests/stages/audio/test_agent_ready_mono_slice.py b/tests/stages/audio/test_agent_ready_mono_slice.py new file mode 100644 index 0000000000..c64476b01a --- /dev/null +++ b/tests/stages/audio/test_agent_ready_mono_slice.py @@ -0,0 +1,139 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Canonical agent-ready vertical slice: ``MonoConversionStage``. + +Shows that the two foundation modules work together end to end: + +* ``_agent_ready.py`` — the stage inherits ``AgentReady`` and declares a + ``StageContract`` via ``describe()`` (reads/writes, accepted audio forms, + gates, configurable ``*_key`` fields). +* ``_residency.py`` — ``process()`` resolves audio from either an in-memory + waveform or a file, and materialized temp files are cleaned up. + +``choose_input_form`` below is a *minimal consumer* that reads only the +declared contract to decide file- vs waveform-based execution. It stands in +for the registry/orchestrator (which lands in a later PR); the equivalent +production path is ``_agent_registry.build_contract(stage)`` + +``_planning.validate_pipeline`` reading the same ``StageContract``. + +Pseudocode for the real boundary:: + + contract = build_contract(stage) # describe() + auto params/roles/dispatch + # build_contract also fills registry-derived fields not set by describe(): + # contract.params (from the constructor signature) + # contract.key_roles (from the *_key field names) + # contract.dispatch (per-item vs batched, from process/process_batch) + if "waveform" in contract.reads.accepts and item.has(stage.waveform_key): + run_on_waveform(...) + elif "file" in contract.reads.accepts and item.has(stage.audio_filepath_key): + run_on_file(...) +""" + +from pathlib import Path + +import numpy as np +import soundfile as sf +import torch + +from nemo_curator.stages.audio._residency import cleanup_temp_files, resolve_audio_path +from nemo_curator.stages.audio.preprocessing.mono_conversion import MonoConversionStage +from nemo_curator.tasks import AudioTask + +SR = 48000 + + +def _write_stereo_wav(path: Path, seconds: float = 0.1) -> None: + """Write a small real stereo WAV so the file path actually resolves.""" + data = np.random.randn(int(SR * seconds), 2).astype("float32") + sf.write(path, data, SR) + + +def choose_input_form(stage: MonoConversionStage, item: dict) -> str: + """Minimal contract consumer: pick execution form from the declared contract. + + Uses ONLY ``stage.describe()`` (the agent-facing contract) plus the item's + keys — no knowledge of the stage internals. + """ + contract = stage.describe() + accepts = set(contract.reads.accepts or []) + if "waveform" in accepts and item.get(stage.waveform_key) is not None: + return "waveform" + if "file" in accepts and item.get(stage.audio_filepath_key): + return "file" + return "none" + + +class TestMonoConversionAgentSlice: + def test_contract_declares_both_forms(self) -> None: + contract = MonoConversionStage().describe() + # Accepted audio residencies are declared, so a consumer can branch. + assert set(contract.reads.accepts) == {"file", "waveform"} + # Configurable output keys are surfaced on the contract. + assert "is_mono" in contract.writes.data_keys + + def test_dispatch_is_per_item(self) -> None: + # Dispatch: MonoConversion runs per item via process() (it is not + # batch-only). The concrete ``dispatch`` field on StageContract is + # filled by the registry (build_contract / static_contract) in the + # orchestrator PR; here we assert the property it is derived from. + assert getattr(MonoConversionStage, "BATCH_ONLY", False) is False + + def test_consumer_picks_file_vs_waveform(self, tmp_path: Path) -> None: + stage = MonoConversionStage(output_sample_rate=SR) + wav = tmp_path / "a.wav" + _write_stereo_wav(wav) + assert choose_input_form(stage, {"audio_filepath": wav.as_posix()}) == "file" + assert choose_input_form(stage, {"waveform": torch.randn(2, SR), "sample_rate": SR}) == "waveform" + + def test_runs_on_file_input(self, tmp_path: Path) -> None: + wav = tmp_path / "stereo.wav" + _write_stereo_wav(wav) + stage = MonoConversionStage(output_sample_rate=SR) + out = stage.process(AudioTask(data={"audio_filepath": wav.as_posix()})) + assert isinstance(out, AudioTask) + assert out.data["is_mono"] is True + assert out.data["waveform"].shape[0] == 1 + assert out.data["sample_rate"] == SR + + def test_runs_on_waveform_input(self) -> None: + stage = MonoConversionStage(output_sample_rate=SR) + out = stage.process(AudioTask(data={"waveform": torch.randn(2, SR), "sample_rate": SR})) + assert isinstance(out, AudioTask) + assert out.data["is_mono"] is True + assert out.data["waveform"].shape[0] == 1 + + def test_both_forms_give_equivalent_shape(self, tmp_path: Path) -> None: + stage = MonoConversionStage(output_sample_rate=SR) + wav = tmp_path / "s.wav" + _write_stereo_wav(wav, seconds=0.1) + from_file = stage.process(AudioTask(data={"audio_filepath": wav.as_posix()})) + from_wave = stage.process(AudioTask(data={"waveform": torch.randn(2, int(SR * 0.1)), "sample_rate": SR})) + assert from_file.data["waveform"].shape[0] == from_wave.data["waveform"].shape[0] == 1 + assert from_file.data["is_mono"] == from_wave.data["is_mono"] is True + + def test_waveform_temp_file_is_created_then_cleaned_up(self) -> None: + # resolve_audio_path materializes a temp WAV from an in-memory waveform; + # cleanup_temp_files removes it (the temp-file lifecycle used by stages + # that need a real path, e.g. resample/ASR/diarization). + temp_paths: list[str] = [] + item = {"waveform": np.random.randn(SR).astype("float32"), "sample_rate": SR} + path = resolve_audio_path(item, residency="waveform", register_temp=temp_paths) + assert path is not None + assert path.endswith(".wav") + assert Path(path).exists() + assert temp_paths == [path] + + cleanup_temp_files(temp_paths) + assert not Path(path).exists()