documentation
sign in→

Adapting Your Code

Most programs need two small changes to run well on Squadron: save your work periodically, and pick up where you left off on restart.

Jobs run on idle machines. If someone sits down at their workstation while your job is running, the job will go back into the queue and be reassigned elsewhere. That next machine will mount whatever the previous run wrote to /artifacts, so your code can read those files and resume, but only if you wrote something to resume from.

If your program already writes output files and can restart from partial results, you may not need to change anything.

1. Save results to /artifacts

Write checkpoints to /artifacts at regular intervals. Anything you put there will be uploaded automatically while the job runs.

# Save a checkpoint every N steps
if step % SAVE_EVERY == 0:
    checkpoint = {
        "step": step,
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
    }
    torch.save(checkpoint, f"/artifacts/checkpoint-{step}.pt")

Save whatever your workload needs to pick up again. For a training script that's usually model weights and optimizer state; for a data pipeline it might be nothing more than the index of the last record processed.

2. Resume from /artifacts on startup

When your program starts, check whether a previous run left anything in /artifacts. If it did, load it and continue from there instead of starting from scratch.

from pathlib import Path

checkpoints = sorted(Path("/artifacts").glob("checkpoint-*.pt"))
if checkpoints:
    state = torch.load(checkpoints[-1])
    model.load_state_dict(state["model"])
    optimizer.load_state_dict(state["optimizer"])
    start_step = state["step"]
    print(f"Resuming from step {start_step}")
else:
    start_step = 0

Detecting the container environment

The environment variable INSIDE_SQUADRON_CONTAINER is set to 1 inside every job container. You can use it to switch behavior depending on whether your code is running locally or on a remote machine.

import os

if os.environ.get("INSIDE_SQUADRON_CONTAINER"):
    save_dir = "/artifacts"
else:
    save_dir = "./outputs"

Next steps