Fix job resume on non-preemption type failures

[rayg1234]

See T192775876,

Slurm jobs that are pre-empted currently call a checkpoint callback and get requeued automatically by slurm. Given a timestamp_id (unique job id) and a checkpoint, the job is able to resume correctly AND resume the same run on wandb.

However, if jobs experience a "node failure" or other methods of failure which does not trigger the "checkpoint" callback. The slurm job is missing the timestamp_id and checkpoint info to be able to resume properly. So it ends up starting a new job (and a new wandb run) instead.

Failure example

Example runs that have failed and restarted with new job from scratch when they should've resumed from last checkpoint:

• https://fairwandb.org/fairchem/fm_testing/runs/2024-08-13-01-42-24-aws_1.3B_oc20x22
• https://fairwandb.org/fairchem/fm_testing/runs/2024-08-12-22-19-44-aws_1.3B_oc20x22

This PR modifies

• the task checkpoint loading logic to check for a previously saved checkpoint in the checkpoint dir. This will allow the job to pick up from last checkpoint automatically even if there was a catastrophic failure
• writes the timestamp id in the slurm submission pickle manually. This allows failed jobs to pickup the timestamp_id

Testing

We can use scontrol requeue <job_id> to simulate a "node failure":

Single-node job requeue

Example job that has been requeued where it did not trigger the checkpoint callback:
https://fairwandb.org/fairchem/fm_testing/runs/2024-08-13-00-08-32-test_resume_requeue

Multi-node job requeue

https://fairwandb.org/fairchem/fm_testing/runs/2024-08-13-20-20-16-test_resume_requeue_5

**Note in this case if the last checkpoint was saved at step N and the most recent step written to wandb is N+K, then the job will resume from step N and overwrite K steps in the logs, it will look like the following:
198 wandb: WARNING (User provided step: 544 is less than current step: 676. Dropping entry: {'train/grad_norm': 10.882038116455078, '_timestamp': 1723508783.0449212}).

It is very important to have deterministic training so that when we repeat these steps the jobs end up in the same future state

Pre-emption

Manual preemptions can be triggered by starting a job on a low priority partition (ie: scavenge) and then start a new job targeting the node that the first job was running (ie: using the srun -w <node_id> ...)

Example job that was premepted where the checkpoint callback was triggered:
https://fairwandb.org/fairchem/fm_testing/runs/2024-08-12-23-40-48-test_resume_node_preemption?nw=nwuserrgao

[codecov]

Codecov Report

Attention: Patch coverage is 41.66667% with 14 lines in your changes missing coverage. Please review.

Files	Patch %	Lines
src/fairchem/core/common/slurm.py	35.71%	9 Missing ⚠️
src/fairchem/core/tasks/task.py	66.66%	2 Missing ⚠️
src/fairchem/core/trainers/base_trainer.py	33.33%	2 Missing ⚠️
src/fairchem/core/_cli.py	0.00%	1 Missing ⚠️

Files	Coverage Δ
src/fairchem/core/_cli.py	63.49% <0.00%> (-1.03%)	⬇️
src/fairchem/core/tasks/task.py	66.03% <66.66%> (-1.31%)	⬇️
src/fairchem/core/trainers/base_trainer.py	89.17% <33.33%> (-0.39%)	⬇️
src/fairchem/core/common/slurm.py	35.71% <35.71%> (ø)