feat(audio-advanced): agent-ready AudioDataFilter composite pipeline#2181
feat(audio-advanced): agent-ready AudioDataFilter composite pipeline#2181shubhamNvidia 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).
…ready (describe/StageContract + residency)
Greptile SummaryThis PR makes
Confidence Score: 3/5Safe to merge after addressing the batch_only contract mismatch in PreserveByValueStage and the temp-file leak in _residency.py; AudioDataFilterStage itself is unchanged at runtime. PreserveByValueStage.describe() returns a contract that says process() is callable, but calling it raises NotImplementedError — any agent or tool that consults describe() directly (rather than going through the registry) will act on wrong information. The _residency.py temp-file leak on sf.write() failure and the _as_soundfile_array shape fall-through for short clips are real defects, though both are narrow error paths unlikely to hit in typical usage. The core AudioDataFilterStage decomposition logic is untouched and the additive nature of the change limits blast radius. nemo_curator/stages/audio/common.py (PreserveByValueStage.describe) and nemo_curator/stages/audio/_residency.py (resolve_audio_path and _as_soundfile_array) Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Agent
participant Stage as AudioDataFilterStage / other AgentReady stage
participant describe as stage.describe()
participant registry as _agent_registry.build_contract()
Agent->>Stage: stage.describe()
Stage->>describe: "returns StageContract(wrappable=False, ...)"
describe-->>Agent: "StageContract (batch_only=False by default)"
Note over Agent,Stage: PreserveByValueStage: BATCH_ONLY=True not reflected here
Agent->>registry: build_contract(stage)
registry->>Stage: reads cls.BATCH_ONLY
registry-->>Agent: "StageContract (batch_only=True correctly set)"
Agent->>Stage: stage.process(task)
Note over Stage: PreserveByValueStage.process() raises NotImplementedError!
Agent->>Stage: stage.process_batch(tasks)
Stage-->>Agent: list[AudioTask]
%%{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 AudioDataFilterStage / other AgentReady stage
participant describe as stage.describe()
participant registry as _agent_registry.build_contract()
Agent->>Stage: stage.describe()
Stage->>describe: "returns StageContract(wrappable=False, ...)"
describe-->>Agent: "StageContract (batch_only=False by default)"
Note over Agent,Stage: PreserveByValueStage: BATCH_ONLY=True not reflected here
Agent->>registry: build_contract(stage)
registry->>Stage: reads cls.BATCH_ONLY
registry-->>Agent: "StageContract (batch_only=True correctly set)"
Agent->>Stage: stage.process(task)
Note over Stage: PreserveByValueStage.process() raises NotImplementedError!
Agent->>Stage: stage.process_batch(tasks)
Stage-->>Agent: list[AudioTask]
Reviews (1): Last reviewed commit: "feat(audio-advanced_pipelines): make adv..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def describe(self) -> StageContract: | ||
| return StageContract( | ||
| reads=IOSpec(data_keys=[self.input_value_key]), | ||
| writes=IOSpec(data_keys=[self.input_value_key]), | ||
| cardinality="filter", | ||
| ) |
There was a problem hiding this comment.
PreserveByValueStage.describe() returns StageContract with batch_only at its default of False, yet BATCH_ONLY = True is set on the class and process() immediately raises NotImplementedError. Any agent or tool that calls stage.describe() directly (instead of going through the registry's build_contract) will receive a contract that incorrectly advertises process() as callable, which would raise at runtime. The fix is to propagate the class flag into the returned contract.
| def describe(self) -> StageContract: | |
| return StageContract( | |
| reads=IOSpec(data_keys=[self.input_value_key]), | |
| writes=IOSpec(data_keys=[self.input_value_key]), | |
| cardinality="filter", | |
| ) | |
| def describe(self) -> StageContract: | |
| return StageContract( | |
| reads=IOSpec(data_keys=[self.input_value_key]), | |
| writes=IOSpec(data_keys=[self.input_value_key]), | |
| cardinality="filter", | |
| batch_only=True, | |
| ) |
| 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 |
There was a problem hiding this comment.
sf.write() raises (e.g. the waveform dtype is unsupported, or the array shape is wrong after _as_soundfile_array), the temp file created by mkstemp is left on disk but never appended to register_temp. The caller therefore never learns about it and cleanup_temp_files will not remove it. Wrapping the write in a try/except and cleaning up on failure closes the leak.
| 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 | |
| fd, tmp = tempfile.mkstemp(suffix=".wav", dir=temp_dir) | |
| os.close(fd) | |
| try: | |
| sf.write(tmp, _as_soundfile_array(waveform), int(sample_rate)) | |
| except Exception: | |
| with contextlib.suppress(OSError): | |
| os.remove(tmp) | |
| raise | |
| if register_temp is not None: | |
| register_temp.append(tmp) | |
| return tmp |
| 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 |
There was a problem hiding this comment.
_as_soundfile_array falls through to return waveform when channels >= samples (i.e. a 2-D array whose first dimension is not smaller than the second). SoundFile expects audio in (samples, channels) shape for multi-channel data, so returning the un-transposed (channels, samples) array in that branch would write audio with channels and samples swapped, producing a corrupt WAV. In practice audio files have far more samples than channels, but a very short clip (e.g. a 1-sample test fixture) would silently trigger this.
1 participant
Summary
Makes the composite
AudioDataFilterStageagent-ready. Additive only.advanced_pipelines/audio_data_filter/audio_data_filter.py— the composite (which decomposes into the underlying filter stages) now subclassesAgentReadyand exposesdescribe()/StageContract. Decomposition and runtime behavior unchanged.Notes for reviewers
Test plan
pytest tests/stages/audio/advanced_pipelines -q