feat(audio-metrics): agent-ready WER/SQUIM/bandwidth metrics (contract hygiene)#2175
feat(audio-metrics): agent-ready WER/SQUIM/bandwidth metrics (contract hygiene)#2175shubhamNvidia 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).
…ontract + residency)
Greptile SummaryThis PR wires the
Confidence Score: 4/5Safe to merge after fixing the temp file leak in resolve_audio_path. The hasattr-to-in corrections and contract declarations are well-matched to the actual validate_input and process logic. The one defect is in the new _residency.py: mkstemp creates a file before sf.write runs, and if sf.write raises the temp file is abandoned on disk with no path for the caller to clean it up. nemo_curator/stages/audio/_residency.py — the temp file cleanup path in resolve_audio_path needs a try/except guard around sf.write. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Agent
participant Stage as AudioStage (AgentReady)
participant Registry as _agent_registry
Agent->>Stage: describe()
Stage-->>Agent: StageContract(reads_one_of, writes, gates, ...)
Agent->>Agent: validate chain compatibility
Agent->>Stage: validate_input(task)
Note over Stage: key in task.data (fixed from hasattr bug)
Stage-->>Agent: True / False
alt BATCH_ONLY stage
Agent->>Stage: process_batch(tasks)
Stage-->>Agent: processed tasks
else normal stage
Agent->>Stage: process(task)
Stage-->>Agent: processed task
end
Agent->>Registry: describe_static(StageClass)
Registry-->>Agent: "StageContract(batch_only=True, ...)"
%%{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 AudioStage (AgentReady)
participant Registry as _agent_registry
Agent->>Stage: describe()
Stage-->>Agent: StageContract(reads_one_of, writes, gates, ...)
Agent->>Agent: validate chain compatibility
Agent->>Stage: validate_input(task)
Note over Stage: key in task.data (fixed from hasattr bug)
Stage-->>Agent: True / False
alt BATCH_ONLY stage
Agent->>Stage: process_batch(tasks)
Stage-->>Agent: processed tasks
else normal stage
Agent->>Stage: process(task)
Stage-->>Agent: processed task
end
Agent->>Registry: describe_static(StageClass)
Registry-->>Agent: "StageContract(batch_only=True, ...)"
Reviews (2): Last reviewed commit: "fix(audio): match _residency.py to found..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def describe(self) -> StageContract: | ||
| # process_batch always loads from the audio file; segments only refine | ||
| # WHERE metrics are attached. The filepath is therefore required in | ||
| # every read shape (segments alone are NOT sufficient). | ||
| return StageContract( | ||
| reads_one_of=[ | ||
| IOSpec(data_keys=[self.audio_filepath_key, self.segments_key], accepts=["file"]), | ||
| IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"]), | ||
| ], | ||
| writes=IOSpec(data_keys=[self.metrics_key], segment_data_keys=[self.metrics_key]), | ||
| gates=Gates(requires_gpu=self.resources.gpus > 0, requires_internet_first_run=True), | ||
| ) |
There was a problem hiding this comment.
batch_only not propagated into describe() contract
TorchSquimQualityMetricsStage sets BATCH_ONLY = True and its process() raises NotImplementedError, but describe() returns a StageContract with the default batch_only=False. An agent relying on the instance-level contract from describe() would believe process() is safe to call, and would get a hard NotImplementedError at runtime. The same issue exists in common.py's PreserveByValueStage.describe().
The fix is to pass batch_only=True in the returned StageContract for every stage that sets BATCH_ONLY = True (or to add a helper that merges the ClassVar into the manually-built contract). The _agent_registry auto-derives this field, but direct callers of stage.describe() bypass the registry entirely.
| 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.
batch_only not propagated into describe() contract
PreserveByValueStage sets BATCH_ONLY = True and its process() raises NotImplementedError, but describe() returns a StageContract with the default batch_only=False. An agent using stage.describe() would see batch_only=False and try to call process(), hitting the hard NotImplementedError. The _agent_registry auto-derives batch_only from the BATCH_ONLY ClassVar, but agents or tools that call stage.describe() directly skip the registry path entirely.
Adding batch_only=True to the returned StageContract would fix this — same as the issue in TorchSquimQualityMetricsStage.describe().
| 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
If sf.write(tmp, ...) throws (e.g. invalid waveform data, disk full, unsupported sample rate), the file created by mkstemp is already on disk but register_temp.append(tmp) is never reached. The caller receives the exception with no knowledge of tmp's existence, so cleanup_temp_files never removes it — leaving an unreachable temp file behind on every write failure.
The registration should happen in a try/except block so tmp is either tracked for cleanup or immediately removed on failure.
…pts/reads helpers)
1 participant
Summary
Makes the metrics stages agent-ready and corrects the WER read contract.
metrics/wer.py(ComputeWERStage) —describe()now declares BOTH supported input shapes (it previously under-declared): per-segment WER over asegmentslist, OR top-level WER overhypothesis+referencetext keys on the record. This matches whatvalidate_input/processalready accept. No metric-math change.metrics/squim.py(TorchSquimQualityMetricsStage),metrics/bandwidth.py(BandwidthEstimationStage) — adddescribe()declaring filepath-based reads (plussegments/durationwhere used) and writing results under themetricsdict.squimisBATCH_ONLY(GPU/batched —process()raises,process_batch()runs).Notes for reviewers
describe()matches the keys the stage actually reads/writes.Test plan
pytest tests/stages/audio/metrics -q