Fuse document iterate and extract stages

[sarahyurick]

Closes #780.

This change allows the following improvements:

• When tested with limited RAM (--memory=200g), the Common Crawl download and extract pipeline without fusion OOM'ed even when scaling all the way down to 16 CPUs. With fusion, the pipeline was able to succeed with 32 CPUs.
• When removing the RAM limit, the pipeline was faster with fusion than without it. See the following benchmark (using https://github.com/NVIDIA-NeMo/Curator/blob/main/benchmarking/scripts/common_crawl_benchmark.py).

Configuration	Runtime
Fusion + Ray Data executor	27m 49.1s
No Fusion + Ray Data executor	29m 33.1s
Fusion + Xenna executor	54m 4.2s
No Fusion + Xenna executor	66m 6.8s

Common Crawl benchmarks with url_limit=16 and num_cpus=8

[sarahyurick]

Open to opinions about where this class should live and/or any file restructuring ideas.

[greptile-apps]

_{4 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{4 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{4 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{4 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{4 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{4 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[ayushdg]

The tiny stories tutorial (and possibly some other small places) may still have references to the iterate and extract stage separately. beyond that the PR should be good to go

[sarahyurick]

The tiny stories tutorial (and possibly some other small places) may still have references to the iterate and extract stage separately. beyond that the PR should be good to go

@ayushdg hi thanks, I checked the tutorials and found that they do not need updated. This is because the TinyStories tutorial, etc. only access the abstract classes DocumentIterator and DocumentExtractor which are not modified by this PR, and do not interact with DocumentIterateStage or DocumentExtractStage at all.

Lmk if I am missing anything.

[ayushdg]

Lmk if I am missing anything.

Nope I missed that. In that case we should be good to go! Thanks a lot @sarahyurick

[greptile-apps]

_{14 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Record counts wrong
record_limit is enforced based on record_count, but record_count is only incremented after a record is kept (i.e., after extract() and after extracted is None filtering). With an extractor that filters heavily, this will iterate/extract far more than record_limit input records per file (and potentially do a lot more work/memory) before record_count reaches the limit. This is triggered when extractor is set and can return None (e.g., content filters).

[sarahyurick]

A record is only added to the result after extraction. This means that the record count is correct.

[greptile-apps]

Additional Comments (2)

nemo_curator/stages/text/download/base/iterator.py
Empty batch lacks schema
When records is empty (empty task.data, iterate() returns None, all files error, or all records filtered), pd.DataFrame(records) creates a DataFrame with no columns. Downstream stages/tests that rely on declared outputs (e.g., expecting the filename column or extractor output columns) will hit KeyError when accessing missing columns. Consider constructing an empty DataFrame with the expected columns from self.outputs()[1] when records is empty.

---

docs/about/concepts/text/data-acquisition-concepts.md
Incorrect API description
This section says the iterator’s yielded list[dict] is converted to a DataFrame and “passed to Extractor”, and later states “DocumentExtractor works on a Pandas DataFrame”. In the current implementation, DocumentIterateExtractStage calls extractor.extract(record_dict) per-record (a dict), not on a DataFrame. As written, this doc is describing a different API/behavior than users will actually get.