Skip to content

GPU snapshots (Beta)

View Markdown

Beta: this feature may change before general availability

A GPU class spends most of its cold start in @nodus.enter(): loading weights onto the GPU, compiling kernels, warming caches. With gpu_snapshot=True, the first worker snapshots its process and GPU memory once its enter hooks finish. A later worker of the same class on a machine that holds the snapshot restores it and resumes right after enter, with the model already in GPU memory.

examples/functions/gpu_snapshot/app.py
"""A GPU class whose model loads once, then is snapshotted so later workers restore it instead (Beta).
Run it with `nodus run examples/functions/gpu_snapshot/app.py`. The first worker runs `load` and offers a snapshot
of itself; a later worker on a machine that holds the snapshot resumes right after `load`, with the model already on
the GPU. Every worker answers the same, snapshot or not.
"""
import nodus
app = nodus.App("gpu-snapshot")
image = nodus.Image.debian_slim().pip_install("torch==2.6.0", "transformers==4.53.3")
MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
@app.cls(gpu="A10", image=image, gpu_snapshot=True, scaledown_window="30s", max_cost=1)
class Generator:
@nodus.enter()
def load(self) -> None:
# Runs once per worker that starts cold; a restored worker resumes after it.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
self.tok = AutoTokenizer.from_pretrained(MODEL)
self.model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.float16).to("cuda").eval()
@nodus.method()
def generate(self, prompt: str, tokens: int = 24) -> str:
import torch
ids = self.tok(prompt, return_tensors="pt").to("cuda")
with torch.no_grad():
out = self.model.generate(**ids, max_new_tokens=tokens, do_sample=False)
return self.tok.decode(out[0][ids["input_ids"].shape[1] :], skip_special_tokens=True)
@app.local_entrypoint()
def main() -> None:
print("generated:", Generator().generate.remote("The capital of France is"))
Terminal window
$ nodus run examples/functions/gpu_snapshot/app.py
generated: Paris. It is the largest city in France and ...

gpu_snapshot is accepted only on an @app.cls class that requests a gpu. It sets the Function’s spec.snapshot.gpuMemory; turning it on or off restarts the workers once.

A worker restores a snapshot only when everything that shaped it is the same:

  • the same organization and the same Function;
  • the same code, image digest, Python version, environment and resources;
  • the same secrets delivered to the worker, at the same versions, so a rotated secret starts a fresh snapshot;
  • the same GPU model and driver version;
  • a machine that holds the snapshot: the machine the earlier worker ran on, which a new worker reuses while it is kept warm, or, with restores on other machines enabled for your organization, any of your organization’s machines with the same GPU, driver and machine software.

Anything else, and any failure on the way, starts the worker normally. If a restore fails after it began, that worker is replaced by one that starts normally, and the snapshot is not tried again on that machine.

A restore reads the whole snapshot back, so the machine gives it time in proportion to the snapshot’s size, but never longer than the cold start it replaces. A restore that runs out of that time fails, that worker is replaced by one that starts normally, and the snapshot is not tried again on that machine. A snapshot that could not be read back within 10 minutes is not kept (too_large) and the worker serves without one; on most machines every snapshot fits.

With it enabled, a worker’s snapshot is also stored, encrypted, in Nodus storage once it is taken, and a later worker of the same class that lands on another of your organization’s machines downloads and restores it instead of running its enter hooks. Nodus downloads a snapshot only when the download and the restore are expected to take well under the cold start it replaces, so it pays off for classes whose enter hooks spend their time compiling kernels and capturing CUDA graphs, such as vLLM engines, and not for classes that mostly read weights from disk, which a download cannot beat.

  • The same machine software. The other machine must have the same GPU model and driver version, and the same kernel, CPU features and snapshot tools. Machines of one offering usually match; a worker on a machine that does not starts normally.
  • What it costs. A stored snapshot counts as your organization’s storage and is billed as storage of the Function, at the storage rate. Nothing is stored when it would take your organization’s storage (stored snapshots, checkpoints and Volumes together) past its storage quota, or while your account is restricted or unpaid; nothing is stored either for a snapshot that could never restore faster than the class starts cold. Snapshots on the machine keep working.
  • When it is deleted. A stored snapshot is deleted after 7 days without a restore, when a newer version of the class replaces it, and when the Function is deleted. When a secret it was taken with is rotated, its key is destroyed at once, so it can no longer be read, and its data is deleted within about 2 hours; a deleted secret’s snapshots go the same way within about 2 hours.
  • How many are kept. A Function keeps at most 16 stored snapshots and 256 GiB of them; past either, the ones used least recently (stored or restored) are deleted first.
  • What to expect. On two A10 machines, a vLLM engine with its weights and KV cache kept out of the snapshot made a 5.39 GB snapshot. A worker on the other machine downloaded and restored it, onto another GPU, and served 16.3 to 17.0 s after it started, against 91.5 to 95.5 s for a cold start, with the same output. The snapshot was stored in a MinIO test store, not in Nodus storage, so downloads from Nodus storage can take a different time.

