Add supporting max_calls to the ray data executor

[ayushdg]

Description

This PR
Fixes #1533

• Adds support for stages to set max_calls that are propagated if it's a task based stage (not actor based).
• Exposes it for the iterate extract stage so that we can override it for justext extraction. Currently exposed as a constructor arg but open to cleaner ideas. Discussed a bit in Explore making stage_spec overridable during stage creation #1532.
• Adds tests including some tests that ensure tasks with max_calls are not fused with other tasks and lose this information.

Usage

# Add snippet demonstrating usage

Checklist

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

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[ayushdg]

/ok to test d7cf34d

[greptile-apps]

Greptile Summary

This PR adds support for the max_calls parameter in Ray Data executor to control worker process lifecycle, addressing memory fragmentation issues in long-running extraction tasks.

Key Changes:

• Added MAX_CALLS_PER_WORKER to RayStageSpecKeys enum for stage configuration
• Modified Ray Data adapter to extract and apply max_calls for task-based stages only (not actor-based stages)
• Added max_calls_per_worker parameter to DocumentIterateExtractStage to enable per-stage control
• Exposed extractor_max_calls_per_worker in DocumentDownloadExtractStage and CommonCrawlDownloadExtractStage
• Set intelligent default of max_calls=2 for JusText extractor with informative logging about memory fragmentation risks
• Comprehensive test suite verifies PID recycling behavior and ensures stages with max_calls are not fused by Ray Data's optimizer

Implementation Quality:

• Correctly applies max_calls only to task-based stages, not actor stages
• Tests verify both functional behavior (PID recycling) and Ray Data execution plan structure (operator fusion prevention)
• Good documentation with inline comments explaining the purpose

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• The implementation is well-designed with proper separation between task and actor stages, comprehensive test coverage including edge cases, and intelligent defaults for known problematic extractors. All changes are additive and backwards-compatible.
• No files require special attention

Important Files Changed

Filename	Overview
nemo_curator/backends/experimental/ray_data/adapter.py	Correctly extracts and applies max_calls from stage spec for task-based stages only
nemo_curator/stages/text/download/base/iterator.py	Adds max_calls_per_worker parameter and implements ray_stage_spec() method correctly
nemo_curator/stages/text/download/common_crawl/stage.py	Sets intelligent default of 2 for JusText extractor with helpful logging
tests/backends/experimental/ray_data/test_max_calls_pid.py	Comprehensive test suite verifying PID recycling and operator fusion behavior

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[ProcessingStage with max_calls_per_worker] --> B{Is Actor Stage?}
    B -->|Yes| C[Create Actor<br/>Ignore max_calls]
    B -->|No| D[Create Task<br/>Extract max_calls from ray_stage_spec]
    D --> E{max_calls != None?}
    E -->|Yes| F[Add max_calls to concurrency_kwargs<br/>Prevents operator fusion]
    E -->|No| G[No max_calls set<br/>Allow operator fusion]
    F --> H[Ray Data map_batches with max_calls]
    G --> I[Ray Data map_batches without max_calls]
    H --> J[Worker restarts every N tasks<br/>Mitigates memory fragmentation]
    I --> K[Worker reused indefinitely<br/>May fuse with other stages]

Loading

_{Last reviewed commit: cd138cb}

[greptile-apps]

_{10 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[sarahyurick]

How did you decide on 2?

[ayushdg]

Somewhat arbitrary but verified it works a snapshot.
On average the higher this value the higher a chance for an OOM through memory fragmentation.
A lower value like 1 will rotate the PID every warc file which has some overhead to kill and spawn a new process.

[sarahyurick]

LGTM!