Add support for Slurm arrays

[sarahyurick]

TODO:

• Add Slurm array parameters to FilePartitioningStage
• Propagate Slurm array parameters through JsonlReader, ParquetReader, etc.
• Add retry support
• Add FailedTask support
• Add a tutorial
• Add nemo-curator-slurm-cli (not planned for this PR)
• Address case when SLURM_ARRAY_TASK_COUNT > cluster limit
• Add tests

[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.

[sarahyurick]

Discussed with @abhinavg4 . For now the PR keeps track of FailedTask instances by looking for a user-set FAILED_TASKS_DIR_ENV_VAR = "NEMO_CURATOR_FAILED_TASKS_DIR" and writing a JSON file per failed task in the specified directory.

I did the environment variable and write approach because it seems more reliable than trying to handle a global Python variable, etc. And the reason it is an environment variable is so that BaseStageAdapter does not have to propagate an additional parameter for every single stage (which I think would involve having to update the executors as well?). Open to other suggestions.

[praateekmahajan]

Took a super quick look, here are some general thoughts

1. Instead of adding the same 3/4 fields to every "source" stage, can we have a base class and inherit that?
2. Alternatively (or maybe in addition), pipeline.build iirc now dynamically sets the first stage as is_source_stage=True, so can we just rely on those? If we do then inside backends/base.py we can say "if this is a source stage AND slurm is enabled then just use task_id as my key and decide which shard it belongs to"... this is something @abhinavg4 and I had discussed, this reduces the number of changes needed across curator code base, and also generalizes, since source_stage have task_id which is likely assigned using get_determenistic_task_id which is a hash(metadat['source_files'])

[sarahyurick]

Took a super quick look, here are some general thoughts

1. Instead of adding the same 3/4 fields to every "source" stage, can we have a base class and inherit that?
2. Alternatively (or maybe in addition), pipeline.build iirc now dynamically sets the first stage as is_source_stage=True, so can we just rely on those? If we do then inside backends/base.py we can say "if this is a source stage AND slurm is enabled then just use task_id as my key and decide which shard it belongs to"... this is something @abhinavg4 and I had discussed, this reduces the number of changes needed across curator code base, and also generalizes, since source_stage have task_id which is likely assigned using get_determenistic_task_id which is a hash(metadat['source_files'])

For 1, sure.

For 2, we could but it makes this PR dependent on the resumability PR, which is what we were trying to avoid I thought... also, I guess it is not immediately obvious to me how it can work for source stages that are not a FilePartitioningStage. I get the general idea I guess but I am not convinced that it could always work, it sounds to me like how it would probably have to work is convert all unselected tasks to NoneTask maybe?