Skip to content

Commit 5b49591

Browse files
committed
snapshot-agent tests: exercise the Helm chart and Makefile install paths
Both integration phases now run against a real install path instead of a test-only build: - standalone: run.sh builds `make standalone` in the test-runner pod and TestStandalone runs the agent from those artifacts (bin/snapshot-agent + bin/cuda-checkpoint) on a plain Debian base with the driver mounts — the standalone user's path. The test-only agent image and Dockerfile.snapshot-agent-runtime are gone; cloudbuild-image.yaml now defaults to the official Dockerfile. - k8s: run.sh installs the official snapshot-agent chart (pinned to TEST_NODE, image and agent port overridden) and TestK8s runs its unchanged scenarios against the chart's DaemonSet, so every backend is exercised through the chart exactly once. run.sh is the single entrypoint (--phase both = standalone + k8s, the default). --build has Cloud Build produce the agent image from the working directory, tagged with the commit so reruns don't fight the node's image cache. Chart fixtures are installed before go test starts and fail fast with a chart-attributed error if the pod never becomes Ready. TEST_NODE is required for the k8s phase so the suite stays off other workloads on shared clusters. Signed-off-by: Aishu Kamal <aishuk@google.com>
1 parent 6de6796 commit 5b49591

8 files changed

Lines changed: 376 additions & 117 deletions

File tree

Dockerfile.snapshot-agent-runtime

Lines changed: 0 additions & 20 deletions
This file was deleted.

cloudbuild-image.yaml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,15 @@ steps:
33
args:
44
- 'build'
55
- '-f'
6-
- 'Dockerfile.snapshot-agent-runtime'
6+
- '${_DOCKERFILE}'
77
- '-t'
88
- '${_IMAGE}'
99
- '.'
1010
images:
1111
- '${_IMAGE}'
1212
substitutions:
1313
_IMAGE: 'gcr.io/${PROJECT_ID}/snapshot-agent:dev'
14+
_DOCKERFILE: 'docker/snapshot-agent/Dockerfile'
1415
timeout: '600s'
1516
options:
1617
machineType: 'E2_HIGHCPU_8'

tests/integration/snapshot-agent/README.md

Lines changed: 65 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -1,75 +1,99 @@
11
# Snapshot Agent Integration Tests
22

3-
End-to-end tests for snapshot-agent backends on real GPU hardware, in both standalone and K8s deployment modes.
3+
End-to-end tests for snapshot-agent backends on real GPU hardware. Every
4+
phase exercises a real install path — the way a user actually deploys the
5+
component, not a test-only build:
46

5-
The test suite is written in Go and runs inside the cluster: `run.sh` deploys a test-runner pod, copies the repo source into it, and executes `go test` there. The Go harness deploys the snapshot-agent and inference engine pods itself — one engine at a time, so a single free GPU is enough.
7+
| Phase | Test | Fixture |
8+
|---|---|---|
9+
| `standalone` | `TestStandalone` | `make standalone` artifacts (`bin/snapshot-agent` + `bin/cuda-checkpoint`), built in-cluster and run from a plain Debian base image |
10+
| `k8s` | `TestK8s` | **official snapshot-agent Helm chart** (DaemonSet pinned to `TEST_NODE`) |
11+
12+
The k8s phase runs against the chart the way a user deploys it, so every
13+
backend (CUDA checkpoint via watcher-discovered PIDs, vLLM sleep/wake, SGLang
14+
release/resume) is exercised through the official chart exactly once.
15+
`run.sh` installs the chart fixture before starting `go test` and fails fast
16+
with a chart-attributed error if the chart pod does not become Ready — a
17+
broken chart never shows up as a test failure.
18+
19+
The test suite is written in Go and runs inside the cluster: `run.sh` deploys a test-runner pod, copies the repo source into it, and executes `go test` there. The Go harness deploys the agent (standalone phase) and inference engine pods itself — one engine at a time, so a single free GPU is enough.
620

721
All snapshot/restore calls go through the **Python client** (`timeslice.snapshot_agent`, invoked via `agentctl.py`), so the entire client layer is covered.
822

9-
- `run.sh` — launcher (deploy runner, copy source, install the Python client, `go test`, cleanup)
23+
- `run.sh` — launcher (build image, install chart fixture, deploy runner, copy source, build `make standalone`, install the Python client, `go test`, cleanup)
1024
- `runner.yaml` — test-runner pod + RBAC
1125
- `harness.go` / `engines.go` — harness: pod lifecycle, exec/HTTP helpers, pod specs
1226
- `agentctl.py` — thin CLI over the Python client; builds `BackendConfig` protos in Python from primitive flags
1327
- `standalone_test.go` / `k8s_test.go` — the test cases
1428

