feat(audio-tagging): agent-ready tagging stages (resample temp-path fix + configurable keys)#2178
feat(audio-tagging): agent-ready tagging stages (resample temp-path fix + configurable keys)#2178shubhamNvidia 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).
…ontract + residency)
shubhamNvidia
requested review from
oyilmaz-nvidia
and removed request for
a team
Greptile SummaryThis PR makes audio tagging stages agent-ready by introducing a
Confidence Score: 3/5Safe to merge for most paths, but the write_to_disk=False code path leaks a temp file on post-conversion exceptions before tests exercise it. The bulk of the changes are clean refactors (key-string → field, @staticmethod → method, describe() additions) with no logic mutations. The residency and agent-contract infrastructure is well-structured. The one concrete defect is in the new write_to_disk=False branch of ResampleAudioStage: the output mkstemp file is outside the temp_paths cleanup list, so a failure in load_audio_file or get_audio_duration after ffmpeg succeeds permanently leaks a temp file on the worker. This path is new code, not guarded by a finally block, and the test plan's scope is not yet confirmed to cover it. nemo_curator/stages/audio/tagging/resample_audio.py — specifically the write_to_disk=False branch and the output temp file lifecycle between lines 160 and 215. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Caller
participant ResampleAudioStage
participant resolve_audio_path
participant ffmpeg
participant load_audio_file
participant cleanup_temp_files
Caller->>ResampleAudioStage: process(task)
ResampleAudioStage->>resolve_audio_path: "residency=file/waveform/auto"
alt waveform in task
resolve_audio_path->>resolve_audio_path: mkstemp → write WAV
resolve_audio_path-->>ResampleAudioStage: tmp_input_path (registered in temp_paths)
else file path
resolve_audio_path-->>ResampleAudioStage: local_path
end
alt "write_to_disk=True"
ResampleAudioStage->>ResampleAudioStage: "output = resampled_audio_dir/id.fmt"
else "write_to_disk=False"
ResampleAudioStage->>ResampleAudioStage: mkstemp → output_audio_path (NOT in temp_paths)
end
ResampleAudioStage->>ffmpeg: convert input → output
alt ffmpeg fails
ResampleAudioStage->>cleanup_temp_files: temp_paths (input WAV)
ResampleAudioStage-->>Caller: raise RuntimeError
end
ResampleAudioStage->>cleanup_temp_files: temp_paths (input WAV)
opt keep_waveform_in_task
ResampleAudioStage->>load_audio_file: output_audio_path
Note over ResampleAudioStage: exception here leaks output temp
load_audio_file-->>ResampleAudioStage: waveform, sample_rate
end
ResampleAudioStage->>ResampleAudioStage: get_audio_duration(output_audio_path)
opt not write_to_disk
ResampleAudioStage->>ResampleAudioStage: os.remove(output_audio_path)
end
ResampleAudioStage-->>Caller: task with updated data_entry
%%{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 Caller
participant ResampleAudioStage
participant resolve_audio_path
participant ffmpeg
participant load_audio_file
participant cleanup_temp_files
Caller->>ResampleAudioStage: process(task)
ResampleAudioStage->>resolve_audio_path: "residency=file/waveform/auto"
alt waveform in task
resolve_audio_path->>resolve_audio_path: mkstemp → write WAV
resolve_audio_path-->>ResampleAudioStage: tmp_input_path (registered in temp_paths)
else file path
resolve_audio_path-->>ResampleAudioStage: local_path
end
alt "write_to_disk=True"
ResampleAudioStage->>ResampleAudioStage: "output = resampled_audio_dir/id.fmt"
else "write_to_disk=False"
ResampleAudioStage->>ResampleAudioStage: mkstemp → output_audio_path (NOT in temp_paths)
end
ResampleAudioStage->>ffmpeg: convert input → output
alt ffmpeg fails
ResampleAudioStage->>cleanup_temp_files: temp_paths (input WAV)
ResampleAudioStage-->>Caller: raise RuntimeError
end
ResampleAudioStage->>cleanup_temp_files: temp_paths (input WAV)
opt keep_waveform_in_task
ResampleAudioStage->>load_audio_file: output_audio_path
Note over ResampleAudioStage: exception here leaks output temp
load_audio_file-->>ResampleAudioStage: waveform, sample_rate
end
ResampleAudioStage->>ResampleAudioStage: get_audio_duration(output_audio_path)
opt not write_to_disk
ResampleAudioStage->>ResampleAudioStage: os.remove(output_audio_path)
end
ResampleAudioStage-->>Caller: task with updated data_entry
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| @@ -135,14 +183,36 @@ def process(self, task: AudioTask) -> AudioTask: | |||
| try: | |||
| subprocess.run(cmd, check=True, capture_output=True, text=True) # noqa: S603 | |||
| except subprocess.CalledProcessError as e: | |||
| cleanup_temp_files(temp_paths) | |||
| msg = f"Error converting {input_audio_path}: {e}" | |||
| raise RuntimeError(msg) from e | |||
|
|
|||
| # Update metadata — preserve original URL for cloud paths | |||
| data_entry[self.audio_filepath_key] = original_audio_filepath | |||
| data_entry[self.resampled_audio_filepath_key] = output_audio_path | |||
| # Input temp WAV (materialized from a waveform) is no longer needed after conversion. | |||
| cleanup_temp_files(temp_paths) | |||
|
|
|||
| # Update metadata — preserve original URL for cloud paths. | |||
| if original_audio_filepath is not None: | |||
| data_entry[self.audio_filepath_key] = original_audio_filepath | |||
| if self.write_to_disk: | |||
| data_entry[self.resampled_audio_filepath_key] = output_audio_path | |||
| if self.update_audio_filepath: | |||
| produce_audio_filepath( | |||
| data_entry, | |||
| output_audio_path, | |||
| key=self.audio_filepath_key, | |||
| original_key=self.original_audio_filepath_key, | |||
| ) | |||
| if self.keep_waveform_in_task: | |||
| waveform, sample_rate = load_audio_file(output_audio_path, mono=False) | |||
| data_entry[self.waveform_key] = waveform | |||
| data_entry[self.sample_rate_key] = sample_rate | |||
| duration = get_audio_duration(output_audio_path) | |||
| data_entry[self.duration_key] = duration | |||
| if not self.write_to_disk: | |||
| try: | |||
| os.remove(output_audio_path) | |||
| except OSError: | |||
| pass | |||
There was a problem hiding this comment.
When write_to_disk=False a temp file is created with mkstemp at line 160, but it is never added to temp_paths (which only tracks the input materialized WAV). If load_audio_file() (line 206) or get_audio_duration() (line 209) raises — e.g., a corrupt conversion result or an I/O error — execution never reaches the os.remove(output_audio_path) block at line 213, so the temp file is permanently leaked. The fix is to wrap the waveform-load + duration-get + remove block in a try/finally, or to register output_audio_path in a second temp list just like the input.
| audio_item_id_key: str = "audio_item_id" | ||
| original_audio_filepath_key: str = "original_audio_filepath" | ||
|
|
||
| input_residency: str = "file" |
There was a problem hiding this comment.
input_residency is declared as str but resolve_audio_path expects InputResidency = Literal["file", "waveform", "auto"]. The workaround # type: ignore[arg-type] at the call site masks the mismatch; an agent or user passing an invalid value (e.g. "disk") would get a silent misroute inside resolve_audio_path rather than an early error.
| input_residency: str = "file" | |
| input_residency: InputResidency = "file" |
1 participant
Summary
Makes the tagging stages agent-ready and fixes temp-path safety in resampling.
tagging/resample_audio.py— temp-path fixmkstemppaths registered viaregister_tempand removed bycleanup_temp_files, so parallel workers / repeated calls no longer collide on a shared temp path. Residency-aware input; optional in-memory output (write_to_disk), with the original path preserved when not updating.Others (contract + configurable keys, logic unchanged)
nemo_asr_align.py,merge_alignment_diarization.py,prepare_module_segments.py,split.py,text/itn.py,text/chinese_conversion.py— adddescribe()/StageContractand replace hardcoded data-key strings with configurable fields.Notes for reviewers
NeMoASRAlignerStage.process()delegates toprocess_batch([task])(works per item, so intentionally noBATCH_ONLY).resample_audio.py(creation + cleanup on both success and error).Test plan
pytest tests/stages/audio/tagging -q