snapshot-agent: add workload channel integration tests (#105)

Design doc: https://docs.google.com/document/d/1QPi1NzlCq2_OXmT2-_bFdGF4eSE__KzkqjHjQKGGEzQ/

## Summary

End-to-end integration tests for the app-channel backend, covering the full production chain: caller → Python client → agent → workload channel → client-library adapter → vLLM Python API.

- **Channel workload pod**: vLLM embedded through its Python API (no HTTP server), registered with the agent via `register_workload()`. The client library under test and the workload script are mounted from a ConfigMap built by the harness, so pods run the exact code from the branch. Deterministic generations are driven through a file protocol; the readiness probe fires once the model is loaded and the workload registered.
- **Tests** (both suites address the workload by job ID alone — no endpoints):
  - `VLLMChannelSleepWake` (standalone + k8s): generate → snapshot over the channel → VRAM freed → restore → identical generation.
  - `VLLMChannelCompound` (standalone): channel suspend, then CUDA checkpoint of the suspended process, restored in reverse order, generation intact.
- `agentctl.py` grows `--backend channel`; harness gains ConfigMap/trigger helpers; runner RBAC adds `configmaps`.

Two workload-pod environment notes encoded in the pod spec: the workload script must not run from the ConfigMap mount (the client package's `types.py` would shadow the stdlib module), and `protobuf` must be upgraded past the image's pin to match the client's generated code.

## Testing

Full suite on H100 (image `snapshot-agent-app-aware:chan-b` = this branch + #88): **12/12 pass** — the 9 existing backend tests plus the 3 channel tests.

```
PASS TestK8s/CUDAWatcherDiscoveredPIDs, VLLMSleepWake, SGLangReleaseResume, VLLMChannelSleepWake
PASS TestStandalone/CUDACheckpointRestore, VLLMSleepWake, VLLMCompound, VLLMSuspendDiscard,
     SGLangReleaseResume, SGLangCompound, VLLMChannelSleepWake, VLLMChannelCompound
```

Signed-off-by: Aishu Kamal <aishuk@google.com>

Author: aishukamal

tests/integration/harness/harness.go

@@ -277,3 +277,13 @@ func (c *Cluster) PodVRAMMiB(t *testing.T, pod, container string, timeout time.D
 	}
 	return mib
 }
+
+// DeleteConfigMap deletes a ConfigMap by name in the cluster's namespace.
+// NotFound errors are ignored (idempotent cleanup).
+func (c *Cluster) DeleteConfigMap(name string) error {
+	err := c.Client.CoreV1().ConfigMaps(c.Namespace).Delete(context.Background(), name, metav1.DeleteOptions{})
+	if apierrors.IsNotFound(err) {
+		return nil
+	}
+	return err
+}

tests/integration/runner.yaml

@@ -68,6 +68,9 @@ rules:
   - apiGroups: [""]
     resources: ["pods/exec"]
     verbs: ["create"]
+  - apiGroups: [""]
+    resources: ["configmaps"]
+    verbs: ["get", "create", "delete"]
   - apiGroups: [""]
     resources: ["pods/log"]
     verbs: ["get"]

tests/integration/snapshot-agent/agentctl.py

@@ -7,10 +7,11 @@
 constructed here, in Python, the same way a real workload would build them.
 
 Usage:
-  agentctl.py --agent HOST:PORT snapshot|restore --job-id ID --backend cuda|app
+  agentctl.py --agent HOST:PORT snapshot|restore --job-id ID
+              --backend cuda|app|channel
               [--pids 1,2,3]                       (cuda)
               [--app vllm|sglang] [--endpoints URL,URL]
-              [--mode offload|discard] [--tags a,b] (app)
+              [--mode offload|discard] [--tags a,b] (app, channel)
 
 Exits 0 when the operation completes, 1 otherwise.
 """
@@ -51,6 +52,12 @@ def build_config(args: argparse.Namespace) -> snapshot_agent_pb2.BackendConfig:
             app_endpoint.tags.extend(args.tags.split(","))
         return snapshot_agent_pb2.BackendConfig(app_endpoint=app_endpoint)
 
+    if args.backend == "channel":
+        app_channel = snapshot_agent_pb2.AppChannelConfig(mode=MODES[args.mode])
+        if args.tags:
+            app_channel.tags.extend(args.tags.split(","))
+        return snapshot_agent_pb2.BackendConfig(app_channel=app_channel)
+
     raise ValueError(f"unknown backend {args.backend!r}")
 
 
@@ -60,7 +67,7 @@ def main() -> int:
     parser.add_argument("--agent", required=True, help="agent endpoint HOST:PORT")
     parser.add_argument("--job-id", required=True)
     parser.add_argument("--group", default="test")
-    parser.add_argument("--backend", required=True, choices=["cuda", "app"])
+    parser.add_argument("--backend", required=True, choices=["cuda", "app", "channel"])
     parser.add_argument("--pids", default="", help="comma-separated PIDs (cuda)")
     parser.add_argument("--app", default="", choices=["", "vllm", "sglang"], help="application (app backend)")
     parser.add_argument("--endpoints", default="", help="comma-separated application URLs")

tests/integration/snapshot-agent/channel_workload.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+"""In-pod test workload for the app-channel backend.
+
+Embeds vLLM through its Python API (no HTTP server) and registers with the
+node's snapshot-agent over the workload channel — the integration path for
+Python-API workloads such as RL samplers. The harness drives deterministic
+generations through a file protocol:
+
+  touch STATE/trigger  ->  completion text written to STATE/result
+
+STATE/ready is created once the model is loaded and the workload is
+registered; the pod's readiness probe watches it.
+"""
+
+import os
+import time
+from pathlib import Path
+
+from timeslice.snapshot_agent import register_workload
+from vllm import LLM, SamplingParams
+
+STATE = Path("/workload-state")
+PROMPT = "The capital of France is"
+
+
+def main() -> None:
+    STATE.mkdir(exist_ok=True)
+    llm = LLM(
+        model=os.environ["MODEL"],
+        enable_sleep_mode=True,
+        gpu_memory_utilization=0.5,
+    )
+    params = SamplingParams(temperature=0, max_tokens=15)
+
+    job_id = os.environ["TIME_SLICE_JOB_ID"]
+    handle = register_workload(
+        os.environ["SNAPSHOT_AGENT_ADDR"],
+        job_id=job_id,
+        group=os.environ.get("TIME_SLICE_GROUP", "test"),
+        workload=llm,
+    )
+    print(f"registered workload for job {job_id}", flush=True)
+
+    (STATE / "ready").write_text("ok")
+    trigger = STATE / "trigger"
+    result = STATE / "result"
+    try:
+        while True:
+            if trigger.exists():
+                outputs = llm.generate([PROMPT], params)
+                result.write_text(outputs[0].outputs[0].text)
+                trigger.unlink()
+            time.sleep(0.5)
+    finally:
+        handle.close()
+
+
+if __name__ == "__main__":
+    main()

tests/integration/snapshot-agent/engines.go

@@ -3,6 +3,8 @@
 package integration
 
 import (
+	"fmt"
+
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -161,3 +163,71 @@ func gpuTolerations() []corev1.Toleration {
 		Effect:   corev1.TaintEffectNoSchedule,
 	}}
 }
+
+// channelWorkloadPod runs vLLM through its Python API (no HTTP server) and
+// registers with the agent over the workload channel. The client library and
+// the workload script are mounted from a ConfigMap built by the harness; the
+// readiness probe fires once the model is loaded and the workload registered.
+func channelWorkloadPod(h *Harness, jobID string) *corev1.Pod {
+	labels := map[string]string{
+		"app":        channelPodName,
+		"test-suite": "snapshot-agent-integration",
+	}
+	if h.Mode == "k8s" {
+		labels["timeslice.io/job-id"] = jobID
+		labels["timeslice.io/group"] = "test"
+	}
+	// The script must not run from /opt/src: the client package's types.py
+	// would shadow the stdlib types module via the script-dir sys.path entry.
+	startup := "pip install -q --upgrade grpcio protobuf && " +
+		"mkdir -p /opt/pkg/timeslice/snapshot_agent && " +
+		"cp /opt/src/*.py /opt/pkg/timeslice/snapshot_agent/ && " +
+		": > /opt/pkg/timeslice/__init__.py && " +
+		"cp /opt/src/channel_workload.py /opt/channel_workload.py && " +
+		"PYTHONPATH=/opt/pkg exec python3 /opt/channel_workload.py"
+	return &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      channelPodName,
+			Namespace: namespace,
+			Labels:    labels,
+		},
+		Spec: corev1.PodSpec{
+			ServiceAccountName: "snapshot-agent-test",
+			RestartPolicy:      corev1.RestartPolicyNever,
+			NodeName:           h.Node,
+			Tolerations:        gpuTolerations(),
+			Volumes: []corev1.Volume{{
+				Name: "src",
+				VolumeSource: corev1.VolumeSource{
+					ConfigMap: &corev1.ConfigMapVolumeSource{
+						LocalObjectReference: corev1.LocalObjectReference{Name: channelConfigMapName},
+					},
+				},
+			}},
+			Containers: []corev1.Container{{
+				Name:    channelContainer,
+				Image:   "vllm/vllm-openai:v0.25.1@sha256:e4f88a835143cd22aee2397a26ec6bb80b3a4a6fe0c882bcbc63822904766089",
+				Command: []string{"sh", "-c", startup},
+				Env: []corev1.EnvVar{
+					{Name: "MODEL", Value: h.Model},
+					{Name: "SNAPSHOT_AGENT_ADDR", Value: fmt.Sprintf("%s:%d", h.AgentIP, h.AgentPort)},
+					{Name: "TIME_SLICE_JOB_ID", Value: jobID},
+					{Name: "TIME_SLICE_GROUP", Value: "test"},
+				},
+				VolumeMounts: []corev1.VolumeMount{{Name: "src", MountPath: "/opt/src"}},
+				ReadinessProbe: &corev1.Probe{
+					ProbeHandler: corev1.ProbeHandler{
+						Exec: &corev1.ExecAction{Command: []string{"test", "-f", "/workload-state/ready"}},
+					},
+					InitialDelaySeconds: 20,
+					PeriodSeconds:       5,
+					FailureThreshold:    60,
+				},
+				Resources: corev1.ResourceRequirements{
+					Limits:   corev1.ResourceList{"nvidia.com/gpu": resource.MustParse("1")},
+					Requests: corev1.ResourceList{"nvidia.com/gpu": resource.MustParse("1")},
+				},
+			}},
+		},
+	}
+}

tests/integration/snapshot-agent/harness.go

@@ -29,11 +29,13 @@ import (
 	"net/http"
 	"os"
 	"os/exec"
+	"path/filepath"
 	"strconv"
 	"strings"
 	"testing"
 	"time"
 
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/llm-d-incubation/llm-d-rl-time-slicing/tests/integration/harness"
@@ -53,6 +55,15 @@ const (
 	opTimeout = 120 * time.Second
 	// vramFreedMiB is the threshold below which we consider GPU memory freed.
 	vramFreedMiB = 5000
+
+	// Channel workload pod pieces (see channelWorkloadPod and
+	// WithChannelWorkload).
+	channelPodName       = "channel-workload-test"
+	channelConfigMapName = "channel-workload-src"
+	channelContainer     = "workload"
+	// channelClientSrcDir is the Python client package, relative to this
+	// package (go test's working directory), mounted into the workload pod.
+	channelClientSrcDir = "../../../pkg/client/python/timeslice/snapshot_agent"
 )
 
 // Harness manages the test stack for one deployment mode.
@@ -301,6 +312,16 @@ func appConfig(app, endpoint, mode string) BackendArgs {
 	return args
 }
 
+// channelConfig targets the workload registered on the job's channel.
+// mode may be "" (workload's registered default), "offload", or "discard".
+func channelConfig(mode string) BackendArgs {
+	args := BackendArgs{"--backend", "channel"}
+	if mode != "" {
+		args = append(args, "--mode", mode)
+	}
+	return args
+}
+
 // SnapshotOK snapshots via the Python client and fails the test if the
 // operation does not complete.
 func (h *Harness) SnapshotOK(t *testing.T, jobID string, cfg BackendArgs) {
@@ -381,3 +402,126 @@ func RequireFreedAndCorrect(t *testing.T, vramWhileAsleep int, before, after str
 		t.Errorf("inference changed after restore: before=%q after=%q", before, after)
 	}
 }
+
+
+// --- Channel workload helpers ---
+
+// ChannelWorkload is a running Python-API workload registered with the agent
+// over the workload channel.
+type ChannelWorkload struct {
+	PodName string
+	JobID   string
+	PID     int32 // standalone mode only
+}
+
+// WithChannelWorkload deploys the channel workload pod (vLLM via the Python
+// API, registered through the client library), waits until it is registered,
+// runs fn, and deletes the pod (freeing the GPU).
+func (h *Harness) WithChannelWorkload(t *testing.T, fn func(t *testing.T, w *ChannelWorkload)) {
+	t.Helper()
+	jobID := "chan-standalone"
+	if h.Mode == "k8s" {
+		jobID = "chan-k8s"
+	}
+
+	h.createChannelSourceConfigMap(t)
+	defer func() {
+		if err := h.DeleteConfigMap(channelConfigMapName); err != nil {
+			t.Logf("warning: failed to delete ConfigMap %s: %v", channelConfigMapName, err)
+		}
+	}()
+
+	h.DeletePodAndWait(t, channelPodName)
+	pod := channelWorkloadPod(h, jobID)
+	if _, err := h.Client.CoreV1().Pods(namespace).Create(context.Background(), pod, metav1.CreateOptions{}); err != nil {
+		t.Fatalf("creating channel workload pod: %v", err)
+	}
+	defer h.DeletePodAndWait(t, channelPodName)
+
+	// The readiness probe covers model load and channel registration.
+	h.WaitPodReady(t, channelPodName, podTimeout)
+	t.Logf("channel workload ready, registered as job %s", jobID)
+
+	w := &ChannelWorkload{PodName: channelPodName, JobID: jobID}
+	if h.Mode == "standalone" {
+		w.PID = h.findPID(t, "channel_workload")
+		t.Logf("channel workload PID: %d", w.PID)
+	} else {
+		t.Log("waiting 10s for watcher to register the job...")
+		time.Sleep(10 * time.Second)
+	}
+
+	fn(t, w)
+}
+
+// createChannelSourceConfigMap packages the Python client library and the
+// workload script into a ConfigMap mounted by the workload pod, so the pod
+// runs the exact client code under test.
+func (h *Harness) createChannelSourceConfigMap(t *testing.T) {
+	t.Helper()
+	files := map[string]string{}
+	entries, err := os.ReadDir(channelClientSrcDir)
+	if err != nil {
+		t.Fatalf("reading client source dir: %v", err)
+	}
+	for _, entry := range entries {
+		if !strings.HasSuffix(entry.Name(), ".py") {
+			continue
+		}
+		data, err := os.ReadFile(filepath.Join(channelClientSrcDir, entry.Name()))
+		if err != nil {
+			t.Fatalf("reading %s: %v", entry.Name(), err)
+		}
+		files[entry.Name()] = string(data)
+	}
+	script, err := os.ReadFile("channel_workload.py")
+	if err != nil {
+		t.Fatalf("reading channel_workload.py: %v", err)
+	}
+	files["channel_workload.py"] = string(script)
+
+	if err := h.DeleteConfigMap(channelConfigMapName); err != nil {
+		t.Logf("warning: pre-create ConfigMap cleanup failed: %v", err)
+	}
+	cm := &corev1.ConfigMap{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      channelConfigMapName,
+			Namespace: namespace,
+			Labels:    map[string]string{"test-suite": "snapshot-agent-integration"},
+		},
+		Data: files,
+	}
+	if _, err := h.Client.CoreV1().ConfigMaps(namespace).Create(context.Background(), cm, metav1.CreateOptions{}); err != nil {
+		t.Fatalf("creating ConfigMap %s: %v", channelConfigMapName, err)
+	}
+}
+
+
+// TriggerGenerate asks the workload for a deterministic generation through
+// its file protocol and returns the completion text.
+func (h *Harness) TriggerGenerate(t *testing.T, w *ChannelWorkload) string {
+	t.Helper()
+	_, err := h.ExecPod(w.PodName, channelContainer, opTimeout, "sh", "-c",
+		"rm -f /workload-state/result && touch /workload-state/trigger")
+	if err != nil {
+		t.Fatalf("triggering generation on %s: %v", w.PodName, err)
+	}
+	deadline := time.Now().Add(2 * time.Minute)
+	for time.Now().Before(deadline) {
+		out, err := h.ExecPod(w.PodName, channelContainer, opTimeout, "sh", "-c",
+			"cat /workload-state/result 2>/dev/null")
+		if err == nil && out != "" {
+			return out
+		}
+		time.Sleep(2 * time.Second)
+	}
+	t.Fatalf("timeout waiting for generation from %s", w.PodName)
+	return ""
+}
+
+// WorkloadVRAMMiB returns the GPU memory used (MiB) as seen from the channel
+// workload pod.
+func (h *Harness) WorkloadVRAMMiB(t *testing.T, w *ChannelWorkload) int {
+	t.Helper()
+	return h.PodVRAMMiB(t, w.PodName, channelContainer, opTimeout)
+}

tests/integration/snapshot-agent/k8s_test.go

@@ -57,4 +57,18 @@ func TestK8s(t *testing.T) {
 			RequireFreedAndCorrect(t, vramReleased, before, after)
 		})
 	})
+
+	// The channel workload's pod carries the job label the watcher discovers;
+	// the workload registers over the channel under the same job ID.
+	h.WithChannelWorkload(t, func(t *testing.T, w *ChannelWorkload) {
+		t.Run("VLLMChannelSleepWake", func(t *testing.T) {
+			before := h.TriggerGenerate(t, w)
+			h.SnapshotOK(t, w.JobID, channelConfig(""))
+			vramAsleep := h.WorkloadVRAMMiB(t, w)
+			t.Logf("VRAM after channel snapshot: %d MiB", vramAsleep)
+			h.RestoreOK(t, w.JobID, channelConfig(""))
+			after := h.TriggerGenerate(t, w)
+			RequireFreedAndCorrect(t, vramAsleep, before, after)
+		})
+	})
 }