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 the need for manual ray_stage_spec() overrides in ProcessingStage subclasses by inferring is_fanout_stage from the return-type annotation of process() at runtime, using get_type_hints plus a recursive _annotation_contains_list helper that covers both parameterized list[T] and union types (X | list[T], Union[X, list[T]]). The URLGenerationStage override — which was the canonical example — is removed as a direct consequence.

• base.py: Adds is_fanout_stage() and _annotation_contains_list() to ProcessingStage; ray_stage_spec() now always emits {"is_fanout_stage": <bool>} instead of {}. Stages that called return super().ray_stage_spec() (e.g., identification.py, identify_duplicates.py) will now propagate this key, but since all Ray backend consumers use .get(key, False), runtime behavior is unchanged.
• url_generation.py: Manual override deleted; URLGenerationStage.process annotated as list[FileGroupTask] so the base inference naturally yields True.
• test_base.py: Two new mock stages and two positive-case tests for list and union-list return types; the non-fanout (False) path is not yet tested.

Confidence Score: 4/5

Safe to merge; the inference logic is correct for the annotated cases in the codebase, and all Ray backend consumers already use .get(key, False) so the added is_fanout_stage key in the base dict does not change runtime dispatch for existing stages.

The core change is a clean, well-scoped addition. The only gaps are a missing non-fanout regression test and silent swallowing of hint-resolution failures with no diagnostic output — both quality concerns rather than correctness defects on the current code paths.

The new exception handler in base.py (lines 310–313) deserves a second look: if a stage annotates process() with a type only available under TYPE_CHECKING, the stage will silently behave as non-fanout with no log message, which may be hard to diagnose in production pipelines.

Important Files Changed

Filename	Overview
nemo_curator/stages/base.py	Adds is_fanout_stage() inference from process() return-type annotation, changes ray_stage_spec() default from {} to {"is_fanout_stage": ...}; the silent fallback on NameError/TypeError has no diagnostic logging
nemo_curator/stages/text/download/base/url_generation.py	Removes manual ray_stage_spec() override now that base class infers is_fanout_stage=True from the list[FileGroupTask] return annotation; correct and clean removal
tests/stages/common/test_base.py	Adds MockFanoutStage and MockOptionalFanoutStage with focused tests for list and union-list return types; missing a regression test for the non-fanout (False) path

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["ray_stage_spec() called on ProcessingStage"] --> B["is_fanout_stage()"]
    B --> C["get_type_hints(type(self).process)"]
    C --> D{Exception?}
    D -- "AttributeError / NameError / TypeError" --> E["return False (silent)"]
    D -- "Success" --> F["process_hints.get('return')"]
    F --> G["_annotation_contains_list(annotation)"]
    G --> H{annotation is None?}
    H -- Yes --> I["return False"]
    H -- No --> J{get_origin is list?}
    J -- Yes --> K["return True fanout"]
    J -- No --> L{get_origin in UnionType or Union?}
    L -- Yes --> M["recurse on each Union arg"]
    M --> G
    L -- No --> N["return False"]
    K --> O["ray_stage_spec returns is_fanout_stage True"]
    E --> P["ray_stage_spec returns is_fanout_stage False"]
    I --> P
    N --> P

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 on ProcessingStage"] --> B["is_fanout_stage()"]
    B --> C["get_type_hints(type(self).process)"]
    C --> D{Exception?}
    D -- "AttributeError / NameError / TypeError" --> E["return False (silent)"]
    D -- "Success" --> F["process_hints.get('return')"]
    F --> G["_annotation_contains_list(annotation)"]
    G --> H{annotation is None?}
    H -- Yes --> I["return False"]
    H -- No --> J{get_origin is list?}
    J -- Yes --> K["return True fanout"]
    J -- No --> L{get_origin in UnionType or Union?}
    L -- Yes --> M["recurse on each Union arg"]
    M --> G
    L -- No --> N["return False"]
    K --> O["ray_stage_spec returns is_fanout_stage True"]
    E --> P["ray_stage_spec returns is_fanout_stage False"]
    I --> P
    N --> P

Loading

_{Reviews (1): Last reviewed commit: "feat: infer fanout stages from process a..." | 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.