feat(audio-alm): make ALM stages agent-ready

nemo_curator/stages/audio/alm/alm_data_builder.py

@@ -26,6 +26,7 @@
 from dataclasses import asdict, dataclass, field
 from typing import Any
 
+from nemo_curator.stages.audio._agent._agent_ready import AgentReady, ConditionalWrite, Gates, IOSpec, StageContract
 from nemo_curator.stages.base import ProcessingStage
 from nemo_curator.tasks import AudioTask
 
@@ -124,7 +125,7 @@ def _record_window_loss(  # noqa: PLR0913
 
 
 @dataclass
-class ALMDataBuilderStage(ProcessingStage[AudioTask, AudioTask]):
+class ALMDataBuilderStage(AgentReady, ProcessingStage[AudioTask, AudioTask]):
     """Build ALM training windows from audio segments.
 
     Filters segments by sample rate, bandwidth, speaker count, and duration
@@ -148,6 +149,13 @@ class ALMDataBuilderStage(ProcessingStage[AudioTask, AudioTask]):
 
     # Top-level fields to drop from output entry (comma-separated)
     drop_fields_top_level: str = "words,segments"
+    audio_filepath_key: str = "audio_filepath"
+    segments_key: str = "segments"
+    audio_sample_rate_key: str = "audio_sample_rate"
+    windows_key: str = "windows"
+    stats_key: str = "stats"
+    truncation_events_key: str = "truncation_events"
+    swift_audio_filepath_key: str = "swift_audio_filepath"
 
     def __post_init__(self) -> None:
         """Compute derived parameters - EXACT match to SDP."""
@@ -157,19 +165,77 @@ def __post_init__(self) -> None:
         self.max_duration = self.target_window_duration + tol
         self._drop_fields_set = {f.strip() for f in self.drop_fields.split(",") if f.strip()}
         self._drop_fields_top_level_set = {f.strip() for f in self.drop_fields_top_level.split(",") if f.strip()}
+        if self.segments_key != "segments" and "segments" in self._drop_fields_top_level_set:
+            self._drop_fields_top_level_set.remove("segments")
+            self._drop_fields_top_level_set.add(self.segments_key)
+        generated_keys = [self.windows_key, self.stats_key, self.truncation_events_key]
+        protected_keys = [
+            self.audio_filepath_key,
+            self.segments_key,
+            self.audio_sample_rate_key,
+            self.swift_audio_filepath_key,
+        ]
+        required_input_keys = [self.audio_filepath_key, self.segments_key, self.audio_sample_rate_key]
+        if len(set(required_input_keys)) != len(required_input_keys):
+            msg = "required input keys must be distinct"
+            raise ValueError(msg)
+        if len(set(generated_keys)) != len(generated_keys) or set(generated_keys) & set(protected_keys):
+            msg = "generated output keys must be distinct from each other and from input keys"
+            raise ValueError(msg)
 
     def inputs(self) -> tuple[list[str], list[str]]:
-        return [], ["audio_filepath", "segments", "audio_sample_rate"]
+        return [], [self.audio_filepath_key, self.segments_key, self.audio_sample_rate_key]
+
+    def outputs(self) -> tuple[list[str], list[str]]:
+        data_keys = [self.windows_key, self.stats_key, self.truncation_events_key]
+        if self.audio_filepath_key not in self._drop_fields_top_level_set:
+            data_keys.append(self.audio_filepath_key)
+        return [], data_keys
+
+    def describe(self) -> StageContract:
+        _, output_keys = self.outputs()
+        generated_keys = {self.windows_key, self.stats_key, self.truncation_events_key}
+        removals = self._drop_fields_top_level_set - generated_keys - {self.audio_filepath_key}
+        conditional_writes = []
+        if self.audio_filepath_key in self._drop_fields_top_level_set:
+            conditional_writes.append(
+                ConditionalWrite(
+                    writes=IOSpec(data_keys=[self.audio_filepath_key]),
+                    condition=(
+                        f"'{self.audio_sample_rate_key}' is below min_sample_rate, so the legacy low-rate branch "
+                        f"preserves '{self.audio_filepath_key}'"
+                    ),
+                    value_origin="upstream_same_key",
+                )
+            )
+        return StageContract(
+            reads=IOSpec(data_keys=[self.audio_filepath_key, self.segments_key, self.audio_sample_rate_key]),
+            optional_reads=IOSpec(data_keys=[self.swift_audio_filepath_key]),
+            writes=IOSpec(data_keys=output_keys),
+            conditional_writes=conditional_writes,
+            # process() rebuilds task.data selectively (words/segments stripped;
+            # the min_sample_rate branch keeps only a four-key subset).
+            preserves_upstream_keys=False,
+            # Named, not merely implied by preserves_upstream_keys: that flag says the entry is
+            # rebuilt, while this says WHICH keys do not survive. Without the names, a downstream
+            # sink check still counts a waveform this stage has already dropped, and refuses a
+            # pipeline that would have run -- the kind of wrong hard gate that pushes a caller
+            # into faking a value to get past it. Derived from the configured parameter so the
+            # declaration tracks whatever the caller actually set.
+            removes_keys=sorted(removals),
+            # Windows are built by sliding over this row's own segments.
+            gates=Gates(per_row_independent=True),
+        )
 
     def process(self, task: AudioTask) -> AudioTask:
         t0 = time.perf_counter()
-        num_segments = len(task.data.get("segments", []))
+        num_segments = len(task.data.get(self.segments_key, []))
         result = self._process_single_entry(task.data)
         task.data.clear()
         task.data.update(result)
         process_time = time.perf_counter() - t0
 
-        num_windows = len(task.data.get("windows", []))
+        num_windows = len(task.data.get(self.windows_key, []))
         self._log_metrics(
             {
                 "process_time": process_time,
@@ -184,26 +250,27 @@ def _process_single_entry(self, entry_data: dict[str, Any]) -> dict[str, Any]:
         """Process a single entry and extract valid training windows."""
         total_truncation_events = 0
 
-        audio_file = entry_data.get("audio_filepath")
-        segments = entry_data.get("segments", [])
+        audio_file = entry_data.get(self.audio_filepath_key)
+        segments = entry_data.get(self.segments_key, [])
         total_dur = sum(seg["end"] - seg["start"] for seg in segments)
 
         stat = BuilderStats(
             total_segments=len(segments),
             total_dur=total_dur,
-            swift_path=entry_data.get("swift_audio_filepath", ""),
-            audio_sample_rate=entry_data.get("audio_sample_rate", 0),
+            swift_path=entry_data.get(self.swift_audio_filepath_key, ""),
+            audio_sample_rate=entry_data.get(self.audio_sample_rate_key, 0),
         )
 
-        if entry_data.get("audio_sample_rate", 0) < self.min_sample_rate:
+        if entry_data.get(self.audio_sample_rate_key, 0) < self.min_sample_rate:
             stat.lost_sr = len(segments)
             stat.dur_lost_sr = total_dur
-            return {
-                "audio_filepath": audio_file,
-                "windows": [],
-                "stats": stat.to_dict(),
-                "truncation_events": total_truncation_events,
+            result = {
+                self.windows_key: [],
+                self.stats_key: stat.to_dict(),
+                self.truncation_events_key: total_truncation_events,
             }
+            result[self.audio_filepath_key] = audio_file
+            return result
 
         valid_windows: list[dict[str, Any]] = []
 
@@ -284,8 +351,8 @@ def _process_single_entry(self, entry_data: dict[str, Any]) -> dict[str, Any]:
             )
 
         result = {k: v for k, v in entry_data.items() if k not in self._drop_fields_top_level_set}
-        result["windows"] = valid_windows
-        result["stats"] = stat.to_dict()
-        result["truncation_events"] = total_truncation_events
+        result[self.windows_key] = valid_windows
+        result[self.stats_key] = stat.to_dict()
+        result[self.truncation_events_key] = total_truncation_events
 
         return result

nemo_curator/stages/audio/alm/alm_data_overlap.py

@@ -24,8 +24,9 @@
 
 import time
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, ClassVar
 
+from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, IOSpec, StageContract
 from nemo_curator.stages.base import ProcessingStage
 from nemo_curator.tasks import AudioTask
 
@@ -149,7 +150,7 @@ def _get_filepath_from_stats(stats: dict[str, Any] | None, key: str) -> str | No
 
 
 @dataclass
-class ALMDataOverlapStage(ProcessingStage[AudioTask, AudioTask]):
+class ALMDataOverlapStage(AgentReady, ProcessingStage[AudioTask, AudioTask]):
     """Filter overlapping ALM windows.
 
     Removes windows with overlap exceeding the threshold, keeping
@@ -161,6 +162,30 @@ class ALMDataOverlapStage(ProcessingStage[AudioTask, AudioTask]):
     # Processing parameters (EXACT match to SDP)
     overlap_percentage: int = 0
     target_duration: float = 120.0
+    windows_key: str = "windows"
+    filtered_windows_key: str = "filtered_windows"
+    stats_key: str = "stats"
+    total_dur_window_key: str = "total_dur_window"
+    total_dur_list_window_key: str = "total_dur_list_window"
+    total_dur_list_window_timestamps_key: str = "total_dur_list_window_timestamps"
+    filtered_key: str = "filtered"
+    filtered_dur_key: str = "filtered_dur"
+    filtered_dur_list_key: str = "filtered_dur_list"
+    manifest_filepath_key: str = "manifest_filepath"
+    swift_filepath_key: str = "swift_filepath"
+
+    INTERNAL_KEY_FIELDS: ClassVar[frozenset[str]] = frozenset(
+        {
+            "total_dur_window_key",
+            "total_dur_list_window_key",
+            "total_dur_list_window_timestamps_key",
+            "filtered_key",
+            "filtered_dur_key",
+            "filtered_dur_list_key",
+            "manifest_filepath_key",
+            "swift_filepath_key",
+        }
+    )
 
     def __post_init__(self) -> None:
         """Validate parameters."""
@@ -170,20 +195,52 @@ def __post_init__(self) -> None:
         if self.target_duration <= 0:
             msg = "target_duration must be positive"
             raise ValueError(msg)
+        output_keys = self.outputs()[1]
+        for key in output_keys:
+            if not isinstance(key, str) or not key.strip():
+                msg = "ALMDataOverlapStage output keys must be non-empty strings"
+                raise ValueError(msg)
+        duplicates = sorted({key for key in output_keys if output_keys.count(key) > 1})
+        if duplicates:
+            msg = f"ALMDataOverlapStage output keys must be distinct: {duplicates}"
+            raise ValueError(msg)
 
     def inputs(self) -> tuple[list[str], list[str]]:
-        return [], ["windows"]
+        return [], [self.windows_key]
+
+    def outputs(self) -> tuple[list[str], list[str]]:
+        return [], [
+            self.total_dur_window_key,
+            self.total_dur_list_window_key,
+            self.total_dur_list_window_timestamps_key,
+            self.filtered_key,
+            self.filtered_windows_key,
+            self.filtered_dur_key,
+            self.filtered_dur_list_key,
+            self.manifest_filepath_key,
+            self.swift_filepath_key,
+        ]
+
+    def describe(self) -> StageContract:
+        _, output_keys = self.outputs()
+        return StageContract(
+            reads=IOSpec(data_keys=[self.windows_key]),
+            optional_reads=IOSpec(data_keys=[self.stats_key]),
+            writes=IOSpec(data_keys=output_keys),
+            # Overlap is measured between windows of the same row.
+            gates=Gates(per_row_independent=True),
+        )
 
     def process(self, task: AudioTask) -> AudioTask:
         t0 = time.perf_counter()
-        input_windows = len(task.data.get("windows", []))
+        input_windows = len(task.data.get(self.windows_key, []))
         result = self._filter_overlaps(task.data)
         task.data.clear()
         task.data.update(result)
         filter_time = time.perf_counter() - t0
 
-        output_windows = len(task.data.get("filtered_windows", []))
-        filtered_dur = task.data.get("filtered_dur", 0.0)
+        output_windows = len(task.data.get(self.filtered_windows_key, []))
+        filtered_dur = task.data.get(self.filtered_dur_key, 0.0)
         self._log_metrics(
             {
                 "filter_time": filter_time,
@@ -199,18 +256,18 @@ def _filter_overlaps(self, entry: dict[str, Any]) -> dict[str, Any]:
         """Filter overlapping windows from entry."""
         threshold = self.overlap_percentage / MAX_OVERLAP_PERCENTAGE
 
-        windows = entry.get("windows", [])
+        windows = entry.get(self.windows_key, [])
         if not windows:
             result = entry.copy()
-            result.setdefault("filtered_windows", [])
-            result.setdefault("filtered_dur", 0.0)
-            result.setdefault("filtered_dur_list", [])
-            result.setdefault("total_dur_window", 0.0)
-            result.setdefault("total_dur_list_window", [])
-            result.setdefault("total_dur_list_window_timestamps", [])
-            result.setdefault("filtered", [])
-            result.setdefault("manifest_filepath", None)
-            result.setdefault("swift_filepath", None)
+            result.setdefault(self.filtered_windows_key, [])
+            result.setdefault(self.filtered_dur_key, 0.0)
+            result.setdefault(self.filtered_dur_list_key, [])
+            result.setdefault(self.total_dur_window_key, 0.0)
+            result.setdefault(self.total_dur_list_window_key, [])
+            result.setdefault(self.total_dur_list_window_timestamps_key, [])
+            result.setdefault(self.filtered_key, [])
+            result.setdefault(self.manifest_filepath_key, None)
+            result.setdefault(self.swift_filepath_key, None)
             return result
 
         total_dur_window = _calculate_total_dur(windows)
@@ -225,18 +282,18 @@ def _filter_overlaps(self, entry: dict[str, Any]) -> dict[str, Any]:
         filtered_dur = _process_filtered_dur(filtered_timestamps)
         filtered_dur_list = _process_filtered_dur_list(filtered_timestamps)
 
-        stats = entry.get("stats", {})
+        stats = entry.get(self.stats_key)
         manifest_filepath = _get_filepath_from_stats(stats, "manifest_path")
         swift_filepath = _get_filepath_from_stats(stats, "swift_path")
 
         result = entry.copy()
-        result["total_dur_window"] = total_dur_window
-        result["total_dur_list_window"] = total_dur_list_window
-        result["total_dur_list_window_timestamps"] = total_dur_list_window_timestamps
-        result["filtered"] = filtered_timestamps
-        result["filtered_windows"] = filtered_windows
-        result["filtered_dur"] = filtered_dur
-        result["filtered_dur_list"] = filtered_dur_list
-        result["manifest_filepath"] = manifest_filepath
-        result["swift_filepath"] = swift_filepath
+        result[self.total_dur_window_key] = total_dur_window
+        result[self.total_dur_list_window_key] = total_dur_list_window
+        result[self.total_dur_list_window_timestamps_key] = total_dur_list_window_timestamps
+        result[self.filtered_key] = filtered_timestamps
+        result[self.filtered_windows_key] = filtered_windows
+        result[self.filtered_dur_key] = filtered_dur
+        result[self.filtered_dur_list_key] = filtered_dur_list
+        result[self.manifest_filepath_key] = manifest_filepath
+        result[self.swift_filepath_key] = swift_filepath
         return result