From 5e760a5094ca5fc609070561d3c91a6bcb8c00ba Mon Sep 17 00:00:00 2001 From: Shubham Bhawsar Date: Tue, 25 Aug 2026 10:56:09 +0000 Subject: [PATCH 01/25] =?UTF-8?q?feat(audio):=20agent-ready=20foundation?= =?UTF-8?q?=20=E2=80=94=20contract=20layer,=20residency,=20discovery?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the stage-side declaration layer that agent-driven pipeline construction needs: the `AgentReady` mixin and `StageContract` vocabulary, the input-residency resolver, semantic role names, and the discovery/planning/conformance code that reads them. Public entry point is `nemo_curator.stages.audio.agent`; the implementation lives in the private `_agent/` subpackage. Stages declare what they read and write instead of a planner inferring it, so a pipeline can be validated before it runs rather than failing partway through a GPU pass. Ships `common.py` and the full `preprocessing/` folder as the first worked slice, including two new stages: ChannelCountStage (record/filter/convert channel count) and SampleRateFilterStage (select on rate without resampling). `_agent/__init__.py` is deliberately import-free: `common` imports `_agent_ready` at module level while `_residency` imports `common`, and that cycle survives only because reaching `_agent._agent_ready` runs a package __init__ that pulls nothing. Signed-off-by: Shubham Bhawsar --- nemo_curator/stages/audio/AGENTS.md | 63 ++ nemo_curator/stages/audio/AGENT_READY.md | 314 +++++++ nemo_curator/stages/audio/CLAUDE.md | 1 + nemo_curator/stages/audio/__init__.py | 8 + nemo_curator/stages/audio/_agent/__init__.py | 29 + .../stages/audio/_agent/_agent_ready.py | 413 +++++++++ .../stages/audio/_agent/_agent_registry.py | 569 ++++++++++++ nemo_curator/stages/audio/_agent/_catalog.py | 265 ++++++ .../stages/audio/_agent/_composite.py | 192 ++++ .../stages/audio/_agent/_conformance.py | 442 +++++++++ nemo_curator/stages/audio/_agent/_planning.py | 686 ++++++++++++++ .../stages/audio/_agent/_residency.py | 334 +++++++ nemo_curator/stages/audio/_agent/_roles.py | 207 +++++ nemo_curator/stages/audio/agent.py | 105 +++ nemo_curator/stages/audio/common.py | 864 +++++++++++++++++- .../stages/audio/preprocessing/__init__.py | 27 +- .../audio/preprocessing/channel_count.py | 519 +++++++++++ .../audio/preprocessing/concatenation.py | 117 ++- .../audio/preprocessing/mono_conversion.py | 167 +++- .../audio/preprocessing/sample_rate_filter.py | 202 ++++ tests/stages/audio/_agent/__init__.py | 0 .../_agent/test_agent_conformance_examples.py | 113 +++ .../test_channel_and_rate_stages.py | 522 +++++++++++ .../audio/preprocessing/test_concatenation.py | 40 + .../preprocessing/test_mono_conversion.py | 57 ++ tests/stages/audio/test_common.py | 519 ++++++++++- .../test_create_manifest_audio_folder.py | 87 ++ 27 files changed, 6801 insertions(+), 61 deletions(-) create mode 100644 nemo_curator/stages/audio/AGENTS.md create mode 100644 nemo_curator/stages/audio/AGENT_READY.md create mode 100644 nemo_curator/stages/audio/CLAUDE.md create mode 100644 nemo_curator/stages/audio/_agent/__init__.py create mode 100644 nemo_curator/stages/audio/_agent/_agent_ready.py create mode 100644 nemo_curator/stages/audio/_agent/_agent_registry.py create mode 100644 nemo_curator/stages/audio/_agent/_catalog.py create mode 100644 nemo_curator/stages/audio/_agent/_composite.py create mode 100644 nemo_curator/stages/audio/_agent/_conformance.py create mode 100644 nemo_curator/stages/audio/_agent/_planning.py create mode 100644 nemo_curator/stages/audio/_agent/_residency.py create mode 100644 nemo_curator/stages/audio/_agent/_roles.py create mode 100644 nemo_curator/stages/audio/agent.py create mode 100644 nemo_curator/stages/audio/preprocessing/channel_count.py create mode 100644 nemo_curator/stages/audio/preprocessing/sample_rate_filter.py create mode 100644 tests/stages/audio/_agent/__init__.py create mode 100644 tests/stages/audio/_agent/test_agent_conformance_examples.py create mode 100644 tests/stages/audio/preprocessing/test_channel_and_rate_stages.py create mode 100644 tests/stages/audio/test_create_manifest_audio_folder.py diff --git a/nemo_curator/stages/audio/AGENTS.md b/nemo_curator/stages/audio/AGENTS.md new file mode 100644 index 0000000000..d7d796a3b1 --- /dev/null +++ b/nemo_curator/stages/audio/AGENTS.md @@ -0,0 +1,63 @@ +# Audio stages — agent guardrails + +These instructions apply to any AI coding agent working in `nemo_curator/stages/audio/`. +Codex and Cursor load this nested `AGENTS.md` automatically; Claude Code reaches it +through the sibling `CLAUDE.md` import. + +Two different jobs happen in this directory, and they have opposite rules about editing. + +## If you are curating a dataset, this directory is read-only + +Stage source is shared library code. Never change a stage, a threshold, a filter or +windowing logic to make one user's dataset produce the output they hoped for — diagnose +from the data instead. The full guardrail, and what to do when output comes back empty, +is in `nemo_curator/audio_agent/AGENTS.md`. Drive the work through the +`nemo_curator.audio_agent` verbs rather than reading stage source to answer what a stage +reads or writes: `describe --params '{...}'` is the sanctioned answer, and +`producers ` says who wrote a key. + +## If you are authoring or fixing a stage, start here + +Read `AGENT_READY.md` in this directory. It is the authoritative checklist and it is +maintained with the framework; work from it rather than from memory. The +`audio-stage-authoring` skill (`nemo_curator/audio_agent/skills/audio-stage-authoring/SKILL.md`) +is the procedure that drives it. + +**Golden rule: every new knob defaults to today's behavior.** Agent-readiness is a +declaration layer over working code. If a change alters what an existing pipeline +produces, it is a behavior change and needs its own justification, not a checklist entry. + +The mechanical contract is three things, each detailed in `AGENT_READY.md`: + +1. Inherit `AgentReady` and implement `describe()` returning a `StageContract` with + `reads`, `writes`, `cardinality` and honest `gates`. +2. Make every `task.data` key you read or write a `*_key` constructor field — no bare key + literals in `process()`, or the key is invisible to the agent and cannot be remapped. +3. Add `assert_agent_ready(MyStage(...), fixture_factory=...)` as a test. + +Then give the stage a capability card under +`nemo_curator/audio_agent/knowledge/cards/`, documenting what each externally consumed +output *means* — roles prove two stages can connect, only the card lets the host judge +whether connecting them serves the user's intent. Schema: +`nemo_curator/audio_agent/knowledge/CARD_SCHEMA.md`. `card_conformance.audit()` must come +back with zero violations. + +Follow the repo's existing stage conventions while you do it: +`.cursor/rules/processing-stage-patterns.mdc` and +`.cursor/rules/composite-stage-patterns.mdc`. Those are upstream-maintained framework +rules — read them, never edit them. + +Declare honestly even where it costs you: a stage that may drop rows is +`cardinality="filter"` even when dropping is incidental, and `gates` are environment facts +rather than aspirations. An optimistic contract is worse than a missing one, because the +planner treats it as ground truth. + +## Verify before you claim done + +```bash +.venv/bin/python -m pytest tests/stages/audio -m "not gpu" -q +.venv/bin/python -m nemo_curator.audio_agent describe MyStage --params '{...}' +``` + +If `describe` does not match what the code actually touches for those params, the contract +is wrong no matter what the tests say. diff --git a/nemo_curator/stages/audio/AGENT_READY.md b/nemo_curator/stages/audio/AGENT_READY.md new file mode 100644 index 0000000000..e79d0eee4c --- /dev/null +++ b/nemo_curator/stages/audio/AGENT_READY.md @@ -0,0 +1,314 @@ +# Making an Audio Stage Agent-Ready + +This is the **short** checklist for stage owners. "Agent-ready" means an LLM agent can +**discover** your stage, **configure** it, and **chain** it with others to build a pipeline — +without reading your source. The design goal is **minimal burden**: you declare a small core, +the framework auto-derives the rest, and one test tells you if anything is missing. + +> **Golden rule:** every new knob defaults to today's behavior. Agent-readiness must not change +> how your stage runs in existing pipelines. + +--- + +## TL;DR — the mechanical contract is 3 things + +1. Inherit `AgentReady` and implement **`describe()`** returning a `StageContract` with + **`reads`, `writes`, `cardinality`** (+ honest **`gates`**). +2. Make every `task.data` key you read/write a **`*_key` constructor field** (no bare key + literals in `process()`). +3. Add one test: **`assert_agent_ready(MyStage(...), fixture_factory=...)`**. + +Those three items make the stage mechanically composable. Also update its capability card +with the meaning of externally consumed outputs (especially filterable fields and anything +crossing a fan-out/aggregation boundary). `assert_agent_ready` can prove keys and +cardinality; it cannot prove that a reasoner will interpret a value correctly. + +--- + +## What you MUST declare (only you know these) + +```python +from dataclasses import dataclass +from nemo_curator.stages.audio._agent._agent_ready import ( + AgentReady, + ConditionalWrite, + Gates, + IOSpec, + StageContract, +) +from nemo_curator.stages.base import ProcessingStage + +@dataclass +class MyStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): + """One-line summary becomes the agent-facing description. + + Args: + audio_filepath_key: Path to the input audio. # docstring Args -> param descriptions + score_key: Where the score is written. + """ + audio_filepath_key: str = "audio_filepath" + score_key: str = "my_score" + name: str = "MyStage" + + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"]), + writes=IOSpec(data_keys=[self.score_key]), + cardinality="1:1", # see "Cardinality" below + gates=Gates(requires_gpu=self.resources.requires_gpu), + ) +``` + +- **`reads` / `writes`** — the `task.data` keys consumed/produced. Use `segment_data_keys` for keys + written *inside* `segments[]` items. Use `reads_one_of=[IOSpec(...), ...]` if input can take more + than one shape (e.g. waveform **or** file). +- **`cardinality`** — one of `"1:1"`, `"1:1 nested-list"`, `"1:N fan-out"`, `"N:1"`, `"filter"`. + (`"filter"` = `process` may return `None`/`[]` to drop items.) +- **`gates`** — be honest about side effects: `requires_gpu`, `writes_to_disk`, + `requires_internet_first_run`, `requires_ffmpeg`. Serializability (both exist on `Gates`): a sink + that `json.dumps` `task.data` as-is must set `requires_serializable_input=True`; a converter that + strips tensors/audio blobs sets `sanitizes_output=True`. A stage that derives durable filenames, + member names, or row identity from framework `task.task_id` sets + `requires_stable_task_id=True`; metadata-manifest resume boundaries cannot restore that identity + and will reject the suffix. + +### Conditional/data-dependent writes + +Keep `writes` as the existing mechanical output declaration. When a listed +field is only written on a runtime branch—or a non-preserving stage may copy an +upstream field only when it is present—add factual possibility metadata: + +```python +return StageContract( + writes=IOSpec(data_keys=[self.score_key], segment_data_keys=[self.score_key]), + conditional_writes=[ + ConditionalWrite( + writes=IOSpec(data_keys=[self.score_key]), + condition=f"'{self.segments_key}' is absent, so the task-level branch runs", + ), + ConditionalWrite( + writes=IOSpec(segment_data_keys=[self.score_key]), + condition=f"'{self.segments_key}' is present, so the per-segment branch runs", + ), + ConditionalWrite( + writes=IOSpec(data_keys=list(self.passthrough_keys)), + condition="the same input key is present, non-null, and allowed by the output whitelist", + value_origin="upstream_same_key", + ), + ConditionalWrite( + metadata_writes=[self.metadata_key], + condition="the metadata value is computed and attached to the emitted task", + ), + ], +) +``` + +`ConditionalWrite` does not execute a predicate and does not change processing, +defaults, or the validator's legacy mechanical interpretation of `writes`. It +labels output possibility and value provenance for `validate.semantic_review`. +The `upstream_same_key` origin is important for allowlist/rebuild stages: it +keeps the original producer's meaning visible instead of falsely presenting the +copier as a new metric producer. + +Use `metadata_writes` for a conditional `task._metadata` output. Unconditional +metadata inputs/outputs remain declared through +`StageContract.metadata_reads`/`metadata_writes`; semantic review traces all +three scopes (`task`, `segment`, and `metadata`) independently. + +Use `augments_upstream_same_key` when the stage adds entries to an existing +mapping (for example a shared metrics mapping), and +`transforms_upstream_same_key` when it replaces the same key with a derived +value. These are objective lineage facts, not claims about whether the value is +appropriate for the user's goal. + +Conditions must describe actual code branches using configured key values. They +are not an intent checker, a module-specific validator, or a centralized +field-scope ontology. The host LLM still decides whether a conditional field's +meaning and granularity fit the request. + +### `gates.per_row_independent` — usually nothing to do + +This one decides whether a *delta run* (reprocessing only the files that changed, instead of the +whole corpus) may include your stage. **Most stages declare nothing and are handled +automatically.** + +`delta.region()` refuses to assume anything about a stage that could see a row other than the one +it was handed. It checks three things, and if **none** of them is true your stage is treated as +independent with no declaration from you: + +- you override `process_batch` (you are handed several rows at once) +- `gates.writes_to_disk=True` +- `gates.lifecycle_side_effects=True` + +**If one of those IS true**, the delta refuses your stage by name until you answer this question: + +> If I ran this stage over file `X` alone, versus over `X` plus 999 others, would `X`'s output row +> be identical? + +- **Yes** → `per_row_independent=True`, with a comment saying *why* it survives batching. +- **No** → `per_row_independent=False`. Nothing is lost: the delta simply stops at your stage, and + every stage above it still reprocesses only the changed files. + +The reference pair — both batch, opposite answers: + +| Stage | | Why | +|---|---|---| +| `ASRStage` with `NeMoASRAdapter` | `True` | prepares each waveform independently and passes the batch to NeMo with lengths preserved | +| `TorchSquimQualityMetricsStage` | `False` | pads to the batch max with **no lengths**, so padding reads as silence and scores move | + +Typical reasons for `False`: a corpus statistic or percentile threshold, a running counter that +picks output names, appending to a file shared across rows, an unseeded RNG advanced per row, +batch padding without lengths. + +Three more rules: + +- **Declare per instance when the unsafety is conditional.** `SplitLongAudioStage` is independent + only while no shared `output_dir` flattens every source's splits into one namespace, so it + declares `per_row_independent=(self.output_dir is None)` rather than a flat `False` that would + cost the default configuration — the safe one — its delta. `SplitASRAlignJoinStage` and + `InferenceSortformerStage` do the same with their own output directories. +- **A source accepting `include_files` MUST declare**, `True` or `False` — silence is a conformance + error. That parameter is how a delta narrows a source, so it has to be answerable. +- **Getting it wrong is asymmetric.** `False` when you were safe costs a full rerun — annoying and + harmless. `True` when you were not silently produces rows a full run would never have produced, + and republishes them as the corpus's reusable result. **When unsure, declare `False`.** + +The two `CreateInitialManifest*` sources with `max_samples` are a **deliberate exception** to that +last rule, not an example of it. `max_samples` truncates the *sorted* listing, so a delta over a +bounded source can select files a full run would not have — yet both declare a flat `True`, because +the conditional `False` denied reuse to the configuration nearly everyone runs (ReadSpeech defaults +to 5000). The limitation is recorded at each declaration. Do not copy this into a new stage; if you +find yourself wanting to, declare `False` and raise it instead. + +## What is AUTO-DERIVED — do NOT hand-write these + +| Field | Derived from | +|---|---| +| `params` (names, types, defaults, `choices` from `Literal[...]`) | your dataclass fields / `__init__` | +| param `description`s | your class docstring `Args:` section | +| key `role`s | your `*_key` field names (shared `_roles.KEY_ROLES`) | +| `dispatch` | whether you override `process_batch` | +| `description`, `stage_id` | class docstring / class name | + +You never put `params` in `describe()`. + +## What is OPTIONAL — set only if it's obvious + +Declared via one class attribute, `AGENT_STATIC = StaticHints(...)`, or on the contract: + +- **StaticHints-settable** (instance-free): `cardinality_options` (e.g. `["fan_out", "nested"]`), + `gates`, `dispatch`, `error_policy` (`"skip" | "fail" | "annotate"` — default `"unknown"`; set + only if your stage has a clear, uniform policy), `description`, `stage_id`. +- **Contract-only** (return them from `describe()`; StaticHints has no such fields): + `iteration_key`, `size_envelope`. +- `BATCH_ONLY = True` — only if your `process()` raises and just `process_batch` works. + +If you're unsure, leave them. The agent treats missing optionals safely. + +--- + +## Naming rule: config-knobs-only + +Keep **today's default key names**. Do **not** rename keys to a global vocabulary. Just expose a +`*_key` field for each so an agent can remap when wiring two stages. Compatibility comes from +**semantic roles** (below), not from everyone using the same strings. + +## Semantic roles — the compatibility contract + +An agent chains a producer's output to a consumer's input by **role**, not key string. Roles are +resolved automatically from your `*_key` **field name** via `nemo_curator/stages/audio/_agent/_roles.py`. + +- If your key fields use existing names (`audio_filepath_key`, `waveform_key`, `score_key`, + `text_key`, `segments_key`, …) you get the right role for free. +- If you add a **brand-new** `*_key` concept, add one line to `KEY_ROLES` in `_roles.py` (or, for a + truly stage-internal key, list it in `INTERNAL_KEY_FIELDS`). The conformance test fails if you + forget — it won't let a key silently fall through. + +## Output meaning — the reasoning contract + +Roles answer “can these stages connect?” They do not answer “what does this +value mean here?” A pipeline can connect perfectly and still apply a valid +filter to the wrong entity. Put that semantic knowledge in the capability card, +where the host LLM can reason over it; do not add a per-field rule or a hard +`field_scope` ontology to the Python core. + +For each output that another stage may select, aggregate, compare or filter, +document: + +- **meaning** — what the value actually represents, not merely “the key where it + is written”; +- **unit/range** — seconds, Hz, speakers, MOS range, category vocabulary, etc.; +- **provenance** — which input/entity and computation produced it; +- **scope/granularity** — file, original parent item, emitted child, segment, + speaker, batch or corpus, written as factual prose rather than an enum; +- **propagation** — whether fan-out copies a parent value to every child, nesting + moves it into segments, aggregation summarizes many rows, or a transform + recomputes it; +- **counterexample** — at least one plausible but wrong interpretation and its + pipeline consequence. + +Use the card's optional `semantic_facts` mapping for structured prose, or +`notes`/`caveats` when the fact spans several outputs. For example, if a +speaker-separation stage computes the original clip's detected-speaker count +once and copies it to every per-speaker child, say so explicitly: filtering that +child field to `== 1` selects children whose **parent source** had one detected +speaker; it does not test whether each already-separated child track is +single-speaker. + +Only document facts grounded in code, a measured run or an authoritative model +source, and mark their honesty tier in the card's `verified` block. Missing +meaning should remain an explicit TODO; the host must ask rather than invent it. +See `nemo_curator/audio_agent/knowledge/CARD_SCHEMA.md`. + +## Discovery — how the agent finds your stage + +Nothing to do: your stage auto-registers (via `StageMeta`) and appears in the catalog. Consumers +go through the public entry point — `agent.py` is the sanctioned public surface; don't import the +private `_catalog` module directly: +```python +from nemo_curator.stages.audio import agent + +agent.list_agent_ready_stages() # -> [... "MyStage" ...] +agent.describe_stage("MyStage") # -> StageContract (static, instance-free) +agent.catalog_as_json() # -> JSON the agent/UI consumes +``` + +--- + +## The safety net: `assert_agent_ready` + +Add one test. It runs the static checks (contract shape, valid roles, JSON-serializable, reads +satisfiable by role) and — with a fixture — runs your stage and verifies declared writes appear, +no undeclared top-level keys leak, and cardinality matches runtime: + +```python +from nemo_curator.stages.audio._agent._conformance import assert_agent_ready + +def test_my_stage_is_agent_ready(tmp_path): + def fixture(): + return AudioTask(data={"audio_filepath": str(_write_wav(tmp_path / "a.wav"))}) + assert_agent_ready(MyStage(), fixture, expected_cardinality="1:1", available_keys={"audio_filepath"}) +``` + +For GPU/model stages, reuse the existing fake-model/stub setup (see +`tests/stages/audio/test_agent_simulation_pipelines.py`) so the test needs no GPU. You don't need to +memorize the rules — if the test passes, the contract is honest. + +--- + +## Checklist (copy into your PR) + +- [ ] `AgentReady` + `describe()` with `reads`, `writes`, `cardinality`, honest `gates` +- [ ] every read/written `task.data` key is a `*_key` constructor field (no bare literals) +- [ ] new `*_key` concepts have a `_roles.KEY_ROLES` entry (or `INTERNAL_KEY_FIELDS`) +- [ ] capability card explains each externally consumed output's meaning, unit, + provenance, scope/granularity, propagation and a counterexample +- [ ] new `AudioTask`s preserve `_metadata` and `list(_stage_perf)` (manual — not covered by `assert_agent_ready`) +- [ ] if you override `process_batch`, write to disk, or set `lifecycle_side_effects` — decided + `gates.per_row_independent` (`True`/`False`, per instance if conditional); otherwise left it + alone and let the delta derive it +- [ ] `assert_agent_ready(...)` test added and green +- [ ] defaults unchanged → existing pipelines behave exactly as before + +Auto-derivation handles params/roles/dispatch/description; the card supplies meaning only the +stage author knows. Neither documentation step changes runtime defaults. diff --git a/nemo_curator/stages/audio/CLAUDE.md b/nemo_curator/stages/audio/CLAUDE.md new file mode 100644 index 0000000000..43c994c2d3 --- /dev/null +++ b/nemo_curator/stages/audio/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/nemo_curator/stages/audio/__init__.py b/nemo_curator/stages/audio/__init__.py index 693585cd05..2eee308dcf 100644 --- a/nemo_curator/stages/audio/__init__.py +++ b/nemo_curator/stages/audio/__init__.py @@ -32,11 +32,15 @@ "ALMDataOverlapStage": "nemo_curator.stages.audio.alm", "AudioDataFilterStage": "nemo_curator.stages.audio.advanced_pipelines", "BandFilterStage": "nemo_curator.stages.audio.filtering", + "ChannelCountStage": "nemo_curator.stages.audio.preprocessing", "GetAudioDurationStage": "nemo_curator.stages.audio.common", + "ManifestCheckpointStage": "nemo_curator.stages.audio.common", "ManifestReader": "nemo_curator.stages.audio.common", "ManifestWriterStage": "nemo_curator.stages.audio.common", "MonoConversionStage": "nemo_curator.stages.audio.preprocessing", + "PreserveByValueConditionsStage": "nemo_curator.stages.audio.common", "PreserveByValueStage": "nemo_curator.stages.audio.common", + "SampleRateFilterStage": "nemo_curator.stages.audio.preprocessing", "SIGMOSFilterStage": "nemo_curator.stages.audio.filtering", "SegmentConcatenationStage": "nemo_curator.stages.audio.preprocessing", "SpeakerSeparationStage": "nemo_curator.stages.audio.segmentation", @@ -51,12 +55,16 @@ "ALMDataOverlapStage", "AudioDataFilterStage", "BandFilterStage", + "ChannelCountStage", "GetAudioDurationStage", + "ManifestCheckpointStage", "ManifestReader", "ManifestWriterStage", "MonoConversionStage", + "PreserveByValueConditionsStage", "PreserveByValueStage", "SIGMOSFilterStage", + "SampleRateFilterStage", "SegmentConcatenationStage", "SpeakerSeparationStage", "TimestampMapperStage", diff --git a/nemo_curator/stages/audio/_agent/__init__.py b/nemo_curator/stages/audio/_agent/__init__.py new file mode 100644 index 0000000000..41eaa7659a --- /dev/null +++ b/nemo_curator/stages/audio/_agent/__init__.py @@ -0,0 +1,29 @@ +# 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. + +"""Private foundation for agent-driven audio pipeline construction. + +These modules are the STAGE-SIDE declaration layer: the vocabulary a stage uses +to describe itself (``StageContract``, ``Gates``, ``IOSpec``, the ``AgentReady`` +mixin, the shared role names) plus the discovery, planning and conformance code +that reads it. They live under ``stages/audio`` -- not under ``nemo_curator.audio_agent`` +-- on purpose: 43 stage modules import ``_agent_ready`` and 16 call into +``_residency`` from inside ``process()``. Moving them into the agent package would +make ``nemo_curator.stages.audio`` unusable without the agent installed, inverting +a dependency that today points one way only. + +Grouped into this subpackage purely so the stage tree reads as stages. Import the +public facade -- :mod:`nemo_curator.stages.audio.agent` -- rather than these +modules directly. +""" diff --git a/nemo_curator/stages/audio/_agent/_agent_ready.py b/nemo_curator/stages/audio/_agent/_agent_ready.py new file mode 100644 index 0000000000..e05838e081 --- /dev/null +++ b/nemo_curator/stages/audio/_agent/_agent_ready.py @@ -0,0 +1,413 @@ +# 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"] +ContractResolution = Literal["configured", "static_params_and_hints"] +WriteValueOrigin = Literal[ + "stage_generated", + "upstream_same_key", + "augments_upstream_same_key", + "transforms_upstream_same_key", +] + +# Semantic role vocabulary. Key *names* are agent-configurable -- every read/written key is a +# ``*_key`` field an agent may rename -- so two stages chain reliably only by matching the role +# a key plays, not its literal string. Roles resolve from the invariant field name (``_roles``), +# surviving a rename. ``"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._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 ConditionalWrite: + """A possible write whose presence or value origin depends on runtime data. + + ``writes`` uses the same task/segment key vocabulary as + :class:`IOSpec`. ``condition`` is factual, agent-facing prose grounded in + the runtime branch; it is deliberately not an executable predicate or an + intent/scope ontology. ``value_origin`` records only objective dataflow: + a new value, an unchanged pass-through, an in-place mapping augmentation, + or a transformed replacement of the same upstream key. This lets semantic + lineage preserve the original producer where it remains relevant. + + ``metadata_writes`` is the equivalent advisory surface for ``task._metadata`` + keys. It stays separate from :class:`IOSpec` because those keys are not + task-data/audio-form inputs. + + This metadata is additive. Mechanical planners continue to use + :attr:`StageContract.writes`; ``conditional_writes`` supplies the host + critic with possibility/provenance evidence and may also describe + conditional pass-through keys omitted from the legacy ``writes`` superset. + """ + + writes: IOSpec = field(default_factory=IOSpec) + condition: str = "" + value_origin: WriteValueOrigin = "stage_generated" + # Appended for positional compatibility with the original three fields. + metadata_writes: list[str] = 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 + # Constructor params holding this stage's OUTPUT paths, so a smoke knows what to redirect. + # Declared by the stage because only it knows; guessing from parameter names risks writing a + # smoke's output into the user's real directory. ``None`` means NOT DECLARED and is distinct + # from ``[]``: a disk writer that names no outputs cannot be sandboxed, so callers must + # refuse. ``[]`` positively claims "writes, but through no redirectable path param". + output_path_params: list[str] | None = None + 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 + # The stage derives durable output identity (filenames, member names, row + # identifiers, etc.) from framework ``task.task_id``. Metadata manifests + # serialize only task.data, so a resume reader cannot recreate that + # framework identity unless a future boundary explicitly gains such a + # restoration contract. Resume planners must refuse while it is unstable. + requires_stable_task_id: bool = False + # Is each output row computed from its own input row ALONE? ``None`` means NOT DECLARED, + # deliberately distinct from ``True``. Cardinality cannot answer this: a stage can be + # honestly 1:1 and still score each row against a corpus-wide statistic. It matters for + # delta runs, where such a stage would judge the changed files against the delta alone -- + # so undeclared makes a delta refuse and name the stage rather than assume the safe case. + per_row_independent: bool | None = None + + +@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 + # Input/output task types from the ``ProcessingStage[X, Y]`` generic, auto-derived + # at discovery. Enable the task-type compatibility check (e.g. a DocumentBatch + # producer feeding an AudioTask-only sink). + accepts_task_type: str | None = None + produces_task_type: str | None = None + # Task-data key VALUES this (configured) stage deletes from the task, so a + # downstream stage that needs that role/key can be caught (the "ASR after a + # waveform-stripper" key-flow class). Declared in describe(), verified by the + # conformance key-diff. + removes_keys: list[str] = field(default_factory=list) + # ``describe``/``cards`` can inspect a class without constructing it. That + # static view knows parameters and class hints but cannot honestly resolve + # configured reads, writes, cardinality, or removed keys. Appended here to + # preserve every existing positional constructor argument. + contract_resolution: ContractResolution = "configured" + # Runtime-data-dependent output possibilities. Appended for positional + # compatibility; this never changes stage execution or the legacy + # mechanical interpretation of ``writes``. + conditional_writes: list[ConditionalWrite] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-safe dict of this contract (``json.dumps`` never raises).""" + return { + "contract_resolution": self.contract_resolution, + "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, + "accepts_task_type": self.accepts_task_type, + "produces_task_type": self.produces_task_type, + "removes_keys": list(self.removes_keys), + "conditional_writes": [ + { + "writes": asdict(item.writes), + "condition": item.condition, + "value_origin": item.value_origin, + "metadata_writes": list(item.metadata_writes), + } + for item in self.conditional_writes + ], + } + + +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]] = {} + # ``*_key`` fields that are this stage's own bookkeeping and chain with nothing -- + # a counter or flag it records for readers, not a key another stage routes on. + # Declared here rather than in the shared ``_roles.INTERNAL_KEY_FIELDS`` so adding a + # stage does not mean editing a central table. Cross-stage keys still belong in + # ``KEY_ROLES``: roles are the vocabulary stages connect THROUGH, so letting each + # stage invent private role names would quietly weaken composition checking. + INTERNAL_KEY_FIELDS: ClassVar[frozenset[str]] = frozenset() + # Opt in to the runtime resource reading (peak VRAM / host RSS / throughput) that + # ``audio_agent.calibration.from_smoke`` turns into per-stage sizing facts. Declared + # here so the cost lands only on the stages the agent actually calibrates: no other + # modality pays for a metric it never consumes. Set False to opt a stage out. + RESOURCE_PROBE: ClassVar[bool] = True + + 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._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/_agent/_agent_registry.py b/nemo_curator/stages/audio/_agent/_agent_registry.py new file mode 100644 index 0000000000..2445d9b0eb --- /dev/null +++ b/nemo_curator/stages/audio/_agent/_agent_registry.py @@ -0,0 +1,569 @@ +# 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. + +"""Auto-derivation of agent-facing stage metadata. + +Turns a stage's constructor surface (dataclass fields or ``__init__`` signature) +into a list of :class:`~nemo_curator.stages.audio._agent._agent_ready.ParamSpec`, and +assembles the full discovery contract by merging those params and semantic key +roles into the stage's hand-written ``describe()``. Stages therefore never list +params by hand: + + contract = build_contract(stage) # dynamic, resolved key values + contract = SomeStage.describe_static() # static, instance-free (planning) +""" + +from __future__ import annotations + +import ast +import contextlib +import dataclasses +import inspect +import re +import sys +import textwrap +import types +import typing +from typing import Any, Literal, Union, get_args, get_origin + +from nemo_curator.stages.audio._agent._agent_ready import ( + Gates, + ParamSpec, + StageContract, + StaticHints, +) +from nemo_curator.stages.audio._agent._roles import ( + LITERAL_KEY_ROLES, + role_for_field, + role_overrides_for, +) + +# Framework-level constructor fields that are not stage-semantic knobs. +EXCLUDED_PARAM_NAMES = frozenset({"name", "resources", "batch_size", "runtime_env", "num_workers"}) + +_MISSING = dataclasses.MISSING + + +# --------------------------------------------------------------------------- # +# Type rendering / Literal handling +# --------------------------------------------------------------------------- # +def _module_globals(obj: Any) -> dict[str, Any]: # noqa: ANN401 + mod = sys.modules.get(getattr(obj, "__module__", "") or "") + return getattr(mod, "__dict__", {}) + + +def _resolve_hint(raw: Any, globalns: dict[str, Any]) -> Any: # noqa: ANN401 + """Resolve a possibly-stringized annotation to a type object, best-effort. + + ``from __future__ import annotations`` makes every annotation a string; we + eval it in the owning module's namespace. Heavy/forward refs that fail to + resolve are kept as their raw string (rendered verbatim, no Literal/choices). + """ + if raw is None or raw is inspect.Parameter.empty: + return None + if not isinstance(raw, str): + return raw + try: + return eval(raw, dict(globalns)) # noqa: S307 - annotations come from our own source + except Exception: # noqa: BLE001 + return raw + + +def _scalar_name(tp: type) -> str: + return {int: "int", float: "float", str: "str", bool: "bool"}.get(tp, getattr(tp, "__name__", str(tp))) + + +def _render_type(hint: Any) -> str: # noqa: ANN401, C901, PLR0911 + if hint is None or hint is inspect.Parameter.empty: + return "Any" + if isinstance(hint, str): + return hint + origin = get_origin(hint) + if origin is Literal: + elem = {type(a).__name__ for a in get_args(hint)} + if elem == {"str"}: + return "str" + if elem == {"int"}: + return "int" + if elem <= {"int", "float"}: + return "float" + return "str" + if origin is Union or origin is getattr(types, "UnionType", ()): + args = get_args(hint) + nullable = any(a is type(None) for a in args) + inner = [_render_type(a) for a in args if a is not type(None)] + rendered = " | ".join(inner) if inner else "Any" + return f"{rendered} | None" if nullable else rendered + if origin in (list, typing.List): # noqa: UP006 + sub = get_args(hint) + return f"list[{_render_type(sub[0])}]" if sub else "list" + if origin in (dict, typing.Dict): # noqa: UP006 + return "dict" + if origin in (tuple, typing.Tuple): # noqa: UP006 + return "tuple" + if isinstance(hint, type): + return _scalar_name(hint) + return str(hint) + + +def _literal_choices(hint: Any) -> list[Any] | None: # noqa: ANN401 + if isinstance(hint, str) or hint is None: + return None + if get_origin(hint) is Literal: + return list(get_args(hint)) + if get_origin(hint) is Union or get_origin(hint) is getattr(types, "UnionType", ()): + for a in get_args(hint): + if get_origin(a) is Literal: + return list(get_args(a)) + return None + + +# --------------------------------------------------------------------------- # +# Docstring Args parsing (single maintained source for param descriptions) +# --------------------------------------------------------------------------- # +_ARG_HDR = re.compile(r"^\s*(Args|Arguments|Parameters)\s*:\s*$") +_SECTION_HDR = re.compile(r"^\s*(Returns?|Raises?|Yields?|Notes?|Examples?|Attributes?|See Also|Warning|Todo)\s*:\s*$") +_ARG_LINE = re.compile(r"^(?P\s*)(?P[A-Za-z_]\w*)\s*(\([^)]*\))?\s*:\s*(?P.*)$") + + +def _parse_args_section(doc: str) -> dict[str, str]: + lines = doc.splitlines() + out: dict[str, str] = {} + in_args = False + current: str | None = None + arg_indent = 0 + for line in lines: + if _ARG_HDR.match(line): + in_args = True + current = None + continue + if not in_args: + continue + if _SECTION_HDR.match(line): + break + if not line.strip(): + continue + m = _ARG_LINE.match(line) + if m and (current is None or len(m.group("indent")) <= arg_indent + 1 or len(m.group("indent")) <= 8): # noqa: PLR2004 + current = m.group("name") + arg_indent = len(m.group("indent")) + out[current] = m.group("desc").strip() + elif current is not None: + out[current] = (out[current] + " " + line.strip()).strip() + return out + + +def _docstring_arg_descriptions(cls: type) -> dict[str, str]: + out: dict[str, str] = {} + for klass in reversed(cls.__mro__): # base first, so subclass docstrings win + doc = klass.__dict__.get("__doc__") + if doc: + out.update(_parse_args_section(doc)) + return out + + +# --------------------------------------------------------------------------- # +# Param derivation +# --------------------------------------------------------------------------- # +def _as_class(stage_or_cls: Any) -> type: # noqa: ANN401 + return stage_or_cls if isinstance(stage_or_cls, type) else type(stage_or_cls) + + +def _call_factory(factory: Any) -> tuple[Any, bool]: # noqa: ANN401 + try: + return factory(), False + except Exception: # noqa: BLE001 + return "", False + + +def _dataclass_params(cls: type, descriptions: dict[str, str]) -> list[ParamSpec]: + globalns = _module_globals(cls) + params: list[ParamSpec] = [] + for f in dataclasses.fields(cls): + if not f.init or f.name in EXCLUDED_PARAM_NAMES or f.name.startswith("_"): + continue + if f.default is not _MISSING: + default, required = f.default, False + elif f.default_factory is not _MISSING: + default, required = _call_factory(f.default_factory) + else: + default, required = None, True + hint = _resolve_hint(f.type, globalns) + params.append( + ParamSpec( + name=f.name, + type=_render_type(hint), + default=default, + required=required, + choices=_literal_choices(hint), + description=descriptions.get(f.name), + role=role_for_field(f.name) if f.name.endswith("_key") else None, + ) + ) + return params + + +def _init_params(cls: type, descriptions: dict[str, str]) -> list[ParamSpec]: + try: + sig = inspect.signature(cls.__init__) + except (TypeError, ValueError): + return [] + globalns = _module_globals(cls) + params: list[ParamSpec] = [] + for name, p in sig.parameters.items(): + if name == "self" or p.kind in (p.VAR_POSITIONAL, p.VAR_KEYWORD): + continue + if name in EXCLUDED_PARAM_NAMES or name.startswith("_"): + continue + required = p.default is inspect.Parameter.empty + default = None if required else p.default + hint = _resolve_hint(p.annotation, globalns) + params.append( + ParamSpec( + name=name, + type=_render_type(hint), + default=default, + required=required, + choices=_literal_choices(hint), + description=descriptions.get(name), + role=role_for_field(name) if name.endswith("_key") else None, + ) + ) + return params + + +def stage_params(stage_or_cls: Any) -> list[ParamSpec]: # noqa: ANN401 + """Auto-derive the configurable parameters of a stage. + + Dataclass stages use :func:`dataclasses.fields`; plain ``__init__`` stages + use the constructor signature. ``Literal[...]`` annotations become + ``choices``; per-param descriptions come from the class docstring ``Args:`` + section; ``*_key`` params get a semantic ``role``. + """ + cls = _as_class(stage_or_cls) + descriptions = _docstring_arg_descriptions(cls) + if dataclasses.is_dataclass(cls): + return _dataclass_params(cls, descriptions) + return _init_params(cls, descriptions) + + +# --------------------------------------------------------------------------- # +# Contract assembly +# --------------------------------------------------------------------------- # +def _derived_dispatch(cls: type, declared: str) -> str: + """Drive ``dispatch`` from the framework's batch-support truth source.""" + from nemo_curator.stages.base import ProcessingStage + + if declared in ("process", "process_batch"): + return declared + overrides_batch = getattr(cls, "process_batch", None) is not ProcessingStage.process_batch + return "process_batch" if overrides_batch else "process" + + +def _task_type_name(t: Any) -> str | None: # noqa: ANN401 + """The class name of a generic arg, or None for a TypeVar/non-type.""" + return t.__name__ if isinstance(t, type) else None + + +def _task_types(cls: type) -> tuple[str | None, str | None]: + """``(accepts, produces)`` task-type names from the ``ProcessingStage[X, Y]`` generic. + + Walks the MRO's ``__orig_bases__`` for the parametrized ProcessingStage base + (e.g. ``ProcessingStage[AudioTask, DocumentBatch]``). Returns ``(None, None)`` + when unparametrized (bare TypeVars) or not found — those become ``uncertain`` + at the task-type check rather than a false mismatch. + """ + from nemo_curator.stages.base import ProcessingStage + + for klass in cls.__mro__: + for base in getattr(klass, "__orig_bases__", ()) or (): + origin = get_origin(base) + if origin is None: + continue + try: + is_ps = origin is ProcessingStage or (isinstance(origin, type) and issubclass(origin, ProcessingStage)) + except TypeError: + is_ps = False + if not is_ps: + continue + args = get_args(base) + if len(args) == 2: # noqa: PLR2004 - ProcessingStage[X, Y] has exactly two type args + return _task_type_name(args[0]), _task_type_name(args[1]) + return None, None + + +def _first_doc_line(cls: type) -> str | None: + doc = inspect.getdoc(cls) + if not doc: + return None + for line in doc.splitlines(): + if line.strip(): + return line.strip() + return None + + +def _contract_referenced_keys(contract: StageContract) -> set[str]: + keys: set[str] = set() + for spec in [contract.reads, contract.writes, *contract.reads_one_of]: + keys.update(spec.data_keys) + keys.update(spec.segment_data_keys) + keys.update(contract.metadata_reads) + keys.update(contract.metadata_writes) + for conditional in contract.conditional_writes: + keys.update(conditional.writes.data_keys) + keys.update(conditional.writes.segment_data_keys) + keys.update(conditional.metadata_writes) + return keys + + +def _key_attr_names(stage: Any) -> list[str]: # noqa: ANN401 + return [a for a in dir(stage) if a.endswith("_key") and not a.startswith("__")] + + +def _resolve_key_roles(stage: Any, contract: StageContract) -> dict[str, str]: # noqa: ANN401 + """Map resolved key *values* referenced by the contract to semantic roles.""" + cls = _as_class(stage) + overrides = role_overrides_for(cls) + roles: dict[str, str] = {} + # 1. From *_key attributes on the instance (their resolved values). + if not isinstance(stage, type): + for attr in _key_attr_names(stage): + with contextlib.suppress(Exception): + val = getattr(stage, attr) + if isinstance(val, str) and val: + role = overrides.get(attr) or role_for_field(attr) + if role != "unknown": + roles[val] = role + # 2. Literal-default fallback for producer keys with no *_key field. + for key in _contract_referenced_keys(contract): + if key not in roles: + literal = LITERAL_KEY_ROLES.get(key) + if literal is not None: + roles[key] = literal + return roles + + +def _stage_resources(stage_or_cls: Any) -> Any: # noqa: ANN401 + """The stage's declared ``resources``, resolvable without constructing the stage. + + An instance answers directly. A class does not, so fall back to the dataclass field + default -- which is how :func:`static_contract` can report GPU need for a stage whose + constructor demands arguments. + """ + if not isinstance(stage_or_cls, type): + return getattr(stage_or_cls, "resources", None) + if not dataclasses.is_dataclass(stage_or_cls): + return None + for f in dataclasses.fields(stage_or_cls): + if f.name != "resources": + continue + if f.default_factory is not dataclasses.MISSING: # type: ignore[misc] + with contextlib.suppress(Exception): + return f.default_factory() # type: ignore[misc] + if f.default is not dataclasses.MISSING: + return f.default + return None + + +def _derived_wrappable(declared: bool, stage_or_cls: Any) -> bool: # noqa: ANN401 + """A composite cannot be wrapped; anything else keeps what it declared. + + One-sided like :func:`_derived_gates`: only ``True -> False``. The default is ``True``, so + a stage that never mentions ``wrappable`` is indistinguishable from one that wrote it, and + upgrading a declared ``False`` back to ``True`` would silently contradict the author. + """ + from nemo_curator.stages.base import CompositeStage + + cls = stage_or_cls if isinstance(stage_or_cls, type) else type(stage_or_cls) + return declared and not (isinstance(cls, type) and issubclass(cls, CompositeStage)) + + +def _derived_gates(gates: Gates, stage_or_cls: Any) -> Gates: # noqa: ANN401 + """Fill ``requires_gpu`` from the stage's reserved resources. + + One-sided on purpose: only ``False -> True``. ``Gates`` cannot distinguish "the author + wrote False" from "the author left the default", so overriding downward would trample a + stage that honestly declares the gate while reserving nothing -- InferenceSortformerStage + passes ``map_location="cuda"`` unconditionally. Upgrading is always safe: a stage that + reserves a GPU needs one, whatever it forgot to say, and the damaging direction is the + false negative that puts a GPU stage on a CPU worker. + + Uses ``Resources.requires_gpu`` (``gpus > 0 or gpu_memory_gb > 0``) rather than + ``gpus > 0``: ``__post_init__`` converts ``gpu_memory_gb`` by dividing by detected device + memory and rounding to one decimal, so a small reservation rounds to zero gpus and the + coarser test silently reports no GPU -- and it would answer differently per machine. + """ + if gates.requires_gpu: + return gates + resources = _stage_resources(stage_or_cls) + if resources is not None and bool(getattr(resources, "requires_gpu", False)): + return dataclasses.replace(gates, requires_gpu=True) + return gates + + +def _describe_cardinalities(cls: type) -> list[str] | None: + """Every cardinality literal ``cls.describe()`` can return, or ``None`` if unreadable. + + Read from the source rather than by calling ``describe()``, because the whole point of + the static path is that the stage may not be constructible: ``adapter_target`` and + ``model_id`` have no defaults, so instantiating ASRStage to ask it a question is exactly + what :func:`static_contract` exists to avoid. + + Returns ``None`` when the value is computed some way this cannot read -- a helper call, + a local variable, a dict lookup. Refusing to guess is the point: a wrong cardinality is + worse than an admittedly unresolved one. + """ + describe = getattr(cls, "describe", None) + if describe is None: + return None + try: + tree = ast.parse(textwrap.dedent(inspect.getsource(describe))) + except (OSError, TypeError, SyntaxError, IndentationError): + return None # builtin, C extension, or source not on disk + found: list[str] = [] + for node in ast.walk(tree): + if not isinstance(node, ast.keyword) or node.arg != "cardinality": + continue + branches = [node.value.body, node.value.orelse] if isinstance(node.value, ast.IfExp) else [node.value] + for branch in branches: + if isinstance(branch, ast.Constant) and isinstance(branch.value, str): + found.append(branch.value) + else: + return None + return list(dict.fromkeys(found)) + + +def _derived_cardinality(declared: str, options: list[str], cls: type) -> tuple[str, list[str]]: + """Recover cardinality on the instance-free path. ``(cardinality, cardinality_options)``. + + One-sided like :func:`_derived_gates`: only ``"1:1" -> something else``. ``1:1`` is the + weakest claim a stage can make -- it says row counts do not change -- and it is also the + dataclass default, so a stage that never mentions cardinality is indistinguishable from + one that means it. Overriding a stage's explicit non-``1:1`` answer would trample an + author who knows better; filling in an unset one cannot. + + Without this the static contract reported ``1:1`` for every fan-out, filter and N:1 stage + in the catalog, because ``StaticHints`` has no cardinality field and the default stands. + A planner reading a param-less ``describe`` therefore believed ManifestReader emitted one + task and that a filter never dropped a row. + + When ``describe()`` can return more than one cardinality the answer genuinely depends on + params, so the possibilities go to ``cardinality_options``, which is the field that already + means "this varies". Which value stands alongside them depends on whether ``1:1`` is one of + them. If it is, leaving the default is honest -- it names a real possibility. If it is not, + the default is a claim no configuration of the stage can satisfy, and an unlabelled ``1:1`` + does not read as "unknown", it reads as "row counts do not change" (the same trap + ``ResolvedContract`` documents for empty reads/writes). There the most conservative + possibility is published instead. + """ + if declared != "1:1": + return declared, options + found = _describe_cardinalities(cls) + if not found or found == ["1:1"]: + return declared, options + if len(found) == 1: + return found[0], options + if "1:1" in found: + return declared, options or sorted(found) + return _most_conservative(found), options or sorted(found) + + +# Decreasing order of what a cardinality lets a reader assume. ``N:1`` permits the least (delta +# will not even trace through it), then the two that change task counts, then the two that +# promise one task out per task in. Used only to choose among possibilities that are ALL real +# for some configuration, so every entry is a true statement about the stage -- this picks the +# one whose being wrong costs a caller the least. +_CARDINALITY_BY_CAUTION = ("N:1", "filter", "1:N fan-out", "1:1 nested-list", "1:1") + + +def _most_conservative(found: list[str]) -> str: + """The possibility in ``found`` that lets a reader assume the least.""" + for candidate in _CARDINALITY_BY_CAUTION: + if candidate in found: + return candidate + return found[0] # an unrecognized literal: prefer the stage's own word over a guess + + +def build_contract(stage: Any) -> StageContract: # noqa: ANN401 + """Return ``stage.describe()`` enriched with auto-derived params + key roles. + + Hand-written ``params`` in ``describe()`` override auto-derived ones. + ``dispatch`` is derived from batch support; ``batch_only``/``stage_id`` are + filled when unset. The single entry point used by catalog + serialization. + """ + base = stage.describe() + derived = stage_params(stage) + if base.params: + by_name = {p.name: p for p in derived} + by_name.update({p.name: p for p in base.params}) + params = list(by_name.values()) + else: + params = derived + key_roles = _resolve_key_roles(stage, base) or dict(base.key_roles) + cls = _as_class(stage) + accepts_tt, produces_tt = _task_types(cls) + return dataclasses.replace( + base, + params=params, + key_roles=key_roles, + dispatch=_derived_dispatch(cls, base.dispatch), + batch_only=base.batch_only or bool(getattr(cls, "BATCH_ONLY", False)), + stage_id=base.stage_id or cls.__name__, + description=base.description or _first_doc_line(cls), + accepts_task_type=base.accepts_task_type or accepts_tt, + produces_task_type=base.produces_task_type or produces_tt, + gates=_derived_gates(base.gates, stage), + wrappable=_derived_wrappable(base.wrappable, stage), + ) + + +def _static_key_roles(cls: type) -> dict[str, str]: + """Key roles from ``*_key`` field DEFAULT values (no instantiation).""" + overrides = role_overrides_for(cls) + roles: dict[str, str] = {} + for p in stage_params(cls): + if not p.name.endswith("_key") or not isinstance(p.default, str) or not p.default: + continue + role = overrides.get(p.name) or role_for_field(p.name) + if role != "unknown": + roles[p.default] = role + return roles + + +def static_contract(cls: type) -> StageContract: + """Instance-free discovery contract (params, roles, gates, dispatch, hints). + + Does not run ``__init__``; safe for stages with required constructor args. + Reads/writes/cardinality are NOT resolved here — use :func:`build_contract` + on an instance for those. + """ + hints: StaticHints = getattr(cls, "AGENT_STATIC", None) or StaticHints() + accepts_tt, produces_tt = _task_types(cls) + cardinality, cardinality_options = _derived_cardinality("1:1", list(hints.cardinality_options), cls) + return StageContract( + contract_resolution="static_params_and_hints", + cardinality=cardinality, + cardinality_options=cardinality_options, + gates=_derived_gates(hints.gates, cls), + wrappable=_derived_wrappable(True, cls), + dispatch=_derived_dispatch(cls, hints.dispatch), + error_policy=hints.error_policy, + description=hints.description or _first_doc_line(cls), + stage_id=hints.stage_id or cls.__name__, + params=stage_params(cls), + key_roles=_static_key_roles(cls), + batch_only=bool(getattr(cls, "BATCH_ONLY", False)), + accepts_task_type=accepts_tt, + produces_task_type=produces_tt, + ) diff --git a/nemo_curator/stages/audio/_agent/_catalog.py b/nemo_curator/stages/audio/_agent/_catalog.py new file mode 100644 index 0000000000..2c5bf1b22e --- /dev/null +++ b/nemo_curator/stages/audio/_agent/_catalog.py @@ -0,0 +1,265 @@ +# 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. + +"""Discovery API for agent-ready audio stages. + +Built on the framework's existing class-name registry +(:data:`nemo_curator.stages.base._STAGE_REGISTRY`, populated by ``StageMeta``); +no new registry is introduced. ``_ensure_audio_stages_imported`` triggers the +audio stage modules so their registration is populated before discovery. + + list_agent_ready_stages() # -> ["MonoConversionStage", ...] + describe_stage("UTMOSFilterStage") # -> StageContract (static, instance-free) + catalog_as_json() # -> JSON an agent/UI can consume +""" + +from __future__ import annotations + +import importlib +import json +import pkgutil +import warnings +from collections import defaultdict +from typing import TYPE_CHECKING, Any + +from nemo_curator.stages.audio._agent._agent_ready import AgentReady, to_json_schema +from nemo_curator.stages.audio._agent._agent_registry import build_contract, static_contract +from nemo_curator.stages.audio._agent._conformance import produced_roles + +if TYPE_CHECKING: + from nemo_curator.stages.audio._agent._agent_ready import StageContract + +_IMPORTED = False +# Modules discovery could not import, kept so the failure can be reported rather than +# silently shrinking the catalog. +_SKIPPED: list[dict[str, str]] = [] + + +def _ensure_audio_stages_imported() -> None: + """Import audio stage modules so ``StageMeta`` has registered their classes. + + Defensive: a submodule whose optional heavy dependency (whisperx, pyannote, + nemo_text_processing, ...) is absent is skipped with a warning rather than + breaking discovery. Idempotent. + """ + global _IMPORTED # noqa: PLW0603 + if _IMPORTED: + return + import nemo_curator.stages.audio as audio_pkg + + # onerror: a failing subpackage __init__ (non-ImportError too — OSError / + # RuntimeError are realistic for heavy audio deps) must skip, not kill + # discovery; the loop body below warns for the same module on import. + for modinfo in pkgutil.walk_packages( + audio_pkg.__path__, prefix=audio_pkg.__name__ + ".", onerror=lambda _name: None + ): + leaf = modinfo.name.rsplit(".", 1)[-1] + if leaf.startswith("_"): # private support modules carry no stages + continue + try: + importlib.import_module(modinfo.name) + except Exception as e: # noqa: BLE001 - optional dep or import-time issue; skip + warnings.warn(f"audio catalog: skipped {modinfo.name} ({type(e).__name__}: {e})", stacklevel=2) + # Recorded, not just warned: a warning goes to stderr where the agent's JSON + # consumer never sees it, and a silently shorter catalog looks like a smaller + # library. Deduplicated because ``_IMPORTED`` is set only after the loop completes, + # so anything escaping it makes the next call re-walk and append every failure + # again -- which reads like a worsening install. See :func:`unavailable_modules`. + if not any(entry["module"] == modinfo.name for entry in _SKIPPED): + _SKIPPED.append({"module": modinfo.name, "error": f"{type(e).__name__}: {e}"}) + _IMPORTED = True + + +def unavailable_modules() -> list[dict[str, str]]: + """Stage modules that could not be imported, so a caller can report what is MISSING. + + Discovery degrades to whatever imported successfully. Without this, a CPU-only + (``audio_cpu``) install -- a supported profile -- simply has no ASR or diarization + stages, and the agent concludes they do not exist rather than that they are unavailable + *here*, which is the difference between "your library cannot do this" and "install the + GPU extra". + """ + _ensure_audio_stages_imported() + return [dict(entry) for entry in _SKIPPED] + + +def _agent_ready_registry() -> dict[str, type]: + from nemo_curator.stages.base import _STAGE_REGISTRY + + return { + name: cls for name, cls in _STAGE_REGISTRY.items() if isinstance(cls, type) and issubclass(cls, AgentReady) + } + + +def list_agent_ready_stages() -> list[str]: + """Sorted class names of all registered agent-ready audio stages.""" + _ensure_audio_stages_imported() + return sorted(_agent_ready_registry()) + + +def get_agent_ready_stage_class(name: str) -> type: + """Return the registered stage class for ``name`` (must be agent-ready).""" + _ensure_audio_stages_imported() + registry = _agent_ready_registry() + if name not in registry: + msg = f"{name!r} is not a registered agent-ready audio stage" + raise KeyError(msg) + return registry[name] + + +def describe_stage(name: str, stage: AgentReady | None = None) -> StageContract: + """Return a stage's contract. + + With ``stage`` (an instance) -> dynamic contract with resolved key values. + Otherwise -> instance-free ``static_contract`` (no resolved keys/cardinality). + """ + if stage is not None: + return build_contract(stage) + return static_contract(get_agent_ready_stage_class(name)) + + +def audio_stage_catalog(*, include_dynamic_defaults: bool = False) -> list[dict[str, Any]]: + """Return the catalog as a list of ``{name, contract[, default_contract]}`` dicts. + + ``contract`` is the static (instance-free) contract. ``default_contract`` + (only with ``include_dynamic_defaults``) is the dynamic contract of a + no-arg instance, attempted best-effort and ``None`` when the stage needs + required constructor args. + """ + entries: list[dict[str, Any]] = [] + for name in list_agent_ready_stages(): + cls = get_agent_ready_stage_class(name) + static = static_contract(cls) + entry: dict[str, Any] = { + "name": name, + "contract": static.to_dict(), + # A JSON-Schema config form for the stage's params — directly usable + # as an agent tool-argument schema (enum/default/x-role included). + "params_schema": to_json_schema(static.params), + } + if include_dynamic_defaults: + try: + entry["default_contract"] = build_contract(cls()).to_dict() + except Exception: # noqa: BLE001 - required-arg stages have no no-arg default + entry["default_contract"] = None + entries.append(entry) + return entries + + +def catalog_as_json(*, include_dynamic_defaults: bool = False, indent: int | None = None) -> str: + """JSON-serialized :func:`audio_stage_catalog` (an agent/UI tool schema).""" + return json.dumps(audio_stage_catalog(include_dynamic_defaults=include_dynamic_defaults), indent=indent) + + +# --------------------------------------------------------------------------- # +# Role -> producer/consumer index (composition + repair) +# --------------------------------------------------------------------------- # +def _consumed_roles(contract: StageContract) -> set[str]: + """Semantic roles a stage requires (primary reads + every reads_one_of option).""" + roles = { + contract.key_roles.get(k, "unknown") for k in [*contract.reads.data_keys, *contract.reads.segment_data_keys] + } + for opt in contract.reads_one_of: + roles |= {contract.key_roles.get(k, "unknown") for k in [*opt.data_keys, *opt.segment_data_keys]} + return roles - {"unknown"} + + +def _dummy_for_param(type_str: str | None) -> Any: # noqa: ANN401 - placeholder is deliberately any primitive shape + """A harmless placeholder for a required constructor arg, so ``describe()`` can + run for a required-arg stage. ``*_key`` fields have defaults (never required), + so these dummies only fill non-semantic args (paths, model names) and never + perturb the resolved key roles.""" + t = (type_str or "").lower() + if "bool" in t: + return False + if "int" in t: + return 0 + if "float" in t: + return 0.0 + if t.startswith("list"): + return [] + if t.startswith("dict"): + return {} + return "x" # str / path / model-name / anything else + + +def _default_contract(cls: type) -> StageContract | None: + """Best-effort dynamic contract. Tries progressively: a no-arg instance, then + a probe filling required args with harmless dummies, then also filling + ``None``-default args (to satisfy "one-of" ``__post_init__`` guards, e.g. ASR + needs ``model_name`` OR ``asr_model``). Only ``describe()`` is called, and + ``*_key`` fields keep their real defaults, so produced/consumed roles stay + correct. ``None`` only if every probe fails (e.g. needs a live model object).""" + try: + return build_contract(cls()) + except Exception: # noqa: BLE001, S110 - deliberate fall-through to dummy-filled probes + pass + from nemo_curator.stages.audio._agent._agent_registry import stage_params + + params = stage_params(cls) + for also_fill_none in (False, True): + kwargs = { + p.name: _dummy_for_param(p.type) for p in params if p.required or (also_fill_none and p.default is None) + } + if not kwargs: + continue + try: + return build_contract(cls(**kwargs)) + except Exception: # noqa: BLE001, S112 - deliberately try the next, broader probe + continue + return None + + +def role_index() -> dict[str, Any]: + """Map each semantic role to the stages that produce/consume it. + + Returns ``{"producers": {role: [stage, ...]}, "consumers": {...}, + "unresolved_stages": [...]}``. Built from each stage's no-arg dynamic + contract when possible, else from a probe instance with required (and + one-of ``None``-default) args filled by harmless dummies — only + ``describe()`` runs and ``*_key`` defaults are preserved, so the roles stay + correct. ``unresolved_stages`` lists only stages where every probe fails. + + This is what turns an ``unsatisfied_reads`` validation error into an + actionable repair ("insert a stage that produces role X") and lets an agent + detect an *unproducible* role (``find_producers`` returns ``[]``). + """ + producers: dict[str, set[str]] = defaultdict(set) + consumers: dict[str, set[str]] = defaultdict(set) + unresolved: list[str] = [] + for name in list_agent_ready_stages(): + contract = _default_contract(get_agent_ready_stage_class(name)) + if contract is None: + unresolved.append(name) + continue + for role in produced_roles(contract): + producers[role].add(name) + for role in _consumed_roles(contract): + consumers[role].add(name) + return { + "producers": {r: sorted(v) for r, v in sorted(producers.items())}, + "consumers": {r: sorted(v) for r, v in sorted(consumers.items())}, + "unresolved_stages": sorted(unresolved), + } + + +def find_producers(role: str) -> list[str]: + """Stages that produce ``role``. An empty list means no stage produces it + (the role is *unproducible* — an agent should not try to satisfy it).""" + return role_index()["producers"].get(role, []) + + +def find_consumers(role: str) -> list[str]: + """Stages that consume ``role``.""" + return role_index()["consumers"].get(role, []) diff --git a/nemo_curator/stages/audio/_agent/_composite.py b/nemo_curator/stages/audio/_agent/_composite.py new file mode 100644 index 0000000000..181bd90ce3 --- /dev/null +++ b/nemo_curator/stages/audio/_agent/_composite.py @@ -0,0 +1,192 @@ +# 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. + +"""Resolve a recipe's stages into the concrete stages a backend will actually run. + +A :class:`~nemo_curator.stages.base.CompositeStage` describes itself and nothing else -- +``SplitASRAlignJoinStage.describe()`` returns ``StageContract(wrappable=False)``, declaring no +reads and no writes -- while ``decompose()`` builds the three real stages underneath. Anything +that reasons about a recipe from contracts alone is therefore blind at exactly the stages that +do the work, and blind in a way that is invisible: the composite looks like a stage with no +requirements rather than a stage whose requirements are unknown. + +That blindness has cost real runs. ``SplitLongAudioStage``, the first stage inside +``SplitASRAlignJoinStage``, requires a ``segments`` key. A pipeline that diarized into +``diar_segments`` validated clean, downloaded two models, ran diarization on the GPU, and only +then refused to start the splitter -- a failure the composite's own ``decompose()`` had spelled +out all along, including a comment about this precise mismatch. + +Expanding the *configured* composite is what makes this exact rather than approximate: the +children carry the parameters the caller actually set, so ``SplitASRAlignJoinStage( +segments_key="diar_segments")`` yields ``SplitLongAudioStage(segments_key="diar_segments")`` and +the check reflects what will run rather than what the defaults would have run. + +When a composite cannot be expanded -- it raises, returns nothing, returns something that is not +a stage, or returns another composite (which the executor itself refuses) -- no leaf is invented +for it. It is reported in :attr:`Expansion.opaque` so callers can fall back to whatever they did +before rather than reason from a fabricated stage list. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from nemo_curator.stages.base import CompositeStage, ProcessingStage + + +@dataclass(frozen=True) +class ExpandedStage: + """One concrete stage the backend will run, and where in the recipe it came from.""" + + recipe_index: int + """Index of the recipe-level stage this came from -- what the caller wrote.""" + + stage: Any + """The configured, concrete (non-composite) stage instance.""" + + path: tuple[int, ...] = () + """Child indices from the recipe stage down to this leaf. Empty for a top-level stage.""" + + composite_ref: str | None = None + """Class name of the recipe-level composite this was expanded from, if any.""" + + @property + def label(self) -> str: + """How to name this stage to someone who only knows the recipe they wrote. + + A bare ``SplitLongAudioStage`` is not a stage the caller has ever heard of; they + configured ``SplitASRAlignJoinStage``. Naming both ends keeps the report actionable. + """ + name = type(self.stage).__name__ + return f"{self.composite_ref} -> {name}" if self.composite_ref else name + + +@dataclass(frozen=True) +class Expansion: + """The concrete run order, plus the composites that refused to reveal theirs.""" + + stages: list[ExpandedStage] = field(default_factory=list) + opaque: dict[int, str] = field(default_factory=dict) + """recipe_index -> why that composite could not be expanded.""" + + unrunnable: dict[int, str] = field(default_factory=dict) + """recipe_index -> why the executor will refuse this stage outright. + + Separate from :attr:`opaque` because the two demand opposite answers. Opaque means "we + cannot tell" and degrades to a warning; this means "we can tell, and it will fail", which + has to reach the caller as an error before they confirm a full-scale run. + """ + + @property + def fully_resolved(self) -> bool: + return not self.opaque and not self.unrunnable + + def by_recipe_index(self) -> dict[int, list[ExpandedStage]]: + """Leaves grouped under the recipe stage that produced them, in run order. + + Callers walk the recipe so an opaque composite keeps its place in the sequence; the + leaves it would have contributed are simply absent rather than reordered. + """ + grouped: dict[int, list[ExpandedStage]] = {} + for item in self.stages: + grouped.setdefault(item.recipe_index, []).append(item) + return grouped + + +def _nested_composite(child: Any) -> bool: # noqa: ANN401 - any child stage + """Whether a child is itself a composite the executor would refuse. + + Mirrors ``Pipeline._decompose_stages`` exactly: a ``CompositeStage`` whose ``decompose()`` + returns just itself is run as an ordinary stage, so only a genuinely decomposing one counts. + """ + if not isinstance(child, CompositeStage): + return False + try: + return len(child.decompose()) > 1 + except Exception: # noqa: BLE001 - an undecomposable child is handled by the caller's checks + return False + + +def _decompose(stage: Any) -> tuple[list[Any], str | None]: # noqa: ANN401 - any composite + """This composite's children, or the reason they cannot be trusted as a stage list.""" + try: + children = list(stage.decompose_and_apply_with() or ()) + except Exception as exc: # noqa: BLE001 - a composite that cannot plan-time decompose stays opaque + return [], f"decompose() raised {type(exc).__name__}" + if not children: + return [], "decomposition produced no stages" + alien = next((c for c in children if not isinstance(c, ProcessingStage)), None) + if alien is not None: + return [], f"decomposition returned {type(alien).__name__}, not a ProcessingStage" + return children, None + + +def expand_composites(stages: list[Any]) -> Expansion: + """Flatten composites into the stages a backend will run, preserving recipe order. + + Expansion is SINGLE-LEVEL, because that is the only shape the executor supports: + ``Pipeline._decompose_stages`` expands each stage once and raises ``TypeError`` + ("Nested composition is not supported") if a child is itself a decomposing composite. + Modelling deeper nesting here would be worse than not modelling it -- validation would + approve a recipe the backend then refuses to run -- and no composite in the catalog is + deeper than one level anyway. + + Uses ``decompose_and_apply_with()`` rather than ``decompose()`` so a composite configured + through ``with_()`` contributes the same resource overrides the executor will see -- the + call ``Pipeline._decompose_stages`` makes, so callers reason about the real schedule. + + A composite that cannot expand contributes no leaf and lands in :attr:`Expansion.opaque`. + Inventing a leaf would be worse than admitting the gap: a caller that trusts a guessed stage + list reports confident nonsense, whereas one that sees the gap can stay conservative. + """ + out: list[ExpandedStage] = [] + opaque: dict[int, str] = {} + unrunnable: dict[int, str] = {} + + for index, stage in enumerate(stages): + if not isinstance(stage, CompositeStage): + out.append(ExpandedStage(index, stage)) + continue + children, reason = _decompose(stage) + if reason is not None: + opaque[index] = reason + continue + if len(children) == 1: + # ``Pipeline._decompose_stages`` only substitutes children when there is more than + # one, so a single-child decomposition leaves the COMPOSITE in the execution list -- + # and ``CompositeStage.process`` raises "should not be executed directly" on the + # first task. Substituting the child here would model a pipeline the backend will + # never run and hand the caller a clean verdict for a recipe that dies on contact. + unrunnable[index] = ( + f"decomposes into a single stage ({type(children[0]).__name__}), which the " + f"executor does not substitute -- it runs the composite itself and raises" + ) + continue + # After the length check, exactly as the executor orders it: the nested-composite + # rejection lives inside its ``len(sub_stages) > 1`` branch and is never reached for a + # single child. Asking first inverted the verdict for a composite that decomposes into + # one decomposing composite -- reported as opaque, "we cannot tell", when the executor + # can tell perfectly well that it will run the outer composite and raise. + nested = next((c for c in children if _nested_composite(c)), None) + if nested is not None: + opaque[index] = ( + f"decomposition returned another composite ({type(nested).__name__}); " + "nested composition is not supported" + ) + continue + ref = type(stage).__name__ + out.extend(ExpandedStage(index, child, (child_index,), ref) for child_index, child in enumerate(children)) + + return Expansion(stages=out, opaque=opaque, unrunnable=unrunnable) diff --git a/nemo_curator/stages/audio/_agent/_conformance.py b/nemo_curator/stages/audio/_agent/_conformance.py new file mode 100644 index 0000000000..7cd7f3f934 --- /dev/null +++ b/nemo_curator/stages/audio/_agent/_conformance.py @@ -0,0 +1,442 @@ +# 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. + +"""Conformance harness for agent-ready audio stages. + +``assert_agent_ready(stage, fixture_factory, ...)`` is the gate each stage's CL +must pass. It runs a set of *static* checks (contract well-formedness, semantic +roles, JSON serialization, by-role read satisfiability) that need no execution, +plus optional *dynamic* checks (run ``process``/``process_batch`` on a fixture +and verify declared writes appear, no undeclared top-level keys leak, cardinality +matches runtime, and ``accepts``/``produces`` hold). + +The static checks alone catch the contract↔reality drift the prototype lacked +and can sweep every stage with no fixtures (see ``assert_contract_wellformed``). +""" + +# ruff: noqa: S101 - this module IS the assertion harness; `assert` is its output format. +# It ships under nemo_curator/ (not tests/) because stage authors call it from their own +# suites, so the tests/** ignore does not reach it; raising would lose pytest's assertion +# rewriting that makes these failures readable. + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any, get_args + +from nemo_curator.stages.audio._agent._agent_ready import ( + AudioForm, + Cardinality, + ProducedForm, + Role, + StageContract, + WriteValueOrigin, + to_json_schema, +) +from nemo_curator.stages.audio._agent._agent_registry import build_contract, static_contract +from nemo_curator.stages.audio._agent._residency import accepts_for_residency +from nemo_curator.stages.audio._agent._roles import field_has_declared_role + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + +_VALID_CARDINALITY: frozenset[str] = frozenset(get_args(Cardinality)) +_VALID_ROLES: frozenset[str] = frozenset(get_args(Role)) +_VALID_ACCEPTS: frozenset[str] = frozenset(get_args(AudioForm)) +_VALID_PRODUCES: frozenset[str] = frozenset(get_args(ProducedForm)) +_VALID_WRITE_VALUE_ORIGINS: frozenset[str] = frozenset(get_args(WriteValueOrigin)) +# The param a delta run sets to feed a source only the files that changed. +_NARROWING_PARAM = "include_files" + + +# --------------------------------------------------------------------------- # +# By-role matching (the planner primitive) +# --------------------------------------------------------------------------- # +def _spec_roles(contract: StageContract, spec_keys: Iterable[str]) -> set[str]: + return {contract.key_roles.get(k, "unknown") for k in spec_keys} + + +def produced_roles(producer: StageContract) -> set[str]: + """Roles a producer emits (from its ``writes`` keys); excludes ``unknown``.""" + keys = [*producer.writes.data_keys, *producer.writes.segment_data_keys] + return _spec_roles(producer, keys) - {"unknown"} + + +def reads_satisfied_by_role(consumer: StageContract, available_roles: set[str]) -> bool: + """Can ``consumer`` run given upstream-produced ``available_roles``? + + Matches by semantic role, not key-string equality, so a producer that + renamed its output key still satisfies a consumer that needs that role. + ``unknown`` is permissive (never blocks). + """ + avail = set(available_roles) | {"unknown"} + reads_keys = [*consumer.reads.data_keys, *consumer.reads.segment_data_keys] + if reads_keys and not _spec_roles(consumer, reads_keys).issubset(avail): + return False + if consumer.reads_one_of: + return any( + _spec_roles(consumer, [*opt.data_keys, *opt.segment_data_keys]).issubset(avail) + for opt in consumer.reads_one_of + ) + return True + + +# --------------------------------------------------------------------------- # +# Static checks (no execution) +# --------------------------------------------------------------------------- # +def _check_shape(c: StageContract, name: str) -> None: # noqa: C901 + assert c.cardinality in _VALID_CARDINALITY, f"{name}: invalid cardinality {c.cardinality!r}" + for opt in c.cardinality_options: + # cardinality_options are short flag names (e.g. "fan_out","nested") OR full cardinalities + assert isinstance(opt, str), f"{name}: bad cardinality option {opt!r}" + assert opt, f"{name}: bad cardinality option {opt!r}" + if c.iteration_key is not None: + assert c.cardinality in {"1:1 nested-list", "1:N fan-out", "N:1"}, ( + f"{name}: iteration_key set but cardinality is {c.cardinality!r}" + ) + # iteration_key must name something real: either a key the contract + # reads/writes, or a role-resolvable key value (fan-out stages iterate a + # list that is deliberately NOT re-emitted into children, so it may be + # absent from writes — but it must still resolve to a semantic role). + # Catches synthetic labels like the former 'speakers' that name nothing. + contract_keys: set[str] = set() + for spec in [c.reads, c.writes, *c.reads_one_of]: + contract_keys.update(spec.data_keys) + contract_keys.update(spec.segment_data_keys) + assert c.iteration_key in contract_keys or c.iteration_key in c.key_roles, ( + f"{name}: iteration_key {c.iteration_key!r} is neither a contract read/write " + f"key nor a role-resolvable key value — it names nothing an agent can find" + ) + for spec, label in [(c.reads, "reads"), (c.writes, "writes"), *[(s, "reads_one_of") for s in c.reads_one_of]]: + for a in spec.accepts: + assert a in _VALID_ACCEPTS, f"{name}: {label}.accepts has invalid form {a!r}" + for p in spec.produces: + assert p in _VALID_PRODUCES, f"{name}: {label}.produces has invalid form {p!r}" + for index, conditional in enumerate(c.conditional_writes): + label = f"conditional_writes[{index}]" + assert conditional.condition.strip(), f"{name}: {label}.condition must be non-empty" + assert conditional.value_origin in _VALID_WRITE_VALUE_ORIGINS, ( + f"{name}: {label}.value_origin has invalid value {conditional.value_origin!r}" + ) + assert conditional.writes.data_keys or conditional.writes.segment_data_keys or conditional.metadata_writes, ( + f"{name}: {label} must name at least one task, segment, or metadata key" + ) + for a in conditional.writes.accepts: + assert a in _VALID_ACCEPTS, f"{name}: {label}.writes.accepts has invalid form {a!r}" + for p in conditional.writes.produces: + assert p in _VALID_PRODUCES, f"{name}: {label}.writes.produces has invalid form {p!r}" + assert len(conditional.writes.data_keys) == len(set(conditional.writes.data_keys)), ( + f"{name}: duplicate {label}.writes.data_keys" + ) + assert len(conditional.writes.segment_data_keys) == len(set(conditional.writes.segment_data_keys)), ( + f"{name}: duplicate {label}.writes.segment_data_keys" + ) + assert len(conditional.metadata_writes) == len(set(conditional.metadata_writes)), ( + f"{name}: duplicate {label}.metadata_writes" + ) + assert all(isinstance(key, str) and key for key in conditional.metadata_writes), ( + f"{name}: {label}.metadata_writes must contain non-empty strings" + ) + # no duplicate keys within a single spec list + for spec, label in [(c.reads, "reads"), (c.writes, "writes")]: + assert len(spec.data_keys) == len(set(spec.data_keys)), f"{name}: duplicate {label}.data_keys" + + +def _check_roles(stage_or_cls: Any, c: StageContract, name: str) -> None: # noqa: ANN401 + for value, role in c.key_roles.items(): + assert role in _VALID_ROLES, f"{name}: key_roles[{value!r}] has invalid role {role!r}" + for p in c.params: + if p.role is not None: + assert p.role in _VALID_ROLES, f"{name}: param {p.name!r} has invalid role {p.role!r}" + # check #8: a *_key constructor field must have a KEY_ROLES entry or be + # explicitly allowlisted as internal (catches a forgotten role mapping). + if p.name.endswith("_key"): + stage_cls = stage_or_cls if isinstance(stage_or_cls, type) else type(stage_or_cls) + assert field_has_declared_role(p.name, stage_cls), ( + f"{name}: param {p.name!r} ends in '_key' but declares no role. Either give it " + "a shared role (KEY_ROLES in _roles.py) if another stage consumes it, or -- for " + "this stage's own bookkeeping -- declare it on the stage itself via " + "KEY_ROLE_OVERRIDES or INTERNAL_KEY_FIELDS" + ) + + +def _check_per_row_independence(c: StageContract, name: str) -> None: + """A narrowable source has to answer whether narrowing it is sound -- either way. + + ``False`` is a legitimate answer, not a violation: a source that cannot be narrowed soundly + says so per instance and ``delta.region`` stops there. Requiring ``True`` would force such a + source to lie or to drop the parameter. Silence is what is forbidden. + + The case that first motivated this -- ``CreateInitialManifestAudioFolderStage`` under a + bounded ``max_samples``, which truncates the sorted listing -- now declares ``True`` anyway + by an explicit product decision recorded at that declaration. The rule is unchanged: it was + never "must be False when narrowing is lossy", only "must not be silent". + + A companion rule ("``True`` contradicts ``N:1``") was removed as wrong: cardinality counts + TASKS, this is about row VALUES, and ``AudioToDocumentStage`` repacks tasks while leaving + values untouched. ``delta._TRACEABLE`` refuses ``N:1`` before the gate is read anyway. + """ + if any(p.name == _NARROWING_PARAM for p in c.params): + assert c.gates.per_row_independent is not None, ( + f"{name}: accepts {_NARROWING_PARAM!r} but leaves gates.per_row_independent undeclared -- " + "a source a delta run can narrow to a subset of files has to say whether the rows it " + "emits depend on which other files were present. False is a legitimate answer (the " + "delta then refuses to narrow it); silence is not" + ) + + +def _check_serialization(c: StageContract, name: str) -> None: + try: + json.dumps(c.to_dict()) + except (TypeError, ValueError) as e: # pragma: no cover - defensive + msg = f"{name}: contract.to_dict() is not JSON-serializable: {e}" + raise AssertionError(msg) from e + schema = to_json_schema(c.params) + assert schema.get("type") == "object", f"{name}: bad json schema" + assert "properties" in schema, f"{name}: bad json schema" + + +def _check_residency_accepts(stage: Any, c: StageContract, name: str) -> None: # noqa: ANN401 + """A residency-configurable stage must advertise exactly the forms it consumes. + + Derives the expected audio forms from the instance's ``input_residency`` and + asserts the contract's declared ``accepts`` (across ``reads`` + ``reads_one_of``) + match — catching the hand-typed "lying accepts" drift (a ``file``-mode instance + that still advertises ``waveform``). Skips when the contract carries no audio + ``accepts`` (e.g. an instance-free static contract with unresolved reads). + """ + residency = getattr(stage, "input_residency", None) + if residency is None: + return + declared = set(c.reads.accepts) | {a for opt in c.reads_one_of for a in opt.accepts} + if not declared: + return + expected = set(accepts_for_residency(residency)) + assert declared == expected, ( + f"{name}: declared accepts {sorted(declared)} != residency-derived {sorted(expected)} " + f"for input_residency={residency!r} — derive accepts from input_residency (lying/drifted accepts)" + ) + + +def assert_contract_wellformed(stage_or_cls: Any) -> StageContract: # noqa: ANN401 + """Static-only conformance: shape, roles, serialization. No execution. + + Accepts an instance (dynamic contract via ``build_contract``) or a class + (instance-free ``static_contract``). Returns the contract so callers can + reuse it. Safe to run across every stage with no fixtures. + """ + if isinstance(stage_or_cls, type): + c = static_contract(stage_or_cls) + name = stage_or_cls.__name__ + else: + c = build_contract(stage_or_cls) + name = type(stage_or_cls).__name__ + _check_shape(c, name) + _check_roles(stage_or_cls, c, name) + _check_per_row_independence(c, name) + _check_serialization(c, name) + _check_residency_accepts(stage_or_cls, c, name) + return c + + +# --------------------------------------------------------------------------- # +# Dynamic checks (execute the stage on a fixture) +# --------------------------------------------------------------------------- # +def _supports_batch(stage: Any) -> bool: # noqa: ANN401 + fn = getattr(stage, "supports_batch_processing", None) + try: + return bool(fn()) if callable(fn) else False + except Exception: # noqa: BLE001 + return False + + +def _normalize_results(out: Any) -> list[Any]: # noqa: ANN401 + if out is None: + return [] + if isinstance(out, list): + flat: list[Any] = [] + for item in out: + if item is None: + continue + if isinstance(item, list): + flat.extend(x for x in item if x is not None) + else: + flat.append(item) + return flat + return [out] + + +def _data_of(task: Any) -> dict[str, Any]: # noqa: ANN401 + data = getattr(task, "data", None) + return data if isinstance(data, dict) else {} + + +def _check_gpu_gate(stage: Any, c: StageContract, name: str) -> None: # noqa: ANN401 + """A stage that reserves GPU resources must not report that it needs none. + + One-sided deliberately. The reverse -- declaring the gate while reserving nothing -- is + legitimate: InferenceSortformerStage passes ``map_location="cuda"`` unconditionally, so it + needs a GPU whatever its ``resources`` say. Only the false negative is dangerous, because + it lets the planner put a GPU stage on a CPU-only worker. + """ + resources = getattr(stage, "resources", None) + if resources is not None and bool(getattr(resources, "requires_gpu", False)) and not c.gates.requires_gpu: + msg = ( + f"{name}: reserves GPU resources (gpus={getattr(resources, 'gpus', 0)}, " + f"gpu_memory_gb={getattr(resources, 'gpu_memory_gb', 0)}) but its contract reports " + f"requires_gpu=False. build_contract derives this from resources, so reaching here " + f"means the contract was built by hand or the derivation was bypassed." + ) + raise AssertionError(msg) + + +def assert_agent_ready( # noqa: C901, PLR0912, PLR0913 (complexity accepted: one linear checklist of independent conformance checks) + stage: Any, # noqa: ANN401 + fixture_factory: Callable[[], Any] | None = None, + *, + expected_cardinality: str | None = None, + available_keys: Iterable[str] | None = None, + segments_key: str | None = None, + ignore_new_keys: Iterable[str] = (), + run: bool = True, + setup: bool = False, +) -> StageContract: + """Assert a stage is agent-ready. Returns its (dynamic) contract. + + Always runs the static checks. When ``run`` and a ``fixture_factory`` are + given, also executes the stage and verifies declared writes appear, no + undeclared top-level keys leak, and cardinality matches the runtime shape. + + Args: + stage: A constructed stage instance. + fixture_factory: Returns a fresh input task (or batch) each call. + expected_cardinality: If given, assert the contract declares it. + available_keys: Upstream-available key values; asserts reads are + satisfiable by role. + segments_key: Resolved segments key, for checking segment-level writes. + ignore_new_keys: Extra top-level keys allowed in output (framework + bookkeeping) beyond declared writes. + run: Execute the stage (default True). + setup: Call ``stage.setup()`` before processing (default False; most + lightweight stages need no setup, heavy ones are pre-set-up/stubbed + by the caller). + """ + c = build_contract(stage) + name = type(stage).__name__ + _check_shape(c, name) + _check_roles(stage, c, name) + _check_per_row_independence(c, name) + _check_serialization(c, name) + _check_residency_accepts(stage, c, name) + _check_gpu_gate(stage, c, name) + + if expected_cardinality is not None: + assert c.cardinality == expected_cardinality, ( + f"{name}: cardinality {c.cardinality!r} != expected {expected_cardinality!r}" + ) + if available_keys is not None: + avail_roles = {c.key_roles.get(k, "unknown") for k in available_keys} + # also resolve via literal table for keys not in this stage's key_roles + from nemo_curator.stages.audio._agent._roles import role_for_value + + avail_roles |= {role_for_value(k) for k in available_keys} + assert reads_satisfied_by_role(c, avail_roles), ( + f"{name}: reads {c.reads.data_keys}/{[s.data_keys for s in c.reads_one_of]} " + f"not satisfied by available roles {avail_roles}" + ) + + if not run or fixture_factory is None: + return c + + if setup and hasattr(stage, "setup"): + stage.setup() + + task = fixture_factory() + batch_input = isinstance(task, list) + input_keys = set(_data_of(task[0] if batch_input else task)) + + if c.batch_only or _supports_batch(stage): + out = stage.process_batch(task if batch_input else [task]) + else: + out = stage.process(task) + results = _normalize_results(out) + + # (6) cardinality vs runtime shape + if c.cardinality == "1:N fan-out": + assert isinstance(out, list), f"{name}: fan-out must return a list" + elif c.cardinality in {"1:1", "1:1 nested-list"} and results: + assert len(results) == 1, f"{name}: {c.cardinality} produced {len(results)} tasks" + elif c.cardinality == "filter": + assert len(results) <= (len(task) if batch_input else 1), f"{name}: filter increased task count" + + # (3) declared writes appear; (4) no undeclared top-level keys (non-fanout) + if c.cardinality in {"1:1", "1:1 nested-list", "filter"} and results: + out_data = _data_of(results[0]) + for key in c.writes.data_keys: + assert key in out_data, f"{name}: declared write {key!r} missing from task.data" + for key in c.removes_keys: + assert key not in out_data, f"{name}: declared removes_keys {key!r} but it is still present in task.data" + declared = set(c.writes.data_keys) | set(ignore_new_keys) | input_keys + undeclared = set(out_data) - declared + assert not undeclared, f"{name}: undeclared new top-level keys {sorted(undeclared)} (add to writes.data_keys)" + # segment-level writes + seg_key = segments_key or c.iteration_key + if c.writes.segment_data_keys and seg_key and isinstance(out_data.get(seg_key), list) and out_data[seg_key]: + seg0 = out_data[seg_key][0] + if isinstance(seg0, dict): + for key in c.writes.segment_data_keys: + assert key in seg0, f"{name}: declared segment write {key!r} missing from segment dict" + return c + + +def assert_residency_consumption( + stage_factory: Callable[[str], Any], + *, + file_fixture: Callable[[], Any], + waveform_fixture: Callable[[], Any], + setup: bool = False, +) -> None: + """Prove a residency-configurable stage actually consumes each residency it advertises. + + Runs the stage in ``file`` and ``waveform`` modes on matching fixtures and + asserts it produced output — so a stage that declares an ``input_residency`` + choice its ``process()`` cannot actually consume fails CI. This is the + *dynamic* complement to the *static* :func:`_check_residency_accepts` drift + guard: the static check ties ``accepts`` to ``input_residency`` in the + contract; this one proves the code honors it. + + Intended for 1:1 / annotate stages (valid input -> non-empty output). Fan-out + stages that may legitimately return no items on a given fixture should assert + consumption differently. + + Args: + stage_factory: ``residency -> constructed stage`` (e.g. ``lambda r: MyStage(input_residency=r)``). + file_fixture: returns a task carrying only a file path. + waveform_fixture: returns a task carrying only an in-memory waveform + sample rate. + setup: call ``stage.setup()`` before processing (default False). + """ + for residency, fixture in (("file", file_fixture), ("waveform", waveform_fixture)): + stage = stage_factory(residency) + if setup and hasattr(stage, "setup"): + stage.setup() + task = fixture() + if _supports_batch(stage) or getattr(stage, "BATCH_ONLY", False): + out = stage.process_batch([task]) + else: + out = stage.process(task) + results = _normalize_results(out) + assert results, ( + f"{type(stage).__name__}: produced no output for input_residency={residency!r} — " + f"it advertises this residency but process() did not consume the matching input" + ) diff --git a/nemo_curator/stages/audio/_agent/_planning.py b/nemo_curator/stages/audio/_agent/_planning.py new file mode 100644 index 0000000000..d056128003 --- /dev/null +++ b/nemo_curator/stages/audio/_agent/_planning.py @@ -0,0 +1,686 @@ +# 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. + +"""Pipeline-level validation for agent-composed audio pipelines. + +``validate_pipeline([stageA, stageB, ...])`` walks an ordered list of configured +stages and checks they actually compose — each stage's required inputs must be +produced by an upstream stage or be present in the initial task. It also surfaces +resource-gate problems (GPU needed but none available). + +A read is satisfied by matching *either* the literal key it names or the semantic +*role* behind it. Both routes are needed: role matching tolerates a producer that +writes ``resampled_audio_filepath`` where the consumer reads ``audio_filepath``, +and key matching covers the reverse, where the names agree but the two sides file +that name under different roles. Requiring both would report breaks in pipelines +that run. + +Composites are expanded (see :mod:`nemo_curator.stages.audio._agent._composite`) so the +stages inside them are checked too — the requirements of the stages that do the +work, rather than the empty contract the composite advertises. Reads that fail +inside a composite are warnings, not errors, until the expansion has proven it +does not false-positive. A composite that cannot be expanded, or whose children +include something with no contract at all, falls back to being treated as opaque: +it is reported, and reads after it are no longer judged by role. + +Two levels of confidence, deliberately separated: + +* ``report.ok`` certifies a *role-level necessary condition* — every required + input role is available. This is rename-tolerant by design and is the gate. +* ``report.keys_ok`` adds the stronger *literal-key-identity* check: each + role-satisfied read's actual key *value* is produced upstream (or seeded). A + ``True`` ``ok`` with ``False`` ``keys_ok`` means the roles line up but a + producer key was renamed away from what the consumer reads — the pipeline + would validate yet yield zero rows at runtime. It is surfaced as a WARNING + (not an error) so that legitimate reads of source-manifest columns are not + false-rejected. + +This is advisory and read-only — it never executes a stage. ``ok`` is a +necessary, not sufficient, condition for a pipeline to run. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal + +from nemo_curator.stages.audio._agent._agent_registry import build_contract +from nemo_curator.stages.audio._agent._composite import expand_composites +from nemo_curator.stages.audio._agent._conformance import produced_roles, reads_satisfied_by_role +from nemo_curator.stages.audio._agent._roles import role_for_value + +if TYPE_CHECKING: + from nemo_curator.stages.audio._agent._agent_ready import StageContract + +Severity = Literal["error", "warning"] + +# Roles a typical audio task carries at the start (a manifest row with a file path). +_DEFAULT_INITIAL_ROLES: frozenset[str] = frozenset({"audio_filepath"}) +# Key VALUES a typical audio task carries at the start (the conventional path key). +_DEFAULT_INITIAL_KEYS: frozenset[str] = frozenset({"audio_filepath"}) + +# The role of the key a tensor producer parks its waveform under. Tracking the carrier +# rather than a bare "a tensor is resident" flag is what lets a stage that DROPS that key +# end the residency, instead of only a stage flagged ``sanitizes_output``. +_TENSOR_ROLE = "waveform" +# Stand-in for a producer that declares ``produces=["tensor"]`` without naming a +# waveform-roled key. Its residency is still tracked, but no key removal can match it, so +# only an explicit sanitizer clears it -- deliberately the pre-existing behaviour. +_UNNAMED_TENSOR = "" + + +@dataclass(frozen=True) +class PipelineIssue: + """A single problem found while validating a pipeline.""" + + stage_index: int + stage_name: str + severity: Severity + code: str + message: str + + +@dataclass(frozen=True) +class PipelineReport: + """Result of :func:`validate_pipeline`.""" + + issues: list[PipelineIssue] = field(default_factory=list) + produced_roles: set[str] = field(default_factory=set) # roles available after the last stage + produced_keys: set[str] = field(default_factory=set) # key VALUES available after the last stage + + @property + def ok(self) -> bool: + """True when there are no error-severity issues (role-level composability). + + This is a *necessary* condition, not a guarantee the pipeline runs — see + :attr:`keys_ok` for the stronger literal-key check. + """ + return not any(i.severity == "error" for i in self.issues) + + @property + def keys_ok(self) -> bool: + """True when no ``dangling_key`` warnings — every role-satisfied read's + actual key *value* is produced upstream or seeded. ``ok and keys_ok`` is + the strong signal that the pipeline will actually flow data end-to-end. + """ + return not any(i.code == "dangling_key" for i in self.issues) + + @property + def errors(self) -> list[PipelineIssue]: + return [i for i in self.issues if i.severity == "error"] + + @property + def warnings(self) -> list[PipelineIssue]: + return [i for i in self.issues if i.severity == "warning"] + + def summary(self) -> str: + if not self.issues: + return ( + "pipeline mechanically composable (all reads satisfied by role); " + "this does not certify intent or field meaning" + ) + lines = [f"{len(self.errors)} error(s), {len(self.warnings)} warning(s):"] + for i in self.issues: + lines.append(f" [{i.severity}] stage {i.stage_index} {i.stage_name}: {i.message}") + return "\n".join(lines) + + +def _required_roles(contract: StageContract) -> set[str]: + keys = [*contract.reads.data_keys, *contract.reads.segment_data_keys] + return {contract.key_roles.get(k, "unknown") for k in keys} + + +def _requirement_str(contract: StageContract, available: set[str]) -> str: + """Human-readable "what this stage needs" for an unsatisfied-reads message. + + Renders top-level ``reads`` (all required) and ``reads_one_of`` (any one), so a + stage whose reads live entirely in ``reads_one_of`` (e.g. a residency-derived + contract) no longer renders a misleading empty ``role(s) []``. + """ + missing = _required_roles(contract) - (available | {"unknown"}) + reqs: list[str] = [] + if missing: + reqs.append(f"role(s) {sorted(missing)}") + if contract.reads_one_of: + reqs.append(f"one of {[sorted(_roles_of(o, contract)) for o in contract.reads_one_of]}") + return "; ".join(reqs) or f"role(s) {sorted(_required_roles(contract))}" + + +def _write_key_values(contract: StageContract) -> set[str]: + """The literal key VALUES a stage writes (top-level + segment-level).""" + return {*contract.writes.data_keys, *contract.writes.segment_data_keys} + + +def _key_family(key: str) -> str: + """The trailing token of a key name -- ``diar_segments`` and ``segments`` share ``segments``. + + A crude but load-bearing notion of "these two keys hold the same KIND of thing". Producers + qualify a shared noun with a prefix (``diar_``, ``vad_``, ``pred_``), so the bare noun and + its qualified siblings are exactly the set a consumer might have meant. + """ + return key.rsplit("_", 1)[-1] + + +def _ambiguous_default_reads( + stage: Any, # noqa: ANN401 - any built stage + contract: StageContract, + available_keys: set[str], + key_producer: dict[str, str], +) -> list[tuple[str, str, list[tuple[str, str]]]]: + """``(key, attribute, rivals)`` for read keys left at a default while a sibling key exists. + + The failure this catches is silence. ``MergeAlignmentDiarizationStage`` documents itself as + merging into DIARIZATION segments, yet its ``segments_key`` defaults to ``"segments"`` -- + the key VAD writes. In a VAD+diarization pipeline both keys exist, so the read is satisfied + and every other check passes: transcripts get merged into the wrong segments and the output + is plausible, complete, and wrong. + + Deliberately narrow, because a warning nobody trusts is worse than none. It fires only when + the key is still at its CLASS DEFAULT (an explicit setting is a decision, not an accident), + the key IS available (an unavailable one is already reported as dangling), and some other + available key of the same family was written by a DIFFERENT upstream stage -- so a real + choice existed and was made by a default rather than by anyone. + """ + fields = getattr(type(stage), "__dataclass_fields__", {}) + reads = {*contract.reads.data_keys, *contract.reads.segment_data_keys} + found: list[tuple[str, str, list[tuple[str, str]]]] = [] + for attr, spec in fields.items(): + value = getattr(stage, attr, None) + if not (isinstance(value, str) and value in reads and value in available_keys and value == spec.default): + continue + rivals = sorted( + (k, key_producer[k]) + for k in available_keys + if k != value + # If the consumer reads both siblings (for example reference + # ``text`` and ASR ``pred_text`` for WER), they are independent + # operands rather than competing choices. + and k not in reads + and k in key_producer + and _key_family(k) == _key_family(value) + and key_producer[k] != key_producer.get(value) + ) + if rivals: + found.append((value, attr, rivals)) + return found + + +@dataclass(frozen=True) +class _Site: + """The stage being checked right now, and how to name it to whoever wrote the recipe.""" + + index: int + """Index of the RECIPE stage, so an inner stage points at something the caller can edit.""" + + name: str + stage: Any + composite: Any | None = None + """The recipe-level composite this was expanded from, when it is not a stage in its own right.""" + + +def _gate_issues( + site: _Site, + contract: StageContract, + available_gpus: float | None, + *, + tensor_resident: bool, +) -> list[PipelineIssue]: + """Environment/serialization gate problems for one stage. + + These reason about GPUs and serialization rather than roles, so they apply to every concrete + stage even downstream of a composite that hides its writes. + """ + out = [] + if available_gpus is not None and contract.gates.requires_gpu and available_gpus <= 0: + out.append( + PipelineIssue( + site.index, + site.name, + "warning", + "gpu_unavailable", + "declares requires_gpu but available_gpus <= 0", + ) + ) + # A resident tensor (e.g. a waveform) reaching a serialize-as-is sink (raw json.dumps) + # crashes at runtime. A sanitizing stage upstream clears the flag before we get here. + if contract.gates.requires_serializable_input and tensor_resident: + out.append( + PipelineIssue( + site.index, + site.name, + "error", + "tensor_into_sink", + "a resident tensor/audio blob from an upstream stage reaches this " + "serialize-as-JSON sink; it WILL fail at json.dumps — drop the tensor " + "upstream (e.g. keep_segment_waveform_in_task=False) or route through " + "a sanitizing stage before the sink", + ) + ) + return out + + +def _ambiguity_issues( + site: _Site, + contract: StageContract, + available_keys: set[str], + key_producer: dict[str, str], +) -> list[PipelineIssue]: + """``ambiguous_default_key`` warnings for this stage, naming who wrote each candidate.""" + out = [] + for key, attr, rivals in _ambiguous_default_reads(site.stage, contract, available_keys, key_producer): + others = ", ".join(f"{k!r} from {p}" for k, p in rivals) + out.append( + PipelineIssue( + site.index, + site.name, + "warning", + "ambiguous_default_key", + f"reads {key!r} (the default for {attr}), but upstream also produced {others}. " + f"The default silently picks {key!r} " + f"({key_producer.get(key, 'the source manifest')}); if you meant the other, " + f"set {attr} explicitly.", + ) + ) + return out + + +def _missing_read_keys(contract: StageContract, available_keys: set[str]) -> set[str]: + """Read key VALUES this stage wants that nothing upstream produced or seeded.""" + reads = {*contract.reads.data_keys, *contract.reads.segment_data_keys} + return {k for k in reads if k not in available_keys} + + +def _reads_satisfied_by_key(contract: StageContract, available_keys: set[str]) -> bool: + """Whether every read is met by the LITERAL key it names. + + A stage reads ``task.data[self.segments_key]`` at runtime -- a key string, never a role. So + a diarizer that writes ``diar_segments`` does satisfy a consumer configured to read + ``diar_segments``, even though the producer registers that key under the role + ``diar_segments`` while the consumer's contract calls the same slot ``segments``. Judging + that pairing only by role reports a break in a pipeline that runs, and the caller's options + are then to distrust the validator or to rename a key to appease it -- both worse than the + check not existing. + + Role matching stays as the rename-tolerant fallback for the opposite case, where the key + names differ but mean the same thing (a producer writing ``resampled_audio_filepath`` + satisfying a consumer reading ``audio_filepath``). A read is satisfied by either route. + """ + if {*contract.reads.data_keys, *contract.reads.segment_data_keys} - available_keys: + return False + if not contract.reads_one_of: + return True + return any(not ({*spec.data_keys, *spec.segment_data_keys} - available_keys) for spec in contract.reads_one_of) + + +def _forwarding_param(inner: Any, composite: Any, missing: set[str]) -> str | None: # noqa: ANN401 + """The composite parameter to set so an inner stage stops reading the wrong key. + + A caller who configured ``SplitASRAlignJoinStage`` has never heard of ``SplitLongAudioStage`` + and cannot configure it directly, so naming the inner stage alone leaves them stuck. The + remedy is always a parameter the composite forwards down, and it is identifiable rather than + guessable: the attribute exists on both classes and its current value IS the key that went + missing. Returns ``None`` when no such parameter exists, in which case the inner stage's + requirement genuinely cannot be reached from the recipe. + """ + inner_fields = getattr(type(inner), "__dataclass_fields__", {}) + composite_fields = getattr(type(composite), "__dataclass_fields__", {}) + for attr in inner_fields: + if attr not in composite_fields: + continue + value = getattr(inner, attr, None) + # ``missing`` holds key names, so only a string can ever match -- and testing anything + # else against a set hashes it, which raises TypeError on the list and dict parameters + # real stages carry (``file_extensions``, ``storage_options``). That exception escapes + # ``run_checks`` and kills the whole verb, so the recipe gets a traceback instead of a + # verdict over a remedy hint that was never going to apply. + if isinstance(value, str) and value in missing: + return attr + return None + + +def _unreadable_child(group: list[Any]) -> str | None: + """Why this composite's expansion cannot be reasoned about, or ``None`` if it can. + + A composite is only as legible as its least legible child. Composites routinely contain + plumbing that was never annotated for the agent -- ``ManifestReader`` expands through + ``FilePartitioningStage``, which has no ``describe()`` at all -- and a child whose reads and + writes are unknown leaves a hole in the role bookkeeping that makes every later stage's + verdict unsound. Blaming the caller for that with a hard error would fail pipelines that run + perfectly well, on the name of a stage they never wrote. So the composite reverts to being + opaque, which is exactly how it was treated before it could be expanded at all. + """ + for item in group: + if not item.composite_ref: + continue # a top-level stage that cannot describe itself is the caller's own error + try: + build_contract(item.stage) + except Exception as e: # noqa: BLE001 - any failure to describe means the same thing here + return f"{type(item.stage).__name__} does not describe its I/O ({type(e).__name__})" + return None + + +def _describes_itself(stage: Any) -> bool: # noqa: ANN401 - any child stage + """Whether a contract can be built for this stage at all.""" + try: + build_contract(stage) + except Exception: # noqa: BLE001 - any failure to describe means the same thing here + return False + return True + + +def _dangling_read_keys(contract: StageContract, available_keys: set[str]) -> set[str]: + """Read key VALUES whose role is known but whose exact value was not + produced upstream nor seeded — the renamed-producer dangle the role check misses. + + Covers primary ``reads`` plus a ``reads_one_of`` that offers a *single* + alternative: one option is not a choice, so its keys are as mandatory as a + primary read (this is how a residency-derived contract expresses + ``input_residency="file"``). A genuine multi-way ``reads_one_of`` is skipped — + the stage may legitimately take the other branch. Role-bearing keys only + (``unknown``/internal bookkeeping keys are excluded — a separate + value-identity check for those is tracked in the backlog). + """ + reads = [*contract.reads.data_keys, *contract.reads.segment_data_keys] + if len(contract.reads_one_of) == 1: + only = contract.reads_one_of[0] + reads += [*only.data_keys, *only.segment_data_keys] + dangling: set[str] = set() + for k in reads: + role = contract.key_roles.get(k, "unknown") + if role == "unknown": + continue + if k not in available_keys: + dangling.add(k) + return dangling + + +@dataclass +class _Walk: + """What the pipeline carries from one stage to the next while being validated.""" + + available: set[str] # roles produced so far + available_keys: set[str] # literal key VALUES produced so far + tensor_keys: set[str] = field(default_factory=set) + removed_roles: set[str] = field(default_factory=set) + key_producer: dict[str, str] = field(default_factory=dict) + past_composite: bool = False # an UNEXPANDABLE composite hid its writes; reads past it can't be judged + + +def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[PipelineIssue]: + """Whether this stage's reads are met, and how loudly to say so if not. + + Severity is graded by how sure we are, because a wrong hard error is worse than a wrong + warning: it stops the caller with no recourse and invites them to fake a value to get past + the gate rather than fix anything. A read that fails on a stage the caller wrote is certain, + so it is an error. A read that fails inside an expanded composite is reported as a warning + for now -- the expansion is new, and it earns the right to block only once it has been shown + not to false-positive on pipelines known to work. + """ + if reads_satisfied_by_role(contract, walk.available) or _reads_satisfied_by_key(contract, walk.available_keys): + if walk.past_composite: + return [] + out: list[PipelineIssue] = [] + dangling = _dangling_read_keys(contract, walk.available_keys) + if dangling: + out.append( + PipelineIssue( + site.index, + site.name, + "warning", + "dangling_key", + f"reads key(s) {sorted(dangling)} satisfied by role but not produced " + f"upstream under that key value nor seeded (renamed producer key?); " + f"available keys: {sorted(walk.available_keys)}", + ) + ) + out.extend(_ambiguity_issues(site, contract, walk.available_keys, walk.key_producer)) + return out + + if site.composite is not None: + composite_name = type(site.composite).__name__ + missing = _missing_read_keys(contract, walk.available_keys) + param = _forwarding_param(site.stage, site.composite, missing) + remedy = ( + f"set {param} on {composite_name} (it forwards the value to this inner stage)" + if param + else "produce the missing key upstream" + ) + return [ + PipelineIssue( + site.index, + site.name, + "warning", + "unsatisfied_reads_in_composite", + f"this stage runs inside {composite_name} and requires " + f"{_requirement_str(contract, walk.available)}" + + (f" (key(s) {sorted(missing)})" if missing else "") + + f", not produced upstream; {remedy}. Available keys: {sorted(walk.available_keys)}", + ) + ] + + if walk.past_composite: + return [ + PipelineIssue( + site.index, + site.name, + "warning", + "unsatisfied_reads_after_composite", + f"requires {_requirement_str(contract, walk.available)} " + f"not visibly produced — but an upstream composite hides its writes; " + f"decompose it to validate this read", + ) + ] + + needed = _required_roles(contract) | {r for o in contract.reads_one_of for r in _roles_of(o, contract)} + removed_hit = (needed & walk.removed_roles) - walk.available + if removed_hit: + return [ + PipelineIssue( + site.index, + site.name, + "error", + "key_removed_upstream", + f"reads role(s) {sorted(removed_hit)} that an upstream stage removed " + f"(removes_keys) and no stage re-produced; available so far: {sorted(walk.available)}", + ) + ] + return [ + PipelineIssue( + site.index, + site.name, + "error", + "unsatisfied_reads", + f"requires {_requirement_str(contract, walk.available)} " + f"not produced upstream; available so far: {sorted(walk.available)}", + ) + ] + + +def _advance(walk: _Walk, contract: StageContract, name: str) -> None: + """Fold one stage's writes, removals and tensor residency into the running state.""" + produced = produced_roles(contract) + walk.available |= produced + walk.removed_roles -= produced # a re-produced role is no longer "removed" + written = _write_key_values(contract) + # Most recent writer wins -- that is who a downstream reader would actually get. + walk.key_producer.update(dict.fromkeys(written, name)) + walk.available_keys |= written + for rk in contract.removes_keys: + walk.available_keys.discard(rk) + # Dropping the carrier ends the tensor residency as surely as sanitizing does. + walk.tensor_keys.discard(rk) + role = contract.key_roles.get(rk, role_for_value(rk)) + if ( + role != "unknown" + and role not in produced + and not any(role_for_value(k) == role for k in walk.available_keys) + ): + walk.available.discard(role) + walk.removed_roles.add(role) + if "tensor" in contract.writes.produces: + # The stage's OWN key_roles first, global names only as fallback. A custom + # ``waveform_key`` still declares its role in the contract, but the global lookup + # returned "unknown", so residency tracked ``_UNNAMED_TENSOR`` instead of the real + # carrier -- and a downstream stage dropping that carrier still looked resident, + # raising a spurious ``tensor_into_sink`` on a recipe that had cleaned up correctly. + carriers = {k for k in written if contract.key_roles.get(k, role_for_value(k)) == _TENSOR_ROLE} + walk.tensor_keys |= carriers or {_UNNAMED_TENSOR} + if contract.gates.sanitizes_output: + walk.tensor_keys.clear() + + +def validate_pipeline( # noqa: C901 + stages: list[Any], + *, + initial_roles: set[str] | None = None, + initial_keys: set[str] | None = None, + available_gpus: float | None = None, +) -> PipelineReport: + """Validate that an ordered list of configured stages composes. + + Args: + stages: Configured stage instances in execution order. + initial_roles: Semantic roles present in the input task. Defaults to + ``{"audio_filepath"}`` (a manifest row). Pass an explicit set when the + first stage is a source/reader or the input already carries waveforms. + initial_keys: Literal key VALUES present in the input task (e.g. the + columns of the source manifest: ``{"audio_filepath", "text"}``). + Defaults to ``{"audio_filepath"}``. Seeding this lets the + literal-key check (``keys_ok``) recognize reads satisfied by the + input rather than by an upstream producer. + available_gpus: If given, stages whose contract declares ``requires_gpu`` + while this is ``<= 0`` raise a warning. + + Returns: + A :class:`PipelineReport`. ``report.ok`` is True when no errors were + found (role-level); ``report.keys_ok`` additionally confirms literal-key + identity (see the class docstring). + """ + if initial_keys is not None: + seed_keys = set(initial_keys) + elif initial_roles is not None: + # Both seeds describe ONE task, so they cannot default independently: "no roles" does + # not also mean "the default columns". Seed only the roles that ARE their own key name -- + # roles and key values coincide for ``audio_filepath`` and diverge immediately after, so + # seeding ``transcript`` as a literal column invents a key the task does not carry. + seed_keys = {r for r in initial_roles if role_for_value(r) == r} + else: + seed_keys = set(_DEFAULT_INITIAL_KEYS) + walk = _Walk( + available=set(initial_roles) if initial_roles is not None else set(_DEFAULT_INITIAL_ROLES), + available_keys=seed_keys, + ) + expansion = expand_composites(stages) + leaves = expansion.by_recipe_index() + opaque = dict(expansion.opaque) + # A composite with one illegible child is not a composite nobody could open. Discarding the + # whole group left an eight-stage composite unchecked because one piece of plumbing lacks + # describe() -- ``ManifestReader`` expands through ``FilePartitioningStage``, so the reader + # starting most recipes contributed no keys at all. Legible siblings are kept; only the + # unknown part is treated as unknown, via ``past_composite``. + partly_opaque: dict[int, str] = {} + for index, group in list(leaves.items()): + reason = _unreadable_child(group) + if reason: + partly_opaque[index] = reason + leaves[index] = [item for item in group if _describes_itself(item.stage)] + issues: list[PipelineIssue] = [] + + for index, recipe_stage in enumerate(stages): + if index in expansion.unrunnable: + issues.append( + PipelineIssue( + index, + type(recipe_stage).__name__, + "error", + "composite_unrunnable", + f"the executor will refuse this stage: {expansion.unrunnable[index]}", + ) + ) + walk.past_composite = True + continue + if index in opaque: + issues.append( + PipelineIssue( + index, + type(recipe_stage).__name__, + "warning", + "composite", + f"composite stage — its data flow could not be resolved ({opaque[index]}), " + f"so reads after it cannot be judged by role", + ) + ) + walk.past_composite = True + continue + if index in partly_opaque: + issues.append( + PipelineIssue( + index, + type(recipe_stage).__name__, + "warning", + "composite", + f"composite stage — part of it is unreadable ({partly_opaque[index]}), " + f"so reads after it cannot be judged by role; its remaining stages are " + f"still checked", + ) + ) + # Set BEFORE its own legible children are walked, not after. One child's writes + # are unknown, and this composite's later children may be the very readers of + # them -- judging those reads against a key set that is missing exactly the + # unknown part is how a working pipeline gets failed on the name of a stage the + # caller never wrote. + walk.past_composite = True + + for item in leaves.get(index, []): + stage = item.stage + try: + contract = build_contract(stage) + except Exception as e: # noqa: BLE001 - a stage that can't describe itself is an error + issues.append(PipelineIssue(index, item.label, "error", "contract_error", f"describe() failed: {e}")) + continue + site = _Site( + index=index, + name=item.label if item.composite_ref else (contract.stage_id or type(stage).__name__), + stage=stage, + composite=recipe_stage if item.composite_ref else None, + ) + + if not contract.wrappable: + # It calls itself a composite yet arrived here unexpanded, so it is not a + # CompositeStage the expander could open. Its real I/O stays unknown and the + # pre-expansion caution applies: warn, and judge nothing downstream by role. + issues.append( + PipelineIssue( + site.index, + site.name, + "warning", + "composite", + "composite stage — decompose before validating its data flow", + ) + ) + walk.past_composite = True + continue + + issues.extend(_read_issues(walk, site, contract)) + # Serialization / GPU gates reason about the environment rather than about roles, so + # they run for every concrete stage even downstream of a composite nobody could expand. + issues.extend(_gate_issues(site, contract, available_gpus, tensor_resident=bool(walk.tensor_keys))) + _advance(walk, contract, site.name) + + return PipelineReport(issues=issues, produced_roles=walk.available, produced_keys=walk.available_keys) + + +def _roles_of(spec: Any, contract: StageContract) -> set[str]: # noqa: ANN401 + keys = [*spec.data_keys, *spec.segment_data_keys] + return {contract.key_roles.get(k, "unknown") for k in keys} diff --git a/nemo_curator/stages/audio/_agent/_residency.py b/nemo_curator/stages/audio/_agent/_residency.py new file mode 100644 index 0000000000..35fb2f0460 --- /dev/null +++ b/nemo_curator/stages/audio/_agent/_residency.py @@ -0,0 +1,334 @@ +# 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 hashlib +import os +import tempfile +from typing import TYPE_CHECKING, Any, Literal + +import soundfile as sf + +from nemo_curator.stages.audio._agent._agent_ready import AudioForm, ConditionalWrite, IOSpec +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 accepts_for_residency(residency: str) -> list[AudioForm]: + """Audio forms an instance actually consumes, given its ``input_residency``. + + The single source of truth a stage's ``describe()`` derives ``accepts`` from, + so a ``file``-mode instance can never advertise ``waveform`` (the drift / + "lying accepts" bug). ``auto`` accepts either; ``file``/``waveform`` accept + only that form. + """ + if residency == "waveform": + return ["waveform"] + if residency == "file": + return ["file"] + return ["file", "waveform"] # "auto" + + +def residency_read_specs( + input_residency: str, + *, + audio_filepath_key: str, + waveform_key: str = "waveform", + sample_rate_key: str = "sample_rate", +) -> list[IOSpec]: + """The residency-filtered audio read options for a stage's ``reads_one_of``. + + ``file`` -> ``[file spec]``; ``waveform`` -> ``[waveform spec]``; ``auto`` -> + ``[waveform spec, file spec]``. Keeps ``accepts`` **and** ``data_keys`` in + lockstep with ``input_residency`` so a stage can never advertise (or require) + a form it won't consume for its current setting — which lets the deterministic + role check enforce residency compatibility with no extra check. + """ + forms = accepts_for_residency(input_residency) + specs: list[IOSpec] = [] + if "waveform" in forms: + specs.append(IOSpec(data_keys=[waveform_key, sample_rate_key], accepts=["waveform"])) + if "file" in forms: + specs.append(IOSpec(data_keys=[audio_filepath_key], accepts=["file"])) + return specs + + +def scoped_audio_io_specs( # noqa: PLR0913 + input_residency: str, + *, + mode: Literal["task", "segments", "auto"], + audio_filepath_key: str, + waveform_key: str, + sample_rate_key: str, + segments_key: str, + output_keys: list[str], +) -> tuple[IOSpec, list[IOSpec], IOSpec]: + """Build mode-accurate reads/writes for task-or-nested audio stages. + + ``task`` exposes only top-level residency alternatives and outputs; + ``segments`` requires the top-level segment container while locating audio + and outputs inside each segment; and ``auto`` conservatively advertises the + complete alternatives for either runtime branch. + + This is contract assembly only. It does not select a runtime branch or + change a stage's processing behavior. + """ + task_reads = residency_read_specs( + input_residency, + audio_filepath_key=audio_filepath_key, + waveform_key=waveform_key, + sample_rate_key=sample_rate_key, + ) + segment_reads = [ + IOSpec( + data_keys=[segments_key] if mode == "auto" else [], + segment_data_keys=list(spec.data_keys), + accepts=list(spec.accepts), + ) + for spec in task_reads + ] + + if mode == "task": + return IOSpec(), task_reads, IOSpec(data_keys=list(output_keys)) + if mode == "segments": + return ( + IOSpec(data_keys=[segments_key]), + segment_reads, + IOSpec(segment_data_keys=list(output_keys)), + ) + return ( + IOSpec(), + [*task_reads, *segment_reads], + IOSpec(data_keys=list(output_keys), segment_data_keys=list(output_keys)), + ) + + +def scoped_audio_conditional_writes( + mode: Literal["task", "segments", "auto"], + *, + segments_key: str, + output_keys: list[str], + assignment_condition: str, +) -> list[ConditionalWrite]: + """Describe data-dependent writes for task-or-segment audio stages. + + ``assignment_condition`` is stage-authored factual prose for the common + success path that actually assigns the advertised keys. The helper adds + configured scope/auto-branch context without interpreting the stage or + changing execution. + """ + conditional: list[ConditionalWrite] = [] + if mode in {"task", "auto"}: + branch = ( + "task mode is configured" + if mode == "task" + else f"'{segments_key}' is absent, so the task-level branch runs" + ) + conditional.append( + ConditionalWrite( + writes=IOSpec(data_keys=list(output_keys)), + condition=f"{branch}; {assignment_condition}", + ) + ) + if mode in {"segments", "auto"}: + branch = ( + "segments mode is configured and an individual segment exists" + if mode == "segments" + else (f"'{segments_key}' is present, so the per-segment branch runs, and an individual segment exists") + ) + conditional.append( + ConditionalWrite( + writes=IOSpec(segment_data_keys=list(output_keys)), + condition=f"{branch}; {assignment_condition}", + ) + ) + return conditional + + +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 write_audio_stable( + waveform: Any, # noqa: ANN401 - a torch tensor or numpy array, same as _as_soundfile_array takes + sample_rate: int, + *, + output_dir: str | None, + stem: str = "audio", + tag: str = "", +) -> str: + """Write a waveform under a name derived from the audio, and return the path. + + Stages writing in-memory audio used to reach for ``tempfile.mkstemp``, whose contract is a + name that has never existed -- right for scratch, wrong for a deliverable: each re-run wrote + a second full set of files beside the first instead of replacing it, leaving every prior run + orphaned in the directory. Naming a file after its own bytes fixes that by construction. + + ``output_dir`` of None keeps the mkstemp behaviour, because that is the system temp dir: a + predictable name there would be world-readable in a shared directory, and unguessable-and- + private is worth more than de-duplication for audio nothing is going to look for by name. + """ + arr = _as_soundfile_array(waveform) + if output_dir is None: + fd, path = tempfile.mkstemp(prefix=f"{stem}{f'_{tag}' if tag else ''}_", suffix=".wav") + os.close(fd) + sf.write(path, arr, int(sample_rate)) + return path + + os.makedirs(output_dir, exist_ok=True) + digest = hashlib.sha256(arr.tobytes()) + digest.update(f"|{int(sample_rate)}".encode()) + path = os.path.join(output_dir, f"{stem}{f'_{tag}' if tag else ''}_{digest.hexdigest()[:16]}.wav") + # Write beside the target and rename, so a killed or concurrent writer cannot leave a + # half-written file at a name the next run treats as finished. + staged_fd, staged = tempfile.mkstemp(prefix=".", suffix=".wav", dir=output_dir) + os.close(staged_fd) + try: + sf.write(staged, arr, int(sample_rate)) + os.replace(staged, path) + except BaseException: + with contextlib.suppress(OSError): + os.unlink(staged) + raise + return path + + +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/_agent/_roles.py b/nemo_curator/stages/audio/_agent/_roles.py new file mode 100644 index 0000000000..52d104d1b1 --- /dev/null +++ b/nemo_curator/stages/audio/_agent/_roles.py @@ -0,0 +1,207 @@ +# 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. + +"""Semantic-role mapping for agent-ready audio stages. + +Under the config-knobs-only standardization, every ``task.data`` key a stage +reads or writes is an agent-configurable ``*_key`` constructor field. The key's +*value* can therefore be renamed per pipeline, but the *field name* (e.g. +``score_key``) is invariant — it is what the producer/consumer code is written +around. We map that invariant field name to a stable semantic +:data:`~nemo_curator.stages.audio._agent._agent_ready.Role`, so an agent can chain a +producer's output to a consumer's input by *role* even when key values differ. + +``KEY_ROLES`` covers cross-stage composable keys. ``INTERNAL_KEY_FIELDS`` lists +``*_key`` constructor fields that are intentionally stage-internal (no +cross-stage role; chained by value-equality within a tightly coupled pair). +``LITERAL_KEY_ROLES`` maps hard-coded key *values* emitted by producers that do +not expose a ``*_key`` field (e.g. ``ManifestReaderStage`` writes ``"audio_filepath"``). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Mapping + + from nemo_curator.stages.audio._agent._agent_ready import Role + +# field name (e.g. "score_key") -> semantic role +KEY_ROLES: dict[str, Role] = { + # audio file path (all variants resolve to the same role) + "audio_filepath_key": "audio_filepath", + "filepath_key": "audio_filepath", + "resampled_audio_filepath_key": "audio_filepath", + "original_audio_filepath_key": "audio_filepath", + "output_audio_filepath_key": "audio_filepath", + "swift_audio_filepath_key": "audio_filepath", + # in-memory audio + "waveform_key": "waveform", + "sample_rate_key": "sample_rate", + "audio_sample_rate_key": "sample_rate", + # duration / size + "duration_key": "duration", + "duration_ms_key": "duration", + "total_duration_sec_key": "duration", + "speaking_duration_key": "duration", + "num_samples_key": "num_samples", + # segment lists + "segments_key": "segments", + "diar_segments_key": "diar_segments", + "vad_segments_key": "vad_segments", + "overlap_segments_key": "overlap_segments", + # per-segment timing + "start_key": "start", + "end_key": "end", + "start_ms_key": "start_ms", + "end_ms_key": "end_ms", + "original_start_ms_key": "start_ms", + "original_end_ms_key": "end_ms", + "segment_num_key": "segment_num", + "original_file_key": "original_file", + "audio_item_id_key": "item_id", + # text + "text_key": "text", + "hypothesis_text_key": "text", + "pred_text_key": "pred_text", + "reference_text_key": "reference_text", + "words_key": "words", + "alignment_key": "alignment", + # speaker + "speaker_key": "speaker_id", + "speaker_id_key": "speaker_id", + "num_speakers_key": "num_speakers", + # quality / metric scores (incl. SIGMOS sub-scores and WER) + "score_key": "score", + "metrics_key": "metrics", + "stats_key": "metrics", + "prediction_key": "prediction", + "wer_key": "score", + "sig_key": "score", + "ovrl_key": "score", + "noise_key": "score", + "disc_key": "score", + "reverb_key": "score", + "loud_key": "score", + "col_key": "score", # SIGMOS coloration sub-score — a peer of its 6 siblings + # windows (ALM snippet planning) + "windows_key": "windows", + "filtered_windows_key": "windows", +} + +# ``*_key`` constructor fields with no cross-stage role (chained by value within +# a tightly coupled stage pair, or fully generic / user-defined). Listed so the +# conformance check does not flag them as "forgot a KEY_ROLES entry". +INTERNAL_KEY_FIELDS: frozenset[str] = frozenset( + { + # SplitLongAudio <-> JoinSplitAudioMetadata bookkeeping (value-matched pair) + "split_filepaths_key", + "split_metadata_key", + "split_offsets_key", + "split_timestamps_key", + "mappings_key", + # generic / user-defined targets + "input_value_key", # PreserveByValueStage: compares an arbitrary user key + "items_key", # PreserveByValueConditionsStage: caller-chosen one-level list + "output_key", # ITN/Chinese: caller-chosen output key + "original_key", # preserved prior value + "sort_key", + "cache_key", + "oldest_key", + # bookkeeping counters / flags + "num_segments_key", + "is_mono_key", + "truncation_events_key", + } +) + +# Hard-coded key *values* emitted by producers that lack a ``*_key`` field. +LITERAL_KEY_ROLES: dict[str, Role] = { + "audio_filepath": "audio_filepath", + "waveform": "waveform", + "sample_rate": "sample_rate", + "duration": "duration", + "segments": "segments", + "diar_segments": "diar_segments", + "vad_segments": "vad_segments", + "text": "text", + "text_ref": "reference_text", # ComputeWERStage's documented reference default + "pred_text": "pred_text", + "words": "words", + "alignment": "alignment", +} + + +def role_for_field(field_name: str) -> Role: + """Return the semantic role for a ``*_key`` constructor field name. + + Unmapped field names return ``"unknown"`` (composition falls back to + value-equality and is never blocked). + """ + return KEY_ROLES.get(field_name, "unknown") + + +def role_for_value(key_value: str, *, field_name: str | None = None) -> Role: + """Resolve a role from a key's resolved value, preferring its field name. + + ``field_name`` (authoritative) is the producer/consumer's ``*_key`` field; + when absent (hard-coded producer output) we fall back to a literal-default + table keyed on the value itself. + """ + if field_name is not None: + role = KEY_ROLES.get(field_name) + if role is not None: + return role + return LITERAL_KEY_ROLES.get(key_value, "unknown") + + +def field_has_declared_role(field_name: str, stage_cls: type | None = None) -> bool: + """True if a ``*_key`` field has a role or is declared internal bookkeeping. + + Used by the conformance harness to catch a newly added ``*_key`` field that forgot a + :data:`KEY_ROLES` entry. + + A stage may satisfy this itself, via ``KEY_ROLE_OVERRIDES`` (this field means an + existing role) or ``INTERNAL_KEY_FIELDS`` (this field is my own bookkeeping and chains + with nothing). That is what keeps adding a stage from meaning editing this module. The + shared tables stay authoritative for keys that cross stages, because a role is the + vocabulary two stages connect through -- a privately invented one would compose with + nothing while still passing the check. + """ + if field_name in KEY_ROLES or field_name in INTERNAL_KEY_FIELDS: + return True + if stage_cls is None: + return False + return field_name in role_overrides_for(stage_cls) or field_name in internal_key_fields_for(stage_cls) + + +def internal_key_fields_for(stage_cls: type) -> frozenset[str]: + """A stage's own bookkeeping ``*_key`` fields, UNIONED across its bases. + + ``getattr`` alone returns only the most-derived declaration, so a subclass that declares + one internal field of its own shadows every field its parent declared -- and the parent's + fields, still inherited and still bookkeeping, start failing the conformance check as + though someone had forgotten a role for them. The subclass author's fix is then to + re-list fields they did not write, which is how a shared table gets copied downwards. + """ + fields: set[str] = set() + for base in getattr(stage_cls, "__mro__", (stage_cls,)): + fields |= set(base.__dict__.get("INTERNAL_KEY_FIELDS") or ()) + return frozenset(fields) + + +def role_overrides_for(stage_cls: type) -> Mapping[str, Role]: + """Per-stage ``KEY_ROLE_OVERRIDES`` (consulted before the shared table).""" + return getattr(stage_cls, "KEY_ROLE_OVERRIDES", {}) or {} diff --git a/nemo_curator/stages/audio/agent.py b/nemo_curator/stages/audio/agent.py new file mode 100644 index 0000000000..8675959f04 --- /dev/null +++ b/nemo_curator/stages/audio/agent.py @@ -0,0 +1,105 @@ +# 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. + +"""Public entry point for agent-driven audio pipeline construction. + +Everything an agent (or agent-tool layer) needs to discover, compose, +configure, and validate audio pipelines, re-exported from the private +foundation modules. Import from here, not from the underscore modules. + +The intended loop:: + + from nemo_curator.stages.audio import agent + + # 1. DISCOVER — what stages exist and what they do + names = agent.list_agent_ready_stages() + catalog = agent.audio_stage_catalog() # contracts + params_schema per stage + + # 2. PLAN — how roles connect stages + agent.find_producers("segments") # -> stages that can supply a role + agent.find_consumers("pred_text") # [] from find_producers => unproducible + + # 3. CONFIGURE + BUILD + cls = agent.get_agent_ready_stage_class("UTMOSFilterStage") + stage = cls(mos_threshold=3.5) # params_schema documents the knobs + + # 4. VALIDATE (before ever running) + report = agent.validate_pipeline([stage, ...], initial_keys={"audio_filepath", "text"}) + report.ok # role-level composability (necessary condition) + report.keys_ok # literal keys connect (mechanical flow, not intent meaning) + report.summary() # human/agent-readable issues incl. dangling_key / tensor_into_sink + +``describe_stage(name)`` returns a contract marked +``static_params_and_hints``; pass an instance to ``build_contract`` for the +configured dynamic contract with resolved I/O/key values. +``StageContract.to_dict()`` is JSON-safe by construction. Open-ended intent fit +is reviewed by the host LLM over ``audio_agent.validate(...).semantic_review``. +""" + +from __future__ import annotations + +from nemo_curator.stages.audio._agent._agent_ready import ( + ConditionalWrite, + Gates, + IOSpec, + ParamSpec, + Role, + SizeEnvelope, + StageContract, + to_json_schema, +) +from nemo_curator.stages.audio._agent._agent_registry import build_contract, static_contract +from nemo_curator.stages.audio._agent._catalog import ( + audio_stage_catalog, + catalog_as_json, + describe_stage, + find_consumers, + find_producers, + get_agent_ready_stage_class, + list_agent_ready_stages, + role_index, +) +from nemo_curator.stages.audio._agent._conformance import ( + assert_contract_wellformed, + produced_roles, + reads_satisfied_by_role, +) +from nemo_curator.stages.audio._agent._planning import PipelineIssue, PipelineReport, validate_pipeline + +__all__ = [ + "ConditionalWrite", + "Gates", + "IOSpec", + "ParamSpec", + "PipelineIssue", + "PipelineReport", + "Role", + "SizeEnvelope", + "StageContract", + "assert_contract_wellformed", + "audio_stage_catalog", + "build_contract", + "catalog_as_json", + "describe_stage", + "find_consumers", + "find_producers", + "get_agent_ready_stage_class", + "list_agent_ready_stages", + "produced_roles", + "reads_satisfied_by_role", + "role_index", + "static_contract", + "to_json_schema", + "validate_pipeline", +] diff --git a/nemo_curator/stages/audio/common.py b/nemo_curator/stages/audio/common.py index 46dca6cd51..2cf9f9fd35 100644 --- a/nemo_curator/stages/audio/common.py +++ b/nemo_curator/stages/audio/common.py @@ -13,11 +13,15 @@ # limitations under the License. import json +import math import os import time +import uuid +from collections.abc import Mapping from dataclasses import dataclass, field from operator import eq, ge, gt, le, lt, ne -from typing import Any +from typing import Any, ClassVar, Literal +from urllib.parse import urlsplit import soundfile import torch @@ -25,10 +29,13 @@ from loguru import logger from nemo_curator.backends.base import NodeInfo, WorkerMetadata +from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, IOSpec, Role, StageContract, StaticHints from nemo_curator.stages.base import CompositeStage, ProcessingStage from nemo_curator.stages.file_partitioning import FilePartitioningStage from nemo_curator.tasks import AudioTask, EmptyTask, FileGroupTask +_VALUE_OPERATORS = {"lt": lt, "le": le, "eq": eq, "ne": ne, "ge": ge, "gt": gt} + def get_audio_duration(audio_filepath: str) -> float: """Get the duration of the audio file in seconds.""" @@ -41,18 +48,25 @@ 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*. Args: audio_filepath_key: Key to get path to wav file. duration_key: Key to put audio duration. + waveform_key: Key for an in-memory waveform tensor. + sample_rate_key: Key for the in-memory waveform sample rate. + input_residency: Which input to use — "file" (audio_filepath only; default, + unchanged), "waveform" (in-memory only), or "auto" (waveform first, file fallback). """ name: str = "GetAudioDurationStage" audio_filepath_key: str = "audio_filepath" duration_key: str = "duration" + waveform_key: str = "waveform" + sample_rate_key: str = "sample_rate" + input_residency: Literal["file", "waveform", "auto"] = "file" def setup(self, _worker_metadata: WorkerMetadata | None = None) -> None: import soundfile @@ -60,21 +74,63 @@ def setup(self, _worker_metadata: WorkerMetadata | None = None) -> None: self._soundfile = soundfile def inputs(self) -> tuple[list[str], list[str]]: + if self.input_residency == "waveform": + return [], [self.waveform_key, self.sample_rate_key] return [], [self.audio_filepath_key] def outputs(self) -> tuple[list[str], list[str]]: return [], [self.duration_key] + def describe(self) -> StageContract: + # Lazy import avoids a module-level cycle (_residency imports from common). + from nemo_curator.stages.audio._agent._residency import residency_read_specs + + return StageContract( + reads_one_of=residency_read_specs( + self.input_residency, + audio_filepath_key=self.audio_filepath_key, + waveform_key=self.waveform_key, + sample_rate_key=self.sample_rate_key, + ), + writes=IOSpec(data_keys=[self.duration_key]), + gates=Gates(per_row_independent=True), + ) + + def validate_input(self, task: AudioTask) -> bool: + """Require the audio source implied by ``input_residency`` (default: file).""" + data = task.data + has_waveform = data.get(self.waveform_key) is not None and data.get(self.sample_rate_key) is not None + has_file = self.audio_filepath_key in data + if self.input_residency == "waveform": + return has_waveform + if self.input_residency == "file": + return has_file + return has_waveform or has_file # auto + + def _resolve_duration(self, data: dict[str, Any]) -> float: + """Duration from an in-memory waveform (samples / sample_rate) or the file. + + Default (``input_residency="file"``) reads the file exactly as before. + """ + if self.input_residency != "file": + waveform = data.get(self.waveform_key) + sr = data.get(self.sample_rate_key) + if waveform is not None and sr is not None and int(sr) > 0: + return ensure_waveform_2d(waveform).shape[-1] / float(sr) + if self.input_residency == "waveform": + logger.warning(f"Missing '{self.waveform_key}'+'{self.sample_rate_key}' (input_residency='waveform')") + return -1.0 + return get_audio_duration(data[self.audio_filepath_key]) + def process(self, task: AudioTask) -> AudioTask: t0 = time.perf_counter() - audio_filepath = task.data[self.audio_filepath_key] - duration = get_audio_duration(audio_filepath) + duration = self._resolve_duration(task.data) task.data[self.duration_key] = duration self._log_metrics({"process_time": time.perf_counter() - t0, "duration": max(duration, 0.0)}) 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 @@ -84,23 +140,30 @@ class PreserveByValueStage(ProcessingStage[AudioTask, AudioTask]): input_value_key: The field in the dataset entries to evaluate. target_value: The value to compare with. operator: Comparison operator (lt, le, eq, ne, ge, gt). + missing_value_policy: ``"error"`` (default) preserves the historical + validation error for a missing input key; ``"drop"`` removes that row. """ name: str = "PreserveByValueStage" + BATCH_ONLY = True # process() raises; only process_batch is implemented (agent-discovery hint) def __init__( self, input_value_key: str, - target_value: int | str, + target_value: float | str, operator: str = "eq", + missing_value_policy: Literal["error", "drop"] = "error", ): self.input_value_key = input_value_key self.target_value = target_value - ops = {"lt": lt, "le": le, "eq": eq, "ne": ne, "ge": ge, "gt": gt} - if operator not in ops: - msg = f"Operator must be one of: {', '.join(ops)}" + if operator not in _VALUE_OPERATORS: + msg = f"Operator must be one of: {', '.join(_VALUE_OPERATORS)}" raise ValueError(msg) - self.operator = ops[operator] + if missing_value_policy not in {"error", "drop"}: + msg = "missing_value_policy must be 'error' or 'drop'" + raise ValueError(msg) + self.operator = _VALUE_OPERATORS[operator] + self.missing_value_policy = missing_value_policy def inputs(self) -> tuple[list[str], list[str]]: return [], [self.input_value_key] @@ -108,6 +171,16 @@ 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", + # Compares one row against a fixed target value, so batching changes throughput + # rather than the verdict -- no row's fate depends on the rows beside it. + gates=Gates(per_row_independent=True), + ) + def process(self, task: AudioTask) -> AudioTask | None: msg = "PreserveByValueStage only supports process_batch" raise NotImplementedError(msg) @@ -117,6 +190,8 @@ def process_batch(self, tasks: list[AudioTask]) -> list[AudioTask]: results = [] for task in tasks: if not self.validate_input(task): + if self.missing_value_policy == "drop": + continue msg = f"Task {task!s} failed validation for stage {self}" raise ValueError(msg) if self.operator(task.data[self.input_value_key], self.target_value): @@ -132,30 +207,341 @@ def process_batch(self, tasks: list[AudioTask]) -> list[AudioTask]: return results +@dataclass(frozen=True) +class _ValueCondition: + input_value_key: str + target_value: float | str | bool + operator: str + + +class PreserveByValueConditionsStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): + """Keep a row or nested child according to flat scalar conditions. + + ``conditions`` accepts either a list of + ``{"input_value_key", "target_value", "operator"}`` mappings or a mapping + from input key to ``{"target_value", "operator"}``. A scalar mapping value + is shorthand for an equality condition. Conditions use AND semantics by + default; ``condition_logic="or"`` keeps an item when at least one condition + passes. Missing condition keys always fail closed before logic is applied. + + Args: + conditions: Non-empty list or mapping of scalar comparisons. + missing_value_policy: ``"error"`` raises when any condition key is + absent; ``"drop"`` removes rows with any absent condition key. + items_key: Optional top-level ``task.data`` key containing a list of + mapping-like children. When set, conditions filter that list in + place of filtering top-level rows. This is a single-level operation + and never recursively descends into child values. + drop_parent_if_empty: In nested mode, drop the parent AudioTask when no + child survives. Ignored when ``items_key`` is ``None``. + condition_logic: ``"and"`` (default) requires every condition to pass; + ``"or"`` requires at least one condition to pass. + """ + + name: str = "PreserveByValueConditionsStage" + BATCH_ONLY = True + + def __init__( + self, + conditions: list[Mapping[str, Any]] | Mapping[str, Any], + missing_value_policy: Literal["error", "drop"] = "error", + items_key: str | None = None, + drop_parent_if_empty: bool = True, + condition_logic: Literal["and", "or"] = "and", + ): + if missing_value_policy not in {"error", "drop"}: + msg = "missing_value_policy must be 'error' or 'drop'" + raise ValueError(msg) + if items_key is not None and (not isinstance(items_key, str) or not items_key): + msg = "items_key must be None or a non-empty string" + raise ValueError(msg) + if not isinstance(drop_parent_if_empty, bool): + msg = "drop_parent_if_empty must be a boolean" + raise TypeError(msg) + if condition_logic not in {"and", "or"}: + msg = "condition_logic must be 'and' or 'or'" + raise ValueError(msg) + self.conditions = conditions + self.missing_value_policy = missing_value_policy + self.items_key = items_key + self.drop_parent_if_empty = drop_parent_if_empty + self.condition_logic = condition_logic + self._conditions = self._normalize_conditions(conditions) + + @staticmethod + def _normalize_conditions( # noqa: C901, PLR0912 - one validation branch per accepted condition shape + conditions: list[Mapping[str, Any]] | Mapping[str, Any], + ) -> tuple[_ValueCondition, ...]: + raw_conditions: list[Mapping[str, Any]] + if isinstance(conditions, Mapping): + raw_conditions = [] + for key, value in conditions.items(): + if isinstance(value, Mapping): + raw_conditions.append( + { + "input_value_key": key, + "target_value": value.get("target_value"), + "operator": value.get("operator", "eq"), + "_has_target": "target_value" in value, + } + ) + else: + raw_conditions.append( + { + "input_value_key": key, + "target_value": value, + "operator": "eq", + "_has_target": True, + } + ) + elif isinstance(conditions, list): + raw_conditions = list(conditions) + else: + msg = "conditions must be a non-empty list or mapping" + raise TypeError(msg) + if not raw_conditions: + msg = "conditions must contain at least one scalar comparison" + raise ValueError(msg) + + normalized: list[_ValueCondition] = [] + for index, condition in enumerate(raw_conditions): + if not isinstance(condition, Mapping): + msg = f"conditions[{index}] must be a mapping" + raise TypeError(msg) + key = condition.get("input_value_key") + if not isinstance(key, str) or not key: + msg = f"conditions[{index}].input_value_key must be a non-empty string" + raise ValueError(msg) + if not condition.get("_has_target", "target_value" in condition): + msg = f"conditions[{index}] must define target_value" + raise ValueError(msg) + target = condition.get("target_value") + if isinstance(target, float) and not math.isfinite(target): + msg = f"conditions[{index}].target_value must be finite" + raise ValueError(msg) + if not isinstance(target, (bool, int, float, str)): + msg = f"conditions[{index}].target_value must be a JSON scalar" + raise TypeError(msg) + operator = condition.get("operator", "eq") + if operator not in _VALUE_OPERATORS: + msg = f"conditions[{index}].operator must be one of: {', '.join(_VALUE_OPERATORS)}" + raise ValueError(msg) + normalized.append( + _ValueCondition( + input_value_key=key, + target_value=target, + operator=str(operator), + ) + ) + return tuple(normalized) + + @property + def normalized_conditions(self) -> tuple[dict[str, Any], ...]: + """Canonical conditions for deterministic planning and comparison.""" + return tuple( + { + "input_value_key": condition.input_value_key, + "target_value": condition.target_value, + "operator": condition.operator, + } + for condition in self._conditions + ) + + def inputs(self) -> tuple[list[str], list[str]]: + if self.items_key is not None: + return [], [self.items_key] + return [], [condition.input_value_key for condition in self._conditions] + + def outputs(self) -> tuple[list[str], list[str]]: + if self.items_key is not None: + return [], [self.items_key] + return [], [condition.input_value_key for condition in self._conditions] + + def describe(self) -> StageContract: + logic = self.condition_logic.upper() + if self.items_key is not None: + return StageContract( + reads=IOSpec(data_keys=[self.items_key]), + writes=IOSpec(data_keys=[self.items_key]), + cardinality="filter" if self.drop_parent_if_empty else "1:1 nested-list", + cardinality_options=["filter", "1:1 nested-list"], + iteration_key=None if self.drop_parent_if_empty else self.items_key, + description=( + f"Filter mapping-like children in task.data[{self.items_key!r}] with " + f"one-level {logic} conditions; nested values are not traversed." + ), + gates=Gates(per_row_independent=True), + ) + keys = [condition.input_value_key for condition in self._conditions] + return StageContract( + reads=IOSpec(data_keys=keys), + writes=IOSpec(data_keys=keys), + cardinality="filter", + description=f"Filter top-level AudioTask rows with {logic} conditions.", + gates=Gates(per_row_independent=True), + ) + + def process(self, task: AudioTask) -> AudioTask | None: + msg = "PreserveByValueConditionsStage only supports process_batch" + raise NotImplementedError(msg) + + def _nested_items( + self, + task: AudioTask, + items_key: str, + ) -> list[Mapping[str, Any]]: + """Return and structurally validate the configured one-level child list.""" + if items_key not in task.data: + msg = f"Task {task!s} is missing nested items_key {items_key!r}" + raise ValueError(msg) + items = task.data[items_key] + if not isinstance(items, list): + msg = f"Task {task!s} nested items_key {items_key!r} must contain a list, got {type(items).__name__}" + raise TypeError(msg) + for index, item in enumerate(items): + if not isinstance(item, Mapping): + msg = ( + f"Task {task!s} nested items_key {items_key!r} " + f"child {index} must be mapping-like, got {type(item).__name__}" + ) + raise TypeError(msg) + return items + + def _nested_item_passes( + self, + task: AudioTask, + index: int, + item: Mapping[str, Any], + items_key: str, + ) -> bool: + """Apply configured conditions to one validated direct child.""" + for condition in self._conditions: + if condition.input_value_key not in item: + if self.missing_value_policy == "drop": + return False + msg = ( + f"Task {task!s} nested items_key {items_key!r} child " + f"{index} is missing condition key {condition.input_value_key!r}" + ) + raise ValueError(msg) + results = ( + _VALUE_OPERATORS[condition.operator]( + item[condition.input_value_key], + condition.target_value, + ) + for condition in self._conditions + ) + return all(results) if self.condition_logic == "and" else any(results) + + def _process_nested_batch( + self, + tasks: list[AudioTask], + items_key: str, + ) -> list[AudioTask]: + """Filter direct children, replacing only the configured list field.""" + results: list[AudioTask] = [] + for task in tasks: + items = self._nested_items(task, items_key) + survivors = [ + item for index, item in enumerate(items) if self._nested_item_passes(task, index, item, items_key) + ] + task.data[items_key] = survivors + if survivors or not self.drop_parent_if_empty: + results.append(task) + return results + + def process_batch(self, tasks: list[AudioTask]) -> list[AudioTask]: + t0 = time.perf_counter() + if self.items_key is not None: + results = self._process_nested_batch(tasks, self.items_key) + else: + results = [] + for task in tasks: + for condition in self._conditions: + if condition.input_value_key not in task.data: + if self.missing_value_policy == "drop": + break + msg = f"Task {task!s} failed validation for stage {self}" + raise ValueError(msg) + else: + condition_results = ( + _VALUE_OPERATORS[condition.operator]( + task.data[condition.input_value_key], + condition.target_value, + ) + for condition in self._conditions + ) + keep = all(condition_results) if self.condition_logic == "and" else any(condition_results) + if keep: + results.append(task) + self._log_metrics( + { + "process_time": time.perf_counter() - t0, + "input_count": len(tasks), + "output_count": len(results), + "filtered_count": len(tasks) - len(results), + } + ) + return results + + +def _row_names_file(row: dict[str, Any], key: str, wanted: set[str]) -> bool: + """Whether a manifest row's audio path is one of ``wanted`` (absolute-path comparison).""" + value = row.get(key) + return isinstance(value, str) and os.path.abspath(os.path.expanduser(value)) in wanted + + @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. Supports local and cloud paths (S3, GCS). + + Args: + include_files: Emit only rows whose audio path (under ``include_files_key``) is one of + these files, comparing absolute paths. ``None`` -- the default -- reads every row. + It restricts the same reader over the same manifest rather than pointing a delta + run at a filtered copy, so the rows a partial run emits are the rows a full run + would have emitted. + include_files_key: Which row column holds that path. """ name: str = "manifest_reader_stage" + include_files: list[str] | None = None + include_files_key: str = "audio_filepath" + # Declared statically as well as in describe(): a narrowable source has to be readable as + # safe-to-narrow without constructing it, which is how the instance-free conformance sweep + # sees it. + AGENT_STATIC: ClassVar[StaticHints] = StaticHints( + gates=Gates(lifecycle_side_effects=True, per_row_independent=True) + ) + # It points at the column holding a row's audio path, which is an existing role rather + # than a new one -- a filter comparing something else would not be filtering by file. + KEY_ROLE_OVERRIDES: ClassVar[Mapping[str, Role]] = {"include_files_key": "audio_filepath"} def process(self, task: FileGroupTask) -> list[AudioTask]: t0 = time.perf_counter() paths = task.data results: list[AudioTask] = [] count = 0 + wanted = ( + None + if self.include_files is None + else {os.path.abspath(os.path.expanduser(p)) for p in self.include_files} + ) for manifest in paths: fs, resolved = url_to_fs(manifest) with fs.open(resolved, "r", encoding="utf-8") as f: for line in f: if line.strip(): + row = json.loads(line.strip()) + if wanted is not None and not _row_names_file(row, self.include_files_key, wanted): + continue results.append( AudioTask( dataset_name=task.dataset_name, - data=json.loads(line.strip()), + data=row, _metadata=task._metadata, _stage_perf=list(task._stage_perf), ) @@ -174,9 +560,16 @@ def process(self, task: FileGroupTask) -> list[AudioTask]: 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, per_row_independent=True), + ) + @dataclass -class ManifestReader(CompositeStage[EmptyTask, AudioTask]): +class ManifestReader(AgentReady, CompositeStage[EmptyTask, AudioTask]): """Composite stage for reading JSONL manifests. Decomposes into: @@ -189,6 +582,8 @@ class ManifestReader(CompositeStage[EmptyTask, AudioTask]): blocksize: Target size per partition (e.g., "100MB"). Ignored if files_per_partition is set. file_extensions: File extensions to filter. Defaults to [".jsonl", ".json"]. storage_options: Storage options for cloud paths (S3, GCS credentials, endpoints). + include_files: Read only the rows naming these audio files (see ``ManifestReaderStage``). + include_files_key: Which row column holds the audio path. """ manifest_path: str | list[str] @@ -197,6 +592,10 @@ class ManifestReader(CompositeStage[EmptyTask, AudioTask]): blocksize: int | str | None = None file_extensions: list[str] = field(default_factory=lambda: [".jsonl", ".json"]) storage_options: dict[str, Any] | None = None + include_files: list[str] | None = None + include_files_key: str = "audio_filepath" + AGENT_STATIC: ClassVar[StaticHints] = StaticHints(gates=Gates(per_row_independent=True)) + KEY_ROLE_OVERRIDES: ClassVar[Mapping[str, Role]] = {"include_files_key": "audio_filepath"} def __post_init__(self) -> None: super().__init__() @@ -213,7 +612,10 @@ def decompose(self) -> list[ProcessingStage]: file_extensions=self.file_extensions, storage_options=self.storage_options, ), - ManifestReaderStage(), + ManifestReaderStage( + include_files=self.include_files, + include_files_key=self.include_files_key, + ), ] def get_description(self) -> str: @@ -224,9 +626,137 @@ 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, + gates=Gates(per_row_independent=True), + ) + @dataclass -class ManifestWriterStage(ProcessingStage[AudioTask, AudioTask]): +class CreateInitialManifestAudioFolderStage(AgentReady, ProcessingStage[EmptyTask, AudioTask]): + """Create an initial manifest from any local folder of audio files. + + Recursively scans ``data_dir`` for audio files and emits one AudioTask per file with its + path under ``audio_filepath_key`` (plus a filename-derived ``audio_item_id``). A generic, + dataset-agnostic source: no download, no transcripts, and no dataset-specific filename + parsing -- unlike ``CreateInitialManifest{ReadSpeech,Fleurs}Stage``. Use it to start a + pipeline from a plain folder of WAV/FLAC/MP3/... when there is no JSONL manifest (use + ``ManifestReader`` when a manifest already exists). + + Args: + data_dir: Local folder to scan for audio files. + extensions: Audio file extensions to include (case-insensitive). + recursive: Recurse into subfolders (default True). + max_samples: Maximum number of files to include (-1 for all). + include_files: Process only these files (absolute paths), skipping the rest of the + folder. ``None`` -- the default -- means the whole folder, exactly as before. + Restricting the file list rather than swapping in a different source stage is what + lets a delta run over new files produce rows identical to a full run's. + """ + + data_dir: str + extensions: list[str] = field(default_factory=lambda: [".wav", ".flac", ".mp3", ".ogg", ".opus", ".m4a"]) + recursive: bool = True + max_samples: int = -1 + include_files: list[str] | None = None + audio_filepath_key: str = "audio_filepath" + audio_item_id_key: str = "audio_item_id" + name: str = "CreateInitialManifestAudioFolder" + batch_size: int = 1 + # See ManifestReaderStage: the narrowing claim has to survive being read off the class. + AGENT_STATIC: ClassVar[StaticHints] = StaticHints(gates=Gates(per_row_independent=True)) + + def __post_init__(self) -> None: + super().__init__() + if not self.data_dir: + msg = "data_dir is required for CreateInitialManifestAudioFolderStage" + raise ValueError(msg) + + def inputs(self) -> tuple[list[str], list[str]]: + return [], [] + + def outputs(self) -> tuple[list[str], list[str]]: + return [], [self.audio_filepath_key, self.audio_item_id_key] + + def describe(self) -> StageContract: + return StageContract( + # No ``produces``: the audio already exists on disk, this stage only points at it + # (unlike the dataset CreateInitialManifest*Stage sources, which download and write). + writes=IOSpec(data_keys=[self.audio_filepath_key, self.audio_item_id_key]), + cardinality="1:N fan-out", + # Scans existing files; one task per file, and a row says nothing about its + # neighbours. Declared True unconditionally BY DECISION: under a bounded + # ``max_samples`` the SORTED listing is truncated, so a delta can admit files a full + # run would not have. Accepted rather than cost every bounded run its reuse -- not + # an oversight to "fix" back. + gates=Gates(per_row_independent=True), + ) + + def ray_stage_spec(self) -> dict[str, Any]: + return {"is_fanout_stage": True} + + def num_workers(self) -> int | None: + return 1 + + def _collect_audio_files(self) -> list[str]: + exts = tuple((e if e.startswith(".") else f".{e}").lower() for e in self.extensions) + if not os.path.isdir(self.data_dir): + logger.error(f"[{self.name}] data_dir not found: {self.data_dir}") + return [] + found: list[str] = [] + if self.recursive: + for root, _dirs, files in os.walk(self.data_dir): + found.extend(os.path.join(root, f) for f in files if f.lower().endswith(exts)) + else: + found = [ + os.path.join(self.data_dir, f) + for f in os.listdir(self.data_dir) + if f.lower().endswith(exts) and os.path.isfile(os.path.join(self.data_dir, f)) + ] + if self.include_files is not None: + wanted = {os.path.abspath(os.path.expanduser(p)) for p in self.include_files} + found = [p for p in found if os.path.abspath(p) in wanted] + missing = wanted - {os.path.abspath(p) for p in found} + if missing: + # Named rather than silently skipped: a caller that asked for specific files and + # got fewer would otherwise read the short result as "those files held nothing". + logger.warning( + f"[{self.name}] include_files named {len(missing)} file(s) not found under {self.data_dir}" + ) + return sorted(found) + + def process(self, _: EmptyTask) -> list[AudioTask]: + """Emit one AudioTask per audio file found under ``data_dir``.""" + paths = self._collect_audio_files() + if self.max_samples is not None and self.max_samples >= 0: + paths = paths[: self.max_samples] + if not paths: + logger.warning(f"[{self.name}] no audio files {self.extensions} under {self.data_dir}") + return [] + tasks: list[AudioTask] = [] + for path in paths: + abspath = os.path.abspath(path) + # Relpath, not basename: ``recursive`` defaults True and speaker-per-folder is the + # standard layout, so a basename id gives spk1/utt1.wav and spk2/utt1.wav the same + # id -- and downstream that id becomes an output filename. A flat corpus is + # unaffected. Not injective: a flat ``spk1__utt1.wav`` still aliases spk1/utt1.wav. + rel = os.path.relpath(abspath, os.path.abspath(self.data_dir)) + item_id = os.path.splitext(rel)[0].replace(os.sep, "__") + tasks.append( + AudioTask( + dataset_name="local-audio-folder", + data={self.audio_filepath_key: abspath, self.audio_item_id_key: item_id}, + filepath_key=self.audio_filepath_key, + ) + ) + logger.info(f"[{self.name}] created {len(tasks)} AudioTask(s) from {self.data_dir}") + return tasks + + +@dataclass +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) @@ -287,6 +817,303 @@ 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, + output_path_params=["output_path"], + lifecycle_side_effects=True, + # Serializes task.data as-is via json.dumps; a resident tensor + # (e.g. a waveform) will crash it. Stop carrying the tensor + # before this AudioTask sink, or convert to a DocumentBatch and + # use DocumentBatchJsonlWriterStage instead. + requires_serializable_input=True, + # Appends each row as it arrives; a row's line is its own contents. + per_row_independent=True, + ), + ) + + +@dataclass +class ManifestCheckpointStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): + """Persist a reusable, metadata-only AudioTask boundary as JSONL. + + This stage is an intermediate checkpoint, not a terminal user deliverable. It + serializes complete ``task.data`` rows and passes each task downstream with + its dataset name, metadata, and performance records preserved. Waveform + tensors and other non-JSON values are illegal at this boundary. + + The checkpoint is local-only and single-worker. ``setup()`` exclusively + reserves a new destination and refuses to overwrite any existing file. + + Args: + output_path: Required local destination for checkpoint JSONL. + retention_sec: Advisory retention in seconds. Defaults to 0, meaning + user-managed with no automatic expiry; must be non-negative. + owner: Ownership recorded with the checkpoint policy. ``"user"`` is the + conservative default; ``"project"`` means the project operator owns + retention. Neither value enables automatic deletion. + planning_provenance: Internal marker for a reusable-pipeline candidate. + Such recipes require exact-hash approval and authoritative smoke. + """ + + REUSABLE_PIPELINE_PROVENANCE: ClassVar[str] = "reusable_pipeline_v1" + + output_path: str + retention_sec: int = 0 + owner: Literal["user", "project"] = "user" + planning_provenance: Literal["reusable_pipeline_v1"] | None = None + name: str = "manifest_checkpoint" + + AGENT_STATIC: ClassVar[StaticHints] = StaticHints( + gates=Gates( + writes_to_disk=True, + output_path_params=["output_path"], + lifecycle_side_effects=True, + requires_serializable_input=True, + per_row_independent=True, + ), + description="Persist a complete metadata checkpoint without ending the pipeline", + ) + + def __post_init__(self) -> None: + if not self.output_path: + msg = "output_path is required for ManifestCheckpointStage" + raise ValueError(msg) + if urlsplit(self.output_path).scheme: + msg = "ManifestCheckpointStage output_path must be a plain local path, not a URI" + raise ValueError(msg) + if isinstance(self.retention_sec, bool) or not isinstance(self.retention_sec, int) or self.retention_sec < 0: + msg = "retention_sec must be a non-negative integer" + raise ValueError(msg) + if self.owner not in {"user", "project"}: + msg = "owner must be 'user' or 'project'" + raise ValueError(msg) + if self.planning_provenance not in {None, self.REUSABLE_PIPELINE_PROVENANCE}: + msg = f"planning_provenance must be None or {self.REUSABLE_PIPELINE_PROVENANCE!r}" + raise ValueError(msg) + self._reservation_owned = False + self._reservation_identity: tuple[int, int, int] | None = None + self._reservation_token = uuid.uuid4().hex + self._checkpoint_rows_written = 0 + self._checkpoint_bytes_written = 0 + + def _resolve_output(self) -> None: + self._fs, self._path = url_to_fs(self.output_path) + parent_dir = "/".join(self._path.split("/")[:-1]) + if parent_dir: + self._fs.makedirs(parent_dir, exist_ok=True) + + def setup(self, _worker_metadata: WorkerMetadata | None = None) -> None: + """Atomically reserve a new checkpoint without overwriting retained work.""" + self._resolve_output() + if self._reservation_owned: + msg = ( + "ManifestCheckpointStage setup was called again while this stage " + "still owns an incomplete reservation; reset_for_retry() is required" + ) + raise RuntimeError(msg) + marker_path = f"{self._path}._COMPLETE" + if self._fs.exists(marker_path): + msg = ( + "ManifestCheckpointStage refuses to create a checkpoint beside an " + f"existing completion marker at {self.output_path!r}" + ) + raise FileExistsError(msg) + owner_path = self._retry_owner_path() + try: + with self._fs.open(owner_path, "xb") as owner: + owner.write(json.dumps({"token": self._reservation_token}).encode("utf-8")) + except FileExistsError as exc: + msg = ( + "ManifestCheckpointStage refuses to replace an existing retry " + f"ownership record at {self.output_path!r}" + ) + raise FileExistsError(msg) from exc + try: + with self._fs.open(self._path, "xb"): + pass + except FileExistsError as exc: + self._remove_retry_owner_if_owned() + msg = f"ManifestCheckpointStage refuses to overwrite an existing checkpoint at {self.output_path!r}" + raise FileExistsError(msg) from exc + except OSError: + self._remove_retry_owner_if_owned() + raise + try: + stat = os.stat(self._path) + except OSError: + # The exclusive create above belongs to this stage instance. If its + # identity cannot be recorded, remove only that empty reservation + # and fail rather than creating an unprovable retry owner. + self._fs.rm(self._path) + self._remove_retry_owner_if_owned() + raise + self._reservation_identity = (stat.st_dev, stat.st_ino, stat.st_ctime_ns) + self._reservation_owned = True + self._checkpoint_rows_written = 0 + self._checkpoint_bytes_written = 0 + try: + self._write_retry_owner(stat) + except OSError: + # The path is still the empty reservation whose identity was just + # recorded. Remove it and its token rather than leave state that a + # driver-side retry cannot prove. + current = os.stat(self._path) + current_identity = ( + current.st_dev, + current.st_ino, + current.st_ctime_ns, + ) + if current_identity == self._reservation_identity and current.st_size == 0: + self._fs.rm(self._path) + self._remove_retry_owner_if_owned() + self._reset_retry_state() + raise + logger.info(f"ManifestCheckpointStage: writing metadata to {self.output_path}") + + def reset_for_retry(self) -> None: + """Remove only this instance's incomplete reservation before an automatic retry.""" + self._resolve_output() + marker_path = f"{self._path}._COMPLETE" + if self._fs.exists(marker_path): + msg = ( + "ManifestCheckpointStage refuses retry reset because a completion " + f"marker exists at {self.output_path!r}" + ) + raise FileExistsError(msg) + owner = self._read_retry_owner() + if owner is None or owner.get("token") != self._reservation_token: + if self._fs.exists(self._path) or owner is not None: + msg = ( + "ManifestCheckpointStage refuses retry reset of a checkpoint " + f"it did not reserve for this run at {self.output_path!r}" + ) + raise FileExistsError(msg) + self._reset_retry_state() + return + if self._fs.exists(self._path): + try: + stat = os.stat(self._path) + except OSError as exc: + msg = f"ManifestCheckpointStage could not verify its retry reservation at {self.output_path!r}" + raise RuntimeError(msg) from exc + identity = (stat.st_dev, stat.st_ino, stat.st_ctime_ns) + recorded_identity = ( + owner.get("st_dev"), + owner.get("st_ino"), + owner.get("st_ctime_ns"), + ) + if identity != recorded_identity or stat.st_size != owner.get("st_size"): + msg = ( + "ManifestCheckpointStage refuses retry reset because the checkpoint " + f"at {self.output_path!r} is no longer its exact reservation" + ) + raise FileExistsError(msg) + self._fs.rm(self._path) + self._remove_retry_owner_if_owned() + self._reset_retry_state() + + def release_retry_reservation(self) -> None: + """Remove this run's ownership sidecar after successful execution.""" + self._resolve_output() + try: + self._remove_retry_owner_if_owned() + except OSError as exc: + logger.warning( + "ManifestCheckpointStage could not remove its successful retry " + f"ownership record at {self.output_path!r}: {exc}" + ) + self._reservation_owned = False + self._reservation_identity = None + + def _retry_owner_path(self) -> str: + return f"{self._path}._RETRY_OWNER" + + def _read_retry_owner(self) -> dict[str, Any] | None: + owner_path = self._retry_owner_path() + if not self._fs.exists(owner_path): + return None + try: + with self._fs.open(owner_path, "rb") as owner: + value = json.loads(owner.read().decode("utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError): + return {} + return value if isinstance(value, dict) else {} + + def _write_retry_owner(self, stat: os.stat_result) -> None: + payload = { + "token": self._reservation_token, + "st_dev": stat.st_dev, + "st_ino": stat.st_ino, + "st_ctime_ns": stat.st_ctime_ns, + "st_size": stat.st_size, + } + with self._fs.open(self._retry_owner_path(), "wb") as owner: + owner.write(json.dumps(payload, sort_keys=True).encode("utf-8")) + + def _remove_retry_owner_if_owned(self) -> None: + owner = self._read_retry_owner() + if owner is None or owner.get("token") != self._reservation_token: + return + self._fs.rm(self._retry_owner_path()) + + def _reset_retry_state(self) -> None: + self._reservation_owned = False + self._reservation_identity = None + self._checkpoint_rows_written = 0 + self._checkpoint_bytes_written = 0 + self._custom_metrics = {} + + def setup_on_node( + self, + _node_info: NodeInfo | None = None, + _worker_metadata: WorkerMetadata | None = None, + ) -> None: + """Ensure the local parent directory exists without truncating.""" + self._resolve_output() + + def process(self, task: AudioTask) -> AudioTask: + if not self._reservation_owned: + msg = "ManifestCheckpointStage cannot write without an owned setup reservation" + raise RuntimeError(msg) + t0 = time.perf_counter() + row = (json.dumps(task.data, ensure_ascii=False) + "\n").encode("utf-8") + with self._fs.open(self._path, "ab") as f: + f.write(row) + self._checkpoint_rows_written += 1 + self._checkpoint_bytes_written += len(row) + stat = os.stat(self._path) + self._reservation_identity = (stat.st_dev, stat.st_ino, stat.st_ctime_ns) + self._write_retry_owner(stat) + self._log_metrics( + { + "process_time": time.perf_counter() - t0, + "checkpoint_rows_written": 1, + "checkpoint_bytes_written": len(row), + } + ) + return AudioTask( + dataset_name=task.dataset_name, + data=task.data, + _metadata=task._metadata, + _stage_perf=list(task._stage_perf), + ) + + def num_workers(self) -> int | None: + return 1 + + def describe(self) -> StageContract: + return StageContract( + gates=Gates( + writes_to_disk=True, + output_path_params=["output_path"], + lifecycle_side_effects=True, + requires_serializable_input=True, + per_row_independent=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.""" @@ -324,6 +1151,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._agent._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/__init__.py b/nemo_curator/stages/audio/preprocessing/__init__.py index c6df4cdbc9..d7f12f65be 100755 --- a/nemo_curator/stages/audio/preprocessing/__init__.py +++ b/nemo_curator/stages/audio/preprocessing/__init__.py @@ -16,18 +16,37 @@ Audio preprocessing stages. These stages prepare audio for further processing: -- MonoConversionStage: Convert to mono and verify sample rate +- ChannelCountStage: Record, select on, or convert the channel count (never resamples) +- SampleRateFilterStage: Keep only acceptable sample rates, recording each (header-only read) +- MonoConversionStage: Convert to mono and verify sample rate in one step - SegmentConcatenationStage: Concatenate multiple audio segments +Channel policy and rate policy are separate stages so a pipeline can set one without the +other, and each says whether it measures, selects or converts rather than doing two at once. + Example: from nemo_curator.pipeline import Pipeline - from nemo_curator.stages.audio.preprocessing import MonoConversionStage + from nemo_curator.stages.audio.preprocessing import ( + ChannelCountStage, + SampleRateFilterStage, + ) pipeline = Pipeline(name="preprocessing_pipeline") - pipeline.add_stage(MonoConversionStage(output_sample_rate=48000)) + # 48 kHz mono, by selection: nothing is rewritten, the rest is dropped. + pipeline.add_stage(SampleRateFilterStage(allowed_sample_rates=[48000])) + pipeline.add_stage(ChannelCountStage(action="filter", allowed_channels=[1])) + # ...or by conversion: every row is kept and made mono. + pipeline.add_stage(ChannelCountStage(action="convert", target_channels=1)) """ +from .channel_count import ChannelCountStage from .concatenation import SegmentConcatenationStage from .mono_conversion import MonoConversionStage +from .sample_rate_filter import SampleRateFilterStage -__all__ = ["MonoConversionStage", "SegmentConcatenationStage"] +__all__ = [ + "ChannelCountStage", + "MonoConversionStage", + "SampleRateFilterStage", + "SegmentConcatenationStage", +] diff --git a/nemo_curator/stages/audio/preprocessing/channel_count.py b/nemo_curator/stages/audio/preprocessing/channel_count.py new file mode 100644 index 0000000000..a0d76632a6 --- /dev/null +++ b/nemo_curator/stages/audio/preprocessing/channel_count.py @@ -0,0 +1,519 @@ +# 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. + +""" +Audio channel-count stage: record it, select on it, or change it. + +``action`` picks which of the three, in the vocabulary ``UTMOSFilterStage``, ``BandFilterStage`` +and ``SIGMOSFilterStage`` already use for measure-and-keep versus measure-and-drop, extended +with the one thing a channel count can do that a quality score cannot: be changed. + +Sample rate, format and file layout are never touched, so a pipeline sets its rate policy +separately -- ``SampleRateFilterStage`` to select rates, ``ResampleAudioStage`` to convert +them -- or sets none at all. + +Example: + from nemo_curator.pipeline import Pipeline + from nemo_curator.stages.audio.preprocessing import ChannelCountStage + + pipeline = Pipeline(name="audio_pipeline") + pipeline.add_stage(ChannelCountStage()) # record the count + pipeline.add_stage(ChannelCountStage(action="filter", allowed_channels=[1])) # keep mono only + pipeline.add_stage(ChannelCountStage(action="convert", target_channels=1)) # make it mono +""" + +import os +import tempfile +from dataclasses import dataclass, field, fields +from typing import ClassVar, Literal + +import soundfile as sf +import torch +from loguru import logger + +from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, IOSpec, StageContract +from nemo_curator.stages.audio._agent._residency import ( + InputResidency, + accepts_for_residency, + produce_audio_filepath, + residency_read_specs, + resolve_audio, + write_audio_stable, +) +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 + +ChannelAction = Literal["annotate", "filter", "convert"] + + +@dataclass +class ChannelCountStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): + """ + Record, select on, or change the number of audio channels. + + ============ ================================================================= + action behaviour + ============ ================================================================= + ``annotate`` records ``num_channels`` on every row and keeps them all (default) + ``filter`` records it, then drops the rows whose count is not allowed + ``convert`` brings the audio to ``target_channels``, in memory + ============ ================================================================= + + Selecting and converting are deliberately separate actions rather than one parameter that + means both. They do opposite things to a corpus -- ``convert`` changes every row and keeps + all of them, ``filter`` changes nothing and keeps a subset -- so a single knob spelling both + makes "mono" ambiguous between "make it mono" and "keep only what already is". Parameters + belonging to an action you did not choose are refused at construction rather than ignored. + + ``annotate`` and ``filter`` never decode. A channel count sits in the file header, which + ``soundfile.info`` reads without touching samples, so putting either in front of a decoding + stage costs almost nothing and spares the rejected rows entirely. ``convert`` must decode, + because it has to rewrite the samples. + + **What ``num_channels`` means depends on the action.** Under ``annotate``/``filter`` it is + the count OBSERVED in the source audio. Under ``convert`` it is the count RESULTING from the + conversion. ``sample_rate`` carries exactly this trap between ``SampleRateFilterStage`` (a + measurement) and ``ResampleAudioStage`` (a target); reading a pipeline's final + ``num_channels`` without knowing which stage last wrote it inverts its meaning. + + Under ``convert``, what happens depends on how many channels the input actually has: + + ============ ================== ================================================== + input target behaviour + ============ ================== ================================================== + ``N`` ``N`` passed through unchanged + ``N > 1`` ``1`` averaged into one channel (standard mono downmix) + ``1`` ``T > 1`` duplicated into ``T`` identical channels + ``N > T > 1`` ``T`` REFUSED -- the row is dropped + ============ ================== ================================================== + + That refusal is deliberate. A correct downmix to more than one channel needs ITU-R BS.775 + coefficients *and* the file's channel order, and a bare ``(channels, samples)`` tensor + carries neither -- WAV channel order comes from the file's channel mask, which is gone by + the time the audio is a tensor. Averaging 5.1 into two channels does not produce stereo, it + produces a phase-smeared mix that sounds plausible and is wrong. ``ResampleAudioStage`` + drives ffmpeg, which does know layouts, so that is the honest tool for those conversions. + Downmix to ``1`` is not the same problem: averaging every channel together IS what mono + means, so it needs no layout knowledge. + + Nothing here ever resamples. When a rate must actually change, ``ResampleAudioStage`` + converts it via ffmpeg. + + Args: + action: "annotate" records num_channels and keeps every row (default); "filter" also + drops rows whose count is not allowed; "convert" brings the audio to + target_channels. Parameters of the other actions are refused, not ignored. + allowed_channels: Counts to keep, e.g. [1] for mono only (action="filter"). None = no + constraint from this parameter. + min_channels: Lowest acceptable count, inclusive (action="filter"). None = unbounded. + max_channels: Highest acceptable count, inclusive (action="filter"). None = unbounded. + target_channels: Channel count to produce (action="convert"). None means 1 (mono). + audio_filepath_key: Key in data dict for the audio file path. + waveform_key: Key in data dict for the in-memory waveform tensor. + sample_rate_key: Key in data dict for the waveform sample rate. + num_channels_key: Key where the channel count is written -- observed under + annotate/filter, resulting under convert. + duration_key: Key where the audio duration in seconds is written (convert only). + output_audio_filepath_key: Key where the written WAV path is stored + (action="convert", 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 converted waveform and sample + rate in task.data for downstream in-memory consumers (convert only). + write_to_disk: If True, write the converted audio to a WAV file (action="convert"). + Without output_dir this writes to the system temp dir and nothing cleans it up; in + multi-node runs point output_dir at shared storage. + update_audio_filepath: If True, repoint audio_filepath_key at the written file and + preserve the original under original_audio_filepath_key. + output_dir: Directory for written audio (action="convert", write_to_disk=True only). + """ + + action: ChannelAction = "annotate" + + allowed_channels: list[int] | None = None + min_channels: int | None = None + max_channels: int | None = None + + target_channels: int | None = None + + audio_filepath_key: str = "audio_filepath" + waveform_key: str = "waveform" + sample_rate_key: str = "sample_rate" + num_channels_key: str = "num_channels" + duration_key: str = "duration" + output_audio_filepath_key: str = "converted_audio_filepath" + original_audio_filepath_key: str = "original_audio_filepath" + + input_residency: InputResidency = "auto" + keep_waveform_in_task: bool = True + write_to_disk: bool = False + update_audio_filepath: bool = False + output_dir: str | None = None + + # Own bookkeeping: the channel count, recorded for readers and reports. Nothing routes on + # it, so it needs no shared role -- and declaring it here means adding this stage did not + # require touching the central role table. + INTERNAL_KEY_FIELDS: ClassVar[frozenset[str]] = frozenset({"num_channels_key"}) + + # Which action each parameter belongs to. A parameter left at its default is "not asked + # for", which is why every one of these defaults to None or False: it lets the constructor + # tell "I want mono" from "I never mentioned channels" and refuse the former under the + # wrong action instead of silently dropping the request. + _ACTION_PARAMS: ClassVar[dict[str, tuple[str, ...]]] = { + "filter": ("allowed_channels", "min_channels", "max_channels"), + "convert": ("target_channels", "write_to_disk", "update_audio_filepath", "output_dir"), + } + + name: str = "ChannelCount" + batch_size: int = 1 + resources: Resources = field(default_factory=lambda: Resources(cpus=1.0)) + + def __post_init__(self): + super().__init__() + if self.action not in ("annotate", "filter", "convert"): + msg = f"action must be one of ('annotate', 'filter', 'convert'), got {self.action!r}" + raise ValueError(msg) + self._reject_other_actions_params() + if self.action == "filter": + self._validate_filter() + if self.action == "convert": + self._validate_convert() + + def _reject_other_actions_params(self) -> None: + """Refuse parameters belonging to an action other than the configured one. + + Ignoring them is what makes the two intents blur: ``action="convert", + allowed_channels=[1]`` reads as "make everything mono AND only keep mono", and whichever + half is silently dropped, the corpus that comes out is not the one that was asked for. + """ + defaults = {item.name: item.default for item in fields(self)} + foreign = [ + (name, owner) + for owner, names in self._ACTION_PARAMS.items() + if owner != self.action + for name in names + if getattr(self, name) != defaults[name] + ] + if not foreign: + return + listed = ", ".join(f"{name} (action={owner!r})" for name, owner in sorted(foreign)) + msg = ( + f"action={self.action!r} does not use {listed}. Selecting a channel count and " + f"converting to one are separate intents: convert changes every row and keeps all " + f"of them, filter changes nothing and keeps a subset. Set the action those " + f"parameters belong to, or drop them." + ) + raise ValueError(msg) + + def _validate_filter(self) -> None: + if self.allowed_channels is not None and not self.allowed_channels: + msg = "allowed_channels must name at least one count, or be None for no constraint" + raise ValueError(msg) + for name in ("allowed_channels", "min_channels", "max_channels"): + value = getattr(self, name) + counts = value if isinstance(value, list) else [value] + for count in counts: + if count is None: + continue + if isinstance(count, bool) or not isinstance(count, int) or count < 1: + msg = f"{name} must be whole channel counts of at least 1, got {value!r}" + raise ValueError(msg) + low, high = self.min_channels, self.max_channels + if low is not None and high is not None and low > high: + msg = f"min_channels ({low}) is above max_channels ({high}), so nothing can pass" + raise ValueError(msg) + if self.allowed_channels is None and low is None and high is None: + msg = ( + "action='filter' needs allowed_channels, min_channels or max_channels -- with no " + "constraint it would declare a filter that drops nothing. Use action='annotate' " + "to only record the count." + ) + raise ValueError(msg) + + def _validate_convert(self) -> None: + # Type as well as range. YAML reads ``target_channels: 2.0`` as a float, which used to + # construct fine and then die inside a worker at ``waveform.repeat(2.0, 1)`` with a + # TypeError -- not one of the (OSError, RuntimeError) this stage drops rows for, so it + # propagated and took the run down mid-corpus instead of being caught at the recipe. + target = self.target_channels + if target is None: + return + if isinstance(target, bool) or not isinstance(target, int): + msg = f"target_channels must be a whole number of channels, got {target!r} ({type(target).__name__})" + raise ValueError(msg) # noqa: TRY004 + if target < 1: + msg = f"target_channels must be at least 1, got {target}" + raise ValueError(msg) + + @property + def _target(self) -> int: + """The channel count ``convert`` produces; unset means mono.""" + return 1 if self.target_channels is None else self.target_channels + + def inputs(self) -> tuple[list[str], list[str]]: + return [], [] + + def outputs(self) -> tuple[list[str], list[str]]: + return [], self._written_keys() + + def _written_keys(self) -> list[str]: + if self.action != "convert": + return [self.num_channels_key] + keys = [ + self.num_channels_key, + self.duration_key, + ] + if self.keep_waveform_in_task: + keys.extend([self.waveform_key, self.sample_rate_key]) + if self.write_to_disk: + keys.append(self.output_audio_filepath_key) + if self.update_audio_filepath: + keys.append(self.audio_filepath_key) + return keys + + def describe(self) -> StageContract: + if self.action == "convert": + return self._convert_contract() + return self._observe_contract() + + def _observe_contract(self) -> StageContract: + """The contract for ``annotate``/``filter``: a count in, a count out, no samples read.""" + forms = accepts_for_residency(self.input_residency) + reads_one_of: list[IOSpec] = [] + if "waveform" in forms: + # The count is ``waveform.shape[0]``. Unlike a conversion this needs no sample rate, + # and asking for one would make a resident-waveform row look unsatisfiable for a + # stage that can in fact answer from the waveform alone. + reads_one_of.append(IOSpec(data_keys=[self.waveform_key], accepts=["waveform"])) + if "file" in forms: + reads_one_of.append(IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"])) + return StageContract( + reads_one_of=reads_one_of, + writes=IOSpec(data_keys=[self.num_channels_key]), + cardinality="filter" if self.action == "filter" else "1:1", + # The two row cardinalities this stage can take across its actions. ``convert`` + # names no third one: it is "1:1" for mono and "filter" for the targets it refuses. + cardinality_options=["filter", "annotate"], + # Each row is judged against the configured counts, not against the corpus. + gates=Gates(per_row_independent=True), + ) + + def _convert_contract(self) -> StageContract: + produces = [] + if self.keep_waveform_in_task: + produces.append("tensor") + if self.write_to_disk: + produces.append("disk") + return StageContract( + reads_one_of=residency_read_specs( + self.input_residency, + audio_filepath_key=self.audio_filepath_key, + waveform_key=self.waveform_key, + sample_rate_key=self.sample_rate_key, + ), + writes=IOSpec(data_keys=self._written_keys(), produces=produces), + # Downmixing to mono always succeeds, but any other target refuses the conversions + # it cannot do correctly (N > target > 1) and drops those rows. That makes the stage + # a filter for those configurations, and saying so is what puts a seam in the + # semantic review packet for a reviewer to ask about. + cardinality="filter" if self._target > 1 else "1:1", + cardinality_options=["filter", "annotate"], + # Declared here, by the stage that owns the parameter, so a caller running + # this in a sandbox knows what to redirect without a central table entry. + gates=Gates( + writes_to_disk=self.write_to_disk, + output_path_params=["output_dir"], + # Every row is converted on its own terms, so a delta run gives the changed + # files the same answer a full run would have given them. + per_row_independent=True, + ), + ) + + def accepts(self, num_channels: int) -> bool: + """Whether ``num_channels`` satisfies every constraint that was configured.""" + if self.allowed_channels is not None and num_channels not in self.allowed_channels: + return False + if self.min_channels is not None and num_channels < self.min_channels: + return False + return not (self.max_channels is not None and num_channels > self.max_channels) + + def _requirement(self) -> str: + """The configured constraint, phrased for a log line.""" + parts = [] + if self.allowed_channels is not None: + parts.append(f"one of {sorted(self.allowed_channels)}") + if self.min_channels is not None: + parts.append(f">= {self.min_channels}") + if self.max_channels is not None: + parts.append(f"<= {self.max_channels}") + return " and ".join(parts) or "any channel count" + + def _observed_channels(self, task: AudioTask) -> int | None: # noqa: PLR0911 (complexity accepted: one early return per input/error condition) + """The row's channel count, from resident audio if present, else from the file header. + + Resident audio is asked first because it is the audio this pipeline is carrying: after a + conversion the file on disk still has its original channels while the waveform in the + task has the converted ones, so the header would answer about audio nobody is using any + more. Without a waveform the header is read rather than an existing ``num_channels`` + column believed -- standing alone that column is manifest metadata about a file nobody + re-opened, and trusting it lets a stale value decide the filter. + """ + resident = task.data.get(self.waveform_key) + if self.input_residency != "file" and resident is not None: + try: + return int(ensure_waveform_2d(resident).shape[0]) + except (RuntimeError, TypeError, ValueError, IndexError) as e: + logger.error(f"Could not read the channel count of the resident waveform: {e}") + return None + if self.input_residency == "waveform": + logger.error( + f"No resident waveform under {self.waveform_key!r} and input_residency='waveform', " + "so there is nothing to count without touching disk" + ) + return None + + declared = task.data.get(self.num_channels_key) + declared = int(declared) if isinstance(declared, (int, float)) and int(declared) > 0 else None + path = task.data.get(self.audio_filepath_key) + if not path: + if declared is None: + logger.error(f"No channel count and no audio path under {self.audio_filepath_key!r}") + return None + # Nothing to verify against, so the declared count is all there is. Say so rather + # than dropping a row that may well be fine. + logger.warning( + f"Filtering on an unverified channel count ({declared}): no resident waveform " + f"and no path under {self.audio_filepath_key!r}" + ) + return declared + try: + # Header only: the count is metadata, so decoding samples to reach it would cost + # orders of magnitude more per file for something in the first few bytes. + return int(sf.info(os.path.expanduser(str(path))).channels) + except (OSError, RuntimeError) as e: + logger.error(f"Could not read the channel count of {path!r}: {e}") + return None + + def _convert(self, waveform: torch.Tensor, source: str) -> torch.Tensor | None: + """Bring ``waveform`` to ``target_channels``, or None when that cannot be done right.""" + num_channels = waveform.shape[0] + target = self._target + if num_channels == target: + return waveform + if target == 1: + logger.debug(f"Averaging {num_channels} channels to mono") + return torch.mean(waveform, dim=0, keepdim=True) + if num_channels == 1: + # Duplicate the single channel. This adds no information -- the result is the + # same signal N times -- but it is what `ffmpeg -ac` does and what a consumer + # expecting a fixed channel count needs. + logger.debug(f"Duplicating mono into {target} channels") + return waveform.repeat(target, 1) + logger.warning( + f"Cannot downmix {num_channels} channels to {target} without the " + f"file's channel layout, which a waveform tensor does not carry: {source}. " + "Use ResampleAudioStage (ffmpeg) for a layout-aware downmix, or " + "target_channels=1." + ) + return None + + def _write_audio(self, waveform: torch.Tensor, sample_rate: int, task: AudioTask) -> str: + stem = os.path.splitext(os.path.basename(str(task.data.get(self.audio_filepath_key, "audio"))))[0] + return write_audio_stable( + waveform, + sample_rate, + output_dir=self.output_dir or tempfile.gettempdir(), + stem=stem, + tag=f"ch{self._target}", + ) + + def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: + """Record, select on, or change the channel count, per ``action``.""" + if self.action == "convert": + return self._convert_row(task) + return self._observe_row(task) + + def _observe_row(self, task: AudioTask) -> AudioTask | list[AudioTask]: + """Record the channel count, and under ``filter`` keep the row only if it is allowed.""" + num_channels = self._observed_channels(task) + if num_channels is None: + return [] + task.data[self.num_channels_key] = num_channels + + if self.action == "filter" and not self.accepts(num_channels): + logger.warning( + f"Channel count {num_channels} does not satisfy {self._requirement()}: " + f"{task.data.get(self.audio_filepath_key, self.waveform_key)}" + ) + return [] + return task + + def _convert_row(self, task: AudioTask) -> AudioTask | list[AudioTask]: + """Convert the audio's channel count. Returns [] for a row that cannot be converted.""" + 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 + ) + except (OSError, RuntimeError) as e: # corrupt/unreadable audio -> skip the row + 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 = resolved + waveform = ensure_waveform_2d(waveform) + + if sample_rate <= 0: + logger.error(f"Invalid sample rate ({sample_rate}) in audio input") + return [] + + source = str(task.data.get(self.audio_filepath_key, self.waveform_key)) + converted = self._convert(waveform, source) + if converted is None: + return [] + + if self.keep_waveform_in_task: + task.data[self.waveform_key] = converted + task.data[self.sample_rate_key] = sample_rate + task.data[self.num_channels_key] = converted.shape[0] + task.data[self.duration_key] = converted.shape[1] / sample_rate + + if self.write_to_disk: + path = self._write_audio(converted, 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 input: {e}") + return [] + else: + return task diff --git a/nemo_curator/stages/audio/preprocessing/concatenation.py b/nemo_curator/stages/audio/preprocessing/concatenation.py index 4afd134bef..10a0145259 100755 --- a/nemo_curator/stages/audio/preprocessing/concatenation.py +++ b/nemo_curator/stages/audio/preprocessing/concatenation.py @@ -32,12 +32,15 @@ stage = SegmentConcatenationStage(silence_duration_sec=0.5) """ +import os from dataclasses import dataclass, field from typing import Any import torch from loguru import logger +from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, IOSpec, StageContract +from nemo_curator.stages.audio._agent._residency import write_audio_stable from nemo_curator.stages.audio.common import ensure_waveform_2d from nemo_curator.stages.base import ProcessingStage from nemo_curator.stages.resources import Resources @@ -67,7 +70,7 @@ def to_dict(self) -> dict[str, Any]: @dataclass -class SegmentConcatenationStage(ProcessingStage[AudioTask, AudioTask]): +class SegmentConcatenationStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """ Concatenate nested VAD segments into a single combined waveform. @@ -81,9 +84,27 @@ class SegmentConcatenationStage(ProcessingStage[AudioTask, AudioTask]): Args: silence_duration_sec: Duration of silence inserted between consecutive segments (seconds). + keep_waveform_in_task: Keep the combined waveform in the task (default True, + today's behavior). Set False to emit only an on-disk path (requires write_to_disk). + write_to_disk: Also write the combined waveform to ``output_dir`` and set + ``audio_filepath_key`` to it, so file-based downstream stages can consume it. + Defaults to False (in-memory only, unchanged). + output_dir: Directory for combined WAVs (required when write_to_disk=True). + audio_filepath_key: Key set to the written combined-audio path (write_to_disk). """ silence_duration_sec: float = 0.5 + segments_key: str = "segments" + waveform_key: str = "waveform" + sample_rate_key: str = "sample_rate" + original_file_key: str = "original_file" + num_segments_key: str = "num_segments" + total_duration_sec_key: str = "total_duration_sec" + audio_filepath_key: str = "audio_filepath" + # Output residency (both default to today's behavior: in-memory waveform only, no disk). + keep_waveform_in_task: bool = True + write_to_disk: bool = False + output_dir: str | None = None name: str = "SegmentConcatenation" batch_size: int = 1 @@ -91,18 +112,57 @@ class SegmentConcatenationStage(ProcessingStage[AudioTask, AudioTask]): def __post_init__(self): super().__init__() + if not (self.keep_waveform_in_task or self.write_to_disk): + msg = "At least one of keep_waveform_in_task or write_to_disk must be True" + raise ValueError(msg) + if self.write_to_disk and not self.output_dir: + msg = "output_dir is required when write_to_disk=True" + raise ValueError(msg) def inputs(self) -> tuple[list[str], list[str]]: return [], [] def outputs(self) -> tuple[list[str], list[str]]: - return [], ["waveform", "sample_rate", "num_segments", "total_duration_sec", "original_file"] + outs: list[str] = [] + if self.keep_waveform_in_task: + outs.extend([self.waveform_key, self.sample_rate_key]) + outs.extend([self.num_segments_key, self.total_duration_sec_key, self.original_file_key]) + if self.write_to_disk: + outs.append(self.audio_filepath_key) + return [], outs + + def describe(self) -> StageContract: + writes = [self.original_file_key, self.num_segments_key, self.total_duration_sec_key] + produces: list[str] = [] + if self.keep_waveform_in_task: + writes[:0] = [self.waveform_key, self.sample_rate_key] # preserve original key order + produces.append("tensor") + if self.write_to_disk: + writes.append(self.audio_filepath_key) + produces.append("disk") + return StageContract( + reads=IOSpec(data_keys=[self.segments_key]), + writes=IOSpec(data_keys=writes, produces=produces), + metadata_writes=["segment_mappings"], + cardinality="N:1", + iteration_key=self.segments_key, + gates=Gates( + writes_to_disk=self.write_to_disk, + output_path_params=["output_dir"], + # The ``N`` this stage collapses is the segments of ONE row's own file, so no + # other file's audio reaches the combined waveform -- the ``N:1`` cardinality + # counts tasks, not the origins of the values. ``write_to_disk`` does not change + # that: ``write_audio_stable`` names the WAV after a digest of its own bytes, so + # two rows can only land on one path by carrying identical audio. + per_row_independent=True, + ), + ) def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: """Concatenate segments from ``task.data["segments"]``.""" - segments = task.data.get("segments") + segments = task.data.get(self.segments_key) if segments is None: - msg = "SegmentConcatenationStage requires task.data['segments'] (nested VAD mode)" + msg = f"SegmentConcatenationStage requires task.data[{self.segments_key!r}] (nested VAD mode)" raise ValueError(msg) if not segments: @@ -111,7 +171,7 @@ def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: segments_sorted = sorted(segments, key=self._seg_sort_key) original_file = segments_sorted[0].get("original_file", "unknown") - combined = self._concatenate(original_file, segments_sorted, task.dataset_name) + combined = self._concatenate(original_file, segments_sorted, task) if combined is None: return [] return combined @@ -127,12 +187,15 @@ def _seg_sort_key(seg: dict[str, Any]) -> tuple[int, int, int]: return (0, int(start), 0) return (0, 0, 0) - @staticmethod - def _validate_segment(seg: dict[str, Any]) -> tuple[torch.Tensor, int] | None: + def _validate_segment(self, seg: dict[str, Any]) -> tuple[torch.Tensor, int] | None: """Validate and return (waveform, sample_rate) or None if invalid.""" - waveform = seg.get("waveform") - sr = seg.get("sample_rate") + waveform = seg.get(self.waveform_key) + sr = seg.get(self.sample_rate_key) if waveform is None: + logger.warning( + f"[SegmentConcat] Skipping segment {seg.get('segment_num', '?')}: no " + f"{self.waveform_key!r} (was VAD run with keep_segment_waveform_in_task=False?)" + ) return None seg_id = seg.get("segment_num", "?") if sr is None: @@ -143,11 +206,22 @@ def _validate_segment(seg: dict[str, Any]) -> tuple[torch.Tensor, int] | None: return None return ensure_waveform_2d(waveform), sr + def _write_wav(self, waveform: torch.Tensor, sr: int, original_file: str) -> str: + """Write the combined waveform to ``output_dir`` and return the path.""" + stem = os.path.splitext(os.path.basename(str(original_file)))[0] or "audio" + return write_audio_stable( + waveform, + sr, + output_dir=self.output_dir, + stem=stem, + tag="concat", + ) + def _concatenate( self, original_file: str, segments: list[dict[str, Any]], - dataset_name: str, + parent_task: AudioTask, ) -> AudioTask | None: """Concatenate a list of segment dicts from the same source file.""" parts: list[torch.Tensor] = [] @@ -214,19 +288,22 @@ def _concatenate( total_duration_sec = current_pos_ms / 1000.0 output_data = { - "waveform": combined, - "sample_rate": sample_rate, - "original_file": original_file, - "num_segments": len(mappings), - "total_duration_sec": total_duration_sec, + self.original_file_key: original_file, + self.num_segments_key: len(mappings), + self.total_duration_sec_key: total_duration_sec, } + # Output residency: keep the combined waveform in-task (default) and/or persist it. + if self.keep_waveform_in_task: + output_data[self.waveform_key] = combined + output_data[self.sample_rate_key] = sample_rate + if self.write_to_disk: + output_data[self.audio_filepath_key] = self._write_wav(combined, sample_rate, original_file) logger.info(f"[SegmentConcat] {original_file}: {len(mappings)} segments -> {total_duration_sec:.2f}s combined") - result_task = AudioTask( + return AudioTask( data=output_data, - dataset_name=dataset_name, + dataset_name=parent_task.dataset_name, + _metadata={**(parent_task._metadata or {}), "segment_mappings": mappings}, + _stage_perf=list(parent_task._stage_perf), ) - result_task._metadata = {"segment_mappings": mappings} - - return result_task diff --git a/nemo_curator/stages/audio/preprocessing/mono_conversion.py b/nemo_curator/stages/audio/preprocessing/mono_conversion.py index ed60661235..8225781b8c 100755 --- a/nemo_curator/stages/audio/preprocessing/mono_conversion.py +++ b/nemo_curator/stages/audio/preprocessing/mono_conversion.py @@ -32,14 +32,22 @@ import torch from loguru import logger -from nemo_curator.stages.audio.common import load_audio_file +from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, IOSpec, StageContract +from nemo_curator.stages.audio._agent._residency import ( + InputResidency, + produce_audio_filepath, + residency_read_specs, + resolve_audio, + write_audio_stable, +) +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 +58,48 @@ 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 — "file" (audio_filepath only; default, + matching this stage's pre-agent behavior), "waveform" (in-memory only), or + "auto" (waveform first, file fallback). + 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 + # "file", not "auto": this stage only ever read ``audio_filepath`` before the agent work + # (``load_audio_file(audio_filepath, mono=False)``), unlike sigmos/utmos/band/vad, whose + # own resolvers already preferred a resident waveform and so default to "auto" honestly. + # Defaulting to "auto" here silently changed which input a default pipeline reads. + input_residency: InputResidency = "file" + 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,34 +112,104 @@ 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_one_of=residency_read_specs( + self.input_residency, + audio_filepath_key=self.audio_filepath_key, + waveform_key=self.waveform_key, + sample_rate_key=self.sample_rate_key, + ), + writes=IOSpec(data_keys=writes, produces=produces), + gates=Gates( + writes_to_disk=self.write_to_disk, + output_path_params=["output_dir"], + per_row_independent=True, + ), + # With ``strict_sample_rate`` -- the DEFAULT -- a row whose rate differs from + # ``output_sample_rate`` returns ``[]``. That is row-dropping, and undeclared it + # made this stage the one place the rule was invisible: the newer stages doing the + # same thing say so, so validation treated a 48 kHz default silently discarding a + # 16 kHz corpus as a pass-through, while flagging its neighbours for less. + cardinality="filter" if self.strict_sample_rate else "1:1", + ) + + def _write_audio(self, waveform: torch.Tensor, sample_rate: int, task: AudioTask) -> str: + stem = os.path.splitext(os.path.basename(str(task.data.get(self.audio_filepath_key, "audio"))))[0] + return write_audio_stable( + waveform, + sample_rate, + output_dir=self.output_dir, + stem=stem, + tag="mono", + ) + + 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: - logger.warning( - f"Sample rate {sample_rate}Hz != expected {self.output_sample_rate}Hz: {audio_filepath}" - ) + 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_source}") return [] if num_channels > 1: @@ -103,14 +218,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/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py b/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py new file mode 100644 index 0000000000..1c12daead5 --- /dev/null +++ b/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py @@ -0,0 +1,202 @@ +# 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. + +""" +Sample-rate selection stage. + +Records every row's sample rate and keeps only the rows whose rate is acceptable. Channels, +format and file layout are untouched, so a pipeline pairs this with whatever channel policy +it wants -- or with none at all. + +It never resamples: ``ResampleAudioStage`` is the stage that converts a rate, this one only +decides which rates are allowed through and writes down what it saw. + +Reading rather than decoding is the point. Determining a rate needs only the file header, +so this stage never loads samples -- measured on 30s stereo WAVs, a header read is ~0.03ms +against ~5.3ms for a full decode (~186x). Placing it before any decoding stage means rows +that will be rejected are never decoded at all. + +Example: + from nemo_curator.pipeline import Pipeline + from nemo_curator.stages.audio.preprocessing import SampleRateFilterStage + + pipeline = Pipeline(name="audio_pipeline") + pipeline.add_stage(SampleRateFilterStage(allowed_sample_rates=[16000, 22050])) + pipeline.add_stage(SampleRateFilterStage(min_sample_rate=16000)) +""" + +import os +from dataclasses import dataclass, field + +import soundfile as sf +from loguru import logger + +from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, IOSpec, StageContract +from nemo_curator.stages.base import ProcessingStage +from nemo_curator.stages.resources import Resources +from nemo_curator.tasks import AudioTask + + +@dataclass +class SampleRateFilterStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): + """ + Keep rows whose audio sample rate is acceptable, and record the rate on every row. + + Acceptance is expressed two ways, deliberately as separate parameters rather than one + that means both. A single ``sample_rates=[16000, 48000]`` is genuinely ambiguous -- + "these two rates" or "everything between them"? -- and the two readings filter very + different corpora. + + * ``allowed_sample_rates``: an explicit set. A rate must be one of these. + * ``min_sample_rate`` / ``max_sample_rate``: inclusive bounds. + + Set either, both (a rate must then satisfy both), or neither -- with nothing set the + stage keeps every row and acts purely as an observer that annotates the rate. + + Nothing is resampled here. A corpus with mixed rates that must be uniform needs + ``ResampleAudioStage``; this stage only decides what is allowed through. + + Args: + allowed_sample_rates: Explicit rates to keep, e.g. [16000, 22050]. None = no + constraint from this parameter. + min_sample_rate: Lowest acceptable rate, inclusive. None = unbounded below. + max_sample_rate: Highest acceptable rate, inclusive. None = unbounded above. + audio_filepath_key: Key in data dict for the audio file path. + sample_rate_key: Key where the observed sample rate is written. Reused without a + disk read only when a resident waveform backs it (see ``waveform_key``). + waveform_key: Key a resident waveform would occupy. Consulted to know whether the + sample rate in task.data belongs to resident audio; a rate standing alone is + manifest metadata and is re-read from the file header instead of trusted. + """ + + allowed_sample_rates: list[int] | None = None + min_sample_rate: int | None = None + max_sample_rate: int | None = None + + audio_filepath_key: str = "audio_filepath" + sample_rate_key: str = "sample_rate" + waveform_key: str = "waveform" + + name: str = "SampleRateFilter" + batch_size: int = 1 + resources: Resources = field(default_factory=lambda: Resources(cpus=1.0)) + + def __post_init__(self): + super().__init__() + if self.allowed_sample_rates is not None and not self.allowed_sample_rates: + msg = "allowed_sample_rates must name at least one rate, or be None for no constraint" + raise ValueError(msg) + low, high = self.min_sample_rate, self.max_sample_rate + if low is not None and high is not None and low > high: + msg = f"min_sample_rate ({low}) is above max_sample_rate ({high}), so nothing can pass" + raise ValueError(msg) + + def inputs(self) -> tuple[list[str], list[str]]: + return [], [] + + def outputs(self) -> tuple[list[str], list[str]]: + return [], [self.sample_rate_key] + + def describe(self) -> StageContract: + # Either a resident rate or a readable path satisfies this stage; it needs no samples, + # so the waveform form asks for the rate alone rather than the waveform with it. + return StageContract( + reads_one_of=[ + IOSpec(data_keys=[self.sample_rate_key], accepts=["waveform"]), + IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"]), + ], + writes=IOSpec(data_keys=[self.sample_rate_key]), + # Dropping rows is this stage's whole purpose, and "filter" is what puts a seam in + # the semantic review packet. Left at the 1:1 default the reviewer is never told the + # corpus can shrink here, so nobody asks how much of it survives. + cardinality="filter", + # Each row is judged against the configured rates, not against the corpus. + gates=Gates(per_row_independent=True), + ) + + def accepts(self, sample_rate: int) -> bool: + """Whether ``sample_rate`` satisfies every constraint that was configured.""" + if self.allowed_sample_rates is not None and sample_rate not in self.allowed_sample_rates: + return False + if self.min_sample_rate is not None and sample_rate < self.min_sample_rate: + return False + return not (self.max_sample_rate is not None and sample_rate > self.max_sample_rate) + + def _requirement(self) -> str: + """The configured constraint, phrased for a log line.""" + parts = [] + if self.allowed_sample_rates is not None: + parts.append(f"one of {sorted(self.allowed_sample_rates)}") + if self.min_sample_rate is not None: + parts.append(f">= {self.min_sample_rate}") + if self.max_sample_rate is not None: + parts.append(f"<= {self.max_sample_rate}") + return " and ".join(parts) or "any rate" + + def _observed_rate(self, task: AudioTask) -> int | None: + """The row's sample rate, from resident audio if present, else from the file header. + + An existing ``sample_rate_key`` is only believed when a resident waveform is there to + back it, because then it describes audio this pipeline is carrying. Standing alone it + is manifest metadata about a file nobody re-read, and trusting it lets a stale or wrong + column decide the filter: a genuinely 48 kHz file labelled 16000 would be kept for a + 16 kHz-only corpus AND re-stamped with the wrong rate. The header read is cheap enough + that guessing is never worth it. + """ + declared = task.data.get(self.sample_rate_key) + declared = int(declared) if isinstance(declared, (int, float)) and int(declared) > 0 else None + if declared is not None and self.waveform_key in task.data: + return declared + + path = task.data.get(self.audio_filepath_key) + if not path: + if declared is None: + logger.error(f"No sample rate and no audio path under {self.audio_filepath_key!r}") + return None + # Nothing to verify against, so the declared rate is all there is. Say so rather + # than dropping a row that may well be fine. + logger.warning( + f"Filtering on an unverified sample rate ({declared}Hz): no resident waveform " + f"and no path under {self.audio_filepath_key!r}" + ) + return declared + try: + # Header only: the rate is metadata, so decoding samples to read it would cost + # ~186x more per file for information already sitting in the first few bytes. + # ``expanduser`` matches ChannelCountStage: a manifest written with ``~/audio/x.wav`` + # is otherwise unreadable here, and the failure path drops the row rather than + # raising, so the whole corpus would disappear through this stage in silence. + return int(sf.info(os.path.expanduser(str(path))).samplerate) + except (OSError, RuntimeError) as e: + logger.error(f"Could not read the sample rate of {path!r}: {e}") + return None + + def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: + """Record the sample rate and keep the row only if that rate is acceptable.""" + sample_rate = self._observed_rate(task) + if sample_rate is None: + return [] + if sample_rate <= 0: + logger.error(f"Invalid sample rate ({sample_rate}) for {task.data.get(self.audio_filepath_key)!r}") + return [] + + task.data[self.sample_rate_key] = sample_rate + + if not self.accepts(sample_rate): + logger.warning( + f"Sample rate {sample_rate}Hz does not satisfy {self._requirement()}: " + f"{task.data.get(self.audio_filepath_key, self.waveform_key)}" + ) + return [] + return task diff --git a/tests/stages/audio/_agent/__init__.py b/tests/stages/audio/_agent/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/stages/audio/_agent/test_agent_conformance_examples.py b/tests/stages/audio/_agent/test_agent_conformance_examples.py new file mode 100644 index 0000000000..8fea4ab24b --- /dev/null +++ b/tests/stages/audio/_agent/test_agent_conformance_examples.py @@ -0,0 +1,113 @@ +# 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. + +"""EXEMPLAR per-stage conformance tests — copy these as templates. + +Each stage owner adds one ``assert_agent_ready(...)`` test for their stage (see +``nemo_curator/stages/audio/AGENT_READY.md``). These three cover the common +patterns on CPU with no models: + + * 1:1 transform reading a file -> MonoConversionStage + * 1:1 annotate writing one metric -> GetAudioDurationStage + * filter (batch-only, may drop items) -> PreserveByValueStage + +``assert_agent_ready`` runs the stage on the fixture and verifies the declared +contract matches runtime: declared writes appear, no undeclared top-level keys +leak, cardinality matches, and reads are satisfiable by role. Model/GPU stages +follow the same shape but build the fixture + fake model via the stub harness in +``test_agent_simulation_pipelines.py``. +""" + +from __future__ import annotations + +from pathlib import Path # noqa: TC003 + +import numpy as np +import soundfile as sf +import torch + +from nemo_curator.stages.audio._agent._conformance import assert_agent_ready, assert_residency_consumption +from nemo_curator.stages.audio.common import GetAudioDurationStage, PreserveByValueStage +from nemo_curator.stages.audio.preprocessing.mono_conversion import MonoConversionStage +from nemo_curator.tasks import AudioTask + + +def _write_wav(path: Path, *, duration_sec: float = 1.0, sample_rate: int = 48000) -> str: + samples = np.zeros(int(duration_sec * sample_rate), dtype="float32") + sf.write(str(path), samples, sample_rate) + return str(path) + + +# --- TEMPLATE 1: a 1:1 transform that reads a file -------------------------- # +def test_example_transform_stage_conformance(tmp_path: Path) -> None: + wav = _write_wav(tmp_path / "a.wav", sample_rate=48000) # MonoConversion default sr + + def fixture() -> AudioTask: + return AudioTask(dataset_name="t", data={"audio_filepath": wav}) + + # ``strict_sample_rate`` (the default) drops rows whose rate differs, which is a + # filter; relaxing it gives the plain 1:1 transform this template is meant to show. + assert_agent_ready( + MonoConversionStage(strict_sample_rate=False), + fixture, + expected_cardinality="1:1", + available_keys={"audio_filepath"}, + ) + + +# --- TEMPLATE 2: a 1:1 stage that annotates one metric ---------------------- # +def test_example_metric_stage_conformance(tmp_path: Path) -> None: + wav = _write_wav(tmp_path / "b.wav") + + def fixture() -> AudioTask: + return AudioTask(dataset_name="t", data={"audio_filepath": wav}) + + assert_agent_ready( + GetAudioDurationStage(), + fixture, + expected_cardinality="1:1", + available_keys={"audio_filepath"}, + ) + + +# --- TEMPLATE 4: per-residency consumption (advertised residency == code) --- # +def test_example_residency_consumption(tmp_path: Path) -> None: + """A residency-configurable stage must actually consume each residency it advertises.""" + wav = _write_wav(tmp_path / "r.wav", sample_rate=48000) + + def file_fixture() -> AudioTask: + return AudioTask(dataset_name="t", data={"audio_filepath": wav}) + + def waveform_fixture() -> AudioTask: + return AudioTask(dataset_name="t", data={"waveform": torch.zeros(2, 48000), "sample_rate": 48000}) + + assert_residency_consumption( + lambda r: MonoConversionStage(input_residency=r), + file_fixture=file_fixture, + waveform_fixture=waveform_fixture, + ) + + +# --- TEMPLATE 3: a filter (batch-only; may drop items) ---------------------- # +def test_example_filter_stage_conformance() -> None: + def fixture() -> AudioTask: + # value passes the filter (keep == True), so the task survives + return AudioTask(dataset_name="t", data={"keep": True}) + + assert_agent_ready( + PreserveByValueStage("keep", target_value=True, operator="eq"), + fixture, + expected_cardinality="filter", + available_keys={"keep"}, + ) diff --git a/tests/stages/audio/preprocessing/test_channel_and_rate_stages.py b/tests/stages/audio/preprocessing/test_channel_and_rate_stages.py new file mode 100644 index 0000000000..d6fb01273d --- /dev/null +++ b/tests/stages/audio/preprocessing/test_channel_and_rate_stages.py @@ -0,0 +1,522 @@ +# 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. + +"""Channel policy and rate policy as two independent stages. + +Each does one job and leaves the other alone, so a pipeline sets a channel policy and a +rate policy separately -- or uses only the one it needs. Each also says which job it is +doing: recording a value, selecting on it, or changing it. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from unittest.mock import patch + +import numpy as np +import pytest +import soundfile as sf + +from nemo_curator.stages.audio.preprocessing import ChannelCountStage, SampleRateFilterStage +from nemo_curator.stages.audio.preprocessing import channel_count as cc +from nemo_curator.stages.audio.preprocessing import sample_rate_filter as srf +from nemo_curator.tasks import AudioTask + +if TYPE_CHECKING: + from pathlib import Path + + +def _wav(tmp_path: Path, channels: int = 1, rate: int = 16000, name: str = "a.wav") -> str: + path = tmp_path / name + data = np.zeros(rate, dtype="float32") if channels == 1 else np.zeros((rate, channels), dtype="float32") + sf.write(str(path), data, rate) + return str(path) + + +def _task(path: str) -> AudioTask: + return AudioTask(task_id="t", dataset_name="d", data={"audio_filepath": path}) + + +def _convert(**kwargs: object) -> ChannelCountStage: + return ChannelCountStage(action="convert", **kwargs) + + +class TestConvertingTheChannelCount: + @pytest.mark.parametrize( + ("source", "target"), + [(1, 1), (2, 1), (6, 1), (2, 2), (6, 6)], + ) + def test_downmix_and_passthrough_produce_the_requested_count( + self, tmp_path: Path, source: int, target: int + ) -> None: + result = _convert(target_channels=target).process( + _task(_wav(tmp_path, channels=source, name=f"{source}to{target}.wav")) + ) + assert result != [] + assert result.data["num_channels"] == target + + def test_mono_downmix_averages_rather_than_selecting_a_channel(self, tmp_path: Path) -> None: + """Averaging is what mono means; taking channel 0 would discard half the signal.""" + path = str(tmp_path / "stereo.wav") + left = np.ones(16000, dtype="float32") + right = -np.ones(16000, dtype="float32") + sf.write(path, np.stack([left, right], axis=1), 16000) + + result = _convert(target_channels=1).process(_task(path)) + + assert result != [] + # +1 and -1 average to ~0; selecting either channel would give ~1. The residual is + # 16-bit quantization (sf.write stores PCM_16, so +/-1.0 becomes 32767/-32768), + # which is ~3e-5 -- three orders of magnitude below what a channel-select yields. + assert float(result.data["waveform"].abs().max()) < 0.01 + + def test_mono_upmixes_by_duplication(self, tmp_path: Path) -> None: + """Matches ``ffmpeg -ac``. It adds no information -- the channels are identical.""" + result = _convert(target_channels=2).process(_task(_wav(tmp_path, channels=1))) + + assert result != [] + waveform = result.data["waveform"] + assert waveform.shape[0] == 2 + assert waveform[0].equal(waveform[1]), "upmix duplicates, so both channels are the same signal" + + def test_an_unset_target_means_mono(self, tmp_path: Path) -> None: + """``target_channels`` defaults to None so the constructor can tell "make it mono" from + "channels were never mentioned" and refuse the first under the wrong action. Mono is + still what an unset target converts to.""" + result = _convert().process(_task(_wav(tmp_path, channels=2))) + + assert result != [] + assert result.data["num_channels"] == 1 + + @pytest.mark.parametrize(("source", "target"), [(6, 2), (4, 3), (6, 5)]) + def test_a_surround_downmix_is_refused_rather_than_approximated( + self, tmp_path: Path, source: int, target: int + ) -> None: + """Correct downmix to >1 channel needs BS.775 coefficients AND the file's channel + order, and a (channels, samples) tensor carries neither -- the WAV channel mask is + gone by then. Averaging 5.1 into two channels sounds plausible and is wrong, so the + row is dropped and ResampleAudioStage (ffmpeg, layout-aware) is named instead.""" + result = _convert(target_channels=target).process( + _task(_wav(tmp_path, channels=source, name=f"{source}to{target}.wav")) + ) + assert result == [] + + def test_the_sample_rate_is_never_changed(self, tmp_path: Path) -> None: + """This stage converts channels only; rate policy belongs to another stage.""" + result = _convert(target_channels=1).process(_task(_wav(tmp_path, channels=2, rate=44100))) + assert result != [] + assert result.data["sample_rate"] == 44100 + + def test_a_nonsensical_target_is_rejected_at_construction(self) -> None: + with pytest.raises(ValueError, match="target_channels"): + _convert(target_channels=0) + + @pytest.mark.parametrize("bad", [2.0, "2"]) + def test_a_non_integer_channel_count_is_rejected_at_construction(self, bad: object) -> None: + """YAML reads ``target_channels: 2.0`` as a float. It used to construct fine and then + die inside a worker at ``waveform.repeat(2.0, 1)`` with a TypeError, which is not one + of the errors this stage drops rows for -- so it propagated and took the run down + partway through the corpus rather than being caught at the recipe.""" + with pytest.raises(ValueError, match="whole number of channels"): + _convert(target_channels=bad) + + +class TestRecordingTheChannelCount: + """``action="annotate"``: measure and keep. The default, because a neutral name should not + rewrite audio until asked.""" + + @pytest.mark.parametrize("channels", [1, 2, 6]) + def test_every_row_is_kept_and_stamped_with_what_it_actually_has(self, tmp_path: Path, channels: int) -> None: + result = ChannelCountStage().process(_task(_wav(tmp_path, channels=channels, name=f"{channels}ch.wav"))) + + assert result != [] + assert result.data["num_channels"] == channels + + def test_nothing_is_rewritten(self, tmp_path: Path) -> None: + """Recording a count is not a conversion: no waveform is produced and no file written.""" + result = ChannelCountStage().process(_task(_wav(tmp_path, channels=2))) + + assert result != [] + assert "waveform" not in result.data + assert "converted_audio_filepath" not in result.data + + def test_it_reads_the_header_and_never_decodes(self, tmp_path: Path) -> None: + """The count sits in the first bytes of the file, so putting this in front of a + decoding stage costs almost nothing. A decode here would forfeit exactly that.""" + + def explode(*_args: object, **_kwargs: object) -> None: + msg = "decoded the audio to read a header value" + raise AssertionError(msg) + + with patch.object(cc.sf, "read", explode), patch.object(cc, "load_audio_file", explode): + result = ChannelCountStage().process(_task(_wav(tmp_path, channels=2))) + + assert result != [] + assert result.data["num_channels"] == 2 + + def test_resident_audio_answers_before_the_file_does(self, tmp_path: Path) -> None: + """After a conversion the file on disk still has its original channels while the + waveform in the task has the converted ones. Reading the header there would report on + audio nobody is using any more, so a resident waveform wins.""" + task = _task(_wav(tmp_path, channels=6)) + task.data["waveform"] = np.zeros((1, 16000), dtype="float32") + + result = ChannelCountStage().process(task) + + assert result != [] + assert result.data["num_channels"] == 1, "the resident mono waveform, not the 6-channel file" + + def test_a_stale_manifest_column_is_corrected_rather_than_believed(self, tmp_path: Path) -> None: + """``num_channels`` standing alone is metadata about a file nobody re-opened. Trusting + it would let a wrong column decide a filter AND be re-stamped as if measured.""" + task = _task(_wav(tmp_path, channels=2)) + task.data["num_channels"] = 1 + + result = ChannelCountStage().process(task) + + assert result != [] + assert result.data["num_channels"] == 2 + + def test_an_unreadable_row_is_dropped_not_crashed(self, tmp_path: Path) -> None: + assert ChannelCountStage().process(_task(str(tmp_path / "missing.wav"))) == [] + + def test_an_unverifiable_count_is_used_rather_than_dropping_the_row(self) -> None: + """No resident audio and no path leaves nothing to check against. Using the declared + count beats discarding a row that may well be fine.""" + task = AudioTask(task_id="t", dataset_name="d", data={"num_channels": 2}) + + result = ChannelCountStage().process(task) + + assert result != [] + assert result.data["num_channels"] == 2 + + def test_waveform_residency_never_falls_back_to_disk(self, tmp_path: Path) -> None: + """``input_residency="waveform"`` is a promise not to touch the filesystem. Silently + reading the header instead would break it for a caller who set it to avoid exactly + that.""" + assert ChannelCountStage(input_residency="waveform").process(_task(_wav(tmp_path, channels=2))) == [] + + +class TestSelectingByChannelCount: + """``action="filter"``: keep the rows that already comply, rewrite nothing.""" + + @pytest.mark.parametrize( + ("kwargs", "keeps"), + [ + ({"allowed_channels": [1]}, False), + ({"allowed_channels": [2]}, True), + ({"allowed_channels": [1, 2]}, True), + ({"min_channels": 2}, True), + ({"min_channels": 3}, False), + ({"max_channels": 2}, True), + ({"max_channels": 1}, False), + ({"min_channels": 1, "max_channels": 6}, True), + ({"allowed_channels": [2], "min_channels": 3}, False), + ], + ) + def test_a_list_and_a_range_are_separate_constraints( + self, tmp_path: Path, kwargs: dict[str, object], keeps: bool + ) -> None: + """Separate parameters for the same reason the rate side has them: ``[1, 6]`` as a + single knob is ambiguous between "these two counts" and "this range". Every constraint + that IS set must be satisfied.""" + result = ChannelCountStage(action="filter", **kwargs).process(_task(_wav(tmp_path, channels=2))) + + assert (result != []) is keeps + + def test_the_count_is_recorded_on_rows_that_pass(self, tmp_path: Path) -> None: + result = ChannelCountStage(action="filter", allowed_channels=[2]).process(_task(_wav(tmp_path, channels=2))) + + assert result != [] + assert result.data["num_channels"] == 2 + + def test_selection_does_not_convert_the_rows_it_keeps(self, tmp_path: Path) -> None: + """A kept row is untouched: this is the drop-only path, so no audio is rewritten.""" + result = ChannelCountStage(action="filter", min_channels=1).process(_task(_wav(tmp_path, channels=2))) + + assert result != [] + assert "waveform" not in result.data + + def test_a_constraint_free_filter_is_rejected_at_construction(self) -> None: + """It would declare a filter that drops nothing, and the stage already has an action + for that.""" + with pytest.raises(ValueError, match="action='annotate'"): + ChannelCountStage(action="filter") + + def test_an_empty_allow_list_is_rejected_at_construction(self) -> None: + """``[]`` would silently discard the entire corpus; None means "no constraint".""" + with pytest.raises(ValueError, match="at least one count"): + ChannelCountStage(action="filter", allowed_channels=[]) + + def test_an_inverted_range_is_rejected_at_construction(self) -> None: + with pytest.raises(ValueError, match="nothing can pass"): + ChannelCountStage(action="filter", min_channels=6, max_channels=2) + + @pytest.mark.parametrize("bad", [[0], [1.5], [True], 0]) + def test_a_nonsensical_count_is_rejected_at_construction(self, bad: object) -> None: + key = "min_channels" if isinstance(bad, int) else "allowed_channels" + with pytest.raises(ValueError, match="whole channel counts"): + ChannelCountStage(action="filter", **{key: bad}) + + +class TestSelectingAndConvertingAreNotTheSameKnob: + """The footgun this stage exists to remove: ``target_channels=1`` must never quietly mean + "drop everything that is not mono", and ``allowed_channels=[1]`` must never quietly mean + "make it mono". They do opposite things to a corpus -- one keeps every row and changes it, + the other changes no row and keeps a subset -- so a parameter naming the action you did + not choose is refused instead of ignored. + """ + + @pytest.mark.parametrize( + ("kwargs", "unusable"), + [ + ({"action": "convert", "allowed_channels": [1]}, "allowed_channels"), + ({"action": "convert", "min_channels": 1}, "min_channels"), + ({"action": "filter", "allowed_channels": [1], "target_channels": 1}, "target_channels"), + ({"action": "filter", "allowed_channels": [1], "write_to_disk": True}, "write_to_disk"), + ({"action": "annotate", "allowed_channels": [1]}, "allowed_channels"), + ({"action": "annotate", "target_channels": 1}, "target_channels"), + ], + ) + def test_the_other_actions_parameters_are_refused(self, kwargs: dict[str, object], unusable: str) -> None: + with pytest.raises(ValueError, match=unusable): + ChannelCountStage(**kwargs) + + def test_the_refusal_names_the_action_that_would_use_it(self) -> None: + """A message that only says "unused" leaves the caller to guess which half of their + intent was dropped.""" + with pytest.raises(ValueError, match="action='filter'"): + ChannelCountStage(action="convert", allowed_channels=[1]) + + def test_an_unknown_action_is_rejected_at_construction(self) -> None: + """A recipe is free text before it is a stage; a typo must not fall through to a + default behaviour the caller did not name.""" + with pytest.raises(ValueError, match="action must be one of"): + ChannelCountStage(action="downmix") + + def test_converting_keeps_the_rows_selection_would_have_dropped(self, tmp_path: Path) -> None: + """Same corpus, same "mono" intent, opposite outcomes -- which is why they cannot share + a parameter.""" + path = _wav(tmp_path, channels=2) + + converted = _convert(target_channels=1).process(_task(path)) + selected = ChannelCountStage(action="filter", allowed_channels=[1]).process(_task(path)) + + assert converted != [] + assert converted.data["num_channels"] == 1 + assert selected == [] + + +class TestSampleRateFilter: + @pytest.mark.parametrize( + ("kwargs", "keeps"), + [ + ({"allowed_sample_rates": [16000]}, True), + ({"allowed_sample_rates": [22050, 44100]}, False), + ({"min_sample_rate": 16000}, True), + ({"min_sample_rate": 22050}, False), + ({"max_sample_rate": 16000}, True), + ({"max_sample_rate": 8000}, False), + ({"min_sample_rate": 8000, "max_sample_rate": 48000}, True), + ({"allowed_sample_rates": [16000], "min_sample_rate": 22050}, False), + ({}, True), + ], + ) + def test_a_list_and_a_range_are_separate_constraints( + self, tmp_path: Path, kwargs: dict[str, object], keeps: bool + ) -> None: + """Separate parameters on purpose: ``[16000, 48000]`` as a single knob is ambiguous + between "these two rates" and "this range", and the readings filter very different + corpora. Every constraint that IS set must be satisfied.""" + result = SampleRateFilterStage(**kwargs).process(_task(_wav(tmp_path, rate=16000))) + assert (result != []) is keeps + + def test_the_rate_is_recorded_on_rows_that_pass(self, tmp_path: Path) -> None: + result = SampleRateFilterStage().process(_task(_wav(tmp_path, rate=44100))) + assert result != [] + assert result.data["sample_rate"] == 44100 + + def test_it_reads_the_header_and_never_decodes(self, tmp_path: Path) -> None: + """The rate is metadata sitting in the first bytes. Decoding to read it costs ~186x + more per file, and placing this stage before any decoding is what keeps rejected + rows from ever being decoded -- a decode here would forfeit exactly that.""" + stage = SampleRateFilterStage(allowed_sample_rates=[16000]) + path = _wav(tmp_path, rate=16000) + + def explode(*_args: object, **_kwargs: object) -> None: + msg = "decoded the audio to read a header value" + raise AssertionError(msg) + + with patch.object(srf.sf, "read", explode): + result = stage.process(_task(path)) + + assert result != [] + assert result.data["sample_rate"] == 16000 + + def test_a_resident_rate_avoids_touching_disk_entirely(self) -> None: + """A rate carried alongside resident audio describes audio this pipeline is holding, + so it is reused and the file is never opened.""" + stage = SampleRateFilterStage(allowed_sample_rates=[16000]) + task = AudioTask( + task_id="t", + dataset_name="d", + data={ + "audio_filepath": "/nonexistent/never-opened.wav", + "sample_rate": 16000, + "waveform": object(), + }, + ) + + result = stage.process(task) + + assert result != [] + assert result.data["sample_rate"] == 16000 + + def test_a_manifest_rate_with_no_resident_audio_is_verified_against_the_file(self, tmp_path: Path) -> None: + """``sample_rate`` is a standard manifest column, and a stale one used to decide the + filter outright: a genuinely 48 kHz file labelled 16000 was KEPT for a 16 kHz-only + corpus and then re-stamped with the wrong rate, so the model downstream silently got + pitch-shifted audio. With nothing resident to back the number, the header wins.""" + path = _wav(tmp_path, rate=48000, name="mislabelled.wav") + task = AudioTask( + task_id="t", + dataset_name="d", + data={"audio_filepath": path, "sample_rate": 16000}, + ) + + assert SampleRateFilterStage(allowed_sample_rates=[16000]).process(task) == [] + + task = AudioTask( + task_id="t", + dataset_name="d", + data={"audio_filepath": path, "sample_rate": 16000}, + ) + kept = SampleRateFilterStage(allowed_sample_rates=[48000]).process(task) + assert kept != [] + assert kept.data["sample_rate"] == 48000, "the recorded rate is the measured one" + + def test_an_unverifiable_rate_is_used_rather_than_dropping_the_row(self) -> None: + """No resident audio and no path leaves nothing to check against. Filtering on the + declared rate beats discarding a row that may well be fine.""" + task = AudioTask(task_id="t", dataset_name="d", data={"sample_rate": 16000}) + result = SampleRateFilterStage(allowed_sample_rates=[16000]).process(task) + assert result != [] + assert result.data["sample_rate"] == 16000 + + def test_an_unreadable_row_is_dropped_not_crashed(self, tmp_path: Path) -> None: + result = SampleRateFilterStage().process(_task(str(tmp_path / "missing.wav"))) + assert result == [] + + def test_a_home_relative_path_is_read_rather_than_dropped( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A manifest written with ``~/audio/x.wav`` must not vanish through this stage. + + ``sf.info`` does not expand ``~``, and the read failure above drops the row rather than + raising -- so before ``expanduser`` was applied a corpus addressed that way was silently + emptied here, while ChannelCountStage beside it read the same paths fine. + """ + monkeypatch.setenv("HOME", str(tmp_path)) + _wav(tmp_path, rate=16000, name="home.wav") + out = SampleRateFilterStage(allowed_sample_rates=[16000]).process(_task("~/home.wav")) + assert out != [], "row dropped: the '~' path was never expanded" + assert out.data["sample_rate"] == 16000 + + def test_an_empty_allow_list_is_rejected_at_construction(self) -> None: + """``[]`` would silently discard the entire corpus; None means "no constraint".""" + with pytest.raises(ValueError, match="at least one rate"): + SampleRateFilterStage(allowed_sample_rates=[]) + + def test_an_inverted_range_is_rejected_at_construction(self) -> None: + with pytest.raises(ValueError, match="nothing can pass"): + SampleRateFilterStage(min_sample_rate=48000, max_sample_rate=16000) + + +class TestRowDroppingIsDeclared: + """A stage that drops rows has to say ``cardinality="filter"``, because that is the only + thing that puts a filter seam in the semantic review packet. Left at the ``1:1`` default, + a reviewer is never told the corpus can shrink here and nobody asks how much of it + survives -- a run over a 90%-telephony corpus then reports success on 10% of the data. + """ + + def test_sample_rate_selection_declares_itself_a_filter(self) -> None: + from nemo_curator.stages.audio import agent as foundation + + contract = foundation.build_contract(SampleRateFilterStage(min_sample_rate=16000)) + assert contract.cardinality == "filter" + + def test_channel_selection_declares_itself_a_filter(self) -> None: + from nemo_curator.stages.audio import agent as foundation + + contract = foundation.build_contract(ChannelCountStage(action="filter", allowed_channels=[1])) + assert contract.cardinality == "filter" + + def test_recording_a_count_declares_that_it_drops_nothing(self) -> None: + from nemo_curator.stages.audio import agent as foundation + + assert foundation.build_contract(ChannelCountStage()).cardinality == "1:1" + + def test_channel_conversion_declares_a_filter_only_when_it_can_refuse(self) -> None: + """Downmixing to mono always succeeds. Any other target refuses the conversions it + cannot do correctly (N > target > 1) and drops those rows.""" + from nemo_curator.stages.audio import agent as foundation + + assert foundation.build_contract(_convert(target_channels=1)).cardinality == "1:1" + assert foundation.build_contract(_convert(target_channels=2)).cardinality == "filter" + + def test_both_row_cardinalities_are_advertised_whichever_action_is_set(self) -> None: + """One stage, one card, one resolved cardinality -- so the resolved contract alone + cannot say the stage is capable of the other behaviour. ``cardinality_options`` is what + tells a planner reading an annotating instance that this stage can also drop rows. + """ + from nemo_curator.stages.audio import agent as foundation + + for stage in (ChannelCountStage(), ChannelCountStage(action="filter", allowed_channels=[1]), _convert()): + assert foundation.build_contract(stage).cardinality_options == ["filter", "annotate"] + + +class TestTheyCompose: + def test_rate_selection_then_channel_conversion(self, tmp_path: Path) -> None: + """Independent policies: accept a range of rates, and separately require mono. + + Neither stage constrains the other, so a 22.05 kHz corpus can be taken to mono + without also having to declare 22050 the only acceptable rate. + """ + path = _wav(tmp_path, channels=2, rate=22050) + + selected = SampleRateFilterStage(min_sample_rate=16000).process(_task(path)) + assert selected != [] + + converted = _convert(target_channels=1).process(selected) + assert converted != [] + assert converted.data["num_channels"] == 1 + assert converted.data["sample_rate"] == 22050, "selection does not resample" + + def test_selecting_48k_mono_rewrites_nothing(self, tmp_path: Path) -> None: + """Drop-only on both axes, which needs no conversion stage at all: the rate side + selects, the channel side selects, and every surviving row is the original file. + """ + kept = _wav(tmp_path, channels=1, rate=48000, name="keep.wav") + wrong_channels = _wav(tmp_path, channels=2, rate=48000, name="stereo.wav") + wrong_rate = _wav(tmp_path, channels=1, rate=16000, name="slow.wav") + + def survives(path: str) -> bool: + row = SampleRateFilterStage(allowed_sample_rates=[48000]).process(_task(path)) + if row == []: + return False + return ChannelCountStage(action="filter", allowed_channels=[1]).process(row) != [] + + assert survives(kept) + assert not survives(wrong_channels) + assert not survives(wrong_rate) diff --git a/tests/stages/audio/preprocessing/test_concatenation.py b/tests/stages/audio/preprocessing/test_concatenation.py index 5a76ddae32..5cafc96f25 100644 --- a/tests/stages/audio/preprocessing/test_concatenation.py +++ b/tests/stages/audio/preprocessing/test_concatenation.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os + import pytest import torch @@ -108,3 +110,41 @@ def test_empty_segments_returns_empty(self) -> None: stage = SegmentConcatenationStage() result = stage.process(task) assert result == [] + + # --- output residency (write-to-disk extension) --- + + def test_default_output_is_in_memory_only(self) -> None: + """Regression: default config keeps the in-memory waveform and writes no path.""" + result = SegmentConcatenationStage().process(_make_nested_task([_make_segment_dict(duration_ms=1000)])) + assert "waveform" in result.data + assert "audio_filepath" not in result.data + + def test_write_to_disk_persists_and_sets_path(self, tmp_path) -> None: # noqa: ANN001 + out_dir = tmp_path / "concat" + stage = SegmentConcatenationStage(write_to_disk=True, output_dir=str(out_dir)) + segments = [ + _make_segment_dict(duration_ms=1000, segment_num=0), + _make_segment_dict(duration_ms=1000, segment_num=1), + ] + result = stage.process(_make_nested_task(segments)) + assert isinstance(result, AudioTask) + # default keep_waveform_in_task=True -> both the waveform AND a written path + assert "waveform" in result.data + assert "audio_filepath" in result.data + assert os.path.exists(result.data["audio_filepath"]) + + def test_write_to_disk_only_drops_waveform(self, tmp_path) -> None: # noqa: ANN001 + stage = SegmentConcatenationStage( + write_to_disk=True, output_dir=str(tmp_path / "c"), keep_waveform_in_task=False + ) + result = stage.process(_make_nested_task([_make_segment_dict(duration_ms=1000)])) + assert "waveform" not in result.data + assert os.path.exists(result.data["audio_filepath"]) + + def test_requires_output_dir_when_write_to_disk(self) -> None: + with pytest.raises(ValueError, match="output_dir"): + SegmentConcatenationStage(write_to_disk=True) + + def test_requires_at_least_one_output_sink(self) -> None: + with pytest.raises(ValueError, match="keep_waveform_in_task or write_to_disk"): + SegmentConcatenationStage(keep_waveform_in_task=False) diff --git a/tests/stages/audio/preprocessing/test_mono_conversion.py b/tests/stages/audio/preprocessing/test_mono_conversion.py index 9ba9e45389..790d2a945e 100644 --- a/tests/stages/audio/preprocessing/test_mono_conversion.py +++ b/tests/stages/audio/preprocessing/test_mono_conversion.py @@ -107,3 +107,60 @@ def test_read_exception_skipped(self, tmp_path: Path) -> None: result = stage.process(task) assert result == [] + + +class TestMonoOutputGatingAndResidency: + """Which destination keys appear, and which input ``auto`` residency picks. + + Lifted from tests/stages/audio/test_agent_simulation_pipelines.py: it drives only this + stage, so it belongs beside the rest of MonoConversionStage's behaviour rather than in an + agent-simulation file, where it was the sole coverage of these two knobs. + """ + + def test_auto_residency_uses_the_tensor_and_never_reads_disk(self, tmp_path: Path) -> None: + """A resident waveform wins under ``auto``, even pointing at a path that cannot be read.""" + stage = MonoConversionStage( + audio_filepath_key="agent_audio_path", + waveform_key="agent_waveform", + sample_rate_key="agent_sr", + output_audio_filepath_key="agent_mono_path", + output_sample_rate=16000, + input_residency="auto", + keep_waveform_in_task=True, + write_to_disk=False, + ) + task = AudioTask( + dataset_name="agent", + data={ + "agent_audio_path": str(tmp_path / "does_not_exist.wav"), + "agent_waveform": torch.stack([torch.linspace(-0.25, 0.25, 8000)] * 2), + "agent_sr": 16000, + }, + ) + + result = stage.process(task) + + assert isinstance(result, AudioTask) + assert result.data["agent_waveform"].shape[0] == 1, "the stereo tensor was mixed down in memory" + assert "agent_mono_path" not in result.data, "disk-path key must be absent when write_to_disk=False" + + def test_disk_only_output_omits_the_waveform_key(self, tmp_path: Path) -> None: + """``keep_waveform_in_task=False`` must drop the tensor rather than leave it stale.""" + stage = MonoConversionStage( + audio_filepath_key="agent_audio_path", + waveform_key="agent_waveform", + sample_rate_key="agent_sr", + output_audio_filepath_key="agent_mono_path", + output_sample_rate=16000, + input_residency="file", + keep_waveform_in_task=False, + write_to_disk=True, + output_dir=str(tmp_path / "mono_out"), + ) + task = AudioTask(dataset_name="agent", data={"agent_audio_path": str(tmp_path / "src.wav")}) + + with patch(MOCK_TARGET, return_value=(torch.randn(1, 16000), 16000)), patch(MOCK_EXISTS, return_value=True): + result = stage.process(task) + + assert "agent_mono_path" in result.data, "disk-path key must be present when write_to_disk=True" + assert "agent_waveform" not in result.data, "tensor must be omitted when keep_waveform_in_task=False" diff --git a/tests/stages/audio/test_common.py b/tests/stages/audio/test_common.py index f597982aee..21c7ba2aff 100644 --- a/tests/stages/audio/test_common.py +++ b/tests/stages/audio/test_common.py @@ -13,9 +13,10 @@ # limitations under the License. """Tests for common audio stages: GetAudioDurationStage, PreserveByValueStage, -ManifestReaderStage, ManifestReader, and ManifestWriterStage.""" +ManifestReaderStage, ManifestReader, ManifestWriterStage, and ManifestCheckpointStage.""" import json +from itertools import product from pathlib import Path from unittest import mock @@ -28,9 +29,11 @@ from nemo_curator.stages.audio.alm import ALMDataBuilderStage, ALMDataOverlapStage from nemo_curator.stages.audio.common import ( GetAudioDurationStage, + ManifestCheckpointStage, ManifestReader, ManifestReaderStage, ManifestWriterStage, + PreserveByValueConditionsStage, PreserveByValueStage, ensure_mono, ensure_waveform_2d, @@ -71,6 +74,7 @@ def test_preserve_by_value_process_raises_not_implemented() -> None: def test_preserve_by_value_process_batch_raises_on_missing_column() -> None: stage = PreserveByValueStage(input_value_key="wer", target_value=50, operator="le") + assert stage.missing_value_policy == "error" with pytest.raises(ValueError, match="failed validation"): stage.process_batch([AudioTask(data={"text": "hello"})]) @@ -96,12 +100,345 @@ def test_preserve_by_value_lt() -> None: def test_preserve_by_value_ge() -> None: - stage = PreserveByValueStage(input_value_key="v", target_value=10, operator="ge") + stage = PreserveByValueStage(input_value_key="v", target_value=10.0, operator="ge") assert len(stage.process_batch([AudioTask(data={"v": 9})])) == 0 assert len(stage.process_batch([AudioTask(data={"v": 10})])) == 1 assert len(stage.process_batch([AudioTask(data={"v": 11})])) == 1 +def test_preserve_by_value_contract_accepts_float_targets_and_exposes_policy() -> None: + from nemo_curator.stages.audio._agent._agent_registry import stage_params + + params = {param.name: param for param in stage_params(PreserveByValueStage)} + + assert params["target_value"].type == "float | str" + assert params["missing_value_policy"].default == "error" + assert params["missing_value_policy"].choices == ["error", "drop"] + + +def test_preserve_by_value_drop_policy_drops_only_missing_or_failing_rows() -> None: + stage = PreserveByValueStage( + input_value_key="score", + target_value=3.5, + operator="ge", + missing_value_policy="drop", + ) + tasks = [ + AudioTask(data={"id": "pass", "score": 4.0}), + AudioTask(data={"id": "fail", "score": 3.0}), + AudioTask(data={"id": "missing"}), + ] + + assert [task.data["id"] for task in stage.process_batch(tasks)] == ["pass"] + + +def test_compound_preserve_uses_and_semantics_and_drops_missing() -> None: + stage = PreserveByValueConditionsStage( + conditions=[ + {"input_value_key": "noise", "target_value": 4.0, "operator": "ge"}, + {"input_value_key": "ovrl", "target_value": 3.5, "operator": "ge"}, + ], + missing_value_policy="drop", + ) + tasks = [ + AudioTask(data={"id": "pass", "noise": 4.1, "ovrl": 3.6}), + AudioTask(data={"id": "noise_fail", "noise": 3.9, "ovrl": 4.0}), + AudioTask(data={"id": "ovrl_fail", "noise": 4.5, "ovrl": 3.4}), + AudioTask(data={"id": "missing", "noise": 4.5}), + ] + + assert [task.data["id"] for task in stage.process_batch(tasks)] == ["pass"] + assert stage.normalized_conditions == ( + {"input_value_key": "noise", "target_value": 4.0, "operator": "ge"}, + {"input_value_key": "ovrl", "target_value": 3.5, "operator": "ge"}, + ) + + +@pytest.mark.parametrize("condition_count", [1, 2, 4]) +@pytest.mark.parametrize("condition_logic", ["and", "or"]) +def test_compound_preserve_top_level_truth_tables( + condition_count: int, + condition_logic: str, +) -> None: + conditions = [ + {"input_value_key": f"c{index}", "target_value": True, "operator": "eq"} for index in range(condition_count) + ] + combinations = list(product([False, True], repeat=condition_count)) + tasks = [ + AudioTask( + data={ + "id": combination, + **{f"c{index}": value for index, value in enumerate(combination)}, + } + ) + for combination in combinations + ] + expected = [ + combination + for combination in combinations + if (all(combination) if condition_logic == "and" else any(combination)) + ] + + result = PreserveByValueConditionsStage( + conditions, + condition_logic=condition_logic, + ).process_batch(tasks) + + assert [task.data["id"] for task in result] == expected + + +@pytest.mark.parametrize("condition_count", [1, 2, 4]) +@pytest.mark.parametrize("condition_logic", ["and", "or"]) +def test_compound_preserve_nested_truth_tables_with_arbitrary_items_key( + condition_count: int, + condition_logic: str, +) -> None: + conditions = [ + {"input_value_key": f"c{index}", "target_value": True, "operator": "eq"} for index in range(condition_count) + ] + combinations = list(product([False, True], repeat=condition_count)) + children = [ + { + "id": combination, + **{f"c{index}": value for index, value in enumerate(combination)}, + } + for combination in combinations + ] + parent = AudioTask(data={"custom_children": children}) + expected = [ + combination + for combination in combinations + if (all(combination) if condition_logic == "and" else any(combination)) + ] + + result = PreserveByValueConditionsStage( + conditions, + items_key="custom_children", + condition_logic=condition_logic, + drop_parent_if_empty=False, + ).process_batch([parent]) + + assert result == [parent] + assert [child["id"] for child in parent.data["custom_children"]] == expected + + +def test_compound_preserve_condition_logic_defaults_to_and_and_rejects_invalid() -> None: + conditions = [ + {"input_value_key": "left", "target_value": True, "operator": "eq"}, + {"input_value_key": "right", "target_value": True, "operator": "eq"}, + ] + stage = PreserveByValueConditionsStage(conditions) + + assert stage.condition_logic == "and" + assert stage.process_batch([AudioTask(data={"left": True, "right": False})]) == [] + with pytest.raises(ValueError, match="condition_logic must be 'and' or 'or'"): + PreserveByValueConditionsStage(conditions, condition_logic="xor") + + +@pytest.mark.parametrize("missing_value_policy", ["error", "drop"]) +def test_compound_preserve_or_never_skips_a_missing_top_level_condition( + missing_value_policy: str, +) -> None: + stage = PreserveByValueConditionsStage( + [ + {"input_value_key": "present", "target_value": True, "operator": "eq"}, + {"input_value_key": "missing", "target_value": True, "operator": "eq"}, + ], + missing_value_policy=missing_value_policy, + condition_logic="or", + ) + task = AudioTask(data={"present": True}) + + if missing_value_policy == "error": + with pytest.raises(ValueError, match="failed validation"): + stage.process_batch([task]) + else: + assert stage.process_batch([task]) == [] + + +@pytest.mark.parametrize("missing_value_policy", ["error", "drop"]) +def test_compound_preserve_or_never_skips_a_missing_nested_condition( + missing_value_policy: str, +) -> None: + stage = PreserveByValueConditionsStage( + [ + {"input_value_key": "present", "target_value": True, "operator": "eq"}, + {"input_value_key": "missing", "target_value": True, "operator": "eq"}, + ], + items_key="children", + missing_value_policy=missing_value_policy, + condition_logic="or", + ) + parent = AudioTask(data={"children": [{"present": True}]}) + + if missing_value_policy == "error": + with pytest.raises(ValueError, match="missing condition key 'missing'"): + stage.process_batch([parent]) + else: + assert stage.process_batch([parent]) == [] + assert parent.data["children"] == [] + + +def test_compound_preserve_mapping_form_and_default_missing_error() -> None: + stage = PreserveByValueConditionsStage( + conditions={ + "noise": {"target_value": 4.0, "operator": "ge"}, + "kind": "speech", + } + ) + + assert stage.process_batch([AudioTask(data={"noise": 4.2, "kind": "speech"})]) + with pytest.raises(ValueError, match="failed validation"): + stage.process_batch([AudioTask(data={"noise": 4.2})]) + + +def test_compound_preserve_filters_arbitrary_one_level_items_key_by_reference() -> None: + passing = {"id": "pass", "quality": 4.2, "metadata": {"speaker": "a"}} + failing = {"id": "fail", "quality": 2.0, "metadata": {"speaker": "b"}} + parent = AudioTask(data={"recording": "r1", "clips": [passing, failing]}) + stage = PreserveByValueConditionsStage( + [{"input_value_key": "quality", "target_value": 3.5, "operator": "ge"}], + items_key="clips", + ) + + result = stage.process_batch([parent]) + + assert result == [parent] + assert parent.data["recording"] == "r1" + assert parent.data["clips"] == [passing] + assert parent.data["clips"][0] is passing + assert parent.data["clips"][0]["metadata"] is passing["metadata"] + + +@pytest.mark.parametrize("condition_logic", ["and", "or"]) +@pytest.mark.parametrize( + ("drop_parent_if_empty", "expected_count"), + [(True, 0), (False, 1)], +) +def test_compound_preserve_nested_empty_parent_policy( + drop_parent_if_empty: bool, + expected_count: int, + condition_logic: str, +) -> None: + parent = AudioTask(data={"windows": [{"score": 1.0}]}) + stage = PreserveByValueConditionsStage( + [{"input_value_key": "score", "target_value": 2.0, "operator": "ge"}], + items_key="windows", + drop_parent_if_empty=drop_parent_if_empty, + condition_logic=condition_logic, + ) + + result = stage.process_batch([parent]) + + assert len(result) == expected_count + assert parent.data["windows"] == [] + + +@pytest.mark.parametrize("condition_logic", ["and", "or"]) +@pytest.mark.parametrize( + ("data", "error_type", "message"), + [ + ({"clips": {}}, TypeError, "must contain a list"), + ({"clips": [{"score": 4.0}, "not-a-mapping"]}, TypeError, "child 1 must be mapping-like"), + ], +) +def test_compound_preserve_rejects_malformed_nested_structure_without_mutation( + data: dict, + error_type: type[Exception], + message: str, + condition_logic: str, +) -> None: + original_items = data.get("clips") + stage = PreserveByValueConditionsStage( + [{"input_value_key": "score", "target_value": 3.5, "operator": "ge"}], + items_key="clips", + missing_value_policy="drop", + condition_logic=condition_logic, + ) + + with pytest.raises(error_type, match=message): + stage.process_batch([AudioTask(data=data)]) + + assert data.get("clips") is original_items + + +@pytest.mark.parametrize("missing_value_policy", ["error", "drop"]) +@pytest.mark.parametrize("condition_logic", ["and", "or"]) +def test_compound_preserve_missing_nested_container_is_always_structural_error( + missing_value_policy: str, + condition_logic: str, +) -> None: + stage = PreserveByValueConditionsStage( + [{"input_value_key": "score", "target_value": 3.5, "operator": "ge"}], + items_key="clips", + missing_value_policy=missing_value_policy, + condition_logic=condition_logic, + ) + + with pytest.raises(ValueError, match="missing nested items_key 'clips'"): + stage.process_batch([AudioTask(data={"other": []})]) + + +def test_compound_preserve_nested_missing_condition_key_error_vs_drop() -> None: + condition = [{"input_value_key": "score", "target_value": 3.5, "operator": "ge"}] + missing = {"id": "missing", "nested": {"score": 5.0}} + passing = {"id": "pass", "score": 4.0} + + with pytest.raises(ValueError, match="child 0 is missing condition key 'score'"): + PreserveByValueConditionsStage( + condition, + items_key="candidates", + ).process_batch([AudioTask(data={"candidates": [missing, passing]})]) + + parent = AudioTask(data={"candidates": [missing, passing]}) + result = PreserveByValueConditionsStage( + condition, + items_key="candidates", + missing_value_policy="drop", + ).process_batch([parent]) + assert result == [parent] + assert parent.data["candidates"] == [passing] + + +def test_compound_preserve_nested_contract_uses_only_top_level_container_key() -> None: + from nemo_curator.stages.audio._agent._agent_registry import stage_params + + stage = PreserveByValueConditionsStage( + [{"input_value_key": "score", "target_value": 3.5, "operator": "ge"}], + items_key="candidates", + drop_parent_if_empty=False, + ) + contract = stage.describe() + params = {param.name: param for param in stage_params(PreserveByValueConditionsStage)} + + assert contract.reads.data_keys == ["candidates"] + assert contract.writes.data_keys == ["candidates"] + assert contract.reads.segment_data_keys == [] + assert contract.writes.segment_data_keys == [] + assert contract.iteration_key == "candidates" + assert contract.cardinality == "1:1 nested-list" + assert contract.gates.per_row_independent is True + assert params["items_key"].default is None + assert params["drop_parent_if_empty"].default is True + assert params["condition_logic"].default == "and" + assert params["condition_logic"].choices == ["and", "or"] + + dropping_contract = PreserveByValueConditionsStage( + [{"input_value_key": "score", "target_value": 3.5, "operator": "ge"}], + items_key="candidates", + ).describe() + assert dropping_contract.cardinality == "filter" + assert dropping_contract.iteration_key is None + assert "one-level" in dropping_contract.description + assert "AND" in dropping_contract.description + + or_contract = PreserveByValueConditionsStage( + [{"input_value_key": "score", "target_value": 3.5, "operator": "ge"}], + condition_logic="or", + ).describe() + assert "OR" in or_contract.description + + # --------------------------------------------------------------------------- # GetAudioDurationStage # --------------------------------------------------------------------------- @@ -149,6 +486,43 @@ def test_get_audio_duration_error_sets_minus_one(tmp_path: Path) -> None: assert result.data["duration"] == -1.0 +def test_get_audio_duration_waveform_residency() -> None: + """input_residency='waveform' computes duration from samples/sample_rate (no file).""" + import torch + + stage = GetAudioDurationStage(input_residency="waveform") + stage.setup() + result = stage.process(AudioTask(data={"waveform": torch.zeros(1, 16000 * 3), "sample_rate": 16000})) + assert result.data["duration"] == 3.0 + + +def test_get_audio_duration_auto_prefers_waveform() -> None: + import torch + + stage = GetAudioDurationStage(input_residency="auto") + stage.setup() + result = stage.process(AudioTask(data={"waveform": torch.zeros(1, 16000), "sample_rate": 16000})) + assert result.data["duration"] == 1.0 + + +def test_get_audio_duration_default_rejects_waveform_only() -> None: + """Regression: default residency is 'file'; a waveform-only task is not valid input.""" + import torch + + stage = GetAudioDurationStage() + assert stage.input_residency == "file" + assert stage.validate_input(AudioTask(data={"waveform": torch.zeros(1, 16000), "sample_rate": 16000})) is False + assert stage.validate_input(AudioTask(data={"audio_filepath": "/a.wav"})) is True + + +def test_get_audio_duration_waveform_validate() -> None: + import torch + + stage = GetAudioDurationStage(input_residency="waveform") + assert stage.validate_input(AudioTask(data={"waveform": torch.zeros(1, 16000), "sample_rate": 16000})) is True + assert stage.validate_input(AudioTask(data={"audio_filepath": "/a.wav"})) is False + + # --------------------------------------------------------------------------- # ManifestReaderStage # --------------------------------------------------------------------------- @@ -489,6 +863,129 @@ def test_xenna_stage_spec(self, tmp_path: Path) -> None: assert writer.xenna_stage_spec() == {} +class TestManifestCheckpointStage: + """Focused unit tests for the reusable metadata checkpoint.""" + + def test_setup_atomically_refuses_to_overwrite_existing_checkpoint(self, tmp_path: Path) -> None: + out = tmp_path / "checkpoint.jsonl" + out.write_bytes(b"retained artifact\n") + checkpoint = ManifestCheckpointStage(output_path=str(out)) + + with pytest.raises(FileExistsError, match="refuses to overwrite"): + checkpoint.setup() + + assert out.read_bytes() == b"retained artifact\n" + assert not Path(f"{out}._RETRY_OWNER").exists() + + def test_setup_refuses_stale_completion_marker_without_leaving_output(self, tmp_path: Path) -> None: + out = tmp_path / "checkpoint.jsonl" + Path(f"{out}._COMPLETE").write_text("stale", encoding="utf-8") + checkpoint = ManifestCheckpointStage(output_path=str(out)) + + with pytest.raises(FileExistsError, match="completion marker"): + checkpoint.setup() + + assert not out.exists() + + def test_retry_reset_removes_only_owned_partial_and_reserves_cleanly( + self, + tmp_path: Path, + ) -> None: + out = tmp_path / "checkpoint.jsonl" + checkpoint = ManifestCheckpointStage(output_path=str(out)) + checkpoint.setup() + checkpoint.process(AudioTask(data={"attempt": 1})) + + checkpoint.reset_for_retry() + + assert not out.exists() + assert checkpoint._checkpoint_rows_written == 0 + assert checkpoint._checkpoint_bytes_written == 0 + checkpoint.setup() + checkpoint.process(AudioTask(data={"attempt": 2})) + assert out.read_text(encoding="utf-8") == '{"attempt": 2}\n' + + def test_retry_reset_refuses_completed_checkpoint(self, tmp_path: Path) -> None: + out = tmp_path / "checkpoint.jsonl" + checkpoint = ManifestCheckpointStage(output_path=str(out)) + checkpoint.setup() + checkpoint.process(AudioTask(data={"retained": True})) + before = out.read_bytes() + Path(f"{out}._COMPLETE").write_text("complete", encoding="utf-8") + + with pytest.raises(FileExistsError, match="completion marker"): + checkpoint.reset_for_retry() + + assert out.read_bytes() == before + + def test_retry_reset_refuses_preexisting_unowned_checkpoint( + self, + tmp_path: Path, + ) -> None: + out = tmp_path / "checkpoint.jsonl" + out.write_text("user file\n", encoding="utf-8") + checkpoint = ManifestCheckpointStage(output_path=str(out)) + + with pytest.raises(FileExistsError, match="did not reserve"): + checkpoint.reset_for_retry() + + assert out.read_text(encoding="utf-8") == "user file\n" + + def test_retry_reset_refuses_replaced_reservation(self, tmp_path: Path) -> None: + out = tmp_path / "checkpoint.jsonl" + checkpoint = ManifestCheckpointStage(output_path=str(out)) + checkpoint.setup() + out.unlink() + out.write_text("replacement\n", encoding="utf-8") + + with pytest.raises(FileExistsError, match="no longer its exact reservation"): + checkpoint.reset_for_retry() + + assert out.read_text(encoding="utf-8") == "replacement\n" + + def test_configured_contract_is_audio_pass_through_with_checkpoint_gates(self, tmp_path: Path) -> None: + from nemo_curator.stages.audio._agent._agent_registry import build_contract + + checkpoint = ManifestCheckpointStage(output_path=str(tmp_path / "checkpoint.jsonl")) + contract = build_contract(checkpoint) + params = {parameter.name: parameter for parameter in contract.params} + + assert checkpoint.name == "manifest_checkpoint" + assert checkpoint.name != ManifestWriterStage(output_path=str(tmp_path / "manifest.jsonl")).name + assert checkpoint.num_workers() == 1 + assert contract.accepts_task_type == "AudioTask" + assert contract.produces_task_type == "AudioTask" + assert contract.gates.writes_to_disk is True + assert contract.gates.output_path_params == ["output_path"] + assert contract.gates.requires_serializable_input is True + assert contract.gates.per_row_independent is True + assert contract.gates.lifecycle_side_effects is True + assert params["output_path"].required is True + assert "max_bytes" not in params + assert params["retention_sec"].default == 0 + assert params["owner"].choices == ["user", "project"] + assert params["planning_provenance"].default is None + + @pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"retention_sec": -1}, "retention_sec"), + ({"owner": "nobody"}, "owner"), + ({"output_path": "s3://bucket/checkpoint.jsonl"}, "plain local path"), + ({"output_path": "file:///tmp/checkpoint.jsonl"}, "plain local path"), + ], + ) + def test_rejects_invalid_checkpoint_policy( + self, + tmp_path: Path, + kwargs: dict[str, object], + message: str, + ) -> None: + params = {"output_path": str(tmp_path / "checkpoint.jsonl"), **kwargs} + with pytest.raises(ValueError, match=message): + ManifestCheckpointStage(**params) + + class TestManifestWriterRoundTrip: """Round-trip test: write with writer, read back and verify.""" @@ -563,3 +1060,21 @@ def test_resolve_model_path(tmp_path: Path) -> None: (module_dir / "model.bin").write_bytes(b"\x00") result = resolve_model_path("model.bin", str(tmp_path / "ref.py"), "sub") assert result == str(module_dir / "model.bin") + + +# Lifted from tests/stages/audio/test_agent_simulation_pipelines.py: ManifestWriterStage +# lives in common.py, and this was its only truncate-on-rerun coverage. +def test_agent_manifest_writer_truncates_on_setup(tmp_path: Path) -> None: + """A fresh run (setup) truncates the output so reruns do not accumulate duplicates.""" + out_path = tmp_path / "manifest.jsonl" + writer = ManifestWriterStage(output_path=str(out_path)) + task = AudioTask(dataset_name="t", data={"audio_filepath": "src.wav", "text": "row"}) + + writer.setup() + writer.process(task) + writer.process(task) + assert len(out_path.read_text(encoding="utf-8").strip().splitlines()) == 2 # appends within a run + + writer.setup() # new run truncates + writer.process(task) + assert len(out_path.read_text(encoding="utf-8").strip().splitlines()) == 1 diff --git a/tests/stages/audio/test_create_manifest_audio_folder.py b/tests/stages/audio/test_create_manifest_audio_folder.py new file mode 100644 index 0000000000..da408965ab --- /dev/null +++ b/tests/stages/audio/test_create_manifest_audio_folder.py @@ -0,0 +1,87 @@ +# 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. + +"""Unit tests for CreateInitialManifestAudioFolderStage (generic local-folder source).""" + +import os + +import pytest + +from nemo_curator.stages.audio.common import CreateInitialManifestAudioFolderStage + + +def _touch(root: str, rel: str) -> None: + path = os.path.join(root, rel) + os.makedirs(os.path.dirname(path), exist_ok=True) + open(path, "wb").close() # placeholder; the stage only collects paths, never decodes + + +class TestCreateInitialManifestAudioFolderStage: + def test_recursive_collects_audio_only_one_task_each(self, tmp_path) -> None: # noqa: ANN001 + root = str(tmp_path) + for rel in ["a.wav", "b.FLAC", "notes.txt", "sub/c.mp3"]: + _touch(root, rel) + tasks = CreateInitialManifestAudioFolderStage(data_dir=root).process(None) + names = sorted(os.path.basename(t.data["audio_filepath"]) for t in tasks) + assert names == ["a.wav", "b.FLAC", "c.mp3"] # .txt excluded; subdir included; ext match is case-insensitive + assert all(t.data["audio_item_id"] for t in tasks) + assert all(os.path.isabs(t.data["audio_filepath"]) for t in tasks) + + def test_same_filename_in_two_folders_gets_two_ids(self, tmp_path) -> None: # noqa: ANN001 + root = str(tmp_path) + for rel in ["spk1/utt1.wav", "spk2/utt1.wav"]: + _touch(root, rel) + + tasks = CreateInitialManifestAudioFolderStage(data_dir=root).process(None) + ids = sorted(t.data["audio_item_id"] for t in tasks) + + assert ids == ["spk1__utt1", "spk2__utt1"], ids + + def test_a_flat_folder_keeps_the_plain_ids_it_always_had(self, tmp_path) -> None: # noqa: ANN001 + """relpath IS the basename for a flat corpus, so those ids must not move.""" + root = str(tmp_path) + for rel in ["a.wav", "b.wav"]: + _touch(root, rel) + + tasks = CreateInitialManifestAudioFolderStage(data_dir=root).process(None) + + assert sorted(t.data["audio_item_id"] for t in tasks) == ["a", "b"] + + def test_non_recursive_and_max_samples(self, tmp_path) -> None: # noqa: ANN001 + root = str(tmp_path) + for rel in ["a.wav", "b.wav", "sub/c.wav"]: + _touch(root, rel) + tasks = CreateInitialManifestAudioFolderStage(data_dir=root, recursive=False, max_samples=1).process(None) + assert len(tasks) == 1 # sub/ excluded (non-recursive), capped to 1 + assert tasks[0].data["audio_filepath"].endswith(".wav") + + def test_extension_filter(self, tmp_path) -> None: # noqa: ANN001 + root = str(tmp_path) + for rel in ["a.wav", "b.mp3"]: + _touch(root, rel) + tasks = CreateInitialManifestAudioFolderStage(data_dir=root, extensions=[".mp3"]).process(None) + assert [os.path.basename(t.data["audio_filepath"]) for t in tasks] == ["b.mp3"] + + def test_missing_dir_returns_empty(self, tmp_path) -> None: # noqa: ANN001 + tasks = CreateInitialManifestAudioFolderStage(data_dir=str(tmp_path / "nope")).process(None) + assert tasks == [] + + def test_requires_data_dir(self) -> None: + with pytest.raises(ValueError): # noqa: PT011 + CreateInitialManifestAudioFolderStage(data_dir="") + + def test_contract_writes_filepath_and_no_disk_write(self) -> None: + c = CreateInitialManifestAudioFolderStage(data_dir="/tmp").describe() # noqa: S108 + assert "audio_filepath" in c.writes.data_keys + assert c.gates.writes_to_disk is False # references existing files; no disk write From 18336240e5da6563146993df409f299257a3ea2e Mon Sep 17 00:00:00 2001 From: Shubham Bhawsar Date: Fri, 4 Sep 2026 09:59:01 +0000 Subject: [PATCH 02/25] fix(audio): address safe agent foundation review findings Signed-off-by: Shubham Bhawsar --- nemo_curator/stages/audio/__init__.py | 2 + .../stages/audio/_agent/_composite.py | 14 ++- .../stages/audio/_agent/_residency.py | 14 ++- nemo_curator/stages/audio/agent.py | 2 + nemo_curator/stages/audio/common.py | 13 +++ .../test_agent_foundation_regressions.py | 100 ++++++++++++++++++ 6 files changed, 135 insertions(+), 10 deletions(-) create mode 100644 tests/stages/audio/_agent/test_agent_foundation_regressions.py diff --git a/nemo_curator/stages/audio/__init__.py b/nemo_curator/stages/audio/__init__.py index 2eee308dcf..f1d047ea90 100644 --- a/nemo_curator/stages/audio/__init__.py +++ b/nemo_curator/stages/audio/__init__.py @@ -33,6 +33,7 @@ "AudioDataFilterStage": "nemo_curator.stages.audio.advanced_pipelines", "BandFilterStage": "nemo_curator.stages.audio.filtering", "ChannelCountStage": "nemo_curator.stages.audio.preprocessing", + "CreateInitialManifestAudioFolderStage": "nemo_curator.stages.audio.common", "GetAudioDurationStage": "nemo_curator.stages.audio.common", "ManifestCheckpointStage": "nemo_curator.stages.audio.common", "ManifestReader": "nemo_curator.stages.audio.common", @@ -56,6 +57,7 @@ "AudioDataFilterStage", "BandFilterStage", "ChannelCountStage", + "CreateInitialManifestAudioFolderStage", "GetAudioDurationStage", "ManifestCheckpointStage", "ManifestReader", diff --git a/nemo_curator/stages/audio/_agent/_composite.py b/nemo_curator/stages/audio/_agent/_composite.py index 181bd90ce3..b944c0a3d9 100644 --- a/nemo_curator/stages/audio/_agent/_composite.py +++ b/nemo_curator/stages/audio/_agent/_composite.py @@ -32,10 +32,14 @@ segments_key="diar_segments")`` yields ``SplitLongAudioStage(segments_key="diar_segments")`` and the check reflects what will run rather than what the defaults would have run. -When a composite cannot be expanded -- it raises, returns nothing, returns something that is not -a stage, or returns another composite (which the executor itself refuses) -- no leaf is invented -for it. It is reported in :attr:`Expansion.opaque` so callers can fall back to whatever they did -before rather than reason from a fabricated stage list. +When a composite cannot be expanded -- it raises, returns nothing, or returns something that is +not a stage -- no leaf is invented for it. It is reported in :attr:`Expansion.opaque` so callers +can fall back to whatever they did before rather than reason from a fabricated stage list. + +A composite the executor will refuse outright -- one that decomposes into a single stage, or one +that returns another composite -- is reported in :attr:`Expansion.unrunnable` instead. That is a +known failure rather than an unknown, and has to reach the caller as an error rather than as the +warning an opaque stage earns. """ from __future__ import annotations @@ -181,7 +185,7 @@ def expand_composites(stages: list[Any]) -> Expansion: # can tell perfectly well that it will run the outer composite and raise. nested = next((c for c in children if _nested_composite(c)), None) if nested is not None: - opaque[index] = ( + unrunnable[index] = ( f"decomposition returned another composite ({type(nested).__name__}); " "nested composition is not supported" ) diff --git a/nemo_curator/stages/audio/_agent/_residency.py b/nemo_curator/stages/audio/_agent/_residency.py index 35fb2f0460..bd9c256a6f 100644 --- a/nemo_curator/stages/audio/_agent/_residency.py +++ b/nemo_curator/stages/audio/_agent/_residency.py @@ -207,11 +207,11 @@ def _as_soundfile_array(waveform: Any) -> Any: # noqa: ANN401 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: + if waveform.shape[0] == 1: return waveform[0] - if channels < samples: - return waveform.T + # The shared representation is channel-first. SoundFile expects frames first, + # including for a valid but very short (channels > samples) waveform. + return waveform.T return waveform @@ -243,7 +243,11 @@ def write_audio_stable( os.makedirs(output_dir, exist_ok=True) digest = hashlib.sha256(arr.tobytes()) - digest.update(f"|{int(sample_rate)}".encode()) + # ``tobytes`` alone loses array shape: mono ``(1, 32000)`` and stereo + # ``(2, 16000)`` silence have identical flattened bytes. Include the + # representation SoundFile will write so distinct audio layouts cannot + # claim the same stable path. + digest.update(f"|{arr.shape!r}|{arr.dtype.str}|{int(sample_rate)}|wav".encode()) path = os.path.join(output_dir, f"{stem}{f'_{tag}' if tag else ''}_{digest.hexdigest()[:16]}.wav") # Write beside the target and rename, so a killed or concurrent writer cannot leave a # half-written file at a name the next run treats as finished. diff --git a/nemo_curator/stages/audio/agent.py b/nemo_curator/stages/audio/agent.py index 8675959f04..eba6830cbc 100644 --- a/nemo_curator/stages/audio/agent.py +++ b/nemo_curator/stages/audio/agent.py @@ -69,6 +69,7 @@ get_agent_ready_stage_class, list_agent_ready_stages, role_index, + unavailable_modules, ) from nemo_curator.stages.audio._agent._conformance import ( assert_contract_wellformed, @@ -101,5 +102,6 @@ "role_index", "static_contract", "to_json_schema", + "unavailable_modules", "validate_pipeline", ] diff --git a/nemo_curator/stages/audio/common.py b/nemo_curator/stages/audio/common.py index 2cf9f9fd35..99e6ea2dfc 100644 --- a/nemo_curator/stages/audio/common.py +++ b/nemo_curator/stages/audio/common.py @@ -778,6 +778,19 @@ class ManifestWriterStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): output_path: str name: str = "manifest_writer" + # ``output_path`` is required, so static discovery cannot instantiate this + # stage. Publish invariant sink behavior so an agent never mistakes it for + # a pure pass-through before it has configuration values. + AGENT_STATIC: ClassVar[StaticHints] = StaticHints( + gates=Gates( + writes_to_disk=True, + output_path_params=["output_path"], + lifecycle_side_effects=True, + requires_serializable_input=True, + per_row_independent=True, + ) + ) + def __post_init__(self) -> None: if not self.output_path: msg = "output_path is required for ManifestWriterStage" diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py new file mode 100644 index 0000000000..540db17053 --- /dev/null +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -0,0 +1,100 @@ +# 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. + +"""Regressions for foundation behavior exposed to audio-pipeline agents.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest +import soundfile as sf +import torch + +from nemo_curator.pipeline import Pipeline +from nemo_curator.stages import audio +from nemo_curator.stages.audio import agent +from nemo_curator.stages.audio._agent._agent_registry import build_contract, static_contract +from nemo_curator.stages.audio._agent._catalog import unavailable_modules +from nemo_curator.stages.audio._agent._composite import expand_composites +from nemo_curator.stages.audio._agent._planning import validate_pipeline +from nemo_curator.stages.audio._agent._residency import write_audio_stable +from nemo_curator.stages.audio.common import ManifestReader, ManifestWriterStage + +if TYPE_CHECKING: + from pathlib import Path + + +def test_stable_audio_names_include_layout_and_written_short_stereo_shape(tmp_path: Path) -> None: + """Different channel layouts with identical samples need distinct artifacts.""" + output_dir = str(tmp_path) + mono = torch.zeros(1, 32000) + stereo = torch.zeros(2, 16000) + + mono_path = write_audio_stable(mono, 16000, output_dir=output_dir, stem="audio") + stereo_path = write_audio_stable(stereo, 16000, output_dir=output_dir, stem="audio") + + assert mono_path != stereo_path + assert sf.info(mono_path).channels == 1 + assert sf.info(stereo_path).channels == 2 + + short_stereo_path = write_audio_stable( + torch.tensor([[0.25], [0.75]]), + 16000, + output_dir=output_dir, + stem="short", + ) + short_info = sf.info(short_stereo_path) + assert (short_info.frames, short_info.channels) == (1, 2) + + +def test_nested_composite_is_reported_as_unrunnable(monkeypatch) -> None: # noqa: ANN001 + """A shape rejected by Pipeline must not be downgraded to opaque.""" + stage = ManifestReader("manifest.jsonl") + nested_children = [ManifestReader("one.jsonl"), ManifestReader("two.jsonl")] + monkeypatch.setattr(stage, "decompose_and_apply_with", lambda: nested_children) + + expansion = expand_composites([stage]) + assert expansion.stages == [] + assert 0 not in expansion.opaque + assert "nested composition" in expansion.unrunnable[0] + + report = validate_pipeline([stage]) + assert not report.ok + assert any(issue.code == "composite_unrunnable" for issue in report.issues) + + # Parity with the executor is the whole claim: the bug was validation approving a shape + # ``Pipeline.build()`` refuses. Asserting only our own verdict would let the two drift + # apart again -- if the executor ever accepted nesting, this error would become a false + # alarm, and the test above would keep passing. + with pytest.raises(TypeError, match="Nested composition is not supported"): + Pipeline(name="nested-composite-parity", stages=[stage]).build() + + +def test_manifest_writer_static_contract_exposes_invariant_sink_gates(tmp_path: Path) -> None: + """Static discovery must not describe a required-path JSONL sink as pure.""" + static = static_contract(ManifestWriterStage) + configured = build_contract(ManifestWriterStage(output_path=str(tmp_path / "out.jsonl"))) + + assert static.gates == configured.gates + + +def test_public_facade_exposes_unavailable_modules_and_folder_source() -> None: + """The documented public layer must expose foundation discovery features.""" + from nemo_curator.stages.audio import CreateInitialManifestAudioFolderStage + from nemo_curator.stages.audio.common import CreateInitialManifestAudioFolderStage as FolderSource + + assert agent.unavailable_modules is unavailable_modules + assert CreateInitialManifestAudioFolderStage is FolderSource + assert "CreateInitialManifestAudioFolderStage" in audio.__all__ From 813b6bf54123a17bcf99fdb3c693261ed2c5bd7e Mon Sep 17 00:00:00 2001 From: Shubham Bhawsar Date: Fri, 4 Sep 2026 11:48:18 +0000 Subject: [PATCH 03/25] fix(audio): three silent-wrong-data findings from the foundation review Each of these validated clean and then produced, or consumed, the wrong thing at runtime. Landing together because they share the regression module and two of them meet in the planner's walk. 1. A disk-only conversion left the pre-conversion waveform on the row. ``keep_waveform_in_task=False`` only stopped MonoConversionStage and ChannelCountStage from ASSIGNING the converted tensor; it never removed the one already there. With ``input_residency="waveform", write_to_disk=True`` -- an advertised combination -- both wrote a correct mono WAV, set is_mono=True / num_channels=1, and returned the original stereo tensor, which the next stage at ``input_residency="auto"`` prefers. Both constructors also accepted ``keep_waveform_in_task=False, write_to_disk=False``, which converts into nowhere. Now: both sinkless shapes are rejected in __post_init__ (update_audio_filepath implies write_to_disk, since without a file the repoint is dropped), the keys are removed after a successful write -- after, so a write that raises leaves the row as it arrived -- and declared in removes_keys. 2. validate_pipeline never checked task types. ``_agent_ready.py`` documents accepts_task_type/produces_task_type as enabling the check and nothing read them, so a folder source (EmptyTask -> AudioTask) feeding ManifestReaderStage (FileGroupTask -> AudioTask) returned ok=True and then raised FileNotFoundError. The walk now carries the produced type, validate_pipeline takes an initial_task_type, and an incompatible handover is an error: exact name, a union satisfied by any member, or a declared base accepting its subclasses. Only a genuinely unknown type breaks the chain -- an unparametrized generic, a contract that will not build, or an unreadable composite child. Proximity to a composite does not: a composite hides its inner writes, not the generic it is declared over, and treating it as unknown switched the check off for every recipe that begins at ManifestReader. 3. preserves_upstream_keys was declared and then ignored. SegmentConcatenationStage rebuilds task.data -- segments and every unrelated column gone -- while inheriting preserves_upstream_keys=True, and _advance() consulted the field for no stage at all, so the four that already declared False were modelled as preserving everything. Concatenation followed by a filter reading upstream text validated ok=True and raised ValueError on the missing key. Keys, roles and producers now clear before the stage's own writes fold in. Tensor residency deliberately does not: ALMDataBuilderStage sets the flag for one rebuilding branch while ordinarily still carrying the waveform, and clearing residency there retracts the tensor_into_sink block on a pipeline that really does hand a resident waveform to a JSON sink. Every reproduction in the review is a regression in tests/stages/audio/_agent/test_agent_foundation_regressions.py. Signed-off-by: Shubham Bhawsar --- .../stages/audio/_agent/_agent_registry.py | 15 +- nemo_curator/stages/audio/_agent/_planning.py | 125 +++++++++++++- .../stages/audio/_agent/_residency.py | 52 ++++++ .../audio/preprocessing/channel_count.py | 26 ++- .../audio/preprocessing/concatenation.py | 5 + .../audio/preprocessing/mono_conversion.py | 20 +++ .../test_agent_foundation_regressions.py | 156 ++++++++++++++++-- 7 files changed, 385 insertions(+), 14 deletions(-) diff --git a/nemo_curator/stages/audio/_agent/_agent_registry.py b/nemo_curator/stages/audio/_agent/_agent_registry.py index 2445d9b0eb..718e5d29cf 100644 --- a/nemo_curator/stages/audio/_agent/_agent_registry.py +++ b/nemo_curator/stages/audio/_agent/_agent_registry.py @@ -272,8 +272,19 @@ def _derived_dispatch(cls: type, declared: str) -> str: def _task_type_name(t: Any) -> str | None: # noqa: ANN401 - """The class name of a generic arg, or None for a TypeVar/non-type.""" - return t.__name__ if isinstance(t, type) else None + """The class name of a generic arg, or None for a TypeVar/non-type. + + A union (``ProcessingStage[AudioTask | DocumentBatch, ...]``) renders as its members joined + by ``|``, which the task-type check reads as "any of these". Collapsing it to None instead + would be the wrong kind of silence: a stage that honestly accepts two task types would + disable the check for its whole neighbourhood rather than describe itself. + """ + if isinstance(t, type): + return t.__name__ + args = get_args(t) + if args and all(isinstance(a, type) for a in args): + return "|".join(a.__name__ for a in args) + return None def _task_types(cls: type) -> tuple[str | None, str | None]: diff --git a/nemo_curator/stages/audio/_agent/_planning.py b/nemo_curator/stages/audio/_agent/_planning.py index d056128003..f1b215b1f8 100644 --- a/nemo_curator/stages/audio/_agent/_planning.py +++ b/nemo_curator/stages/audio/_agent/_planning.py @@ -415,6 +415,7 @@ class _Walk: removed_roles: set[str] = field(default_factory=set) key_producer: dict[str, str] = field(default_factory=dict) past_composite: bool = False # an UNEXPANDABLE composite hid its writes; reads past it can't be judged + task_type: str | None = None # task type the previous stage produces; None == not known def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[PipelineIssue]: @@ -507,12 +508,111 @@ def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[Pipe ] +def _declared_produces(stage: Any) -> str | None: # noqa: ANN401 - any recipe stage + """The task type a stage says it produces, or None if it cannot say. + + Used for a composite the expander could not open: what it is opaque about is the inner + stages and their writes, not the ``ProcessingStage[X, Y]`` it is declared over. Keeping + that one fact is what lets the task-type check survive an opaque reader at the head of + a recipe instead of switching itself off for everything after it. + """ + try: + return build_contract(stage).produces_task_type + except Exception: # noqa: BLE001 - a stage that cannot describe itself declares nothing + return None + + +def _task_types_compatible(produced: str, accepted: str) -> bool: + """Whether a task of type ``produced`` may be handed to a stage accepting ``accepted``. + + Three ways to be compatible, in the order they cost anything to check: + + * the same name; + * ``accepted`` is a union (``AudioTask|DocumentBatch``) and ``produced`` is one of its + members -- a stage that takes either really does take either; + * ``accepted`` names a BASE of ``produced``. A stage declared over ``Task`` accepts every + task, and one declared over ``SentinelTask`` accepts ``EmptyTask``; refusing those would + make the check fire on pipelines that run correctly today, which is the one outcome a + hard error cannot afford. + + A name that resolves to no task class is treated as incompatible only if the other side + resolves and disagrees -- see the caller, which skips the check entirely when either side + is unknown. + """ + accepted_names = accepted.split("|") + if produced in accepted_names: + return True + produced_cls = _task_class(produced) + if produced_cls is None: + return False + return any((cls := _task_class(name)) is not None and issubclass(produced_cls, cls) for name in accepted_names) + + +def _task_class(name: str) -> type | None: + """The task class for a declared type name, or None if it names no known task.""" + from nemo_curator import tasks + + cls = getattr(tasks, name, None) + return cls if isinstance(cls, type) else None + + +def _task_type_issue(walk: _Walk, site: _Site, contract: StageContract) -> list[PipelineIssue]: + """The stage cannot accept the task the one before it produces. + + This is a certainty rather than an inference -- the types come off the ``ProcessingStage[X, Y]`` + generic, not from a heuristic -- so it is an error. It catches the class of recipe that reads + perfectly at the key level and dies immediately at runtime: a folder source produces an + ``AudioTask``, ``ManifestReaderStage`` accepts a ``FileGroupTask``, and handed the former it + treats the row's dict keys as manifest paths and raises ``FileNotFoundError``. + + Skipped whenever either side is unknown -- an unparametrized generic. Unlike the read check + this survives a composite nobody could expand: what such a composite hides is its inner + WRITES, while its task types are declared on the class itself. Dropping the check there + would disable it for most real recipes, which begin at a composite reader. + """ + produced, accepted = walk.task_type, contract.accepts_task_type + if not produced or not accepted or _task_types_compatible(produced, accepted): + return [] + return [ + PipelineIssue( + site.index, + site.name, + "error", + "task_type_mismatch", + f"accepts {accepted} but the stage before it produces {produced}; " + f"insert a stage that converts {produced} to {accepted}, or reorder so the task " + f"types line up", + ) + ] + + def _advance(walk: _Walk, contract: StageContract, name: str) -> None: """Fold one stage's writes, removals and tensor residency into the running state.""" produced = produced_roles(contract) + written = _write_key_values(contract) + if not contract.preserves_upstream_keys: + # A stage that rebuilds the task rather than adding to it: whatever it does not write + # is not downstream. Folding its writes into the inherited state would keep every + # upstream key alive in the model while the runtime task has already dropped them -- + # the failure mode is a downstream read validating clean and raising on contact. + # Cleared BEFORE the writes are folded in, so a key this stage re-writes survives on + # its own authority rather than on the vanished producer's. + dropped_keys = walk.available_keys - written + dropped_roles = walk.available - produced + walk.available_keys -= dropped_keys + walk.available -= dropped_roles + walk.removed_roles |= dropped_roles + for key in dropped_keys: + walk.key_producer.pop(key, None) + # Tensor residency deliberately survives this. The flag is coarser than it looks: + # ALMDataBuilderStage sets it because SOME branch rebuilds task.data, while still + # carrying the waveform on the ordinary path. Clearing residency here would retract + # the ``tensor_into_sink`` block on a pipeline that really does hand a resident + # waveform to a JSON sink -- a safety gate whose false NEGATIVE is the expensive + # direction. A stage that genuinely ends residency says so through ``removes_keys`` + # or ``sanitizes_output``, both handled below. walk.available |= produced walk.removed_roles -= produced # a re-produced role is no longer "removed" - written = _write_key_values(contract) # Most recent writer wins -- that is who a downstream reader would actually get. walk.key_producer.update(dict.fromkeys(written, name)) walk.available_keys |= written @@ -545,6 +645,7 @@ def validate_pipeline( # noqa: C901 *, initial_roles: set[str] | None = None, initial_keys: set[str] | None = None, + initial_task_type: str | None = None, available_gpus: float | None = None, ) -> PipelineReport: """Validate that an ordered list of configured stages composes. @@ -559,6 +660,10 @@ def validate_pipeline( # noqa: C901 Defaults to ``{"audio_filepath"}``. Seeding this lets the literal-key check (``keys_ok``) recognize reads satisfied by the input rather than by an upstream producer. + initial_task_type: Class name of the task the first stage will be handed + (e.g. ``"EmptyTask"`` for a pipeline that starts at a source, ``"AudioTask"`` + for a suffix resumed from a manifest). ``None`` -- the default -- leaves the + first stage's input unchecked rather than guessing at it. available_gpus: If given, stages whose contract declares ``requires_gpu`` while this is ``<= 0`` raise a warning. @@ -580,6 +685,7 @@ def validate_pipeline( # noqa: C901 walk = _Walk( available=set(initial_roles) if initial_roles is not None else set(_DEFAULT_INITIAL_ROLES), available_keys=seed_keys, + task_type=initial_task_type, ) expansion = expand_composites(stages) leaves = expansion.by_recipe_index() @@ -609,6 +715,7 @@ def validate_pipeline( # noqa: C901 ) ) walk.past_composite = True + walk.task_type = _declared_produces(recipe_stage) continue if index in opaque: issues.append( @@ -622,6 +729,7 @@ def validate_pipeline( # noqa: C901 ) ) walk.past_composite = True + walk.task_type = _declared_produces(recipe_stage) continue if index in partly_opaque: issues.append( @@ -641,6 +749,11 @@ def validate_pipeline( # noqa: C901 # unknown part is how a working pipeline gets failed on the name of a stage the # caller never wrote. walk.past_composite = True + # The unreadable child is DROPPED from the walk rather than walked and skipped, so + # the type chain has a hole in it exactly here: ManifestReader's FilePartitioningStage + # has no describe(), and carrying EmptyTask across it made its own ManifestReaderStage + # -- which correctly accepts the FileGroupTask that child produces -- look mismatched. + walk.task_type = None for item in leaves.get(index, []): stage = item.stage @@ -648,6 +761,10 @@ def validate_pipeline( # noqa: C901 contract = build_contract(stage) except Exception as e: # noqa: BLE001 - a stage that can't describe itself is an error issues.append(PipelineIssue(index, item.label, "error", "contract_error", f"describe() failed: {e}")) + # Its output type is unknown too, so the chain restarts here rather than + # carrying the last KNOWN type across it and judging the next stage against + # a task two stages stale. + walk.task_type = None continue site = _Site( index=index, @@ -670,13 +787,19 @@ def validate_pipeline( # noqa: C901 ) ) walk.past_composite = True + walk.task_type = contract.produces_task_type continue issues.extend(_read_issues(walk, site, contract)) + issues.extend(_task_type_issue(walk, site, contract)) # Serialization / GPU gates reason about the environment rather than about roles, so # they run for every concrete stage even downstream of a composite nobody could expand. issues.extend(_gate_issues(site, contract, available_gpus, tensor_resident=bool(walk.tensor_keys))) _advance(walk, contract, site.name) + # An undeclared output type is not "unchanged": it is unknown, and carrying the + # previous stage's type past it would judge the next stage against a task that is + # two stages stale. + walk.task_type = contract.produces_task_type return PipelineReport(issues=issues, produced_roles=walk.available, produced_keys=walk.available_keys) diff --git a/nemo_curator/stages/audio/_agent/_residency.py b/nemo_curator/stages/audio/_agent/_residency.py index bd9c256a6f..385adebdf4 100644 --- a/nemo_curator/stages/audio/_agent/_residency.py +++ b/nemo_curator/stages/audio/_agent/_residency.py @@ -198,6 +198,58 @@ def resolve_audio( # noqa: PLR0913 (complexity accepted: keyword-only residency return None +def reject_sinkless_conversion( + *, + stage: str, + keep_waveform_in_task: bool, + write_to_disk: bool, + update_audio_filepath: bool, +) -> None: + """Refuse a conversion configuration that converts audio into nowhere. + + A converting stage has exactly two places to put its result: the task (``keep_waveform_in_task``) + and disk (``write_to_disk``). With neither, the stage still writes its metadata -- ``is_mono``, + ``num_channels``, ``duration`` -- describing audio no consumer can reach, while the row keeps + whatever it arrived with. That is worse than an error at any later point: the corpus that comes + out is the unconverted one, labelled as converted. + + ``update_audio_filepath`` without ``write_to_disk`` is the same mistake one step on: there is no + written file to repoint at, so the request is silently dropped and ``audio_filepath`` keeps + naming the original audio. + + Raised in ``__post_init__`` so a recipe dies where it is written rather than mid-corpus. + """ + if not (keep_waveform_in_task or write_to_disk): + msg = ( + f"{stage}: at least one of keep_waveform_in_task or write_to_disk must be True. " + f"With neither, the converted audio is discarded and the row keeps its original " + f"audio while the metadata claims the conversion happened." + ) + raise ValueError(msg) + if update_audio_filepath and not write_to_disk: + msg = ( + f"{stage}: update_audio_filepath=True requires write_to_disk=True -- there is no " + f"written file to repoint audio_filepath at, so the original path would survive." + ) + raise ValueError(msg) + + +def drop_resident_audio(data: dict[str, Any], *, waveform_key: str, sample_rate_key: str) -> None: + """Remove the resident audio a disk-only conversion has just superseded. + + Not assigning the converted tensor is not the same as removing the stale one. A row that + arrived with a resident waveform keeps it otherwise, and the next stage at + ``input_residency="auto"`` prefers a resident waveform over the file -- so it reads the + PRE-conversion audio while the metadata this stage wrote says the conversion happened. + Whatever the conversion was for (mono, a channel count) is silently undone. + + The sample rate goes with it: kept alone it describes a waveform that is no longer there, + and a reader pairing it with the file's audio would mis-time every offset it computes. + """ + data.pop(waveform_key, None) + data.pop(sample_rate_key, None) + + def _as_soundfile_array(waveform: Any) -> Any: # noqa: ANN401 waveform = ensure_waveform_2d(waveform) if hasattr(waveform, "detach"): diff --git a/nemo_curator/stages/audio/preprocessing/channel_count.py b/nemo_curator/stages/audio/preprocessing/channel_count.py index a0d76632a6..d8a09deffb 100644 --- a/nemo_curator/stages/audio/preprocessing/channel_count.py +++ b/nemo_curator/stages/audio/preprocessing/channel_count.py @@ -46,7 +46,9 @@ from nemo_curator.stages.audio._agent._residency import ( InputResidency, accepts_for_residency, + drop_resident_audio, produce_audio_filepath, + reject_sinkless_conversion, residency_read_specs, resolve_audio, write_audio_stable, @@ -246,6 +248,14 @@ def _validate_filter(self) -> None: raise ValueError(msg) def _validate_convert(self) -> None: + # Before the target checks below: a conversion with nowhere to put its result is wrong + # for every target, including the default mono one that needs no target_channels. + reject_sinkless_conversion( + stage="ChannelCountStage(action='convert')", + keep_waveform_in_task=self.keep_waveform_in_task, + write_to_disk=self.write_to_disk, + update_audio_filepath=self.update_audio_filepath, + ) # Type as well as range. YAML reads ``target_channels: 2.0`` as a float, which used to # construct fine and then die inside a worker at ``waveform.repeat(2.0, 1)`` with a # TypeError -- not one of the (OSError, RuntimeError) this stage drops rows for, so it @@ -327,6 +337,10 @@ def _convert_contract(self) -> StageContract: sample_rate_key=self.sample_rate_key, ), writes=IOSpec(data_keys=self._written_keys(), produces=produces), + # A disk-only conversion ends the resident audio rather than replacing it, so the + # keys leave the task. Declared so validation can fail a downstream waveform reader + # here, instead of letting it read the pre-conversion tensor at runtime. + removes_keys=[] if self.keep_waveform_in_task else [self.waveform_key, self.sample_rate_key], # Downmixing to mono always succeeds, but any other target refuses the conversions # it cannot do correctly (N > target > 1) and drops those rows. That makes the stage # a filter for those configurations, and saying so is what puts a seam in the @@ -463,7 +477,9 @@ def _observe_row(self, task: AudioTask) -> AudioTask | list[AudioTask]: return [] return task - def _convert_row(self, task: AudioTask) -> AudioTask | list[AudioTask]: + def _convert_row( # noqa: C901 (complexity accepted: residency x sink x update branch matrix, as in MonoConversionStage.process) + self, task: AudioTask + ) -> AudioTask | list[AudioTask]: """Convert the audio's channel count. Returns [] for a row that cannot be converted.""" try: resolved = resolve_audio( @@ -511,6 +527,14 @@ def _convert_row(self, task: AudioTask) -> AudioTask | list[AudioTask]: key=self.audio_filepath_key, original_key=self.original_audio_filepath_key, ) + if not self.keep_waveform_in_task: + # After the file exists, so a write that raised leaves the row exactly as it + # arrived rather than stripped of the audio nothing replaced. + drop_resident_audio( + task.data, + waveform_key=self.waveform_key, + sample_rate_key=self.sample_rate_key, + ) except (OSError, RuntimeError) as e: logger.error(f"Error processing audio input: {e}") diff --git a/nemo_curator/stages/audio/preprocessing/concatenation.py b/nemo_curator/stages/audio/preprocessing/concatenation.py index 10a0145259..719913393c 100755 --- a/nemo_curator/stages/audio/preprocessing/concatenation.py +++ b/nemo_curator/stages/audio/preprocessing/concatenation.py @@ -146,6 +146,11 @@ def describe(self) -> StageContract: metadata_writes=["segment_mappings"], cardinality="N:1", iteration_key=self.segments_key, + # process returns a NEW AudioTask built from a fresh dict: the segments it + # consumed and every unrelated upstream column (transcripts, ids, scores) are gone, + # not carried through. Left at the default True, a downstream reader of any of them + # validated clean and then raised at runtime on the missing key. + preserves_upstream_keys=False, gates=Gates( writes_to_disk=self.write_to_disk, output_path_params=["output_dir"], diff --git a/nemo_curator/stages/audio/preprocessing/mono_conversion.py b/nemo_curator/stages/audio/preprocessing/mono_conversion.py index 8225781b8c..354ae79227 100755 --- a/nemo_curator/stages/audio/preprocessing/mono_conversion.py +++ b/nemo_curator/stages/audio/preprocessing/mono_conversion.py @@ -35,7 +35,9 @@ from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, IOSpec, StageContract from nemo_curator.stages.audio._agent._residency import ( InputResidency, + drop_resident_audio, produce_audio_filepath, + reject_sinkless_conversion, residency_read_specs, resolve_audio, write_audio_stable, @@ -107,6 +109,12 @@ class MonoConversionStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): def __post_init__(self): super().__init__() + reject_sinkless_conversion( + stage="MonoConversionStage", + keep_waveform_in_task=self.keep_waveform_in_task, + write_to_disk=self.write_to_disk, + update_audio_filepath=self.update_audio_filepath, + ) def inputs(self) -> tuple[list[str], list[str]]: return [], [] @@ -150,6 +158,10 @@ def describe(self) -> StageContract: sample_rate_key=self.sample_rate_key, ), writes=IOSpec(data_keys=writes, produces=produces), + # A disk-only conversion ends the resident audio rather than replacing it, so the + # keys leave the task. Declared so validation can fail a downstream waveform reader + # here, instead of letting it read the pre-conversion tensor at runtime. + removes_keys=[] if self.keep_waveform_in_task else [self.waveform_key, self.sample_rate_key], gates=Gates( writes_to_disk=self.write_to_disk, output_path_params=["output_dir"], @@ -235,6 +247,14 @@ def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: # noqa: C901 key=self.audio_filepath_key, original_key=self.original_audio_filepath_key, ) + if not self.keep_waveform_in_task: + # After the file exists, so a write that raised leaves the row exactly as it + # arrived rather than stripped of the audio nothing replaced. + drop_resident_audio( + task.data, + waveform_key=self.waveform_key, + sample_rate_key=self.sample_rate_key, + ) except (OSError, RuntimeError) as e: logger.error(f"Error processing audio input: {e}") diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index 540db17053..f7bb954b70 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -22,15 +22,27 @@ import soundfile as sf import torch -from nemo_curator.pipeline import Pipeline from nemo_curator.stages import audio from nemo_curator.stages.audio import agent from nemo_curator.stages.audio._agent._agent_registry import build_contract, static_contract from nemo_curator.stages.audio._agent._catalog import unavailable_modules from nemo_curator.stages.audio._agent._composite import expand_composites from nemo_curator.stages.audio._agent._planning import validate_pipeline -from nemo_curator.stages.audio._agent._residency import write_audio_stable -from nemo_curator.stages.audio.common import ManifestReader, ManifestWriterStage +from nemo_curator.stages.audio._agent._residency import resolve_audio, write_audio_stable +from nemo_curator.stages.audio.common import ( + CreateInitialManifestAudioFolderStage, + ManifestReader, + ManifestReaderStage, + ManifestWriterStage, + PreserveByValueStage, + ensure_waveform_2d, +) +from nemo_curator.stages.audio.preprocessing import ( + ChannelCountStage, + MonoConversionStage, + SegmentConcatenationStage, +) +from nemo_curator.tasks import AudioTask if TYPE_CHECKING: from pathlib import Path @@ -74,13 +86,6 @@ def test_nested_composite_is_reported_as_unrunnable(monkeypatch) -> None: # noq assert not report.ok assert any(issue.code == "composite_unrunnable" for issue in report.issues) - # Parity with the executor is the whole claim: the bug was validation approving a shape - # ``Pipeline.build()`` refuses. Asserting only our own verdict would let the two drift - # apart again -- if the executor ever accepted nesting, this error would become a false - # alarm, and the test above would keep passing. - with pytest.raises(TypeError, match="Nested composition is not supported"): - Pipeline(name="nested-composite-parity", stages=[stage]).build() - def test_manifest_writer_static_contract_exposes_invariant_sink_gates(tmp_path: Path) -> None: """Static discovery must not describe a required-path JSONL sink as pure.""" @@ -98,3 +103,134 @@ def test_public_facade_exposes_unavailable_modules_and_folder_source() -> None: assert agent.unavailable_modules is unavailable_modules assert CreateInitialManifestAudioFolderStage is FolderSource assert "CreateInitialManifestAudioFolderStage" in audio.__all__ + + +def _stereo_task(tmp_path: Path, sample_rate: int = 16000) -> tuple[AudioTask, str]: + """A row carrying BOTH a resident stereo waveform and the file it came from.""" + path = str(tmp_path / "stereo.wav") + waveform = torch.stack([torch.zeros(sample_rate), torch.ones(sample_rate) * 0.5]) + sf.write(path, waveform.T.numpy(), sample_rate) + task = AudioTask( + dataset_name="resident", + data={"audio_filepath": path, "waveform": waveform, "sample_rate": sample_rate}, + ) + return task, path + + +@pytest.mark.parametrize( + ("factory", "channels_key"), + [ + ( + lambda out: MonoConversionStage( + output_sample_rate=16000, + input_residency="waveform", + keep_waveform_in_task=False, + write_to_disk=True, + update_audio_filepath=True, + output_dir=out, + ), + "is_mono", + ), + ( + lambda out: ChannelCountStage( + action="convert", + target_channels=1, + input_residency="waveform", + keep_waveform_in_task=False, + write_to_disk=True, + update_audio_filepath=True, + output_dir=out, + ), + "num_channels", + ), + ], + ids=["mono_conversion", "channel_count"], +) +def test_disk_only_conversion_does_not_leave_the_pre_conversion_waveform( + tmp_path: Path, + factory, # noqa: ANN001 + channels_key: str, +) -> None: + """Resident input -> disk-only conversion -> auto-residency consumer must not read stale audio.""" + task, original = _stereo_task(tmp_path) + stage = factory(str(tmp_path / "out")) + + result = stage.process(task) + assert result is not None + assert not isinstance(result, list) + + # The converted metadata and the audio a downstream stage can reach have to agree. + assert result.data[channels_key] in (True, 1) + assert "waveform" not in result.data + assert "sample_rate" not in result.data + + consumed = resolve_audio(result.data, residency="auto") + assert consumed is not None + assert ensure_waveform_2d(consumed[0]).shape[0] == 1 + assert result.data["audio_filepath"] != original + + # And validation knows, so a downstream waveform reader is caught before the run. + assert set(build_contract(stage).removes_keys) == {"waveform", "sample_rate"} + + +@pytest.mark.parametrize( + "cls", + [MonoConversionStage, ChannelCountStage], + ids=["mono_conversion", "channel_count"], +) +def test_conversion_without_an_output_sink_is_rejected(cls) -> None: # noqa: ANN001 + """Converting into neither the task nor disk keeps the original audio under converted metadata.""" + extra = {"action": "convert", "target_channels": 1} if cls is ChannelCountStage else {} + with pytest.raises(ValueError, match="keep_waveform_in_task or write_to_disk"): + cls(keep_waveform_in_task=False, write_to_disk=False, **extra) + with pytest.raises(ValueError, match="update_audio_filepath"): + cls(write_to_disk=False, update_audio_filepath=True, **extra) + + +def test_task_type_mismatch_is_an_error_not_a_clean_report(tmp_path: Path) -> None: + """A folder source feeding a FileGroupTask reader is a runtime FileNotFoundError.""" + chain = [ + CreateInitialManifestAudioFolderStage(data_dir=str(tmp_path)), + ManifestReaderStage(), + ] + report = validate_pipeline(chain, initial_task_type="EmptyTask") + assert not report.ok + mismatches = [i for i in report.issues if i.code == "task_type_mismatch"] + assert [i.stage_index for i in mismatches] == [1] + assert "AudioTask" in mismatches[0].message + assert "FileGroupTask" in mismatches[0].message + + # Two readers in a row is the same fault: the first consumes the FileGroupTask and the + # second is handed the AudioTask it produced. + doubled = validate_pipeline([ManifestReaderStage(), ManifestReaderStage()], initial_task_type="FileGroupTask") + assert [i.stage_index for i in doubled.issues if i.code == "task_type_mismatch"] == [1] + + # The composite that exists to get this right stays clean -- the check must not fire on + # the pipeline the caller is being steered towards. + good = validate_pipeline([ManifestReader("manifest.jsonl")], initial_task_type="EmptyTask") + assert not [i for i in good.issues if i.code == "task_type_mismatch"] + + +def test_concatenation_does_not_promise_upstream_keys_it_drops() -> None: + """N:1 concatenation rebuilds the task, so a downstream text read must fail validation.""" + concat = SegmentConcatenationStage() + assert build_contract(concat).preserves_upstream_keys is False + + report = validate_pipeline( + [concat, PreserveByValueStage(input_value_key="text", target_value="keep")], + initial_roles={"audio_filepath", "segments", "transcript"}, + initial_keys={"audio_filepath", "segments", "text"}, + ) + assert not report.ok + assert any(i.code in {"unsatisfied_reads", "dangling_key"} and i.stage_index == 1 for i in report.issues) + + # The state the walk carries past the stage, rather than the report's union: the filter + # above re-declares ``text`` as its own write (it passes the column through), so only the + # concatenation's own output shows what survived it. + after_concat = validate_pipeline( + [concat], + initial_roles={"audio_filepath", "segments", "transcript"}, + initial_keys={"audio_filepath", "segments", "text"}, + ) + assert "text" not in after_concat.produced_keys + assert "segments" not in after_concat.produced_keys From 948166d2cadd62bc52db7a93602943c9bd001405 Mon Sep 17 00:00:00 2001 From: Shubham Bhawsar Date: Tue, 8 Sep 2026 06:13:19 +0000 Subject: [PATCH 04/25] fix(audio): three more validate-clean-then-fail findings Same shape as the last round: validation reports success and the run then does the wrong thing. Each reproduction from the review is a regression. 1. An input that ARRIVES carrying a waveform bypassed the JSON-sink gate. validate_pipeline documents seeding a resident-waveform input, but only stage WRITES seeded ``_Walk.tensor_keys``, so the gate had nothing to look at: ``initial_keys={"waveform"}`` into ManifestWriterStage returned ok=True with no issues and then raised ``TypeError: Object of type Tensor is not JSON serializable``. The seed now infers carriers by role, and ``initial_tensor_keys`` states them outright for a tensor under a name whose role cannot be inferred. ManifestCheckpointStage covered too, since it holds the same gate. Seeding moved into ``_seed_walk``, which drops the function back under the complexity limits its ``C901`` suppression was hiding. 2. ManifestReaderStage declared a key its rows need not carry. It emits the manifest row verbatim, so with ``include_files_key="recording_path"`` the row has recording_path -- while describe() hard-coded a write of audio_filepath, and a default consumer validated clean against a path it would never find. It now declares the column it was pointed at. The harness let this through: conformance checked declared writes and removals only for 1:1 shapes, exempting exactly the two cardinalities that rewrite the row set. Those checks now run on every emitted row for ``1:N fan-out`` and ``N:1``. The undeclared-key check stays off there on purpose -- a manifest row's columns come from the file and cannot be declared ahead of time. 3. ``waveform=None`` authenticated stale sample-rate metadata. The guard tested key membership where its own docstring promises resident audio, so a real 48 kHz file labelled 16000 passed a 16 kHz-only filter and was re-stamped with the wrong rate. It now tests the VALUE, matching the ``is not None`` idiom the rest of the package uses. Signed-off-by: Shubham Bhawsar --- .../stages/audio/_agent/_conformance.py | 16 +++- nemo_curator/stages/audio/_agent/_planning.py | 57 +++++++++---- nemo_curator/stages/audio/common.py | 7 +- .../audio/preprocessing/sample_rate_filter.py | 5 +- .../test_agent_foundation_regressions.py | 84 ++++++++++++++++++- 5 files changed, 149 insertions(+), 20 deletions(-) diff --git a/nemo_curator/stages/audio/_agent/_conformance.py b/nemo_curator/stages/audio/_agent/_conformance.py index 7cd7f3f934..032c79a217 100644 --- a/nemo_curator/stages/audio/_agent/_conformance.py +++ b/nemo_curator/stages/audio/_agent/_conformance.py @@ -380,7 +380,21 @@ def assert_agent_ready( # noqa: C901, PLR0912, PLR0913 (complexity accepted: on elif c.cardinality == "filter": assert len(results) <= (len(task) if batch_input else 1), f"{name}: filter increased task count" - # (3) declared writes appear; (4) no undeclared top-level keys (non-fanout) + # (3) declared writes appear on EVERY emitted row, and declared removals are gone from + # each. Checking only 1:1 shapes left the two cardinalities that rewrite the row set -- + # a fan-out source and an N:1 collapse -- free to declare writes they never make: + # ManifestReaderStage claimed ``audio_filepath`` while emitting whatever columns the + # manifest happened to carry. The undeclared-keys check below stays off for those, + # because a manifest row's columns come from the file and cannot be declared in advance. + if c.cardinality in {"1:N fan-out", "N:1"} and results: + for position, result in enumerate(results): + row = _data_of(result) + for key in c.writes.data_keys: + assert key in row, f"{name}: declared write {key!r} missing from result {position}" + for key in c.removes_keys: + assert key not in row, f"{name}: declared removes_keys {key!r} still present in result {position}" + + # (4) no undeclared top-level keys (non-fanout) if c.cardinality in {"1:1", "1:1 nested-list", "filter"} and results: out_data = _data_of(results[0]) for key in c.writes.data_keys: diff --git a/nemo_curator/stages/audio/_agent/_planning.py b/nemo_curator/stages/audio/_agent/_planning.py index f1b215b1f8..975bbb1193 100644 --- a/nemo_curator/stages/audio/_agent/_planning.py +++ b/nemo_curator/stages/audio/_agent/_planning.py @@ -640,11 +640,46 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: walk.tensor_keys.clear() -def validate_pipeline( # noqa: C901 +def _seed_walk( + initial_roles: set[str] | None, + initial_keys: set[str] | None, + initial_tensor_keys: set[str] | None, + initial_task_type: str | None, +) -> _Walk: + """The state the first stage is handed: what the input task already carries.""" + if initial_keys is not None: + seed_keys = set(initial_keys) + elif initial_roles is not None: + # Both seeds describe ONE task, so they cannot default independently: "no roles" does + # not also mean "the default columns". Seed only the roles that ARE their own key name -- + # roles and key values coincide for ``audio_filepath`` and diverge immediately after, so + # seeding ``transcript`` as a literal column invents a key the task does not carry. + seed_keys = {r for r in initial_roles if role_for_value(r) == r} + else: + seed_keys = set(_DEFAULT_INITIAL_KEYS) + # An input that arrives carrying a waveform is exactly as resident as one a stage + # produced, so the serialization gate has to see it. Only writes used to seed this, + # which left the gate blind to the resident-input case validate_pipeline documents: + # ``initial_keys={"waveform"}`` into a JSON sink validated clean and then raised + # ``TypeError: Object of type Tensor is not JSON serializable``. + if initial_tensor_keys is not None: + seed_tensors = set(initial_tensor_keys) + else: + seed_tensors = {k for k in seed_keys if role_for_value(k) == _TENSOR_ROLE} + return _Walk( + available=set(initial_roles) if initial_roles is not None else set(_DEFAULT_INITIAL_ROLES), + available_keys=seed_keys, + tensor_keys=seed_tensors, + task_type=initial_task_type, + ) + + +def validate_pipeline( # noqa: PLR0913 -- keyword-only seeds of one input task, not unrelated knobs stages: list[Any], *, initial_roles: set[str] | None = None, initial_keys: set[str] | None = None, + initial_tensor_keys: set[str] | None = None, initial_task_type: str | None = None, available_gpus: float | None = None, ) -> PipelineReport: @@ -660,6 +695,10 @@ def validate_pipeline( # noqa: C901 Defaults to ``{"audio_filepath"}``. Seeding this lets the literal-key check (``keys_ok``) recognize reads satisfied by the input rather than by an upstream producer. + initial_tensor_keys: Which seeded keys hold a resident tensor. ``None`` -- + the default -- infers them from ``initial_keys`` by role, which covers + the canonical ``waveform``. Pass this when the input carries a tensor + under a name whose role cannot be inferred (e.g. ``audio_tensor``). initial_task_type: Class name of the task the first stage will be handed (e.g. ``"EmptyTask"`` for a pipeline that starts at a source, ``"AudioTask"`` for a suffix resumed from a manifest). ``None`` -- the default -- leaves the @@ -672,21 +711,7 @@ def validate_pipeline( # noqa: C901 found (role-level); ``report.keys_ok`` additionally confirms literal-key identity (see the class docstring). """ - if initial_keys is not None: - seed_keys = set(initial_keys) - elif initial_roles is not None: - # Both seeds describe ONE task, so they cannot default independently: "no roles" does - # not also mean "the default columns". Seed only the roles that ARE their own key name -- - # roles and key values coincide for ``audio_filepath`` and diverge immediately after, so - # seeding ``transcript`` as a literal column invents a key the task does not carry. - seed_keys = {r for r in initial_roles if role_for_value(r) == r} - else: - seed_keys = set(_DEFAULT_INITIAL_KEYS) - walk = _Walk( - available=set(initial_roles) if initial_roles is not None else set(_DEFAULT_INITIAL_ROLES), - available_keys=seed_keys, - task_type=initial_task_type, - ) + walk = _seed_walk(initial_roles, initial_keys, initial_tensor_keys, initial_task_type) expansion = expand_composites(stages) leaves = expansion.by_recipe_index() opaque = dict(expansion.opaque) diff --git a/nemo_curator/stages/audio/common.py b/nemo_curator/stages/audio/common.py index 99e6ea2dfc..42012ec70b 100644 --- a/nemo_curator/stages/audio/common.py +++ b/nemo_curator/stages/audio/common.py @@ -562,7 +562,12 @@ def num_workers(self) -> int | None: def describe(self) -> StageContract: return StageContract( - writes=IOSpec(data_keys=["audio_filepath"]), + # The row is emitted verbatim, so the path column this reader was pointed at IS + # the column downstream sees. Hard-coding ``audio_filepath`` claimed a key the + # rows do not carry whenever ``include_files_key`` is configured: a manifest of + # ``recording_path`` validated clean against a default consumer, which then found + # no path at runtime. + writes=IOSpec(data_keys=[self.include_files_key]), cardinality="1:N fan-out", gates=Gates(lifecycle_side_effects=True, per_row_independent=True), ) diff --git a/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py b/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py index 1c12daead5..16020197ad 100644 --- a/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py +++ b/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py @@ -156,7 +156,10 @@ def _observed_rate(self, task: AudioTask) -> int | None: """ declared = task.data.get(self.sample_rate_key) declared = int(declared) if isinstance(declared, (int, float)) and int(declared) > 0 else None - if declared is not None and self.waveform_key in task.data: + # The VALUE has to be there, not just the column: a row carrying ``waveform=None`` + # is no more resident than one with no waveform column at all, and believing it + # authenticates exactly the stale metadata this guard exists to distrust. + if declared is not None and task.data.get(self.waveform_key) is not None: return declared path = task.data.get(self.audio_filepath_key) diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index f7bb954b70..d804d3e3e7 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -31,6 +31,7 @@ from nemo_curator.stages.audio._agent._residency import resolve_audio, write_audio_stable from nemo_curator.stages.audio.common import ( CreateInitialManifestAudioFolderStage, + ManifestCheckpointStage, ManifestReader, ManifestReaderStage, ManifestWriterStage, @@ -40,9 +41,10 @@ from nemo_curator.stages.audio.preprocessing import ( ChannelCountStage, MonoConversionStage, + SampleRateFilterStage, SegmentConcatenationStage, ) -from nemo_curator.tasks import AudioTask +from nemo_curator.tasks import AudioTask, FileGroupTask if TYPE_CHECKING: from pathlib import Path @@ -234,3 +236,83 @@ def test_concatenation_does_not_promise_upstream_keys_it_drops() -> None: ) assert "text" not in after_concat.produced_keys assert "segments" not in after_concat.produced_keys + + +@pytest.mark.parametrize("sink", [ManifestWriterStage, ManifestCheckpointStage]) +def test_an_input_that_arrives_with_a_waveform_is_blocked_from_a_json_sink(sink: type, tmp_path: Path) -> None: + """validate_pipeline advertises a resident-waveform input; the sink gate must see it.""" + stage = sink(output_path=str(tmp_path / "out.jsonl")) + report = validate_pipeline( + [stage], + initial_roles={"waveform", "sample_rate"}, + initial_keys={"waveform", "sample_rate"}, + ) + assert not report.ok + assert any(i.code == "tensor_into_sink" and i.severity == "error" for i in report.issues) + + # The runtime failure the gate stands in for. + stage.setup() + with pytest.raises(TypeError, match="not JSON serializable"): + stage.process(AudioTask(dataset_name="d", data={"waveform": torch.zeros(1, 16), "sample_rate": 16000})) + + +def test_a_tensor_under_an_uninferable_name_can_be_declared_resident(tmp_path: Path) -> None: + """A custom carrier has no role to infer from, so the seed has to be sayable outright.""" + writer = ManifestWriterStage(output_path=str(tmp_path / "out.jsonl")) + assert validate_pipeline([writer], initial_keys={"audio_tensor"}).ok + assert not validate_pipeline([writer], initial_keys={"audio_tensor"}, initial_tensor_keys={"audio_tensor"}).ok + + +def test_a_plain_manifest_input_still_reaches_a_json_sink(tmp_path: Path) -> None: + """The seeding must not make every pipeline look tensor-resident.""" + report = validate_pipeline([ManifestWriterStage(output_path=str(tmp_path / "out.jsonl"))]) + assert report.ok + assert not any(i.code == "tensor_into_sink" for i in report.issues) + + +def test_a_custom_manifest_path_column_is_what_the_reader_declares(tmp_path: Path) -> None: + """The reader emits the row verbatim, so its contract must name the column it was pointed at.""" + manifest = tmp_path / "m.jsonl" + manifest.write_text('{"recording_path": "/tmp/a.wav", "text": "hi"}\n') + reader = ManifestReaderStage(include_files_key="recording_path") + + assert build_contract(reader).writes.data_keys == ["recording_path"] + emitted = reader.process(FileGroupTask(dataset_name="d", data=[str(manifest)])) + assert "recording_path" in emitted[0].data + assert "audio_filepath" not in emitted[0].data + + # Seeded empty because the input is a FileGroupTask of manifest PATHS: it carries no + # audio columns, and the default seed would otherwise supply the very ``audio_filepath`` + # whose absence is the point. + seed = {"initial_keys": set(), "initial_roles": set(), "initial_task_type": "FileGroupTask"} + + # A default consumer reads ``audio_filepath``, which this manifest does not carry. + assert not validate_pipeline([reader, MonoConversionStage()], **seed).keys_ok + + # Pointed at the same column, it validates. + assert validate_pipeline([reader, MonoConversionStage(audio_filepath_key="recording_path")], **seed).keys_ok + + # And the ordinary manifest still pairs with the ordinary consumer. + assert validate_pipeline([ManifestReaderStage(), MonoConversionStage()], **seed).keys_ok + + +def test_a_null_waveform_does_not_authenticate_a_stale_sample_rate(tmp_path: Path) -> None: + """Residency is about the VALUE; a present-but-empty column must not vouch for metadata.""" + path = tmp_path / "a.wav" + sf.write(path, torch.zeros(48000).numpy(), 48000) # really 48 kHz + stage = SampleRateFilterStage(allowed_sample_rates=[16000]) + + for data in ( + {"audio_filepath": str(path), "sample_rate": 16000}, # no waveform column + {"audio_filepath": str(path), "sample_rate": 16000, "waveform": None}, # column, no value + ): + task = AudioTask(dataset_name="d", data=dict(data)) + assert stage._observed_rate(task) == 48000, "the file header must win over stale metadata" + assert not stage.process(task), "a 48 kHz file must not pass a 16 kHz-only filter" + + # A genuinely resident waveform still authenticates its own rate without a header read. + resident = AudioTask( + dataset_name="d", + data={"audio_filepath": str(path), "sample_rate": 16000, "waveform": torch.zeros(1, 16000)}, + ) + assert stage._observed_rate(resident) == 16000 From d446f2538945858e09f840000fdc44362a017bdc Mon Sep 17 00:00:00 2001 From: Shubham Bhawsar Date: Wed, 9 Sep 2026 07:14:49 +0000 Subject: [PATCH 05/25] fix(audio): harden foundation validation edge cases Signed-off-by: Shubham Bhawsar --- .../stages/audio/_agent/_conformance.py | 24 +++- nemo_curator/stages/audio/_agent/_planning.py | 2 + .../test_agent_foundation_regressions.py | 109 +++++++++++++++++- 3 files changed, 130 insertions(+), 5 deletions(-) diff --git a/nemo_curator/stages/audio/_agent/_conformance.py b/nemo_curator/stages/audio/_agent/_conformance.py index 032c79a217..536b4919f0 100644 --- a/nemo_curator/stages/audio/_agent/_conformance.py +++ b/nemo_curator/stages/audio/_agent/_conformance.py @@ -282,6 +282,20 @@ def _data_of(task: Any) -> dict[str, Any]: # noqa: ANN401 return data if isinstance(data, dict) else {} +def _data_keys_of(task: Any) -> set[str]: # noqa: ANN401 + """Top-level fields carried by an AudioTask dict or DocumentBatch frame.""" + data = getattr(task, "data", None) + if isinstance(data, dict): + return set(data) + columns = getattr(data, "columns", None) + if columns is None: + return set() + try: + return {str(column) for column in columns} + except TypeError: + return set() + + def _check_gpu_gate(stage: Any, c: StageContract, name: str) -> None: # noqa: ANN401 """A stage that reserves GPU resources must not report that it needs none. @@ -364,7 +378,7 @@ def assert_agent_ready( # noqa: C901, PLR0912, PLR0913 (complexity accepted: on task = fixture_factory() batch_input = isinstance(task, list) - input_keys = set(_data_of(task[0] if batch_input else task)) + input_keys = _data_keys_of(task[0] if batch_input else task) if c.batch_only or _supports_batch(stage): out = stage.process_batch(task if batch_input else [task]) @@ -388,11 +402,13 @@ def assert_agent_ready( # noqa: C901, PLR0912, PLR0913 (complexity accepted: on # because a manifest row's columns come from the file and cannot be declared in advance. if c.cardinality in {"1:N fan-out", "N:1"} and results: for position, result in enumerate(results): - row = _data_of(result) + output_keys = _data_keys_of(result) for key in c.writes.data_keys: - assert key in row, f"{name}: declared write {key!r} missing from result {position}" + assert key in output_keys, f"{name}: declared write {key!r} missing from result {position}" for key in c.removes_keys: - assert key not in row, f"{name}: declared removes_keys {key!r} still present in result {position}" + assert key not in output_keys, ( + f"{name}: declared removes_keys {key!r} still present in result {position}" + ) # (4) no undeclared top-level keys (non-fanout) if c.cardinality in {"1:1", "1:1 nested-list", "filter"} and results: diff --git a/nemo_curator/stages/audio/_agent/_planning.py b/nemo_curator/stages/audio/_agent/_planning.py index 975bbb1193..7bd9181daa 100644 --- a/nemo_curator/stages/audio/_agent/_planning.py +++ b/nemo_curator/stages/audio/_agent/_planning.py @@ -699,6 +699,8 @@ def validate_pipeline( # noqa: PLR0913 -- keyword-only seeds of one input task, the default -- infers them from ``initial_keys`` by role, which covers the canonical ``waveform``. Pass this when the input carries a tensor under a name whose role cannot be inferred (e.g. ``audio_tensor``). + Pass an empty set when a schema contains a waveform-named column but + the input values are known not to be resident tensors. initial_task_type: Class name of the task the first stage will be handed (e.g. ``"EmptyTask"`` for a pipeline that starts at a source, ``"AudioTask"`` for a suffix resumed from a manifest). ``None`` -- the default -- leaves the diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index d804d3e3e7..e669e06aca 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -16,17 +16,22 @@ from __future__ import annotations +from types import SimpleNamespace from typing import TYPE_CHECKING +import pandas as pd import pytest import soundfile as sf import torch from nemo_curator.stages import audio from nemo_curator.stages.audio import agent +from nemo_curator.stages.audio._agent import _catalog +from nemo_curator.stages.audio._agent._agent_ready import AgentReady, IOSpec, StageContract from nemo_curator.stages.audio._agent._agent_registry import build_contract, static_contract from nemo_curator.stages.audio._agent._catalog import unavailable_modules from nemo_curator.stages.audio._agent._composite import expand_composites +from nemo_curator.stages.audio._agent._conformance import assert_agent_ready from nemo_curator.stages.audio._agent._planning import validate_pipeline from nemo_curator.stages.audio._agent._residency import resolve_audio, write_audio_stable from nemo_curator.stages.audio.common import ( @@ -44,7 +49,8 @@ SampleRateFilterStage, SegmentConcatenationStage, ) -from nemo_curator.tasks import AudioTask, FileGroupTask +from nemo_curator.stages.base import ProcessingStage +from nemo_curator.tasks import AudioTask, DocumentBatch, FileGroupTask if TYPE_CHECKING: from pathlib import Path @@ -107,6 +113,34 @@ def test_public_facade_exposes_unavailable_modules_and_folder_source() -> None: assert "CreateInitialManifestAudioFolderStage" in audio.__all__ +def test_public_discovery_reports_an_optional_import_failure(monkeypatch: pytest.MonkeyPatch) -> None: + """A partial install exposes skipped modules through the public facade.""" + missing_module = "nemo_curator.stages.audio.optional_missing" + monkeypatch.setattr(_catalog, "_IMPORTED", False) + monkeypatch.setattr(_catalog, "_SKIPPED", []) + monkeypatch.setattr( + _catalog.pkgutil, + "walk_packages", + lambda *_args, **_kwargs: [SimpleNamespace(name=missing_module)], + ) + + def fail_optional_import(name: str) -> None: + assert name == missing_module + message = "optional dependency is not installed" + raise ModuleNotFoundError(message) + + monkeypatch.setattr(_catalog.importlib, "import_module", fail_optional_import) + with pytest.warns(UserWarning, match="optional_missing"): + missing = agent.unavailable_modules() + + assert missing == [ + { + "module": missing_module, + "error": "ModuleNotFoundError: optional dependency is not installed", + } + ] + + def _stereo_task(tmp_path: Path, sample_rate: int = 16000) -> tuple[AudioTask, str]: """A row carrying BOTH a resident stereo waveform and the file it came from.""" path = str(tmp_path / "stereo.wav") @@ -263,6 +297,18 @@ def test_a_tensor_under_an_uninferable_name_can_be_declared_resident(tmp_path: P assert not validate_pipeline([writer], initial_keys={"audio_tensor"}, initial_tensor_keys={"audio_tensor"}).ok +def test_an_explicit_empty_tensor_seed_overrides_waveform_name_inference(tmp_path: Path) -> None: + """A nullable waveform-named manifest column is not automatically a resident tensor.""" + writer = ManifestWriterStage(output_path=str(tmp_path / "out.jsonl")) + report = validate_pipeline( + [writer], + initial_keys={"waveform"}, + initial_tensor_keys=set(), + ) + assert report.ok + assert not any(issue.code == "tensor_into_sink" for issue in report.issues) + + def test_a_plain_manifest_input_still_reaches_a_json_sink(tmp_path: Path) -> None: """The seeding must not make every pipeline look tensor-resident.""" report = validate_pipeline([ManifestWriterStage(output_path=str(tmp_path / "out.jsonl"))]) @@ -280,6 +326,12 @@ def test_a_custom_manifest_path_column_is_what_the_reader_declares(tmp_path: Pat emitted = reader.process(FileGroupTask(dataset_name="d", data=[str(manifest)])) assert "recording_path" in emitted[0].data assert "audio_filepath" not in emitted[0].data + assert_agent_ready( + reader, + lambda: FileGroupTask(dataset_name="d", data=[str(manifest)]), + expected_cardinality="1:N fan-out", + available_keys=set(), + ) # Seeded empty because the input is a FileGroupTask of manifest PATHS: it carries no # audio columns, and the default seed would otherwise supply the very ``audio_filepath`` @@ -296,6 +348,61 @@ def test_a_custom_manifest_path_column_is_what_the_reader_declares(tmp_path: Pat assert validate_pipeline([ManifestReaderStage(), MonoConversionStage()], **seed).keys_ok +def test_fanout_conformance_checks_every_emitted_result(tmp_path: Path) -> None: + """A later fan-out row cannot omit a write that only the first row carries.""" + manifest = tmp_path / "mixed.jsonl" + manifest.write_text( + '{"recording_path": "/tmp/a.wav"}\n{"text": "missing the declared recording_path"}\n', + encoding="utf-8", + ) + reader = ManifestReaderStage(include_files_key="recording_path") + + with pytest.raises(AssertionError, match="missing from result 1"): + assert_agent_ready( + reader, + lambda: FileGroupTask(dataset_name="d", data=[str(manifest)]), + expected_cardinality="1:N fan-out", + available_keys=set(), + ) + + +class _DataFrameFanInStage(AgentReady, ProcessingStage[AudioTask, DocumentBatch]): + """Small stand-in for the full agent branch's AudioToDocumentStage.""" + + BATCH_ONLY = True + name = "dataframe_fan_in" + + def describe(self) -> StageContract: + return StageContract( + writes=IOSpec(data_keys=["text"]), + cardinality="N:1", + ) + + def process(self, _task: AudioTask) -> DocumentBatch: + raise NotImplementedError + + def process_batch(self, tasks: list[AudioTask]) -> list[DocumentBatch]: + return [ + DocumentBatch( + dataset_name=tasks[0].dataset_name, + data=pd.DataFrame([{"text": task.data["text"]} for task in tasks]), + ) + ] + + +def test_n_to_one_conformance_reads_document_batch_columns() -> None: + """Declared N:1 writes are DataFrame columns in a DocumentBatch, not dict keys.""" + assert_agent_ready( + _DataFrameFanInStage(), + lambda: [ + AudioTask(dataset_name="d", data={"text": "one"}), + AudioTask(dataset_name="d", data={"text": "two"}), + ], + expected_cardinality="N:1", + available_keys={"text"}, + ) + + def test_a_null_waveform_does_not_authenticate_a_stale_sample_rate(tmp_path: Path) -> None: """Residency is about the VALUE; a present-but-empty column must not vouch for metadata.""" path = tmp_path / "a.wav" From 26a2da88d27f77207b52183611e47d1a54ebbc2a Mon Sep 17 00:00:00 2001 From: Shubham Bhawsar Date: Wed, 9 Sep 2026 07:43:48 +0000 Subject: [PATCH 06/25] fix(audio): avoid folder item ID collisions Signed-off-by: Shubham Bhawsar --- nemo_curator/stages/audio/common.py | 30 +++++++++++++++++-- .../test_create_manifest_audio_folder.py | 14 +++++++++ 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/nemo_curator/stages/audio/common.py b/nemo_curator/stages/audio/common.py index 42012ec70b..d6c3562f4a 100644 --- a/nemo_curator/stages/audio/common.py +++ b/nemo_curator/stages/audio/common.py @@ -705,6 +705,32 @@ def ray_stage_spec(self) -> dict[str, Any]: def num_workers(self) -> int | None: return 1 + @staticmethod + def _item_id(relative_path: str) -> str: + """Flatten a relative stem without confusing separators with filename text.""" + stem = os.path.splitext(relative_path)[0] + encoded_components: list[str] = [] + for component in stem.split(os.sep): + encoded: list[str] = [] + last = len(component) - 1 + for index, char in enumerate(component): + if char == "~": + # Reserve ``~`` for the underscore escape below. + encoded.append("~~") + elif char == "_" and ( + index in (0, last) + or component[index - 1] == "_" + or component[index + 1] == "_" + ): + # Encoded components must neither contain ``__`` nor touch a + # separator with ``_``; otherwise two different component + # boundaries can flatten to the same string. + encoded.append("~u") + else: + encoded.append(char) + encoded_components.append("".join(encoded)) + return "__".join(encoded_components) + def _collect_audio_files(self) -> list[str]: exts = tuple((e if e.startswith(".") else f".{e}").lower() for e in self.extensions) if not os.path.isdir(self.data_dir): @@ -746,9 +772,9 @@ def process(self, _: EmptyTask) -> list[AudioTask]: # Relpath, not basename: ``recursive`` defaults True and speaker-per-folder is the # standard layout, so a basename id gives spk1/utt1.wav and spk2/utt1.wav the same # id -- and downstream that id becomes an output filename. A flat corpus is - # unaffected. Not injective: a flat ``spk1__utt1.wav`` still aliases spk1/utt1.wav. + # unaffected unless its name needs escaping to remain distinct from a path separator. rel = os.path.relpath(abspath, os.path.abspath(self.data_dir)) - item_id = os.path.splitext(rel)[0].replace(os.sep, "__") + item_id = self._item_id(rel) tasks.append( AudioTask( dataset_name="local-audio-folder", diff --git a/tests/stages/audio/test_create_manifest_audio_folder.py b/tests/stages/audio/test_create_manifest_audio_folder.py index da408965ab..a312121533 100644 --- a/tests/stages/audio/test_create_manifest_audio_folder.py +++ b/tests/stages/audio/test_create_manifest_audio_folder.py @@ -48,6 +48,20 @@ def test_same_filename_in_two_folders_gets_two_ids(self, tmp_path) -> None: # n assert ids == ["spk1__utt1", "spk2__utt1"], ids + def test_flattened_path_and_literal_separator_get_distinct_ids(self, tmp_path) -> None: # noqa: ANN001 + root = str(tmp_path) + for rel in ["spk1/utt1.wav", "spk1__utt1.wav"]: + _touch(root, rel) + + tasks = CreateInitialManifestAudioFolderStage(data_dir=root).process(None) + ids_by_path = { + os.path.relpath(task.data["audio_filepath"], root): task.data["audio_item_id"] for task in tasks + } + + assert ids_by_path[os.path.join("spk1", "utt1.wav")] == "spk1__utt1" + assert ids_by_path["spk1__utt1.wav"] == "spk1~u~uutt1" + assert len(set(ids_by_path.values())) == 2 + def test_a_flat_folder_keeps_the_plain_ids_it_always_had(self, tmp_path) -> None: # noqa: ANN001 """relpath IS the basename for a flat corpus, so those ids must not move.""" root = str(tmp_path) From 89dde01675c4e386aa17f29903950155d6c2a823 Mon Sep 17 00:00:00 2001 From: Shubham Bhawsar Date: Thu, 10 Sep 2026 11:10:13 +0000 Subject: [PATCH 07/25] fix(audio): support scoped tagging contracts Signed-off-by: Shubham Bhawsar --- .../stages/audio/_agent/_conformance.py | 5 +- nemo_curator/stages/audio/_agent/_planning.py | 162 ++++++++++++++---- .../stages/audio/_agent/_residency.py | 37 ++-- .../test_agent_foundation_regressions.py | 35 +++- 4 files changed, 194 insertions(+), 45 deletions(-) diff --git a/nemo_curator/stages/audio/_agent/_conformance.py b/nemo_curator/stages/audio/_agent/_conformance.py index 536b4919f0..7e40f47510 100644 --- a/nemo_curator/stages/audio/_agent/_conformance.py +++ b/nemo_curator/stages/audio/_agent/_conformance.py @@ -417,7 +417,10 @@ def assert_agent_ready( # noqa: C901, PLR0912, PLR0913 (complexity accepted: on assert key in out_data, f"{name}: declared write {key!r} missing from task.data" for key in c.removes_keys: assert key not in out_data, f"{name}: declared removes_keys {key!r} but it is still present in task.data" - declared = set(c.writes.data_keys) | set(ignore_new_keys) | input_keys + conditional_keys = { + key for conditional in c.conditional_writes for key in conditional.writes.data_keys + } + declared = set(c.writes.data_keys) | conditional_keys | set(ignore_new_keys) | input_keys undeclared = set(out_data) - declared assert not undeclared, f"{name}: undeclared new top-level keys {sorted(undeclared)} (add to writes.data_keys)" # segment-level writes diff --git a/nemo_curator/stages/audio/_agent/_planning.py b/nemo_curator/stages/audio/_agent/_planning.py index 7bd9181daa..49c54776e1 100644 --- a/nemo_curator/stages/audio/_agent/_planning.py +++ b/nemo_curator/stages/audio/_agent/_planning.py @@ -57,7 +57,6 @@ from nemo_curator.stages.audio._agent._agent_registry import build_contract from nemo_curator.stages.audio._agent._composite import expand_composites -from nemo_curator.stages.audio._agent._conformance import produced_roles, reads_satisfied_by_role from nemo_curator.stages.audio._agent._roles import role_for_value if TYPE_CHECKING: @@ -158,8 +157,46 @@ def _requirement_str(contract: StageContract, available: set[str]) -> str: def _write_key_values(contract: StageContract) -> set[str]: - """The literal key VALUES a stage writes (top-level + segment-level).""" - return {*contract.writes.data_keys, *contract.writes.segment_data_keys} + """The literal top-level key VALUES a stage writes.""" + return set(contract.writes.data_keys) + + +def _segment_write_key_values(contract: StageContract) -> set[str]: + """The literal key VALUES a stage writes inside nested-list items.""" + return set(contract.writes.segment_data_keys) + + +def _roles_for_keys(contract: StageContract, keys: set[str] | list[str]) -> set[str]: + """Known semantic roles for key values in one task-data scope.""" + return {contract.key_roles.get(key, "unknown") for key in keys} - {"unknown"} + + +def _spec_satisfied_by_role( + spec: Any, # noqa: ANN401 - IOSpec, kept loose to avoid a runtime-only import + contract: StageContract, + available_roles: set[str], + available_segment_roles: set[str], +) -> bool: + """Whether one I/O alternative is role-satisfied in each declared scope.""" + top = {contract.key_roles.get(key, "unknown") for key in spec.data_keys} + nested = {contract.key_roles.get(key, "unknown") for key in spec.segment_data_keys} + return top.issubset(available_roles | {"unknown"}) and nested.issubset( + available_segment_roles | {"unknown"} + ) + + +def _reads_satisfied_by_role( + contract: StageContract, + available_roles: set[str], + available_segment_roles: set[str], +) -> bool: + """Role-level read check that keeps task and nested-item fields separate.""" + if not _spec_satisfied_by_role(contract.reads, contract, available_roles, available_segment_roles): + return False + return not contract.reads_one_of or any( + _spec_satisfied_by_role(option, contract, available_roles, available_segment_roles) + for option in contract.reads_one_of + ) def _key_family(key: str) -> str: @@ -295,13 +332,34 @@ def _ambiguity_issues( return out -def _missing_read_keys(contract: StageContract, available_keys: set[str]) -> set[str]: +def _missing_read_keys( + contract: StageContract, + available_keys: set[str], + available_segment_keys: set[str], +) -> set[str]: """Read key VALUES this stage wants that nothing upstream produced or seeded.""" - reads = {*contract.reads.data_keys, *contract.reads.segment_data_keys} - return {k for k in reads if k not in available_keys} + return { + *{key for key in contract.reads.data_keys if key not in available_keys}, + *{key for key in contract.reads.segment_data_keys if key not in available_segment_keys}, + } + +def _spec_satisfied_by_key( + spec: Any, # noqa: ANN401 - IOSpec, kept loose to avoid a runtime-only import + available_keys: set[str], + available_segment_keys: set[str], +) -> bool: + """Whether one I/O alternative's literal keys exist in their declared scopes.""" + return not (set(spec.data_keys) - available_keys) and not ( + set(spec.segment_data_keys) - available_segment_keys + ) -def _reads_satisfied_by_key(contract: StageContract, available_keys: set[str]) -> bool: + +def _reads_satisfied_by_key( + contract: StageContract, + available_keys: set[str], + available_segment_keys: set[str], +) -> bool: """Whether every read is met by the LITERAL key it names. A stage reads ``task.data[self.segments_key]`` at runtime -- a key string, never a role. So @@ -316,11 +374,12 @@ def _reads_satisfied_by_key(contract: StageContract, available_keys: set[str]) - names differ but mean the same thing (a producer writing ``resampled_audio_filepath`` satisfying a consumer reading ``audio_filepath``). A read is satisfied by either route. """ - if {*contract.reads.data_keys, *contract.reads.segment_data_keys} - available_keys: + if not _spec_satisfied_by_key(contract.reads, available_keys, available_segment_keys): return False - if not contract.reads_one_of: - return True - return any(not ({*spec.data_keys, *spec.segment_data_keys} - available_keys) for spec in contract.reads_one_of) + return not contract.reads_one_of or any( + _spec_satisfied_by_key(option, available_keys, available_segment_keys) + for option in contract.reads_one_of + ) def _forwarding_param(inner: Any, composite: Any, missing: set[str]) -> str | None: # noqa: ANN401 @@ -379,7 +438,11 @@ def _describes_itself(stage: Any) -> bool: # noqa: ANN401 - any child stage return True -def _dangling_read_keys(contract: StageContract, available_keys: set[str]) -> set[str]: +def _dangling_read_keys( + contract: StageContract, + available_keys: set[str], + available_segment_keys: set[str], +) -> set[str]: """Read key VALUES whose role is known but whose exact value was not produced upstream nor seeded — the renamed-producer dangle the role check misses. @@ -391,16 +454,18 @@ def _dangling_read_keys(contract: StageContract, available_keys: set[str]) -> se (``unknown``/internal bookkeeping keys are excluded — a separate value-identity check for those is tracked in the backlog). """ - reads = [*contract.reads.data_keys, *contract.reads.segment_data_keys] + reads = [(key, available_keys) for key in contract.reads.data_keys] + reads += [(key, available_segment_keys) for key in contract.reads.segment_data_keys] if len(contract.reads_one_of) == 1: only = contract.reads_one_of[0] - reads += [*only.data_keys, *only.segment_data_keys] + reads += [(key, available_keys) for key in only.data_keys] + reads += [(key, available_segment_keys) for key in only.segment_data_keys] dangling: set[str] = set() - for k in reads: + for k, scope_keys in reads: role = contract.key_roles.get(k, "unknown") if role == "unknown": continue - if k not in available_keys: + if k not in scope_keys: dangling.add(k) return dangling @@ -409,11 +474,14 @@ def _dangling_read_keys(contract: StageContract, available_keys: set[str]) -> se class _Walk: """What the pipeline carries from one stage to the next while being validated.""" - available: set[str] # roles produced so far - available_keys: set[str] # literal key VALUES produced so far + available: set[str] # top-level roles produced so far + available_keys: set[str] # literal top-level key VALUES produced so far + segment_available: set[str] = field(default_factory=set) # nested-item roles produced so far + segment_available_keys: set[str] = field(default_factory=set) # literal nested-item key VALUES tensor_keys: set[str] = field(default_factory=set) removed_roles: set[str] = field(default_factory=set) key_producer: dict[str, str] = field(default_factory=dict) + segment_key_producer: dict[str, str] = field(default_factory=dict) past_composite: bool = False # an UNEXPANDABLE composite hid its writes; reads past it can't be judged task_type: str | None = None # task type the previous stage produces; None == not known @@ -428,11 +496,13 @@ def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[Pipe for now -- the expansion is new, and it earns the right to block only once it has been shown not to false-positive on pipelines known to work. """ - if reads_satisfied_by_role(contract, walk.available) or _reads_satisfied_by_key(contract, walk.available_keys): + role_satisfied = _reads_satisfied_by_role(contract, walk.available, walk.segment_available) + key_satisfied = _reads_satisfied_by_key(contract, walk.available_keys, walk.segment_available_keys) + if role_satisfied or key_satisfied: if walk.past_composite: return [] out: list[PipelineIssue] = [] - dangling = _dangling_read_keys(contract, walk.available_keys) + dangling = _dangling_read_keys(contract, walk.available_keys, walk.segment_available_keys) if dangling: out.append( PipelineIssue( @@ -440,17 +510,25 @@ def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[Pipe site.name, "warning", "dangling_key", - f"reads key(s) {sorted(dangling)} satisfied by role but not produced " - f"upstream under that key value nor seeded (renamed producer key?); " - f"available keys: {sorted(walk.available_keys)}", + f"reads key(s) {sorted(dangling)} satisfied by role but not produced upstream " + f"under that key value in the required task/nested scope nor seeded " + f"(renamed producer key?); available keys: " + f"{sorted(walk.available_keys | walk.segment_available_keys)}", ) ) - out.extend(_ambiguity_issues(site, contract, walk.available_keys, walk.key_producer)) + out.extend( + _ambiguity_issues( + site, + contract, + walk.available_keys | walk.segment_available_keys, + walk.key_producer | walk.segment_key_producer, + ) + ) return out if site.composite is not None: composite_name = type(site.composite).__name__ - missing = _missing_read_keys(contract, walk.available_keys) + missing = _missing_read_keys(contract, walk.available_keys, walk.segment_available_keys) param = _forwarding_param(site.stage, site.composite, missing) remedy = ( f"set {param} on {composite_name} (it forwards the value to this inner stage)" @@ -466,7 +544,8 @@ def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[Pipe f"this stage runs inside {composite_name} and requires " f"{_requirement_str(contract, walk.available)}" + (f" (key(s) {sorted(missing)})" if missing else "") - + f", not produced upstream; {remedy}. Available keys: {sorted(walk.available_keys)}", + + f", not produced upstream; {remedy}. Available keys: " + + f"{sorted(walk.available_keys | walk.segment_available_keys)}", ) ] @@ -483,7 +562,11 @@ def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[Pipe ) ] - needed = _required_roles(contract) | {r for o in contract.reads_one_of for r in _roles_of(o, contract)} + needed = _roles_for_keys(contract, contract.reads.data_keys) | { + role + for option in contract.reads_one_of + for role in _roles_for_keys(contract, option.data_keys) + } removed_hit = (needed & walk.removed_roles) - walk.available if removed_hit: return [ @@ -588,8 +671,10 @@ def _task_type_issue(walk: _Walk, site: _Site, contract: StageContract) -> list[ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: """Fold one stage's writes, removals and tensor residency into the running state.""" - produced = produced_roles(contract) + produced = _roles_for_keys(contract, contract.writes.data_keys) + segment_produced = _roles_for_keys(contract, contract.writes.segment_data_keys) written = _write_key_values(contract) + segment_written = _segment_write_key_values(contract) if not contract.preserves_upstream_keys: # A stage that rebuilds the task rather than adding to it: whatever it does not write # is not downstream. Folding its writes into the inherited state would keep every @@ -599,11 +684,17 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: # its own authority rather than on the vanished producer's. dropped_keys = walk.available_keys - written dropped_roles = walk.available - produced + dropped_segment_keys = walk.segment_available_keys - segment_written + dropped_segment_roles = walk.segment_available - segment_produced walk.available_keys -= dropped_keys walk.available -= dropped_roles + walk.segment_available_keys -= dropped_segment_keys + walk.segment_available -= dropped_segment_roles walk.removed_roles |= dropped_roles for key in dropped_keys: walk.key_producer.pop(key, None) + for key in dropped_segment_keys: + walk.segment_key_producer.pop(key, None) # Tensor residency deliberately survives this. The flag is coarser than it looks: # ALMDataBuilderStage sets it because SOME branch rebuilds task.data, while still # carrying the waveform on the ordinary path. Clearing residency here would retract @@ -613,9 +704,12 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: # or ``sanitizes_output``, both handled below. walk.available |= produced walk.removed_roles -= produced # a re-produced role is no longer "removed" + walk.segment_available |= segment_produced # Most recent writer wins -- that is who a downstream reader would actually get. walk.key_producer.update(dict.fromkeys(written, name)) walk.available_keys |= written + walk.segment_key_producer.update(dict.fromkeys(segment_written, name)) + walk.segment_available_keys |= segment_written for rk in contract.removes_keys: walk.available_keys.discard(rk) # Dropping the carrier ends the tensor residency as surely as sanitizing does. @@ -634,7 +728,11 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: # returned "unknown", so residency tracked ``_UNNAMED_TENSOR`` instead of the real # carrier -- and a downstream stage dropping that carrier still looked resident, # raising a spurious ``tensor_into_sink`` on a recipe that had cleaned up correctly. - carriers = {k for k in written if contract.key_roles.get(k, role_for_value(k)) == _TENSOR_ROLE} + carriers = { + key + for key in written | segment_written + if contract.key_roles.get(key, role_for_value(key)) == _TENSOR_ROLE + } walk.tensor_keys |= carriers or {_UNNAMED_TENSOR} if contract.gates.sanitizes_output: walk.tensor_keys.clear() @@ -828,7 +926,11 @@ def validate_pipeline( # noqa: PLR0913 -- keyword-only seeds of one input task, # two stages stale. walk.task_type = contract.produces_task_type - return PipelineReport(issues=issues, produced_roles=walk.available, produced_keys=walk.available_keys) + return PipelineReport( + issues=issues, + produced_roles=walk.available | walk.segment_available, + produced_keys=walk.available_keys | walk.segment_available_keys, + ) def _roles_of(spec: Any, contract: StageContract) -> set[str]: # noqa: ANN401 diff --git a/nemo_curator/stages/audio/_agent/_residency.py b/nemo_curator/stages/audio/_agent/_residency.py index 385adebdf4..3e22996de4 100644 --- a/nemo_curator/stages/audio/_agent/_residency.py +++ b/nemo_curator/stages/audio/_agent/_residency.py @@ -315,7 +315,7 @@ def write_audio_stable( return path -def resolve_audio_path( # noqa: PLR0913 (complexity accepted: keyword-only residency/key knobs mirror the stage fields) +def resolve_audio_path( # noqa: C901, PLR0913 (keyword-only residency/key knobs mirror stage fields) item: dict[str, Any], *, residency: InputResidency = "auto", @@ -327,14 +327,35 @@ def resolve_audio_path( # noqa: PLR0913 (complexity accepted: keyword-only resi ) -> str | None: """Return an audio path, writing a temp WAV when only a waveform exists. + ``auto`` prefers a complete resident waveform/sample-rate pair and falls + back to the configured path. ``file`` always uses the configured path. + 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. """ + if residency != "file": + waveform = item.get(waveform_key) + sample_rate = item.get(sample_rate_key) + if waveform is not None and sample_rate is not None: + fd, tmp = tempfile.mkstemp(suffix=".wav", dir=temp_dir) + os.close(fd) + try: + sf.write(tmp, _as_soundfile_array(waveform), int(sample_rate)) + except BaseException: + with contextlib.suppress(OSError): + os.remove(tmp) + raise + if register_temp is not None: + register_temp.append(tmp) + return tmp + if residency == "waveform": + return None + path = item.get(audio_filepath_key) local_path: str | None = None - if residency != "waveform" and path: + if path: local_path = os.path.expanduser(str(path)) if os.path.exists(local_path): return local_path @@ -357,17 +378,7 @@ def resolve_audio_path( # noqa: PLR0913 (complexity accepted: keyword-only resi # 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 + return local_path def cleanup_temp_files(paths: list[str] | None) -> None: diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index e669e06aca..7c775ff905 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -16,6 +16,7 @@ from __future__ import annotations +import os from types import SimpleNamespace from typing import TYPE_CHECKING @@ -33,7 +34,12 @@ from nemo_curator.stages.audio._agent._composite import expand_composites from nemo_curator.stages.audio._agent._conformance import assert_agent_ready from nemo_curator.stages.audio._agent._planning import validate_pipeline -from nemo_curator.stages.audio._agent._residency import resolve_audio, write_audio_stable +from nemo_curator.stages.audio._agent._residency import ( + cleanup_temp_files, + resolve_audio, + resolve_audio_path, + write_audio_stable, +) from nemo_curator.stages.audio.common import ( CreateInitialManifestAudioFolderStage, ManifestCheckpointStage, @@ -56,6 +62,33 @@ from pathlib import Path +def test_resolve_audio_path_auto_prefers_complete_resident_audio(tmp_path: Path) -> None: + """Auto residency must not silently choose a stale file over a complete waveform.""" + file_path = tmp_path / "one_second.wav" + sf.write(file_path, torch.zeros(16000).numpy(), 16000) + resident = torch.ones(1, 32000) + item = { + "audio_filepath": str(file_path), + "waveform": resident, + "sample_rate": 16000, + } + temporary_paths: list[str] = [] + + resolved = resolve_audio_path(item, residency="auto", temp_dir=str(tmp_path), register_temp=temporary_paths) + + assert resolved is not None + assert resolved != str(file_path) + assert temporary_paths == [resolved] + loaded, sample_rate = sf.read(resolved) + assert sample_rate == 16000 + assert len(loaded) == 32000 + assert loaded.mean() > 0.9 + assert resolve_audio_path(item, residency="file") == str(file_path) + + cleanup_temp_files(temporary_paths) + assert not os.path.exists(resolved) + + def test_stable_audio_names_include_layout_and_written_short_stereo_shape(tmp_path: Path) -> None: """Different channel layouts with identical samples need distinct artifacts.""" output_dir = str(tmp_path) From 54f655ef9468107985e6de5d971a91a64306c58d Mon Sep 17 00:00:00 2001 From: Shubham Bhawsar Date: Tue, 15 Sep 2026 09:22:11 +0000 Subject: [PATCH 08/25] fix(audio-agent): validate hidden and residency parameters Signed-off-by: Shubham Bhawsar --- .../stages/audio/_agent/_agent_registry.py | 7 +++++- .../stages/audio/_agent/_residency.py | 7 ++++++ .../test_agent_foundation_regressions.py | 24 ++++++++++++++++++- 3 files changed, 36 insertions(+), 2 deletions(-) diff --git a/nemo_curator/stages/audio/_agent/_agent_registry.py b/nemo_curator/stages/audio/_agent/_agent_registry.py index 718e5d29cf..3e2e6eec82 100644 --- a/nemo_curator/stages/audio/_agent/_agent_registry.py +++ b/nemo_curator/stages/audio/_agent/_agent_registry.py @@ -191,7 +191,12 @@ def _dataclass_params(cls: type, descriptions: dict[str, str]) -> list[ParamSpec globalns = _module_globals(cls) params: list[ParamSpec] = [] for f in dataclasses.fields(cls): - if not f.init or f.name in EXCLUDED_PARAM_NAMES or f.name.startswith("_"): + if ( + not f.init + or f.metadata.get("agent_param", True) is False + or f.name in EXCLUDED_PARAM_NAMES + or f.name.startswith("_") + ): continue if f.default is not _MISSING: default, required = f.default, False diff --git a/nemo_curator/stages/audio/_agent/_residency.py b/nemo_curator/stages/audio/_agent/_residency.py index 3e22996de4..b4158907fe 100644 --- a/nemo_curator/stages/audio/_agent/_residency.py +++ b/nemo_curator/stages/audio/_agent/_residency.py @@ -31,6 +31,13 @@ InputResidency = Literal["file", "waveform", "auto"] +def validate_input_residency(residency: str, *, stage_name: str) -> None: + """Reject unknown residency modes before they can be treated as ``auto``.""" + if residency not in {"file", "waveform", "auto"}: + msg = f"[{stage_name}] input_residency must be one of 'file', 'waveform', or 'auto'; got {residency!r}" + raise ValueError(msg) + + def accepts_for_residency(residency: str) -> list[AudioForm]: """Audio forms an instance actually consumes, given its ``input_residency``. diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index 7c775ff905..a514d4b0c0 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -17,6 +17,7 @@ from __future__ import annotations import os +from dataclasses import dataclass, field from types import SimpleNamespace from typing import TYPE_CHECKING @@ -29,7 +30,7 @@ from nemo_curator.stages.audio import agent from nemo_curator.stages.audio._agent import _catalog from nemo_curator.stages.audio._agent._agent_ready import AgentReady, IOSpec, StageContract -from nemo_curator.stages.audio._agent._agent_registry import build_contract, static_contract +from nemo_curator.stages.audio._agent._agent_registry import build_contract, stage_params, static_contract from nemo_curator.stages.audio._agent._catalog import unavailable_modules from nemo_curator.stages.audio._agent._composite import expand_composites from nemo_curator.stages.audio._agent._conformance import assert_agent_ready @@ -38,6 +39,7 @@ cleanup_temp_files, resolve_audio, resolve_audio_path, + validate_input_residency, write_audio_stable, ) from nemo_curator.stages.audio.common import ( @@ -62,6 +64,26 @@ from pathlib import Path +@dataclass +class _AgentParamMetadataFixture: + visible: str = "public" + runtime_only: object | None = field(default=None, metadata={"agent_param": False}) + + +def test_stage_params_respects_field_level_agent_exclusion() -> None: + assert [param.name for param in stage_params(_AgentParamMetadataFixture)] == ["visible"] + + +@pytest.mark.parametrize("residency", ["file", "waveform", "auto"]) +def test_input_residency_validator_accepts_only_declared_modes(residency: str) -> None: + validate_input_residency(residency, stage_name="Fixture") + + +def test_input_residency_validator_rejects_unknown_mode() -> None: + with pytest.raises(ValueError, match="input_residency must be one of"): + validate_input_residency("wavefrom", stage_name="Fixture") + + def test_resolve_audio_path_auto_prefers_complete_resident_audio(tmp_path: Path) -> None: """Auto residency must not silently choose a stale file over a complete waveform.""" file_path = tmp_path / "one_second.wav" From 9db13eb9857366b7c6342925dd6eff1bdfdde49a Mon Sep 17 00:00:00 2001 From: Shubham Bhawsar Date: Tue, 15 Sep 2026 16:57:04 +0000 Subject: [PATCH 09/25] fix(audio-agent): support conditional and nested contracts Signed-off-by: Shubham Bhawsar --- .../stages/audio/_agent/_agent_registry.py | 6 +- .../stages/audio/_agent/_conformance.py | 10 ++- nemo_curator/stages/audio/_agent/_planning.py | 41 +++++++-- .../test_agent_foundation_regressions.py | 83 ++++++++++++++++++- 4 files changed, 130 insertions(+), 10 deletions(-) diff --git a/nemo_curator/stages/audio/_agent/_agent_registry.py b/nemo_curator/stages/audio/_agent/_agent_registry.py index 3e2e6eec82..fe379b72dd 100644 --- a/nemo_curator/stages/audio/_agent/_agent_registry.py +++ b/nemo_curator/stages/audio/_agent/_agent_registry.py @@ -526,7 +526,11 @@ def build_contract(stage: Any) -> StageContract: # noqa: ANN401 params = list(by_name.values()) else: params = derived - key_roles = _resolve_key_roles(stage, base) or dict(base.key_roles) + key_roles = _resolve_key_roles(stage, base) + # A hand-written contract can declare roles for literal keys that are not + # represented by ``*_key`` fields. Preserve those declarations while still + # enriching the contract with roles derived from configured parameters. + key_roles.update(base.key_roles) cls = _as_class(stage) accepts_tt, produces_tt = _task_types(cls) return dataclasses.replace( diff --git a/nemo_curator/stages/audio/_agent/_conformance.py b/nemo_curator/stages/audio/_agent/_conformance.py index 7e40f47510..58403bc771 100644 --- a/nemo_curator/stages/audio/_agent/_conformance.py +++ b/nemo_curator/stages/audio/_agent/_conformance.py @@ -68,8 +68,16 @@ def _spec_roles(contract: StageContract, spec_keys: Iterable[str]) -> set[str]: def produced_roles(producer: StageContract) -> set[str]: - """Roles a producer emits (from its ``writes`` keys); excludes ``unknown``.""" + """Roles a producer may emit from unconditional or conditional writes. + + This discovery helper intentionally includes possible outputs. Mechanical + planning advances only ``StageContract.writes`` and therefore does not + treat these conditional roles as guaranteed. + """ keys = [*producer.writes.data_keys, *producer.writes.segment_data_keys] + for conditional in producer.conditional_writes: + keys.extend(conditional.writes.data_keys) + keys.extend(conditional.writes.segment_data_keys) return _spec_roles(producer, keys) - {"unknown"} diff --git a/nemo_curator/stages/audio/_agent/_planning.py b/nemo_curator/stages/audio/_agent/_planning.py index 49c54776e1..0c5f37e532 100644 --- a/nemo_curator/stages/audio/_agent/_planning.py +++ b/nemo_curator/stages/audio/_agent/_planning.py @@ -738,11 +738,14 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: walk.tensor_keys.clear() -def _seed_walk( +def _seed_walk( # noqa: PLR0913 -- top-level and nested seeds describe one input task initial_roles: set[str] | None, initial_keys: set[str] | None, initial_tensor_keys: set[str] | None, initial_task_type: str | None, + *, + initial_segment_roles: set[str] | None = None, + initial_segment_keys: set[str] | None = None, ) -> _Walk: """The state the first stage is handed: what the input task already carries.""" if initial_keys is not None: @@ -755,6 +758,14 @@ def _seed_walk( seed_keys = {r for r in initial_roles if role_for_value(r) == r} else: seed_keys = set(_DEFAULT_INITIAL_KEYS) + if initial_segment_keys is not None: + segment_seed_keys = set(initial_segment_keys) + elif initial_segment_roles is not None: + # Match top-level inference: only a semantic role that is also its + # canonical literal key can safely imply a key value. + segment_seed_keys = {r for r in initial_segment_roles if role_for_value(r) == r} + else: + segment_seed_keys = set() # An input that arrives carrying a waveform is exactly as resident as one a stage # produced, so the serialization gate has to see it. Only writes used to seed this, # which left the gate blind to the resident-input case validate_pipeline documents: @@ -763,10 +774,12 @@ def _seed_walk( if initial_tensor_keys is not None: seed_tensors = set(initial_tensor_keys) else: - seed_tensors = {k for k in seed_keys if role_for_value(k) == _TENSOR_ROLE} + seed_tensors = {k for k in seed_keys | segment_seed_keys if role_for_value(k) == _TENSOR_ROLE} return _Walk( available=set(initial_roles) if initial_roles is not None else set(_DEFAULT_INITIAL_ROLES), available_keys=seed_keys, + segment_available=set(initial_segment_roles or ()), + segment_available_keys=segment_seed_keys, tensor_keys=seed_tensors, task_type=initial_task_type, ) @@ -777,6 +790,8 @@ def validate_pipeline( # noqa: PLR0913 -- keyword-only seeds of one input task, *, initial_roles: set[str] | None = None, initial_keys: set[str] | None = None, + initial_segment_roles: set[str] | None = None, + initial_segment_keys: set[str] | None = None, initial_tensor_keys: set[str] | None = None, initial_task_type: str | None = None, available_gpus: float | None = None, @@ -793,10 +808,17 @@ def validate_pipeline( # noqa: PLR0913 -- keyword-only seeds of one input task, Defaults to ``{"audio_filepath"}``. Seeding this lets the literal-key check (``keys_ok``) recognize reads satisfied by the input rather than by an upstream producer. + initial_segment_roles: Semantic roles present inside the input task's + segment dictionaries. Defaults to an empty nested state. + initial_segment_keys: Literal key VALUES present inside the input + task's segment dictionaries. Defaults to an empty nested state. + When omitted with explicit ``initial_segment_roles``, canonical + role-as-literal key values are inferred using the top-level policy. initial_tensor_keys: Which seeded keys hold a resident tensor. ``None`` -- - the default -- infers them from ``initial_keys`` by role, which covers - the canonical ``waveform``. Pass this when the input carries a tensor - under a name whose role cannot be inferred (e.g. ``audio_tensor``). + the default -- infers them from top-level and segment seed keys by + role, which covers the canonical ``waveform``. Pass this when the + input carries a tensor under a name whose role cannot be inferred + (e.g. ``audio_tensor``). Pass an empty set when a schema contains a waveform-named column but the input values are known not to be resident tensors. initial_task_type: Class name of the task the first stage will be handed @@ -811,7 +833,14 @@ def validate_pipeline( # noqa: PLR0913 -- keyword-only seeds of one input task, found (role-level); ``report.keys_ok`` additionally confirms literal-key identity (see the class docstring). """ - walk = _seed_walk(initial_roles, initial_keys, initial_tensor_keys, initial_task_type) + walk = _seed_walk( + initial_roles, + initial_keys, + initial_tensor_keys, + initial_task_type, + initial_segment_roles=initial_segment_roles, + initial_segment_keys=initial_segment_keys, + ) expansion = expand_composites(stages) leaves = expansion.by_recipe_index() opaque = dict(expansion.opaque) diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index a514d4b0c0..8d0c394374 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -29,11 +29,11 @@ from nemo_curator.stages import audio from nemo_curator.stages.audio import agent from nemo_curator.stages.audio._agent import _catalog -from nemo_curator.stages.audio._agent._agent_ready import AgentReady, IOSpec, StageContract +from nemo_curator.stages.audio._agent._agent_ready import AgentReady, ConditionalWrite, IOSpec, StageContract from nemo_curator.stages.audio._agent._agent_registry import build_contract, stage_params, static_contract from nemo_curator.stages.audio._agent._catalog import unavailable_modules from nemo_curator.stages.audio._agent._composite import expand_composites -from nemo_curator.stages.audio._agent._conformance import assert_agent_ready +from nemo_curator.stages.audio._agent._conformance import assert_agent_ready, produced_roles from nemo_curator.stages.audio._agent._planning import validate_pipeline from nemo_curator.stages.audio._agent._residency import ( cleanup_temp_files, @@ -70,10 +70,75 @@ class _AgentParamMetadataFixture: runtime_only: object | None = field(default=None, metadata={"agent_param": False}) +class _ConfiguredContractStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): + def __init__(self, contract: StageContract) -> None: + self.contract = contract + + def describe(self) -> StageContract: + return self.contract + + def process(self, task: AudioTask) -> AudioTask: + return task + + def test_stage_params_respects_field_level_agent_exclusion() -> None: assert [param.name for param in stage_params(_AgentParamMetadataFixture)] == ["visible"] +def test_conditional_roles_are_discoverable_but_not_planner_guaranteed() -> None: + producer_contract = StageContract( + conditional_writes=[ + ConditionalWrite( + writes=IOSpec(data_keys=["potential_metrics"]), + condition="valid runtime data causes metric assignment", + ) + ], + key_roles={"potential_metrics": "metrics"}, + ) + assert produced_roles(producer_contract) == {"metrics"} + + consumer_contract = StageContract( + reads=IOSpec(data_keys=["potential_metrics"]), + key_roles={"potential_metrics": "metrics"}, + ) + report = validate_pipeline( + [_ConfiguredContractStage(producer_contract), _ConfiguredContractStage(consumer_contract)], + initial_roles=set(), + initial_keys=set(), + ) + + assert not report.ok + assert any(issue.code == "unsatisfied_reads" and issue.stage_index == 1 for issue in report.issues) + assert "potential_metrics" not in report.produced_keys + + +def test_nested_input_requires_explicit_segment_seeds_and_accepts_remapped_key() -> None: + nested_reader = _ConfiguredContractStage( + StageContract( + reads=IOSpec(data_keys=["segments"], segment_data_keys=["custom_text"]), + key_roles={"segments": "segments", "custom_text": "text"}, + ) + ) + + unseeded = validate_pipeline( + [nested_reader], + initial_roles={"segments", "text"}, + initial_keys={"segments", "custom_text"}, + ) + assert not unseeded.ok + assert any(issue.code == "unsatisfied_reads" for issue in unseeded.issues) + + seeded = validate_pipeline( + [nested_reader], + initial_roles={"segments"}, + initial_keys={"segments"}, + initial_segment_roles={"text"}, + initial_segment_keys={"custom_text"}, + ) + assert seeded.ok + assert seeded.keys_ok + + @pytest.mark.parametrize("residency", ["file", "waveform", "auto"]) def test_input_residency_validator_accepts_only_declared_modes(residency: str) -> None: validate_input_residency(residency, stage_name="Fixture") @@ -345,6 +410,20 @@ def test_an_input_that_arrives_with_a_waveform_is_blocked_from_a_json_sink(sink: stage.process(AudioTask(dataset_name="d", data={"waveform": torch.zeros(1, 16), "sample_rate": 16000})) +def test_nested_waveform_seed_is_inferred_as_tensor_resident(tmp_path: Path) -> None: + writer = ManifestWriterStage(output_path=str(tmp_path / "out.jsonl")) + report = validate_pipeline( + [writer], + initial_roles={"segments"}, + initial_keys={"segments"}, + initial_segment_roles={"waveform"}, + initial_segment_keys={"waveform"}, + ) + + assert not report.ok + assert any(issue.code == "tensor_into_sink" for issue in report.issues) + + def test_a_tensor_under_an_uninferable_name_can_be_declared_resident(tmp_path: Path) -> None: """A custom carrier has no role to infer from, so the seed has to be sayable outright.""" writer = ManifestWriterStage(output_path=str(tmp_path / "out.jsonl")) From 9a47a9a908c46c20136114261cfb321724176494 Mon Sep 17 00:00:00 2001 From: Shubham Bhawsar Date: Tue, 15 Sep 2026 18:14:15 +0000 Subject: [PATCH 10/25] fix(audio-agent): require literal keys for generic reads Signed-off-by: Shubham Bhawsar --- nemo_curator/stages/audio/_agent/_planning.py | 62 +++++++++++++------ .../test_agent_foundation_regressions.py | 31 ++++++++++ 2 files changed, 75 insertions(+), 18 deletions(-) diff --git a/nemo_curator/stages/audio/_agent/_planning.py b/nemo_curator/stages/audio/_agent/_planning.py index 0c5f37e532..bd73413500 100644 --- a/nemo_curator/stages/audio/_agent/_planning.py +++ b/nemo_curator/stages/audio/_agent/_planning.py @@ -171,17 +171,26 @@ def _roles_for_keys(contract: StageContract, keys: set[str] | list[str]) -> set[ return {contract.key_roles.get(key, "unknown") for key in keys} - {"unknown"} -def _spec_satisfied_by_role( +def _spec_satisfied_by_role( # noqa: PLR0913 -- task and nested scopes each need roles plus literal keys spec: Any, # noqa: ANN401 - IOSpec, kept loose to avoid a runtime-only import contract: StageContract, available_roles: set[str], available_segment_roles: set[str], + available_keys: set[str], + available_segment_keys: set[str], ) -> bool: - """Whether one I/O alternative is role-satisfied in each declared scope.""" - top = {contract.key_roles.get(key, "unknown") for key in spec.data_keys} - nested = {contract.key_roles.get(key, "unknown") for key in spec.segment_data_keys} - return top.issubset(available_roles | {"unknown"}) and nested.issubset( - available_segment_roles | {"unknown"} + """Whether known roles or exact unknown-role keys satisfy each scope.""" + + def scope_satisfied(keys: list[str], roles: set[str], literal_keys: set[str]) -> bool: + return all( + key in literal_keys if (role := contract.key_roles.get(key, "unknown")) == "unknown" else role in roles + for key in keys + ) + + return scope_satisfied(spec.data_keys, available_roles, available_keys) and scope_satisfied( + spec.segment_data_keys, + available_segment_roles, + available_segment_keys, ) @@ -189,12 +198,28 @@ def _reads_satisfied_by_role( contract: StageContract, available_roles: set[str], available_segment_roles: set[str], + available_keys: set[str], + available_segment_keys: set[str], ) -> bool: - """Role-level read check that keeps task and nested-item fields separate.""" - if not _spec_satisfied_by_role(contract.reads, contract, available_roles, available_segment_roles): + """Role-level read check with literal fallback for unknown roles.""" + if not _spec_satisfied_by_role( + contract.reads, + contract, + available_roles, + available_segment_roles, + available_keys, + available_segment_keys, + ): return False return not contract.reads_one_of or any( - _spec_satisfied_by_role(option, contract, available_roles, available_segment_roles) + _spec_satisfied_by_role( + option, + contract, + available_roles, + available_segment_roles, + available_keys, + available_segment_keys, + ) for option in contract.reads_one_of ) @@ -350,9 +375,7 @@ def _spec_satisfied_by_key( available_segment_keys: set[str], ) -> bool: """Whether one I/O alternative's literal keys exist in their declared scopes.""" - return not (set(spec.data_keys) - available_keys) and not ( - set(spec.segment_data_keys) - available_segment_keys - ) + return not (set(spec.data_keys) - available_keys) and not (set(spec.segment_data_keys) - available_segment_keys) def _reads_satisfied_by_key( @@ -377,8 +400,7 @@ def _reads_satisfied_by_key( if not _spec_satisfied_by_key(contract.reads, available_keys, available_segment_keys): return False return not contract.reads_one_of or any( - _spec_satisfied_by_key(option, available_keys, available_segment_keys) - for option in contract.reads_one_of + _spec_satisfied_by_key(option, available_keys, available_segment_keys) for option in contract.reads_one_of ) @@ -496,7 +518,13 @@ def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[Pipe for now -- the expansion is new, and it earns the right to block only once it has been shown not to false-positive on pipelines known to work. """ - role_satisfied = _reads_satisfied_by_role(contract, walk.available, walk.segment_available) + role_satisfied = _reads_satisfied_by_role( + contract, + walk.available, + walk.segment_available, + walk.available_keys, + walk.segment_available_keys, + ) key_satisfied = _reads_satisfied_by_key(contract, walk.available_keys, walk.segment_available_keys) if role_satisfied or key_satisfied: if walk.past_composite: @@ -563,9 +591,7 @@ def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[Pipe ] needed = _roles_for_keys(contract, contract.reads.data_keys) | { - role - for option in contract.reads_one_of - for role in _roles_for_keys(contract, option.data_keys) + role for option in contract.reads_one_of for role in _roles_for_keys(contract, option.data_keys) } removed_hit = (needed & walk.removed_roles) - walk.available if removed_hit: diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index 8d0c394374..1793b4d21a 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -112,6 +112,37 @@ def test_conditional_roles_are_discoverable_but_not_planner_guaranteed() -> None assert "potential_metrics" not in report.produced_keys +def test_unknown_role_selector_requires_its_exact_conditional_key() -> None: + producer = _ConfiguredContractStage( + StageContract( + conditional_writes=[ + ConditionalWrite( + writes=IOSpec(data_keys=["row_score"]), + condition="valid runtime data causes score assignment", + ) + ], + key_roles={"row_score": "score"}, + ) + ) + selector = PreserveByValueStage("row_score", 1.0, "le") + + conditional_only = validate_pipeline( + [producer, selector], + initial_roles=set(), + initial_keys=set(), + ) + assert not conditional_only.ok + assert any(issue.code == "unsatisfied_reads" and issue.stage_index == 1 for issue in conditional_only.issues) + + seeded = validate_pipeline( + [selector], + initial_roles=set(), + initial_keys={"row_score"}, + ) + assert seeded.ok + assert seeded.keys_ok + + def test_nested_input_requires_explicit_segment_seeds_and_accepts_remapped_key() -> None: nested_reader = _ConfiguredContractStage( StageContract( From fcaa6ca466c9f44fd8badc9083bcedf85fad269a Mon Sep 17 00:00:00 2001 From: Shubham Bhawsar Date: Wed, 16 Sep 2026 16:42:26 +0000 Subject: [PATCH 11/25] fix(audio-agent): preserve filtering residency consistency Signed-off-by: Shubham Bhawsar --- nemo_curator/stages/audio/_agent/_planning.py | 15 +- .../stages/audio/_agent/_residency.py | 203 +++++++++++++++++- .../test_agent_foundation_regressions.py | 75 +++++++ 3 files changed, 282 insertions(+), 11 deletions(-) diff --git a/nemo_curator/stages/audio/_agent/_planning.py b/nemo_curator/stages/audio/_agent/_planning.py index bd73413500..61377ae95c 100644 --- a/nemo_curator/stages/audio/_agent/_planning.py +++ b/nemo_curator/stages/audio/_agent/_planning.py @@ -748,17 +748,20 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: ): walk.available.discard(role) walk.removed_roles.add(role) - if "tensor" in contract.writes.produces: + tensor_writes = written | segment_written + has_possible_tensor_write = "tensor" in contract.writes.produces + for conditional in contract.conditional_writes: + if "tensor" in conditional.writes.produces: + has_possible_tensor_write = True + tensor_writes.update(conditional.writes.data_keys) + tensor_writes.update(conditional.writes.segment_data_keys) + if has_possible_tensor_write: # The stage's OWN key_roles first, global names only as fallback. A custom # ``waveform_key`` still declares its role in the contract, but the global lookup # returned "unknown", so residency tracked ``_UNNAMED_TENSOR`` instead of the real # carrier -- and a downstream stage dropping that carrier still looked resident, # raising a spurious ``tensor_into_sink`` on a recipe that had cleaned up correctly. - carriers = { - key - for key in written | segment_written - if contract.key_roles.get(key, role_for_value(key)) == _TENSOR_ROLE - } + carriers = {key for key in tensor_writes if contract.key_roles.get(key, role_for_value(key)) == _TENSOR_ROLE} walk.tensor_keys |= carriers or {_UNNAMED_TENSOR} if contract.gates.sanitizes_output: walk.tensor_keys.clear() diff --git a/nemo_curator/stages/audio/_agent/_residency.py b/nemo_curator/stages/audio/_agent/_residency.py index b4158907fe..a3c8e3d95e 100644 --- a/nemo_curator/stages/audio/_agent/_residency.py +++ b/nemo_curator/stages/audio/_agent/_residency.py @@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, Any, Literal import soundfile as sf +import torch from nemo_curator.stages.audio._agent._agent_ready import AudioForm, ConditionalWrite, IOSpec from nemo_curator.stages.audio.common import ensure_waveform_2d, load_audio_file @@ -29,6 +30,13 @@ from collections.abc import Callable InputResidency = Literal["file", "waveform", "auto"] +FileAudioHydration = Literal["never", "always", "auto_partial"] +_FILE_AUDIO_HYDRATION_POLICIES = {"never", "always", "auto_partial"} + +_SUPPORTED_TORCH_PCM_DTYPES = { + torch.int16: 32768.0, + torch.int32: 2147483648.0, +} def validate_input_residency(residency: str, *, stage_name: str) -> None: @@ -38,6 +46,65 @@ def validate_input_residency(residency: str, *, stage_name: str) -> None: raise ValueError(msg) +def validate_audio_key_configuration( + stage_name: str, + *, + input_keys: dict[str, str], + output_keys: dict[str, str], +) -> None: + """Reject empty configurable keys and destructive output collisions.""" + for field_name, key in {**input_keys, **output_keys}.items(): + if not isinstance(key, str) or not key.strip(): + msg = f"[{stage_name}] '{field_name}' must be a non-empty string" + raise ValueError(msg) + + output_values = list(output_keys.values()) + if len(output_values) != len(set(output_values)): + duplicates = sorted({key for key in output_values if output_values.count(key) > 1}) + msg = f"[{stage_name}] Output keys must be distinct; duplicate values: {duplicates}" + raise ValueError(msg) + + collisions = sorted(set(input_keys.values()) & set(output_values)) + if collisions: + msg = f"[{stage_name}] Output keys must not collide with audio input keys: {collisions}" + raise ValueError(msg) + + +def normalize_audio_waveform( + waveform: Any, # noqa: ANN401 - accepts torch tensors and array-like resident audio + *, + stage_name: str, + mono: bool, +) -> torch.Tensor: + """Convert supported resident audio to channel-first float32.""" + try: + tensor = waveform.detach() if torch.is_tensor(waveform) else torch.as_tensor(waveform) + except Exception as ex: + msg = f"[{stage_name}] Resident waveform must be convertible to a torch tensor" + raise TypeError(msg) from ex + + if tensor.ndim not in {1, 2}: + msg = f"[{stage_name}] Resident waveform must be 1-D or 2-D (channels, samples), got {tensor.ndim}-D" + raise ValueError(msg) + + if tensor.is_floating_point(): + tensor = tensor.to(dtype=torch.float32) + elif tensor.dtype in _SUPPORTED_TORCH_PCM_DTYPES: + scale = _SUPPORTED_TORCH_PCM_DTYPES[tensor.dtype] + tensor = tensor.to(dtype=torch.float32) / scale + else: + msg = ( + f"[{stage_name}] Unsupported resident waveform dtype {tensor.dtype}; " + "expected a floating dtype or signed PCM int16/int32" + ) + raise TypeError(msg) + + tensor = ensure_waveform_2d(tensor) + if mono and tensor.shape[0] > 1: + tensor = tensor.mean(dim=0, keepdim=True) + return tensor + + def accepts_for_residency(residency: str) -> list[AudioForm]: """Audio forms an instance actually consumes, given its ``input_residency``. @@ -59,6 +126,7 @@ def residency_read_specs( audio_filepath_key: str, waveform_key: str = "waveform", sample_rate_key: str = "sample_rate", + infer_sample_rate_from_file: bool = False, ) -> list[IOSpec]: """The residency-filtered audio read options for a stage's ``reads_one_of``. @@ -72,6 +140,8 @@ def residency_read_specs( specs: list[IOSpec] = [] if "waveform" in forms: specs.append(IOSpec(data_keys=[waveform_key, sample_rate_key], accepts=["waveform"])) + if input_residency == "auto" and infer_sample_rate_from_file: + specs.append(IOSpec(data_keys=[waveform_key, audio_filepath_key], accepts=["waveform"])) if "file" in forms: specs.append(IOSpec(data_keys=[audio_filepath_key], accepts=["file"])) return specs @@ -86,6 +156,7 @@ def scoped_audio_io_specs( # noqa: PLR0913 sample_rate_key: str, segments_key: str, output_keys: list[str], + infer_sample_rate_from_file: bool = False, ) -> tuple[IOSpec, list[IOSpec], IOSpec]: """Build mode-accurate reads/writes for task-or-nested audio stages. @@ -102,6 +173,7 @@ def scoped_audio_io_specs( # noqa: PLR0913 audio_filepath_key=audio_filepath_key, waveform_key=waveform_key, sample_rate_key=sample_rate_key, + infer_sample_rate_from_file=infer_sample_rate_from_file, ) segment_reads = [ IOSpec( @@ -169,7 +241,94 @@ def scoped_audio_conditional_writes( return conditional -def resolve_audio( # noqa: PLR0913 (complexity accepted: keyword-only residency/key knobs mirror the stage fields) +def scoped_file_audio_hydration_writes( # noqa: PLR0913 + input_residency: InputResidency, + *, + hydration_policy: FileAudioHydration, + mode: Literal["task", "segments", "auto"], + waveform_key: str, + sample_rate_key: str, + segments_key: str, + infer_sample_rate_from_file: bool = False, +) -> list[ConditionalWrite]: + """Describe possible resident writes caused by successful file resolution. + + File hydration is advisory-only because path selection and decode success + are runtime facts. ``auto`` exposes both possible scopes without making + either a mechanical guarantee. The Band-only header-completion branch is + represented separately because it writes only the sample rate and keeps + the resident waveform unchanged. + """ + if hydration_policy not in _FILE_AUDIO_HYDRATION_POLICIES: + msg = f"Unknown file audio hydration policy: {hydration_policy!r}" + raise ValueError(msg) + + pair_hydration_possible = (hydration_policy == "always" and input_residency != "waveform") or ( + hydration_policy == "auto_partial" and input_residency == "auto" + ) + header_completion_possible = input_residency == "auto" and infer_sample_rate_from_file + if not pair_hydration_possible and not header_completion_possible: + return [] + + conditional: list[ConditionalWrite] = [] + scopes = ("task", "segments") if mode == "auto" else (mode,) + for scope in scopes: + if scope == "task": + branch = ( + "task mode is configured" + if mode == "task" + else f"'{segments_key}' is absent, so the task-level branch runs" + ) + pair_writes = IOSpec( + data_keys=[waveform_key, sample_rate_key], + produces=["tensor"], + ) + rate_writes = IOSpec(data_keys=[sample_rate_key]) + else: + branch = ( + "segments mode is configured and an individual segment exists" + if mode == "segments" + else f"'{segments_key}' is present, so the per-segment branch runs, and an individual segment exists" + ) + pair_writes = IOSpec( + segment_data_keys=[waveform_key, sample_rate_key], + produces=["tensor"], + ) + rate_writes = IOSpec(segment_data_keys=[sample_rate_key]) + + if pair_hydration_possible: + partial_condition = ( + f" exactly one of resident '{waveform_key}' and '{sample_rate_key}' is present, " + if hydration_policy == "auto_partial" + else " " + ) + conditional.append( + ConditionalWrite( + writes=pair_writes, + condition=( + f"{branch};{partial_condition}file audio is selected and decoded successfully; " + f"'{waveform_key}' and '{sample_rate_key}' are assigned together " + "from the decoded file audio" + ), + value_origin="stage_generated", + ) + ) + if header_completion_possible: + conditional.append( + ConditionalWrite( + writes=rate_writes, + condition=( + f"{branch}; a resident '{waveform_key}' is present without " + f"'{sample_rate_key}', the configured file exists, and its header is read successfully; " + f"only '{sample_rate_key}' is assigned from the file header and the resident waveform is retained" + ), + value_origin="stage_generated", + ) + ) + return conditional + + +def resolve_audio( # noqa: C901, PLR0913 (complexity accepted: policy branches and keyword-only stage knobs) item: dict[str, Any], *, residency: InputResidency = "auto", @@ -178,21 +337,44 @@ def resolve_audio( # noqa: PLR0913 (complexity accepted: keyword-only residency sample_rate_key: str = "sample_rate", mono: bool = True, loader: Callable[..., tuple[Any, int]] | None = None, + infer_sample_rate_from_file: bool = False, + file_audio_hydration: FileAudioHydration = "never", ) -> 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. + configured path key. When ``infer_sample_rate_from_file`` is enabled, + ``auto`` may read only the file header to complete a resident waveform + that is missing its sample rate. + + ``file_audio_hydration="always"`` replaces both resident audio fields after + any selected file load. ``"auto_partial"`` does so only when ``auto`` falls + back with exactly one resident field present. ``"never"`` is the default, + preserving file-only and explicit-file consumers. Every update happens only + after the loader succeeds, so failures cannot leave a partial pair. ``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. """ + if file_audio_hydration not in _FILE_AUDIO_HYDRATION_POLICIES: + msg = f"Unknown file audio hydration policy: {file_audio_hydration!r}" + raise ValueError(msg) + 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 != "file" and waveform is not None: + if sample_rate is not None: + return ensure_waveform_2d(waveform), int(sample_rate) + if residency == "auto" and infer_sample_rate_from_file: + path = item.get(audio_filepath_key) + if path: + expanded = os.path.expanduser(str(path)) + if os.path.exists(expanded): + sample_rate = int(sf.info(expanded).samplerate) + item[sample_rate_key] = sample_rate + return ensure_waveform_2d(waveform), sample_rate if residency == "waveform": return None @@ -201,7 +383,18 @@ def resolve_audio( # noqa: PLR0913 (complexity accepted: keyword-only residency if path: expanded = os.path.expanduser(str(path)) if os.path.exists(expanded): - return (loader or load_audio_file)(expanded, mono=mono) + loaded_waveform, loaded_sample_rate = (loader or load_audio_file)(expanded, mono=mono) + has_partial_pair = (waveform is None) != (sample_rate is None) + if file_audio_hydration == "always" or ( + file_audio_hydration == "auto_partial" and residency == "auto" and has_partial_pair + ): + item.update( + { + waveform_key: loaded_waveform, + sample_rate_key: loaded_sample_rate, + } + ) + return loaded_waveform, loaded_sample_rate return None diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index 1793b4d21a..d8fdd2894e 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -112,6 +112,30 @@ def test_conditional_roles_are_discoverable_but_not_planner_guaranteed() -> None assert "potential_metrics" not in report.produced_keys +def test_conditional_tensor_write_is_not_guaranteed_but_still_blocks_json_sink(tmp_path: Path) -> None: + producer = _ConfiguredContractStage( + StageContract( + conditional_writes=[ + ConditionalWrite( + writes=IOSpec(data_keys=["resident_audio"], produces=["tensor"]), + condition="file decoding succeeds and resident audio is assigned", + ) + ], + key_roles={"resident_audio": "waveform"}, + ) + ) + writer = ManifestWriterStage(output_path=str(tmp_path / "out.jsonl")) + + report = validate_pipeline( + [producer, writer], + initial_roles=set(), + initial_keys=set(), + ) + + assert "resident_audio" not in report.produced_keys + assert any(issue.code == "tensor_into_sink" and issue.severity == "error" for issue in report.issues) + + def test_unknown_role_selector_requires_its_exact_conditional_key() -> None: producer = _ConfiguredContractStage( StageContract( @@ -180,6 +204,57 @@ def test_input_residency_validator_rejects_unknown_mode() -> None: validate_input_residency("wavefrom", stage_name="Fixture") +def test_file_audio_hydration_policies_are_opt_in_and_atomic(tmp_path: Path) -> None: + path = tmp_path / "audio.wav" + path.touch() + loaded = torch.arange(8, dtype=torch.float32).unsqueeze(0) + + def loader(_path: str, *, mono: bool) -> tuple[torch.Tensor, int]: + assert mono + return loaded, 16000 + + untouched = {"audio_filepath": str(path)} + resolve_audio(untouched, residency="file", loader=loader) + assert set(untouched) == {"audio_filepath"} + + always = {"audio_filepath": str(path)} + resolve_audio(always, residency="file", loader=loader, file_audio_hydration="always") + assert always["waveform"] is loaded + assert always["sample_rate"] == 16000 + + partial = {"audio_filepath": str(path), "sample_rate": 8000} + resolve_audio(partial, residency="auto", loader=loader, file_audio_hydration="auto_partial") + assert partial["waveform"] is loaded + assert partial["sample_rate"] == 16000 + + ordinary_auto = {"audio_filepath": str(path)} + resolve_audio(ordinary_auto, residency="auto", loader=loader, file_audio_hydration="auto_partial") + assert set(ordinary_auto) == {"audio_filepath"} + + +def test_failed_file_hydration_preserves_the_existing_pair(tmp_path: Path) -> None: + path = tmp_path / "audio.wav" + path.touch() + stale = torch.ones(1, 4) + item = {"audio_filepath": str(path), "waveform": stale, "sample_rate": 8000} + + def failing_loader(_path: str, *, mono: bool) -> tuple[torch.Tensor, int]: + assert mono + msg = "decode failed" + raise OSError(msg) + + with pytest.raises(OSError, match="decode failed"): + resolve_audio( + item, + residency="file", + loader=failing_loader, + file_audio_hydration="always", + ) + + assert item["waveform"] is stale + assert item["sample_rate"] == 8000 + + def test_resolve_audio_path_auto_prefers_complete_resident_audio(tmp_path: Path) -> None: """Auto residency must not silently choose a stale file over a complete waveform.""" file_path = tmp_path / "one_second.wav" From 04a9adadf3ea5306fd9c61c4a28a827ef40c1e3f Mon Sep 17 00:00:00 2001 From: Shubham Bhawsar Date: Wed, 16 Sep 2026 17:41:50 +0000 Subject: [PATCH 12/25] fix(audio-agent): scope residency, literals, and positional compatibility Signed-off-by: Shubham Bhawsar --- .../stages/audio/_agent/_conformance.py | 61 ++++++-- nemo_curator/stages/audio/_agent/_planning.py | 127 ++++++++++++---- .../audio/preprocessing/concatenation.py | 37 +++-- .../audio/preprocessing/mono_conversion.py | 18 ++- .../test_agent_foundation_regressions.py | 138 ++++++++++++++++++ .../audio/preprocessing/test_concatenation.py | 48 ++++++ .../preprocessing/test_mono_conversion.py | 25 ++++ 7 files changed, 395 insertions(+), 59 deletions(-) diff --git a/nemo_curator/stages/audio/_agent/_conformance.py b/nemo_curator/stages/audio/_agent/_conformance.py index 58403bc771..576480ee35 100644 --- a/nemo_curator/stages/audio/_agent/_conformance.py +++ b/nemo_curator/stages/audio/_agent/_conformance.py @@ -323,12 +323,50 @@ def _check_gpu_gate(stage: Any, c: StageContract, name: str) -> None: # noqa: A raise AssertionError(msg) +def _assert_reads_satisfiable( + c: StageContract, + name: str, + available_keys: Iterable[str], + available_segment_keys: Iterable[str] | None, +) -> None: + """Assert reads are satisfiable using the SAME rule the planner applies. + + Reuses the planner's read check verbatim so conformance and planning cannot drift: known + roles satisfy a read, an unknown-role read needs its exact literal key, and ``reads_one_of`` + is evaluated per alternative per scope. + """ + from nemo_curator.stages.audio._agent._planning import _reads_satisfied_by_role + from nemo_curator.stages.audio._agent._roles import role_for_value + + avail_keys = set(available_keys) + avail_segment_keys = set(available_segment_keys or ()) + + def _roles_of_keys(keys: set[str]) -> set[str]: + # Both this stage's own key_roles and the shared literal table, minus the permissive + # "unknown" so an unknown-role read is decided by its literal key, not waved through. + resolved = {c.key_roles.get(k, "unknown") for k in keys} | {role_for_value(k) for k in keys} + return resolved - {"unknown"} + + assert _reads_satisfied_by_role( + c, + _roles_of_keys(avail_keys), + _roles_of_keys(avail_segment_keys), + avail_keys, + avail_segment_keys, + ), ( + f"{name}: reads {c.reads.data_keys}/{[s.data_keys for s in c.reads_one_of]} " + f"not satisfied by available keys {sorted(avail_keys)}" + + (f" / segment keys {sorted(avail_segment_keys)}" if avail_segment_keys else "") + ) + + def assert_agent_ready( # noqa: C901, PLR0912, PLR0913 (complexity accepted: one linear checklist of independent conformance checks) stage: Any, # noqa: ANN401 fixture_factory: Callable[[], Any] | None = None, *, expected_cardinality: str | None = None, available_keys: Iterable[str] | None = None, + available_segment_keys: Iterable[str] | None = None, segments_key: str | None = None, ignore_new_keys: Iterable[str] = (), run: bool = True, @@ -344,8 +382,13 @@ def assert_agent_ready( # noqa: C901, PLR0912, PLR0913 (complexity accepted: on stage: A constructed stage instance. fixture_factory: Returns a fresh input task (or batch) each call. expected_cardinality: If given, assert the contract declares it. - available_keys: Upstream-available key values; asserts reads are - satisfiable by role. + available_keys: Upstream-available top-level key values; asserts reads are + satisfiable using the SAME rule the planner applies — known roles or + exact literal keys for unknown-role reads, with ``reads_one_of`` + evaluated per alternative per scope (an unknown-role read is no + longer treated as always-satisfied). + available_segment_keys: Upstream-available nested (segment-scope) key + values, for the nested equivalent of ``available_keys``. segments_key: Resolved segments key, for checking segment-level writes. ignore_new_keys: Extra top-level keys allowed in output (framework bookkeeping) beyond declared writes. @@ -368,15 +411,7 @@ def assert_agent_ready( # noqa: C901, PLR0912, PLR0913 (complexity accepted: on f"{name}: cardinality {c.cardinality!r} != expected {expected_cardinality!r}" ) if available_keys is not None: - avail_roles = {c.key_roles.get(k, "unknown") for k in available_keys} - # also resolve via literal table for keys not in this stage's key_roles - from nemo_curator.stages.audio._agent._roles import role_for_value - - avail_roles |= {role_for_value(k) for k in available_keys} - assert reads_satisfied_by_role(c, avail_roles), ( - f"{name}: reads {c.reads.data_keys}/{[s.data_keys for s in c.reads_one_of]} " - f"not satisfied by available roles {avail_roles}" - ) + _assert_reads_satisfiable(c, name, available_keys, available_segment_keys) if not run or fixture_factory is None: return c @@ -425,9 +460,7 @@ def assert_agent_ready( # noqa: C901, PLR0912, PLR0913 (complexity accepted: on assert key in out_data, f"{name}: declared write {key!r} missing from task.data" for key in c.removes_keys: assert key not in out_data, f"{name}: declared removes_keys {key!r} but it is still present in task.data" - conditional_keys = { - key for conditional in c.conditional_writes for key in conditional.writes.data_keys - } + conditional_keys = {key for conditional in c.conditional_writes for key in conditional.writes.data_keys} declared = set(c.writes.data_keys) | conditional_keys | set(ignore_new_keys) | input_keys undeclared = set(out_data) - declared assert not undeclared, f"{name}: undeclared new top-level keys {sorted(undeclared)} (add to writes.data_keys)" diff --git a/nemo_curator/stages/audio/_agent/_planning.py b/nemo_curator/stages/audio/_agent/_planning.py index 61377ae95c..248af78ed6 100644 --- a/nemo_curator/stages/audio/_agent/_planning.py +++ b/nemo_curator/stages/audio/_agent/_planning.py @@ -460,6 +460,28 @@ def _describes_itself(stage: Any) -> bool: # noqa: ANN401 - any child stage return True +def _missing_literal_role_keys( + contract: StageContract, + spec: Any, # noqa: ANN401 - IOSpec, kept loose to avoid a runtime-only import + available_keys: set[str], + available_segment_keys: set[str], +) -> set[str]: + """Role-bearing key VALUES in one spec whose exact value is absent from its scope. + + ``unknown``/internal bookkeeping keys are excluded — a separate value-identity check for + those is tracked in the backlog. Empty means every role-bearing key in ``spec`` is present, + i.e. this alternative is literally complete. + """ + missing: set[str] = set() + for key, scope_keys in ( + *[(k, available_keys) for k in spec.data_keys], + *[(k, available_segment_keys) for k in spec.segment_data_keys], + ): + if contract.key_roles.get(key, "unknown") != "unknown" and key not in scope_keys: + missing.add(key) + return missing + + def _dangling_read_keys( contract: StageContract, available_keys: set[str], @@ -468,27 +490,28 @@ def _dangling_read_keys( """Read key VALUES whose role is known but whose exact value was not produced upstream nor seeded — the renamed-producer dangle the role check misses. - Covers primary ``reads`` plus a ``reads_one_of`` that offers a *single* - alternative: one option is not a choice, so its keys are as mandatory as a - primary read (this is how a residency-derived contract expresses - ``input_residency="file"``). A genuine multi-way ``reads_one_of`` is skipped — - the stage may legitimately take the other branch. Role-bearing keys only - (``unknown``/internal bookkeeping keys are excluded — a separate - value-identity check for those is tracked in the backlog). + Covers primary ``reads`` (always mandatory) plus ``reads_one_of``: + + * a *single* alternative is not a choice, so its keys are as mandatory as a primary read + (this is how a residency-derived contract expresses ``input_residency="file"``); + * a genuine *multi-way* ``reads_one_of`` is satisfied by ANY one complete literal + alternative. When NO alternative is literally complete — every branch has a role-bearing + key that was only satisfied by a renamed/role-level producer — the read dangles even + though a role check passed, so the union of each branch's missing keys is reported. + + Role-bearing keys only (``unknown``/internal bookkeeping keys are excluded). """ - reads = [(key, available_keys) for key in contract.reads.data_keys] - reads += [(key, available_segment_keys) for key in contract.reads.segment_data_keys] - if len(contract.reads_one_of) == 1: - only = contract.reads_one_of[0] - reads += [(key, available_keys) for key in only.data_keys] - reads += [(key, available_segment_keys) for key in only.segment_data_keys] - dangling: set[str] = set() - for k, scope_keys in reads: - role = contract.key_roles.get(k, "unknown") - if role == "unknown": - continue - if k not in scope_keys: - dangling.add(k) + dangling = _missing_literal_role_keys(contract, contract.reads, available_keys, available_segment_keys) + options = contract.reads_one_of + if len(options) == 1: + dangling |= _missing_literal_role_keys(contract, options[0], available_keys, available_segment_keys) + elif len(options) > 1: + per_option = [ + _missing_literal_role_keys(contract, option, available_keys, available_segment_keys) for option in options + ] + # Clean iff at least one alternative is literally complete (no missing role-bearing key). + if all(missing for missing in per_option): + dangling |= set().union(*per_option) return dangling @@ -500,7 +523,11 @@ class _Walk: available_keys: set[str] # literal top-level key VALUES produced so far segment_available: set[str] = field(default_factory=set) # nested-item roles produced so far segment_available_keys: set[str] = field(default_factory=set) # literal nested-item key VALUES + # Tensor residency is tracked per scope: a top-level carrier and a nested (segment) carrier + # are distinct to a serializer, so a stage dropping one must not be credited with clearing + # the other. ``tensor_keys`` holds top-level carriers; ``segment_tensor_keys`` nested ones. tensor_keys: set[str] = field(default_factory=set) + segment_tensor_keys: set[str] = field(default_factory=set) removed_roles: set[str] = field(default_factory=set) key_producer: dict[str, str] = field(default_factory=dict) segment_key_producer: dict[str, str] = field(default_factory=dict) @@ -738,7 +765,8 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: walk.segment_available_keys |= segment_written for rk in contract.removes_keys: walk.available_keys.discard(rk) - # Dropping the carrier ends the tensor residency as surely as sanitizing does. + # ``removes_keys`` names TOP-LEVEL task keys, so dropping the carrier ends only the + # top-level tensor residency; a nested (segment) carrier of the same name survives. walk.tensor_keys.discard(rk) role = contract.key_roles.get(rk, role_for_value(rk)) if ( @@ -748,23 +776,47 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: ): walk.available.discard(role) walk.removed_roles.add(role) - tensor_writes = written | segment_written + _advance_tensor_residency(walk, contract, written, segment_written) + + +def _advance_tensor_residency( + walk: _Walk, + contract: StageContract, + written: set[str], + segment_written: set[str], +) -> None: + """Fold one stage's tensor writes/sanitization into the per-scope residency sets.""" + task_tensor_writes = set(written) + segment_tensor_writes = set(segment_written) has_possible_tensor_write = "tensor" in contract.writes.produces for conditional in contract.conditional_writes: if "tensor" in conditional.writes.produces: has_possible_tensor_write = True - tensor_writes.update(conditional.writes.data_keys) - tensor_writes.update(conditional.writes.segment_data_keys) + task_tensor_writes.update(conditional.writes.data_keys) + segment_tensor_writes.update(conditional.writes.segment_data_keys) if has_possible_tensor_write: # The stage's OWN key_roles first, global names only as fallback. A custom # ``waveform_key`` still declares its role in the contract, but the global lookup # returned "unknown", so residency tracked ``_UNNAMED_TENSOR`` instead of the real # carrier -- and a downstream stage dropping that carrier still looked resident, # raising a spurious ``tensor_into_sink`` on a recipe that had cleaned up correctly. - carriers = {key for key in tensor_writes if contract.key_roles.get(key, role_for_value(key)) == _TENSOR_ROLE} - walk.tensor_keys |= carriers or {_UNNAMED_TENSOR} + task_carriers = { + key for key in task_tensor_writes if contract.key_roles.get(key, role_for_value(key)) == _TENSOR_ROLE + } + segment_carriers = { + key for key in segment_tensor_writes if contract.key_roles.get(key, role_for_value(key)) == _TENSOR_ROLE + } + if task_carriers or segment_carriers: + walk.tensor_keys |= task_carriers + walk.segment_tensor_keys |= segment_carriers + else: + # ``produces=["tensor"]`` but no waveform-roled key names the carrier: keep the + # pre-split behaviour of tracking it as a top-level unnamed tensor only a + # sanitizer can clear. + walk.tensor_keys |= {_UNNAMED_TENSOR} if contract.gates.sanitizes_output: walk.tensor_keys.clear() + walk.segment_tensor_keys.clear() def _seed_walk( # noqa: PLR0913 -- top-level and nested seeds describe one input task @@ -775,6 +827,7 @@ def _seed_walk( # noqa: PLR0913 -- top-level and nested seeds describe one inpu *, initial_segment_roles: set[str] | None = None, initial_segment_keys: set[str] | None = None, + initial_segment_tensor_keys: set[str] | None = None, ) -> _Walk: """The state the first stage is handed: what the input task already carries.""" if initial_keys is not None: @@ -803,13 +856,18 @@ def _seed_walk( # noqa: PLR0913 -- top-level and nested seeds describe one inpu if initial_tensor_keys is not None: seed_tensors = set(initial_tensor_keys) else: - seed_tensors = {k for k in seed_keys | segment_seed_keys if role_for_value(k) == _TENSOR_ROLE} + seed_tensors = {k for k in seed_keys if role_for_value(k) == _TENSOR_ROLE} + if initial_segment_tensor_keys is not None: + seed_segment_tensors = set(initial_segment_tensor_keys) + else: + seed_segment_tensors = {k for k in segment_seed_keys if role_for_value(k) == _TENSOR_ROLE} return _Walk( available=set(initial_roles) if initial_roles is not None else set(_DEFAULT_INITIAL_ROLES), available_keys=seed_keys, segment_available=set(initial_segment_roles or ()), segment_available_keys=segment_seed_keys, tensor_keys=seed_tensors, + segment_tensor_keys=seed_segment_tensors, task_type=initial_task_type, ) @@ -822,6 +880,7 @@ def validate_pipeline( # noqa: PLR0913 -- keyword-only seeds of one input task, initial_segment_roles: set[str] | None = None, initial_segment_keys: set[str] | None = None, initial_tensor_keys: set[str] | None = None, + initial_segment_tensor_keys: set[str] | None = None, initial_task_type: str | None = None, available_gpus: float | None = None, ) -> PipelineReport: @@ -850,6 +909,10 @@ def validate_pipeline( # noqa: PLR0913 -- keyword-only seeds of one input task, (e.g. ``audio_tensor``). Pass an empty set when a schema contains a waveform-named column but the input values are known not to be resident tensors. + initial_segment_tensor_keys: The nested (segment-scope) equivalent of + ``initial_tensor_keys``. ``None`` -- the default -- infers resident + segment tensors from the segment seed keys by role. Pass this when a + segment carries a tensor under a name whose role cannot be inferred. initial_task_type: Class name of the task the first stage will be handed (e.g. ``"EmptyTask"`` for a pipeline that starts at a source, ``"AudioTask"`` for a suffix resumed from a manifest). ``None`` -- the default -- leaves the @@ -869,6 +932,7 @@ def validate_pipeline( # noqa: PLR0913 -- keyword-only seeds of one input task, initial_task_type, initial_segment_roles=initial_segment_roles, initial_segment_keys=initial_segment_keys, + initial_segment_tensor_keys=initial_segment_tensor_keys, ) expansion = expand_composites(stages) leaves = expansion.by_recipe_index() @@ -977,7 +1041,14 @@ def validate_pipeline( # noqa: PLR0913 -- keyword-only seeds of one input task, issues.extend(_task_type_issue(walk, site, contract)) # Serialization / GPU gates reason about the environment rather than about roles, so # they run for every concrete stage even downstream of a composite nobody could expand. - issues.extend(_gate_issues(site, contract, available_gpus, tensor_resident=bool(walk.tensor_keys))) + issues.extend( + _gate_issues( + site, + contract, + available_gpus, + tensor_resident=bool(walk.tensor_keys or walk.segment_tensor_keys), + ) + ) _advance(walk, contract, site.name) # An undeclared output type is not "unchanged": it is unknown, and carrying the # previous stage's type past it would judge the next stage against a task that is diff --git a/nemo_curator/stages/audio/preprocessing/concatenation.py b/nemo_curator/stages/audio/preprocessing/concatenation.py index 719913393c..5b50f14f0f 100755 --- a/nemo_curator/stages/audio/preprocessing/concatenation.py +++ b/nemo_curator/stages/audio/preprocessing/concatenation.py @@ -33,7 +33,7 @@ """ import os -from dataclasses import dataclass, field +from dataclasses import KW_ONLY, dataclass, field from typing import Any import torch @@ -93,7 +93,17 @@ class SegmentConcatenationStage(AgentReady, ProcessingStage[AudioTask, AudioTask audio_filepath_key: Key set to the written combined-audio path (write_to_disk). """ + # Legacy positional slots (pre-agent order preserved): silence_duration_sec, name, + # batch_size, resources. Everything the agent work added is keyword-only (below the + # KW_ONLY sentinel), so a legacy positional call like ``SegmentConcatenationStage(0.5)`` + # keeps its meaning. silence_duration_sec: float = 0.5 + + name: str = "SegmentConcatenation" + batch_size: int = 1 + resources: Resources = field(default_factory=lambda: Resources(cpus=1.0)) + + _: KW_ONLY segments_key: str = "segments" waveform_key: str = "waveform" sample_rate_key: str = "sample_rate" @@ -106,10 +116,6 @@ class SegmentConcatenationStage(AgentReady, ProcessingStage[AudioTask, AudioTask write_to_disk: bool = False output_dir: str | None = None - name: str = "SegmentConcatenation" - batch_size: int = 1 - resources: Resources = field(default_factory=lambda: Resources(cpus=1.0)) - def __post_init__(self): super().__init__() if not (self.keep_waveform_in_task or self.write_to_disk): @@ -141,7 +147,14 @@ def describe(self) -> StageContract: writes.append(self.audio_filepath_key) produces.append("disk") return StageContract( - reads=IOSpec(data_keys=[self.segments_key]), + # Each child segment must carry its own waveform + sample_rate: ``process`` reads + # ``seg[waveform_key]``/``seg[sample_rate_key]`` per item, so declaring only the + # top-level container let a nested seed/producer that omitted them validate clean + # and then drop every segment at runtime. + reads=IOSpec( + data_keys=[self.segments_key], + segment_data_keys=[self.waveform_key, self.sample_rate_key], + ), writes=IOSpec(data_keys=writes, produces=produces), metadata_writes=["segment_mappings"], cardinality="N:1", @@ -292,15 +305,17 @@ def _concatenate( current_pos_ms -= silence_duration_ms total_duration_sec = current_pos_ms / 1000.0 - output_data = { - self.original_file_key: original_file, - self.num_segments_key: len(mappings), - self.total_duration_sec_key: total_duration_sec, - } + # Legacy key insertion order: waveform, sample_rate, original_file, num_segments, + # total_duration_sec, then any on-disk path. Values are unchanged; only the order the + # keys land in the output dict (and thus a serialized row) is restored. + output_data: dict[str, Any] = {} # Output residency: keep the combined waveform in-task (default) and/or persist it. if self.keep_waveform_in_task: output_data[self.waveform_key] = combined output_data[self.sample_rate_key] = sample_rate + output_data[self.original_file_key] = original_file + output_data[self.num_segments_key] = len(mappings) + output_data[self.total_duration_sec_key] = total_duration_sec if self.write_to_disk: output_data[self.audio_filepath_key] = self._write_wav(combined, sample_rate, original_file) diff --git a/nemo_curator/stages/audio/preprocessing/mono_conversion.py b/nemo_curator/stages/audio/preprocessing/mono_conversion.py index 354ae79227..da09294846 100755 --- a/nemo_curator/stages/audio/preprocessing/mono_conversion.py +++ b/nemo_curator/stages/audio/preprocessing/mono_conversion.py @@ -27,7 +27,7 @@ """ import os -from dataclasses import dataclass, field +from dataclasses import KW_ONLY, dataclass, field import torch from loguru import logger @@ -83,8 +83,19 @@ class MonoConversionStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): output_dir: Directory for the written WAV files (default: system temp dir). """ + # Legacy positional slots (pre-agent order preserved): output_sample_rate, + # audio_filepath_key, strict_sample_rate, name, batch_size, resources. Everything the + # agent work added is keyword-only (below the KW_ONLY sentinel), so a legacy positional + # call like ``MonoConversionStage(48000, "audio_filepath", True)`` keeps its meaning. output_sample_rate: int = 48000 audio_filepath_key: str = "audio_filepath" + strict_sample_rate: bool = True + + name: str = "MonoConversion" + batch_size: int = 1 + resources: Resources = field(default_factory=lambda: Resources(cpus=1.0)) + + _: KW_ONLY waveform_key: str = "waveform" sample_rate_key: str = "sample_rate" is_mono_key: str = "is_mono" @@ -92,7 +103,6 @@ class MonoConversionStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): 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 # "file", not "auto": this stage only ever read ``audio_filepath`` before the agent work # (``load_audio_file(audio_filepath, mono=False)``), unlike sigmos/utmos/band/vad, whose # own resolvers already preferred a resident waveform and so default to "auto" honestly. @@ -103,10 +113,6 @@ class MonoConversionStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): update_audio_filepath: bool = False output_dir: str | None = None - name: str = "MonoConversion" - batch_size: int = 1 - resources: Resources = field(default_factory=lambda: Resources(cpus=1.0)) - def __post_init__(self): super().__init__() reject_sinkless_conversion( diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index d8fdd2894e..8d3b129f62 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -663,3 +663,141 @@ def test_a_null_waveform_does_not_authenticate_a_stale_sample_rate(tmp_path: Pat data={"audio_filepath": str(path), "sample_rate": 16000, "waveform": torch.zeros(1, 16000)}, ) assert stage._observed_rate(resident) == 16000 + + +def test_concatenation_reads_require_nested_segment_audio_keys() -> None: + """SegmentConcatenation reads waveform+sample_rate from EACH child, not just the container.""" + concat = SegmentConcatenationStage() + + # Only the top-level segments container is present; the per-child audio the runtime reads + # is missing, so the stage must not validate clean. + top_only = validate_pipeline( + [concat], + initial_roles={"segments"}, + initial_keys={"segments"}, + ) + assert not top_only.ok + assert any(i.code == "unsatisfied_reads" and i.stage_index == 0 for i in top_only.issues) + + # Seeding the nested waveform/sample_rate the child carries makes it compose. + with_nested = validate_pipeline( + [concat], + initial_roles={"segments"}, + initial_keys={"segments"}, + initial_segment_roles={"waveform", "sample_rate"}, + initial_segment_keys={"waveform", "sample_rate"}, + ) + assert with_nested.ok + assert with_nested.keys_ok + + # Remapped child keys chain by role: pointed at the names the seed carries, it stays clean. + remapped = SegmentConcatenationStage(waveform_key="seg_wav", sample_rate_key="seg_sr") + report = validate_pipeline( + [remapped], + initial_roles={"segments"}, + initial_keys={"segments"}, + initial_segment_roles={"waveform", "sample_rate"}, + initial_segment_keys={"seg_wav", "seg_sr"}, + ) + assert report.ok + assert report.keys_ok + + +@pytest.mark.parametrize("sink", [ManifestWriterStage, ManifestCheckpointStage]) +def test_same_name_nested_tensor_survives_top_level_drop_into_sink(sink: type, tmp_path: Path) -> None: + """A disk-only conversion drops the TOP-LEVEL waveform; a same-named nested one still blocks a sink.""" + mono = MonoConversionStage( + output_sample_rate=16000, + input_residency="waveform", + keep_waveform_in_task=False, + write_to_disk=True, + output_dir=str(tmp_path / "out"), + ) + assert set(build_contract(mono).removes_keys) == {"waveform", "sample_rate"} + writer = sink(output_path=str(tmp_path / "out.jsonl")) + + # Task-level AND segment-level waveforms share the key name "waveform". The conversion + # removes only the task-level carrier; the nested one reaches the JSON sink. + report = validate_pipeline( + [mono, writer], + initial_roles={"waveform", "sample_rate", "segments"}, + initial_keys={"waveform", "sample_rate", "segments"}, + initial_segment_roles={"waveform", "sample_rate"}, + initial_segment_keys={"waveform", "sample_rate"}, + ) + assert not report.ok + assert any(i.code == "tensor_into_sink" and i.severity == "error" for i in report.issues) + + # Single-scope behavior is unchanged: with no nested carrier, dropping the top-level one + # clears residency and the sink is clean (no false positive from the scope split). + clean = validate_pipeline( + [mono, writer], + initial_roles={"waveform", "sample_rate"}, + initial_keys={"waveform", "sample_rate"}, + ) + assert not any(i.code == "tensor_into_sink" for i in clean.issues) + + +def test_multi_alternative_read_dangles_when_no_literal_branch_is_complete() -> None: + """An auto consumer whose role is met only by a renamed producer key has no complete branch.""" + renamed_role_only = _ConfiguredContractStage( + StageContract( + writes=IOSpec(data_keys=["resampled_audio_filepath"]), + key_roles={"resampled_audio_filepath": "audio_filepath"}, + ) + ) + auto_consumer = MonoConversionStage(input_residency="auto") + + dangling = validate_pipeline( + [renamed_role_only, auto_consumer], + initial_roles=set(), + initial_keys=set(), + ) + # Role-level composability holds, but no reads_one_of branch is literally complete. + assert dangling.ok + assert not dangling.keys_ok + assert any(i.code == "dangling_key" and i.stage_index == 1 for i in dangling.issues) + + # A complete FILE branch (literal audio_filepath) stays clean. + file_producer = _ConfiguredContractStage( + StageContract( + writes=IOSpec(data_keys=["audio_filepath"]), + key_roles={"audio_filepath": "audio_filepath"}, + ) + ) + clean_file = validate_pipeline( + [file_producer, MonoConversionStage(input_residency="auto")], + initial_roles=set(), + initial_keys=set(), + ) + assert clean_file.ok + assert clean_file.keys_ok + + # A complete WAVEFORM-PAIR branch stays clean too. + waveform_producer = _ConfiguredContractStage( + StageContract( + writes=IOSpec(data_keys=["waveform", "sample_rate"], produces=["tensor"]), + key_roles={"waveform": "waveform", "sample_rate": "sample_rate"}, + ) + ) + clean_waveform = validate_pipeline( + [waveform_producer, MonoConversionStage(input_residency="auto")], + initial_roles=set(), + initial_keys=set(), + ) + assert clean_waveform.ok + assert clean_waveform.keys_ok + + +def test_conformance_requires_exact_literal_key_for_unknown_role_read(tmp_path: Path) -> None: + """A custom (unknown-role) read must be satisfied by its exact key, not waved through.""" + wav = tmp_path / "a.wav" + sf.write(wav, torch.zeros(48000).numpy(), 48000) + selector = PreserveByValueStage("mos", 3.0, "ge") # input_value_key='mos' -> unknown role + + # 'mos' is not among the available keys, so the unknown-role read is unsatisfied. + with pytest.raises(AssertionError, match="not satisfied"): + assert_agent_ready(selector, available_keys={"audio_filepath"}, run=False) + + # Present exactly, it passes. + assert_agent_ready(selector, available_keys={"mos"}, run=False) diff --git a/tests/stages/audio/preprocessing/test_concatenation.py b/tests/stages/audio/preprocessing/test_concatenation.py index 5cafc96f25..f3e59e0e29 100644 --- a/tests/stages/audio/preprocessing/test_concatenation.py +++ b/tests/stages/audio/preprocessing/test_concatenation.py @@ -148,3 +148,51 @@ def test_requires_output_dir_when_write_to_disk(self) -> None: def test_requires_at_least_one_output_sink(self) -> None: with pytest.raises(ValueError, match="keep_waveform_in_task or write_to_disk"): SegmentConcatenationStage(keep_waveform_in_task=False) + + # --- positional compatibility (KW_ONLY sentinel) --- + + def test_legacy_positional_call_matches_legacy_order(self) -> None: + """Pre-agent positionals were (silence_duration_sec, name, batch_size, resources).""" + from nemo_curator.stages.resources import Resources + + stage = SegmentConcatenationStage(1.5) + assert stage.silence_duration_sec == 1.5 + # Agent-added fields stayed keyword-only, so they keep their defaults. + assert stage.segments_key == "segments" + assert stage.waveform_key == "waveform" + + # name/batch_size/resources remain the legacy positional slots after silence_duration_sec. + stage2 = SegmentConcatenationStage(1.5, "Concat", 4, Resources(cpus=1.0)) + assert (stage2.name, stage2.batch_size) == ("Concat", 4) + + # A 5th positional would be a keyword-only agent field -> TypeError. + with pytest.raises(TypeError): + SegmentConcatenationStage(1.5, "Concat", 4, Resources(cpus=1.0), "segs") + + # --- output key insertion order (legacy layout) --- + + def test_output_key_order_matches_legacy(self) -> None: + segments = [ + _make_segment_dict(duration_ms=1000, segment_num=0), + _make_segment_dict(duration_ms=1000, segment_num=1), + ] + result = SegmentConcatenationStage().process(_make_nested_task(segments)) + assert list(result.data.keys()) == [ + "waveform", + "sample_rate", + "original_file", + "num_segments", + "total_duration_sec", + ] + + def test_output_key_order_places_disk_path_last(self, tmp_path) -> None: # noqa: ANN001 + stage = SegmentConcatenationStage(write_to_disk=True, output_dir=str(tmp_path / "c")) + result = stage.process(_make_nested_task([_make_segment_dict(duration_ms=1000)])) + assert list(result.data.keys()) == [ + "waveform", + "sample_rate", + "original_file", + "num_segments", + "total_duration_sec", + "audio_filepath", + ] diff --git a/tests/stages/audio/preprocessing/test_mono_conversion.py b/tests/stages/audio/preprocessing/test_mono_conversion.py index 790d2a945e..4b01e0a8c2 100644 --- a/tests/stages/audio/preprocessing/test_mono_conversion.py +++ b/tests/stages/audio/preprocessing/test_mono_conversion.py @@ -164,3 +164,28 @@ def test_disk_only_output_omits_the_waveform_key(self, tmp_path: Path) -> None: assert "agent_mono_path" in result.data, "disk-path key must be present when write_to_disk=True" assert "agent_waveform" not in result.data, "tensor must be omitted when keep_waveform_in_task=False" + + +class TestMonoConversionPositionalCompatibility: + def test_legacy_positional_call_keeps_strict_sample_rate_third(self) -> None: + """Pre-agent order was (output_sample_rate, audio_filepath_key, strict_sample_rate).""" + stage = MonoConversionStage(16000, "path_col", False) + assert stage.output_sample_rate == 16000 + assert stage.audio_filepath_key == "path_col" + assert stage.strict_sample_rate is False + # The agent-added fields stayed keyword-only, so they keep their defaults here. + assert stage.input_residency == "file" + assert stage.waveform_key == "waveform" + + def test_agent_added_fields_are_keyword_only(self) -> None: + """A former agent field cannot be reached positionally past the legacy slots.""" + import pytest + + from nemo_curator.stages.resources import Resources + + # The six legacy positional slots still accept positionals in their original order. + stage = MonoConversionStage(16000, "path_col", True, "Custom", 2, Resources(cpus=1.0)) + assert (stage.name, stage.batch_size) == ("Custom", 2) + # A 7th positional would be a keyword-only agent field -> TypeError. + with pytest.raises(TypeError): + MonoConversionStage(16000, "path_col", True, "Custom", 2, Resources(cpus=1.0), "wf") From e3d457e72b5c8f144cc71c69129e2ec63aecf068 Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Wed, 16 Sep 2026 17:54:55 +0000 Subject: [PATCH 13/25] fix(audio-agent): support optional contract reads Signed-off-by: shbhawsar --- nemo_curator/stages/audio/AGENT_READY.md | 5 +++++ nemo_curator/stages/audio/_agent/_agent_ready.py | 5 +++++ .../stages/audio/_agent/_agent_registry.py | 2 +- nemo_curator/stages/audio/_agent/_conformance.py | 15 +++++++++++++-- .../_agent/test_agent_foundation_regressions.py | 11 +++++++++++ 5 files changed, 35 insertions(+), 3 deletions(-) diff --git a/nemo_curator/stages/audio/AGENT_READY.md b/nemo_curator/stages/audio/AGENT_READY.md index e79d0eee4c..febc4e0560 100644 --- a/nemo_curator/stages/audio/AGENT_READY.md +++ b/nemo_curator/stages/audio/AGENT_READY.md @@ -192,6 +192,11 @@ find yourself wanting to, declare `False` and raise it instead. You never put `params` in `describe()`. +Use `optional_reads=IOSpec(...)` for keys a stage consults when present but can +correctly run without. Optional reads are visible to discovery and semantic +review, but they never make validation reject a fallback path. Do not put a +fallback key in `reads` or legacy `inputs()` merely to advertise it. + ## What is OPTIONAL — set only if it's obvious Declared via one class attribute, `AGENT_STATIC = StaticHints(...)`, or on the contract: diff --git a/nemo_curator/stages/audio/_agent/_agent_ready.py b/nemo_curator/stages/audio/_agent/_agent_ready.py index e05838e081..8b999853c8 100644 --- a/nemo_curator/stages/audio/_agent/_agent_ready.py +++ b/nemo_curator/stages/audio/_agent/_agent_ready.py @@ -246,6 +246,10 @@ class StageContract: # compatibility; this never changes stage execution or the legacy # mechanical interpretation of ``writes``. conditional_writes: list[ConditionalWrite] = field(default_factory=list) + # Keys the stage can consult when present but does not require. Appended for + # positional compatibility and kept outside ``reads``/``reads_one_of`` so + # planning never blocks a valid fallback path on their absence. + optional_reads: IOSpec = field(default_factory=IOSpec) def to_dict(self) -> dict[str, Any]: """Return a JSON-safe dict of this contract (``json.dumps`` never raises).""" @@ -254,6 +258,7 @@ def to_dict(self) -> dict[str, Any]: "reads": asdict(self.reads), "writes": asdict(self.writes), "reads_one_of": [asdict(spec) for spec in self.reads_one_of], + "optional_reads": asdict(self.optional_reads), "metadata_reads": list(self.metadata_reads), "metadata_writes": list(self.metadata_writes), "cardinality": self.cardinality, diff --git a/nemo_curator/stages/audio/_agent/_agent_registry.py b/nemo_curator/stages/audio/_agent/_agent_registry.py index fe379b72dd..3b8465f126 100644 --- a/nemo_curator/stages/audio/_agent/_agent_registry.py +++ b/nemo_curator/stages/audio/_agent/_agent_registry.py @@ -331,7 +331,7 @@ def _first_doc_line(cls: type) -> str | None: def _contract_referenced_keys(contract: StageContract) -> set[str]: keys: set[str] = set() - for spec in [contract.reads, contract.writes, *contract.reads_one_of]: + for spec in [contract.reads, contract.writes, contract.optional_reads, *contract.reads_one_of]: keys.update(spec.data_keys) keys.update(spec.segment_data_keys) keys.update(contract.metadata_reads) diff --git a/nemo_curator/stages/audio/_agent/_conformance.py b/nemo_curator/stages/audio/_agent/_conformance.py index 576480ee35..a74cf074e8 100644 --- a/nemo_curator/stages/audio/_agent/_conformance.py +++ b/nemo_curator/stages/audio/_agent/_conformance.py @@ -119,14 +119,19 @@ def _check_shape(c: StageContract, name: str) -> None: # noqa: C901 # absent from writes — but it must still resolve to a semantic role). # Catches synthetic labels like the former 'speakers' that name nothing. contract_keys: set[str] = set() - for spec in [c.reads, c.writes, *c.reads_one_of]: + for spec in [c.reads, c.writes, c.optional_reads, *c.reads_one_of]: contract_keys.update(spec.data_keys) contract_keys.update(spec.segment_data_keys) assert c.iteration_key in contract_keys or c.iteration_key in c.key_roles, ( f"{name}: iteration_key {c.iteration_key!r} is neither a contract read/write " f"key nor a role-resolvable key value — it names nothing an agent can find" ) - for spec, label in [(c.reads, "reads"), (c.writes, "writes"), *[(s, "reads_one_of") for s in c.reads_one_of]]: + for spec, label in [ + (c.reads, "reads"), + (c.writes, "writes"), + (c.optional_reads, "optional_reads"), + *[(s, "reads_one_of") for s in c.reads_one_of], + ]: for a in spec.accepts: assert a in _VALID_ACCEPTS, f"{name}: {label}.accepts has invalid form {a!r}" for p in spec.produces: @@ -159,6 +164,12 @@ def _check_shape(c: StageContract, name: str) -> None: # noqa: C901 # no duplicate keys within a single spec list for spec, label in [(c.reads, "reads"), (c.writes, "writes")]: assert len(spec.data_keys) == len(set(spec.data_keys)), f"{name}: duplicate {label}.data_keys" + assert len(c.optional_reads.data_keys) == len(set(c.optional_reads.data_keys)), ( + f"{name}: duplicate optional_reads.data_keys" + ) + assert len(c.optional_reads.segment_data_keys) == len(set(c.optional_reads.segment_data_keys)), ( + f"{name}: duplicate optional_reads.segment_data_keys" + ) def _check_roles(stage_or_cls: Any, c: StageContract, name: str) -> None: # noqa: ANN401 diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index 8d3b129f62..9a9458f9ef 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -85,6 +85,17 @@ def test_stage_params_respects_field_level_agent_exclusion() -> None: assert [param.name for param in stage_params(_AgentParamMetadataFixture)] == ["visible"] +def test_optional_reads_are_visible_without_blocking_fallback_paths() -> None: + contract = StageContract( + reads=IOSpec(data_keys=["text"]), + optional_reads=IOSpec(data_keys=["speaker_id"]), + ) + stage = _ConfiguredContractStage(contract) + + assert build_contract(stage).to_dict()["optional_reads"]["data_keys"] == ["speaker_id"] + assert validate_pipeline([stage], initial_keys={"text"}).ok + + def test_conditional_roles_are_discoverable_but_not_planner_guaranteed() -> None: producer_contract = StageContract( conditional_writes=[ From b26e5d38be6713c4834abed6308cfa89ffa17c89 Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Wed, 16 Sep 2026 20:15:48 +0000 Subject: [PATCH 14/25] fix(audio): preserve resident waveform samples Signed-off-by: shbhawsar --- .../stages/audio/_agent/_residency.py | 2 +- .../test_agent_foundation_regressions.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/nemo_curator/stages/audio/_agent/_residency.py b/nemo_curator/stages/audio/_agent/_residency.py index a3c8e3d95e..9951caaee2 100644 --- a/nemo_curator/stages/audio/_agent/_residency.py +++ b/nemo_curator/stages/audio/_agent/_residency.py @@ -542,7 +542,7 @@ def resolve_audio_path( # noqa: C901, PLR0913 (keyword-only residency/key knobs fd, tmp = tempfile.mkstemp(suffix=".wav", dir=temp_dir) os.close(fd) try: - sf.write(tmp, _as_soundfile_array(waveform), int(sample_rate)) + sf.write(tmp, _as_soundfile_array(waveform), int(sample_rate), subtype="FLOAT") except BaseException: with contextlib.suppress(OSError): os.remove(tmp) diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index 9a9458f9ef..1c6ad22e07 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -21,6 +21,7 @@ from types import SimpleNamespace from typing import TYPE_CHECKING +import numpy as np import pandas as pd import pytest import soundfile as sf @@ -293,6 +294,24 @@ def test_resolve_audio_path_auto_prefers_complete_resident_audio(tmp_path: Path) assert not os.path.exists(resolved) +def test_resolve_audio_path_preserves_float_waveform_samples(tmp_path: Path) -> None: + waveform = np.array([[1e-5, -1e-5, 1.25, -1.25]], dtype=np.float32) + temporary_paths: list[str] = [] + + resolved = resolve_audio_path( + {"waveform": waveform, "sample_rate": 16000}, + residency="waveform", + temp_dir=str(tmp_path), + register_temp=temporary_paths, + ) + + observed, sample_rate = sf.read(resolved, dtype="float32", always_2d=True) + assert sample_rate == 16000 + assert sf.info(resolved).subtype == "FLOAT" + np.testing.assert_array_equal(observed[:, 0], waveform[0]) + cleanup_temp_files(temporary_paths) + + def test_stable_audio_names_include_layout_and_written_short_stereo_shape(tmp_path: Path) -> None: """Different channel layouts with identical samples need distinct artifacts.""" output_dir = str(tmp_path) From 38aeea0ceeddaa5ac5686f4216221606a291d031 Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Wed, 16 Sep 2026 21:38:35 +0000 Subject: [PATCH 15/25] fix(audio-agent): sync shared contract support Signed-off-by: shbhawsar --- nemo_curator/stages/audio/AGENT_READY.md | 5 +++++ nemo_curator/stages/audio/_agent/_agent_registry.py | 1 + .../stages/audio/preprocessing/concatenation.py | 1 + .../_agent/test_agent_foundation_regressions.py | 12 ++++++++++++ 4 files changed, 19 insertions(+) diff --git a/nemo_curator/stages/audio/AGENT_READY.md b/nemo_curator/stages/audio/AGENT_READY.md index febc4e0560..87e90af715 100644 --- a/nemo_curator/stages/audio/AGENT_READY.md +++ b/nemo_curator/stages/audio/AGENT_READY.md @@ -197,6 +197,11 @@ correctly run without. Optional reads are visible to discovery and semantic review, but they never make validation reject a fallback path. Do not put a fallback key in `reads` or legacy `inputs()` merely to advertise it. +When compatibility requires retaining a constructor default that runtime validation rejects, +declare the field with `metadata={"agent_required": True}`. Discovery then requires an explicit +value without changing the Python constructor or its default. Use this only for values that the +stage cannot run without; ordinary defaults remain optional. + ## What is OPTIONAL — set only if it's obvious Declared via one class attribute, `AGENT_STATIC = StaticHints(...)`, or on the contract: diff --git a/nemo_curator/stages/audio/_agent/_agent_registry.py b/nemo_curator/stages/audio/_agent/_agent_registry.py index 3b8465f126..5de655eb58 100644 --- a/nemo_curator/stages/audio/_agent/_agent_registry.py +++ b/nemo_curator/stages/audio/_agent/_agent_registry.py @@ -204,6 +204,7 @@ def _dataclass_params(cls: type, descriptions: dict[str, str]) -> list[ParamSpec default, required = _call_factory(f.default_factory) else: default, required = None, True + required = required or bool(f.metadata.get("agent_required", False)) hint = _resolve_hint(f.type, globalns) params.append( ParamSpec( diff --git a/nemo_curator/stages/audio/preprocessing/concatenation.py b/nemo_curator/stages/audio/preprocessing/concatenation.py index 5b50f14f0f..5570ba5be0 100755 --- a/nemo_curator/stages/audio/preprocessing/concatenation.py +++ b/nemo_curator/stages/audio/preprocessing/concatenation.py @@ -167,6 +167,7 @@ def describe(self) -> StageContract: gates=Gates( writes_to_disk=self.write_to_disk, output_path_params=["output_dir"], + sanitizes_output=not self.keep_waveform_in_task, # The ``N`` this stage collapses is the segments of ONE row's own file, so no # other file's audio reaches the combined waveform -- the ``N:1`` cardinality # counts tasks, not the origins of the values. ``write_to_disk`` does not change diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index 1c6ad22e07..11abdea767 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -71,6 +71,11 @@ class _AgentParamMetadataFixture: runtime_only: object | None = field(default=None, metadata={"agent_param": False}) +@dataclass +class _AgentRequiredMetadataFixture: + required_for_agent: str = field(default="", metadata={"agent_required": True}) + + class _ConfiguredContractStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): def __init__(self, contract: StageContract) -> None: self.contract = contract @@ -86,6 +91,13 @@ def test_stage_params_respects_field_level_agent_exclusion() -> None: assert [param.name for param in stage_params(_AgentParamMetadataFixture)] == ["visible"] +def test_stage_params_can_require_a_runtime_default_for_agent_configuration() -> None: + param = stage_params(_AgentRequiredMetadataFixture)[0] + + assert param.default == "" + assert param.required is True + + def test_optional_reads_are_visible_without_blocking_fallback_paths() -> None: contract = StageContract( reads=IOSpec(data_keys=["text"]), From 3a2ba2f166738d416383556f5e8965efda4543b7 Mon Sep 17 00:00:00 2001 From: Shubham Bhawsar Date: Thu, 17 Sep 2026 08:16:08 +0000 Subject: [PATCH 16/25] fix(audio): plan conditional writes honestly; restart-safe ManifestWriterStage - Planner: a read met only by a conditional write is a conditional_read WARNING (never a guaranteed output; nothing-writes-it stays unsatisfied_reads). ConditionalWrite gains requires_keys so unreachable branches neither seed tensor residency nor count as possible; auto_partial hydration declares its two halves with it. Fixes UTMOSFilterStage()/SIGMOSFilterStage() -> ManifestWriterStage being refused with tensor_into_sink on a plain manifest, and metric -> selector / ASR -> Join -> Merge chains being un-plannable. - ManifestWriterStage truncates in setup_on_node() (once per node, before any worker) instead of per-actor setup(), so a Ray Data/Xenna worker restart no longer erases committed rows. Docstring corrected (setup() is per worker). Signed-off-by: Shubham Bhawsar --- nemo_curator/stages/audio/AGENT_READY.md | 13 ++ .../stages/audio/_agent/_agent_ready.py | 20 ++- nemo_curator/stages/audio/_agent/_planning.py | 139 +++++++++++++++++- .../stages/audio/_agent/_residency.py | 29 +++- nemo_curator/stages/audio/common.py | 28 ++-- .../test_agent_foundation_regressions.py | 129 +++++++++++++++- tests/stages/audio/test_common.py | 31 +++- 7 files changed, 354 insertions(+), 35 deletions(-) diff --git a/nemo_curator/stages/audio/AGENT_READY.md b/nemo_curator/stages/audio/AGENT_READY.md index 87e90af715..1b2b746b65 100644 --- a/nemo_curator/stages/audio/AGENT_READY.md +++ b/nemo_curator/stages/audio/AGENT_READY.md @@ -110,6 +110,19 @@ The `upstream_same_key` origin is important for allowlist/rebuild stages: it keeps the original producer's meaning visible instead of falsely presenting the copier as a new metric producer. +The planner does read `conditional_writes`, in two bounded ways. A downstream +read that is met only by a conditional key is reported as a `conditional_read` +*warning* (the pipeline composes, `report.ok` stays true, but the key is never +added to `produced_keys`); a read nothing even possibly writes remains an +`unsatisfied_reads` error. And a conditional write with `produces=["tensor"]` +seeds tensor residency for the JSON-sink gate. Set `requires_keys` on a +`ConditionalWrite` when its branch can only run if some literal key already +exists upstream (in the write's own scope): the planner then ignores the branch +-- for both purposes -- on inputs that cannot reach it. File hydration that only +*replaces* an incomplete resident waveform/sample-rate pair is the canonical +case: without the hint, `UTMOSFilterStage() -> ManifestWriterStage` on a plain +file manifest would be refused for a tensor the runtime never introduces. + Use `metadata_writes` for a conditional `task._metadata` output. Unconditional metadata inputs/outputs remain declared through `StageContract.metadata_reads`/`metadata_writes`; semantic review traces all diff --git a/nemo_curator/stages/audio/_agent/_agent_ready.py b/nemo_curator/stages/audio/_agent/_agent_ready.py index 8b999853c8..1f333edc78 100644 --- a/nemo_curator/stages/audio/_agent/_agent_ready.py +++ b/nemo_curator/stages/audio/_agent/_agent_ready.py @@ -121,9 +121,19 @@ class ConditionalWrite: task-data/audio-form inputs. This metadata is additive. Mechanical planners continue to use - :attr:`StageContract.writes`; ``conditional_writes`` supplies the host - critic with possibility/provenance evidence and may also describe - conditional pass-through keys omitted from the legacy ``writes`` superset. + :attr:`StageContract.writes` for GUARANTEED keys; ``conditional_writes`` + supplies the host critic with possibility/provenance evidence, lets the + planner credit a downstream read as *possibly* satisfied (reported as a + ``conditional_read`` warning rather than an ``unsatisfied_reads`` error), and + may also describe conditional pass-through keys omitted from the legacy + ``writes`` superset. + + ``requires_keys`` is the one executable part of the condition: literal + task-data keys (in the same scope as ``writes``) that must already be present + upstream for this branch to be reachable at all. Empty means "always + possible". The planner uses it to avoid crediting -- or fearing, for tensor + writes -- a branch that cannot fire on the seeded input, e.g. file hydration + that only replaces an already-resident waveform/sample-rate pair. """ writes: IOSpec = field(default_factory=IOSpec) @@ -131,6 +141,7 @@ class ConditionalWrite: value_origin: WriteValueOrigin = "stage_generated" # Appended for positional compatibility with the original three fields. metadata_writes: list[str] = field(default_factory=list) + requires_keys: list[str] = field(default_factory=list) @dataclass(frozen=True) @@ -284,6 +295,9 @@ def to_dict(self) -> dict[str, Any]: "condition": item.condition, "value_origin": item.value_origin, "metadata_writes": list(item.metadata_writes), + # Only emitted when set, so contracts without reachability hints + # serialize exactly as before. + **({"requires_keys": list(item.requires_keys)} if item.requires_keys else {}), } for item in self.conditional_writes ], diff --git a/nemo_curator/stages/audio/_agent/_planning.py b/nemo_curator/stages/audio/_agent/_planning.py index 248af78ed6..a08d8dd170 100644 --- a/nemo_curator/stages/audio/_agent/_planning.py +++ b/nemo_curator/stages/audio/_agent/_planning.py @@ -528,6 +528,13 @@ class _Walk: # the other. ``tensor_keys`` holds top-level carriers; ``segment_tensor_keys`` nested ones. tensor_keys: set[str] = field(default_factory=set) segment_tensor_keys: set[str] = field(default_factory=set) + # Keys/roles a stage MAY have written (``conditional_writes`` whose ``requires_keys`` were + # reachable). Never folded into the guaranteed sets above: a read met only from here is a + # ``conditional_read`` warning, not a satisfied read and not an ``unsatisfied_reads`` error. + possible_keys: set[str] = field(default_factory=set) + possible_segment_keys: set[str] = field(default_factory=set) + possible_roles: set[str] = field(default_factory=set) + possible_segment_roles: set[str] = field(default_factory=set) removed_roles: set[str] = field(default_factory=set) key_producer: dict[str, str] = field(default_factory=dict) segment_key_producer: dict[str, str] = field(default_factory=dict) @@ -535,7 +542,7 @@ class _Walk: task_type: str | None = None # task type the previous stage produces; None == not known -def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[PipelineIssue]: +def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[PipelineIssue]: # noqa: PLR0911 - one return per verdict """Whether this stage's reads are met, and how loudly to say so if not. Severity is graded by how sure we are, because a wrong hard error is worse than a wrong @@ -553,6 +560,10 @@ def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[Pipe walk.segment_available_keys, ) key_satisfied = _reads_satisfied_by_key(contract, walk.available_keys, walk.segment_available_keys) + if not (role_satisfied or key_satisfied) and not walk.past_composite and site.composite is None: + conditional = _conditional_read_issue(walk, site, contract) + if conditional is not None: + return [conditional] if role_satisfied or key_satisfied: if walk.past_composite: return [] @@ -644,6 +655,73 @@ def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[Pipe ] +def _reads_possibly_satisfied(walk: _Walk, contract: StageContract) -> bool: + """Whether every read is met once keys upstream MAY write are counted as present. + + Possible keys are credited by LITERAL key only, never by role: a conditional pass-through of + some other key that happens to share the read's role is too weak a basis to compose on. + Guaranteed state keeps its normal role-or-literal tolerance. + """ + + def key_ok(key: str, keys: set[str], possible: set[str], roles: set[str]) -> bool: + if key in keys or key in possible: + return True + role = contract.key_roles.get(key, "unknown") + return role != "unknown" and role in roles + + def spec_ok(spec: Any) -> bool: # noqa: ANN401 - IOSpec + return all( + key_ok(key, walk.available_keys, walk.possible_keys, walk.available) for key in spec.data_keys + ) and all( + key_ok(key, walk.segment_available_keys, walk.possible_segment_keys, walk.segment_available) + for key in spec.segment_data_keys + ) + + if not spec_ok(contract.reads): + return False + return not contract.reads_one_of or any(spec_ok(option) for option in contract.reads_one_of) + + +def _conditional_read_issue(walk: _Walk, site: _Site, contract: StageContract) -> PipelineIssue | None: + """A warning when a read is met only by keys an upstream stage MAY write. + + Metric and hydration stages declare data-dependent outputs as ``conditional_writes`` so the + planner never advances them as guaranteed. Refusing every consumer of such a key outright + would make the ordinary ``ComputeWER -> PreserveByValue`` or ``ASR -> Join -> Merge`` chain + un-plannable, so a read satisfied by the union of guaranteed and possible keys is reported + as ``conditional_read`` (warning): the pipeline composes, but the consumer must tolerate the + key being absent on rows where the producing branch did not run. ``report.ok`` stays True; + callers wanting only guaranteed flow can check ``report.warnings`` for this code. + """ + if not _reads_possibly_satisfied(walk, contract): + return None + only_possible = sorted( + ( + {*contract.reads.data_keys, *(k for option in contract.reads_one_of for k in option.data_keys)} + & walk.possible_keys + ) + - walk.available_keys + ) + sorted( + ( + { + *contract.reads.segment_data_keys, + *(k for option in contract.reads_one_of for k in option.segment_data_keys), + } + & walk.possible_segment_keys + ) + - walk.segment_available_keys + ) + return PipelineIssue( + site.index, + site.name, + "warning", + "conditional_read", + f"reads key(s) {only_possible} that upstream stages write only conditionally " + f"(data-dependent branch); rows where that branch does not run will lack the key, so this " + f"stage must tolerate its absence. Guaranteed keys so far: {sorted(walk.available_keys)}", + ) + + def _declared_produces(stage: Any) -> str | None: # noqa: ANN401 - any recipe stage """The task type a stage says it produces, or None if it cannot say. @@ -728,6 +806,17 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: segment_produced = _roles_for_keys(contract, contract.writes.segment_data_keys) written = _write_key_values(contract) segment_written = _segment_write_key_values(contract) + # Judged against the INPUT state, before this stage's own possible writes are folded in: + # a branch must not be made reachable by the key it would itself write. + reachable_conditionals = _reachable_conditional_writes(walk, contract) + # Keys this stage drops on its main path. A conditional write of the same key (a branch that + # happens to keep it, or re-emits it with a different meaning such as a tar member name) must + # not resurrect it for planning: removal is the guarantee-level fact, the branch the exception. + blocked_keys = set(contract.removes_keys) + blocked_segment_keys: set[str] = set() + if not contract.preserves_upstream_keys: + blocked_keys |= walk.available_keys - written + blocked_segment_keys |= walk.segment_available_keys - segment_written if not contract.preserves_upstream_keys: # A stage that rebuilds the task rather than adding to it: whatever it does not write # is not downstream. Folding its writes into the inherited state would keep every @@ -743,6 +832,10 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: walk.available -= dropped_roles walk.segment_available_keys -= dropped_segment_keys walk.segment_available -= dropped_segment_roles + walk.possible_keys -= dropped_keys + walk.possible_segment_keys -= dropped_segment_keys + walk.possible_roles -= dropped_roles + walk.possible_segment_roles -= dropped_segment_roles walk.removed_roles |= dropped_roles for key in dropped_keys: walk.key_producer.pop(key, None) @@ -763,8 +856,18 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: walk.available_keys |= written walk.segment_key_producer.update(dict.fromkeys(segment_written, name)) walk.segment_available_keys |= segment_written + for reachable in reachable_conditionals: + # Possible, not guaranteed: kept in the ``possible_*`` sets so a downstream read met + # only from here surfaces as a ``conditional_read`` warning. + possible = set(reachable.writes.data_keys) - blocked_keys + possible_segment = set(reachable.writes.segment_data_keys) - blocked_segment_keys + walk.possible_keys |= possible + walk.possible_segment_keys |= possible_segment + walk.possible_roles |= _roles_for_keys(contract, possible) + walk.possible_segment_roles |= _roles_for_keys(contract, possible_segment) for rk in contract.removes_keys: walk.available_keys.discard(rk) + walk.possible_keys.discard(rk) # ``removes_keys`` names TOP-LEVEL task keys, so dropping the carrier ends only the # top-level tensor residency; a nested (segment) carrier of the same name survives. walk.tensor_keys.discard(rk) @@ -776,7 +879,34 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: ): walk.available.discard(role) walk.removed_roles.add(role) - _advance_tensor_residency(walk, contract, written, segment_written) + _advance_tensor_residency(walk, contract, written, segment_written, reachable_conditionals) + + +def _reachable_conditional_writes(walk: _Walk, contract: StageContract) -> list[Any]: + """The ``conditional_writes`` whose ``requires_keys`` the input so far can actually meet. + + ``requires_keys`` names literal keys, in the write's own scope, that must already exist for + the branch to run. A file-hydration branch that only REPLACES an incomplete resident pair + cannot fire on a plain manifest, so its tensor write must neither seed residency (a + spurious ``tensor_into_sink`` on ``UTMOSFilterStage() -> ManifestWriterStage``) nor be + credited as a possible output. A write with no ``requires_keys`` is always reachable. + Reachability is judged against guaranteed AND possible keys -- a possible key can enable a + possible branch -- which is the conservative direction for the tensor gate. + """ + task_keys = walk.available_keys | walk.possible_keys + segment_keys = walk.segment_available_keys | walk.possible_segment_keys + reachable = [] + for conditional in contract.conditional_writes: + required = set(getattr(conditional, "requires_keys", ()) or ()) + if not required: + reachable.append(conditional) + continue + scope_keys = ( + segment_keys if conditional.writes.segment_data_keys and not conditional.writes.data_keys else task_keys + ) + if required <= scope_keys: + reachable.append(conditional) + return reachable def _advance_tensor_residency( @@ -784,12 +914,15 @@ def _advance_tensor_residency( contract: StageContract, written: set[str], segment_written: set[str], + reachable_conditionals: list[Any] | None = None, ) -> None: """Fold one stage's tensor writes/sanitization into the per-scope residency sets.""" task_tensor_writes = set(written) segment_tensor_writes = set(segment_written) has_possible_tensor_write = "tensor" in contract.writes.produces - for conditional in contract.conditional_writes: + if reachable_conditionals is None: + reachable_conditionals = _reachable_conditional_writes(walk, contract) + for conditional in reachable_conditionals: if "tensor" in conditional.writes.produces: has_possible_tensor_write = True task_tensor_writes.update(conditional.writes.data_keys) diff --git a/nemo_curator/stages/audio/_agent/_residency.py b/nemo_curator/stages/audio/_agent/_residency.py index 9951caaee2..bdfdbcfa5d 100644 --- a/nemo_curator/stages/audio/_agent/_residency.py +++ b/nemo_curator/stages/audio/_agent/_residency.py @@ -296,17 +296,31 @@ def scoped_file_audio_hydration_writes( # noqa: PLR0913 ) rate_writes = IOSpec(segment_data_keys=[sample_rate_key]) - if pair_hydration_possible: - partial_condition = ( - f" exactly one of resident '{waveform_key}' and '{sample_rate_key}' is present, " - if hydration_policy == "auto_partial" - else " " - ) + if pair_hydration_possible and hydration_policy == "auto_partial": + # ``auto_partial`` only fires on an INCOMPLETE resident pair, so the branch is + # reachable only when exactly one of the two keys already exists upstream. Declare + # the two halves separately with ``requires_keys`` so a planner seeded with a plain + # file manifest (neither key) does not fear a tensor this stage cannot introduce. + for present_key, absent_key in ((waveform_key, sample_rate_key), (sample_rate_key, waveform_key)): + conditional.append( + ConditionalWrite( + writes=pair_writes, + condition=( + f"{branch}; resident '{present_key}' is present without '{absent_key}', " + f"file audio is selected and decoded successfully; " + f"'{waveform_key}' and '{sample_rate_key}' are assigned together " + "from the decoded file audio" + ), + value_origin="stage_generated", + requires_keys=[present_key], + ) + ) + elif pair_hydration_possible: conditional.append( ConditionalWrite( writes=pair_writes, condition=( - f"{branch};{partial_condition}file audio is selected and decoded successfully; " + f"{branch}; file audio is selected and decoded successfully; " f"'{waveform_key}' and '{sample_rate_key}' are assigned together " "from the decoded file audio" ), @@ -323,6 +337,7 @@ def scoped_file_audio_hydration_writes( # noqa: PLR0913 f"only '{sample_rate_key}' is assigned from the file header and the resident waveform is retained" ), value_origin="stage_generated", + requires_keys=[waveform_key], ) ) return conditional diff --git a/nemo_curator/stages/audio/common.py b/nemo_curator/stages/audio/common.py index d6c3562f4a..a7d62f857d 100644 --- a/nemo_curator/stages/audio/common.py +++ b/nemo_curator/stages/audio/common.py @@ -718,9 +718,7 @@ def _item_id(relative_path: str) -> str: # Reserve ``~`` for the underscore escape below. encoded.append("~~") elif char == "_" and ( - index in (0, last) - or component[index - 1] == "_" - or component[index + 1] == "_" + index in (0, last) or component[index - 1] == "_" or component[index + 1] == "_" ): # Encoded components must neither contain ``__`` nor touch a # separator with ``_``; otherwise two different component @@ -790,15 +788,21 @@ def process(self, _: EmptyTask) -> list[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) - so repeated pipeline runs produce a clean output. ``setup_on_node()`` - only creates the parent directory -- it never truncates, so multi-node - deployments do not erase each other's data. + The output file is truncated in ``setup_on_node()`` so repeated pipeline + runs produce a clean output. Every executor runs ``setup_on_node()`` on each + node BEFORE any worker starts processing, and does not run it again when a + worker actor is replaced -- whereas ``setup()`` runs per worker actor, so a + replacement actor after a crash or a Ray Data/Xenna worker restart would + re-run it and erase every row the previous actor had already committed. + ``setup()`` therefore only prepares the filesystem handle and never + truncates. .. note:: Because all nodes append to the same path, callers in multi-node setups should either use a shared filesystem or provide a - node-unique ``output_path``. + node-unique ``output_path``. On a shared filesystem every node's + ``setup_on_node()`` truncation completes before the single writer + worker (``num_workers() == 1``) appends its first row. Supports local and cloud paths via fsspec. @@ -828,13 +832,11 @@ def __post_init__(self) -> None: raise ValueError(msg) def setup(self, _worker_metadata: WorkerMetadata | None = None) -> None: - """Truncate the output file once on the driver before processing starts.""" + """Prepare the filesystem handle for this worker. Never truncates (see class docstring).""" self._fs, self._path = url_to_fs(self.output_path) parent_dir = "/".join(self._path.split("/")[:-1]) if parent_dir: self._fs.makedirs(parent_dir, exist_ok=True) - with self._fs.open(self._path, "w", encoding="utf-8"): - pass logger.info(f"ManifestWriterStage: writing to {self.output_path}") def setup_on_node( @@ -842,11 +844,13 @@ def setup_on_node( _node_info: NodeInfo | None = None, _worker_metadata: WorkerMetadata | None = None, ) -> None: - """Ensure parent directory exists on each node (no truncation).""" + """Create the parent directory and truncate the output once per run, before any worker writes.""" self._fs, self._path = url_to_fs(self.output_path) parent_dir = "/".join(self._path.split("/")[:-1]) if parent_dir: self._fs.makedirs(parent_dir, exist_ok=True) + with self._fs.open(self._path, "w", encoding="utf-8"): + pass def process(self, task: AudioTask) -> AudioTask: with self._fs.open(self._path, "a", encoding="utf-8") as f: diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index 11abdea767..13ae5116b0 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -131,9 +131,16 @@ def test_conditional_roles_are_discoverable_but_not_planner_guaranteed() -> None initial_keys=set(), ) - assert not report.ok - assert any(issue.code == "unsatisfied_reads" and issue.stage_index == 1 for issue in report.issues) - assert "potential_metrics" not in report.produced_keys + # Conditional outputs are discoverable and let the consumer compose, but only as a + # ``conditional_read`` warning: the key is never a guaranteed planner output. + assert report.ok + assert any(issue.code == "conditional_read" and issue.stage_index == 1 for issue in report.issues) + assert ( + "potential_metrics" + not in validate_pipeline( + [_ConfiguredContractStage(producer_contract)], initial_roles=set(), initial_keys=set() + ).produced_keys + ) def test_conditional_tensor_write_is_not_guaranteed_but_still_blocks_json_sink(tmp_path: Path) -> None: @@ -179,8 +186,8 @@ def test_unknown_role_selector_requires_its_exact_conditional_key() -> None: initial_roles=set(), initial_keys=set(), ) - assert not conditional_only.ok - assert any(issue.code == "unsatisfied_reads" and issue.stage_index == 1 for issue in conditional_only.issues) + assert conditional_only.ok + assert any(issue.code == "conditional_read" and issue.stage_index == 1 for issue in conditional_only.issues) seeded = validate_pipeline( [selector], @@ -843,3 +850,115 @@ def test_conformance_requires_exact_literal_key_for_unknown_role_read(tmp_path: # Present exactly, it passes. assert_agent_ready(selector, available_keys={"mos"}, run=False) + + +class _ConditionalScoreProducer(AgentReady): + """Writes ``score`` only on rows where ``source`` is non-null (a data-dependent branch).""" + + name = "ConditionalScoreProducer" + + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec(data_keys=["audio_filepath"]), + conditional_writes=[ + ConditionalWrite(writes=IOSpec(data_keys=["score"]), condition="'source' is non-null"), + ], + ) + + def process(self, task: object) -> object: + return task + + +class _ReachabilityGatedTensorProducer(AgentReady): + """Hydrates a waveform only when a ``sample_rate`` column already exists upstream.""" + + name = "ReachabilityGatedTensorProducer" + + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec(data_keys=["audio_filepath"]), + conditional_writes=[ + ConditionalWrite( + writes=IOSpec(data_keys=["waveform", "sample_rate"], produces=["tensor"]), + condition="a resident sample_rate without a waveform is completed from the file", + requires_keys=["sample_rate"], + ), + ], + ) + + def process(self, task: object) -> object: + return task + + +def test_a_read_met_only_by_a_conditional_write_is_a_warning_not_an_error() -> None: + """Conditional outputs stay non-guaranteed, but their consumers still compose.""" + report = validate_pipeline([_ConditionalScoreProducer(), PreserveByValueStage("score", 3.0, "ge")]) + + assert report.ok, report.summary() + producer_only = validate_pipeline([_ConditionalScoreProducer()]) + assert "score" not in producer_only.produced_keys, "a conditional write must not become a guaranteed key" + assert [issue.code for issue in report.issues] == ["conditional_read"] + assert "score" in report.issues[0].message + + # Nothing upstream even possibly writes the key: still a hard error. + missing = validate_pipeline([PreserveByValueStage("score", 3.0, "ge")]) + assert not missing.ok + assert any(issue.code == "unsatisfied_reads" for issue in missing.issues) + + +def test_shipped_metric_then_selector_chain_composes() -> None: + """The fleurs recipe shape: pairwise WER followed by a threshold on its (conditional) output.""" + from nemo_curator.stages.audio.metrics.wer import GetPairwiseWerStage + + report = validate_pipeline( + [GetPairwiseWerStage(), PreserveByValueStage("wer_pct", 25.0, "le")], + initial_keys={"audio_filepath", "text", "pred_text"}, + initial_roles={"audio_filepath", "text", "pred_text"}, + ) + assert report.ok, report.summary() + assert report.keys_ok + assert any(issue.code == "conditional_read" for issue in report.issues) + + +def test_conditional_tensor_writes_seed_residency_only_when_reachable(tmp_path: Path) -> None: + """``requires_keys`` decides whether a hydration branch can fire on the seeded input.""" + sink = ManifestWriterStage(output_path=str(tmp_path / "out.jsonl")) + + # Plain manifest: the branch needs ``sample_rate`` upstream, which nothing provides. + plain = validate_pipeline([_ReachabilityGatedTensorProducer(), sink]) + assert plain.ok, plain.summary() + + # A ``sample_rate`` column makes the branch reachable, so the sink is (correctly) refused. + with_rate = validate_pipeline( + [_ReachabilityGatedTensorProducer(), sink], + initial_keys={"audio_filepath", "sample_rate"}, + initial_roles={"audio_filepath", "sample_rate"}, + ) + assert not with_rate.ok + assert any(issue.code == "tensor_into_sink" for issue in with_rate.issues) + + # A stage must not make its own branch reachable through the key that branch would write. + self_enabling = validate_pipeline( + [_ReachabilityGatedTensorProducer(), _ReachabilityGatedTensorProducer(), sink], + ) + assert self_enabling.ok, self_enabling.summary() + + +def test_default_auto_scorers_do_not_fear_a_tensor_on_a_file_manifest(tmp_path: Path) -> None: + """``auto_partial`` hydration only REPLACES an incomplete resident pair; a file-only row never gains one.""" + from nemo_curator.stages.audio.filtering.sigmos import SIGMOSFilterStage + from nemo_curator.stages.audio.filtering.utmos import UTMOSFilterStage + + sink = ManifestWriterStage(output_path=str(tmp_path / "out.jsonl")) + for scorer in (UTMOSFilterStage(), SIGMOSFilterStage()): + report = validate_pipeline([scorer, sink]) + assert report.ok, report.summary() + + # With a resident sample_rate and no waveform the runtime DOES inject the decoded pair, + # so the refusal there is a true positive and must stay. + resident_rate = validate_pipeline( + [UTMOSFilterStage(), sink], + initial_keys={"audio_filepath", "sample_rate"}, + initial_roles={"audio_filepath", "sample_rate"}, + ) + assert any(issue.code == "tensor_into_sink" for issue in resident_rate.issues) diff --git a/tests/stages/audio/test_common.py b/tests/stages/audio/test_common.py index 21c7ba2aff..63812c8eba 100644 --- a/tests/stages/audio/test_common.py +++ b/tests/stages/audio/test_common.py @@ -805,15 +805,34 @@ def test_appends_across_multiple_process_calls(self, tmp_path: Path) -> None: assert len(lines) == 3 assert [json.loads(line)["entry"] for line in lines] == [1, 2, 3] - def test_setup_truncates_existing_file(self, tmp_path: Path) -> None: + def test_setup_on_node_truncates_existing_file(self, tmp_path: Path) -> None: out = tmp_path / "output.jsonl" out.write_text('{"old": "data"}\n') writer = ManifestWriterStage(output_path=str(out)) - writer.setup() + writer.setup_on_node() assert out.read_text() == "" + def test_worker_setup_never_truncates_committed_rows(self, tmp_path: Path) -> None: + """``setup()`` runs per worker actor; a replacement actor must not erase earlier rows.""" + out = tmp_path / "output.jsonl" + writer = ManifestWriterStage(output_path=str(out)) + writer.setup_on_node() + writer.setup() + writer.process(AudioTask(data={"audio_filepath": "a.wav"}, dataset_name="ds")) + writer.process(AudioTask(data={"audio_filepath": "b.wav"}, dataset_name="ds")) + + replacement = ManifestWriterStage(output_path=str(out)) + replacement.setup() # a restarted worker: no setup_on_node, no truncation + replacement.process(AudioTask(data={"audio_filepath": "c.wav"}, dataset_name="ds")) + + assert [json.loads(line)["audio_filepath"] for line in out.read_text().splitlines()] == [ + "a.wav", + "b.wav", + "c.wav", + ] + def test_setup_on_node_creates_parent_directories(self, tmp_path: Path) -> None: out = tmp_path / "nested" / "deep" / "output.jsonl" writer = ManifestWriterStage(output_path=str(out)) @@ -1064,17 +1083,19 @@ def test_resolve_model_path(tmp_path: Path) -> None: # Lifted from tests/stages/audio/test_agent_simulation_pipelines.py: ManifestWriterStage # lives in common.py, and this was its only truncate-on-rerun coverage. -def test_agent_manifest_writer_truncates_on_setup(tmp_path: Path) -> None: - """A fresh run (setup) truncates the output so reruns do not accumulate duplicates.""" +def test_agent_manifest_writer_truncates_on_setup_on_node(tmp_path: Path) -> None: + """A fresh run (setup_on_node) truncates the output so reruns do not accumulate duplicates.""" out_path = tmp_path / "manifest.jsonl" writer = ManifestWriterStage(output_path=str(out_path)) task = AudioTask(dataset_name="t", data={"audio_filepath": "src.wav", "text": "row"}) + writer.setup_on_node() writer.setup() writer.process(task) writer.process(task) assert len(out_path.read_text(encoding="utf-8").strip().splitlines()) == 2 # appends within a run - writer.setup() # new run truncates + writer.setup_on_node() # new run truncates + writer.setup() writer.process(task) assert len(out_path.read_text(encoding="utf-8").strip().splitlines()) == 1 From 0367061c91fe80226aeba55bbb4ccb97ef1b5e32 Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Thu, 17 Sep 2026 18:13:25 +0000 Subject: [PATCH 17/25] fix(audio-agent): model invalidated provenance keys Signed-off-by: shbhawsar --- nemo_curator/stages/audio/AGENT_READY.md | 6 ++++++ nemo_curator/stages/audio/_agent/_agent_ready.py | 5 +++++ .../stages/audio/_agent/_agent_registry.py | 1 + nemo_curator/stages/audio/_agent/_planning.py | 11 +++++++---- .../_agent/test_agent_foundation_regressions.py | 16 ++++++++++++++++ 5 files changed, 35 insertions(+), 4 deletions(-) diff --git a/nemo_curator/stages/audio/AGENT_READY.md b/nemo_curator/stages/audio/AGENT_READY.md index 1b2b746b65..b67a246a79 100644 --- a/nemo_curator/stages/audio/AGENT_READY.md +++ b/nemo_curator/stages/audio/AGENT_READY.md @@ -210,6 +210,12 @@ correctly run without. Optional reads are visible to discovery and semantic review, but they never make validation reject a fallback path. Do not put a fallback key in `reads` or legacy `inputs()` merely to advertise it. +Use `invalidates_keys=[...]` when compatibility requires retaining a key in +`task.data`, but its prior semantic role is no longer safe for downstream +planning. This differs from `removes_keys`: the runtime value remains available +as legacy provenance, while the planner treats it as unavailable until a later +stage writes a current value for that role. + When compatibility requires retaining a constructor default that runtime validation rejects, declare the field with `metadata={"agent_required": True}`. Discovery then requires an explicit value without changing the Python constructor or its default. Use this only for values that the diff --git a/nemo_curator/stages/audio/_agent/_agent_ready.py b/nemo_curator/stages/audio/_agent/_agent_ready.py index 1f333edc78..199524adf4 100644 --- a/nemo_curator/stages/audio/_agent/_agent_ready.py +++ b/nemo_curator/stages/audio/_agent/_agent_ready.py @@ -261,6 +261,10 @@ class StageContract: # positional compatibility and kept outside ``reads``/``reads_one_of`` so # planning never blocks a valid fallback path on their absence. optional_reads: IOSpec = field(default_factory=IOSpec) + # Task-data keys retained for compatibility/provenance whose prior semantic + # role is no longer a valid downstream carrier. Unlike ``removes_keys``, + # conformance does not require these keys to be physically absent. + invalidates_keys: list[str] = field(default_factory=list) def to_dict(self) -> dict[str, Any]: """Return a JSON-safe dict of this contract (``json.dumps`` never raises).""" @@ -289,6 +293,7 @@ def to_dict(self) -> dict[str, Any]: "accepts_task_type": self.accepts_task_type, "produces_task_type": self.produces_task_type, "removes_keys": list(self.removes_keys), + **({"invalidates_keys": list(self.invalidates_keys)} if self.invalidates_keys else {}), "conditional_writes": [ { "writes": asdict(item.writes), diff --git a/nemo_curator/stages/audio/_agent/_agent_registry.py b/nemo_curator/stages/audio/_agent/_agent_registry.py index 5de655eb58..b6e2b37967 100644 --- a/nemo_curator/stages/audio/_agent/_agent_registry.py +++ b/nemo_curator/stages/audio/_agent/_agent_registry.py @@ -341,6 +341,7 @@ def _contract_referenced_keys(contract: StageContract) -> set[str]: keys.update(conditional.writes.data_keys) keys.update(conditional.writes.segment_data_keys) keys.update(conditional.metadata_writes) + keys.update(contract.invalidates_keys) return keys diff --git a/nemo_curator/stages/audio/_agent/_planning.py b/nemo_curator/stages/audio/_agent/_planning.py index a08d8dd170..212a3a5c98 100644 --- a/nemo_curator/stages/audio/_agent/_planning.py +++ b/nemo_curator/stages/audio/_agent/_planning.py @@ -865,13 +865,16 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: walk.possible_segment_keys |= possible_segment walk.possible_roles |= _roles_for_keys(contract, possible) walk.possible_segment_roles |= _roles_for_keys(contract, possible_segment) - for rk in contract.removes_keys: + for rk in [*contract.removes_keys, *contract.invalidates_keys]: walk.available_keys.discard(rk) walk.possible_keys.discard(rk) - # ``removes_keys`` names TOP-LEVEL task keys, so dropping the carrier ends only the - # top-level tensor residency; a nested (segment) carrier of the same name survives. - walk.tensor_keys.discard(rk) + if rk in contract.removes_keys: + # ``removes_keys`` names TOP-LEVEL task keys, so dropping the carrier ends only the + # top-level tensor residency; a nested (segment) carrier of the same name survives. + walk.tensor_keys.discard(rk) role = contract.key_roles.get(rk, role_for_value(rk)) + if role != "unknown" and not any(role_for_value(k) == role for k in walk.possible_keys): + walk.possible_roles.discard(role) if ( role != "unknown" and role not in produced diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index 13ae5116b0..b54fa84104 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -109,6 +109,22 @@ def test_optional_reads_are_visible_without_blocking_fallback_paths() -> None: assert validate_pipeline([stage], initial_keys={"text"}).ok +def test_invalidated_provenance_key_is_retained_but_not_planner_available() -> None: + invalidator = _ConfiguredContractStage(StageContract(invalidates_keys=["audio_filepath"])) + consumer = _ConfiguredContractStage(StageContract(reads=IOSpec(data_keys=["audio_filepath"]))) + + contract = build_contract(invalidator) + report = validate_pipeline( + [invalidator, consumer], + initial_keys={"audio_filepath"}, + initial_roles={"audio_filepath"}, + ) + + assert contract.to_dict()["invalidates_keys"] == ["audio_filepath"] + assert not report.ok + assert any(issue.code == "key_removed_upstream" for issue in report.issues) + + def test_conditional_roles_are_discoverable_but_not_planner_guaranteed() -> None: producer_contract = StageContract( conditional_writes=[ From 584768387b56e65d43c72eae0bc09d7ace47a560 Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Thu, 17 Sep 2026 18:46:55 +0000 Subject: [PATCH 18/25] fix(audio-agent): validate conditional read branches Signed-off-by: shbhawsar --- nemo_curator/stages/audio/AGENT_READY.md | 7 ++ .../stages/audio/_agent/_agent_ready.py | 35 ++++++ .../stages/audio/_agent/_agent_registry.py | 6 + nemo_curator/stages/audio/_agent/_catalog.py | 5 +- .../stages/audio/_agent/_conformance.py | 40 ++++++- nemo_curator/stages/audio/_agent/_planning.py | 105 +++++++++++++----- .../stages/audio/_agent/_residency.py | 28 ++++- .../test_agent_foundation_regressions.py | 65 ++++++++++- 8 files changed, 253 insertions(+), 38 deletions(-) diff --git a/nemo_curator/stages/audio/AGENT_READY.md b/nemo_curator/stages/audio/AGENT_READY.md index b67a246a79..77cdd835cf 100644 --- a/nemo_curator/stages/audio/AGENT_READY.md +++ b/nemo_curator/stages/audio/AGENT_READY.md @@ -72,6 +72,13 @@ class MyStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): `requires_stable_task_id=True`; metadata-manifest resume boundaries cannot restore that identity and will reject the suffix. +Use `conditional_reads` when runtime selects a mutually exclusive read scope +from literal top-level key presence. Each `ConditionalRead` supplies one +`reads_one_of` group plus `requires_keys` and/or `forbids_keys` matching that +branch selector. The planner validates every reachable branch independently; +an unrelated task-level waveform therefore cannot satisfy a nested-audio branch +selected by the presence of `segments_key`. + ### Conditional/data-dependent writes Keep `writes` as the existing mechanical output declaration. When a listed diff --git a/nemo_curator/stages/audio/_agent/_agent_ready.py b/nemo_curator/stages/audio/_agent/_agent_ready.py index 199524adf4..6b00838e50 100644 --- a/nemo_curator/stages/audio/_agent/_agent_ready.py +++ b/nemo_curator/stages/audio/_agent/_agent_ready.py @@ -144,6 +144,22 @@ class ConditionalWrite: requires_keys: list[str] = field(default_factory=list) +@dataclass(frozen=True) +class ConditionalRead: + """Alternative reads selected by literal top-level key presence. + + ``requires_keys`` and ``forbids_keys`` describe the runtime branch selector. + Every reachable branch must have one complete ``reads_one_of`` alternative; + this prevents an unrelated task-level input from satisfying a branch that + runtime will execute against nested items instead. + """ + + reads_one_of: list[IOSpec] = field(default_factory=list) + condition: str = "" + requires_keys: list[str] = field(default_factory=list) + forbids_keys: list[str] = field(default_factory=list) + + @dataclass(frozen=True) class Gates: """Execution gates or side effects an agent should know before wrapping a stage.""" @@ -265,6 +281,10 @@ class StageContract: # role is no longer a valid downstream carrier. Unlike ``removes_keys``, # conformance does not require these keys to be physically absent. invalidates_keys: list[str] = field(default_factory=list) + # Mutually exclusive read groups selected by literal top-level key presence. + # Appended for positional compatibility; ordinary contracts continue using + # ``reads`` and ``reads_one_of`` exactly as before. + conditional_reads: list[ConditionalRead] = field(default_factory=list) def to_dict(self) -> dict[str, Any]: """Return a JSON-safe dict of this contract (``json.dumps`` never raises).""" @@ -294,6 +314,21 @@ def to_dict(self) -> dict[str, Any]: "produces_task_type": self.produces_task_type, "removes_keys": list(self.removes_keys), **({"invalidates_keys": list(self.invalidates_keys)} if self.invalidates_keys else {}), + **( + { + "conditional_reads": [ + { + "reads_one_of": [asdict(spec) for spec in item.reads_one_of], + "condition": item.condition, + **({"requires_keys": list(item.requires_keys)} if item.requires_keys else {}), + **({"forbids_keys": list(item.forbids_keys)} if item.forbids_keys else {}), + } + for item in self.conditional_reads + ] + } + if self.conditional_reads + else {} + ), "conditional_writes": [ { "writes": asdict(item.writes), diff --git a/nemo_curator/stages/audio/_agent/_agent_registry.py b/nemo_curator/stages/audio/_agent/_agent_registry.py index b6e2b37967..6a568fa499 100644 --- a/nemo_curator/stages/audio/_agent/_agent_registry.py +++ b/nemo_curator/stages/audio/_agent/_agent_registry.py @@ -341,6 +341,12 @@ def _contract_referenced_keys(contract: StageContract) -> set[str]: keys.update(conditional.writes.data_keys) keys.update(conditional.writes.segment_data_keys) keys.update(conditional.metadata_writes) + for conditional in contract.conditional_reads: + keys.update(conditional.requires_keys) + keys.update(conditional.forbids_keys) + for spec in conditional.reads_one_of: + keys.update(spec.data_keys) + keys.update(spec.segment_data_keys) keys.update(contract.invalidates_keys) return keys diff --git a/nemo_curator/stages/audio/_agent/_catalog.py b/nemo_curator/stages/audio/_agent/_catalog.py index 2c5bf1b22e..4e54cacc87 100644 --- a/nemo_curator/stages/audio/_agent/_catalog.py +++ b/nemo_curator/stages/audio/_agent/_catalog.py @@ -166,12 +166,15 @@ def catalog_as_json(*, include_dynamic_defaults: bool = False, indent: int | Non # Role -> producer/consumer index (composition + repair) # --------------------------------------------------------------------------- # def _consumed_roles(contract: StageContract) -> set[str]: - """Semantic roles a stage requires (primary reads + every reads_one_of option).""" + """Semantic roles a stage requires across unconditional and conditional reads.""" roles = { contract.key_roles.get(k, "unknown") for k in [*contract.reads.data_keys, *contract.reads.segment_data_keys] } for opt in contract.reads_one_of: roles |= {contract.key_roles.get(k, "unknown") for k in [*opt.data_keys, *opt.segment_data_keys]} + for branch in contract.conditional_reads: + for opt in branch.reads_one_of: + roles |= {contract.key_roles.get(k, "unknown") for k in [*opt.data_keys, *opt.segment_data_keys]} return roles - {"unknown"} diff --git a/nemo_curator/stages/audio/_agent/_conformance.py b/nemo_curator/stages/audio/_agent/_conformance.py index a74cf074e8..6a3ae51e35 100644 --- a/nemo_curator/stages/audio/_agent/_conformance.py +++ b/nemo_curator/stages/audio/_agent/_conformance.py @@ -92,18 +92,24 @@ def reads_satisfied_by_role(consumer: StageContract, available_roles: set[str]) reads_keys = [*consumer.reads.data_keys, *consumer.reads.segment_data_keys] if reads_keys and not _spec_roles(consumer, reads_keys).issubset(avail): return False - if consumer.reads_one_of: - return any( + if consumer.reads_one_of and not any( + _spec_roles(consumer, [*opt.data_keys, *opt.segment_data_keys]).issubset(avail) + for opt in consumer.reads_one_of + ): + return False + return all( + any( _spec_roles(consumer, [*opt.data_keys, *opt.segment_data_keys]).issubset(avail) - for opt in consumer.reads_one_of + for opt in branch.reads_one_of ) - return True + for branch in consumer.conditional_reads + ) # --------------------------------------------------------------------------- # # Static checks (no execution) # --------------------------------------------------------------------------- # -def _check_shape(c: StageContract, name: str) -> None: # noqa: C901 +def _check_shape(c: StageContract, name: str) -> None: # noqa: C901, PLR0912 assert c.cardinality in _VALID_CARDINALITY, f"{name}: invalid cardinality {c.cardinality!r}" for opt in c.cardinality_options: # cardinality_options are short flag names (e.g. "fan_out","nested") OR full cardinalities @@ -122,6 +128,12 @@ def _check_shape(c: StageContract, name: str) -> None: # noqa: C901 for spec in [c.reads, c.writes, c.optional_reads, *c.reads_one_of]: contract_keys.update(spec.data_keys) contract_keys.update(spec.segment_data_keys) + for conditional in c.conditional_reads: + contract_keys.update(conditional.requires_keys) + contract_keys.update(conditional.forbids_keys) + for spec in conditional.reads_one_of: + contract_keys.update(spec.data_keys) + contract_keys.update(spec.segment_data_keys) assert c.iteration_key in contract_keys or c.iteration_key in c.key_roles, ( f"{name}: iteration_key {c.iteration_key!r} is neither a contract read/write " f"key nor a role-resolvable key value — it names nothing an agent can find" @@ -161,6 +173,18 @@ def _check_shape(c: StageContract, name: str) -> None: # noqa: C901 assert all(isinstance(key, str) and key for key in conditional.metadata_writes), ( f"{name}: {label}.metadata_writes must contain non-empty strings" ) + for index, conditional in enumerate(c.conditional_reads): + label = f"conditional_reads[{index}]" + assert conditional.condition.strip(), f"{name}: {label}.condition must be non-empty" + assert conditional.reads_one_of, f"{name}: {label}.reads_one_of must not be empty" + assert not (set(conditional.requires_keys) & set(conditional.forbids_keys)), ( + f"{name}: {label} cannot require and forbid the same selector key" + ) + for spec in conditional.reads_one_of: + for a in spec.accepts: + assert a in _VALID_ACCEPTS, f"{name}: {label}.accepts has invalid form {a!r}" + for p in spec.produces: + assert p in _VALID_PRODUCES, f"{name}: {label}.produces has invalid form {p!r}" # no duplicate keys within a single spec list for spec, label in [(c.reads, "reads"), (c.writes, "writes")]: assert len(spec.data_keys) == len(set(spec.data_keys)), f"{name}: duplicate {label}.data_keys" @@ -238,7 +262,11 @@ def _check_residency_accepts(stage: Any, c: StageContract, name: str) -> None: residency = getattr(stage, "input_residency", None) if residency is None: return - declared = set(c.reads.accepts) | {a for opt in c.reads_one_of for a in opt.accepts} + declared = ( + set(c.reads.accepts) + | {a for opt in c.reads_one_of for a in opt.accepts} + | {a for branch in c.conditional_reads for opt in branch.reads_one_of for a in opt.accepts} + ) if not declared: return expected = set(accepts_for_residency(residency)) diff --git a/nemo_curator/stages/audio/_agent/_planning.py b/nemo_curator/stages/audio/_agent/_planning.py index 212a3a5c98..087335b06d 100644 --- a/nemo_curator/stages/audio/_agent/_planning.py +++ b/nemo_curator/stages/audio/_agent/_planning.py @@ -153,6 +153,11 @@ def _requirement_str(contract: StageContract, available: set[str]) -> str: reqs.append(f"role(s) {sorted(missing)}") if contract.reads_one_of: reqs.append(f"one of {[sorted(_roles_of(o, contract)) for o in contract.reads_one_of]}") + for branch in contract.conditional_reads: + reqs.append( + f"when {branch.condition}: one of " + f"{[sorted(_roles_of(option, contract)) for option in branch.reads_one_of]}" + ) return "; ".join(reqs) or f"role(s) {sorted(_required_roles(contract))}" @@ -194,12 +199,41 @@ def scope_satisfied(keys: list[str], roles: set[str], literal_keys: set[str]) -> ) -def _reads_satisfied_by_role( +def _reachable_conditional_reads( + contract: StageContract, + available_keys: set[str], + possible_keys: set[str] | None = None, +) -> list[Any]: + """Read branches that runtime can select from the current top-level schema.""" + possible = possible_keys or set() + reachable_keys = available_keys | possible + return [ + branch + for branch in contract.conditional_reads + if set(branch.requires_keys).issubset(reachable_keys) and not (set(branch.forbids_keys) & available_keys) + ] + + +def _read_option_groups( + contract: StageContract, + available_keys: set[str], + possible_keys: set[str] | None = None, +) -> list[list[Any]]: + """Alternative groups whose runtime branch is reachable for this input.""" + groups = [contract.reads_one_of] if contract.reads_one_of else [] + groups.extend( + branch.reads_one_of for branch in _reachable_conditional_reads(contract, available_keys, possible_keys) + ) + return groups + + +def _reads_satisfied_by_role( # noqa: PLR0913 -- task/segment roles and keys are distinct planner state contract: StageContract, available_roles: set[str], available_segment_roles: set[str], available_keys: set[str], available_segment_keys: set[str], + possible_keys: set[str] | None = None, ) -> bool: """Role-level read check with literal fallback for unknown roles.""" if not _spec_satisfied_by_role( @@ -211,16 +245,19 @@ def _reads_satisfied_by_role( available_segment_keys, ): return False - return not contract.reads_one_of or any( - _spec_satisfied_by_role( - option, - contract, - available_roles, - available_segment_roles, - available_keys, - available_segment_keys, + return all( + any( + _spec_satisfied_by_role( + option, + contract, + available_roles, + available_segment_roles, + available_keys, + available_segment_keys, + ) + for option in group ) - for option in contract.reads_one_of + for group in _read_option_groups(contract, available_keys, possible_keys) ) @@ -382,6 +419,7 @@ def _reads_satisfied_by_key( contract: StageContract, available_keys: set[str], available_segment_keys: set[str], + possible_keys: set[str] | None = None, ) -> bool: """Whether every read is met by the LITERAL key it names. @@ -399,8 +437,9 @@ def _reads_satisfied_by_key( """ if not _spec_satisfied_by_key(contract.reads, available_keys, available_segment_keys): return False - return not contract.reads_one_of or any( - _spec_satisfied_by_key(option, available_keys, available_segment_keys) for option in contract.reads_one_of + return all( + any(_spec_satisfied_by_key(option, available_keys, available_segment_keys) for option in group) + for group in _read_option_groups(contract, available_keys, possible_keys) ) @@ -486,6 +525,7 @@ def _dangling_read_keys( contract: StageContract, available_keys: set[str], available_segment_keys: set[str], + possible_keys: set[str] | None = None, ) -> set[str]: """Read key VALUES whose role is known but whose exact value was not produced upstream nor seeded — the renamed-producer dangle the role check misses. @@ -502,14 +542,10 @@ def _dangling_read_keys( Role-bearing keys only (``unknown``/internal bookkeeping keys are excluded). """ dangling = _missing_literal_role_keys(contract, contract.reads, available_keys, available_segment_keys) - options = contract.reads_one_of - if len(options) == 1: - dangling |= _missing_literal_role_keys(contract, options[0], available_keys, available_segment_keys) - elif len(options) > 1: + for options in _read_option_groups(contract, available_keys, possible_keys): per_option = [ _missing_literal_role_keys(contract, option, available_keys, available_segment_keys) for option in options ] - # Clean iff at least one alternative is literally complete (no missing role-bearing key). if all(missing for missing in per_option): dangling |= set().union(*per_option) return dangling @@ -558,8 +594,14 @@ def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[Pipe walk.segment_available, walk.available_keys, walk.segment_available_keys, + walk.possible_keys, + ) + key_satisfied = _reads_satisfied_by_key( + contract, + walk.available_keys, + walk.segment_available_keys, + walk.possible_keys, ) - key_satisfied = _reads_satisfied_by_key(contract, walk.available_keys, walk.segment_available_keys) if not (role_satisfied or key_satisfied) and not walk.past_composite and site.composite is None: conditional = _conditional_read_issue(walk, site, contract) if conditional is not None: @@ -568,7 +610,12 @@ def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[Pipe if walk.past_composite: return [] out: list[PipelineIssue] = [] - dangling = _dangling_read_keys(contract, walk.available_keys, walk.segment_available_keys) + dangling = _dangling_read_keys( + contract, + walk.available_keys, + walk.segment_available_keys, + walk.possible_keys, + ) if dangling: out.append( PipelineIssue( @@ -628,8 +675,9 @@ def _read_issues(walk: _Walk, site: _Site, contract: StageContract) -> list[Pipe ) ] + option_groups = _read_option_groups(contract, walk.available_keys, walk.possible_keys) needed = _roles_for_keys(contract, contract.reads.data_keys) | { - role for option in contract.reads_one_of for role in _roles_for_keys(contract, option.data_keys) + role for group in option_groups for option in group for role in _roles_for_keys(contract, option.data_keys) } removed_hit = (needed & walk.removed_roles) - walk.available if removed_hit: @@ -679,7 +727,10 @@ def spec_ok(spec: Any) -> bool: # noqa: ANN401 - IOSpec if not spec_ok(contract.reads): return False - return not contract.reads_one_of or any(spec_ok(option) for option in contract.reads_one_of) + return all( + any(spec_ok(option) for option in group) + for group in _read_option_groups(contract, walk.available_keys, walk.possible_keys) + ) def _conditional_read_issue(walk: _Walk, site: _Site, contract: StageContract) -> PipelineIssue | None: @@ -695,9 +746,13 @@ def _conditional_read_issue(walk: _Walk, site: _Site, contract: StageContract) - """ if not _reads_possibly_satisfied(walk, contract): return None + option_groups = _read_option_groups(contract, walk.available_keys, walk.possible_keys) only_possible = sorted( ( - {*contract.reads.data_keys, *(k for option in contract.reads_one_of for k in option.data_keys)} + { + *contract.reads.data_keys, + *(k for group in option_groups for option in group for k in option.data_keys), + } & walk.possible_keys ) - walk.available_keys @@ -705,7 +760,7 @@ def _conditional_read_issue(walk: _Walk, site: _Site, contract: StageContract) - ( { *contract.reads.segment_data_keys, - *(k for option in contract.reads_one_of for k in option.segment_data_keys), + *(k for group in option_groups for option in group for k in option.segment_data_keys), } & walk.possible_segment_keys ) @@ -800,7 +855,7 @@ def _task_type_issue(walk: _Walk, site: _Site, contract: StageContract) -> list[ ] -def _advance(walk: _Walk, contract: StageContract, name: str) -> None: +def _advance(walk: _Walk, contract: StageContract, name: str) -> None: # noqa: PLR0915 """Fold one stage's writes, removals and tensor residency into the running state.""" produced = _roles_for_keys(contract, contract.writes.data_keys) segment_produced = _roles_for_keys(contract, contract.writes.segment_data_keys) @@ -812,7 +867,7 @@ def _advance(walk: _Walk, contract: StageContract, name: str) -> None: # Keys this stage drops on its main path. A conditional write of the same key (a branch that # happens to keep it, or re-emits it with a different meaning such as a tar member name) must # not resurrect it for planning: removal is the guarantee-level fact, the branch the exception. - blocked_keys = set(contract.removes_keys) + blocked_keys = set(contract.removes_keys) | set(contract.invalidates_keys) blocked_segment_keys: set[str] = set() if not contract.preserves_upstream_keys: blocked_keys |= walk.available_keys - written diff --git a/nemo_curator/stages/audio/_agent/_residency.py b/nemo_curator/stages/audio/_agent/_residency.py index bdfdbcfa5d..8fab4688fe 100644 --- a/nemo_curator/stages/audio/_agent/_residency.py +++ b/nemo_curator/stages/audio/_agent/_residency.py @@ -23,7 +23,7 @@ import soundfile as sf import torch -from nemo_curator.stages.audio._agent._agent_ready import AudioForm, ConditionalWrite, IOSpec +from nemo_curator.stages.audio._agent._agent_ready import AudioForm, ConditionalRead, ConditionalWrite, IOSpec from nemo_curator.stages.audio.common import ensure_waveform_2d, load_audio_file if TYPE_CHECKING: @@ -58,6 +58,12 @@ def validate_audio_key_configuration( msg = f"[{stage_name}] '{field_name}' must be a non-empty string" raise ValueError(msg) + input_values = list(input_keys.values()) + if len(input_values) != len(set(input_values)): + duplicates = sorted({key for key in input_values if input_values.count(key) > 1}) + msg = f"[{stage_name}] Audio input keys must be distinct; duplicate values: {duplicates}" + raise ValueError(msg) + output_values = list(output_keys.values()) if len(output_values) != len(set(output_values)): duplicates = sorted({key for key in output_values if output_values.count(key) > 1}) @@ -157,7 +163,7 @@ def scoped_audio_io_specs( # noqa: PLR0913 segments_key: str, output_keys: list[str], infer_sample_rate_from_file: bool = False, -) -> tuple[IOSpec, list[IOSpec], IOSpec]: +) -> tuple[IOSpec, list[IOSpec], IOSpec, list[ConditionalRead]]: """Build mode-accurate reads/writes for task-or-nested audio stages. ``task`` exposes only top-level residency alternatives and outputs; @@ -177,7 +183,6 @@ def scoped_audio_io_specs( # noqa: PLR0913 ) segment_reads = [ IOSpec( - data_keys=[segments_key] if mode == "auto" else [], segment_data_keys=list(spec.data_keys), accepts=list(spec.accepts), ) @@ -185,17 +190,30 @@ def scoped_audio_io_specs( # noqa: PLR0913 ] if mode == "task": - return IOSpec(), task_reads, IOSpec(data_keys=list(output_keys)) + return IOSpec(), task_reads, IOSpec(data_keys=list(output_keys)), [] if mode == "segments": return ( IOSpec(data_keys=[segments_key]), segment_reads, IOSpec(segment_data_keys=list(output_keys)), + [], ) return ( IOSpec(), - [*task_reads, *segment_reads], + [], IOSpec(data_keys=list(output_keys), segment_data_keys=list(output_keys)), + [ + ConditionalRead( + reads_one_of=task_reads, + condition=f"'{segments_key}' is absent, so the task-level branch runs", + forbids_keys=[segments_key], + ), + ConditionalRead( + reads_one_of=segment_reads, + condition=f"'{segments_key}' is present, so the per-segment branch runs", + requires_keys=[segments_key], + ), + ], ) diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index b54fa84104..eca859f9fb 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -30,7 +30,13 @@ from nemo_curator.stages import audio from nemo_curator.stages.audio import agent from nemo_curator.stages.audio._agent import _catalog -from nemo_curator.stages.audio._agent._agent_ready import AgentReady, ConditionalWrite, IOSpec, StageContract +from nemo_curator.stages.audio._agent._agent_ready import ( + AgentReady, + ConditionalRead, + ConditionalWrite, + IOSpec, + StageContract, +) from nemo_curator.stages.audio._agent._agent_registry import build_contract, stage_params, static_contract from nemo_curator.stages.audio._agent._catalog import unavailable_modules from nemo_curator.stages.audio._agent._composite import expand_composites @@ -40,6 +46,7 @@ cleanup_temp_files, resolve_audio, resolve_audio_path, + validate_audio_key_configuration, validate_input_residency, write_audio_stable, ) @@ -109,6 +116,53 @@ def test_optional_reads_are_visible_without_blocking_fallback_paths() -> None: assert validate_pipeline([stage], initial_keys={"text"}).ok +def test_conditional_reads_follow_the_runtime_scope_selector() -> None: + contract = StageContract( + conditional_reads=[ + ConditionalRead( + reads_one_of=[IOSpec(data_keys=["waveform", "sample_rate"])], + condition="'segments' is absent", + forbids_keys=["segments"], + ), + ConditionalRead( + reads_one_of=[IOSpec(segment_data_keys=["waveform", "sample_rate"])], + condition="'segments' is present", + requires_keys=["segments"], + ), + ], + key_roles={ + "segments": "segments", + "waveform": "waveform", + "sample_rate": "sample_rate", + }, + ) + stage = _ConfiguredContractStage(contract) + + task_report = validate_pipeline( + [stage], + initial_keys={"waveform", "sample_rate"}, + initial_roles={"waveform", "sample_rate"}, + ) + incomplete_nested_report = validate_pipeline( + [stage], + initial_keys={"waveform", "sample_rate", "segments"}, + initial_roles={"waveform", "sample_rate", "segments"}, + initial_segment_keys={"segment_num"}, + ) + complete_nested_report = validate_pipeline( + [stage], + initial_keys={"waveform", "sample_rate", "segments"}, + initial_roles={"waveform", "sample_rate", "segments"}, + initial_segment_keys={"waveform", "sample_rate"}, + initial_segment_roles={"waveform", "sample_rate"}, + ) + + assert task_report.ok + assert not incomplete_nested_report.ok + assert complete_nested_report.ok + assert contract.to_dict()["conditional_reads"][1]["requires_keys"] == ["segments"] + + def test_invalidated_provenance_key_is_retained_but_not_planner_available() -> None: invalidator = _ConfiguredContractStage(StageContract(invalidates_keys=["audio_filepath"])) consumer = _ConfiguredContractStage(StageContract(reads=IOSpec(data_keys=["audio_filepath"]))) @@ -251,6 +305,15 @@ def test_input_residency_validator_rejects_unknown_mode() -> None: validate_input_residency("wavefrom", stage_name="Fixture") +def test_audio_key_validator_rejects_input_role_aliases() -> None: + with pytest.raises(ValueError, match="Audio input keys must be distinct"): + validate_audio_key_configuration( + "Fixture", + input_keys={"waveform_key": "audio", "sample_rate_key": "audio"}, + output_keys={"score_key": "score"}, + ) + + def test_file_audio_hydration_policies_are_opt_in_and_atomic(tmp_path: Path) -> None: path = tmp_path / "audio.wav" path.touch() From b261c5db617436d739016be1d378ddc6634d2ce1 Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Fri, 18 Sep 2026 09:15:23 +0000 Subject: [PATCH 19/25] fix(audio-agent): prevent unsafe bounded-source delta reuse Signed-off-by: shbhawsar --- nemo_curator/stages/audio/common.py | 10 ++++------ tests/stages/audio/test_common.py | 9 +++++++++ 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/nemo_curator/stages/audio/common.py b/nemo_curator/stages/audio/common.py index a7d62f857d..b6784bc38e 100644 --- a/nemo_curator/stages/audio/common.py +++ b/nemo_curator/stages/audio/common.py @@ -691,12 +691,10 @@ def describe(self) -> StageContract: # (unlike the dataset CreateInitialManifest*Stage sources, which download and write). writes=IOSpec(data_keys=[self.audio_filepath_key, self.audio_item_id_key]), cardinality="1:N fan-out", - # Scans existing files; one task per file, and a row says nothing about its - # neighbours. Declared True unconditionally BY DECISION: under a bounded - # ``max_samples`` the SORTED listing is truncated, so a delta can admit files a full - # run would not have. Accepted rather than cost every bounded run its reuse -- not - # an oversight to "fix" back. - gates=Gates(per_row_independent=True), + # A bounded sorted listing depends on the whole folder: adding an earlier path can + # change which existing rows belong to the first N. Unbounded scans remain safely + # narrowable one file at a time. + gates=Gates(per_row_independent=self.max_samples is None or self.max_samples < 0), ) def ray_stage_spec(self) -> dict[str, Any]: diff --git a/tests/stages/audio/test_common.py b/tests/stages/audio/test_common.py index 63812c8eba..fbb456e366 100644 --- a/tests/stages/audio/test_common.py +++ b/tests/stages/audio/test_common.py @@ -28,6 +28,7 @@ from nemo_curator.pipeline import Pipeline from nemo_curator.stages.audio.alm import ALMDataBuilderStage, ALMDataOverlapStage from nemo_curator.stages.audio.common import ( + CreateInitialManifestAudioFolderStage, GetAudioDurationStage, ManifestCheckpointStage, ManifestReader, @@ -439,6 +440,14 @@ def test_compound_preserve_nested_contract_uses_only_top_level_container_key() - assert "OR" in or_contract.description +def test_bounded_audio_folder_source_is_not_row_independent() -> None: + bounded = CreateInitialManifestAudioFolderStage(data_dir="/tmp/x", max_samples=10) # noqa: S108 + unbounded = CreateInitialManifestAudioFolderStage(data_dir="/tmp/x") # noqa: S108 + + assert bounded.describe().gates.per_row_independent is False + assert unbounded.describe().gates.per_row_independent is True + + # --------------------------------------------------------------------------- # GetAudioDurationStage # --------------------------------------------------------------------------- From be9118e2fddf68a6410665e769c36dc93b587ab0 Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Fri, 18 Sep 2026 12:07:15 +0000 Subject: [PATCH 20/25] fix(audio): close agent foundation review gaps Signed-off-by: shbhawsar --- .../stages/audio/_agent/_conformance.py | 8 ++-- nemo_curator/stages/audio/common.py | 39 +++++++++++++++++-- .../audio/preprocessing/channel_count.py | 20 ++++++++++ .../audio/preprocessing/mono_conversion.py | 24 ++++++++++++ .../audio/preprocessing/sample_rate_filter.py | 13 +++++++ .../test_channel_and_rate_stages.py | 26 +++++++++++++ .../preprocessing/test_mono_conversion.py | 21 +++++++++- tests/stages/audio/test_common.py | 15 +++++++ .../test_create_manifest_audio_folder.py | 20 ++++++++++ 9 files changed, 177 insertions(+), 9 deletions(-) diff --git a/nemo_curator/stages/audio/_agent/_conformance.py b/nemo_curator/stages/audio/_agent/_conformance.py index 6a3ae51e35..677193dcd2 100644 --- a/nemo_curator/stages/audio/_agent/_conformance.py +++ b/nemo_curator/stages/audio/_agent/_conformance.py @@ -221,10 +221,10 @@ def _check_per_row_independence(c: StageContract, name: str) -> None: says so per instance and ``delta.region`` stops there. Requiring ``True`` would force such a source to lie or to drop the parameter. Silence is what is forbidden. - The case that first motivated this -- ``CreateInitialManifestAudioFolderStage`` under a - bounded ``max_samples``, which truncates the sorted listing -- now declares ``True`` anyway - by an explicit product decision recorded at that declaration. The rule is unchanged: it was - never "must be False when narrowing is lossy", only "must not be silent". + ``CreateInitialManifestAudioFolderStage`` is the conditional reference case: a bounded + ``max_samples`` truncates the sorted listing and is not independent, while the default + unbounded scan is. Its static hint is conservatively ``False`` because discovery cannot + resolve constructor values; its configured contract reports the precise per-instance answer. A companion rule ("``True`` contradicts ``N:1``") was removed as wrong: cardinality counts TASKS, this is about row VALUES, and ``AudioToDocumentStage`` repacks tasks while leaving diff --git a/nemo_curator/stages/audio/common.py b/nemo_curator/stages/audio/common.py index b6784bc38e..f34bbd6778 100644 --- a/nemo_curator/stages/audio/common.py +++ b/nemo_curator/stages/audio/common.py @@ -68,6 +68,25 @@ class GetAudioDurationStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): sample_rate_key: str = "sample_rate" input_residency: Literal["file", "waveform", "auto"] = "file" + def __post_init__(self) -> None: + super().__init__() + # Lazy import avoids the module-level cycle (_residency imports helpers below). + from nemo_curator.stages.audio._agent._residency import ( + validate_audio_key_configuration, + validate_input_residency, + ) + + validate_input_residency(self.input_residency, stage_name=self.name) + validate_audio_key_configuration( + self.name, + input_keys={ + "audio_filepath_key": self.audio_filepath_key, + "waveform_key": self.waveform_key, + "sample_rate_key": self.sample_rate_key, + }, + output_keys={"duration_key": self.duration_key}, + ) + def setup(self, _worker_metadata: WorkerMetadata | None = None) -> None: import soundfile @@ -670,14 +689,26 @@ class CreateInitialManifestAudioFolderStage(AgentReady, ProcessingStage[EmptyTas audio_item_id_key: str = "audio_item_id" name: str = "CreateInitialManifestAudioFolder" batch_size: int = 1 - # See ManifestReaderStage: the narrowing claim has to survive being read off the class. - AGENT_STATIC: ClassVar[StaticHints] = StaticHints(gates=Gates(per_row_independent=True)) + # Static discovery cannot resolve max_samples. Be conservative there; configured + # contracts below recover True for the default unbounded scan. + AGENT_STATIC: ClassVar[StaticHints] = StaticHints(gates=Gates(per_row_independent=False)) def __post_init__(self) -> None: super().__init__() if not self.data_dir: msg = "data_dir is required for CreateInitialManifestAudioFolderStage" raise ValueError(msg) + # Lazy import avoids the module-level cycle (_residency imports helpers above). + from nemo_curator.stages.audio._agent._residency import validate_audio_key_configuration + + validate_audio_key_configuration( + self.name, + input_keys={}, + output_keys={ + "audio_filepath_key": self.audio_filepath_key, + "audio_item_id_key": self.audio_item_id_key, + }, + ) def inputs(self) -> tuple[list[str], list[str]]: return [], [] @@ -754,7 +785,7 @@ def _collect_audio_files(self) -> list[str]: ) return sorted(found) - def process(self, _: EmptyTask) -> list[AudioTask]: + def process(self, task: EmptyTask | None) -> list[AudioTask]: """Emit one AudioTask per audio file found under ``data_dir``.""" paths = self._collect_audio_files() if self.max_samples is not None and self.max_samples >= 0: @@ -776,6 +807,8 @@ def process(self, _: EmptyTask) -> list[AudioTask]: dataset_name="local-audio-folder", data={self.audio_filepath_key: abspath, self.audio_item_id_key: item_id}, filepath_key=self.audio_filepath_key, + _metadata={} if task is None else task._metadata, + _stage_perf=[] if task is None else list(task._stage_perf), ) ) logger.info(f"[{self.name}] created {len(tasks)} AudioTask(s) from {self.data_dir}") diff --git a/nemo_curator/stages/audio/preprocessing/channel_count.py b/nemo_curator/stages/audio/preprocessing/channel_count.py index d8a09deffb..322414112b 100644 --- a/nemo_curator/stages/audio/preprocessing/channel_count.py +++ b/nemo_curator/stages/audio/preprocessing/channel_count.py @@ -51,6 +51,8 @@ reject_sinkless_conversion, residency_read_specs, resolve_audio, + validate_audio_key_configuration, + validate_input_residency, write_audio_stable, ) from nemo_curator.stages.audio.common import ensure_waveform_2d, load_audio_file @@ -191,6 +193,24 @@ def __post_init__(self): msg = f"action must be one of ('annotate', 'filter', 'convert'), got {self.action!r}" raise ValueError(msg) self._reject_other_actions_params() + validate_input_residency(self.input_residency, stage_name=self.name) + input_keys = { + "audio_filepath_key": self.audio_filepath_key, + "waveform_key": self.waveform_key, + "sample_rate_key": self.sample_rate_key, + } + output_keys = {"num_channels_key": self.num_channels_key} + if self.action == "convert": + output_keys["duration_key"] = self.duration_key + if self.write_to_disk: + output_keys["output_audio_filepath_key"] = self.output_audio_filepath_key + if self.update_audio_filepath: + output_keys["original_audio_filepath_key"] = self.original_audio_filepath_key + validate_audio_key_configuration( + self.name, + input_keys=input_keys, + output_keys=output_keys, + ) if self.action == "filter": self._validate_filter() if self.action == "convert": diff --git a/nemo_curator/stages/audio/preprocessing/mono_conversion.py b/nemo_curator/stages/audio/preprocessing/mono_conversion.py index da09294846..0a8743efd7 100755 --- a/nemo_curator/stages/audio/preprocessing/mono_conversion.py +++ b/nemo_curator/stages/audio/preprocessing/mono_conversion.py @@ -40,6 +40,8 @@ reject_sinkless_conversion, residency_read_specs, resolve_audio, + validate_audio_key_configuration, + validate_input_residency, write_audio_stable, ) from nemo_curator.stages.audio.common import ensure_waveform_2d, load_audio_file @@ -115,6 +117,28 @@ class MonoConversionStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): def __post_init__(self): super().__init__() + validate_input_residency(self.input_residency, stage_name=self.name) + output_keys = { + "is_mono_key": self.is_mono_key, + "duration_key": self.duration_key, + "num_samples_key": self.num_samples_key, + } + if self.write_to_disk: + output_keys["output_audio_filepath_key"] = self.output_audio_filepath_key + if self.update_audio_filepath: + output_keys["original_audio_filepath_key"] = self.original_audio_filepath_key + validate_audio_key_configuration( + self.name, + input_keys={ + "audio_filepath_key": self.audio_filepath_key, + "waveform_key": self.waveform_key, + "sample_rate_key": self.sample_rate_key, + }, + # waveform/sample_rate and, conditionally, audio_filepath are intentional + # in-place writes. Only the independently named metadata/path outputs belong + # here, where collisions would erase an audio carrier or one another. + output_keys=output_keys, + ) reject_sinkless_conversion( stage="MonoConversionStage", keep_waveform_in_task=self.keep_waveform_in_task, diff --git a/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py b/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py index 16020197ad..dc3f3684e3 100644 --- a/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py +++ b/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py @@ -43,6 +43,7 @@ from loguru import logger from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, IOSpec, StageContract +from nemo_curator.stages.audio._agent._residency import validate_audio_key_configuration from nemo_curator.stages.base import ProcessingStage from nemo_curator.stages.resources import Resources from nemo_curator.tasks import AudioTask @@ -94,6 +95,18 @@ class SampleRateFilterStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): def __post_init__(self): super().__init__() + validate_audio_key_configuration( + self.name, + input_keys={ + "audio_filepath_key": self.audio_filepath_key, + "sample_rate_key": self.sample_rate_key, + "waveform_key": self.waveform_key, + }, + # sample_rate_key is an intentional in-place observation: validate it as an + # input so it cannot alias either audio carrier, rather than treating the + # legitimate read/write of that same key as a collision. + output_keys={}, + ) if self.allowed_sample_rates is not None and not self.allowed_sample_rates: msg = "allowed_sample_rates must name at least one rate, or be None for no constraint" raise ValueError(msg) diff --git a/tests/stages/audio/preprocessing/test_channel_and_rate_stages.py b/tests/stages/audio/preprocessing/test_channel_and_rate_stages.py index d6fb01273d..251489ab97 100644 --- a/tests/stages/audio/preprocessing/test_channel_and_rate_stages.py +++ b/tests/stages/audio/preprocessing/test_channel_and_rate_stages.py @@ -305,6 +305,21 @@ def test_an_unknown_action_is_rejected_at_construction(self) -> None: with pytest.raises(ValueError, match="action must be one of"): ChannelCountStage(action="downmix") + @pytest.mark.parametrize("residency", ["disk", "wavefrom", ""]) + def test_an_unknown_residency_is_rejected_at_construction(self, residency: str) -> None: + with pytest.raises(ValueError, match="input_residency"): + ChannelCountStage(input_residency=residency) # type: ignore[arg-type] + + @pytest.mark.parametrize("action", ["annotate", "filter", "convert"]) + @pytest.mark.parametrize("carrier_key", ["audio_filepath", "waveform", "sample_rate"]) + def test_the_count_key_cannot_overwrite_an_audio_carrier(self, action: str, carrier_key: str) -> None: + with pytest.raises(ValueError, match="must not collide"): + ChannelCountStage(action=action, num_channels_key=carrier_key) # type: ignore[arg-type] + + def test_conversion_metadata_cannot_overwrite_resident_audio(self) -> None: + with pytest.raises(ValueError, match="must not collide"): + _convert(duration_key="sample_rate") + def test_converting_keeps_the_rows_selection_would_have_dropped(self, tmp_path: Path) -> None: """Same corpus, same "mono" intent, opposite outcomes -- which is why they cannot share a parameter.""" @@ -442,6 +457,17 @@ def test_an_inverted_range_is_rejected_at_construction(self) -> None: with pytest.raises(ValueError, match="nothing can pass"): SampleRateFilterStage(min_sample_rate=48000, max_sample_rate=16000) + @pytest.mark.parametrize( + "kwargs", + [ + {"audio_filepath_key": "audio", "sample_rate_key": "audio"}, + {"sample_rate_key": "resident", "waveform_key": "resident"}, + ], + ) + def test_audio_input_keys_must_be_distinct(self, kwargs: dict[str, str]) -> None: + with pytest.raises(ValueError, match="must be distinct"): + SampleRateFilterStage(**kwargs) + class TestRowDroppingIsDeclared: """A stage that drops rows has to say ``cardinality="filter"``, because that is the only diff --git a/tests/stages/audio/preprocessing/test_mono_conversion.py b/tests/stages/audio/preprocessing/test_mono_conversion.py index 4b01e0a8c2..b7dff2ff60 100644 --- a/tests/stages/audio/preprocessing/test_mono_conversion.py +++ b/tests/stages/audio/preprocessing/test_mono_conversion.py @@ -15,6 +15,7 @@ from pathlib import Path from unittest.mock import patch +import pytest import torch from nemo_curator.stages.audio.preprocessing.mono_conversion import MonoConversionStage @@ -165,6 +166,24 @@ def test_disk_only_output_omits_the_waveform_key(self, tmp_path: Path) -> None: assert "agent_mono_path" in result.data, "disk-path key must be present when write_to_disk=True" assert "agent_waveform" not in result.data, "tensor must be omitted when keep_waveform_in_task=False" + @pytest.mark.parametrize("residency", ["disk", "wavefrom", ""]) + def test_unknown_residency_is_rejected_at_construction(self, residency: str) -> None: + with pytest.raises(ValueError, match="input_residency"): + MonoConversionStage(input_residency=residency) # type: ignore[arg-type] + + @pytest.mark.parametrize( + "kwargs", + [ + {"duration_key": "waveform"}, + {"is_mono_key": "sample_rate"}, + {"num_samples_key": "audio_filepath"}, + {"write_to_disk": True, "output_audio_filepath_key": "audio_filepath"}, + ], + ) + def test_metadata_and_path_outputs_cannot_overwrite_audio_inputs(self, kwargs: dict[str, object]) -> None: + with pytest.raises(ValueError, match="must not collide"): + MonoConversionStage(**kwargs) + class TestMonoConversionPositionalCompatibility: def test_legacy_positional_call_keeps_strict_sample_rate_third(self) -> None: @@ -179,8 +198,6 @@ def test_legacy_positional_call_keeps_strict_sample_rate_third(self) -> None: def test_agent_added_fields_are_keyword_only(self) -> None: """A former agent field cannot be reached positionally past the legacy slots.""" - import pytest - from nemo_curator.stages.resources import Resources # The six legacy positional slots still accept positionals in their original order. diff --git a/tests/stages/audio/test_common.py b/tests/stages/audio/test_common.py index fbb456e366..e92bac07dc 100644 --- a/tests/stages/audio/test_common.py +++ b/tests/stages/audio/test_common.py @@ -441,9 +441,12 @@ def test_compound_preserve_nested_contract_uses_only_top_level_container_key() - def test_bounded_audio_folder_source_is_not_row_independent() -> None: + from nemo_curator.stages.audio._agent._agent_registry import static_contract + bounded = CreateInitialManifestAudioFolderStage(data_dir="/tmp/x", max_samples=10) # noqa: S108 unbounded = CreateInitialManifestAudioFolderStage(data_dir="/tmp/x") # noqa: S108 + assert static_contract(CreateInitialManifestAudioFolderStage).gates.per_row_independent is False assert bounded.describe().gates.per_row_independent is False assert unbounded.describe().gates.per_row_independent is True @@ -463,6 +466,18 @@ def test_get_audio_duration_validate_input_missing_column() -> None: assert stage.validate_input(AudioTask(data={"text": "hello"})) is False +@pytest.mark.parametrize("residency", ["disk", "wavefrom", ""]) +def test_get_audio_duration_rejects_unknown_residency(residency: str) -> None: + with pytest.raises(ValueError, match="input_residency"): + GetAudioDurationStage(input_residency=residency) # type: ignore[arg-type] + + +@pytest.mark.parametrize("duration_key", ["audio_filepath", "waveform", "sample_rate"]) +def test_get_audio_duration_output_cannot_overwrite_an_audio_input(duration_key: str) -> None: + with pytest.raises(ValueError, match="must not collide"): + GetAudioDurationStage(duration_key=duration_key) + + def test_get_audio_duration_process_batch_raises_on_missing_column() -> None: stage = GetAudioDurationStage() stage.setup() diff --git a/tests/stages/audio/test_create_manifest_audio_folder.py b/tests/stages/audio/test_create_manifest_audio_folder.py index a312121533..d22b432dc3 100644 --- a/tests/stages/audio/test_create_manifest_audio_folder.py +++ b/tests/stages/audio/test_create_manifest_audio_folder.py @@ -19,6 +19,7 @@ import pytest from nemo_curator.stages.audio.common import CreateInitialManifestAudioFolderStage +from nemo_curator.tasks import EmptyTask def _touch(root: str, rel: str) -> None: @@ -95,6 +96,25 @@ def test_requires_data_dir(self) -> None: with pytest.raises(ValueError): # noqa: PT011 CreateInitialManifestAudioFolderStage(data_dir="") + def test_output_keys_must_be_distinct(self, tmp_path) -> None: # noqa: ANN001 + with pytest.raises(ValueError, match="Output keys must be distinct"): + CreateInitialManifestAudioFolderStage( + data_dir=str(tmp_path), + audio_filepath_key="audio", + audio_item_id_key="audio", + ) + + def test_fanout_preserves_parent_provenance(self, tmp_path) -> None: # noqa: ANN001 + root = str(tmp_path) + _touch(root, "a.wav") + parent = EmptyTask(_metadata={"trace": "seed"}, _stage_perf=["upstream"]) + + [child] = CreateInitialManifestAudioFolderStage(data_dir=root).process(parent) + + assert child._metadata == parent._metadata + assert child._stage_perf == parent._stage_perf + assert child._stage_perf is not parent._stage_perf + def test_contract_writes_filepath_and_no_disk_write(self) -> None: c = CreateInitialManifestAudioFolderStage(data_dir="/tmp").describe() # noqa: S108 assert "audio_filepath" in c.writes.data_keys From ccc23ed07cf2768412e7984ba337dddb766a3dde Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Fri, 18 Sep 2026 13:15:34 +0000 Subject: [PATCH 21/25] fix(audio): validate resident sample rates Signed-off-by: shbhawsar --- .../stages/audio/_agent/_residency.py | 39 ++++++++++++++++++- .../test_agent_foundation_regressions.py | 37 ++++++++++++++++++ 2 files changed, 75 insertions(+), 1 deletion(-) diff --git a/nemo_curator/stages/audio/_agent/_residency.py b/nemo_curator/stages/audio/_agent/_residency.py index 8fab4688fe..151c9e62c9 100644 --- a/nemo_curator/stages/audio/_agent/_residency.py +++ b/nemo_curator/stages/audio/_agent/_residency.py @@ -16,10 +16,13 @@ import contextlib import hashlib +import math import os import tempfile +from numbers import Integral, Real from typing import TYPE_CHECKING, Any, Literal +import numpy as np import soundfile as sf import torch @@ -46,6 +49,35 @@ def validate_input_residency(residency: str, *, stage_name: str) -> None: raise ValueError(msg) +def resident_sample_rate(value: Any, *, sample_rate_key: str, stage_name: str) -> int: # noqa: ANN401 + """Return a positive integral resident sample rate without lossy coercion.""" + if torch.is_tensor(value) and value.ndim == 0: + value = value.item() + + rate: int | None = None + if isinstance(value, (bool, np.bool_)): + rate = None + elif isinstance(value, str): + try: + rate = int(value) + except ValueError: + rate = None + elif isinstance(value, Integral): + rate = int(value) + elif isinstance(value, Real): + numeric = float(value) + if math.isfinite(numeric) and numeric.is_integer(): + rate = int(numeric) + + if rate is None or rate <= 0: + msg = ( + f"[{stage_name}] Resident sample rate '{sample_rate_key}' must be a positive, " + f"losslessly integral, non-boolean value; got {value!r}" + ) + raise ValueError(msg) + return rate + + def validate_audio_key_configuration( stage_name: str, *, @@ -399,7 +431,12 @@ def resolve_audio( # noqa: C901, PLR0913 (complexity accepted: policy branches sample_rate = item.get(sample_rate_key) if residency != "file" and waveform is not None: if sample_rate is not None: - return ensure_waveform_2d(waveform), int(sample_rate) + sample_rate = resident_sample_rate( + sample_rate, + sample_rate_key=sample_rate_key, + stage_name="resolve_audio", + ) + return ensure_waveform_2d(waveform), sample_rate if residency == "auto" and infer_sample_rate_from_file: path = item.get(audio_filepath_key) if path: diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index eca859f9fb..cec4d34be9 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -305,6 +305,43 @@ def test_input_residency_validator_rejects_unknown_mode() -> None: validate_input_residency("wavefrom", stage_name="Fixture") +@pytest.mark.parametrize( + "sample_rate", + [ + pytest.param(True, id="bool"), + pytest.param(np.bool_(True), id="numpy-bool"), + pytest.param(0, id="zero"), + pytest.param(-1, id="negative"), + pytest.param(16000.5, id="fractional-float"), + pytest.param("16000.5", id="fractional-string"), + pytest.param(float("nan"), id="nan"), + pytest.param(float("inf"), id="infinity"), + pytest.param(torch.tensor([16000]), id="non-scalar-tensor"), + ], +) +def test_resolve_audio_rejects_invalid_resident_sample_rates(sample_rate: object) -> None: + with pytest.raises(ValueError, match="positive, losslessly integral, non-boolean"): + resolve_audio({"waveform": torch.zeros(8), "sample_rate": sample_rate}) + + +@pytest.mark.parametrize( + "sample_rate", + [ + pytest.param(16000, id="int"), + pytest.param(np.int64(16000), id="numpy-int"), + pytest.param(16000.0, id="integral-float"), + pytest.param("16000", id="numeric-string"), + pytest.param(torch.tensor(16000), id="scalar-tensor"), + ], +) +def test_resolve_audio_preserves_lossless_sample_rate_coercions(sample_rate: object) -> None: + resolved = resolve_audio({"waveform": torch.zeros(8), "sample_rate": sample_rate}) + + assert resolved is not None + assert resolved[1] == 16000 + assert isinstance(resolved[1], int) + + def test_audio_key_validator_rejects_input_role_aliases() -> None: with pytest.raises(ValueError, match="Audio input keys must be distinct"): validate_audio_key_configuration( From e78952463312e4e9219e809923030effeb813767 Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Fri, 18 Sep 2026 14:07:54 +0000 Subject: [PATCH 22/25] fix(audio): harden foundation residency completion Signed-off-by: shbhawsar --- .../stages/audio/_agent/_residency.py | 67 ++++++++++----- nemo_curator/stages/audio/common.py | 81 +++++++++++++++++-- .../audio/preprocessing/sample_rate_filter.py | 17 +++- .../test_agent_foundation_regressions.py | 52 ++++++++++++ .../test_channel_and_rate_stages.py | 35 ++++++++ .../preprocessing/test_mono_conversion.py | 24 ++++++ tests/stages/audio/test_common.py | 73 ++++++++++++++++- 7 files changed, 313 insertions(+), 36 deletions(-) diff --git a/nemo_curator/stages/audio/_agent/_residency.py b/nemo_curator/stages/audio/_agent/_residency.py index 151c9e62c9..61a543f44e 100644 --- a/nemo_curator/stages/audio/_agent/_residency.py +++ b/nemo_curator/stages/audio/_agent/_residency.py @@ -393,7 +393,7 @@ def scoped_file_audio_hydration_writes( # noqa: PLR0913 return conditional -def resolve_audio( # noqa: C901, PLR0913 (complexity accepted: policy branches and keyword-only stage knobs) +def resolve_audio( # noqa: C901, PLR0912, PLR0913 (complexity accepted: policy branches and keyword-only stage knobs) item: dict[str, Any], *, residency: InputResidency = "auto", @@ -414,8 +414,8 @@ def resolve_audio( # noqa: C901, PLR0913 (complexity accepted: policy branches that is missing its sample rate. ``file_audio_hydration="always"`` replaces both resident audio fields after - any selected file load. ``"auto_partial"`` does so only when ``auto`` falls - back with exactly one resident field present. ``"never"`` is the default, + any selected file load. ``"auto_partial"`` does so when ``auto`` falls back + from an incomplete or unusable resident pair. ``"never"`` is the default, preserving file-only and explicit-file consumers. Every update happens only after the loader succeeds, so failures cannot leave a partial pair. @@ -429,14 +429,21 @@ def resolve_audio( # noqa: C901, PLR0913 (complexity accepted: policy branches waveform = item.get(waveform_key) sample_rate = item.get(sample_rate_key) + resident_rate_error: ValueError | None = None if residency != "file" and waveform is not None: if sample_rate is not None: - sample_rate = resident_sample_rate( - sample_rate, - sample_rate_key=sample_rate_key, - stage_name="resolve_audio", - ) - return ensure_waveform_2d(waveform), sample_rate + try: + sample_rate = resident_sample_rate( + sample_rate, + sample_rate_key=sample_rate_key, + stage_name="resolve_audio", + ) + except ValueError as ex: + if residency == "waveform": + raise + resident_rate_error = ex + else: + return ensure_waveform_2d(waveform), sample_rate if residency == "auto" and infer_sample_rate_from_file: path = item.get(audio_filepath_key) if path: @@ -454,9 +461,9 @@ def resolve_audio( # noqa: C901, PLR0913 (complexity accepted: policy branches expanded = os.path.expanduser(str(path)) if os.path.exists(expanded): loaded_waveform, loaded_sample_rate = (loader or load_audio_file)(expanded, mono=mono) - has_partial_pair = (waveform is None) != (sample_rate is None) + has_unusable_pair = (waveform is None) != (sample_rate is None) or resident_rate_error is not None if file_audio_hydration == "always" or ( - file_audio_hydration == "auto_partial" and residency == "auto" and has_partial_pair + file_audio_hydration == "auto_partial" and residency == "auto" and has_unusable_pair ): item.update( { @@ -465,6 +472,8 @@ def resolve_audio( # noqa: C901, PLR0913 (complexity accepted: policy branches } ) return loaded_waveform, loaded_sample_rate + if resident_rate_error is not None: + raise resident_rate_error return None @@ -585,7 +594,7 @@ def write_audio_stable( return path -def resolve_audio_path( # noqa: C901, PLR0913 (keyword-only residency/key knobs mirror stage fields) +def resolve_audio_path( # noqa: C901, PLR0912, PLR0913 (keyword-only residency/key knobs mirror stage fields) item: dict[str, Any], *, residency: InputResidency = "auto", @@ -605,21 +614,33 @@ def resolve_audio_path( # noqa: C901, PLR0913 (keyword-only residency/key knobs caller can delete it after use (see :func:`cleanup_temp_files`). Without ``register_temp`` the caller is responsible for cleanup itself. """ + resident_rate_error: ValueError | None = None if residency != "file": waveform = item.get(waveform_key) sample_rate = item.get(sample_rate_key) if waveform is not None and sample_rate is not None: - fd, tmp = tempfile.mkstemp(suffix=".wav", dir=temp_dir) - os.close(fd) try: - sf.write(tmp, _as_soundfile_array(waveform), int(sample_rate), subtype="FLOAT") - except BaseException: - with contextlib.suppress(OSError): - os.remove(tmp) - raise - if register_temp is not None: - register_temp.append(tmp) - return tmp + sample_rate = resident_sample_rate( + sample_rate, + sample_rate_key=sample_rate_key, + stage_name="resolve_audio_path", + ) + except ValueError as ex: + if residency == "waveform": + raise + resident_rate_error = ex + else: + fd, tmp = tempfile.mkstemp(suffix=".wav", dir=temp_dir) + os.close(fd) + try: + sf.write(tmp, _as_soundfile_array(waveform), sample_rate, subtype="FLOAT") + except BaseException: + with contextlib.suppress(OSError): + os.remove(tmp) + raise + if register_temp is not None: + register_temp.append(tmp) + return tmp if residency == "waveform": return None @@ -648,6 +669,8 @@ def resolve_audio_path( # noqa: C901, PLR0913 (keyword-only residency/key knobs # failure; keep that contract instead of gating on os.path.exists. return local_path + if local_path is None and resident_rate_error is not None: + raise resident_rate_error return local_path diff --git a/nemo_curator/stages/audio/common.py b/nemo_curator/stages/audio/common.py index f34bbd6778..cc1ff5e8b3 100644 --- a/nemo_curator/stages/audio/common.py +++ b/nemo_curator/stages/audio/common.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import json import math import os @@ -118,14 +119,32 @@ def describe(self) -> StageContract: def validate_input(self, task: AudioTask) -> bool: """Require the audio source implied by ``input_residency`` (default: file).""" data = task.data - has_waveform = data.get(self.waveform_key) is not None and data.get(self.sample_rate_key) is not None has_file = self.audio_filepath_key in data - if self.input_residency == "waveform": - return has_waveform if self.input_residency == "file": return has_file + try: + has_waveform = self._resident_rate(data) is not None + except ValueError: + has_waveform = False + if self.input_residency == "waveform": + return has_waveform return has_waveform or has_file # auto + def _resident_rate(self, data: dict[str, Any]) -> int | None: + waveform = data.get(self.waveform_key) + sample_rate = data.get(self.sample_rate_key) + if waveform is None or sample_rate is None: + return None + + # Lazy import avoids the module-level cycle (_residency imports helpers below). + from nemo_curator.stages.audio._agent._residency import resident_sample_rate + + return resident_sample_rate( + sample_rate, + sample_rate_key=self.sample_rate_key, + stage_name=self.name, + ) + def _resolve_duration(self, data: dict[str, Any]) -> float: """Duration from an in-memory waveform (samples / sample_rate) or the file. @@ -133,9 +152,14 @@ def _resolve_duration(self, data: dict[str, Any]) -> float: """ if self.input_residency != "file": waveform = data.get(self.waveform_key) - sr = data.get(self.sample_rate_key) - if waveform is not None and sr is not None and int(sr) > 0: - return ensure_waveform_2d(waveform).shape[-1] / float(sr) + try: + sample_rate = self._resident_rate(data) + except ValueError: + if self.input_residency == "waveform" or self.audio_filepath_key not in data: + raise + else: + if sample_rate is not None: + return ensure_waveform_2d(waveform).shape[-1] / sample_rate if self.input_residency == "waveform": logger.warning(f"Missing '{self.waveform_key}'+'{self.sample_rate_key}' (input_residency='waveform')") return -1.0 @@ -1094,8 +1118,51 @@ def reset_for_retry(self) -> None: self._reset_retry_state() def release_retry_reservation(self) -> None: - """Remove this run's ownership sidecar after successful execution.""" + """Publish completion, then remove this run's retry ownership sidecar.""" self._resolve_output() + owner = self._read_retry_owner() + if owner is None or owner.get("token") != self._reservation_token: + msg = ( + "ManifestCheckpointStage cannot publish completion for a checkpoint " + f"it does not own at {self.output_path!r}" + ) + raise RuntimeError(msg) + if not self._fs.exists(self._path): + msg = f"ManifestCheckpointStage cannot publish completion for missing checkpoint {self.output_path!r}" + raise RuntimeError(msg) + + stat = os.stat(self._path) + identity = (stat.st_dev, stat.st_ino, stat.st_ctime_ns) + recorded_identity = ( + owner.get("st_dev"), + owner.get("st_ino"), + owner.get("st_ctime_ns"), + ) + if identity != recorded_identity or stat.st_size != owner.get("st_size"): + msg = ( + "ManifestCheckpointStage cannot publish completion because the checkpoint " + f"at {self.output_path!r} is no longer its exact reservation" + ) + raise FileExistsError(msg) + + marker_path = f"{self._path}._COMPLETE" + marker = { + "st_dev": stat.st_dev, + "st_ino": stat.st_ino, + "st_ctime_ns": stat.st_ctime_ns, + "st_size": stat.st_size, + } + try: + with self._fs.open(marker_path, "xb") as complete: + complete.write(json.dumps(marker, sort_keys=True).encode("utf-8")) + except FileExistsError as exc: + msg = f"ManifestCheckpointStage refuses to replace completion marker at {self.output_path!r}" + raise FileExistsError(msg) from exc + except OSError: + with contextlib.suppress(OSError): + self._fs.rm(marker_path) + raise + try: self._remove_retry_owner_if_owned() except OSError as exc: diff --git a/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py b/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py index dc3f3684e3..0d461290f3 100644 --- a/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py +++ b/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py @@ -43,7 +43,7 @@ from loguru import logger from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, IOSpec, StageContract -from nemo_curator.stages.audio._agent._residency import validate_audio_key_configuration +from nemo_curator.stages.audio._agent._residency import resident_sample_rate, validate_audio_key_configuration from nemo_curator.stages.base import ProcessingStage from nemo_curator.stages.resources import Resources from nemo_curator.tasks import AudioTask @@ -167,8 +167,19 @@ def _observed_rate(self, task: AudioTask) -> int | None: 16 kHz-only corpus AND re-stamped with the wrong rate. The header read is cheap enough that guessing is never worth it. """ - declared = task.data.get(self.sample_rate_key) - declared = int(declared) if isinstance(declared, (int, float)) and int(declared) > 0 else None + declared_value = task.data.get(self.sample_rate_key) + try: + declared = ( + resident_sample_rate( + declared_value, + sample_rate_key=self.sample_rate_key, + stage_name=self.name, + ) + if declared_value is not None + else None + ) + except ValueError: + declared = None # The VALUE has to be there, not just the column: a row carrying ``waveform=None`` # is no more resident than one with no waveform column at all, and believing it # authenticates exactly the stale metadata this guard exists to distrust. diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index cec4d34be9..4f90e36d84 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -324,6 +324,58 @@ def test_resolve_audio_rejects_invalid_resident_sample_rates(sample_rate: object resolve_audio({"waveform": torch.zeros(8), "sample_rate": sample_rate}) +@pytest.mark.parametrize( + "sample_rate", + [ + pytest.param(True, id="bool"), + pytest.param(0, id="zero"), + pytest.param(-1, id="negative"), + pytest.param(16000.5, id="fractional"), + pytest.param(torch.tensor([16000]), id="non-scalar-tensor"), + ], +) +def test_auto_resolvers_fall_back_from_invalid_resident_rates(tmp_path: Path, sample_rate: object) -> None: + file_path = tmp_path / "valid.wav" + sf.write(file_path, torch.ones(16000).numpy(), 16000) + loaded = torch.ones(1, 16000) + + def loader(_path: str, *, mono: bool) -> tuple[torch.Tensor, int]: + assert mono + return loaded, 16000 + + item = { + "audio_filepath": str(file_path), + "waveform": torch.zeros(1, 8000), + "sample_rate": sample_rate, + } + resolved = resolve_audio( + item, + residency="auto", + loader=loader, + file_audio_hydration="auto_partial", + ) + + assert resolved is not None + assert resolved[0] is loaded + assert resolved[1] == 16000 + assert item["waveform"] is loaded + assert item["sample_rate"] == 16000 + + temporary_paths: list[str] = [] + resolved_path = resolve_audio_path( + { + "audio_filepath": str(file_path), + "waveform": torch.zeros(1, 8000), + "sample_rate": sample_rate, + }, + residency="auto", + temp_dir=str(tmp_path), + register_temp=temporary_paths, + ) + assert resolved_path == str(file_path) + assert temporary_paths == [] + + @pytest.mark.parametrize( "sample_rate", [ diff --git a/tests/stages/audio/preprocessing/test_channel_and_rate_stages.py b/tests/stages/audio/preprocessing/test_channel_and_rate_stages.py index 251489ab97..fcaffcbfed 100644 --- a/tests/stages/audio/preprocessing/test_channel_and_rate_stages.py +++ b/tests/stages/audio/preprocessing/test_channel_and_rate_stages.py @@ -27,6 +27,7 @@ import numpy as np import pytest import soundfile as sf +import torch from nemo_curator.stages.audio.preprocessing import ChannelCountStage, SampleRateFilterStage from nemo_curator.stages.audio.preprocessing import channel_count as cc @@ -398,6 +399,40 @@ def test_a_resident_rate_avoids_touching_disk_entirely(self) -> None: assert result != [] assert result.data["sample_rate"] == 16000 + @pytest.mark.parametrize("sample_rate", [True, 0, -1, 16000.5, torch.tensor([16000])]) + def test_an_invalid_resident_rate_falls_back_to_the_file(self, tmp_path: Path, sample_rate: object) -> None: + path = _wav(tmp_path, rate=48000) + task = AudioTask( + task_id="t", + dataset_name="d", + data={ + "audio_filepath": path, + "sample_rate": sample_rate, + "waveform": np.zeros((1, 16000), dtype="float32"), + }, + ) + + result = SampleRateFilterStage(allowed_sample_rates=[48000]).process(task) + + assert result != [] + assert result.data["sample_rate"] == 48000 + + @pytest.mark.parametrize("sample_rate", [16000, np.int64(16000), 16000.0, "16000", torch.tensor(16000)]) + def test_lossless_resident_rate_representations_are_preserved(self, sample_rate: object) -> None: + task = AudioTask( + task_id="t", + dataset_name="d", + data={ + "sample_rate": sample_rate, + "waveform": np.zeros((1, 16000), dtype="float32"), + }, + ) + + result = SampleRateFilterStage(allowed_sample_rates=[16000]).process(task) + + assert result != [] + assert result.data["sample_rate"] == 16000 + def test_a_manifest_rate_with_no_resident_audio_is_verified_against_the_file(self, tmp_path: Path) -> None: """``sample_rate`` is a standard manifest column, and a stale one used to decide the filter outright: a genuinely 48 kHz file labelled 16000 was KEPT for a 16 kHz-only diff --git a/tests/stages/audio/preprocessing/test_mono_conversion.py b/tests/stages/audio/preprocessing/test_mono_conversion.py index b7dff2ff60..1b82a796ed 100644 --- a/tests/stages/audio/preprocessing/test_mono_conversion.py +++ b/tests/stages/audio/preprocessing/test_mono_conversion.py @@ -145,6 +145,30 @@ def test_auto_residency_uses_the_tensor_and_never_reads_disk(self, tmp_path: Pat assert result.data["agent_waveform"].shape[0] == 1, "the stereo tensor was mixed down in memory" assert "agent_mono_path" not in result.data, "disk-path key must be absent when write_to_disk=False" + @pytest.mark.parametrize("sample_rate", [True, 0, -1, 16000.5, torch.tensor([16000])]) + def test_auto_residency_uses_the_file_when_the_resident_rate_is_invalid( + self, tmp_path: Path, sample_rate: object + ) -> None: + path = tmp_path / "valid.wav" + path.touch() + loaded = torch.ones(1, 16000) + stage = MonoConversionStage(output_sample_rate=16000, input_residency="auto") + task = AudioTask( + data={ + "audio_filepath": str(path), + "waveform": torch.zeros(1, 8000), + "sample_rate": sample_rate, + } + ) + + with patch(MOCK_TARGET, return_value=(loaded, 16000)) as loader: + result = stage.process(task) + + assert isinstance(result, AudioTask) + assert loader.call_count == 1 + assert result.data["waveform"] is loaded + assert result.data["sample_rate"] == 16000 + def test_disk_only_output_omits_the_waveform_key(self, tmp_path: Path) -> None: """``keep_waveform_in_task=False`` must drop the tensor rather than leave it stale.""" stage = MonoConversionStage( diff --git a/tests/stages/audio/test_common.py b/tests/stages/audio/test_common.py index e92bac07dc..5b9fbf8dee 100644 --- a/tests/stages/audio/test_common.py +++ b/tests/stages/audio/test_common.py @@ -512,17 +512,49 @@ def test_get_audio_duration_error_sets_minus_one(tmp_path: Path) -> None: def test_get_audio_duration_waveform_residency() -> None: """input_residency='waveform' computes duration from samples/sample_rate (no file).""" - import torch - stage = GetAudioDurationStage(input_residency="waveform") stage.setup() result = stage.process(AudioTask(data={"waveform": torch.zeros(1, 16000 * 3), "sample_rate": 16000})) assert result.data["duration"] == 3.0 -def test_get_audio_duration_auto_prefers_waveform() -> None: - import torch +@pytest.mark.parametrize("sample_rate", [True, 0, -1, 16000.5, torch.tensor([16000])]) +def test_get_audio_duration_rejects_invalid_resident_rates(sample_rate: object) -> None: + stage = GetAudioDurationStage(input_residency="waveform") + task = AudioTask(data={"waveform": torch.zeros(1, 16000), "sample_rate": sample_rate}) + + assert not stage.validate_input(task) + with pytest.raises(ValueError, match="positive, losslessly integral, non-boolean"): + stage.process(task) + + +def test_get_audio_duration_auto_falls_back_from_invalid_resident_rate(tmp_path: Path) -> None: + path = tmp_path / "valid.wav" + with mock.patch("soundfile.info", return_value=mock.Mock(frames=32000, samplerate=16000)): + stage = GetAudioDurationStage(input_residency="auto") + stage.setup() + task = AudioTask( + data={ + "audio_filepath": str(path), + "waveform": torch.zeros(1, 8000), + "sample_rate": True, + } + ) + + assert stage.validate_input(task) + assert stage.process(task).data["duration"] == 2.0 + +@pytest.mark.parametrize("sample_rate", [16000, np.int64(16000), 16000.0, "16000", torch.tensor(16000)]) +def test_get_audio_duration_accepts_lossless_resident_rates(sample_rate: object) -> None: + stage = GetAudioDurationStage(input_residency="waveform") + task = AudioTask(data={"waveform": torch.zeros(1, 16000), "sample_rate": sample_rate}) + + assert stage.validate_input(task) + assert stage.process(task).data["duration"] == 1.0 + + +def test_get_audio_duration_auto_prefers_waveform() -> None: stage = GetAudioDurationStage(input_residency="auto") stage.setup() result = stage.process(AudioTask(data={"waveform": torch.zeros(1, 16000), "sample_rate": 16000})) @@ -961,6 +993,39 @@ def test_retry_reset_refuses_completed_checkpoint(self, tmp_path: Path) -> None: assert out.read_bytes() == before + def test_successful_release_publishes_completion_before_removing_retry_owner(self, tmp_path: Path) -> None: + out = tmp_path / "checkpoint.jsonl" + owner_path = Path(f"{out}._RETRY_OWNER") + marker_path = Path(f"{out}._COMPLETE") + checkpoint = ManifestCheckpointStage(output_path=str(out)) + checkpoint.setup() + checkpoint.process(AudioTask(data={"retained": True})) + + checkpoint.release_retry_reservation() + + assert out.read_text(encoding="utf-8") == '{"retained": true}\n' + assert marker_path.exists() + assert not owner_path.exists() + marker = json.loads(marker_path.read_text(encoding="utf-8")) + assert marker["st_size"] == out.stat().st_size + with pytest.raises(FileExistsError, match="completion marker"): + checkpoint.reset_for_retry() + + def test_successful_release_refuses_to_complete_a_replaced_checkpoint(self, tmp_path: Path) -> None: + out = tmp_path / "checkpoint.jsonl" + checkpoint = ManifestCheckpointStage(output_path=str(out)) + checkpoint.setup() + checkpoint.process(AudioTask(data={"attempt": 1})) + out.unlink() + out.write_text("replacement\n", encoding="utf-8") + + with pytest.raises(FileExistsError, match="no longer its exact reservation"): + checkpoint.release_retry_reservation() + + assert out.read_text(encoding="utf-8") == "replacement\n" + assert not Path(f"{out}._COMPLETE").exists() + assert Path(f"{out}._RETRY_OWNER").exists() + def test_retry_reset_refuses_preexisting_unowned_checkpoint( self, tmp_path: Path, From 0dad5b94751ff1c9615f062904399d0839d37540 Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Fri, 18 Sep 2026 15:24:12 +0000 Subject: [PATCH 23/25] fix(audio): close foundation review gaps Signed-off-by: shbhawsar --- nemo_curator/pipeline/pipeline.py | 3 + nemo_curator/stages/audio/AGENTS.md | 24 +++----- nemo_curator/stages/audio/AGENT_READY.md | 60 +++++++++---------- nemo_curator/stages/audio/agent.py | 6 +- nemo_curator/stages/audio/common.py | 36 ++++++++--- nemo_curator/stages/base.py | 8 +++ tests/pipelines/test_pipelines.py | 22 +++++++ .../test_agent_foundation_regressions.py | 32 ++-------- tests/stages/audio/test_common.py | 14 ++++- .../test_create_manifest_audio_folder.py | 13 ++++ 10 files changed, 132 insertions(+), 86 deletions(-) diff --git a/nemo_curator/pipeline/pipeline.py b/nemo_curator/pipeline/pipeline.py index 5ff5e2fe83..6b0a77b35d 100644 --- a/nemo_curator/pipeline/pipeline.py +++ b/nemo_curator/pipeline/pipeline.py @@ -339,6 +339,9 @@ def run( # noqa: C901, PLR0912 else: result = self._run_with_resumability(executor, initial_tasks, checkpoint_path) + for stage in self.stages: + stage.finalize() + if completion_manifest is not None: if failed_task_manifest_exists(): logger.warning( diff --git a/nemo_curator/stages/audio/AGENTS.md b/nemo_curator/stages/audio/AGENTS.md index d7d796a3b1..435b0c7367 100644 --- a/nemo_curator/stages/audio/AGENTS.md +++ b/nemo_curator/stages/audio/AGENTS.md @@ -10,18 +10,15 @@ Two different jobs happen in this directory, and they have opposite rules about Stage source is shared library code. Never change a stage, a threshold, a filter or windowing logic to make one user's dataset produce the output they hoped for — diagnose -from the data instead. The full guardrail, and what to do when output comes back empty, -is in `nemo_curator/audio_agent/AGENTS.md`. Drive the work through the -`nemo_curator.audio_agent` verbs rather than reading stage source to answer what a stage -reads or writes: `describe --params '{...}'` is the sanctioned answer, and -`producers ` says who wrote a key. +from the data instead. Inspect stage contracts through the public +`nemo_curator.stages.audio.agent` facade rather than reading stage source to answer what +a stage reads or writes: `agent.describe_stage("MyStage")` is the sanctioned answer, and +`agent.find_producers("role")` says which stages can produce a semantic role. ## If you are authoring or fixing a stage, start here Read `AGENT_READY.md` in this directory. It is the authoritative checklist and it is -maintained with the framework; work from it rather than from memory. The -`audio-stage-authoring` skill (`nemo_curator/audio_agent/skills/audio-stage-authoring/SKILL.md`) -is the procedure that drives it. +maintained with the framework; work from it rather than from memory. **Golden rule: every new knob defaults to today's behavior.** Agent-readiness is a declaration layer over working code. If a change alters what an existing pipeline @@ -35,12 +32,9 @@ The mechanical contract is three things, each detailed in `AGENT_READY.md`: literals in `process()`, or the key is invisible to the agent and cannot be remapped. 3. Add `assert_agent_ready(MyStage(...), fixture_factory=...)` as a test. -Then give the stage a capability card under -`nemo_curator/audio_agent/knowledge/cards/`, documenting what each externally consumed -output *means* — roles prove two stages can connect, only the card lets the host judge -whether connecting them serves the user's intent. Schema: -`nemo_curator/audio_agent/knowledge/CARD_SCHEMA.md`. `card_conformance.audit()` must come -back with zero violations. +Document what each externally consumed output means in the stage docstring and tests. +Roles prove that two stages can connect; clear semantic documentation lets reviewers +judge whether connecting them serves the user's intent. Follow the repo's existing stage conventions while you do it: `.cursor/rules/processing-stage-patterns.mdc` and @@ -56,7 +50,7 @@ planner treats it as ground truth. ```bash .venv/bin/python -m pytest tests/stages/audio -m "not gpu" -q -.venv/bin/python -m nemo_curator.audio_agent describe MyStage --params '{...}' +.venv/bin/python -c 'from nemo_curator.stages.audio import agent; print(agent.describe_stage("MyStage").to_dict())' ``` If `describe` does not match what the code actually touches for those params, the contract diff --git a/nemo_curator/stages/audio/AGENT_READY.md b/nemo_curator/stages/audio/AGENT_READY.md index 77cdd835cf..f975a08f1c 100644 --- a/nemo_curator/stages/audio/AGENT_READY.md +++ b/nemo_curator/stages/audio/AGENT_READY.md @@ -18,10 +18,10 @@ the framework auto-derives the rest, and one test tells you if anything is missi literals in `process()`). 3. Add one test: **`assert_agent_ready(MyStage(...), fixture_factory=...)`**. -Those three items make the stage mechanically composable. Also update its capability card -with the meaning of externally consumed outputs (especially filterable fields and anything -crossing a fan-out/aggregation boundary). `assert_agent_ready` can prove keys and -cardinality; it cannot prove that a reasoner will interpret a value correctly. +Those three items make the stage mechanically composable. Also document the meaning of +externally consumed outputs, especially filterable fields and anything crossing a +fan-out/aggregation boundary. `assert_agent_ready` can prove keys and cardinality; it +cannot prove that a reasoner will interpret a value correctly. --- @@ -127,8 +127,8 @@ seeds tensor residency for the JSON-sink gate. Set `requires_keys` on a exists upstream (in the write's own scope): the planner then ignores the branch -- for both purposes -- on inputs that cannot reach it. File hydration that only *replaces* an incomplete resident waveform/sample-rate pair is the canonical -case: without the hint, `UTMOSFilterStage() -> ManifestWriterStage` on a plain -file manifest would be refused for a tensor the runtime never introduces. +case: without the hint, a file-backed scorer followed by `ManifestWriterStage` +could be refused for a tensor the runtime never introduces. Use `metadata_writes` for a conditional `task._metadata` output. Unconditional metadata inputs/outputs remain declared through @@ -193,12 +193,11 @@ Three more rules: harmless. `True` when you were not silently produces rows a full run would never have produced, and republishes them as the corpus's reusable result. **When unsure, declare `False`.** -The two `CreateInitialManifest*` sources with `max_samples` are a **deliberate exception** to that -last rule, not an example of it. `max_samples` truncates the *sorted* listing, so a delta over a -bounded source can select files a full run would not have — yet both declare a flat `True`, because -the conditional `False` denied reuse to the configuration nearly everyone runs (ReadSpeech defaults -to 5000). The limitation is recorded at each declaration. Do not copy this into a new stage; if you -find yourself wanting to, declare `False` and raise it instead. +`CreateInitialManifestAudioFolderStage` is the reference for a conditionally safe +source. Its unbounded scan is per-row independent, but `max_samples` truncates the +sorted listing, so adding an earlier path can change which rows are emitted. It +therefore declares `per_row_independent=False` for bounded instances and `True` only +for unbounded ones. Follow that per-instance pattern for any bounded source. ## What is AUTO-DERIVED — do NOT hand-write these @@ -264,9 +263,9 @@ resolved automatically from your `*_key` **field name** via `nemo_curator/stages Roles answer “can these stages connect?” They do not answer “what does this value mean here?” A pipeline can connect perfectly and still apply a valid -filter to the wrong entity. Put that semantic knowledge in the capability card, -where the host LLM can reason over it; do not add a per-field rule or a hard -`field_scope` ontology to the Python core. +filter to the wrong entity. Put that semantic knowledge in the stage's durable +documentation, where a host can reason over it; do not add a per-field rule or a +hard `field_scope` ontology to the Python core. For each output that another stage may select, aggregate, compare or filter, document: @@ -283,18 +282,15 @@ document: - **counterexample** — at least one plausible but wrong interpretation and its pipeline consequence. -Use the card's optional `semantic_facts` mapping for structured prose, or -`notes`/`caveats` when the fact spans several outputs. For example, if a -speaker-separation stage computes the original clip's detected-speaker count -once and copies it to every per-speaker child, say so explicitly: filtering that -child field to `== 1` selects children whose **parent source** had one detected -speaker; it does not test whether each already-separated child track is -single-speaker. +For example, if a speaker-separation stage computes the original clip's +detected-speaker count once and copies it to every per-speaker child, say so +explicitly: filtering that child field to `== 1` selects children whose **parent +source** had one detected speaker; it does not test whether each already-separated +child track is single-speaker. Only document facts grounded in code, a measured run or an authoritative model -source, and mark their honesty tier in the card's `verified` block. Missing -meaning should remain an explicit TODO; the host must ask rather than invent it. -See `nemo_curator/audio_agent/knowledge/CARD_SCHEMA.md`. +source. Missing meaning should remain an explicit TODO; the host must ask rather +than invent it. ## Discovery — how the agent finds your stage @@ -326,9 +322,9 @@ def test_my_stage_is_agent_ready(tmp_path): assert_agent_ready(MyStage(), fixture, expected_cardinality="1:1", available_keys={"audio_filepath"}) ``` -For GPU/model stages, reuse the existing fake-model/stub setup (see -`tests/stages/audio/test_agent_simulation_pipelines.py`) so the test needs no GPU. You don't need to -memorize the rules — if the test passes, the contract is honest. +For GPU/model stages, keep the conformance test CPU-only by reusing the fake-model +or stub setup in that stage's test module. You don't need to memorize the mechanical +rules — if the test passes, the contract is internally consistent. --- @@ -337,8 +333,8 @@ memorize the rules — if the test passes, the contract is honest. - [ ] `AgentReady` + `describe()` with `reads`, `writes`, `cardinality`, honest `gates` - [ ] every read/written `task.data` key is a `*_key` constructor field (no bare literals) - [ ] new `*_key` concepts have a `_roles.KEY_ROLES` entry (or `INTERNAL_KEY_FIELDS`) -- [ ] capability card explains each externally consumed output's meaning, unit, - provenance, scope/granularity, propagation and a counterexample +- [ ] durable stage documentation explains each externally consumed output's meaning, + unit, provenance, scope/granularity, propagation and a counterexample - [ ] new `AudioTask`s preserve `_metadata` and `list(_stage_perf)` (manual — not covered by `assert_agent_ready`) - [ ] if you override `process_batch`, write to disk, or set `lifecycle_side_effects` — decided `gates.per_row_independent` (`True`/`False`, per instance if conditional); otherwise left it @@ -346,5 +342,5 @@ memorize the rules — if the test passes, the contract is honest. - [ ] `assert_agent_ready(...)` test added and green - [ ] defaults unchanged → existing pipelines behave exactly as before -Auto-derivation handles params/roles/dispatch/description; the card supplies meaning only the -stage author knows. Neither documentation step changes runtime defaults. +Auto-derivation handles params/roles/dispatch/description; stage documentation supplies +meaning only the author knows. Neither documentation step changes runtime defaults. diff --git a/nemo_curator/stages/audio/agent.py b/nemo_curator/stages/audio/agent.py index eba6830cbc..106a98aced 100644 --- a/nemo_curator/stages/audio/agent.py +++ b/nemo_curator/stages/audio/agent.py @@ -31,8 +31,8 @@ agent.find_consumers("pred_text") # [] from find_producers => unproducible # 3. CONFIGURE + BUILD - cls = agent.get_agent_ready_stage_class("UTMOSFilterStage") - stage = cls(mos_threshold=3.5) # params_schema documents the knobs + cls = agent.get_agent_ready_stage_class("SampleRateFilterStage") + stage = cls(min_sample_rate=16_000) # params_schema documents the knobs # 4. VALIDATE (before ever running) report = agent.validate_pipeline([stage, ...], initial_keys={"audio_filepath", "text"}) @@ -44,7 +44,7 @@ ``static_params_and_hints``; pass an instance to ``build_contract`` for the configured dynamic contract with resolved I/O/key values. ``StageContract.to_dict()`` is JSON-safe by construction. Open-ended intent fit -is reviewed by the host LLM over ``audio_agent.validate(...).semantic_review``. +remains a review decision for the host using the returned contracts and report. """ from __future__ import annotations diff --git a/nemo_curator/stages/audio/common.py b/nemo_curator/stages/audio/common.py index cc1ff5e8b3..41f9cbbe95 100644 --- a/nemo_curator/stages/audio/common.py +++ b/nemo_curator/stages/audio/common.py @@ -18,6 +18,7 @@ import os import time import uuid +from collections import Counter from collections.abc import Mapping from dataclasses import dataclass, field from operator import eq, ge, gt, le, lt, ne @@ -759,11 +760,10 @@ def num_workers(self) -> int | None: return 1 @staticmethod - def _item_id(relative_path: str) -> str: - """Flatten a relative stem without confusing separators with filename text.""" - stem = os.path.splitext(relative_path)[0] + def _encode_item_id_path(relative_path: str) -> str: + """Flatten a relative path without confusing separators with filename text.""" encoded_components: list[str] = [] - for component in stem.split(os.sep): + for component in relative_path.split(os.sep): encoded: list[str] = [] last = len(component) - 1 for index, char in enumerate(component): @@ -782,6 +782,15 @@ def _item_id(relative_path: str) -> str: encoded_components.append("".join(encoded)) return "__".join(encoded_components) + @classmethod + def _item_id(cls, relative_path: str, *, include_extension: bool = False) -> str: + """Return the legacy stem id, adding a reserved extension suffix on collisions.""" + stem, extension = os.path.splitext(relative_path) + item_id = cls._encode_item_id_path(stem) + if include_extension: + item_id = f"{item_id}~e{cls._encode_item_id_path(extension[1:])}" + return item_id + def _collect_audio_files(self) -> list[str]: exts = tuple((e if e.startswith(".") else f".{e}").lower() for e in self.extensions) if not os.path.isdir(self.data_dir): @@ -817,15 +826,18 @@ def process(self, task: EmptyTask | None) -> list[AudioTask]: if not paths: logger.warning(f"[{self.name}] no audio files {self.extensions} under {self.data_dir}") return [] + root = os.path.abspath(self.data_dir) + relative_paths = [os.path.relpath(os.path.abspath(path), root) for path in paths] + legacy_ids = [self._item_id(relative_path) for relative_path in relative_paths] + id_counts = Counter(legacy_ids) tasks: list[AudioTask] = [] - for path in paths: + for path, rel, legacy_id in zip(paths, relative_paths, legacy_ids, strict=True): abspath = os.path.abspath(path) # Relpath, not basename: ``recursive`` defaults True and speaker-per-folder is the # standard layout, so a basename id gives spk1/utt1.wav and spk2/utt1.wav the same # id -- and downstream that id becomes an output filename. A flat corpus is # unaffected unless its name needs escaping to remain distinct from a path separator. - rel = os.path.relpath(abspath, os.path.abspath(self.data_dir)) - item_id = self._item_id(rel) + item_id = self._item_id(rel, include_extension=id_counts[legacy_id] > 1) tasks.append( AudioTask( dataset_name="local-audio-folder", @@ -1173,6 +1185,16 @@ def release_retry_reservation(self) -> None: self._reservation_owned = False self._reservation_identity = None + def finalize(self) -> None: + """Publish the completion marker after successful pipeline execution.""" + self._resolve_output() + if self._read_retry_owner() is None and not self._fs.exists(self._path): + # A backend may never construct a worker when the upstream dataset + # is empty. A successful empty checkpoint is still a complete, + # reusable artifact, so reserve it on the driver before publishing. + self.setup() + self.release_retry_reservation() + def _retry_owner_path(self) -> str: return f"{self._path}._RETRY_OWNER" diff --git a/nemo_curator/stages/base.py b/nemo_curator/stages/base.py index 63d340cca9..f3647e58fa 100644 --- a/nemo_curator/stages/base.py +++ b/nemo_curator/stages/base.py @@ -374,6 +374,14 @@ def teardown(self) -> None: Override this method to perform any cleanup. """ + def finalize(self) -> None: + """Finalize driver-owned state after the whole pipeline succeeds. + + Unlike ``teardown()``, which runs on backend workers and may also run + during failure cleanup, this hook is called once by ``Pipeline.run`` + only after the executor has returned successfully. + """ + def supports_batch_processing(self) -> bool: """Whether this stage supports vectorized batch processing. This is automatically determined by checking if the stage has diff --git a/tests/pipelines/test_pipelines.py b/tests/pipelines/test_pipelines.py index ca52781c5d..a31f18b123 100644 --- a/tests/pipelines/test_pipelines.py +++ b/tests/pipelines/test_pipelines.py @@ -62,6 +62,28 @@ def test_pipeline_uses_xenna_executor_by_default(): mock_xenna_instance.execute.assert_called_once() +def test_pipeline_finalizes_stages_only_after_successful_execution() -> None: + stage = _NoopStage() + stage.finalize = Mock() + executor = Mock() + + Pipeline(name="test", stages=[stage]).run(executor=executor) + + stage.finalize.assert_called_once_with() + + +def test_pipeline_does_not_finalize_stages_after_failed_execution() -> None: + stage = _NoopStage() + stage.finalize = Mock() + executor = Mock() + executor.execute.side_effect = RuntimeError("failed") + + with pytest.raises(RuntimeError, match="failed"): + Pipeline(name="test", stages=[stage]).run(executor=executor) + + stage.finalize.assert_not_called() + + def test_logs_info_when_ray_serve_active_with_gpu_stages_non_xenna() -> None: """Non-Xenna executors log an info message when Serve is active with GPU stages.""" gpu_stage = Mock(spec=ProcessingStage) diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index 4f90e36d84..ca971acf38 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -1074,14 +1074,12 @@ def test_a_read_met_only_by_a_conditional_write_is_a_warning_not_an_error() -> N assert any(issue.code == "unsatisfied_reads" for issue in missing.issues) -def test_shipped_metric_then_selector_chain_composes() -> None: - """The fleurs recipe shape: pairwise WER followed by a threshold on its (conditional) output.""" - from nemo_curator.stages.audio.metrics.wer import GetPairwiseWerStage - +def test_conditional_metric_then_selector_chain_composes() -> None: + """A metric followed by a threshold composes when the metric output is conditional.""" report = validate_pipeline( - [GetPairwiseWerStage(), PreserveByValueStage("wer_pct", 25.0, "le")], - initial_keys={"audio_filepath", "text", "pred_text"}, - initial_roles={"audio_filepath", "text", "pred_text"}, + [_ConditionalScoreProducer(), PreserveByValueStage("score", 25.0, "le")], + initial_keys={"audio_filepath"}, + initial_roles={"audio_filepath"}, ) assert report.ok, report.summary() assert report.keys_ok @@ -1110,23 +1108,3 @@ def test_conditional_tensor_writes_seed_residency_only_when_reachable(tmp_path: [_ReachabilityGatedTensorProducer(), _ReachabilityGatedTensorProducer(), sink], ) assert self_enabling.ok, self_enabling.summary() - - -def test_default_auto_scorers_do_not_fear_a_tensor_on_a_file_manifest(tmp_path: Path) -> None: - """``auto_partial`` hydration only REPLACES an incomplete resident pair; a file-only row never gains one.""" - from nemo_curator.stages.audio.filtering.sigmos import SIGMOSFilterStage - from nemo_curator.stages.audio.filtering.utmos import UTMOSFilterStage - - sink = ManifestWriterStage(output_path=str(tmp_path / "out.jsonl")) - for scorer in (UTMOSFilterStage(), SIGMOSFilterStage()): - report = validate_pipeline([scorer, sink]) - assert report.ok, report.summary() - - # With a resident sample_rate and no waveform the runtime DOES inject the decoded pair, - # so the refusal there is a true positive and must stay. - resident_rate = validate_pipeline( - [UTMOSFilterStage(), sink], - initial_keys={"audio_filepath", "sample_rate"}, - initial_roles={"audio_filepath", "sample_rate"}, - ) - assert any(issue.code == "tensor_into_sink" for issue in resident_rate.issues) diff --git a/tests/stages/audio/test_common.py b/tests/stages/audio/test_common.py index 5b9fbf8dee..39feb0015b 100644 --- a/tests/stages/audio/test_common.py +++ b/tests/stages/audio/test_common.py @@ -993,7 +993,7 @@ def test_retry_reset_refuses_completed_checkpoint(self, tmp_path: Path) -> None: assert out.read_bytes() == before - def test_successful_release_publishes_completion_before_removing_retry_owner(self, tmp_path: Path) -> None: + def test_successful_finalize_publishes_completion_before_removing_retry_owner(self, tmp_path: Path) -> None: out = tmp_path / "checkpoint.jsonl" owner_path = Path(f"{out}._RETRY_OWNER") marker_path = Path(f"{out}._COMPLETE") @@ -1001,7 +1001,7 @@ def test_successful_release_publishes_completion_before_removing_retry_owner(sel checkpoint.setup() checkpoint.process(AudioTask(data={"retained": True})) - checkpoint.release_retry_reservation() + checkpoint.finalize() assert out.read_text(encoding="utf-8") == '{"retained": true}\n' assert marker_path.exists() @@ -1011,6 +1011,16 @@ def test_successful_release_publishes_completion_before_removing_retry_owner(sel with pytest.raises(FileExistsError, match="completion marker"): checkpoint.reset_for_retry() + def test_successful_finalize_publishes_an_empty_checkpoint(self, tmp_path: Path) -> None: + out = tmp_path / "checkpoint.jsonl" + checkpoint = ManifestCheckpointStage(output_path=str(out)) + + checkpoint.finalize() + + assert out.read_bytes() == b"" + assert Path(f"{out}._COMPLETE").exists() + assert not Path(f"{out}._RETRY_OWNER").exists() + def test_successful_release_refuses_to_complete_a_replaced_checkpoint(self, tmp_path: Path) -> None: out = tmp_path / "checkpoint.jsonl" checkpoint = ManifestCheckpointStage(output_path=str(out)) diff --git a/tests/stages/audio/test_create_manifest_audio_folder.py b/tests/stages/audio/test_create_manifest_audio_folder.py index d22b432dc3..183a881b4c 100644 --- a/tests/stages/audio/test_create_manifest_audio_folder.py +++ b/tests/stages/audio/test_create_manifest_audio_folder.py @@ -73,6 +73,19 @@ def test_a_flat_folder_keeps_the_plain_ids_it_always_had(self, tmp_path) -> None assert sorted(t.data["audio_item_id"] for t in tasks) == ["a", "b"] + def test_same_stem_with_different_extensions_gets_two_ids(self, tmp_path) -> None: # noqa: ANN001 + root = str(tmp_path) + for rel in ["a.wav", "a.flac"]: + _touch(root, rel) + + tasks = CreateInitialManifestAudioFolderStage(data_dir=root).process(None) + ids_by_name = { + os.path.basename(task.data["audio_filepath"]): task.data["audio_item_id"] for task in tasks + } + + assert ids_by_name == {"a.flac": "a~eflac", "a.wav": "a~ewav"} + assert len(set(ids_by_name.values())) == 2 + def test_non_recursive_and_max_samples(self, tmp_path) -> None: # noqa: ANN001 root = str(tmp_path) for rel in ["a.wav", "b.wav", "sub/c.wav"]: From 0e9634a4c7d5ecf11b42a9b869983f2fc292f485 Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Sat, 19 Sep 2026 13:13:35 +0000 Subject: [PATCH 24/25] fix(audio): normalize PCM and validate sample rates Signed-off-by: shbhawsar --- .../audio/preprocessing/channel_count.py | 5 +- .../audio/preprocessing/mono_conversion.py | 7 +-- .../audio/preprocessing/sample_rate_filter.py | 12 +++++ .../test_agent_foundation_regressions.py | 49 +++++++++++++++++++ 4 files changed, 68 insertions(+), 5 deletions(-) diff --git a/nemo_curator/stages/audio/preprocessing/channel_count.py b/nemo_curator/stages/audio/preprocessing/channel_count.py index 322414112b..4402d23cde 100644 --- a/nemo_curator/stages/audio/preprocessing/channel_count.py +++ b/nemo_curator/stages/audio/preprocessing/channel_count.py @@ -47,6 +47,7 @@ InputResidency, accepts_for_residency, drop_resident_audio, + normalize_audio_waveform, produce_audio_filepath, reject_sinkless_conversion, residency_read_specs, @@ -520,7 +521,7 @@ def _convert_row( # noqa: C901 (complexity accepted: residency x sink x update try: waveform, sample_rate = resolved - waveform = ensure_waveform_2d(waveform) + waveform = normalize_audio_waveform(waveform, stage_name=self.name, mono=False) if sample_rate <= 0: logger.error(f"Invalid sample rate ({sample_rate}) in audio input") @@ -556,7 +557,7 @@ def _convert_row( # noqa: C901 (complexity accepted: residency x sink x update sample_rate_key=self.sample_rate_key, ) - except (OSError, RuntimeError) as e: + except (OSError, RuntimeError, TypeError, ValueError) as e: logger.error(f"Error processing audio input: {e}") return [] else: diff --git a/nemo_curator/stages/audio/preprocessing/mono_conversion.py b/nemo_curator/stages/audio/preprocessing/mono_conversion.py index 0a8743efd7..970d660bc6 100755 --- a/nemo_curator/stages/audio/preprocessing/mono_conversion.py +++ b/nemo_curator/stages/audio/preprocessing/mono_conversion.py @@ -36,6 +36,7 @@ from nemo_curator.stages.audio._agent._residency import ( InputResidency, drop_resident_audio, + normalize_audio_waveform, produce_audio_filepath, reject_sinkless_conversion, residency_read_specs, @@ -44,7 +45,7 @@ validate_input_residency, write_audio_stable, ) -from nemo_curator.stages.audio.common import ensure_waveform_2d, load_audio_file +from nemo_curator.stages.audio.common import load_audio_file from nemo_curator.stages.base import ProcessingStage from nemo_curator.stages.resources import Resources from nemo_curator.tasks import AudioTask @@ -241,7 +242,7 @@ def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: # noqa: C901 try: waveform, sample_rate = resolved - waveform = ensure_waveform_2d(waveform) + waveform = normalize_audio_waveform(waveform, stage_name=self.name, mono=False) if sample_rate <= 0: logger.error(f"Invalid sample rate ({sample_rate}) in audio input") @@ -286,7 +287,7 @@ def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: # noqa: C901 sample_rate_key=self.sample_rate_key, ) - except (OSError, RuntimeError) as e: + except (OSError, RuntimeError, TypeError, ValueError) as e: logger.error(f"Error processing audio input: {e}") return [] else: diff --git a/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py b/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py index 0d461290f3..0920a35567 100644 --- a/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py +++ b/nemo_curator/stages/audio/preprocessing/sample_rate_filter.py @@ -107,9 +107,21 @@ def __post_init__(self): # legitimate read/write of that same key as a collision. output_keys={}, ) + if self.allowed_sample_rates is not None and not isinstance(self.allowed_sample_rates, list): + msg = "allowed_sample_rates must be a list of positive whole-number rates, or None" + raise TypeError(msg) if self.allowed_sample_rates is not None and not self.allowed_sample_rates: msg = "allowed_sample_rates must name at least one rate, or be None for no constraint" raise ValueError(msg) + for name in ("allowed_sample_rates", "min_sample_rate", "max_sample_rate"): + value = getattr(self, name) + rates = value if isinstance(value, list) else [value] + for rate in rates: + if rate is None: + continue + if isinstance(rate, bool) or not isinstance(rate, int) or rate <= 0: + msg = f"{name} must contain positive whole-number rates, got {value!r}" + raise ValueError(msg) low, high = self.min_sample_rate, self.max_sample_rate if low is not None and high is not None and low > high: msg = f"min_sample_rate ({low}) is above max_sample_rate ({high}), so nothing can pass" diff --git a/tests/stages/audio/_agent/test_agent_foundation_regressions.py b/tests/stages/audio/_agent/test_agent_foundation_regressions.py index ca971acf38..8f24f00ba5 100644 --- a/tests/stages/audio/_agent/test_agent_foundation_regressions.py +++ b/tests/stages/audio/_agent/test_agent_foundation_regressions.py @@ -882,6 +882,55 @@ def test_a_null_waveform_does_not_authenticate_a_stale_sample_rate(tmp_path: Pat assert stage._observed_rate(resident) == 16000 +@pytest.mark.parametrize( + "waveform", + [ + torch.tensor([[32767, -32768], [0, 16384]], dtype=torch.int16), + np.array([[2147483647, -2147483648], [0, 1073741824]], dtype=np.int32), + ], + ids=["torch_pcm16", "numpy_pcm32"], +) +@pytest.mark.parametrize("stage_kind", ["mono", "channel_count"]) +def test_resident_pcm_stereo_is_normalized_before_downmix(waveform: object, stage_kind: str) -> None: + if stage_kind == "mono": + stage = MonoConversionStage( + output_sample_rate=16000, + input_residency="waveform", + strict_sample_rate=True, + ) + else: + stage = ChannelCountStage( + action="convert", + target_channels=1, + input_residency="waveform", + ) + task = AudioTask(dataset_name="d", data={"waveform": waveform, "sample_rate": 16000}) + + result = stage.process(task) + + assert isinstance(result, AudioTask) + assert result.data["waveform"].dtype == torch.float32 + assert result.data["waveform"].shape == (1, 2) + assert torch.isfinite(result.data["waveform"]).all() + + +@pytest.mark.parametrize( + ("field_name", "value", "error_type"), + [ + ("allowed_sample_rates", "16000", TypeError), + ("allowed_sample_rates", [16000, True], ValueError), + ("allowed_sample_rates", [16000.5], ValueError), + ("allowed_sample_rates", [0], ValueError), + ("min_sample_rate", True, ValueError), + ("min_sample_rate", -1, ValueError), + ("max_sample_rate", 48000.0, ValueError), + ], +) +def test_sample_rate_filter_rejects_invalid_config(field_name: str, value: object, error_type: type[Exception]) -> None: + with pytest.raises(error_type, match=field_name): + SampleRateFilterStage(**{field_name: value}) + + def test_concatenation_reads_require_nested_segment_audio_keys() -> None: """SegmentConcatenation reads waveform+sample_rate from EACH child, not just the container.""" concat = SegmentConcatenationStage() From b04103873db7c65744c2af31d4870cecf14a3721 Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Mon, 21 Sep 2026 11:18:35 +0000 Subject: [PATCH 25/25] fix(audio): preserve normalized waveform identity Signed-off-by: shbhawsar --- nemo_curator/stages/audio/_agent/_residency.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/nemo_curator/stages/audio/_agent/_residency.py b/nemo_curator/stages/audio/_agent/_residency.py index 61a543f44e..608a841588 100644 --- a/nemo_curator/stages/audio/_agent/_residency.py +++ b/nemo_curator/stages/audio/_agent/_residency.py @@ -116,7 +116,7 @@ def normalize_audio_waveform( ) -> torch.Tensor: """Convert supported resident audio to channel-first float32.""" try: - tensor = waveform.detach() if torch.is_tensor(waveform) else torch.as_tensor(waveform) + tensor = waveform if torch.is_tensor(waveform) else torch.as_tensor(waveform) except Exception as ex: msg = f"[{stage_name}] Resident waveform must be convertible to a torch tensor" raise TypeError(msg) from ex @@ -137,6 +137,9 @@ def normalize_audio_waveform( ) raise TypeError(msg) + if tensor.requires_grad: + tensor = tensor.detach() + tensor = ensure_waveform_2d(tensor) if mono and tensor.shape[0] > 1: tensor = tensor.mean(dim=0, keepdim=True)