feat(audio): agent-ready foundation — contract layer, residency, discovery - #2332
feat(audio): agent-ready foundation — contract layer, residency, discovery#2332shubhamNvidia wants to merge 4 commits into
Conversation
…overy 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 <shbhawsar@nvidia.com>
Greptile SummaryThe PR adds an agent-facing contract, discovery, planning, residency, and conformance layer for audio stages, then applies it to common preprocessing stages.
Confidence Score: 3/5The PR is not yet safe to merge because audio-folder identity collisions can silently reuse the wrong output and successful checkpoints still lack a durable completion state. The current path-to-ID transformation remains non-injective and feeds collision-prone output filenames, while checkpoint finalization removes ownership without creating the completion marker that setup and retry logic consume. Files Needing Attention: nemo_curator/stages/audio/common.py Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart LR
Files[Audio files or manifests] --> Tasks[AudioTask records]
Tasks --> Preprocessing[Audio preprocessing stages]
Preprocessing --> Contracts[Agent-ready stage contracts]
Contracts --> Discovery[Catalog and discovery]
Contracts --> Planning[Pipeline validation]
Preprocessing --> Checkpoint[Manifest checkpoint]
Preprocessing --> Downstream[Inference and filtering stages]
Reviews (5): Last reviewed commit: "fix(audio): three silent-wrong-data find..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| rel = os.path.relpath(abspath, os.path.abspath(self.data_dir)) | ||
| item_id = os.path.splitext(rel)[0].replace(os.sep, "__") |
There was a problem hiding this comment.
When a recursive folder contains both spk1/utt1.wav and spk1__utt1.wav, replacing separators with __ assigns both tasks the same audio_item_id, causing identity-indexed records or downstream output filenames to collide.
Knowledge Base Used: Audio curation stages
| 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 |
There was a problem hiding this comment.
After a checkpoint run succeeds, release_retry_reservation() only removes the ownership sidecar and never creates the _COMPLETE marker consumed by setup and retry logic. Completed and interrupted JSONL files therefore lack a durable completion distinction, so resume tooling cannot safely determine whether all rows were written.
Knowledge Base Used: Curation stage framework
The complete audio agent work as one branch: the same stage-side contracts split across PRs NVIDIA-NeMo#2332-NVIDIA-NeMo#2342, plus the agent that consumes them. Stages (nemo_curator/stages/audio/): - the _agent/ contract layer (AgentReady, StageContract, roles, input residency, discovery, planning, conformance) and the public agent.py entry point - describe() contracts and *_key params across common, preprocessing, io, filtering, metrics, segmentation, tagging, inference, alm, datasets and advanced_pipelines - new stages: ChannelCount, SampleRateFilter, CreateInitialManifestAudioFolder, PreserveByValueConditions, ManifestCheckpoint, ManifestGroupExport, DocumentBatchJsonlWriter Agent (nemo_curator/audio_agent/): - deterministic verbs (discover/describe/validate/smoke/run/report/reuse/delta), knowledge cards, recipes, and the host skills Supporting: resource telemetry used by smoke calibration (utils/performance_utils.py, opt-in via RESOURCE_PROBE in backends/base.py), the nemo-curator-audio console script, and the audio_agent CI matrix entry. Verified: 51/51 stages pass static conformance, the stage-compatibility baseline hash is unchanged, and whole-repo `ruff check` is clean. Signed-off-by: Shubham Bhawsar <shbhawsar@nvidia.com>
The complete audio agent work as one branch: the same stage-side contracts split across PRs NVIDIA-NeMo#2332-NVIDIA-NeMo#2342, plus the agent that consumes them. Stages (nemo_curator/stages/audio/): - the _agent/ contract layer (AgentReady, StageContract, roles, input residency, discovery, planning, conformance) and the public agent.py entry point - describe() contracts and *_key params across common, preprocessing, io, filtering, metrics, segmentation, tagging, inference, alm, datasets and advanced_pipelines - new stages: ChannelCount, SampleRateFilter, CreateInitialManifestAudioFolder, PreserveByValueConditions, ManifestCheckpoint, ManifestGroupExport, DocumentBatchJsonlWriter Agent (nemo_curator/audio_agent/): - deterministic verbs (discover/describe/validate/smoke/run/report/reuse/delta), knowledge cards, recipes, and the host skills Supporting: resource telemetry used by smoke calibration (utils/performance_utils.py, opt-in via RESOURCE_PROBE in backends/base.py), the nemo-curator-audio console script, and the audio_agent CI matrix entry. Verified: 51/51 stages pass static conformance, the stage-compatibility baseline hash is unchanged, and whole-repo `ruff check` is clean. Signed-off-by: Shubham Bhawsar <shbhawsar@nvidia.com>
| 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: |
There was a problem hiding this comment.
nemo_curator/stages/audio/preprocessing/mono_conversion.py:221-237andnemo_curator/stages/audio/preprocessing/channel_count.py:498-513only avoid assigning the converted tensor whenkeep_waveform_in_task=False; neither removes a waveform that was already present. With supported input_residency="waveform", keep_waveform_in_task=False, write_to_disk=True, both stages wrote a mono [8, 1]WAV but returned the original[2, 8]stereo tensor while claimingis_mono=True/num_channels=1. The next stage in input_residency="auto"prefers that stale tensor and silently bypasses the conversion. Both constructors also accept the no-sink combinationkeep_waveform_in_task=False, write_to_disk=False`, which retains the original
audio while emitting converted metadata.
Reject conversion configurations with neither output sink (and require
update_audio_filepath to imply disk output). After a successful disk-only
conversion, remove waveform_key and sample_rate_key, declare those
removes_keys in the configured contract, and add a resident-input ->
disk-only conversion -> auto-residency consumer regression for both stages.
This is a blocker under the audio review rubric because a valid advertised
configuration silently carries the wrong audio.
There was a problem hiding this comment.
Agreed, blocker. The existing test only covers file input, so your resident-waveform case is genuinely uncovered. Implementing the full remedy — reject no-sink configs, update_audio_filepath implies disk output, remove waveform_key/sample_rate_key after disk-only conversion and declare them in removes_keys, plus the resident → disk-only → auto-residency regression for both stages. Landing separately so the behavioural change reviews on its own.
| sf.write(path, arr, int(sample_rate)) | ||
| return path | ||
|
|
||
| os.makedirs(output_dir, exist_ok=True) |
There was a problem hiding this comment.
nemo_curator/stages/audio/_agent/_residency.py:244-254 hashes only
arr.tobytes() and sample rate; array shape/channel count and dtype are absent.
Two seconds of mono silence ([1, 32000]) and one second of stereo silence
([2, 16000]) therefore receive the same filename at 16 kHz. The second atomic
replace changed the first task's path from a 32,000-frame mono file into a
16,000-frame stereo file. This is not a cryptographic SHA collision: the
identity omits audio-defining inputs.
Include the canonical array shape and dtype (and any format/subtype parameters)
in the digest, or hash the fully encoded artifact. Add a regression proving
that equal flattened samples with different channel layouts cannot alias.
Adjacent lines 209-215 should also always transpose the documented
channel-first 2-D representation: the current channels < samples heuristic
writes a valid two-channel/one-sample tensor as two mono frames
There was a problem hiding this comment.
Agreed — the identity was missing audio-defining inputs. The digest now includes canonical shape, dtype and output format alongside sample rate, and lines 209-215 always transpose the documented channel-first layout. Regression added for equal flattened samples with different channel layouts, and for the two-channel/one-sample case.
| # caller never wrote. | ||
| walk.past_composite = True | ||
|
|
||
| for item in leaves.get(index, []): |
There was a problem hiding this comment.
nemo_curator/stages/audio/_agent/_planning.py:645-679 builds contracts and
checks reads/gates but never compares the accepts_task_type and
produces_task_type fields that _agent_ready.py:230-234 says enable this
check. A folder source (EmptyTask -> AudioTask) followed by
ManifestReaderStage (FileGroupTask -> AudioTask) returns ok=True,
keys_ok=True, and no issues; executing the second stage on that AudioTask
raises FileNotFoundError while treating dict keys as manifest paths.
This is reachable through the intended repair flow: public
find_producers("audio_filepath") recommends the internal
ManifestReaderStage rather than the public ManifestReader composite.
Track the current produced task type through the walk, accept an explicit
initial_task_type, and emit an error when the next concrete stage cannot
accept it (including a defined policy for unions/subclasses). Cover this exact
chain and two consecutive ManifestReaderStage instances.
There was a problem hiding this comment.
Agreed. The check is already defined on the agent side (task_type_mismatch, error, exact match on the derived type name, skipped when either side is unknown) — the gap is that it isn't in validate_pipeline(). Moving it down: the walk tracks the produced task type, accepts an explicit initial_task_type, and errors when the next stage can't accept it, with a union satisfied by any member and a declared base accepting subclasses. Covering your folder-source → ManifestReaderStage chain and two consecutive ManifestReaderStage instances.
| "original_file": original_file, | ||
| "num_segments": len(mappings), | ||
| "total_duration_sec": total_duration_sec, | ||
| self.original_file_key: original_file, |
There was a problem hiding this comment.
nemo_curator/stages/audio/preprocessing/concatenation.py:290-305 constructs a
new output_data dict, dropping segments and every unrelated upstream key,
but its contract at lines 143-159 inherits
preserves_upstream_keys=True. Independently,
nemo_curator/stages/audio/_agent/_planning.py:510-540 never consults that
contract field. Consequently concatenation followed by a filter reading
upstream text validates with ok=True/keys_ok=True; runtime output has no
text and the filter raises ValueError.
Set the concatenation contract to preserves_upstream_keys=False, and teach
_advance() to clear unavailable keys, their roles/producers, and tensor
residency before folding the replacement stage's writes. Add a regression that
the downstream text read is rejected during validation.
There was a problem hiding this comment.
Agreed on both halves. Setting preserves_upstream_keys=False and teaching _advance() to clear unavailable keys, their roles/producers and tensor residency before folding in the replacement writes, with the downstream text read rejected during validation.
| # 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: |
There was a problem hiding this comment.
nemo_curator/stages/audio/_agent/_composite.py:182-188 detects exactly the
nested-composite condition that Pipeline._decompose_stages rejects, but stores
it in opaque rather than unrunnable. A two-level composite therefore gets
only [('composite', 'warning')] and report.ok=True; Pipeline.build() raises
TypeError: Nested composition is not supported for the same object.
Assign this branch to unrunnable[index] so validate_pipeline() emits
composite_unrunnable as an error, and add a regression asserting parity with
Pipeline.build().
There was a problem hiding this comment.
Agreed, moved to unrunnable[index] so validate_pipeline() emits composite_unrunnable as an error, with a parity regression against Pipeline.build(). Kept the check after the single-child branch, matching the executor: its nested rejection only fires inside len(sub_stages) > 1.
|
|
||
|
|
||
| @dataclass | ||
| class ManifestWriterStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): |
There was a problem hiding this comment.
nemo_curator/stages/audio/common.py:758-834 makes
ManifestWriterStage agent-ready but supplies no AGENT_STATIC. As a result,
the sanctioned agent.describe_stage("ManifestWriterStage") reports
writes_to_disk=False, no output path, no lifecycle side effect, and no
serialization constraint, while the configured contract reports all four and
setup() truncates the destination. An agent doing the prescribed static
sandbox check can therefore treat a destructive JSON sink as a pure stage.
Add invariant StaticHints(gates=...) matching the configured contract, and a
static-versus-configured regression for every invariant gate.
There was a problem hiding this comment.
Agreed. Added invariant StaticHints(gates=...) for all four gates plus per_row_independent. The regression asserts static equals configured rather than a hardcoded list, so the two can't drift.
| if TYPE_CHECKING: | ||
| from nemo_curator.stages.audio._agent._agent_ready import StageContract | ||
|
|
||
| _IMPORTED = False |
There was a problem hiding this comment.
nemo_curator/stages/audio/_agent/_catalog.py:43-94 explicitly records skipped
modules so a supported CPU install can distinguish “not installed here” from
“Curator cannot do this.” However, nemo_curator/stages/audio/agent.py:63-105
neither exports unavailable_modules() nor includes the failures in catalog
JSON. In the review environment, public discovery returned 12 stages while the
private-only function recorded five missing modules (pyannote, whisperx,
nemo_text_processing twice, and OpenCC). The only supported facade offers no
way to observe that distinction.
Expose unavailable_modules through agent.py and/or include an
unavailable_modules field in the public catalog JSON. Test discovery under a
supported partial dependency set and assert that import failures are visible
without parsing stderr warnings.
There was a problem hiding this comment.
Agreed, unavailable_modules is now exported through nemo_curator.stages.audio.agent and in its all, with a test. Happy to add it to the catalog JSON too if you'd prefer that shape.
| @@ -32,11 +32,15 @@ | |||
| "ALMDataOverlapStage": "nemo_curator.stages.audio.alm", | |||
| "AudioDataFilterStage": "nemo_curator.stages.audio.advanced_pipelines", | |||
| "BandFilterStage": "nemo_curator.stages.audio.filtering", | |||
There was a problem hiding this comment.
nemo_curator/stages/audio/__init__.py:30-74 adds other stages from this PR to
both _LAZY and __all__, but omits
CreateInitialManifestAudioFolderStage. Although the PR presents it as the
user-facing way to start a pipeline from a folder, the conventional import
from nemo_curator.stages.audio import CreateInitialManifestAudioFolderStage
raises ImportError and hasattr(nemo_curator.stages.audio, ...) is false.
Add the class to both public-name tables and cover the top-level import.
There was a problem hiding this comment.
Agreed, oversight. Added to both _LAZY and all, with the top-level import covered.
| 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 |
There was a problem hiding this comment.
nemo_curator/stages/audio/AGENTS.md:14-24,38-43,57-60 tells coding agents they
must use nemo_curator.audio_agent, its CLI, authoring skill, card schema and
card directory; none exists at this head. AGENT_READY.md:293-295 points to a
missing test_agent_simulation_pipelines.py, and the public facade example at
agent.py:34 asks for UTMOSFilterStage, which is not agent-ready in this PR
and raises KeyError. Line 47 then invokes another nonexistent
audio_agent.validate API. Because AGENTS.md is automatically loaded and
uses MUST/sanctioned language, this makes the documented workflow impossible
in the foundation PR itself.
Remove the forward references and use stages/APIs present in this slice, or
land the required package/cards/test harness atomically. Add a smoke/doctest
that executes the public discovery example from a supported install.
There was a problem hiding this comment.
Context: these 11 PRs are a split of one branch that also contains nemo_curator/audio_agent, so the agent package, cards, test_agent_simulation_pipelines.py and the UTMOSFilterStage example all resolve there but not at this head (the latter arrives with #2335). Since the branches are meant to end up equivalent to that branch, I'd rather not rewrite the guardrails for a temporary state.
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| """EXEMPLAR per-stage conformance tests — copy these as templates. |
There was a problem hiding this comment.
tests/stages/audio/_agent/test_agent_conformance_examples.py:15-113 contains
four per-stage exemplar tests, but the test tree contains no reference to
validate_pipeline, composite expansion, audio_stage_catalog,
catalog_as_json, describe_stage, find_producers/find_consumers,
role_index, or the get/list public APIs. The focused CPU slice passes 209
tests, but those tests do not exercise the core paths above; current GitHub
checks also contain no unit-test job. This falls short of CONTRIBUTING.md's
requirement that source changes carry relevant tests and cover at least 80% of
their changes, and the single +6,801-line commit is too large to review safely
without those seams.
Split the framework from stage migrations/new stages where practical, and add
focused test modules mirroring _agent for registry/schema, catalog/partial
imports, composite parity, planner state transitions and residency identity.
At minimum, turn each deterministic reproduction in this review into a
regression before merge.
There was a problem hiding this comment.
Every deterministic reproduction in this review is now a regression in
tests/stages/audio/_agent/test_agent_foundation_regressions.py, including the three
behavioural ones in 813b6bf. The focused modules mirroring _agent exist on the source
branch — say the word and I'll port them into this slice instead of letting them arrive
with the later PRs. Stage migrations are already split across the ten dependent PRs; if
you want the framework cut further, tell me where and I'll restructure.
6dc96a4 to
fbf865b
Compare
Signed-off-by: Shubham Bhawsar <shbhawsar@nvidia.com>
Conflict was an add/add in the audio ``_LAZY`` map: upstream added ``RegexSubstitutionStage`` (new ``text_filtering`` package) at the same sorted slot this branch added ``SampleRateFilterStage``. Both kept. ``__all__`` auto-merged with both names already present. Signed-off-by: Shubham Bhawsar <shbhawsar@nvidia.com>
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 <shbhawsar@nvidia.com>
2 participants
The layer everything else builds on.
Public surface:
agent.py— importnemo_curator.stages.audio.agent. Everything else is private (_agent/), so the supported API is one module.What's inside
_agent/(8 modules, all private):_agent_ready.pyAgentReady(the mixin a stage inherits) andStageContract, plus the closed sets a contract is written in:Cardinality(1:1/1:1 nested-list/1:N fan-out/N:1/filter),AudioForm(file/waveform),ErrorPolicy(skip/fail/annotate/unknown),IOSpec,Gates, andConditionalWrite. Everything else in the package reads these types._residency.pyinput_residency/keep_waveform_in_taskoptions in the module PRs are built on. Also does the stable/atomic write helper the stages use._roles.pypred_textand one calledtextmay both be "the transcript"; roles are how the agent decides two stages actually connect instead of matching raw strings._agent_registry.pydescribe()stays about I/O and behavior. Also builds the resolved contract for a configured stage._catalog.py_composite.pySplitASRAlignJoinStageinstead of treating it as an opaque box._conformance.pyprocess()really touches, and that cardinality matches runtime. It can also sweep every stage with no fixtures, which is what keeps contracts honest as stages change._planning.pyApplied to the first stages
common.py+ the fullpreprocessing/folder made agent-ready, as the first worked slice.MonoConversionStageandSegmentConcatenationStagegainkeep_waveform_in_task/write_to_disk/output_dir(and residency input on mono), so a preprocessing step feeding a GPU stage can hand its audio over in memory instead of writing a file the next stage immediately re-reads. Writing to disk stays the default, and a stage refuses a configuration with no output sink at all rather than silently producing nothing.Five new stages
In
preprocessing/:ChannelCountStage— record the channel count, filter on it, or convert it (e.g. downmix to mono), in one stage.SampleRateFilterStage— select clips by sample rate without resampling them (previously the only option was a full resample).In
common.py:CreateInitialManifestAudioFolderStage— start a pipeline from a plain folder of audio files. Scans a directory and emits one task per file with its path and a derived id. The existing sources either read an already-written JSONL manifest (ManifestReader) or were dataset-specific (CreateInitialManifest{ReadSpeech,Fleurs}Stage); there was nothing for "here is a folder of WAVs, curate it", which is the most common starting point.PreserveByValueConditionsStage— filter on several conditions at once (AND or OR), against top-level rows or a nested per-segment list.PreserveByValueStagecompares exactly one key to one value, so anything like "duration ≥ 5 and MOS ≥ 3.5" needed several stacked stages and still couldn't express OR. Missing keys fail closed rather than silently passing.ManifestCheckpointStage— persist a metadata-only JSONL boundary mid-pipeline and pass rows through unchanged, so a later run can resume from it instead of recomputing expensive GPU work. It refuses to overwrite an existing file, and rejects non-JSON values (e.g. a resident waveform) at the boundary rather than writing something unreadable.Tests:
test_common.pyandtest_create_manifest_audio_folder.pycover thecommon.pystages, thepreprocessing/tests cover the preprocessing folder, and the_agentframework tests live intests/stages/audio/_agent/— so every stage shipped here ships with its tests.Note for reviewers:
_agent/__init__.pyis intentionally import-free (it breaks acommon↔_residencyimport cycle); please don't add re-exports there.