Add support for Slurm arrays

[sarahyurick]

TODO:

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

[abhinavg4]

Ok I think a lot of the util functions are coming because of this feature, and there might be an easier way for this. Continuing in DMs.

[sarahyurick]

The functions are really only:

1. Write info about failed tasks
2. Use task IDs to filter by Slurm array index

but I can move to util scripts if that makes it easier to read.

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

[abhinavg4]

I think I still have some major comments about the algorithm. Especially at scale, how will the writing and reading work. Will continue in the DMs.

[abhinavg4]

Do we need all these variables, or can they be self-inferred? I think the initial design that @praateekmahajan had in mind was we just add --array=1-100 to the slurm submit command, and everything else works OOTB. Currently, it seems like the effort from the user side is a bit more than that?

[sarahyurick]

So the idea is we want to give the user full control if they need to override total shards (needed for reruns) or minimum shard index (needed to get around any Slurm array size limits). Really to enable Slurm array partitioning, the only thing explicitly needed is:

NEMO_CURATOR_SLURM_ARRAY_ENABLED=1

and it can automatically grab the environment variables without any issues. And then the user can override with NEMO_CURATOR_SLURM_ARRAY_SHARD_INDEX, etc. as desired (but not required).

[abhinavg4]

Ok I think a lot of the util functions are coming because of this feature, and there might be an easier way for this. Continuing in DMs.