feat: infer fanout stages from process annotations

[nightcityblade]

Description

closes #1613

• infer is_fanout_stage from the annotated process() return type in the base ProcessingStage
• treat both list[Task] and unions that include list[...] as fanout-capable outputs
• drop the manual ray_stage_spec() override from URLGenerationStage now that the base behavior covers it
• add focused regression tests for the inferred fanout behavior

Usage

class URLGenerationStage(ProcessingStage[_EmptyTask, FileGroupTask]):
    def process(self, task: _EmptyTask) -> list[FileGroupTask]:
        ...

assert URLGenerationStage(...).ray_stage_spec()["is_fanout_stage"] is True

Checklist

• I am familiar with the Contributing Guide.
• New or Existing tests cover these changes.
• The documentation is up to date with these changes.

Validation run locally:

• python3 -m ruff check nemo_curator/stages/base.py nemo_curator/stages/text/download/base/url_generation.py tests/stages/common/test_base.py tests/stages/text/download/base/test_url_generation.py
• python3 -m pytest --noconftest tests/stages/common/test_base.py tests/stages/text/download/base/test_url_generation.py -q (blocked in this environment: cosmos_xenna is not installed for package import; the full test setup also requires ray via tests/conftest.py)

[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 eliminates boilerplate ray_stage_spec() overrides across ~13 stages by teaching the base ProcessingStage to infer is_fanout_stage from the annotated return type of process(). The new is_fanout_stage() method uses get_type_hints and _annotation_contains_list to detect both plain list[T] and union types that include a list arm (e.g. T | list[T]).

• base.py gains is_fanout_stage() and _annotation_contains_list(), and ray_stage_spec() now returns {\"is_fanout_stage\": self.is_fanout_stage()} instead of {}.
• All removed ray_stage_spec() overrides are confirmed to have properly annotated process() methods returning list[...]; tests are added for each affected stage to assert the inferred value.
• ClipTranscodingStage needed its annotation corrected from VideoTask to VideoTask | list[VideoTask] to match its actual branching behavior and trigger the detection.

Confidence Score: 5/5

Safe to merge — no functional regressions identified across the 13 affected stages.

All stages that had ray_stage_spec() overrides removed have properly annotated process() methods returning list[...], and the backend consumer already uses .get(IS_FANOUT_STAGE, False) so the key-absent edge cases (nested VAD, VideoFrameExtraction) are handled safely. The two-line annotation inference logic is straightforward and well-tested.

clip_extraction_stages.py is the only changed fanout stage without a new ray_stage_spec regression test.

Important Files Changed

Filename	Overview
nemo_curator/stages/base.py	Core change: adds is_fanout_stage(), _annotation_contains_list(), and updates ray_stage_spec() to return the inferred key. Logic is correct for list[T], `T
nemo_curator/stages/audio/segmentation/vad_segmentation.py	Replaces hardcoded {IS_FANOUT_STAGE: True} with super().ray_stage_spec() for the non-nested path. The nested=True path still returns {} (intentional, same as before); process() is annotated `AudioTask
nemo_curator/stages/video/clipping/clip_extraction_stages.py	Return type corrected from VideoTask to `VideoTask
tests/stages/common/test_base.py	Adds MockFanoutStage, MockOptionalFanoutStage, and TestProcessingStageFanoutDetection covering plain-list, union-with-list, and single-item return types; both prior reviewer notes are addressed.
nemo_curator/stages/text/download/base/url_generation.py	Removes the manual ray_stage_spec() override (the motivating example from the PR description); process() returns list[FileGroupTask] so inference is correct.
nemo_curator/stages/file_partitioning.py	Removes the manual ray_stage_spec() override; process() returns list[FileGroupTask] so inference produces the same is_fanout_stage: True result.
nemo_curator/stages/audio/alm/pretrain/extraction.py	Removes the manual ray_stage_spec() override; process() returns list[AudioTask] so inference produces the same is_fanout_stage: True result. Test added to confirm.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["ray_stage_spec() called"] --> B["Call is_fanout_stage()"]
    B --> C["get_type_hints(type(self).process)"]
    C -->|"Exception (AttributeError/NameError/TypeError)"| D["logger.debug + return False"]
    C -->|"Success"| E["hints.get('return')"]
    E --> F["_annotation_contains_list(annotation)"]
    F -->|"annotation is None"| G["return False"]
    F -->|"get_origin is list"| H["return True"]
    F -->|"get_origin in (UnionType, Union)"| I["Recurse on each arg via get_args"]
    I -->|"any arg contains list"| H
    I -->|"no arg contains list"| G
    F -->|"anything else"| G
    H --> J["{is_fanout_stage: True}"]
    G --> K["{is_fanout_stage: False}"]
    D --> K

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"}}}%%
flowchart TD
    A["ray_stage_spec() called"] --> B["Call is_fanout_stage()"]
    B --> C["get_type_hints(type(self).process)"]
    C -->|"Exception (AttributeError/NameError/TypeError)"| D["logger.debug + return False"]
    C -->|"Success"| E["hints.get('return')"]
    E --> F["_annotation_contains_list(annotation)"]
    F -->|"annotation is None"| G["return False"]
    F -->|"get_origin is list"| H["return True"]
    F -->|"get_origin in (UnionType, Union)"| I["Recurse on each arg via get_args"]
    I -->|"any arg contains list"| H
    I -->|"no arg contains list"| G
    F -->|"anything else"| G
    H --> J["{is_fanout_stage: True}"]
    G --> K["{is_fanout_stage: False}"]
    D --> K

Loading

_{Reviews (5): Last reviewed commit: "Merge branch 'main' into fix/issue-1613-..." | Re-trigger Greptile}

[greptile-apps]

Silent False return with no diagnostic logging makes misclassification hard to debug. If a stage's process() annotation references a type that's only imported under TYPE_CHECKING (a common pattern), get_type_hints() raises NameError, the exception is swallowed, and the stage silently behaves as non-fanout — with no indication of why. A logger.debug here would make such failures visible.

Suggested change

	try:
	process_hints = get_type_hints(type(self).process)
	except (AttributeError, NameError, TypeError):
	return False
	try:
	process_hints = get_type_hints(type(self).process)
	except (AttributeError, NameError, TypeError) as exc:
	logger.debug(
	"Could not resolve type hints for {}.process; defaulting is_fanout_stage to False: {}",
	type(self).__name__,
	exc,
	)
	return False

[greptile-apps]

There is no test asserting that is_fanout_stage() returns False for a stage whose process() returns a single item, nor that ray_stage_spec() emits {"is_fanout_stage": False} in that case. Without this regression test, a future change that accidentally makes _annotation_contains_list always return True would go undetected.

Suggested change

	class TestProcessingStageFanoutDetection:
	def test_base_ray_stage_spec_marks_list_outputs_as_fanout(self):
	stage = MockFanoutStage()
	class TestProcessingStageFanoutDetection:
	def test_base_ray_stage_spec_marks_non_list_outputs_as_non_fanout(self):
	stage = MockStageA()
	assert stage.is_fanout_stage() is False
	assert stage.ray_stage_spec()["is_fanout_stage"] is False
	def test_base_ray_stage_spec_marks_list_outputs_as_fanout(self):
	stage = MockFanoutStage()

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!

[sarahyurick]

Hi @nightcityblade I will echo previous feedback from other PRs you opened for this issue:

"""
Let's remove the ray_stage_spec function from the actual stages and make sure the tests still pass. The ones I found:

• https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/text/download/base/url_generation.py
• https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/image/io/image_reader.py
• https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/audio/datasets/fleurs/create_initial_manifest.py
• https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/audio/common.py
• https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/video/clipping/clip_extraction_stages.py
• https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/audio/alm/pretrain/extraction.py
• https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/audio/segmentation/speaker_separation.py
• https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/file_partitioning.py
• https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/audio/datasets/readspeech/create_initial_manifest.py
• https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/audio/alm/pretrain/io.py
• https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/deduplication/semantic/pairwise_io.py

• Curator/nemo_curator/stages/audio/segmentation/vad_segmentation.py

  Line 112 in c0a02ad

  return {RayStageSpecKeys.IS_FANOUT_STAGE: True}

Do all of these work with this PR? Are there any that should be considered exceptions/explicitly keep ray_stage_spec for some reason?
"""

[nightcityblade]

Thanks for flagging the broader set of ray_stage_spec() overrides. I’ve reviewed the list and agree the right next step here is to audit each remaining override against the new inference path rather than assuming url_generation.py is the only safe removal.

I haven’t pushed on this PR yet because I want to verify that list carefully and keep the branch focused on overrides that can be removed without changing behavior. I’ll follow up on this PR once I’ve finished that pass.

[nightcityblade]

Thanks — I completed the audit and pushed aa06659.

• Removed the redundant ray_stage_spec() fanout overrides from all stages whose process() return annotations already carry the information, including the newer Omni and PDF partitioning stages added since the original PR.
• Kept the two necessary overrides: VAD still disables fanout dynamically when nested=True, and clip transcoding still adds its Ray CPU setting while delegating fanout detection to the base implementation. I corrected clip transcoding's return annotation to VideoTask | list[VideoTask] to match its actual returns.
• Added the non-fanout regression case and debug logging when type-hint resolution falls back to False.

Validation:

• ruff check passed on all 16 changed files.
• ruff format --check passed on all 16 changed files.
• Targeted tests: 268 passed, 5 deselected, 6 skipped. These were run without the Linux-only shared Ray fixture because the available host is macOS; the one /dev/shm integration case was deselected, while the relevant stage/unit suites all passed.

[sarahyurick]

Thanks @nightcityblade ! Can you also go through the respective pytest files for each stage you changed, see if it already has a test to assert that the stage is a fanout stage, and if not add a test to make sure it is still marked as fanout?

[nightcityblade]

Thanks for catching this. I rebased onto current main and added direct stage.fanout assertions in the existing pytest modules for the seven stages that were missing them: SnippetExtractionStage, ReadLongFormManifestStage, ManifestReaderStage, SpeakerSeparationStage, ClusterWiseFilePartitioningStage, ImageReaderStage, and HFDatasetImageReaderStage.

The other five removed overrides already had direct assertions in their own modules, and the VAD/transcoding dynamic exceptions retain their dedicated fanout coverage.

Validation completed:

• Ruff check and format check pass for all changed test files
• Runnable changed-module suites: 116 passed, 2 platform-specific tests deselected
• Existing Fleurs, ReadSpeech, FilePartitioning, and VAD coverage: 5 passed
• The RAPIDS-only stage assertion also passes when loaded directly; its pytest module cannot import on macOS without RAPIDS

[sarahyurick]

/ok to test 6066959

[sarahyurick]

/ok to test 56f7b1d