1

I often run data processing/machine learning/filesystem-scanning scripts that can take well over 24 hours to complete. For processes with arbitrarily low memory requirements, I can run without using Slurm, but for anything that uses a remotely large amount of memory (>= 1 GB), I have to use Slurm. The exact memory threshold can vary over time, though.

Unfortunately, the way Slurm is configured (and outside of my control), scripts/executables that last over 24 hours are terminated.

With Python, I am aware that most objects can be serialized using the Pickle module, and other forms of data can also be saved/reloaded in different formats. However, many of the scripts I wrote and use are not entirely simple in structure. The stack of function calls (to user-created functions) can be several layers deep, and there are many loops used in the program.

I'm trying to think up an effective solution to this problem that doesn't involve excessive effort in refactoring. So far I just have one:

Solution Part 1: Monitor & re-run with Bash script

Regardless of the changes to the Python code, I think I will have to run the scripts from a bash script in a while loop, check for a specific error code/error output message, and rerun it if it fails due to being terminated by Slurm.

Solution Part 2: Editing Code

Part 2: Checkpoint Objects

I think that one option would be to encapsulate the existing code in a series of if statements along these lines:

from datetime import datetime

checkpoint_1 = Checkpoint(name='checkpoint_1')
checkpoint_1_vals = checkpoint_1.load_locals()
locals().update(**checkpoint_1_vals)
if not checkpoint_1_vals:
    # original code
    t1 = datetime.now()
    # save
    checkpoint_1.save_locals(locals())
print(t1)

and a general implementation

    import dill
    import os
    import pickle
    import types

    # wrapper for checks to avoid overwriting self
    class RDict(dict):
        pass

    class Checkpoint:
        def __init__(
            self, 
            name, 
            checkpoint_folder = 'checkpoints', 
            overwrite=False, 
            logger=None
        ):
            self.name = name
            self.pickle_file = os.path.join(
                checkpoint_folder,
                f'{self.name}.pickle'
            )
            self.overwrite = overwrite
            self.logger=logger
            os.makedirs(checkpoint_folder, exist_ok=True)

        def save_locals(self, temp_locals):
            # for some reason dill.pickles([object of type module]) tends to return True
            # so extra exception has to be made
            valid_locals = RDict({
                k: v 
                for k, v in temp_locals.items()
                if dill.pickles(v) and 
                not isinstance(v, types.ModuleType) and
                not isinstance(v, Checkpoint) and 
                not isinstance(v, RDict)
            })

            # for debugging
            skipped_keys = sorted([k for k in temp_locals.keys() if k not in valid_locals])
            valid_locals['__SKIPPED_KEYS__'] = skipped_keys

            with open(self.pickle_file, 'wb') as f:
                if self.logger:
                    self.logger.info(f'Saving {self.name} locals to {self.pickle_file}')
                pickle.dump(valid_locals, f)

        def load_locals(self):
            if not os.path.isfile(self.pickle_file) or self.overwrite:
                return RDict({})

            with open(self.pickle_file, 'rb') as f:
                temp_locals = RDict({
                    k:v for k,v in pickle.load(f).items()
                })
                if self.logger:
                    self.logger.info(f'Loading {self.name} locals from {self.pickle_file}')
                    
            return temp_locals

For explicitly global variables, I could create a separate file, or make sure that they are not encapsulated within these if statements. nonlocal variables are generally a non-issue, as they are only used within wrappers, and I can usually reference them at the local level if they absolutely need to be part of a checkpoint.

Are there any potential drawbacks to this solution vs. other possible ones?

3

Checkpointing is a great idea because it lets you resume interrupted tasks without much lost time.

But don't be excessively clever about it. Metaprogramming like updating locals() on the fly technically works, but it's non-obvious and can lead to interesting failures. In particular, many objects are not Pickle-able. You're filtering out some types, but that also means resuming a checkpoint will use a slightly different set of locals.

Instead, manually divide your tasks into subtasks that can be checkpointed. I would suggest to put each checkpointable computation into a function that looks somewhat like this:

def some_subtask(args):
  checkpoint_id = derive_id('some_subtask', args)
  try:
    return load_checkpoint(checkpoint_id)
  except CheckpointDoesNotExist:
    pass
  results = some_expensive_computation()
  save_checkpoint(checkpoint_id, results)
  return results

Of course, that memoization pattern can be extracted into a @decorator. The standard library includes @functools.cache() that does a similar thing in-memory, but you will want to persist the data externally, for example in an SQLite database.

Doing checkpointing within the function has the benefit that it's transparent to the callers, and does not impede the flow of your code. Note that you no longer need to pickle all locals, only the results object.

However, it assumes that you can capture all relevant state that affects the expensive computation and include it in the checkpoint ID. In the above example, I assume that some_subtask() is a pure function where the result only depends on the args. This also assumes that the function only computes a results object. If the tasks has externally visible effects such as creating files or sending messages, it might not be possible to checkpoint it meaningfully.

I don't think your first solution (supervise and restart tasks) is directly useful because your scripts might take longer than their deadline – then restarting will just make them run into the deadline again. However, the idea of having a supervisor process that automatically restarts tasks is great. You can integrate this with a checkpoint system so that a job will be restarted as long as it makes progress (i.e. keeps creating new checkpoints).

Not the answer you're looking for? Browse other questions tagged or ask your own question.