A restored worker skips everything its enter hooks did, so you can keep fewer workers warm: lower min_workers and scaledown_window and let a worker that comes back on a warm machine restore instead of loading.

A restore reads the snapshot back from disk, and the snapshot holds the GPU memory the worker filled, so its time grows with the model unless you keep the model out of it. On an A10, for a class whose enter hook loads a model from its image onto the GPU:

Model Snapshot Cold start, enter hook done Restored
0.5B parameters 2.7 GB 5.7 to 5.9 s 3.1 s
3B parameters 8.2 GB 8.0 to 8.5 s 7.1 s
8B parameters 18.2 GB 11.4 to 11.6 s 14.2 to 14.8 s

With the weights in the snapshot, loading weights that sit on a local disk is about as fast as reading a snapshot of them, so such snapshots pay off when your enter hooks do more than that: download the weights, compile kernels, capture CUDA graphs or warm caches. Each restored worker answered greedy decoding exactly as the cold one did.

A large model’s GPU memory makes the snapshot large, and copying it out and back can take longer than loading the model again. Call nodus.exclude_from_gpu_snapshot at the end of an enter hook to keep a model’s weights out of the snapshot:

@app.cls(gpu="A10", image=image, gpu_snapshot=True)
class Generator:
@nodus.enter()
def load(self) -> None:
self.model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.float16).to("cuda").eval()
nodus.exclude_from_gpu_snapshot(self.model)

It takes torch.nn.Module objects, whose parameters and buffers it keeps out, and tensors; only tensors on the GPU count. When the snapshot is taken, the worker writes their memory once to a file on the machine and frees it from the GPU, so the snapshot holds the rest of the worker. The worker that was snapshotted, and every worker restored from it, reads the weights back into the same tensors before it takes a call. Without the call nothing changes, and outside a class with gpu_snapshot=True it does nothing.

  • Eager code only. The weights come back in the same tensors with the same shapes, but at new GPU addresses. Code that kept a raw GPU address breaks: CUDA graphs, torch.compile(mode="reduce-overhead"), and engines such as vLLM that capture CUDA graphs. Leave the weights in the snapshot for those; a vLLM engine can keep them out its own way.
  • What it costs. A copy of the weights stays on the machine’s disk, and in its memory cache while there is room, for as long as the snapshot is kept, and is deleted with it. The worker that is snapshotted writes the weights once and reads them back before it serves. When the weights would not fit on the machine’s disk or in its snapshot space, the worker writes nothing and keeps them on the GPU; a snapshot that then holds them is usually too large as well and is declined (no_space), and the worker serves either way. While one worker of a class is being snapshotted on a machine, others that start beside it skip their snapshot (cached).
  • What to expect. The snapshot holds the CUDA context and the rest of the worker’s memory instead of the weights, so a restore reads the weights from that file instead of moving them through the snapshot. Measure your own class before you rely on it: nodus: gpu snapshot excluded … bytes and nodus: gpu snapshot refilled … in the worker’s log show how much it kept out and how long reading it back took.

On an A10, with a Transformers model loaded from the image, keeping the weights out made a restore at 8B parameters about 2.9× faster than with the weights in the snapshot and about 2.3× faster than a cold start, comparing medians (5.1 s against 14.5 s and 11.8 s). These are three runs of 3, separate from the table above on the same machine, so their times differ a little:

Model Cold start Restored, weights in the snapshot Restored, weights kept out
0.5B parameters 5.5 to 7.1 s 3.1 to 3.3 s 2.5 to 2.8 s
8B parameters 11.2 to 12.4 s 14.2 to 16.3 s 5.0 to 6.2 s

A worker whose weights cannot be read back exits instead of serving without them, and Nodus replaces it. When it was restored, its snapshot is not used again on that machine. If the weights cannot be written, the worker keeps them in the snapshot and logs nodus: gpu snapshot keeps the weights.

A vLLM engine reserves most of the GPU for its KV cache, so its snapshot holds the KV cache’s whole reservation as well as the weights. Pass the engine to nodus.exclude_from_gpu_snapshot to keep both out:

@app.cls(gpu="A10", image=image, gpu_snapshot=True)
class Server:
@nodus.enter()
def load(self) -> None:
from vllm import LLM
self.llm = LLM(MODEL, enable_sleep_mode=True, worker_extension_cls="nodus.runtime.VLLMWorkerExtension")
nodus.exclude_from_gpu_snapshot(self.llm)

On an A10, a vLLM 0.30 engine serving Qwen3-0.6B in a class restored in about 6 s with its weights and KV cache kept out, against about 15 s with them in the snapshot and about 90 s for a cold start, with the same output; its snapshot fell from 17.1 GB to 5.4 GB. That was measured with vllm.LLM; LLMEngine and AsyncLLM take the same path and have not been measured yet.

