Idea for RFC to add task checkpointing so interrupted tasks can resume from checkpoints

[mwiebe]

Background

Many long-running workloads have the ability to save their state and then resume processing from a different host. For example, a physics simulation, a high-sample-count final frame, fine-tuning an AI model, or optimizing a Gaussian splat scene can run for hours or days. If the scheduler doesn't track those checkpoints, an interruption partway through loses all of that progress and the task restarts from the beginning. Checkpointing can improve overall utilization on interruptible hosts like EC2 Spot, and avoid lost work when taking hosts down for maintenance. A task running a chunk of frames under the TASK_CHUNKING extension has a similar gap: it has no way to tell the scheduler that some of the tasks within the chunk are already complete.

These workloads would benefit from a portable way to communicate checkpoint status and data to a scheduler, so that any system running the job can preserve and restore that progress instead of restarting from scratch.

The RFC idea

The idea has two parts: the job template declares that a step's tasks support checkpointing, and the task and scheduler communicate about checkpoints in both directions while the job runs. The task reports checkpoint state and the files it involves, portions of a chunk that are already complete, and output data that is ready before the task finishes — which a worker agent could upload in a pipelined fashion while the task keeps running. The scheduler, in turn, signals the task when it needs to checkpoint, and hands the saved state back when it resumes the task.

Periodic vs signaled checkpoints

There are two models a job could adopt using these mechanisms. It could save checkpoint data and create a checkpoint every once in a while, say every five minutes — this also protects against crashes and lost workers, where no warning arrives. Or, it could use NOTIFY_THEN_TERMINATE in the checkpointCancelation property described below, and save a checkpoint when it receives that signal, capturing progress right up to the moment of interruption. The two compose naturally, as the example below shows.

Shared file systems vs managed file data flow

Compute farms also differ in how file data reaches the worker hosts. When tasks run against a shared file system, the file path metadata the job provides is largely informational — the checkpoint files are already visible to whichever host resumes the task. When a system
explicitly manages the data flow of files, that same metadata defines what the scheduler should save for the checkpoint and provide for resumption. The design is structured so that job templates can simultaneously support both models with little effort.

Messages from the task to the scheduler

The foundation is to let the job template provide status updates and metadata about checkpoints. We define a new set of stdout/stderr messages following the existing stdout/stderr message protocol (openjd_env: et al.), which a running task emits to report its progress:

1. openjd_checkpoint_files: — declares files and directories that are available as part of
   the task's output so far, forming the checkpoint dataset. It accepts lists of file paths and
   directory paths under separate files and dirs keys, with an option to accumulate onto the
   current dataset or to reset it first. Because these messages arrive while the task is still
   running, a worker agent can begin processing the listed files immediately, pipelining output
   transfer with compute.

   openjd_checkpoint_files: {"files": ["/var/session-abc123/output/partial.exr"], "dirs": ["/var/session-abc123/scratch/restart/"], "mode": "ACCUMULATE"}
   

2. openjd_task_complete: — reports that a particular task within a chunk is complete, as
   a key/value pair whose key is the name of the step's
   CHUNK[INT]
   task parameter. The scheduler can mark that task complete, so that if the chunk is
   interrupted, the rescheduled remainder excludes the finished work. For a step whose chunked
   task parameter is named Frame:

   openjd_task_complete: {"Frame": 12}
   

3. openjd_checkpoint_save: — commits a resumable checkpoint, and implicitly includes the
   file dataset declared via openjd_checkpoint_files up to that point. It accepts an optional
   "state" key with a metadata string chosen by the task, and an optional "file" key with
   the absolute path of a checkpoint file — many applications have exactly one such file, and
   requiring a separate openjd_checkpoint_files message for it would be annoying. A task emits
   this message periodically and/or in response to the checkpoint signal; each save supersedes
   the previous one.

   openjd_checkpoint_save: {"state": "3,2,3", "file": "/var/session-abc123/scratch/checkpoint.rman"}
   

4. openjd_checkpoint_resume_status: — valid only within the checkpoint resume action;
   reports whether the task actually used the checkpoint data (RESUMED) or discarded it and
   started from scratch (DISCARDED), so the scheduler knows what the resumed run did.

   openjd_checkpoint_resume_status: {"status": "RESUMED"}
   

The onCheckpointResume action

A step declares that its tasks are resumable by adding an onCheckpointResume action to its <StepActions>, complementing onRun and using the same <Action> schema. When a previously checkpointed task is scheduled again, the scheduler runs onCheckpointResume instead of onRun. A step with no onCheckpointResume is not resumable and behaves exactly as today.

Two new value references are available only within this action's format strings, carrying the values from the task's last openjd_checkpoint_save message (PascalCased like other spec-defined references, e.g. WrappedAction.Command):

• Task.Checkpoint.State — the state metadata string.
• Task.Checkpoint.File — the location where the scheduler restored the checkpoint file, which
  may differ from its original absolute path when the task resumes on a different worker host.
  The scheduler likewise restores the rest of the checkpoint's file dataset before the action
  runs.

The checkpointCancelation property

A new checkpointCancelation property on <Action> is the "suspend" counterpart to the existing cancelation property. When a scheduler interrupts a task with the intent to resume it, it follows this method instead of cancelation. Both properties share the same <CancelationMethod> shape, with its TERMINATE and NOTIFY_THEN_TERMINATE modes. To support this, we would generalize <CancelationMethod> with a new signalToNotify field (default SIGTERM), because applications differ in which signal triggers a checkpoint — one application might checkpoint on SIGINT or SIGTERM, while another expects SIGHUP or SIGUSR1 — and the same flexibility applies equally to ordinary cancelation.

