feat(audio-alm): agent-ready ALM data-builder and pretrain stages

[shubhamNvidia]

Summary

Makes the ALM (audio-language-model) data-pipeline stages agent-ready. Additive only.

• alm/alm_data_builder.py, alm/alm_data_overlap.py, alm/pretrain/extraction.py, alm/pretrain/io.py, alm/pretrain/planning.py — subclass AgentReady and add describe() / StageContract declaring each stage's reads/writes; some hardcoded keys made configurable. Core data-building / planning / IO logic is unchanged.

Notes for reviewers

• Additive, backward compatible; no runtime behavior change expected.
• Worth a look: that each new describe() matches what the stage actually reads/writes.

Test plan

• pytest tests/stages/audio/alm -q

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR makes eight audio pipeline stages agent-ready by subclassing AgentReady and implementing describe() / StageContract, and introduces two new support modules (_agent_ready.py and _residency.py). It also makes previously hardcoded data-key names configurable via constructor fields. Core processing logic is unchanged.

• _agent_ready.py defines the full agent-discovery contract vocabulary (StageContract, IOSpec, Gates, ParamSpec, to_json_schema) and the AgentReady mixin; describe_static() defers to _agent_registry.static_contract which doesn't yet exist, making that entry-point always raise ModuleNotFoundError.
• Two describe() contracts misreport what stages actually read/write: ALMDataBuilderStage omits swift_audio_filepath_key from its reads, and ALMDataOverlapStage omits eight written keys (filtered_dur, manifest_filepath, swift_filepath, etc.) and the stats read that drives them.

Confidence Score: 3/5

The changes are additive and core processing logic is untouched, but the agent-facing contracts have concrete gaps that would produce incorrect stage-chaining decisions.

Two describe() implementations misreport the actual data keys consumed and produced — the gaps in ALMDataOverlapStage are particularly wide (eight undeclared writes). describe_static() will crash unconditionally due to a missing dependency module. These issues directly undermine the stated goal of the PR.

alm_data_overlap.py and alm_data_builder.py have incomplete describe() contracts; _agent_ready.py has a broken describe_static() entrypoint.

Important Files Changed

Filename	Overview
nemo_curator/stages/audio/_agent_ready.py	New file defining the AgentReady mixin and contract dataclasses; describe_static() is broken because _agent_registry.py does not exist, and to_json_schema silently drops None defaults.
nemo_curator/stages/audio/_residency.py	New audio residency helpers; resolve_audio_path leaks a temp file when sf.write raises before the path can be registered for cleanup.
nemo_curator/stages/audio/alm/alm_data_builder.py	ALMDataBuilderStage made AgentReady; describe() reads contract omits swift_audio_filepath_key which _process_single_entry reads.
nemo_curator/stages/audio/alm/alm_data_overlap.py	ALMDataOverlapStage made AgentReady; describe() writes only filtered_windows_key but _filter_overlaps produces 8+ additional keys and reads 'stats' which is undeclared.
nemo_curator/stages/audio/alm/pretrain/extraction.py	SnippetExtractionStage made AgentReady; describe() correctly declares cardinality, iteration_key, gates, and dry_run-conditional disk writes.
nemo_curator/stages/audio/alm/pretrain/io.py	Three stages made AgentReady with correct gates and cardinality declarations.
nemo_curator/stages/audio/alm/pretrain/planning.py	Four stages made AgentReady; contracts otherwise match implementation.
nemo_curator/stages/audio/common.py	Five common stages made AgentReady with correct IOSpec, cardinality, and gates declarations.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Agent
    participant AgentReady
    participant Stage
    participant _agent_registry

    Agent->>AgentReady: describe_static()
    AgentReady->>_agent_registry: static_contract(cls)
    Note over _agent_registry: Module does not exist - ModuleNotFoundError

    Agent->>Stage: stage.describe()
    Stage-->>Agent: StageContract(reads, writes, gates, ...)

    Agent->>Agent: Chain stages by matching IOSpec keys
    Note over Agent: ALMDataBuilderStage omits swift_audio_filepath_key from reads
    Note over Agent: ALMDataOverlapStage omits 8+ writes

    Agent->>Stage: stage.process(task)
    Stage-->>Agent: AudioTask (modified data)

Loading

%%{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 AgentReady
    participant Stage
    participant _agent_registry

    Agent->>AgentReady: describe_static()
    AgentReady->>_agent_registry: static_contract(cls)
    Note over _agent_registry: Module does not exist - ModuleNotFoundError

    Agent->>Stage: stage.describe()
    Stage-->>Agent: StageContract(reads, writes, gates, ...)

    Agent->>Agent: Chain stages by matching IOSpec keys
    Note over Agent: ALMDataBuilderStage omits swift_audio_filepath_key from reads
    Note over Agent: ALMDataOverlapStage omits 8+ writes

    Agent->>Stage: stage.process(task)
    Stage-->>Agent: AudioTask (modified data)

Loading

Comments Outside Diff (4)

1. nemo_curator/stages/audio/alm/alm_data_builder.py, line 533-540 (link)

   describe() reads contract is incomplete — swift_audio_filepath_key not declared

   _process_single_entry reads self.swift_audio_filepath_key (to populate stat.swift_path), but the reads IOSpec only lists audio_filepath_key, segments_key, and audio_sample_rate_key. An agent that chains based solely on the declared contract will not know this key is consumed. The PR author flagged exactly this class of mismatch as worth verifying.

2. nemo_curator/stages/audio/alm/alm_data_overlap.py, line 643-647 (link)

   describe() writes contract severely underreports what _filter_overlaps actually produces

   The contract only declares filtered_windows_key, but _filter_overlaps always writes at least eight additional keys: total_dur_window, total_dur_list_window, total_dur_list_window_timestamps, filtered, filtered_dur, filtered_dur_list, manifest_filepath, and swift_filepath. Agents chaining on any of these will not see them in the contract and may mis-sequence the pipeline. The stage also reads "stats" from the entry which is absent from reads.

3. nemo_curator/stages/audio/_agent_ready.py, line 307-319 (link)

   describe_static() will always raise ModuleNotFoundError — _agent_registry does not exist

   describe_static() lazily imports static_contract from nemo_curator.stages.audio._agent_registry, but that module is not present in the repository. Any call to cls.describe_static() on any AgentReady subclass will fail at runtime. If the registry module is intentionally deferred to a follow-up PR, the method body should guard against the missing import (e.g. raise NotImplementedError) to make the gap explicit rather than silently crashing.

4. nemo_curator/stages/audio/_residency.py, line 462-467 (link)

   Temp file leaked if sf.write raises

   mkstemp creates the file and closes the fd before sf.write is called. If sf.write raises (e.g. bad dtype, soundfile codec error), the temp file remains on disk permanently — it is never appended to register_temp and the caller has no path to clean up. Wrapping the write in a try/except that deletes tmp on error would prevent the leak.

_{Reviews (1): Last reviewed commit: "feat(audio-alm): make alm stages agent-r..." | Re-trigger Greptile}

[greptile-apps]

to_json_schema silently drops default: None from the JSON schema

The guard if p.default is not None skips emitting "default": null for parameters whose true default is None. An agent reading the schema cannot distinguish "no default declared" from "default is None", and may incorrectly treat those parameters as required. Using a sentinel (e.g. _MISSING = object()) in ParamSpec to represent "no default" would preserve the distinction.

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!