[FEAT] Automatic Experiment Rescheduling before SLURM timeout

[Grutschus]

On some SLURM clusters, the maximum runtime of an experiment is limited, requiring experiments to be resubmitted to the SLURM queue.
With this PR, I intend to add a feature to SEML that makes this process much easier.

The following example shows how SEML users may use this feature:

import logging
from time import sleep

from seml import Experiment

ex = Experiment()

@ex.reschedule_hook
def reschedule(step: int):
    logging.info(f'Reschedule triggered at step {step}.')
    return {'checkpoint': step}


@ex.automain
def run(
    n_steps: int,
    checkpoint: int | None = None,
):
    logging.info('Starting experiment with the following parameters:')
    logging.info(f'n_steps: {n_steps}, checkpoint: {checkpoint}')

    if checkpoint is not None:
        logging.info(f'Resuming from checkpoint: {checkpoint}')
        # Load your model/state from the checkpoint here

    # Simulate some processing
    for step in range(checkpoint or 0, n_steps):
        reschedule(step)
        logging.info(f'Processing step {step + 1}/{n_steps}')
        sleep(1.0)

    logging.info('Experiment completed successfully.')
    return

From the user's perspective, there are just two new components:

1. The user must register a reschedule-function that is responsible for saving the experiment state before the SLURM timeout is reached. The function must return a dictionary containing updates to the config of the experiment that are required for the next run to pick up the saved experiment state. In practice, I expect this to be a checkpoint, for example.
2. The user must call the registered function regularly from within the main training loop. In contrast to simply propagating a customer signal into the Python program, this ensures that a step has finished before the state is saved, allowing the user to have maximal control over when the reschedule may take place.

What does this implement?

• The slurm_template is updated to create (and delete) a signal file at the reschedule_file path once it receives the timeout signal.
• Experiments in the MongoDB contain three new fields:
  • slurm.{sbatch_id}.reschedule_file: A template for the path experiments should look for the reschedule signal. This template is filled, regardless of whether the SLURM job successfully claims an experiment.
  • execution.reschedule_file: The actual path at which the experiment should look for the signal.
  • reschedule_config_update: The updates to the config due to the reschedule.
• The Experiment class contains the new reschedule_hook that takes care of
  • Finding the signal reschedule_file
  • Watching for the file
  • Calling the user's reschedule hook
  • Rescheduling the experiment

TODOs and Discussion

• Add configurability to the time at which the signal is triggered
• Update the config of the experiment

Limitations

The reschedule_config_update is leaked into the actual experiment's config due to the way that Sacred stores the config in the database.
This, of course, hurts reproducibility, but I have not found an easy way to prevent this. Open for suggestions.

Note: This PR already contains commits from #154 and #155