Support Slurm Autorequeue for Array Jobs

[Queuecumber]

What does this PR do?

This PR adds support for autorequeueing slurm array jobs.

Array jobs are groups of jobs which are grouped in some meaningful way. Slurm provides facilities to quickly (i.e. with less load on the scheduler) create such jobs.

These jobs behave slightly differently from regular jobs.

Although they have SLURM_JOB_ID set, and it is unique for each job in the array, some commands, notably scontrol, differ in their meaning when the JOB_ID is used vs. the array specific SLURM_ARRAY_JOB_ID and SLURM_ARRAY_TASK_ID. For example, if scontrol is called with the SLURM_JOB_ID of the "parent job" it will apply the scontrol command to all array jobs. In order to requeue (for example) only a specific array job, the string <array job id>_<array_task_id> needs to be used.

For example if a job has JOB_ID 1234, the first array job will have ARRAY_JOB_ID 1234 and TASK_ID 1 and the second job will have JOB_ID 1235 ARRAY_JOB_ID 1234 and TASK_ID 2. In this case scontrol requeue 1234 would requeue both jobs in the array. scontrol requeue 1234_1 would requeue only the first job in the array and scontrol requeue 1234_2 would requeue only the second.

Since lightning was only checking the JOB_ID this was causing it to prematurely requeue some array jobs. This was causing lots of weird issues like failed requeues and timeouts with no requeue attempted.

Fixes #15022

Does your PR introduce any breaking changes? If yes, please list them.

None

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

[Queuecumber]

I'm pretty confident in the test cases I wrote but I'm leaving this in draft mode until I get a chance to run it on one of our clusters. Should be before next week.

[Queuecumber]

This is confirmed working on one of our clusters so I think we're good to go

[Queuecumber]

Not sure why the python 3.7 tests are failing, when I run them locally with python 3.7.14 they work

[awaelchli]

@Queuecumber I fixed the tests.
Thanks for extending this feature!!

[awaelchli]

Do you think we should integrate this into SLURMEnvironment class similar to job_id?

[Queuecumber]

Yeah actually the entire function should probably be in the slurm environment since that would keep all the slurm stuff in one place

[carmocca]

Resolve to merge, but you should follow-up on this

[Queuecumber]

Thanks for merging, I definitely could have included this in the PR if you wanted, although it's somewhat of an invasive change there's probably not much risk to it either