feat(audio): agent-ready foundation — StageContract mixin + residency resolver#2170
feat(audio): agent-ready foundation — StageContract mixin + residency resolver#2170shubhamNvidia wants to merge 2 commits into
Conversation
…idency resolver Shared base imported by the audio stage modules to make them agent-ready: - _agent_ready.py: AgentReady mixin, StageContract/IOSpec/Gates/Role - _residency.py: input residency resolver (file/waveform/auto) - common.py: agent-ready helpers (ensure_mono/ensure_waveform_2d) Excludes the agent orchestration layer (registry/catalog/conformance/planning/agent).
shubhamNvidia
requested review from
praateekmahajan
and removed request for
a team
Greptile SummaryThis PR adds the agent-ready foundation for audio stage discovery: a
Confidence Score: 3/5Two real defects need fixing before merge: a resource leak in the core temp-file materialization path and a stale outputs() declaration that misreports available keys when keep_waveform_in_task=False. The temp-file leak in both resolve_audio_path and _write_audio is a genuine resource management bug: mkstemp creates a file before sf.write is called, and any exception during the write orphans that file with no way to recover it via cleanup_temp_files. In _write_audio the situation is worse because soundfile.SoundFileError escapes the process() except clause entirely, turning a per-item skip into a stage crash when write_to_disk=True. The outputs() mismatch means pipeline metadata says the waveform key is always produced while process() only writes it conditionally. nemo_curator/stages/audio/preprocessing/mono_conversion.py (outputs() mismatch and _write_audio leak) and nemo_curator/stages/audio/_residency.py (resolve_audio_path temp-file leak). Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Agent
participant Stage as MonoConversionStage
participant Residency as _residency.py
participant FS as Filesystem
Agent->>Stage: describe()
Stage-->>Agent: "StageContract(reads.accepts=[file,waveform], gates, writes)"
Agent->>Stage: process(task)
Stage->>Residency: "resolve_audio(item, residency=input_residency)"
alt waveform in item
Residency-->>Stage: (waveform_2d, sample_rate)
else file path exists
Residency->>FS: load_audio_file(path)
FS-->>Residency: (tensor, sr)
Residency-->>Stage: (tensor, sr)
end
Stage->>Stage: mean(channels to mono)
opt "keep_waveform_in_task=True"
Stage->>Stage: "task.data[waveform_key] = mono_waveform"
end
opt "write_to_disk=True"
Stage->>Stage: _write_audio(mono_waveform, sr, task)
Stage->>FS: mkstemp() to temp path
Stage->>FS: sf.write(temp_path, arr, sr)
Note over Stage,FS: if sf.write raises, temp file leaks
Stage->>Stage: "task.data[output_audio_filepath_key] = path"
opt update_audio_filepath
Stage->>Residency: produce_audio_filepath(task.data, path)
end
end
Stage-->>Agent: task (or [])
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant Agent
participant Stage as MonoConversionStage
participant Residency as _residency.py
participant FS as Filesystem
Agent->>Stage: describe()
Stage-->>Agent: "StageContract(reads.accepts=[file,waveform], gates, writes)"
Agent->>Stage: process(task)
Stage->>Residency: "resolve_audio(item, residency=input_residency)"
alt waveform in item
Residency-->>Stage: (waveform_2d, sample_rate)
else file path exists
Residency->>FS: load_audio_file(path)
FS-->>Residency: (tensor, sr)
Residency-->>Stage: (tensor, sr)
end
Stage->>Stage: mean(channels to mono)
opt "keep_waveform_in_task=True"
Stage->>Stage: "task.data[waveform_key] = mono_waveform"
end
opt "write_to_disk=True"
Stage->>Stage: _write_audio(mono_waveform, sr, task)
Stage->>FS: mkstemp() to temp path
Stage->>FS: sf.write(temp_path, arr, sr)
Note over Stage,FS: if sf.write raises, temp file leaks
Stage->>Stage: "task.data[output_audio_filepath_key] = path"
opt update_audio_filepath
Stage->>Residency: produce_audio_filepath(task.data, path)
end
end
Stage-->>Agent: task (or [])
Reviews (2): Last reviewed commit: "feat(audio): add MonoConversionStage as ..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| 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 |
There was a problem hiding this comment.
None | X union type silently maps to no JSON-schema type
replace("|None", "") only strips the suffix |None from a type string. If the type annotation is written as None | str (i.e. None first, PEP 604 style), the stripped string becomes None|str (after space removal), which matches none of the scalar/collection checks and causes _json_type to return None — so the property ends up in the schema with no type field at all. Any stage that uses a None | str-ordered union annotation will silently produce an incomplete schema entry.
| 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 |
There was a problem hiding this comment.
_as_soundfile_array transposes the (channels, samples) array to (samples, channels) only when channels < samples. soundfile.write always expects (samples, channels) for multi-channel data. For any waveform where the channel count is ≥ the sample count (e.g. a 2-channel clip with only 1 sample), the guard channels < samples is False and the untransposed (channels, samples) array is passed directly, producing a WAV with swapped channel/sample axes. Real audio almost never hits this, but it can occur in unit tests with synthetic fixtures or extremely short audio clips.
| def describe(self) -> StageContract: | ||
| return StageContract(cardinality="1:N fan-out", wrappable=False) |
There was a problem hiding this comment.
ManifestReader.describe() silently drops the lifecycle_side_effects gate
ManifestReaderStage.describe() declares Gates(lifecycle_side_effects=True) to signal that reading manifests is not a pure, side-effect-free operation. The wrapping ManifestReader returns a bare StageContract(cardinality="1:N fan-out", wrappable=False) with no gates, so an agent querying the composite's contract would not see the lifecycle_side_effects signal. Composites should either forward the composed stages' gates or explicitly re-declare them.
| def describe(self) -> StageContract: | ||
| return StageContract( | ||
| reads=IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"]), | ||
| writes=IOSpec(data_keys=[self.duration_key]), | ||
| ) |
There was a problem hiding this comment.
Why not just have AgentReady read the inputs and outputs instead of having to repeat code here?
There was a problem hiding this comment.
Fair, the keys here do duplicate inputs()/outputs(). describe() is kept separate because StageContract expresses things those hooks can't: residency (accepts=["file"|"waveform"]), alternative input shapes (reads_one_of), per-segment writes (segment_data_keys), cardinality, and gates. Also inputs()/outputs() default to ([], []) across the framework, so most stages don't populate them reliably — for this stage the only thing beyond them is accepts=["file"].
| def describe(self) -> StageContract: | ||
| return StageContract( | ||
| writes=IOSpec(data_keys=["audio_filepath"]), | ||
| cardinality="1:N fan-out", |
There was a problem hiding this comment.
Similar comment here, AgentReady can just check the ray_stage_spec. Instead of having to manually do it for every audio stage.
There was a problem hiding this comment.
ray_stage_spec is documented as Ray-Data-backend-only, and is_fanout_stage is just a boolean; it can't express the other cardinalities the contract needs (filter, N:1 aggregate, 1:1 nested-list), so it isn't a complete source for cardinality. That's why it's stated explicitly in describe(), which is the single backend-agnostic place that captures the full shape.
mohammadaaftabv
left a comment
mohammadaaftabv
left a comment
There was a problem hiding this comment.
Could we include one representative, end-to-end audio stage migration in this PR to demonstrate how this foundation is actually consumed?
GetAudioDurationStage.describe() shows how a contract is declared, but the PR still leaves it unclear how the complete path is intended to work: which contract fields are inferred versus manually authored, how an agent or registry consumes the contract, and how the new residency helpers participate in execution.
I’m not asking to migrate every audio stage here. One small vertical slice would be enough—preferably an existing stage that can accept either audio_filepath or waveform. The example should show:
- The stage inheriting
AgentReady. - Its
StageContract, including reads/writes, accepted audio forms, dispatch, gates, and configurable keys. - How the supported discovery path obtains that contract.
- How the consumer uses the contract to choose file versus waveform execution.
- A focused test covering both input forms and any temporary-file cleanup.
This would give later audio-stage PRs a canonical implementation to copy and would validate that _agent_ready.py and _residency.py work together as intended, rather than merging two foundational APIs whose integration is demonstrated only in later PRs. If the registry/orchestrator must remain outside this PR, a minimal test consumer or documented pseudocode for that boundary would still make the intended contract much clearer.
… test Demonstrates _agent_ready + _residency working together end to end: - MonoConversionStage inherits AgentReady and declares a StageContract (reads/writes, accepts=[file, waveform], gates, configurable *_key fields). - process() resolves audio from a file OR an in-memory waveform via _residency. - test covers both input forms, per-item dispatch, and temp-WAV materialize- then-cleanup, plus a minimal contract-consumer + pseudocode standing in for the registry/orchestrator (which fills params/key_roles/dispatch).
|
@mohammadaaftabv Contract: the stage inherits AgentReady and its describe() declares reads/writes (IOSpec), accepts=["file","waveform"], cardinality, gates (writes_to_disk when write_to_disk=True), and the configurable *_key fields. This gives later audio-stage PRs a reference implementation to copy and validates that _agent_ready.py and _residency.py work together within this PR. (I also moved mono_conversion.py out of the preprocessing PR #2172 so it isn't duplicated.) |
3 participants
Summary
Shared helpers imported by the audio stage modules.
_agent_ready.py—AgentReadymixin +StageContract/IOSpec/Gates/Role: a read-only, additive discovery contract exposed viadescribe(). No runtime behavior change._residency.py— input-residency resolver (resolve_audio,resolve_audio_path,produce_audio_filepath,cleanup_temp_files) unifying file/waveform/auto audio input handling and temp-file lifecycle.common.py— shared audio helpers (ensure_mono,ensure_waveform_2d) used by the stage modules.Notes for reviewers
describe()is read-only and does not change execution.