15-
## Adding a test
29+
## Testing your changes
1630

17-
Add a `t.Run(...)` inside the engine group that provides the pods it needs, using the harness helpers:
31+
Everything runs from your working directory — uncommitted changes included —
32+
so no commit or merge is needed at any layer:
1833

19-
```go
20-
h.WithEngine(t, VLLM, func(t *testing.T, e *Engine) {
21-
t.Run("MyNewTest", func(t *testing.T) {
22-
before := h.Inference(t, e) // deterministic completion
23-
h.SnapshotOK(t, "my-job", vllmSleepConfig(e.Endpoint(), 1))
24-
vram := h.VRAMMiB(t, e) // GPU memory in use
25-
h.RestoreOK(t, "my-job", vllmWakeConfig(e.Endpoint()))
26-
RequireFreedAndCorrect(t, vram, before, h.Inference(t, e))
27-
})
28-
})
34+
```bash
35+
TEST_NODE=<gpu-node> ./tests/integration/snapshot-agent/run.sh \
36+
--build --project <gcp-project>
2937
```
3038

31-
A new engine is an `EngineSpec` in `engines.go`.
39+
`--build` has Cloud Build produce the agent image from the working directory
40+
(tagged `integ-<commit>` so repeated runs don't collide with node image
41+
caches) and runs both phases against it. `run.sh` then copies the local
42+
workspace into the cluster, so the standalone phase's `make standalone`
43+
build, the Helm chart (installed from local `deploy/`), the Python client,
44+
and the test code all come from the workspace too.
3245

33-
## Prerequisites
34-
35-
- A GKE cluster with at least 1 free GPU
36-
- `gcloud` and `kubectl` on the machine running the tests
37-
(Go and everything else run inside the cluster)
38-
- A snapshot-agent image. Build one from the repo root with:
46+
Alternative (pre-built image — any registry the cluster can pull from):
3947

4048
```bash
4149
gcloud builds submit --config=cloudbuild-image.yaml \
42-
--substitutions=_IMAGE=gcr.io/<your-project>/snapshot-agent:dev .
43-
```
50+
--substitutions=_IMAGE=gcr.io/<project>/snapshot-agent:dev .
4451

45-
This builds from your working directory, so local modifications are included — no commit needed. Requires the Cloud Build API (`gcloud services enable cloudbuild.googleapis.com`) and permission to push to the project's registry; GKE nodes in the same project can pull from `gcr.io/<project>` by default.
52+
TEST_NODE=<gpu-node> ./tests/integration/snapshot-agent/run.sh \
53+
--agent-image gcr.io/<project>/snapshot-agent:dev
54+
```
4655

47-
## Running
56+
## Prerequisites
4857

49-
```bash
50-
./tests/integration/snapshot-agent/run.sh \
51-
--image gcr.io/<your-project>/snapshot-agent:dev \
52-
--project <your-project> \
53-
--cluster <your-cluster> \
54-
--zone <your-zone>
55-
```
58+
- A GKE cluster with at least 1 free GPU
59+
- `gcloud` and `kubectl` on the machine running the tests
60+
(Go and everything else run inside the cluster)
61+
- For the k8s phase: a snapshot-agent image built from the official
62+
Dockerfile (`--build` does this for you; `cloudbuild-image.yaml` defaults
63+
to `docker/snapshot-agent/Dockerfile`). The standalone phase needs no
64+
image — it builds the agent from source in the test runner.
65+
- Cloud Build API enabled (`gcloud services enable cloudbuild.googleapis.com`)
66+
and permission to push to the project's registry; GKE nodes in the same
67+
project can pull from `gcr.io/<project>` by default.
5668

5769
## Options
5870

5971
```text
60-
--image IMAGE Snapshot-agent container image (required)
61-
--project PROJECT GCP project (runs gcloud get-credentials)
72+
--agent-image IMAGE Snapshot-agent image the k8s phase installs via the
73+
official chart (required for k8s/both unless --build)
74+
--build Build the agent image from the working directory via
75+
Cloud Build (requires --project); explicit
76+
--agent-image overrides
77+
--project PROJECT GCP project (image pushes with --build; also used by
78+
gcloud get-credentials with --cluster)
6279
--cluster CLUSTER GKE cluster name
6380
--zone ZONE GKE cluster zone
6481
--model MODEL Model to load (default: Qwen/Qwen2.5-0.5B)
65-
--phase PHASE "standalone", "k8s", or "both" (default: both)
66-
--skip-cleanup Leave the test-runner pod running for debugging
82+
--phase PHASE "standalone", "k8s", or "both" (default)
83+
--skip-cleanup Leave the test-runner pod and chart fixture running
84+
for debugging
6785
```
6886

