feat(audio-segmentation): agent-ready speaker-separation + VAD (provenance fix)#2177
feat(audio-segmentation): agent-ready speaker-separation + VAD (provenance fix)#2177shubhamNvidia 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).
…ibe/StageContract + residency)
Greptile SummaryThis PR makes the audio segmentation stages agent-ready by adding
Confidence Score: 4/5Safe to merge with one fix: fan-out VAD segmentation with keep_segment_waveform_in_task=False silently passes audio_filepath to child tasks, causing downstream auto-mode stages to reload the full original file instead of the extracted segment. The new infrastructure files and speaker-separation provenance fix are clean. The vad_segmentation fan-out path has a real correctness gap: stripping the waveform while leaving audio_filepath intact causes wrong data when any downstream stage resolves audio in auto mode — it will load the full recording for every segment instead of the sliced portion. nemo_curator/stages/audio/segmentation/vad_segmentation.py — specifically _build_segment_item and the missing guard in post_init for fan-out + keep_segment_waveform_in_task=False. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Agent
participant VADSegmentationStage
participant SpeakerSeparationStage
participant resolve_audio
Agent->>VADSegmentationStage: describe()
VADSegmentationStage-->>Agent: "cardinality=1:N fan-out, reads_one_of=[waveform|file]"
Agent->>SpeakerSeparationStage: describe()
SpeakerSeparationStage-->>Agent: "cardinality=1:N fan-out, iteration_key=speaker_id_key"
Agent->>VADSegmentationStage: process(task)
VADSegmentationStage->>resolve_audio: (item, residency, keys)
resolve_audio-->>VADSegmentationStage: (waveform, sample_rate) or None
VADSegmentationStage->>VADSegmentationStage: _build_segment_item x N
VADSegmentationStage-->>Agent: list[AudioTask] per segment
Agent->>SpeakerSeparationStage: process(task)
SpeakerSeparationStage->>resolve_audio: (item, residency, keys)
resolve_audio-->>SpeakerSeparationStage: (waveform, sample_rate) or None
SpeakerSeparationStage->>SpeakerSeparationStage: _build_speaker_tasks (drops audio_filepath, sets original_file)
SpeakerSeparationStage-->>Agent: list[AudioTask] per speaker
%%{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 VADSegmentationStage
participant SpeakerSeparationStage
participant resolve_audio
Agent->>VADSegmentationStage: describe()
VADSegmentationStage-->>Agent: "cardinality=1:N fan-out, reads_one_of=[waveform|file]"
Agent->>SpeakerSeparationStage: describe()
SpeakerSeparationStage-->>Agent: "cardinality=1:N fan-out, iteration_key=speaker_id_key"
Agent->>VADSegmentationStage: process(task)
VADSegmentationStage->>resolve_audio: (item, residency, keys)
resolve_audio-->>VADSegmentationStage: (waveform, sample_rate) or None
VADSegmentationStage->>VADSegmentationStage: _build_segment_item x N
VADSegmentationStage-->>Agent: list[AudioTask] per segment
Agent->>SpeakerSeparationStage: process(task)
SpeakerSeparationStage->>resolve_audio: (item, residency, keys)
resolve_audio-->>SpeakerSeparationStage: (waveform, sample_rate) or None
SpeakerSeparationStage->>SpeakerSeparationStage: _build_speaker_tasks (drops audio_filepath, sets original_file)
SpeakerSeparationStage-->>Agent: list[AudioTask] per speaker
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| writes=IOSpec(data_keys=writes, produces=produces), | ||
| cardinality="1:1 nested-list" if self.nested else "1:N fan-out", | ||
| cardinality_options=["fan_out", "nested"], | ||
| iteration_key=self.segments_key, |
There was a problem hiding this comment.
iteration_key in fan-out mode
iteration_key=self.segments_key is always emitted from describe(), but in fan-out mode (nested=False) no child task ever carries a "segments" key — each child has segment_num, start_ms, etc. An agent that uses this contract to identify which key distinguishes children will look for a non-existent key. The correct value for fan-out is self.segment_num_key; for nested it remains self.segments_key.
| return StageContract( | ||
| reads_one_of=[ | ||
| IOSpec(data_keys=[self.waveform_key, self.sample_rate_key], accepts=["waveform"]), | ||
| IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"]), | ||
| ], | ||
| writes=IOSpec(data_keys=writes, produces=produces), | ||
| cardinality="1:1 nested-list" if self.nested else "1:N fan-out", | ||
| cardinality_options=["fan_out", "nested"], | ||
| iteration_key=self.segments_key, | ||
| gates=Gates(requires_gpu=self.resources.gpus > 0), | ||
| ) |
There was a problem hiding this comment.
"segments" key, so iteration_key should point at segment_num_key instead. Using segments_key here makes the contract inaccurate for agents.
| return StageContract( | |
| reads_one_of=[ | |
| IOSpec(data_keys=[self.waveform_key, self.sample_rate_key], accepts=["waveform"]), | |
| IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"]), | |
| ], | |
| writes=IOSpec(data_keys=writes, produces=produces), | |
| cardinality="1:1 nested-list" if self.nested else "1:N fan-out", | |
| cardinality_options=["fan_out", "nested"], | |
| iteration_key=self.segments_key, | |
| gates=Gates(requires_gpu=self.resources.gpus > 0), | |
| ) | |
| return StageContract( | |
| reads_one_of=[ | |
| IOSpec(data_keys=[self.waveform_key, self.sample_rate_key], accepts=["waveform"]), | |
| IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"]), | |
| ], | |
| writes=IOSpec(data_keys=writes, produces=produces), | |
| cardinality="1:1 nested-list" if self.nested else "1:N fan-out", | |
| cardinality_options=["fan_out", "nested"], | |
| iteration_key=self.segments_key if self.nested else self.segment_num_key, | |
| gates=Gates(requires_gpu=self.resources.gpus > 0), | |
| ) |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| @@ -178,21 +234,39 @@ def _build_speaker_tasks( | |||
| logger.debug(f"Skipping {speaker_id}: duration {result.duration:.2f}s < {self.min_duration}s") | |||
| continue | |||
| spk_waveform, spk_sr = _pydub_to_waveform_sr(result.audio) | |||
| # Drop the parent's file path(s) too: they point at the FULL | |||
| # multi-speaker file, so a file-preferring downstream stage would | |||
| # process the whole file per speaker instead of this speaker's | |||
| # extracted waveform. With the path gone, downstream resolves the | |||
| # per-speaker waveform (input_residency="auto") instead. | |||
| drop_keys = { | |||
| *self._INHERITED_DROP_KEYS, | |||
| self.waveform_key, | |||
| self.duration_key, | |||
| self.audio_filepath_key, | |||
| "audio_filepath", | |||
| } | |||
| speaker_data = { | |||
| **{k: v for k, v in item.items() if k not in self._INHERITED_DROP_KEYS}, | |||
| "waveform": spk_waveform, | |||
| "sample_rate": spk_sr, | |||
| "speaker_id": speaker_id, | |||
| "num_speakers": num_speakers, | |||
| "duration": result.duration, | |||
| "diar_segments": result.diar_segments, | |||
| **{k: v for k, v in item.items() if k not in drop_keys}, | |||
| self.waveform_key: spk_waveform, | |||
| self.sample_rate_key: spk_sr, | |||
| self.speaker_id_key: speaker_id, | |||
| self.num_speakers_key: num_speakers, | |||
| self.duration_key: result.duration, | |||
| self.diar_segments_key: result.diar_segments, | |||
| # Source identity must survive the audio_filepath drop above — | |||
| # TimestampMapper (and any provenance consumer) reads original_file. | |||
| "original_file": item.get("original_file") | |||
| or item.get(self.audio_filepath_key) | |||
| or item.get("audio_filepath") | |||
| or "unknown", | |||
There was a problem hiding this comment.
"original_file" key is hardcoded in SpeakerSeparationStage
VADSegmentationStage exposes original_file_key as a configurable dataclass field (default "original_file"), but SpeakerSeparationStage hardcodes the string "original_file" in both the describe() contract (data_keys=[..., "original_file"]) and the runtime data ("original_file": item.get("original_file") or ...). When these two stages are chained and the user sets a non-default original_file_key in VAD, the speaker-separation stage reads and writes a different key, silently breaking provenance for all downstream consumers.
| audio_result = resolve_audio( | ||
| item, | ||
| residency=self.input_residency, # type: ignore[arg-type] | ||
| audio_filepath_key=self.audio_filepath_key, | ||
| waveform_key=self.waveform_key, | ||
| sample_rate_key=self.sample_rate_key, | ||
| ) | ||
| if audio_result is None: | ||
| return [] |
There was a problem hiding this comment.
The old resolve_waveform_from_item logged an error before returning None. The new resolve_audio call returns None without any log, and the callsite just returns [] with no diagnostic. A task that silently disappears from the pipeline (missing waveform and non-existent file path) will be very hard to debug. Adding a log at the if audio_result is None guard (mirroring _resolve_audio in VADSegmentationStage) would preserve the previous observability contract.
…ker_separation) + match _residency.py
1 participant
Summary
Makes the segmentation stages agent-ready and fixes speaker-separation provenance.
segmentation/speaker_separation.py(1:N fan-out) — provenance fixaudio_filepath(the full mixed file) and carried no source provenance. Now each child drops the parent file path and snapshotsoriginal_file(fromoriginal_file/audio_filepath) before emission, and carries a distinctspeaker_idplus per-speakerwaveform/diar_segments.segmentation/vad_segmentation.pyfan-out(one task per segment) ornested(all segments under one key) modes, configurable keys, anddescribe().Notes for reviewers
"auto"prefers in-memory waveform._build_speaker_tasksinspeaker_separation.py(which parent keys are dropped vs inherited).Test plan
pytest tests/stages/audio/segmentation -q