feat(audio-inference): agent-ready ASR/VAD/diarization (residency + fan-out provenance fix)#2176
feat(audio-inference): agent-ready ASR/VAD/diarization (residency + fan-out provenance fix)#2176shubhamNvidia 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).
…ageContract + residency)
Greptile SummaryThis PR makes the ASR, VAD, and speaker diarization stages agent-ready by routing audio loading through a new
Confidence Score: 3/5Safe to merge for file-based pipelines, but waveform-input paths through PyAnnote with the default write_rttm=True will silently accumulate RTTM files in the OS temp directory on every invocation. The pyannote RTTM write derives its output path from file_path (which is the materialized temp WAV when waveform input is used), but only the WAV is registered in temp_paths; the .rttm file is never cleaned up — a real resource leak on every waveform-input invocation in the default configuration. pyannote.py deserves a close second look around the write_rttm + temp-path interaction; sortformer.py and asr_nemo.py have smaller inconsistencies worth addressing before agentic waveform workloads go to production. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[AudioTask arrives] --> B{residency setting}
B -- file or auto --> C{audio_filepath_key in task.data?}
B -- waveform --> D{waveform + sample_rate in task.data?}
C -- yes, file exists --> E[use real file path]
C -- no --> F{residency == file?}
F -- yes --> G[return expanded path unverified]
F -- no --> H{waveform in task.data?}
H -- yes --> I[write temp WAV, append to temp_paths]
H -- no --> J[return None, raise ValueError]
D -- yes --> I
D -- no --> J
I --> K[pass path to ASR / VAD / diarizer]
E --> K
K --> L[inference runs]
L --> M[finally: cleanup_temp_files removes temp WAVs]
M --> N[return tasks or fan-out children]
subgraph Fan-out provenance
N --> O{fanout == True?}
O -- yes --> P[_segment_child_data per segment]
P --> Q[emit AudioTask per segment with segment_num start end]
O -- no --> R[emit single AudioTask with segments list]
end
%%{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"}}}%%
flowchart TD
A[AudioTask arrives] --> B{residency setting}
B -- file or auto --> C{audio_filepath_key in task.data?}
B -- waveform --> D{waveform + sample_rate in task.data?}
C -- yes, file exists --> E[use real file path]
C -- no --> F{residency == file?}
F -- yes --> G[return expanded path unverified]
F -- no --> H{waveform in task.data?}
H -- yes --> I[write temp WAV, append to temp_paths]
H -- no --> J[return None, raise ValueError]
D -- yes --> I
D -- no --> J
I --> K[pass path to ASR / VAD / diarizer]
E --> K
K --> L[inference runs]
L --> M[finally: cleanup_temp_files removes temp WAVs]
M --> N[return tasks or fan-out children]
subgraph Fan-out provenance
N --> O{fanout == True?}
O -- yes --> P[_segment_child_data per segment]
P --> Q[emit AudioTask per segment with segment_num start end]
O -- no --> R[emit single AudioTask with segments list]
end
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if self.write_rttm: | ||
| logger.info(f"Writing {len(diarization._tracks)} turns to RTTM file") | ||
| rttm_filepath = os.path.splitext(file_path)[0] + ".rttm" | ||
| rttm_fs, rttm_path = url_to_fs(rttm_filepath) | ||
| with rttm_fs.open(rttm_path, "w") as rttm_file: | ||
| diarization.write_rttm(rttm_file) |
There was a problem hiding this comment.
When write_rttm=True (the default) and audio arrives as an in-memory waveform, resolve_audio_path materializes a temp WAV (e.g. /tmp/tmpXXXXXX.wav) and registers it in temp_paths. The RTTM is then written to /tmp/tmpXXXXXX.rttm — derived from file_path via os.path.splitext. The finally block in process() calls cleanup_temp_files(temp_paths), which only contains the .wav entry; the .rttm file is never removed. In long-running agentic pipelines that push waveforms through this stage, this will silently accumulate RTTM files in the temp directory until the disk is full.
The fix is to register the RTTM path alongside the WAV when file_path came from a temp materialization. One approach is to pass temp_paths (or a boolean flag) into _diarize_file and append rttm_path to it, or to guard RTTM writing with file_path not in temp_paths.
| for task in tasks: | ||
| if not self.validate_input(task): | ||
| has_file = self.filepath_key in task.data | ||
| has_waveform = self.waveform_key in task.data and self.sample_rate_key in task.data | ||
| if not (has_file or has_waveform): | ||
| msg = f"Task {task.task_id} missing required columns for {type(self).__name__}: {self.inputs()}" | ||
| raise ValueError(msg) |
There was a problem hiding this comment.
residency="file" silently ignores them
The pre-resolution validation loop accepts a task when has_waveform=True, regardless of input_residency. But when input_residency="file" (the default), resolve_audio_path ignores the waveform entirely and returns None if filepath_key is absent from task.data. The task then hits the if path is None: guard with the error "missing audio input", leaving the caller no indication that the waveform was present but not used.
Either tighten the first-pass check to account for residency (e.g. has_file or (has_waveform and self.input_residency != "file")), or amend the error message to hint that the stage is running in residency="file" mode and requires a path.
1 participant
Summary
Makes the ASR/VAD/diarization stages agent-ready, routes their audio loading through the residency resolver, and fixes fan-out provenance.
Common changes
resolve_audio_path(accepts a file, or an in-memory waveform materialized to a temp WAV) and cleans up temp files in afinallyblock.input_residency="file"by default → same behavior as before. Each addsdescribe(); batched stages (e.g. ASR) areBATCH_ONLY.Fan-out provenance fix (pyannote, sortformer, whisperx_vad)
original_filethrough a fallback chain — configuredoriginal_file_key→audio_filepath_key→ canonicalaudio_filepath→ legacyresampled_audio_filepath→"unknown"— instead of collapsing to"unknown"when audio arrived under the canonical key. Children also get a distinctsegment_numand per-segment timing (start/end,start_ms/end_ms,duration).segmentskey from children (a list of tuples) so it doesn't collide with the semanticsegmentsrole and break dict-shaped consumers.Notes for reviewers
_segment_child_data/_fanout_segmentshelpers inpyannote.py,sortformer.py,whisperx_vad.py(provenance fallback + temp handling) — these are near-identical across the three.Test plan
pytest tests/stages/audio/inference -q