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/advanced_pipelines/audio_data_filter/audio_data_filter.py b/nemo_curator/stages/audio/advanced_pipelines/audio_data_filter/audio_data_filter.py index 0f17f9e085..3b280b155d 100644 --- a/nemo_curator/stages/audio/advanced_pipelines/audio_data_filter/audio_data_filter.py +++ b/nemo_curator/stages/audio/advanced_pipelines/audio_data_filter/audio_data_filter.py @@ -46,6 +46,7 @@ from loguru import logger +from nemo_curator.stages.audio._agent_ready import AgentReady, StageContract from nemo_curator.stages.audio.filtering import BandFilterStage, SIGMOSFilterStage, UTMOSFilterStage from nemo_curator.stages.audio.postprocessing import TimestampMapperStage from nemo_curator.stages.audio.preprocessing import MonoConversionStage, SegmentConcatenationStage @@ -56,7 +57,7 @@ from .config import _deep_merge, get_enabled_stages, load_config -class AudioDataFilterStage(CompositeStage): +class AudioDataFilterStage(AgentReady, CompositeStage): """Complete audio data filtering and curation pipeline (CompositeStage). Decomposes into independent stages that the executor can schedule with @@ -90,6 +91,9 @@ def __init__( if config: self._cfg = _deep_merge(self._cfg, config) + def describe(self) -> StageContract: + return StageContract(wrappable=False) + def decompose(self) -> list[ProcessingStage]: """Build a self-consistent pipeline topology based on enabled features.""" cfg = self._cfg 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")