Skip to content

Commit 3bad43c

Browse files
authored
snapshot-agent: document the app-channel backend in the user guide (#106)
## Summary Documents the app-channel backend in the snapshot-agent user guide: - Backends table gains the `app_channel` row. - New "Application-Aware (app_channel)" section: when to use it (Python-API workloads with no HTTP server), workload-side registration with `register_workload` — recognized engines, the `Snapshottable` object form, and callbacks for workloads without a public C/R API — capabilities and mode resolution, and the caller side (`AppChannelConfig`, empty config = workload's registered default). - Composing-backends section notes both application-aware transports compose with CUDA checkpoint the same way. ## Testing Docs only. The documented flows are exercised by the channel integration tests ([app-channel 4]): 12/12 on H100. Signed-off-by: Aishu Kamal <aishuk@google.com>
1 parent 2c4d4b6 commit 3bad43c

1 file changed

Lines changed: 82 additions & 1 deletion

File tree

guides/snapshot-agent/README.md

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,7 @@ The Snapshot Agent supports multiple backends for different GPU memory managemen
151151
|---------|--------|-------------|------------|-------------|
152152
| CUDA Checkpoint | `cuda` | Process-level CUDA state save/restore via `cuda-checkpoint` | ~100% | ~1-3s |
153153
| Application-Aware | `app_endpoint` | Suspend/resume through the application's own HTTP API (vLLM, SGLang) | ~96% | ~50-100ms |
154+
| Application-Aware | `app_channel` | Suspend/resume pushed over a channel the workload registered (Python-API workloads, no HTTP server) | ~96% | ~50-100ms |
154155

155156
The VRAM Freed and Resume Time figures are illustrative, measured with a small model (Qwen2.5-0.5B) on an H100; actual numbers depend on the model size, hardware, and engine version.
156157

@@ -263,9 +264,89 @@ with SnapshotAgentClient("localhost:9001") as client:
263264
client.restore_and_wait(job_id="my-sglang-job", backend_config=sglang_config)
264265
```
265266

267+
### Application-Aware (app_channel)
268+
269+
For workloads that embed their engine in-process through a Python API — no HTTP
270+
server for the agent to call (e.g. an RL sampler running vLLM via
271+
`AsyncLLMEngine`). The connection is inverted: the workload registers with the
272+
node-local agent once at startup, and the agent pushes suspend/resume commands
273+
over that stream. Callers address the workload by `job_id` alone — no
274+
endpoints, and no knowledge of which application is running.
275+
276+
**Workload side** — register once at startup with `register_workload`:
277+
278+
```python
279+
from timeslice.snapshot_agent import register_workload
280+
281+
engine = AsyncLLMEngine.from_engine_args(args) # enable_sleep_mode=True
282+
283+
handle = register_workload(
284+
"127.0.0.1:9001", # the agent on this node (registration is node-scoped)
285+
job_id="my-sampler", # must match the job_id used in Snapshot/Restore
286+
# requests (in k8s: the timeslice.io/job-id pod label)
287+
group="samplers",
288+
workload=engine, # vLLM engines are recognized by type
289+
)
290+
# ... run; the library services commands in the background ...
291+
handle.close() # on clean shutdown
292+
```
293+
294+
The library owns the stream: a background thread, command dispatch and
295+
acknowledgements, and reconnect with backoff (re-registering after agent
296+
restarts). Recognized engines (vLLM `LLM`/`AsyncLLMEngine`/`AsyncLLM`) need
297+
nothing else.
298+
299+
Workloads with no publicly known C/R API (e.g. hand-rolled FSDP offload) keep
300+
their own mechanics and hand them to the library, either as an object with
301+
`snapshot(mode, tags)`/`restore(tags)` methods:
302+
303+
```python
304+
class TrainerWorkload:
305+
supported_modes = ["offload"] # trainers can't reconstruct dropped state
306+
307+
def snapshot(self, mode, tags):
308+
offload_model_and_optimizer_to_host()
309+
310+
def restore(self, tags):
311+
reload_from_host()
312+
313+
register_workload("127.0.0.1:9001", job_id="my-trainer", group="trainers",
314+
workload=TrainerWorkload())
315+
```
316+
317+
or as plain callbacks:
318+
319+
```python
320+
register_workload("127.0.0.1:9001", job_id="my-trainer", group="trainers",
321+
on_snapshot=lambda mode, tags: trainer.offload(),
322+
on_restore=lambda tags: trainer.reload(),
323+
supported_modes=["offload"])
324+
```
325+
326+
At registration the workload advertises its capabilities (`supported_modes`,
327+
`default_mode`). The agent resolves each request as: explicit request mode →
328+
registered default → `SUSPEND_MODE_OFFLOAD`, and rejects unsupported modes
329+
before any command is sent (e.g. DISCARD against a trainer that only supports
330+
OFFLOAD fails the operation immediately).
331+
332+
**Caller side** — the usual `Snapshot`/`Restore` with an `app_channel` config.
333+
An empty config means "suspend however the workload declared at registration":
334+
335+
```python
336+
channel_config = snapshot.BackendConfig(app_channel=snapshot.AppChannelConfig())
337+
338+
result = client.snapshot_and_wait(job_id="my-sampler", backend_config=channel_config)
339+
340+
result = client.restore_and_wait(job_id="my-sampler", backend_config=channel_config)
341+
```
342+
343+
A request for a job with no registered channel fails fast with
344+
`no workload channel registered for job "..."`. If the workload's suspend
345+
raises, the operation fails with the workload's error text.
346+
266347
### Composing Backends
267348

268-
Application-aware suspend and CUDA checkpoint are separate operations that compose. Suspend first, then checkpoint; restore in reverse order:
349+
Application-aware suspend (either transport) and CUDA checkpoint are separate operations that compose. Suspend first, then checkpoint; restore in reverse order:
269350

270351
```python
271352
app_config = snapshot.BackendConfig(

0 commit comments

Comments
 (0)