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 3 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 ASR, VAD, and diarization stages agent-ready by adding
Confidence Score: 5/5Safe to merge; all changes are additive, default residency preserves existing behaviour, and the temp-cleanup logic is correct for the normal happy path. The residency routing and fan-out provenance fix are well-structured.
Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[AudioTask with file or waveform] --> B{resolve_audio_path}
B -->|residency=file, path exists| C[Return local path]
B -->|residency=file, path absent| D[Return None → ValueError]
B -->|auto/waveform + waveform present| E[sf.write → temp WAV]
E --> F[Register in temp_paths]
F --> G[Return temp path]
B -->|protocol path e.g. s3://| H[Return protocol path as-is]
C --> I{Stage processing}
G --> I
H --> I
I -->|ASR batch| J[NeMo transcribe]
I -->|Pyannote| K[sf.read → diarize → RTTM?]
I -->|Sortformer| L[SortformerEncLabelModel.diarize]
I -->|WhisperX VAD| M[WhisperXVADModel.get_vad_segments]
J --> N[task.data pred_text set]
K -->|fanout=False| O[task.data segments set]
K -->|fanout=True| P[_fanout_segments → list of AudioTask]
L -->|fanout=False| Q[task.data diar_segments set]
L -->|fanout=True| R[_fanout_segments → list of AudioTask]
M -->|fanout=False| S[task.data vad_segments set]
M -->|fanout=True| T[_fanout_segments → list of AudioTask]
N --> U[finally: cleanup_temp_files]
O --> U
P --> U
Q --> U
R --> U
S --> U
T --> U
%%{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 with file or waveform] --> B{resolve_audio_path}
B -->|residency=file, path exists| C[Return local path]
B -->|residency=file, path absent| D[Return None → ValueError]
B -->|auto/waveform + waveform present| E[sf.write → temp WAV]
E --> F[Register in temp_paths]
F --> G[Return temp path]
B -->|protocol path e.g. s3://| H[Return protocol path as-is]
C --> I{Stage processing}
G --> I
H --> I
I -->|ASR batch| J[NeMo transcribe]
I -->|Pyannote| K[sf.read → diarize → RTTM?]
I -->|Sortformer| L[SortformerEncLabelModel.diarize]
I -->|WhisperX VAD| M[WhisperXVADModel.get_vad_segments]
J --> N[task.data pred_text set]
K -->|fanout=False| O[task.data segments set]
K -->|fanout=True| P[_fanout_segments → list of AudioTask]
L -->|fanout=False| Q[task.data diar_segments set]
L -->|fanout=True| R[_fanout_segments → list of AudioTask]
M -->|fanout=False| S[task.data vad_segments set]
M -->|fanout=True| T[_fanout_segments → list of AudioTask]
N --> U[finally: cleanup_temp_files]
O --> U
P --> U
Q --> U
R --> U
S --> U
T --> U
Reviews (2): Last reviewed commit: "fix(audio): residency-derived accepts/re..." | Re-trigger Greptile |
![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.
…x/sortformer/pyannote) + match _residency.py
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