feat(audio-io): make IO stages agent-ready (contract + serialization-safe sink)#2171
feat(audio-io): make IO stages agent-ready (contract + serialization-safe sink)#2171shubhamNvidia 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).
Greptile SummaryThis PR makes audio IO stages agent-ready by adding
Confidence Score: 4/5Safe to merge for the common pipeline paths; the changes are additive and backward-compatible, but the The nemo_curator/stages/audio/io/convert.py — the Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Agent
participant Stage
participant AgentReady
participant StageContract
Agent->>Stage: stage.describe()
Stage->>AgentReady: (inherits describe())
Stage-->>Agent: StageContract(reads, writes, gates, cardinality)
Note over Agent,StageContract: AudioToDocumentStage (N:1 sink)
Agent->>Stage: process_batch(tasks)
Stage->>Stage: _sanitize(task.data) per task
Stage->>Stage: _sanitize_nested(segments) if serialize_segments
Stage-->>Agent: DocumentBatch (tensor-free DataFrame)
Note over Agent,StageContract: SegmentExtractionStage (BATCH_ONLY)
Agent->>Stage: process_batch(tasks)
Stage->>Stage: detect_combo(entries) to 2/3/4
Stage->>Stage: _extract_file_segments(entries, ...)
Stage->>Stage: entry[output_key].append(output_path)
Stage-->>Agent: tasks (with output_key list updated in-place)
%%{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
participant AgentReady
participant StageContract
Agent->>Stage: stage.describe()
Stage->>AgentReady: (inherits describe())
Stage-->>Agent: StageContract(reads, writes, gates, cardinality)
Note over Agent,StageContract: AudioToDocumentStage (N:1 sink)
Agent->>Stage: process_batch(tasks)
Stage->>Stage: _sanitize(task.data) per task
Stage->>Stage: _sanitize_nested(segments) if serialize_segments
Stage-->>Agent: DocumentBatch (tensor-free DataFrame)
Note over Agent,StageContract: SegmentExtractionStage (BATCH_ONLY)
Agent->>Stage: process_batch(tasks)
Stage->>Stage: detect_combo(entries) to 2/3/4
Stage->>Stage: _extract_file_segments(entries, ...)
Stage->>Stage: entry[output_key].append(output_path)
Stage-->>Agent: tasks (with output_key list updated in-place)
Reviews (1): Last reviewed commit: "feat(audio-io): make io stages agent-rea..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def __init__( | ||
| self, | ||
| batch_size: int = 64, | ||
| keep_keys: list[str] | None = None, | ||
| drop_keys: tuple[str, ...] = (), | ||
| serialize_segments: bool = False, | ||
| segments_key: str = "segments", | ||
| ) -> None: | ||
| self.batch_size = batch_size | ||
| self.keep_keys = keep_keys | ||
| self.drop_keys = drop_keys | ||
| self.serialize_segments = serialize_segments | ||
| self.segments_key = segments_key |
There was a problem hiding this comment.
self.keep_keys is a list[str], so k not in keys is an O(n) scan on every iteration through data.items(). For a data dict with many keys and a non-trivial keep_keys whitelist, this becomes O(data_keys × keep_keys) per call to _sanitize. Since this runs for every task in every batch, converting to a frozenset at construction time is worth the one-time cost.
| def __init__( | |
| self, | |
| batch_size: int = 64, | |
| keep_keys: list[str] | None = None, | |
| drop_keys: tuple[str, ...] = (), | |
| serialize_segments: bool = False, | |
| segments_key: str = "segments", | |
| ) -> None: | |
| self.batch_size = batch_size | |
| self.keep_keys = keep_keys | |
| self.drop_keys = drop_keys | |
| self.serialize_segments = serialize_segments | |
| self.segments_key = segments_key | |
| def __init__( | |
| self, | |
| batch_size: int = 64, | |
| keep_keys: list[str] | None = None, | |
| drop_keys: tuple[str, ...] = (), | |
| serialize_segments: bool = False, | |
| segments_key: str = "segments", | |
| ) -> None: | |
| self.batch_size = batch_size | |
| self.keep_keys = frozenset(keep_keys) if keep_keys is not None else None | |
| self.drop_keys = frozenset(drop_keys) | |
| self.serialize_segments = serialize_segments | |
| self.segments_key = segments_key |
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!
| def describe(self) -> StageContract: | ||
| return StageContract( | ||
| reads=IOSpec(data_keys=[]), | ||
| writes=IOSpec(data_keys=[]), | ||
| cardinality="N:1", | ||
| # Strips tensors/audio blobs while building the DataFrame, so its | ||
| # output is serialization-safe — the sanctioned sink to place before | ||
| # a JSON writer when a resident tensor may be present. | ||
| gates=Gates(sanitizes_output=True), | ||
| description="Aggregate AudioTasks into a DocumentBatch, stripping tensors/audio blobs (JSON/disk-safe).", | ||
| ) |
There was a problem hiding this comment.
batch_only not reflected in describe() output
BATCH_ONLY = True is set on the class, but the StageContract returned by describe() keeps batch_only=False (the default). The doc comment on StageContract.batch_only says it is "auto-derived at discovery time", so the registry path is correct — but any caller that invokes stage.describe() directly (e.g. a test, a planner that bypasses the registry, or a stage that calls describe() on a sub-stage) receives a contract that says process() is safe to call when it actually raises NotImplementedError. Consider passing batch_only=True here explicitly; the same gap exists in SegmentExtractionStage.describe() and PreserveByValueStage.describe().
| if k in _NON_SERIALIZABLE_KEYS or k == self.segments_key: | ||
| if k == self.segments_key and self.serialize_segments: | ||
| cleaned[k] = self._sanitize_nested(v) | ||
| continue |
There was a problem hiding this comment.
_DROP_VALUE sentinel returned by _sanitize_nested is stored into cleaned without a guard check. If segments_key's value happens to be a raw torch.Tensor (rather than the expected list-of-dicts), _sanitize_nested returns _DROP_VALUE and cleaned[k] becomes an unserializable Python sentinel object, which pandas will store as an object-dtype cell. Any downstream JSON or Parquet writer will then fail — silently violating the sanitizes_output=True contract declared in describe().
| if k in _NON_SERIALIZABLE_KEYS or k == self.segments_key: | |
| if k == self.segments_key and self.serialize_segments: | |
| cleaned[k] = self._sanitize_nested(v) | |
| continue | |
| if k in _NON_SERIALIZABLE_KEYS or k == self.segments_key: | |
| if k == self.segments_key and self.serialize_segments: | |
| nested = self._sanitize_nested(v) | |
| if nested is not _DROP_VALUE: | |
| cleaned[k] = nested | |
| continue |
| def describe(self) -> StageContract: | ||
| return StageContract( | ||
| reads_one_of=[ | ||
| IOSpec(data_keys=["original_file", "original_start_ms", "original_end_ms"], accepts=["file"]), | ||
| IOSpec(data_keys=["original_file", "diar_segments", "speaker_id"], accepts=["file"]), | ||
| ], | ||
| writes=IOSpec(data_keys=[self.output_key], produces=["disk"]), | ||
| gates=Gates(writes_to_disk=True), | ||
| ) |
There was a problem hiding this comment.
reads_one_of contract only covers Combo 2 (timestamp-based) and Combo 3 (diarization with diar_segments + speaker_id), but misses Combo 4 (speaker + timestamps: speaker_id present, diar_segments absent). An agent planning around this contract would incorrectly infer that no speaker_id input is needed for the timestamp combo, or that diar_segments is always required alongside speaker_id.
| def describe(self) -> StageContract: | |
| return StageContract( | |
| reads_one_of=[ | |
| IOSpec(data_keys=["original_file", "original_start_ms", "original_end_ms"], accepts=["file"]), | |
| IOSpec(data_keys=["original_file", "diar_segments", "speaker_id"], accepts=["file"]), | |
| ], | |
| writes=IOSpec(data_keys=[self.output_key], produces=["disk"]), | |
| gates=Gates(writes_to_disk=True), | |
| ) | |
| def describe(self) -> StageContract: | |
| return StageContract( | |
| reads_one_of=[ | |
| # Combo 2: timestamp-based extraction (no speaker info) | |
| IOSpec(data_keys=["original_file", "original_start_ms", "original_end_ms"], accepts=["file"]), | |
| # Combo 3: diarization-based extraction (diar_segments + speaker_id) | |
| IOSpec(data_keys=["original_file", "diar_segments", "speaker_id"], accepts=["file"]), | |
| # Combo 4: speaker + timestamp extraction (speaker_id present, no diar_segments) | |
| IOSpec(data_keys=["original_file", "original_start_ms", "original_end_ms", "speaker_id"], accepts=["file"]), | |
| ], | |
| writes=IOSpec(data_keys=[self.output_key], produces=["disk"]), | |
| gates=Gates(writes_to_disk=True), | |
| ) |
1 participant
Summary
Makes the audio IO stages agent-ready (they now implement
describe()) and hardens the document sink.io/convert.py(AudioToDocumentStage) — N:1 sink, AudioTask → DocumentBatchkeep_keys(whitelist of columns to emit;None= all, prior behavior),drop_keys(blacklist),serialize_segments+segments_key.segments(a list of per-segment dicts that may embed waveform tensors) was previously dropped wholesale. Withserialize_segments=Trueit is kept but recursively cleaned by the new_sanitize_nestedhelper, which walks nested dicts/lists and stripstorch.Tensor/ audio-blob values (using a_DROP_VALUEsentinel so a legitimateNoneis preserved).describe()setsGates(sanitizes_output=True)to mark this as the tensor-safe boundary before a JSON/parquet writer.io/extract_segments.py(SegmentExtractionStage)describe()with OR-shaped reads covering the extraction combos (timestamp-based vs diarization-based). MarkedBATCH_ONLY: extraction runs per batch and appends to a shared metadata file, soprocess()raises andprocess_batch()does the work.Notes for reviewers
keep_keys=Noneemits all keys as before;segmentsstill dropped unlessserialize_segments=True._sanitize/_sanitize_nestedinconvert.py(recursive tensor stripping) and the combo auto-detection inextract_segments.py.