Whenever a stage does def process(self, task: Task) -> list[Task], it is considered a fanout stage. We should add logic to the base ProcessingStage to automatically do this, instead of the child stages having to do it every time.
For example, in https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/text/download/base/url_generation.py we explicitly do:
def ray_stage_spec(self) -> dict[str, Any]:
return {
"is_fanout_stage": True,
}
but since def process(self, task: _EmptyTask) -> list[FileGroupTask] is returning a list, the base ProcessingStage should just have some extra logic to automatically do it instead. That way, we can remove it from our existing stages like URLGenerationStage, and contributors do not need to remember to explicitly set it when creating new stages.
This could be done with get_type_hints or something similar.
We should make sure we are handling cases like def process(self, task: Task) -> Task | list[Task], etc. appropriately too. When we do IS_FANOUT_STAGE, there is a perf cost to it. Under the hood it runs StreamingRepartition which is an expensive operation at scale.
Whenever a stage does
def process(self, task: Task) -> list[Task], it is considered a fanout stage. We should add logic to the baseProcessingStageto automatically do this, instead of the child stages having to do it every time.For example, in https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/text/download/base/url_generation.py we explicitly do:
but since
def process(self, task: _EmptyTask) -> list[FileGroupTask]is returning a list, the baseProcessingStageshould just have some extra logic to automatically do it instead. That way, we can remove it from our existing stages likeURLGenerationStage, and contributors do not need to remember to explicitly set it when creating new stages.This could be done with
get_type_hintsor something similar.We should make sure we are handling cases like
def process(self, task: Task) -> Task | list[Task], etc. appropriately too. When we doIS_FANOUT_STAGE, there is a perf cost to it. Under the hood it runsStreamingRepartitionwhich is an expensive operation at scale.