Semantics

• cancelation and checkpointCancelation express different intents. cancelation means "stop
  and abandon this work"; checkpointCancelation means "stop, but preserve progress so the task
  can continue later." The scheduler chooses which to invoke based on why it is interrupting.
• Each openjd_checkpoint_save message supersedes the previous checkpoint as the resume point.
  A scheduler that is capturing the file datasets explicitly could offer a history of
  checkpoints to resume from, not just the latest. Periodic saves are what make crash recovery
  possible — if the worker is lost without warning, the task can still resume from its last
  periodic checkpoint.
• openjd_task_complete works independently of resumability. A chunked step with no
  checkpointCancelation or onCheckpointResume still benefits: when a chunk is interrupted,
  the scheduler reschedules only the tasks not yet reported complete, running the remainder via
  the ordinary onRun action.
• The signalToNotify values name POSIX signals. On Windows, where
  <CancelationMethodNotifyThenTerminate>
  delivers its notification as a console control event rather than a named signal, the RFC
  should define how signalToNotify maps onto that mechanism.

Example: a resumable Fibonacci task

This toy job computes a Fibonacci number slowly, one step every 5 seconds, and checkpoints its (n, v0, v1) state as a compact "n,v0,v1" string — both every few minutes and when it receives SIGUSR1.

specificationVersion: 'jobtemplate-2023-09'
extensions: [CHECKPOINT]  # A new extension, per RFC 0002
name: Fibonacci checkpointing job

parameterDefinitions:
- name: "N"
  description: Which Fibonacci number to calculate (0-based index).
  type: INT
  default: 1000

steps:
- name: ResumableFibonacci
  script:
    actions:
      onRun:
        command: python
        args: ['{{Task.File.Fibonacci}}', '--number', '{{Param.N}}',
               '--checkpoint-interval', '5']
        # An ordinary cancel abandons the work, exactly as today.
        cancelation:
          mode: TERMINATE
        # NEW: how to interrupt this task when the intent is to resume it later.
        # The scheduler sends SIGUSR1; the process saves a checkpoint and exits.
        # If it hasn't exited within notifyPeriodInSeconds, it is terminated.
        checkpointCancelation:
          mode: NOTIFY_THEN_TERMINATE
          notifyPeriodInSeconds: 300
          signalToNotify: "SIGUSR1"
      # NEW: run instead of onRun when the task is resumed from a checkpoint.
      # {{Task.Checkpoint.State}} carries the "state" value from the last
      # openjd_checkpoint_save message emitted before the interruption.
      onCheckpointResume:
        command: python
        args: ['{{Task.File.Fibonacci}}', '--number', '{{Param.N}}',
               '--checkpoint-interval', '5',
               '--resume', '{{Task.Checkpoint.State}}']
        cancelation:
          mode: TERMINATE
        # A resumed task can itself be checkpointed and resumed again.
        checkpointCancelation:
          mode: NOTIFY_THEN_TERMINATE
          notifyPeriodInSeconds: 300
          signalToNotify: "SIGUSR1"
    embeddedFiles:
    - name: Fibonacci
      filename: fibonacci.py
      type: TEXT
      data: |
        import argparse, json, signal, sys, time

        parser = argparse.ArgumentParser(prog="fibonacci.py")
        parser.add_argument("--number", type=int, required=True)
        parser.add_argument("--resume", type=str)
        parser.add_argument("--checkpoint-interval", type=float, default=5)
        args = parser.parse_args()

        if args.resume:
            # Restore state from the checkpoint, and tell the scheduler it was honored.
            n, v0, v1 = (int(value) for value in args.resume.split(","))
            print('openjd_checkpoint_resume_status: {"status": "RESUMED"}')
        else:
            n, v0, v1 = 0, 0, 1

        def save_checkpoint(exit_after=False):
            global last_save
            # NEW: commit a checkpoint by reporting it to the scheduler on stdout.
            print("openjd_checkpoint_save: " + json.dumps({"state": f"{n},{v0},{v1}"}))
            if exit_after:
                sys.exit(1)
            last_save = time.monotonic()

        # checkpointCancelation delivers SIGUSR1: save a checkpoint, then exit.
        signal.signal(signal.SIGUSR1, lambda sig, frame: save_checkpoint(exit_after=True))

        last_save = time.monotonic()
        while n < args.number:
            # Periodic checkpoints also protect against crashes and lost workers.
            if time.monotonic() - last_save > args.checkpoint_interval * 60:
                save_checkpoint()
            print(f"Fibonacci number {n} is {v0}")
            n, v0, v1 = n + 1, v1, v0 + v1
            time.sleep(5)

        print(f"Final answer: Fibonacci number {args.number} is {v0}")

When the scheduler interrupts this task via checkpointCancelation, the log contains, e.g.:

openjd_checkpoint_save: {"state": "3,2,3"}

and when the task is later resumed:

openjd_checkpoint_resume_status: {"status": "RESUMED"}

An application that checkpoints to disk instead — the way RenderMan or Arnold write partial results — would name its checkpoint file in the save message, declaring any additional output files as it produces them. The scheduler takes responsibility for preserving those files and restoring them before onCheckpointResume runs, where {{Task.Checkpoint.File}} resolves to the restored checkpoint file's location:

openjd_checkpoint_files: {"files": ["/var/session-abc123/output/partial.exr"], "mode": "ACCUMULATE"}
openjd_checkpoint_save: {"file": "/var/session-abc123/scratch/checkpoint.rman"}