Preemption & checkpoints

@bladekit.on_preempt

Registers a single callback to run when the container is being preempted. Usable bare — no parentheses — as @bladekit.on_preempt on a free function, or on an @app.cls method (self is supplied, so the callback can read model/optimizer state off the live instance). Re-exported as @app.on_preempt.

@bladekit.on_preempt          # free function
def checkpoint_fn(): ...

class Trainer:
    @app.on_preempt        # or @bladekit.on_preempt — same hook, method form
    def checkpoint_fn(self): ...

fn Callable[..., Any]

Zero-arg callback (plus self for an @app.cls method). Marked with the __fabric_on_preempt__ attribute so the runner can discover it.

save_checkpoint / load_checkpoint

The idempotent-resume primitives: serialize state to a file on a mounted bladekit.Volume inside the on_preempt hook, and read it back at the start of the function or @bladekit.enter() hook on the next attempt.

def bladekit.save_checkpoint(obj: Any, path: str) -> None
def bladekit.load_checkpoint(path: str, default: Optional[Any] = None) -> Any

obj Any

Object to serialize (via cloudpickle) — passed to save_checkpoint.

path str

File path the checkpoint is written to / read from. Should point at a file inside a volumes= mount — the Volume, not this helper, owns durability.

default Optional[Any]Default: None

Returned by load_checkpoint when path doesn't exist yet (no checkpoint written so far — the first attempt).

gpu_class, best_effort, max_interruptions

Three keyword arguments on @app.function(...) and @app.cls(...) control how a function's containers behave under GPU preemption. A .map()/.starmap() call on that function automatically carries the same settings — they aren't repeated as map() keyword arguments.

@app.function(
    gpu="RTX-A4500",
    gpu_class="spot",       # "spot" (default) | "reserved"
    best_effort=False,
    max_interruptions=None,
)
def train(...): ...

gpu_class strDefault: "spot"

Capacity class: spot = interruptible, cheap (default); reserved = non-interruptible, no preemption. Validated client-side at decoration time; any other value raises InvalidArgumentError.

best_effort boolDefault: False

Opts the run into best-effort scheduling. Sent to the control plane only when True — a default False keeps the request body unchanged.

max_interruptions Optional[int]Default: None

The preemption budget for the run — caps how many times the control plane may preempt this function's containers. Sent only when set; None leaves the platform default in effect.

PreemptedError

Signals that a run or call was preempted by a higher-priority job. PreemptedError is a RunError subclass — an infrastructure-level failure, never confused with an exception raised by user code.

class bladekit.PreemptedError(RunError):
    code = "PREEMPTED"

code strDefault: "PREEMPTED"

Stable machine code, inherited from RunError.

run_id Optional[str]

The preempted run's id, when available (inherited from RunError.__init__).

details dictDefault: {}

Structured error payload (inherited from RunError.__init__).

Checkpoint/resume pattern

Putting the pieces together: interruptible capacity, a mounted Volume for durable state, @bladekit.enter() to resume, and @bladekit.on_preempt to checkpoint inside the grace window.

import bladekit

app = bladekit.app("train-resnet")
ckpt_vol = bladekit.Volume.from_name("resnet-ckpts", create_if_missing=True)


@app.cls(
    gpu="RTX-A4500",
    volumes={"/ckpt": ckpt_vol},
    gpu_class="spot",
    max_interruptions=5,
)
class Trainer:
    @bladekit.enter()
    def load(self):
        self.model, self.step = bladekit.load_checkpoint(
            "/ckpt/state.pkl", default=(build_model(), 0)
        )

    @bladekit.on_preempt
    def checkpoint(self):
        bladekit.save_checkpoint((self.model, self.step), "/ckpt/state.pkl")
        ckpt_vol.commit()

    @app.method()
    def train_step(self, batch):
        loss = self.model.step(batch)
        self.step += 1
        return loss

Related reference