69-
Environment: `TEST_NODE=<node-name>` pins the suite to a specific node instead
70-
of the default pick (first node with a free GPU by requests). Use it when the
71-
cluster runs workloads that occupy GPUs without requesting them (time-slicing
72-
experiments), which the default pick cannot see.
87+
Environment:
88+
89+
- `TEST_NODE=<node-name>` — required for the k8s phase (the chart is pinned
90+
to this node); for the standalone phase it pins the suite instead of the
91+
default pick (first node with a free GPU by requests). Use the pin when
92+
the cluster runs workloads that occupy GPUs without requesting them
93+
(time-slicing experiments), which the default pick cannot see.
94+
- `CHART_AGENT_PORT=<port>` — port for the chart-deployed agent (default
95+
9001), so the suite can coexist with an unrelated agent on the default
96+
port (the chart runs on hostNetwork).
7397

7498
## Exit code
7599

tests/integration/snapshot-agent/engines.go

Lines changed: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,13 @@ func enginePod(h *Harness, name, image string, command, args []string, env []cor
9393
}
9494
}
9595

96-
func agentPod(image, node, mode string) *corev1.Pod {
96+
// agentPod runs the standalone agent from the `make standalone` artifacts
97+
// (bin/snapshot-agent, bin/cuda-checkpoint): a plain Debian base image plus
98+
// the binaries the harness copies in — the install path a standalone user
99+
// takes, rather than any container image of the agent. The container waits
100+
// for the harness to copy the artifacts (see installAgentBinaries) and then
101+
// execs the agent.
102+
func agentPod(node string) *corev1.Pod {
97103
privileged := true
98104
hostPathDir := corev1.HostPathDirectory
99105
return &corev1.Pod{
@@ -111,21 +117,14 @@ func agentPod(image, node, mode string) *corev1.Pod {
111117
NodeName: node,
112118
HostPID: true,
113119
Tolerations: gpuTolerations(),
114-
InitContainers: []corev1.Container{{
115-
Name: "install-cuda-checkpoint",
116-
Image: "alpine:latest",
117-
// Pinned to an immutable commit rather than the mutable main ref.
118-
Command: []string{"sh", "-c", "apk add --no-cache wget && wget -qO /opt/bin/cuda-checkpoint https://raw.githubusercontent.com/NVIDIA/cuda-checkpoint/00d5cce84c628088d6caa203fc4af40c1538b6f7/bin/x86_64_Linux/cuda-checkpoint && chmod +x /opt/bin/cuda-checkpoint"},
119-
VolumeMounts: []corev1.VolumeMount{
120-
{Name: "bin-dir", MountPath: "/opt/bin"},
121-
},
122-
}},
123120
Containers: []corev1.Container{{
124121
Name: "snapshot-agent",
125-
Image: image,
126-
Args: []string{"--port=9001", "--deployment-mode=" + mode},
122+
Image: "debian:bookworm-slim",
123+
Command: []string{"sh", "-c",
124+
`until [ -f /opt/rlts/.ready ]; do sleep 1; done; exec /opt/rlts/bin/snapshot-agent --port=9001 --deployment-mode=standalone`},
127125
Env: []corev1.EnvVar{
128-
{Name: "PATH", Value: "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/bin"},
126+
// /opt/rlts/bin provides the make-built cuda-checkpoint.
127+
{Name: "PATH", Value: "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rlts/bin"},
129128
{Name: "NODE_NAME", ValueFrom: &corev1.EnvVarSource{
130129
FieldRef: &corev1.ObjectFieldSelector{FieldPath: "spec.nodeName"},
131130
}},
@@ -136,12 +135,12 @@ func agentPod(image, node, mode string) *corev1.Pod {
136135
Ports: []corev1.ContainerPort{{ContainerPort: agentPort}},
137136
SecurityContext: &corev1.SecurityContext{Privileged: &privileged},
138137
VolumeMounts: []corev1.VolumeMount{
139-
{Name: "bin-dir", MountPath: "/opt/bin"},
138+
{Name: "rlts", MountPath: "/opt/rlts"},
140139
{Name: "nvidia-driver", MountPath: "/usr/local/nvidia", ReadOnly: true},
141140
},
142141
}},
143142
Volumes: []corev1.Volume{
144-
{Name: "bin-dir", VolumeSource: corev1.VolumeSource{
143+
{Name: "rlts", VolumeSource: corev1.VolumeSource{
145144
EmptyDir: &corev1.EmptyDirVolumeSource{},
146145
}},
147146
{Name: "nvidia-driver", VolumeSource: corev1.VolumeSource{

0 commit comments

Comments
 (0)