Create the engine with enable_sleep_mode=True and worker_extension_cls="nodus.runtime.VLLMWorkerExtension". If your engine already has a worker extension, make it a subclass of nodus.runtime.VLLMWorkerExtension. It takes a vllm.LLM, its LLMEngine, or an AsyncLLM (AsyncLLMEngine) created in an async enter hook, one per class, and not together with PyTorch models or tensors. An engine that does not meet these requirements raises nodus.errors.Invalid in the enter hook, with the fix.

When the snapshot is taken, the engine writes its weights once to a file on the machine and goes to sleep, which frees its weights and its KV cache from the GPU, so the snapshot holds the rest of the worker. Going to sleep aborts any request the engine is still running, so let your enter hooks finish their own requests before they return. The worker that was snapshotted, and every worker restored from it, wakes the engine and reads the weights back before it takes a call.

  • The same GPU addresses. vLLM’s sleep mode keeps the GPU addresses of the weights and the KV cache reserved while it sleeps and maps memory back at the same addresses when it wakes, so the addresses the engine’s CUDA graphs captured point at the weights again once they are read back. The weights are read back byte for byte.
  • The KV cache comes back empty. Its prefix cache is reset, so the first request with a cached prefix computes it again. Nothing in the KV cache is written to the machine’s disk.
  • A spawned engine process. vLLM runs the model in an engine process of its own. In a class with gpu_snapshot=True, Nodus sets VLLM_WORKER_MULTIPROC_METHOD=spawn unless you set it yourself, because an engine process that vLLM forks keeps GPU device files open that a snapshot cannot hold. If you set it to fork and vLLM keeps it, taking the snapshot fails and the worker serves as it would without one, whether or not you keep the engine out. A spawned engine process receives its arguments by pickling, so pass the engine settings that can be pickled: a function defined inside your class or a lambda in its arguments makes LLM() raise. Any other library that forks a process after the GPU is in use, such as a PyTorch DataLoader with workers, makes the snapshot fail the same way.
  • What it costs. As for a model kept out: a copy of the weights stays on the machine with the snapshot, the snapshotted worker writes it once, and every worker reads it back. If the machine cannot keep the file, the engine stays awake with everything in the snapshot, and the worker serves either way.
  • What to expect. nodus: gpu snapshot excluded … bytes of vLLM weights and nodus: gpu snapshot refilled … in the worker’s log show how much it kept out and how long waking took. Measure your own class before you rely on it.

A worker whose engine cannot read its weights back exits instead of serving without them, and Nodus replaces it, as for a model kept out. If the engine fails to go to sleep, the worker reads back whatever it freed and serves without a snapshot (nodus: gpu snapshot skipped).

A snapshot is the worker’s process memory and GPU memory after its enter hooks, and the files those hooks wrote to the worker’s filesystem; with exclude_from_gpu_snapshot, the weights it kept out are a file beside it. It can hold anything your hooks read, including secret values, so it is readable only by Nodus, on your organization’s machines, and is deleted on a machine when the Function changes, after an hour unused and when the machine is released. With restores on other machines enabled, it is also stored, encrypted with a key of its own, in Nodus storage; it is restored only onto your organization’s own machines, never downloadable, checked against what was stored before it is used, and its key is destroyed when a secret it holds is rotated.

Keep in mind when you write the enter hooks:

  • Files. A restored worker sees the files your enter hooks wrote outside Volumes, such as compiled kernel caches, as they were when the snapshot was taken. Up to 8 GiB of them is kept; a worker that wrote more is not snapshotted (no_space), so keep downloaded weights in the image or a Volume. Nodus extracts your code again in a restored worker.
  • Connections. A worker that still holds an open network connection when its enter hooks finish is not snapshotted and keeps serving. Close connections after loading, or open them in a method. Listening sockets are fine.
  • Memory. While the snapshot is taken, the worker’s GPU memory is also held in its process memory. A worker whose memory floor cannot hold both is not snapshotted.
  • Randomness. Every restored worker starts from the same memory. Reseed your own random generators in a method if calls must differ between workers; Nodus reseeds its own. A vLLM engine’s sampling state is in its engine process, where your code does not reach, so pass a seed in each request’s sampling parameters if sampled output must differ between restored workers.
  • Engines. PyTorch and Transformers models loaded in your class restore. vllm serve on one GPU, in its own container, restores with the kernels it compiled and its CUDA graphs, since its engine process and its semaphores are part of the snapshot. A vLLM engine created in your class gets a spawned engine process, as vllm serve does, restores too, and can keep its weights and KV cache out. vLLM spread over several GPUs has not been tested with snapshots yet.

Each worker’s log says what happened at its snapshot point: nodus: gpu snapshot taken, restored, or declined with a reason such as not_enabled, cached, connected (the worker kept a connection to Nodus open), memory_limit, memory_unreadable (the machine could not read the memory in use, so the next worker offers again), no_space or too_large (the snapshot could not be read back within 10 minutes). A worker that could have restored another machine’s snapshot and started normally did so for a reason such as too_slow (the download would not have beaten the cold start, or did not finish in time), corrupt (what was downloaded did not match what was stored, and the snapshot is not used again), download_failed, no_space, too_large or incompatible; a snapshot not stored for your storage quota or account is quota.