diff --git a/agent/cmd/agent/config.go b/agent/cmd/agent/config.go index b8d8550c..2e0dbb84 100644 --- a/agent/cmd/agent/config.go +++ b/agent/cmd/agent/config.go @@ -8,7 +8,9 @@ import ( "errors" "fmt" "os" + "sync/atomic" + "github.com/go-logr/logr" "gopkg.in/yaml.v3" "github.com/ai-dynamo/snapshot/agent/internal/types" @@ -33,6 +35,30 @@ func LoadConfig(path string) (*types.AgentConfig, error) { return cfg, nil } +// NewSkipCompatCheckFn returns a per-restore read of the node-wide switch off +// the mounted ConfigMap, so an admin who flips it does not have to roll the +// DaemonSet to be heard. Kubernetes projects ConfigMap updates into the mount +// on its own; nothing here watches or polls. +// +// A read that fails keeps the last value it did get: a restore is never failed, +// and never quietly checked differently, because a config read went wrong. +func NewSkipCompatCheckFn(path string, initial bool, log logr.Logger) func() bool { + var last atomic.Bool + last.Store(initial) + return func() bool { + cfg, err := LoadConfig(path) + if err != nil { + log.Error(err, "Failed to re-read the restore compatibility switch; keeping the last known value", + "skipCompatCheck", last.Load(), + "path", path, + ) + return last.Load() + } + last.Store(cfg.Restore.SkipCompatCheck) + return cfg.Restore.SkipCompatCheck + } +} + // LoadConfigOrDefault loads configuration from a file, falling back to defaults if the file doesn't exist. func LoadConfigOrDefault(path string) (*types.AgentConfig, error) { cfg, err := LoadConfig(path) diff --git a/agent/cmd/agent/config_test.go b/agent/cmd/agent/config_test.go new file mode 100644 index 00000000..616332a6 --- /dev/null +++ b/agent/cmd/agent/config_test.go @@ -0,0 +1,72 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "os" + "path/filepath" + "testing" + + "github.com/go-logr/logr" +) + +func writeConfig(t *testing.T, path, document string) { + t.Helper() + if err := os.WriteFile(path, []byte(document), 0o600); err != nil { + t.Fatalf("write config: %v", err) + } +} + +// The point of re-reading is that flipping the ConfigMap is enough: the kubelet +// updates the mounted file on its own, and the next restore sees the new value +// without the DaemonSet being rolled. +func TestNewSkipCompatCheckFnFollowsTheFile(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.yaml") + writeConfig(t, path, "restore:\n skipCompatCheck: false\n") + skip := NewSkipCompatCheckFn(path, false, logr.Discard()) + + if skip() { + t.Fatal("switch read true from a file that says false") + } + + writeConfig(t, path, "restore:\n skipCompatCheck: true\n") + if !skip() { + t.Fatal("switch did not follow the file being flipped on") + } + + writeConfig(t, path, "restore:\n skipCompatCheck: false\n") + if skip() { + t.Fatal("switch did not follow the file being flipped back off") + } +} + +// A restore is never failed, and never quietly checked differently, because a +// config read went wrong. +func TestNewSkipCompatCheckFnKeepsTheLastGoodValue(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.yaml") + + t.Run("missing file keeps what the agent started with", func(t *testing.T) { + absent := filepath.Join(dir, "absent.yaml") + if !NewSkipCompatCheckFn(absent, true, logr.Discard())() { + t.Fatal("switch lost the startup value when the file was missing") + } + if NewSkipCompatCheckFn(absent, false, logr.Discard())() { + t.Fatal("switch invented a value when the file was missing") + } + }) + + t.Run("malformed file keeps the last value read", func(t *testing.T) { + writeConfig(t, path, "restore:\n skipCompatCheck: true\n") + skip := NewSkipCompatCheckFn(path, false, logr.Discard()) + if !skip() { + t.Fatal("switch read false from a file that says true") + } + + writeConfig(t, path, "restore:\n\tskipCompatCheck: not-a-bool\n") + if !skip() { + t.Fatal("switch dropped the last known value on an unparseable file") + } + }) +} diff --git a/agent/cmd/agent/main.go b/agent/cmd/agent/main.go index 4a1d281c..83700cd3 100644 --- a/agent/cmd/agent/main.go +++ b/agent/cmd/agent/main.go @@ -38,6 +38,13 @@ func main() { if err := cfg.Validate(); err != nil { fatal(agentLog, err, "Invalid configuration") } + // A host fact that cannot be read is unknown, never fatal: the node keeps + // capturing and restoring, and the checks that need it do not apply. + if kernelVersion, err := snapshotruntime.ReadKernelVersion(snapshotruntime.HostProcPath); err != nil { + agentLog.Error(err, "Failed to read the host kernel version; checkpoints taken here will not record it") + } else { + cfg.HostKernelVersion = kernelVersion + } rt, err := snapshotruntime.New(*runtimeType, *runtimeSocket) if err != nil { @@ -59,7 +66,8 @@ func main() { ) // The node controller handles both restore and capture paths. - nodeController, err := controller.NewNodeController(cfg, rt, rootLog.WithName("controller")) + nodeController, err := controller.NewNodeController(cfg, rt, rootLog.WithName("controller"), + NewSkipCompatCheckFn(ConfigMapPath, cfg.Restore.SkipCompatCheck, agentLog)) if err != nil { fatal(agentLog, err, "Failed to create snapshot node controller") } diff --git a/agent/internal/controller/compat.go b/agent/internal/controller/compat.go new file mode 100644 index 00000000..0573e8d5 --- /dev/null +++ b/agent/internal/controller/compat.go @@ -0,0 +1,143 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package controller + +import ( + "context" + "fmt" + "runtime" + + corev1 "k8s.io/api/core/v1" + + "github.com/ai-dynamo/snapshot/agent/internal/types" + "github.com/ai-dynamo/snapshot/api/compat" + snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" +) + +// refuseRestore records a restore this node will not attempt. It is terminal +// like any other restore failure and reports through the same condition, with +// its own reason so an operator can tell a checkpoint that cannot run here from +// one that tried and broke. +func (w *NodeController) refuseRestore(ctx context.Context, pod *corev1.Pod, incompatible *compat.IncompatibleError) bool { + reason := compat.Reasons(incompatible.Mismatches) + w.logRestoreRefusal(pod, incompatible, reason) + return w.finishRestore( + ctx, + pod, + corev1.ConditionFalse, + restoreIncompatibleReason, + reason, + ) != nil +} + +func (w *NodeController) logRestoreRefusal(pod *corev1.Pod, incompatible *compat.IncompatibleError, reason string) { + w.log.Info("Refusing restore; this node cannot run the checkpoint", + "pod", fmt.Sprintf("%s/%s", pod.Namespace, pod.Name), + "gate", string(incompatible.Gate), + "reason", reason, + ) +} + +// reopenedAfterRefusal reports a pod that the gates turned down and that has +// since asked for them to be skipped. Nothing else reopens a terminal restore, +// which is what makes the skip request an escape hatch and not a retry. +func (w *NodeController) reopenedAfterRefusal(pod *corev1.Pod) bool { + condition := findRestoredCondition(pod) + if condition == nil || condition.Status != corev1.ConditionFalse || condition.Reason != restoreIncompatibleReason { + return false + } + return w.skipCompatCheckRequested(pod) +} + +// podFacts reads what one container of a pod runs as and is allowed. It serves +// both sides of a comparison: what a capture records about the source pod, and +// what a restore target offers. +// +// A container that is not in the pod leaves its facts unknown. +func podFacts(pod *corev1.Pod, containerName string) compat.Facts { + if pod == nil { + return compat.Facts{} + } + + facts := compat.Facts{} + for _, container := range pod.Spec.Containers { + if container.Name != containerName { + continue + } + facts.Image = container.Image + facts.CPULimit = limitString(container.Resources.Limits, corev1.ResourceCPU) + facts.MemoryLimit = limitString(container.Resources.Limits, corev1.ResourceMemory) + } + return facts +} + +// limitString keeps an unset limit unset. A missing quantity formats as "0", +// which would otherwise read as a container limited to nothing. +func limitString(limits corev1.ResourceList, name corev1.ResourceName) string { + quantity, ok := limits[name] + if !ok { + return "" + } + return quantity.String() +} + +// skipCompatCheckRequested reports whether this restore was asked to skip +// the compatibility gates, by the pod that is being restored or by the node +// it landed on. +func (w *NodeController) skipCompatCheckRequested(pod *corev1.Pod) bool { + return w.skipCompatCheckFn() || + snapshotv1alpha1.SkipCompatCheckFromAnnotations(pod.Annotations) +} + +// preflightCompatibility runs the pre-flight compatibility gate for one restore. +// A nil error means the restore may be attempted. +func (w *NodeController) preflightCompatibility( + pod *corev1.Pod, + artifact *restoreArtifact, + mappings []snapshotv1alpha1.RestoreContainerMapping, +) error { + log := w.log.WithValues("pod", fmt.Sprintf("%s/%s", pod.Namespace, pod.Name), "container", artifact.SourceContainerName) + if artifact.SkipCompatCheck { + log.Info("Restore compatibility check skipped by request", "gate", string(compat.GatePreflight)) + return nil + } + + manifest, err := types.ReadManifest(artifact.Path) + if err != nil { + // An unreadable manifest is not an incompatibility. The restore path + // reads it again and reports the real error from there, so refusing here + // would relabel a broken artifact as an incompatible one. + log.V(1).Info("Skipping restore compatibility gate; checkpoint manifest is unreadable", + "artifact_path", artifact.Path, + "error", err.Error(), + ) + return nil + } + + sourceFacts := manifest.CompatFacts() + for _, mapping := range mappings { + mismatches := w.compareFn( + compat.GatePreflight, + sourceFacts, + w.preflightTargetFacts(pod, mapping.Destination), + ) + if len(mismatches) != 0 { + return compat.NewIncompatibleError(compat.GatePreflight, mismatches) + } + } + return nil +} + +// preflightTargetFacts describes what this node and this pod offer a restore, as +// far as it is knowable before the placeholder container exists. It is assembled +// per restore from facts the agent already holds, so the gate costs no syscalls +// and no API reads. +func (w *NodeController) preflightTargetFacts(pod *corev1.Pod, containerName string) compat.Facts { + facts := podFacts(pod, containerName) + // The agent's own architecture, which is the node's: this binary could not + // be running here otherwise. + facts.CPUArch = runtime.GOARCH + facts.KernelVersion = w.config.HostKernelVersion + return facts +} diff --git a/agent/internal/controller/compat_harness_test.go b/agent/internal/controller/compat_harness_test.go new file mode 100644 index 00000000..286064a1 --- /dev/null +++ b/agent/internal/controller/compat_harness_test.go @@ -0,0 +1,228 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package controller + +import ( + "context" + "encoding/json" + "fmt" + "os" + "sync" + "testing" + + "github.com/go-logr/logr" + "github.com/stretchr/testify/require" + corev1 "k8s.io/api/core/v1" + "k8s.io/client-go/kubernetes/fake" + + "github.com/ai-dynamo/snapshot/agent/internal/nsmount" + "github.com/ai-dynamo/snapshot/agent/internal/types" + "github.com/ai-dynamo/snapshot/api/compat" + snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" +) + +const gatedRestoreContainer = "main" + +func gatedRestoreMappings() []snapshotv1alpha1.RestoreContainerMapping { + return []snapshotv1alpha1.RestoreContainerMapping{{ + Source: gatedRestoreContainer, + Destination: gatedRestoreContainer, + }} +} + +type comparisonCall struct { + gate compat.Gate + source compat.Facts + target compat.Facts +} + +// comparisonSpy stands in for the policy table so a test can decide the verdict +// while the table itself is still being filled in. +type comparisonSpy struct { + mismatches []compat.Mismatch + calls []comparisonCall +} + +func (s *comparisonSpy) compare(gate compat.Gate, source, target compat.Facts) []compat.Mismatch { + s.calls = append(s.calls, comparisonCall{gate: gate, source: source, target: target}) + return s.mismatches +} + +// gatedRestore is one restore driven through the compatibility gates: a +// controller, the restore pod, the snapshot it names, the artifact it resolves +// to, and a stubbed comparison. Every test about what a gate does starts from +// one of these, whether the verdict it forces lets the restore through or not. +type gatedRestore struct { + controller *NodeController + pod *corev1.Pod + artifact *restoreArtifact + logs *logRecorder + comparison *comparisonSpy +} + +func newGatedRestore(t *testing.T, mismatches ...compat.Mismatch) *gatedRestore { + t.Helper() + pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + snapshot, content := readySnapshotObjects() + w := makeTestController(t, pod, snapshot, content) + logs := &logRecorder{} + w.log = logr.New(&recordingSink{recorder: logs}) + comparison := &comparisonSpy{mismatches: mismatches} + w.compareFn = comparison.compare + + path := writeTestArtifact(t, w.config.Storage.BasePath, string(content.UID), &types.CheckpointManifest{ + Artifact: types.ArtifactManifest{ContentUID: string(content.UID), ContainerName: gatedRestoreContainer}, + }) + + return &gatedRestore{ + controller: w, + pod: pod, + artifact: &restoreArtifact{ + SnapshotName: snapshot.Name, + ContentUID: string(content.UID), + SourceContainerName: gatedRestoreContainer, + Path: path, + }, + logs: logs, + comparison: comparison, + } +} + +// writeTestArtifact creates the artifact directory a restore pod resolves to, +// optionally with a manifest in it. +func writeTestArtifact(t *testing.T, basePath, contentUID string, manifest *types.CheckpointManifest) string { + t.Helper() + path, err := nsmount.ResolveArtifactPath(basePath, contentUID, gatedRestoreContainer) + require.NoError(t, err) + require.NoError(t, os.MkdirAll(path, 0o700)) + if manifest != nil { + require.NoError(t, types.WriteManifest(path, manifest)) + } + return path +} + +// reconcile drives the restore the way the queue worker does. +func (r *gatedRestore) reconcile(t *testing.T) { + t.Helper() + processQueuedRestorePod(t, r.controller, r.pod) +} + +// runRestore drives the restore worker directly, which is where the second gate +// reports from. +func (r *gatedRestore) runRestore(t *testing.T) bool { + t.Helper() + return r.controller.restorePodContainers( + context.Background(), + r.pod, + &restorePlan{artifact: r.artifact, mappings: gatedRestoreMappings()}, + fmt.Sprintf("%s/%s", r.pod.Namespace, r.pod.Name), + ) +} + +func (r *gatedRestore) clientset(t *testing.T) *fake.Clientset { + t.Helper() + clientset, ok := r.controller.clientset.(*fake.Clientset) + require.Truef(t, ok, "controller clientset is %T, want *fake.Clientset", r.controller.clientset) + return clientset +} + +// condition reads the condition off the last status apply, which is how the +// agent publishes a verdict. +func (r *gatedRestore) condition(t *testing.T) corev1.PodCondition { + t.Helper() + var applied struct { + Status struct { + Conditions []corev1.PodCondition `json:"conditions"` + } `json:"status"` + } + require.NoError(t, json.Unmarshal(lastPodStatusApply(t, r.controller).GetPatch(), &applied)) + require.Len(t, applied.Status.Conditions, 1) + return applied.Status.Conditions[0] +} + +// events returns every event emitted under one reason, so a test can assert on +// how many there are and not only that there was one. +func (r *gatedRestore) events(t *testing.T, reason string) []*corev1.Event { + t.Helper() + return eventsForReason(r.clientset(t), reason) +} + +type logRecord struct { + message string + fields map[string]any +} + +// logRecorder captures what the agent logged, so a test can assert on the field +// an operator greps for rather than on a formatted sentence. +type logRecorder struct { + mu sync.Mutex + records []logRecord +} + +func (r *logRecorder) add(message string, inherited, keysAndValues []any) { + r.mu.Lock() + defer r.mu.Unlock() + fields := map[string]any{} + for _, pairs := range [][]any{inherited, keysAndValues} { + for i := 0; i+1 < len(pairs); i += 2 { + if key, ok := pairs[i].(string); ok { + fields[key] = pairs[i+1] + } + } + } + r.records = append(r.records, logRecord{message: message, fields: fields}) +} + +// fieldsOf returns the fields of every record logged under one message, +// including the ones inherited from the logger. +func (r *logRecorder) fieldsOf(message string) []map[string]any { + r.mu.Lock() + defer r.mu.Unlock() + var matched []map[string]any + for _, record := range r.records { + if record.message == message { + matched = append(matched, record.fields) + } + } + return matched +} + +// refusalLog returns the fields of the one refusal the agent logged, which is +// what an operator reads to learn why a restore was turned down. +func (r *gatedRestore) refusalLog(t *testing.T) map[string]any { + t.Helper() + logged := r.logs.fieldsOf("Refusing restore; this node cannot run the checkpoint") + require.Len(t, logged, 1) + return logged[0] +} + +type recordingSink struct { + recorder *logRecorder + values []any +} + +var _ logr.LogSink = (*recordingSink)(nil) + +func (s *recordingSink) Init(logr.RuntimeInfo) {} +func (s *recordingSink) Enabled(int) bool { return true } + +func (s *recordingSink) Info(_ int, message string, keysAndValues ...any) { + s.recorder.add(message, s.values, keysAndValues) +} + +func (s *recordingSink) Error(err error, message string, keysAndValues ...any) { + s.recorder.add(message, s.values, append(keysAndValues, "error", err)) +} + +// WithValues has to accumulate: the gates log through a logger that already +// carries the pod and container, and a test asserting on those must still see +// them on the record. +func (s *recordingSink) WithValues(keysAndValues ...any) logr.LogSink { + values := make([]any, 0, len(s.values)+len(keysAndValues)) + values = append(values, s.values...) + values = append(values, keysAndValues...) + return &recordingSink{recorder: s.recorder, values: values} +} + +func (s *recordingSink) WithName(string) logr.LogSink { return s } diff --git a/agent/internal/controller/compat_test.go b/agent/internal/controller/compat_test.go new file mode 100644 index 00000000..d7f2cb0c --- /dev/null +++ b/agent/internal/controller/compat_test.go @@ -0,0 +1,460 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package controller + +import ( + "context" + "errors" + "runtime" + "testing" + + "github.com/go-logr/logr" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" + + "github.com/ai-dynamo/snapshot/agent/internal/executor" + snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" + "github.com/ai-dynamo/snapshot/agent/internal/types" + "github.com/ai-dynamo/snapshot/api/compat" + snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" +) + +// A checkpoint whose manifest cannot be read is not incompatible, it is broken. +// The restore path reads the manifest again and reports that; refusing here +// would report the wrong outcome and hide the real error. +func TestPreflightCompatibilityAllowsUnreadableManifest(t *testing.T) { + r := newGatedRestore(t, compat.Mismatch{Check: "kernel-version"}) + path := writeTestArtifact(t, r.controller.config.Storage.BasePath, "no-manifest-here", nil) + + err := r.controller.preflightCompatibility(r.pod, &restoreArtifact{ + SourceContainerName: gatedRestoreContainer, + Path: path, + }, gatedRestoreMappings()) + + require.NoError(t, err) + assert.Empty(t, r.comparison.calls, "comparison ran without a manifest") +} + +func TestPreflightCompatibilityComparesRecordedFacts(t *testing.T) { + r := newGatedRestore(t) + path := writeTestArtifact(t, r.controller.config.Storage.BasePath, "mounted-content", &types.CheckpointManifest{ + Artifact: types.ArtifactManifest{ContentUID: "mounted-content", ContainerName: gatedRestoreContainer}, + CRIUDump: types.CRIUDumpManifest{ + ExtMnt: map[string]string{ + "/model-cache": "/model-cache", + "/etc/hosts": "/etc/hosts", + }, + }, + }) + + err := r.controller.preflightCompatibility(r.pod, &restoreArtifact{ + SourceContainerName: gatedRestoreContainer, + Path: path, + }, gatedRestoreMappings()) + + require.NoError(t, err) + require.Len(t, r.comparison.calls, 1) + assert.Equal(t, compat.GatePreflight, r.comparison.calls[0].gate) + assert.Equal(t, []string{"/etc/hosts", "/model-cache"}, r.comparison.calls[0].source.ExternalizedMounts) +} + +// The gate compares the checkpoint against every destination, so each target +// side has to describe this node and its destination container. +func TestPreflightCompatibilityDescribesEveryRestoreTarget(t *testing.T) { + r := newGatedRestore(t) + r.controller.config.HostKernelVersion = "5.15.0-1071-aws" + r.pod.Spec.Containers[0].Image = "nvcr.io/nvidia/tritonserver:24.09-py3" + r.pod.Status.ContainerStatuses[0].ImageID = "sha256:deadbeef" + r.pod.Spec.Containers = append(r.pod.Spec.Containers, corev1.Container{ + Name: "engine-1", + Image: "nvcr.io/nvidia/tritonserver:25.01-py3", + }) + r.pod.Status.ContainerStatuses = append(r.pod.Status.ContainerStatuses, corev1.ContainerStatus{ + Name: "engine-1", + ImageID: "sha256:cafebabe", + }) + + mappings := append(gatedRestoreMappings(), snapshotv1alpha1.RestoreContainerMapping{ + Source: gatedRestoreContainer, + Destination: "engine-1", + }) + require.NoError(t, r.controller.preflightCompatibility(r.pod, r.artifact, mappings)) + + require.Len(t, r.comparison.calls, 2) + assert.Equal(t, compat.Facts{ + CPUArch: runtime.GOARCH, + KernelVersion: "5.15.0-1071-aws", + Image: "nvcr.io/nvidia/tritonserver:24.09-py3", + }, r.comparison.calls[0].target) + assert.Equal(t, compat.Facts{ + CPUArch: runtime.GOARCH, + KernelVersion: "5.15.0-1071-aws", + Image: "nvcr.io/nvidia/tritonserver:25.01-py3", + }, r.comparison.calls[1].target) +} + +// The refusal event is its own reason, so alerting that pages on restore +// failures does not fire on a restore that was never attempted. +func TestRefusalEmitsOneIncompatibleEventAtBothGates(t *testing.T) { + mismatch := compat.Mismatch{Check: "cpu-arch", Source: "amd64", Target: "arm64"} + wantMessage := "cpu-arch: source amd64, target arm64" + + assertOneEvent := func(t *testing.T, r *gatedRestore) { + t.Helper() + events := r.events(t, restoreIncompatibleReason) + require.Len(t, events, 1) + assert.Equal(t, wantMessage, events[0].Message) + assert.Equal(t, corev1.EventTypeWarning, events[0].Type) + assert.Empty(t, r.events(t, restoreFailedReason), "refusal also reported a restore failure") + } + + t.Run("preflight gate", func(t *testing.T) { + r := newGatedRestore(t, mismatch) + + r.reconcile(t) + + assertOneEvent(t, r) + assert.Empty(t, r.events(t, restoreRequestedReason), "refused restore still announced a request") + }) + + t.Run("inspect gate", func(t *testing.T) { + r := newGatedRestore(t) + r.controller.restoreFn = refuseWith(mismatch) + + r.runRestore(t) + + assertOneEvent(t, r) + }) +} + +// The pod carries its own verdict on the condition every other restore outcome +// is published on, so a refusal is visible to anything that reads pods and is +// told apart from a failure by its reason alone. +func TestRefusalPublishesTheRestoredConditionAtBothGates(t *testing.T) { + mismatch := compat.Mismatch{Check: "gpu-model", Source: "Tesla T4", Target: "NVIDIA A100-SXM4-40GB"} + wantMessage := "gpu-model: source Tesla T4, target NVIDIA A100-SXM4-40GB" + + assertPublished := func(t *testing.T, r *gatedRestore) { + t.Helper() + condition := r.condition(t) + assert.Equal(t, corev1.PodConditionType(snapshotv1alpha1.RestoredCondition), condition.Type) + assert.Equal(t, corev1.ConditionFalse, condition.Status) + assert.Equal(t, restoreIncompatibleReason, condition.Reason) + // The same sentence the log line and the event carry, so a reader who + // starts from the pod does not get a different answer. + assert.Equal(t, wantMessage, condition.Message) + assert.Equal(t, wantMessage, r.refusalLog(t)["reason"]) + } + + t.Run("preflight gate", func(t *testing.T) { + r := newGatedRestore(t, mismatch) + + r.reconcile(t) + + assertPublished(t, r) + }) + + t.Run("inspect gate", func(t *testing.T) { + r := newGatedRestore(t) + r.controller.restoreFn = refuseWith(mismatch) + + r.runRestore(t) + + assertPublished(t, r) + }) +} + +func TestRecordRestoreResultsReportsEveryIncompatibleDestination(t *testing.T) { + r := newGatedRestore(t) + requeue := r.controller.recordRestoreResults( + context.Background(), + r.pod, + r.artifact, + []restoreResult{ + {destination: "engine-0", state: restoreResultIncompatible, reason: "gpu-model: source Tesla T4, target NVIDIA A100"}, + {destination: "engine-1", state: restoreResultIncompatible, reason: "gpu-count: source 1, target 0"}, + }, + ) + + assert.False(t, requeue) + condition := r.condition(t) + assert.Equal(t, restoreIncompatibleReason, condition.Reason) + assert.Contains(t, condition.Message, "engine-0: gpu-model") + assert.Contains(t, condition.Message, "engine-1: gpu-count") + assert.Empty(t, r.events(t, restoreSucceededReason)) +} + +// A refusal is terminal like any other restore failure: the pod is not compared +// again on the next resync, and the queue worker says why rather than silently +// dropping it. +func TestReconcileRestorePodLeavesARefusedPodAlone(t *testing.T) { + r := newGatedRestore(t, compat.Mismatch{Check: "cpu-arch", Source: "amd64", Target: "arm64"}) + r.pod.Status.Conditions = append(r.pod.Status.Conditions, corev1.PodCondition{ + Type: corev1.PodConditionType(snapshotv1alpha1.RestoredCondition), + Status: corev1.ConditionFalse, + Reason: restoreIncompatibleReason, + Message: "cpu-arch: source amd64, target arm64", + }) + + r.reconcile(t) + + assert.Empty(t, r.comparison.calls, "already refused restore was compared again") + assert.Empty(t, r.events(t, restoreIncompatibleReason), "already refused restore was refused again") + assert.Len(t, r.events(t, restoreAlreadyFailedReason), 1) + assert.Empty(t, r.controller.inFlight, "already refused restore claimed an attempt") +} + +// The escape hatches: with either one set, neither gate runs, so a checkpoint +// the policy table would turn down is still attempted. +func TestSkipCompatCheckTurnsOffTheGates(t *testing.T) { + mismatch := compat.Mismatch{Check: "cpu-arch", Source: "amd64", Target: "arm64"} + + // Lets the restore start and end quickly, since the point here is only + // whether the gate let it through. + stopEarly := func(r *gatedRestore) { + r.controller.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (int, error) { + return 0, errors.New("test restore stopped") + } + } + + t.Run("the pod annotation turns it off", func(t *testing.T) { + r := newGatedRestore(t, mismatch) + r.pod.Annotations[snapshotv1alpha1.SkipCompatCheckAnnotation] = "true" + stopEarly(r) + + r.reconcile(t) + + assert.Empty(t, r.comparison.calls, "skipped gate compared anyway") + assert.Empty(t, r.events(t, restoreIncompatibleReason), "skipped gate refused the restore") + }) + + // A node with the gate off skips every restore it handles, whether or not + // the pod asked for it. + t.Run("the node config turns it off for an unannotated pod", func(t *testing.T) { + r := newGatedRestore(t, mismatch) + r.controller.config.Restore.SkipCompatCheck = true + stopEarly(r) + + r.reconcile(t) + + assert.Empty(t, r.comparison.calls, "skipped gate compared anyway") + assert.Empty(t, r.events(t, restoreIncompatibleReason), "skipped gate refused the restore") + }) + + // Gate B is inside the executor, past the point where either switch can be + // read again, so the decision travels with the request. Without it, a + // skipped restore would still be refused a few steps later. + t.Run("the decision travels to the second gate", func(t *testing.T) { + for _, tc := range []struct { + name string + set func(*gatedRestore) + want bool + }{ + {name: "checked", set: func(*gatedRestore) {}}, + { + name: "skipped by pod", + set: func(r *gatedRestore) { + r.pod.Annotations[snapshotv1alpha1.SkipCompatCheckAnnotation] = "true" + }, + want: true, + }, + { + name: "skipped by node", + set: func(r *gatedRestore) { r.controller.config.Restore.SkipCompatCheck = true }, + want: true, + }, + } { + t.Run(tc.name, func(t *testing.T) { + r := newGatedRestore(t) + tc.set(r) + var requested executor.RestoreRequest + r.controller.restoreFn = func(_ context.Context, _ snapshotruntime.Runtime, _ logr.Logger, req executor.RestoreRequest, _ executor.RestoreMounter) (int, error) { + requested = req + return 0, errors.New("test restore stopped") + } + + r.reconcile(t) + + assert.Equal(t, tc.want, requested.SkipCompatCheck) + }) + } + }) + + // The node switch is read per restore, not once at startup, which is what + // makes flipping the ConfigMap enough to be heard. + t.Run("the node config is re-read for every restore", func(t *testing.T) { + r := newGatedRestore(t, mismatch) + reads := 0 + r.controller.skipCompatCheckFn = func() bool { + reads++ + return reads > 1 + } + stopEarly(r) + + r.reconcile(t) + require.Len(t, r.comparison.calls, 1, "gate did not run while the switch was off") + + r.pod.Status.Conditions = nil + r.controller.handledRestores.Delete(string(r.pod.UID)) + r.reconcile(t) + + assert.Len(t, r.comparison.calls, 1, "gate ran after the switch was flipped on") + assert.Equal(t, 2, reads) + }) + + // The annotation has to reach a pod the gate already turned down, or the + // only way out of a wrong refusal is deleting and recreating the pod. + t.Run("it reopens a pod that was already refused", func(t *testing.T) { + r := newGatedRestore(t, mismatch) + stopEarly(r) + + r.reconcile(t) + require.Len(t, r.events(t, restoreIncompatibleReason), 1, "the gate did not refuse the restore") + + r.pod.Status.Conditions = append(r.pod.Status.Conditions, corev1.PodCondition{ + Type: corev1.PodConditionType(snapshotv1alpha1.RestoredCondition), + Status: corev1.ConditionFalse, + Reason: restoreIncompatibleReason, + Message: mismatch.Reason(), + }) + r.pod.Annotations[snapshotv1alpha1.SkipCompatCheckAnnotation] = "true" + r.comparison.calls = nil + + r.reconcile(t) + + assert.Empty(t, r.comparison.calls, "the reopened restore was compared anyway") + assert.Len(t, r.events(t, restoreIncompatibleReason), 1, "the reopened restore was refused again") + }) +} + +// Both gates log the same sentence for the same refusal, and each names the gate +// it came from, so an operator greps one field and learns how far the restore got +// before the node turned it down. +func TestRefusalIsLoggedWithTheSameReasonAtBothGates(t *testing.T) { + mismatch := compat.Mismatch{Check: "memory-limit", Source: "32Gi", Target: "1Gi"} + wantReason := "memory-limit: source 32Gi, target 1Gi" + + t.Run("preflight gate", func(t *testing.T) { + r := newGatedRestore(t, mismatch) + + r.reconcile(t) + + refusal := r.refusalLog(t) + assert.Equal(t, wantReason, refusal["reason"]) + assert.Equal(t, string(compat.GatePreflight), refusal["gate"]) + // The refusal names the pod it belongs to, so a reader does not have to + // correlate on time. + assert.Equal(t, "inference/restore-worker", refusal["pod"]) + }) + + t.Run("inspect gate", func(t *testing.T) { + r := newGatedRestore(t) + r.controller.restoreFn = refuseWith(mismatch) + + r.runRestore(t) + + refusal := r.refusalLog(t) + assert.Equal(t, wantReason, refusal["reason"]) + assert.Equal(t, string(compat.GateInspect), refusal["gate"]) + }) +} + +// A refusal from the second gate is not a CRIU failure, so it neither reports +// one nor kills the placeholder: killing it would restart the container straight +// back into the same answer. +func TestRunRestoreTreatsIncompatibleAsTerminal(t *testing.T) { + r := newGatedRestore(t) + rt := &fakeRuntime{} + r.controller.runtime = rt + sentinels := 0 + r.controller.writeControlSentinelFn = func(int, string) error { + sentinels++ + return nil + } + r.controller.restoreFn = refuseWith(compat.Mismatch{Check: "cpu-arch", Source: "amd64", Target: "arm64"}) + + requeue := r.runRestore(t) + + assert.False(t, requeue, "a refusal asked to be driven again") + assert.Empty(t, r.events(t, restoreFailedReason), "refusal reported itself as a restore failure") + assert.Zero(t, sentinels, "refusal released the workload") + assert.Empty(t, rt.resolvedContainerIDs, "refusal reached the placeholder kill path") +} + +// The gate runs in preflight, before the restore is entered at all, so a refusal +// leaves no in-flight entry and no restore worker behind. +func TestReconcileRestorePodRefusesBeforeEnteringRestore(t *testing.T) { + r := newGatedRestore(t, compat.Mismatch{Check: "memory-limit", Source: "32Gi", Target: "1Gi"}) + r.controller.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (int, error) { + t.Error("a refused restore was entered") + return 0, nil + } + + r.reconcile(t) + + assert.Len(t, r.comparison.calls, 1) + assert.Empty(t, r.events(t, restoreRequestedReason), "refused restore still announced a request") + assert.Empty(t, r.controller.inFlight, "refused restore claimed an attempt") +} + +func refuseWith(mismatches ...compat.Mismatch) func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (int, error) { + return func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (int, error) { + return 0, compat.NewIncompatibleError(compat.GateInspect, mismatches) + } +} + +// The facts recorded at capture describe one container, so a multi-container pod +// must not contribute another container's image or limits. +func TestPodFactsReadTheTargetContainer(t *testing.T) { + pod := &corev1.Pod{ + Spec: corev1.PodSpec{Containers: []corev1.Container{ + { + Name: "sidecar", + Image: "busybox:1.36", + Resources: corev1.ResourceRequirements{Limits: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("1"), + corev1.ResourceMemory: resource.MustParse("256Mi"), + }}, + }, + { + Name: "main", + Image: "nvcr.io/nvidia/tritonserver:24.09-py3", + Resources: corev1.ResourceRequirements{Limits: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("4"), + corev1.ResourceMemory: resource.MustParse("16Gi"), + }}, + }, + }}, + Status: corev1.PodStatus{ContainerStatuses: []corev1.ContainerStatus{ + {Name: "sidecar", ImageID: "sha256:sidecar"}, + {Name: "main", ImageID: "docker-pullable://nvcr.io/nvidia/tritonserver@sha256:deadbeef"}, + }}, + } + + assert.Equal(t, compat.Facts{ + Image: "nvcr.io/nvidia/tritonserver:24.09-py3", + CPULimit: "4", + MemoryLimit: "16Gi", + }, podFacts(pod, "main")) +} + +// A fact the pod does not carry stays unknown. An unlimited container is not a +// container limited to zero. +func TestPodFactsLeaveWhatThePodDoesNotSayUnknown(t *testing.T) { + pod := &corev1.Pod{ + Spec: corev1.PodSpec{Containers: []corev1.Container{{ + Name: "main", + Image: "busybox:1.36", + Resources: corev1.ResourceRequirements{Limits: corev1.ResourceList{ + corev1.ResourceMemory: resource.MustParse("16Gi"), + }}, + }}}, + } + + assert.Equal(t, compat.Facts{Image: "busybox:1.36", MemoryLimit: "16Gi"}, podFacts(pod, "main")) + assert.Equal(t, compat.Facts{}, podFacts(pod, "absent"), "a container not in the pod") + assert.Equal(t, compat.Facts{}, podFacts(nil, "main"), "no pod at all") +} diff --git a/agent/internal/controller/controller.go b/agent/internal/controller/controller.go index ddd3dca1..f3d8f6a8 100644 --- a/agent/internal/controller/controller.go +++ b/agent/internal/controller/controller.go @@ -47,6 +47,7 @@ import ( "github.com/ai-dynamo/snapshot/agent/internal/nsmount" snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" "github.com/ai-dynamo/snapshot/agent/internal/types" + "github.com/ai-dynamo/snapshot/api/compat" snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" ) @@ -71,6 +72,12 @@ type NodeController struct { sendSignalFn func(logr.Logger, int, syscall.Signal, string) error restoreQueue workqueue.TypedDelayingInterface[client.ObjectKey] restorePodLister corev1listers.PodLister + compareFn func(compat.Gate, compat.Facts, compat.Facts) []compat.Mismatch + + // skipCompatCheckFn is read once per restore rather than at startup, so the + // node-wide switch can be flipped without a DaemonSet rollout. Injected so + // the controller never learns where the config file lives. + skipCompatCheckFn func() bool inFlight map[string]struct{} inFlightMu sync.Mutex @@ -89,6 +96,9 @@ type restoreArtifact struct { ContentUID string SourceContainerName string Path string + // SkipCompatCheck is decided once in preflight and carried from there, so + // the gate inside the restore reaches the same answer as the one before it. + SkipCompatCheck bool } type restoreTarget struct { @@ -108,11 +118,13 @@ const ( restoreResultPending restoreResultState = iota restoreResultSucceeded restoreResultFailed + restoreResultIncompatible ) type restoreResult struct { destination string state restoreResultState + reason string } type restorePendingError struct { @@ -153,6 +165,10 @@ const ( snapshotEventComponent = "snapshot" restoreSafetyRequeueInterval = 30 * time.Second + // restoreIncompatibleReason marks a restore the node turned down before + // attempting it, because the checkpoint cannot run here. + restoreIncompatibleReason = "RestoreIncompatible" + // snapshotContentResyncInterval re-drives every PodSnapshotContent work order so a // not-yet-Ready source pod is re-checked for quiesce without a busy loop. snapshotContentResyncInterval = 10 * time.Second @@ -162,10 +178,13 @@ const ( var podSnapshotContentGVR = snapshotv1alpha1.GroupVersion.WithResource("podsnapshotcontents") // NewNodeController creates the node-local controller that runs inside snapshot-agent. +// skipCompatCheckFn is read per restore; passing nil pins the switch to the +// configuration the agent started with. func NewNodeController( cfg *types.AgentConfig, rt snapshotruntime.Runtime, log logr.Logger, + skipCompatCheckFn func() bool, ) (*NodeController, error) { restConfig, err := rest.InClusterConfig() if err != nil { @@ -192,7 +211,7 @@ func NewNodeController( } nsm := nsmount.New(log) - return newDefaultController(cfg, clientset, typedClient, dynClient, rt, nsm, log), nil + return newDefaultController(cfg, clientset, typedClient, dynClient, rt, nsm, log, skipCompatCheckFn), nil } func newDefaultController( @@ -203,6 +222,7 @@ func newDefaultController( rt snapshotruntime.Runtime, injector executor.RestoreMounter, log logr.Logger, + skipCompatCheckFn func() bool, ) *NodeController { w := &NodeController{ config: cfg, @@ -223,8 +243,13 @@ func newDefaultController( writeControlSentinelFn: snapshotruntime.WriteControlSentinel, controlSentinelExistsFn: snapshotruntime.ControlSentinelExists, sendSignalFn: snapshotruntime.SendSignalToPID, + compareFn: compat.Compare, } w.checkpointFn = w.executorCheckpoint + w.skipCompatCheckFn = skipCompatCheckFn + if w.skipCompatCheckFn == nil { + w.skipCompatCheckFn = func() bool { return w.config.Restore.SkipCompatCheck } + } return w } @@ -426,11 +451,14 @@ func (w *NodeController) processRestoreQueueItem(ctx context.Context, key client return } pod = pod.DeepCopy() - if w.restoreHandled(pod) { + if w.reopenedAfterRefusal(pod) { + // The skip request is the way back for a pod the gates turned down, so + // it has to clear the in-process marker as well as the condition below. + w.handledRestores.Delete(string(pod.UID)) + } else if w.restoreHandled(pod) { requeue = w.removeRestoreFinalizerWithEvent(ctx, pod) return - } - if isRestoreTerminal(pod) { + } else if isRestoreTerminal(pod) { requeue = w.handleTerminalRestorePod(ctx, pod) return } @@ -511,6 +539,12 @@ func (w *NodeController) preflightRestore(ctx context.Context, pod *corev1.Pod) if err != nil { return nil, err } + // Gate A: the earliest point the checkpoint's own record of what it was + // captured on is readable, and still before any of the restore is attempted. + artifact.SkipCompatCheck = w.skipCompatCheckRequested(pod) + if err := w.preflightCompatibility(pod, artifact, mappings); err != nil { + return nil, err + } if w.config.CRIU.TcpEstablished && pod.Status.PodIP == "" { return nil, newRestorePendingError("PodIPPending", fmt.Sprintf("Waiting for restore Pod %s/%s to receive an IP address", pod.Namespace, pod.Name)) } @@ -717,33 +751,50 @@ func (w *NodeController) restorePodContainers(ctx context.Context, pod *corev1.P // recordRestoreResults publishes the aggregate Pod outcome after every worker // in the current pass has returned. func (w *NodeController) recordRestoreResults(ctx context.Context, pod *corev1.Pod, artifact *restoreArtifact, results []restoreResult) bool { - byState := make(map[restoreResultState][]string, 3) + byState := make(map[restoreResultState][]string, 4) + var incompatibilityReasons []string for _, result := range results { byState[result.state] = append(byState[result.state], result.destination) + if result.state == restoreResultIncompatible { + reason := result.reason + if len(results) > 1 { + reason = fmt.Sprintf("%s: %s", result.destination, reason) + } + incompatibilityReasons = append(incompatibilityReasons, reason) + } } succeeded := byState[restoreResultSucceeded] failed := byState[restoreResultFailed] + incompatible := byState[restoreResultIncompatible] pending := byState[restoreResultPending] if len(pending) != 0 { - message := fmt.Sprintf( - "Restore from PodSnapshot %s remains in progress: %d succeeded, %d failed, %d pending (%s)", - artifact.SnapshotName, len(succeeded), len(failed), len(pending), strings.Join(pending, ", "), - ) + message := fmt.Sprintf("Restore from PodSnapshot %s remains in progress: %d succeeded, %d failed, %d pending (%s)", artifact.SnapshotName, len(succeeded), len(failed), len(pending), strings.Join(pending, ", ")) + if len(incompatible) != 0 { + message = fmt.Sprintf("Restore from PodSnapshot %s remains in progress: %d succeeded, %d failed, %d incompatible, %d pending (%s)", artifact.SnapshotName, len(succeeded), len(failed), len(incompatible), len(pending), strings.Join(pending, ", ")) + } if err := w.applyRestoredCondition(ctx, pod, corev1.ConditionFalse, restoreInProgressReason, message); err != nil { emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, corev1.EventTypeWarning, restoreStatusUpdateFailedReason, err.Error()) } return true } - if len(failed) == 0 { + if len(failed) == 0 && len(incompatible) == 0 { message := fmt.Sprintf("Restored %d destination container(s) from PodSnapshot %s: %s", len(succeeded), artifact.SnapshotName, strings.Join(succeeded, ", ")) return w.finishRestore(ctx, pod, corev1.ConditionTrue, restoreSucceededReason, message) != nil } if len(succeeded) != 0 { - message := fmt.Sprintf("Restored %d of %d destination containers from PodSnapshot %s; failed: %s", len(succeeded), len(results), artifact.SnapshotName, strings.Join(failed, ", ")) + notRestored := append(append([]string{}, failed...), incompatible...) + message := fmt.Sprintf("Restored %d of %d destination containers from PodSnapshot %s; not restored: %s", len(succeeded), len(results), artifact.SnapshotName, strings.Join(notRestored, ", ")) return w.finishRestore(ctx, pod, corev1.ConditionFalse, restorePartiallySucceededReason, message) != nil } + if len(failed) == 0 { + return w.finishRestore(ctx, pod, corev1.ConditionFalse, restoreIncompatibleReason, strings.Join(incompatibilityReasons, "; ")) != nil + } + if len(incompatible) != 0 { + message := fmt.Sprintf("Restore failed for %d destination container(s) and refused %d incompatible destination(s) from PodSnapshot %s", len(failed), len(incompatible), artifact.SnapshotName) + return w.finishRestore(ctx, pod, corev1.ConditionFalse, restoreFailedReason, message) != nil + } message := fmt.Sprintf("Restore failed for all %d destination container(s) from PodSnapshot %s: %s", len(failed), artifact.SnapshotName, strings.Join(failed, ", ")) return w.finishRestore(ctx, pod, corev1.ConditionFalse, restoreFailedReason, message) != nil } @@ -773,6 +824,13 @@ func (w *NodeController) restoreDestination( emitPodEvent(ctx, w.clientset, log, pod, snapshotEventComponent, corev1.EventTypeNormal, restoreRequestedReason, fmt.Sprintf("Restore requested from PodSnapshot %s for destination %s", artifact.SnapshotName, destination)) if err := w.runRestore(ctx, pod, artifact, destination, containerID, startedAt, recovering); err != nil { + var incompatible *compat.IncompatibleError + if errors.As(err, &incompatible) { + result.state = restoreResultIncompatible + result.reason = compat.Reasons(incompatible.Mismatches) + w.logRestoreRefusal(pod, incompatible, result.reason) + return result + } result.state = restoreResultFailed log.Error(err, "Restore controller worker failed") emitPodEvent(ctx, w.clientset, log, pod, snapshotEventComponent, corev1.EventTypeWarning, "RestoreWorkerFailed", err.Error()) @@ -854,6 +912,14 @@ func (w *NodeController) runRestore(ctx context.Context, pod *corev1.Pod, artifa placeholderHostPID, err := op.executeRestore(restoreCtx) if err != nil { + var incompatible *compat.IncompatibleError + if errors.As(err, &incompatible) { + // Nothing was attempted, so there is no half-restored process to + // clean up. The placeholder is deliberately left running: killing it + // would restart the container straight back into the same answer. + return incompatible + } + var cleanupErr *executor.RestoreCleanupError if !errors.As(err, &cleanupErr) { return op.failRestore(ctx, err) @@ -917,6 +983,7 @@ func (op *restoreOperation) executeRestore(ctx context.Context) (int, error) { TargetPodIP: op.pod.Status.PodIP, ArtifactContainerName: op.artifact.SourceContainerName, DestinationContainerName: op.destination, + SkipCompatCheck: op.artifact.SkipCompatCheck, Clientset: w.clientset, } return w.restoreFn(ctx, w.runtime, op.log, req, w.injector) @@ -1073,6 +1140,11 @@ func (w *NodeController) failRestorePod(ctx context.Context, pod *corev1.Pod, ca } func (w *NodeController) handleRestorePreflightError(ctx context.Context, pod *corev1.Pod, cause error) bool { + var incompatible *compat.IncompatibleError + if errors.As(cause, &incompatible) { + return w.refuseRestore(ctx, pod, incompatible) + } + var pending *restorePendingError if !errors.As(cause, &pending) { return w.failRestorePod(ctx, pod, cause) @@ -1240,8 +1312,13 @@ func isRestorePartiallySucceeded(pod *corev1.Pod) bool { func isRestoreTerminal(pod *corev1.Pod) bool { condition := findRestoredCondition(pod) - return isRestoreSucceeded(pod) || isRestorePartiallySucceeded(pod) || - (condition != nil && condition.Status == corev1.ConditionFalse && condition.Reason == restoreFailedReason) + if isRestoreSucceeded(pod) || isRestorePartiallySucceeded(pod) { + return true + } + if condition == nil || condition.Status != corev1.ConditionFalse { + return false + } + return condition.Reason == restoreFailedReason || condition.Reason == restoreIncompatibleReason } func isRestorePodActive(pod *corev1.Pod) bool { diff --git a/agent/internal/controller/controller_test.go b/agent/internal/controller/controller_test.go index 06a23ba7..8143e3a0 100644 --- a/agent/internal/controller/controller_test.go +++ b/agent/internal/controller/controller_test.go @@ -38,6 +38,7 @@ import ( "github.com/ai-dynamo/snapshot/agent/internal/nsmount" snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" "github.com/ai-dynamo/snapshot/agent/internal/types" + "github.com/ai-dynamo/snapshot/api/compat" snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" ) @@ -78,6 +79,10 @@ func (r *fakeRuntime) ResolveContainerByPod(_ context.Context, _, _, _ string) ( return 0, nil, errors.New("not implemented") } +func (r *fakeRuntime) ResolveContainerImageID(_ context.Context, _ string) (string, error) { + return "", errors.New("not implemented") +} + func (r *fakeRuntime) Close() error { return nil } type noopInjector struct{} @@ -106,11 +111,24 @@ func TestNewDefaultControllerSetsDefaultOperations(t *testing.T) { &fakeRuntime{}, noopInjector{}, testr.New(t), + nil, ) t.Cleanup(w.restoreQueue.ShutDown) if w.checkpointFn == nil || w.restoreFn == nil || w.writeControlSentinelFn == nil || w.controlSentinelExistsFn == nil || w.sendSignalFn == nil || w.restoreQueue == nil { t.Fatal("default controller operations must be initialized") } + if w.compareFn == nil { + t.Fatal("default controller must compare restore compatibility") + } + // Without an injected read there is still one to make: the copy the agent + // started with, rather than a nil call on the restore path. + if w.skipCompatCheckFn == nil { + t.Fatal("default controller must resolve the node compatibility switch") + } + w.config.Restore.SkipCompatCheck = true + if !w.skipCompatCheckFn() { + t.Fatal("default controller ignored the configured node compatibility switch") + } } func testScheme(t *testing.T) *runtime.Scheme { @@ -149,11 +167,13 @@ func makeTestController(t *testing.T, pod *corev1.Pod, apiObjects ...runtime.Obj controlSentinelExistsFn: func(int, string) (bool, error) { return false, nil }, sendSignalFn: func(logr.Logger, int, syscall.Signal, string) error { return nil }, restoreQueue: workqueue.NewTypedDelayingQueue[client.ObjectKey](), + compareFn: compat.Compare, log: testr.New(t), holderID: "test-holder", inFlight: make(map[string]struct{}), stopCh: make(chan struct{}), } + w.skipCompatCheckFn = func() bool { return w.config.Restore.SkipCompatCheck } t.Cleanup(w.restoreQueue.ShutDown) return w } @@ -185,6 +205,17 @@ func sawEventReason(clientset *fake.Clientset, reason string) bool { } func eventForReason(clientset *fake.Clientset, reason string) *corev1.Event { + events := eventsForReason(clientset, reason) + if len(events) == 0 { + return nil + } + return events[0] +} + +// eventsForReason returns every event created under one reason, so a test can +// assert on how many there are and not only that there was one. +func eventsForReason(clientset *fake.Clientset, reason string) []*corev1.Event { + var events []*corev1.Event for _, action := range clientset.Actions() { create, ok := action.(clientgotesting.CreateAction) if !ok || create.GetResource().Resource != "events" { @@ -192,10 +223,10 @@ func eventForReason(clientset *fake.Clientset, reason string) *corev1.Event { } event, ok := create.GetObject().(*corev1.Event) if ok && event.Reason == reason { - return event + events = append(events, event) } } - return nil + return events } func pendingRestoreReason(t *testing.T, err error) string { diff --git a/agent/internal/controller/podsnapshotcontent.go b/agent/internal/controller/podsnapshotcontent.go index 0617fc41..059af9e5 100644 --- a/agent/internal/controller/podsnapshotcontent.go +++ b/agent/internal/controller/podsnapshotcontent.go @@ -587,6 +587,7 @@ func (w *NodeController) executorCheckpoint(ctx context.Context, params Checkpoi PodName: params.Pod.Name, PodNamespace: params.Pod.Namespace, PodIP: params.Pod.Status.PodIP, + Pod: podFacts(params.Pod, params.ContainerName), Clientset: w.clientset, } if err := executor.Checkpoint(ctx, w.runtime, log, req, w.config); err != nil { diff --git a/agent/internal/cuda/cuda.go b/agent/internal/cuda/cuda.go index 67430b83..2c2c8483 100644 --- a/agent/internal/cuda/cuda.go +++ b/agent/internal/cuda/cuda.go @@ -19,6 +19,8 @@ import ( "google.golang.org/grpc/credentials/insecure" "k8s.io/client-go/kubernetes" podresourcesv1 "k8s.io/kubelet/pkg/apis/podresources/v1" + + "github.com/ai-dynamo/snapshot/api/compat" ) const ( @@ -30,6 +32,11 @@ const ( // DefaultHelperBinaryPath is the agent-side cuda-checkpoint-helper absolute path. // In the placeholder namespace pass filepath.Join(bundleDir, HelperBinaryName) instead. DefaultHelperBinaryPath = "/usr/local/bin/" + HelperBinaryName + + // nvidiaSMITimeout bounds every nsenter nvidia-smi call. The agent's own + // context carries no deadline, so a hung one would block the worker for good + // and cost the node every restore that followed. + nvidiaSMITimeout = 30 * time.Second ) var podResourcesSocketPath = "/var/lib/kubelet/pod-resources/kubelet.sock" @@ -87,11 +94,15 @@ func GetPodGPUUUIDs(ctx context.Context, podName, podNamespace, containerName st return uuids, nil } -// GetGPUUUIDsViaNvidiaSmi discovers GPU UUIDs by running nvidia-smi inside the -// container's mount and PID namespaces. This is the fallback path when the kubelet -// PodResources API does not report GPU devices (e.g. when GPUs are allocated -// via DRA instead of the NVIDIA device plugin). -func GetGPUUUIDsViaNvidiaSmi(ctx context.Context, hostProcPath string, pid int) ([]string, error) { +// DiscoverVisibleGPUFacts describes the GPUs a container can see, by running +// nvidia-smi inside its mount and PID namespaces. The model and the driver +// version come from the same call as the UUIDs: nothing else on the restore path +// gets to look at the source node's GPUs, so what is not read here cannot be +// compared later. +func DiscoverVisibleGPUFacts(ctx context.Context, hostProcPath string, pid int) (compat.GPUFacts, error) { + ctx, cancel := context.WithTimeout(ctx, nvidiaSMITimeout) + defer cancel() + mountPath := fmt.Sprintf("%s/%d/ns/mnt", strings.TrimRight(hostProcPath, "/"), pid) pidPath := fmt.Sprintf("%s/%d/ns/pid", strings.TrimRight(hostProcPath, "/"), pid) cmd := exec.CommandContext( @@ -100,27 +111,61 @@ func GetGPUUUIDsViaNvidiaSmi(ctx context.Context, hostProcPath string, pid int) fmt.Sprintf("--mount=%s", mountPath), fmt.Sprintf("--pid=%s", pidPath), "--", - "nvidia-smi", "--query-gpu=gpu_uuid", "--format=csv,noheader", + "nvidia-smi", "--query-gpu=gpu_uuid,name,driver_version", "--format=csv,noheader", ) output, err := cmd.Output() if err != nil { - return nil, fmt.Errorf("nvidia-smi via nsenter (pid %d) failed: %w", pid, err) + return compat.GPUFacts{}, fmt.Errorf("nvidia-smi via nsenter (pid %d) failed: %w", pid, err) } - var uuids []string - for _, line := range strings.Split(strings.TrimSpace(string(output)), "\n") { + return parseNvidiaSmiGPUFacts(string(output)), nil +} + +// parseNvidiaSmiGPUFacts reads the unquoted CSV nvidia-smi writes. Splitting on +// commas is safe because nvidia-smi documents name and driver_version as +// alphanumeric strings: https://docs.nvidia.com/deploy/nvidia-smi/index.html +// A row it cannot make sense of still contributes its UUID, because the device +// map is built from UUIDs and must not start failing over a model name. +func parseNvidiaSmiGPUFacts(output string) compat.GPUFacts { + var facts compat.GPUFacts + for _, line := range strings.Split(strings.TrimSpace(output), "\n") { line = strings.TrimSpace(line) - if line != "" { - uuids = append(uuids, line) + if line == "" { + continue + } + fields := strings.SplitN(line, ",", 3) + uuid := strings.TrimSpace(fields[0]) + if uuid == "" { + continue + } + device := compat.GPUDevice{UUID: uuid} + if len(fields) == 3 { + device.ProductName = strings.TrimSpace(fields[1]) + if driverVersion := strings.TrimSpace(fields[2]); driverVersion != "" { + facts.DriverVersion = driverVersion + } } + facts.Devices = append(facts.Devices, device) } - return uuids, nil + return facts } -type visibleGPUDiscovery func(context.Context, string, int) ([]string, error) +type visibleGPUDiscovery func(context.Context, string, int) (compat.GPUFacts, error) // DiscoverGPUUUIDs resolves GPU UUIDs in the container's runtime ordinal order. func DiscoverGPUUUIDs(ctx context.Context, clientset kubernetes.Interface, podName, podNamespace, containerName, hostProcPath string, pid int, log logr.Logger) ([]string, error) { - return discoverGPUUUIDs( + facts, err := DiscoverGPUFacts(ctx, clientset, podName, podNamespace, containerName, hostProcPath, pid, log) + if err != nil { + return nil, err + } + return gpuUUIDsOf(facts), nil +} + +// DiscoverGPUFacts resolves the same GPUs as DiscoverGPUUUIDs, in the same +// order, described by model and driver version wherever nvidia-smi can be +// reached. Whichever path finds the GPUs, the facts come out the same shape, so +// what gets recorded does not depend on how this cluster allocates GPUs. +func DiscoverGPUFacts(ctx context.Context, clientset kubernetes.Interface, podName, podNamespace, containerName, hostProcPath string, pid int, log logr.Logger) (compat.GPUFacts, error) { + return discoverGPUFacts( ctx, clientset, podName, @@ -128,12 +173,12 @@ func DiscoverGPUUUIDs(ctx context.Context, clientset kubernetes.Interface, podNa containerName, hostProcPath, pid, - GetGPUUUIDsViaNvidiaSmi, + DiscoverVisibleGPUFacts, log, ) } -func discoverGPUUUIDs( +func discoverGPUFacts( ctx context.Context, clientset kubernetes.Interface, podName, @@ -143,11 +188,11 @@ func discoverGPUUUIDs( pid int, discoverVisibleGPUs visibleGPUDiscovery, log logr.Logger, -) ([]string, error) { +) (compat.GPUFacts, error) { gpuUUIDs, hasNVIDIADRAAllocation, err := GetGPUUUIDsViaDRAAPI(ctx, clientset, podName, podNamespace, containerName, log) if err != nil { if hasNVIDIADRAAllocation { - return nil, fmt.Errorf("DRA GPU UUID lookup failed: %w", err) + return compat.GPUFacts{}, fmt.Errorf("DRA GPU UUID lookup failed: %w", err) } log.Error( err, @@ -159,43 +204,85 @@ func discoverGPUUUIDs( if hasNVIDIADRAAllocation { if len(gpuUUIDs) == 0 { - return nil, errors.New( + return compat.GPUFacts{}, errors.New( "DRA GPU allocation has no resolvable UUIDs", ) } - visibleGPUUUIDs, err := discoverVisibleGPUs(ctx, hostProcPath, pid) + visible, err := discoverVisibleGPUs(ctx, hostProcPath, pid) if err != nil { - return nil, fmt.Errorf( + return compat.GPUFacts{}, fmt.Errorf( "discover DRA GPUs in container ordinal order: %w", err, ) } - orderedUUIDs, err := orderDRAUUIDsByRuntime(gpuUUIDs, visibleGPUUUIDs) + orderedUUIDs, err := orderDRAUUIDsByRuntime(gpuUUIDs, gpuUUIDsOf(visible)) if err != nil { - return nil, err + return compat.GPUFacts{}, err } log.Info( "resolved DRA GPU UUIDs in container ordinal order", "uuids", orderedUUIDs, ) - return orderedUUIDs, nil + return describeGPUs(orderedUUIDs, visible), nil } gpuUUIDs, err = GetPodGPUUUIDs(ctx, podName, podNamespace, containerName) if err != nil { - return nil, fmt.Errorf("PodResources GPU UUID lookup failed: %w", err) + return compat.GPUFacts{}, fmt.Errorf("PodResources GPU UUID lookup failed: %w", err) } if len(gpuUUIDs) > 0 { - return gpuUUIDs, nil + // This path has its GPUs already and needs nvidia-smi only to describe + // them, so a failure here costs facts, not the checkpoint. + visible, err := discoverVisibleGPUs(ctx, hostProcPath, pid) + if err != nil { + log.V(1).Info("Failed to describe PodResources GPUs; recording their UUIDs alone", + "pid", pid, + "error", err, + ) + return describeGPUs(gpuUUIDs, compat.GPUFacts{}), nil + } + return describeGPUs(gpuUUIDs, visible), nil } log.Info("PodResources API returned no GPU UUIDs, falling back to nvidia-smi", "pid", pid) - gpuUUIDs, err = discoverVisibleGPUs(ctx, hostProcPath, pid) + visible, err := discoverVisibleGPUs(ctx, hostProcPath, pid) if err != nil { - return nil, fmt.Errorf("nvidia-smi GPU UUID fallback failed: %w", err) + return compat.GPUFacts{}, fmt.Errorf("nvidia-smi GPU UUID fallback failed: %w", err) + } + log.Info("nvidia-smi fallback discovered GPU UUIDs", "uuids", gpuUUIDsOf(visible)) + return visible, nil +} + +// describeGPUs keeps the allocated order and fills each UUID in from what +// nvidia-smi reported about it. A UUID nvidia-smi did not report keeps its +// place undescribed rather than dropping out of the set. +func describeGPUs(uuids []string, visible compat.GPUFacts) compat.GPUFacts { + described := make(map[string]compat.GPUDevice, len(visible.Devices)) + for _, device := range visible.Devices { + described[device.UUID] = device + } + facts := compat.GPUFacts{ + DriverVersion: visible.DriverVersion, + Devices: make([]compat.GPUDevice, 0, len(uuids)), + } + for _, uuid := range uuids { + device, ok := described[uuid] + if !ok { + device = compat.GPUDevice{UUID: uuid} + } + facts.Devices = append(facts.Devices, device) + } + return facts +} + +func gpuUUIDsOf(facts compat.GPUFacts) []string { + var uuids []string + for _, device := range facts.Devices { + if device.UUID != "" { + uuids = append(uuids, device.UUID) + } } - log.Info("nvidia-smi fallback discovered GPU UUIDs", "uuids", gpuUUIDs) - return gpuUUIDs, nil + return uuids } func orderDRAUUIDsByRuntime(allocatedUUIDs, visibleUUIDs []string) ([]string, error) { diff --git a/agent/internal/cuda/cuda_test.go b/agent/internal/cuda/cuda_test.go index f486dc76..4a0911ff 100644 --- a/agent/internal/cuda/cuda_test.go +++ b/agent/internal/cuda/cuda_test.go @@ -7,7 +7,9 @@ import ( "context" "errors" "net" + "os" "path/filepath" + "reflect" "strings" "testing" "time" @@ -21,8 +23,118 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/client-go/kubernetes/fake" podresourcesv1 "k8s.io/kubelet/pkg/apis/podresources/v1" + + "github.com/ai-dynamo/snapshot/api/compat" ) +func TestParseNvidiaSmiGPUFacts(t *testing.T) { + tests := []struct { + name string + output string + want compat.GPUFacts + }{ + { + name: "two GPUs on one driver", + output: "GPU-aaa, NVIDIA A100-SXM4-40GB, 580.65.06\nGPU-bbb, NVIDIA A100-SXM4-40GB, 580.65.06\n", + want: compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{ + {UUID: "GPU-aaa", ProductName: "NVIDIA A100-SXM4-40GB"}, + {UUID: "GPU-bbb", ProductName: "NVIDIA A100-SXM4-40GB"}, + }, + }, + }, + { + // The device map is built from UUIDs, so a row that loses its model + // still has to count as a GPU. + name: "a row without a model still reports its GPU", + output: "GPU-aaa\nGPU-bbb, NVIDIA H100 80GB HBM3, 580.65.06\n", + want: compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{ + {UUID: "GPU-aaa"}, + {UUID: "GPU-bbb", ProductName: "NVIDIA H100 80GB HBM3"}, + }, + }, + }, + { + name: "blank lines are not GPUs", + output: "\n\nGPU-aaa, NVIDIA L4, 580.65.06\n\n", + want: compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{{UUID: "GPU-aaa", ProductName: "NVIDIA L4"}}, + }, + }, + { + name: "rows without UUIDs are not GPUs", + output: ", NVIDIA L4, 580.65.06\nGPU-aaa, NVIDIA L4, 580.65.06\n", + want: compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{{UUID: "GPU-aaa", ProductName: "NVIDIA L4"}}, + }, + }, + { + name: "a node with no GPUs reports nothing", + output: "\n", + want: compat.GPUFacts{}, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := parseNvidiaSmiGPUFacts(tc.output); !reflect.DeepEqual(got, tc.want) { + t.Fatalf("parseNvidiaSmiGPUFacts() = %#v, want %#v", got, tc.want) + } + }) + } +} + +func installFakeNSenter(t *testing.T, body string) { + t.Helper() + + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "nsenter"), []byte("#!/bin/sh\nset -eu\n"+body), 0o755); err != nil { + t.Fatalf("write fake nsenter: %v", err) + } + t.Setenv("PATH", dir+string(os.PathListSeparator)+os.Getenv("PATH")) +} + +func TestDiscoverVisibleGPUFacts(t *testing.T) { + installFakeNSenter(t, ` +test "$#" = 6 +test "$1" = "--mount=/host/proc/42/ns/mnt" +test "$2" = "--pid=/host/proc/42/ns/pid" +test "$3" = "--" +test "$4" = "nvidia-smi" +test "$5" = "--query-gpu=gpu_uuid,name,driver_version" +test "$6" = "--format=csv,noheader" +printf '%s\n' 'GPU-a, NVIDIA L4, 580.65.06' +`) + + got, err := DiscoverVisibleGPUFacts(context.Background(), "/host/proc/", 42) + if err != nil { + t.Fatalf("DiscoverVisibleGPUFacts: %v", err) + } + want := compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{{UUID: "GPU-a", ProductName: "NVIDIA L4"}}, + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("DiscoverVisibleGPUFacts() = %#v, want %#v", got, want) + } +} + +func TestDiscoverVisibleGPUFactsReturnsCommandFailure(t *testing.T) { + installFakeNSenter(t, "exit 17\n") + + _, err := DiscoverVisibleGPUFacts(context.Background(), "/host/proc", 42) + if err == nil { + t.Fatal("DiscoverVisibleGPUFacts succeeded after nsenter failed") + } + if !strings.Contains(err.Error(), "pid 42") { + t.Fatalf("DiscoverVisibleGPUFacts error = %q, want pid", err) + } +} + func TestBuildDeviceMap(t *testing.T) { tests := []struct { name string @@ -266,6 +378,21 @@ func TestDiscoverGPUUUIDsUsesPodResourcesForClassicPod(t *testing.T) { } } +func TestDiscoverGPUUUIDsReturnsDiscoveryError(t *testing.T) { + previousSocketPath := podResourcesSocketPath + podResourcesSocketPath = filepath.Join(t.TempDir(), "missing-kubelet.sock") + t.Cleanup(func() { + podResourcesSocketPath = previousSocketPath + }) + + _, err := DiscoverGPUUUIDs( + context.Background(), nil, "test-pod", "default", "main", "/host/proc", 42, logr.Discard(), + ) + if err == nil { + t.Fatal("DiscoverGPUUUIDs succeeded after PodResources lookup failed") + } +} + func TestDiscoverGPUUUIDsFallsBackToPodResourcesAfterDRAAPILookupError(t *testing.T) { installTestPodResourcesServer(t, &podresourcesv1.ListPodResourcesResponse{ PodResources: []*podresourcesv1.PodResources{ @@ -384,7 +511,7 @@ func TestDiscoverGPUUUIDsOrdersDRAPodByContainerOrdinal(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - got, err := discoverGPUUUIDs( + got, err := discoverGPUFacts( ctx, client, podName, @@ -392,22 +519,225 @@ func TestDiscoverGPUUUIDsOrdersDRAPodByContainerOrdinal(t *testing.T) { "main", "/proc", 123, - func(context.Context, string, int) ([]string, error) { - return []string{uuid0, uuid1}, nil + func(context.Context, string, int) (compat.GPUFacts, error) { + return compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{ + {UUID: uuid0, ProductName: "NVIDIA A100-SXM4-40GB"}, + {UUID: uuid1, ProductName: "NVIDIA A100-SXM4-40GB"}, + }, + }, nil }, logr.Discard(), ) if err != nil { - t.Fatalf("DiscoverGPUUUIDs: %v", err) + t.Fatalf("discoverGPUFacts: %v", err) } - want := []string{uuid0, uuid1} - if len(got) != len(want) { - t.Fatalf("got %v, want %v", got, want) + // Ordered by the runtime, and still described: the DRA path used to reduce + // nvidia-smi's answer to an ordering and throw the rest away. + want := compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{ + {UUID: uuid0, ProductName: "NVIDIA A100-SXM4-40GB"}, + {UUID: uuid1, ProductName: "NVIDIA A100-SXM4-40GB"}, + }, } - for i := range want { - if got[i] != want[i] { - t.Fatalf("got %v, want %v", got, want) - } + if !reflect.DeepEqual(got, want) { + t.Fatalf("discoverGPUFacts() = %#v, want %#v", got, want) + } +} + +// The kubelet path has its GPUs without nvidia-smi, so it never used to run it. +// It runs it now for the model and driver version, and a failure there costs +// those facts and nothing else. +func TestDiscoverGPUFactsDescribesPodResourcesGPUs(t *testing.T) { + installTestPodResourcesServer(t, &podresourcesv1.ListPodResourcesResponse{ + PodResources: []*podresourcesv1.PodResources{ + { + Name: "test-pod", + Namespace: "default", + Containers: []*podresourcesv1.ContainerResources{ + { + Name: "main", + Devices: []*podresourcesv1.ContainerDevices{ + { + ResourceName: nvidiaGPUResource, + DeviceIds: []string{"GPU-a", "GPU-b"}, + }, + }, + }, + }, + }, + }, + }) + + tests := []struct { + name string + visible func(context.Context, string, int) (compat.GPUFacts, error) + want compat.GPUFacts + }{ + { + name: "described in the kubelet's order", + visible: func(context.Context, string, int) (compat.GPUFacts, error) { + return compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{ + {UUID: "GPU-b", ProductName: "NVIDIA L4"}, + {UUID: "GPU-a", ProductName: "NVIDIA L4"}, + }, + }, nil + }, + want: compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{ + {UUID: "GPU-a", ProductName: "NVIDIA L4"}, + {UUID: "GPU-b", ProductName: "NVIDIA L4"}, + }, + }, + }, + { + name: "undescribed when nvidia-smi cannot be reached", + visible: func(context.Context, string, int) (compat.GPUFacts, error) { + return compat.GPUFacts{}, errors.New("nsenter unavailable") + }, + want: compat.GPUFacts{ + Devices: []compat.GPUDevice{{UUID: "GPU-a"}, {UUID: "GPU-b"}}, + }, + }, + { + name: "undescribed when nvidia-smi reports other GPUs", + visible: func(context.Context, string, int) (compat.GPUFacts, error) { + return compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{{UUID: "GPU-z", ProductName: "NVIDIA L4"}}, + }, nil + }, + want: compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{{UUID: "GPU-a"}, {UUID: "GPU-b"}}, + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + got, err := discoverGPUFacts( + ctx, nil, "test-pod", "default", "main", "/proc", 123, tc.visible, logr.Discard(), + ) + if err != nil { + t.Fatalf("discoverGPUFacts: %v", err) + } + if !reflect.DeepEqual(got, tc.want) { + t.Fatalf("discoverGPUFacts() = %#v, want %#v", got, tc.want) + } + }) + } +} + +func TestDiscoverGPUFactsUsesVisibleGPUDescriptions(t *testing.T) { + installTestPodResourcesServer(t, &podresourcesv1.ListPodResourcesResponse{ + PodResources: []*podresourcesv1.PodResources{ + { + Name: "test-pod", + Namespace: "default", + Containers: []*podresourcesv1.ContainerResources{ + { + Name: "main", + Devices: []*podresourcesv1.ContainerDevices{ + { + ResourceName: nvidiaGPUResource, + DeviceIds: []string{"GPU-a"}, + }, + }, + }, + }, + }, + }, + }) + installFakeNSenter(t, "printf '%s\\n' 'GPU-a, NVIDIA L4, 580.65.06'\n") + + got, err := DiscoverGPUFacts( + context.Background(), nil, "test-pod", "default", "main", "/host/proc", 42, logr.Discard(), + ) + if err != nil { + t.Fatalf("DiscoverGPUFacts: %v", err) + } + want := compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{{UUID: "GPU-a", ProductName: "NVIDIA L4"}}, + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("DiscoverGPUFacts() = %#v, want %#v", got, want) + } +} + +func TestDiscoverGPUFactsFallsBackToVisibleGPUs(t *testing.T) { + installTestPodResourcesServer(t, &podresourcesv1.ListPodResourcesResponse{}) + want := compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{{UUID: "GPU-a", ProductName: "NVIDIA L4"}}, + } + + got, err := discoverGPUFacts( + context.Background(), + nil, + "test-pod", + "default", + "main", + "/host/proc", + 42, + func(context.Context, string, int) (compat.GPUFacts, error) { + return want, nil + }, + logr.Discard(), + ) + if err != nil { + t.Fatalf("discoverGPUFacts: %v", err) + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("discoverGPUFacts() = %#v, want %#v", got, want) + } +} + +func TestDescribeGPUs(t *testing.T) { + visible := compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{ + {UUID: "GPU-b", ProductName: "NVIDIA H100"}, + {UUID: "GPU-extra", ProductName: "NVIDIA L4"}, + {UUID: "GPU-a", ProductName: "NVIDIA A100"}, + }, + } + + got := describeGPUs([]string{"GPU-a", "GPU-b", "GPU-missing"}, visible) + want := compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{ + {UUID: "GPU-a", ProductName: "NVIDIA A100"}, + {UUID: "GPU-b", ProductName: "NVIDIA H100"}, + {UUID: "GPU-missing"}, + }, + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("describeGPUs() = %#v, want %#v", got, want) + } +} + +func TestGPUUUIDsOf(t *testing.T) { + facts := compat.GPUFacts{ + Devices: []compat.GPUDevice{ + {UUID: "GPU-a"}, + {ProductName: "NVIDIA L4"}, + {UUID: "GPU-b"}, + }, + } + + got := gpuUUIDsOf(facts) + want := []string{"GPU-a", "GPU-b"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("gpuUUIDsOf() = %v, want %v", got, want) } } diff --git a/agent/internal/executor/checkpoint.go b/agent/internal/executor/checkpoint.go index 1f965012..8d554148 100644 --- a/agent/internal/executor/checkpoint.go +++ b/agent/internal/executor/checkpoint.go @@ -22,6 +22,7 @@ import ( "github.com/ai-dynamo/snapshot/agent/internal/nsmount" snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" "github.com/ai-dynamo/snapshot/agent/internal/types" + "github.com/ai-dynamo/snapshot/api/compat" ) // CheckpointRequest holds the content-owned inputs for a checkpoint operation. @@ -35,6 +36,11 @@ type CheckpointRequest struct { PodNamespace string PodIP string Clientset kubernetes.Interface + + // Pod carries the image reference and limits the target container runs with, read from + // the live pod by the caller rather than here: the capture path has no API + // client for the pod, and the reconciler already holds it. + Pod compat.Facts } type checkpointPhaseTimings struct { @@ -78,7 +84,7 @@ func Checkpoint(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger } cudaJobFile := "" if len(state.CUDAHostPIDs) > 0 { - cudaJobFile, err = cuda.StageJobFile(state.RootFS, tmpDir, len(state.GPUUUIDs)) + cudaJobFile, err = cuda.StageJobFile(state.RootFS, tmpDir, len(state.GPUs.Devices)) if err != nil { return err } @@ -138,6 +144,10 @@ func inspectContainer(ctx context.Context, rt snapshotruntime.Runtime, log logr. if err != nil { return nil, 0, fmt.Errorf("failed to resolve container: %w", err) } + imageID, err := rt.ResolveContainerImageID(ctx, containerID) + if err != nil { + return nil, 0, fmt.Errorf("failed to resolve container image ID: %w", err) + } var hostCgroupPath string if cgPath, err := snapshotruntime.ResolveCgroupRootFromHostPID(pid); err == nil && cgPath != "" { @@ -193,11 +203,11 @@ func inspectContainer(ctx context.Context, rt snapshotruntime.Runtime, log logr. if len(cudaHostPIDs) > 0 { log.V(1).Info("Resolved checkpoint CUDA PID mapping", "host_pids", cudaHostPIDs, "namespace_pids", cudaNamespacePIDs) } - var gpuUUIDs []string + var gpus compat.GPUFacts var gpuDeviceMapDuration time.Duration if len(cudaHostPIDs) > 0 { gpuStart := time.Now() - gpuUUIDs, err = cuda.DiscoverGPUUUIDs( + gpus, err = cuda.DiscoverGPUFacts( ctx, req.Clientset, req.PodName, @@ -215,6 +225,7 @@ func inspectContainer(ctx context.Context, rt snapshotruntime.Runtime, log logr. return &types.CheckpointContainerSnapshot{ PID: pid, + ImageID: imageID, RootFS: rootFS, UpperDir: upperDir, OCISpec: ociSpec, @@ -224,7 +235,7 @@ func inspectContainer(ctx context.Context, rt snapshotruntime.Runtime, log logr. HostCgroupPath: hostCgroupPath, CUDAHostPIDs: cudaHostPIDs, CUDANSPIDs: cudaNamespacePIDs, - GPUUUIDs: gpuUUIDs, + GPUs: gpus, }, gpuDeviceMapDuration, nil } @@ -239,16 +250,20 @@ func configureCheckpoint( if err != nil { return nil, nil, err } + podFacts := req.Pod + podFacts.ImageID = state.ImageID m := types.NewCheckpointManifest( req.ContentUID, req.ContainerName, types.NewCRIUDumpManifest(criuOpts, cfg.CRIU), - types.NewSourcePodManifest(req.ContainerID, state.PID, req.NodeName, req.PodName, req.PodNamespace, req.PodIP, state.StdioFDs), + types.NewSourcePodManifest(req.ContainerID, state.PID, req.NodeName, req.PodName, req.PodNamespace, req.PodIP, state.StdioFDs). + WithPodFacts(podFacts), types.NewOverlayManifest(cfg.Overlay, state.UpperDir, state.OCISpec), + types.NewHostManifest(cfg.HostKernelVersion), ) if len(state.CUDANSPIDs) > 0 { - m.CUDA = types.NewCUDAManifest(state.CUDANSPIDs, state.GPUUUIDs) + m.CUDA = types.NewCUDAManifest(state.CUDANSPIDs, state.GPUs) } if err := types.WriteManifest(checkpointDir, m); err != nil { diff --git a/agent/internal/executor/checkpoint_test.go b/agent/internal/executor/checkpoint_test.go index 6674987a..d51bd87e 100644 --- a/agent/internal/executor/checkpoint_test.go +++ b/agent/internal/executor/checkpoint_test.go @@ -16,6 +16,7 @@ import ( "github.com/ai-dynamo/snapshot/agent/internal/nsmount" "github.com/ai-dynamo/snapshot/agent/internal/types" + "github.com/ai-dynamo/snapshot/api/compat" ) type checkpointPathRuntime struct{} @@ -32,8 +33,24 @@ func (checkpointPathRuntime) ResolveContainerByPod(context.Context, string, stri return 0, nil, errors.New("not implemented") } +func (checkpointPathRuntime) ResolveContainerImageID(context.Context, string) (string, error) { + return "", errors.New("not implemented") +} + func (checkpointPathRuntime) Close() error { return nil } +type checkpointImageRuntime struct { + checkpointPathRuntime +} + +func (checkpointImageRuntime) ResolveContainer(context.Context, string) (int, *specs.Spec, error) { + return 1, &specs.Spec{}, nil +} + +func (checkpointImageRuntime) ResolveContainerImageID(context.Context, string) (string, error) { + return "", errors.New("runtime image unavailable") +} + func TestCheckpointPreparesContentArtifactParents(t *testing.T) { cfg := &types.AgentConfig{Storage: types.StorageSpec{BasePath: t.TempDir()}} finalDir, err := nsmount.ResolveArtifactPath(cfg.Storage.BasePath, "content-uid", "main") @@ -47,3 +64,43 @@ func TestCheckpointPreparesContentArtifactParents(t *testing.T) { assert.DirExists(t, filepath.Dir(finalDir)) assert.DirExists(t, filepath.Join(cfg.Storage.BasePath, "artifacts", "content-uid", ".tmp")) } + +func TestInspectContainerRequiresRuntimeImageID(t *testing.T) { + _, _, err := inspectContainer( + context.Background(), + checkpointImageRuntime{}, + logr.Discard(), + CheckpointRequest{ContainerID: "container-id"}, + ) + require.ErrorContains(t, err, "failed to resolve container image ID: runtime image unavailable") +} + +func TestConfigureCheckpointRecordsRuntimeImageID(t *testing.T) { + checkpointDir := t.TempDir() + _, _, err := configureCheckpoint( + logr.Discard(), + &types.CheckpointContainerSnapshot{ + PID: 42, + ImageID: "sha256:runtime-content", + RootFS: "/", + NetNSInode: 7, + }, + CheckpointRequest{ + ContentUID: "content-uid", + ContainerID: "container-id", + ContainerName: "main", + Pod: compat.Facts{ + Image: "registry.example/workload:latest", + ImageID: "sha256:kubelet-alias", + }, + }, + &types.AgentConfig{}, + checkpointDir, + ) + require.NoError(t, err) + + manifest, err := types.ReadManifest(checkpointDir) + require.NoError(t, err) + assert.Equal(t, "registry.example/workload:latest", manifest.K8s.Image) + assert.Equal(t, "sha256:runtime-content", manifest.K8s.ImageID) +} diff --git a/agent/internal/executor/compat.go b/agent/internal/executor/compat.go new file mode 100644 index 00000000..08e93356 --- /dev/null +++ b/agent/internal/executor/compat.go @@ -0,0 +1,42 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package executor + +import ( + "github.com/go-logr/logr" + + "github.com/ai-dynamo/snapshot/agent/internal/types" + "github.com/ai-dynamo/snapshot/api/compat" +) + +// inspectCompatibility runs the inspect gate for one restore, the counterpart of +// the controller's preflightCompatibility. A nil error means the restore may go +// ahead. It gathers the target facts this gate can read, which the earlier gate +// cannot: the runtime image ID, GPUs, and mounts under its rootfs. +func inspectCompatibility( + log logr.Logger, + manifest *types.CheckpointManifest, + targetGPUs compat.GPUFacts, + targetRoot string, + targetImageID string, + skipCompatCheck bool, +) error { + if skipCompatCheck { + log.Info("Restore compatibility check skipped by request", "gate", string(compat.GateInspect)) + return nil + } + + sourceFacts := manifest.CompatFacts() + targetFacts := compat.Facts{ + ImageID: targetImageID, + DriverVersion: targetGPUs.DriverVersion, + GPUDevices: targetGPUs.Devices, + ExistingMounts: existingMounts(targetRoot, sourceFacts.ExternalizedMounts), + } + mismatches := compat.Compare(compat.GateInspect, sourceFacts, targetFacts) + if len(mismatches) == 0 { + return nil + } + return compat.NewIncompatibleError(compat.GateInspect, mismatches) +} diff --git a/agent/internal/executor/restore.go b/agent/internal/executor/restore.go index 8dfde222..192d3a12 100644 --- a/agent/internal/executor/restore.go +++ b/agent/internal/executor/restore.go @@ -24,6 +24,7 @@ import ( "github.com/ai-dynamo/snapshot/agent/internal/nsmount" snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" "github.com/ai-dynamo/snapshot/agent/internal/types" + "github.com/ai-dynamo/snapshot/api/compat" ) // RestoreMounter installs the fixed binary bundle and one validated checkpoint @@ -74,6 +75,11 @@ type RestoreRequest struct { ArtifactContainerName string DestinationContainerName string Clientset kubernetes.Interface + + // SkipCompatCheck carries the decision the caller already made, so the + // second gate cannot reach a different answer than the first: one restore + // is either checked or it is not. + SkipCompatCheck bool } // Restore performs external restore for the given request. @@ -248,20 +254,37 @@ func inspectRestore( } log.V(1).Info("Resolved placeholder container", "pid", placeholderPID) + targetImageID := "" + if manifest.K8s.ImageID != "" { + if req.ContainerID == "" { + return nil, 0, fmt.Errorf("container ID is required to compare the runtime image ID") + } + targetImageID, err = rt.ResolveContainerImageID(ctx, req.ContainerID) + if err != nil { + return nil, 0, fmt.Errorf("failed to resolve placeholder image ID: %w", err) + } + } + cgroupRoot, err := snapshotruntime.ResolveCgroupRootFromHostPID(placeholderPID) if err != nil { log.Error(err, "Failed to resolve placeholder cgroup root; proceeding without explicit cgroup remap") cgroupRoot = "" } - cudaDeviceMap := "" - var gpuDeviceMapDuration time.Duration + targetRoot := fmt.Sprintf("%s/%d/root", snapshotruntime.HostProcPath, placeholderPID) + + var ( + targetGPUs compat.GPUFacts + targetGPUUUIDs []string + discoverDuration time.Duration + deviceMapDuration time.Duration + ) if !manifest.CUDA.IsEmpty() { if len(manifest.CUDA.SourceGPUUUIDs) == 0 { return nil, 0, fmt.Errorf("missing source GPU UUIDs in checkpoint manifest") } - gpuStart := time.Now() - targetGPUUUIDs, err := cuda.DiscoverGPUUUIDs( + discoverStart := time.Now() + targetGPUs, err = cuda.DiscoverGPUFacts( ctx, req.Clientset, req.PodName, @@ -271,14 +294,33 @@ func inspectRestore( placeholderPID, log, ) + discoverDuration = time.Since(discoverStart) if err != nil { return nil, 0, fmt.Errorf("failed to get target GPU UUIDs: %w", err) } - if len(targetGPUUUIDs) == 0 { - return nil, 0, fmt.Errorf("missing target GPU UUIDs for %s/%s container %s", req.PodNamespace, req.PodName, req.DestinationContainerName) + for _, device := range targetGPUs.Devices { + targetGPUUUIDs = append(targetGPUUUIDs, device.UUID) } + } + + // Gate B, once the placeholder is resolved and this node's own facts are + // readable. It runs ahead of BuildDeviceMap, whose positional pairing turns + // a GPU difference into a device-map error that names neither GPU. + if err := inspectCompatibility(log, manifest, targetGPUs, targetRoot, targetImageID, req.SkipCompatCheck); err != nil { + return nil, 0, err + } + + // Behind the gate, which names a target with no GPUs as a count refusal. + // This is what is left when the gate is skipped. + if !manifest.CUDA.IsEmpty() && len(targetGPUUUIDs) == 0 { + return nil, 0, fmt.Errorf("missing target GPU UUIDs for %s/%s container %s", req.PodNamespace, req.PodName, req.DestinationContainerName) + } + + cudaDeviceMap := "" + if len(targetGPUUUIDs) > 0 { + deviceMapStart := time.Now() cudaDeviceMap, err = cuda.BuildDeviceMap(manifest.CUDA.SourceGPUUUIDs, targetGPUUUIDs, log) - gpuDeviceMapDuration = time.Since(gpuStart) + deviceMapDuration = time.Since(deviceMapStart) if err != nil { return nil, 0, fmt.Errorf("failed to build CUDA device map: %w", err) } @@ -291,10 +333,27 @@ func inspectRestore( return &types.RestoreContainerSnapshot{ PlaceholderPID: placeholderPID, - TargetRoot: fmt.Sprintf("%s/%d/root", snapshotruntime.HostProcPath, placeholderPID), + TargetRoot: targetRoot, CgroupRoot: cgroupRoot, CUDADeviceMap: cudaDeviceMap, - }, gpuDeviceMapDuration, nil + }, discoverDuration + deviceMapDuration, nil +} + +// existingMounts reports which of the recorded destinations resolve inside the +// placeholder's rootfs. Only what the checkpoint recorded is looked up, so a gate +// on this path costs one stat per volume the checkpoint actually used. +// +// Only a path that is definitely absent is left out. Any other stat failure is +// this agent failing to look rather than the pod missing a volume, and reporting +// it as missing would refuse a restore that would have worked. +func existingMounts(targetRoot string, destinations []string) []string { + existing := make([]string, 0, len(destinations)) + for _, destination := range destinations { + if _, err := os.Stat(filepath.Join(targetRoot, destination)); !os.IsNotExist(err) { + existing = append(existing, destination) + } + } + return existing } // execNSRestore launches the nsrestore binary inside the placeholder container's diff --git a/agent/internal/executor/restore_test.go b/agent/internal/executor/restore_test.go index 238c067d..d7868e36 100644 --- a/agent/internal/executor/restore_test.go +++ b/agent/internal/executor/restore_test.go @@ -8,6 +8,8 @@ import ( "errors" "fmt" "os" + "path/filepath" + "reflect" "strings" "testing" "time" @@ -17,6 +19,7 @@ import ( "github.com/ai-dynamo/snapshot/agent/internal/nsmount" "github.com/ai-dynamo/snapshot/agent/internal/types" + "github.com/ai-dynamo/snapshot/api/compat" ) // testMountPoint satisfies nsmount.MountPoint for executor unit tests. @@ -31,6 +34,8 @@ type restoreFakeRuntime struct { resolvedID string resolvedByPodContainer string resolveByPodHit bool + imageID string + imageIDError error } func (r *restoreFakeRuntime) ResolveContainer(ctx context.Context, id string) (int, *specs.Spec, error) { @@ -48,6 +53,10 @@ func (r *restoreFakeRuntime) ResolveContainerByPod(ctx context.Context, pod, ns, return 0, nil, errors.New("pod lookup should not be used") } +func (r *restoreFakeRuntime) ResolveContainerImageID(_ context.Context, _ string) (string, error) { + return r.imageID, r.imageIDError +} + func (r *restoreFakeRuntime) Close() error { return nil } func TestInspectRestoreUsesContainerIDWhenProvided(t *testing.T) { @@ -57,6 +66,7 @@ func TestInspectRestoreUsesContainerIDWhenProvided(t *testing.T) { types.CRIUDumpManifest{}, types.NewSourcePodManifest("source-id", 456, "node-1", "source-pod", "default", "10.0.0.11", nil), types.OverlayManifest{}, + types.HostManifest{}, ) rt := &restoreFakeRuntime{} _, _, err := inspectRestore( @@ -84,6 +94,95 @@ func TestInspectRestoreUsesContainerIDWhenProvided(t *testing.T) { } } +func TestInspectRestoreComparesRuntimeImageID(t *testing.T) { + const ( + captured = "sha256:1111111111111111111111111111111111111111111111111111111111111111" + rebuilt = "sha256:2222222222222222222222222222222222222222222222222222222222222222" + ) + tests := []struct { + name string + sourceID string + targetID string + targetError error + want []compat.Mismatch + wantError string + }{ + { + name: "same runtime content", + sourceID: captured, + targetID: captured, + }, + { + name: "different runtime content", + sourceID: captured, + targetID: rebuilt, + want: []compat.Mismatch{{Check: compat.CheckImageDigest, Source: captured, Target: rebuilt}}, + }, + { + name: "artifact without a runtime image ID", + targetID: captured, + }, + { + name: "runtime image ID unavailable", + sourceID: captured, + targetError: errors.New("runtime unavailable"), + wantError: "failed to resolve placeholder image ID: runtime unavailable", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + manifest := types.NewCheckpointManifest( + "content-uid-123", + "main", + types.CRIUDumpManifest{}, + types.NewSourcePodManifest("source-id", 456, "node-1", "source-pod", "default", "10.0.0.11", nil), + types.OverlayManifest{}, + types.HostManifest{}, + ) + manifest.K8s.ImageID = tc.sourceID + rt := &restoreFakeRuntime{imageID: tc.targetID, imageIDError: tc.targetError} + + _, _, err := inspectRestore( + context.Background(), + rt, + testr.New(t), + RestoreRequest{ + ContentUID: "content-uid-123", + ContainerID: "placeholder-id", + PodName: "restore-pod", + PodNamespace: "default", + ArtifactContainerName: "main", + DestinationContainerName: "main", + }, + manifest, + ) + if tc.wantError != "" { + if err == nil || !strings.Contains(err.Error(), tc.wantError) { + t.Fatalf("error = %v, want containing %q", err, tc.wantError) + } + return + } + if len(tc.want) == 0 { + if err != nil { + t.Fatalf("inspectRestore: %v", err) + } + return + } + var incompatible *compat.IncompatibleError + if !errors.As(err, &incompatible) { + t.Fatalf("error = %v, want *compat.IncompatibleError", err) + } + if incompatible.Gate != compat.GateInspect { + t.Fatalf("gate = %q, want %q", incompatible.Gate, compat.GateInspect) + } + if !reflect.DeepEqual(incompatible.Mismatches, tc.want) { + t.Fatalf("mismatches = %+v, want %+v", incompatible.Mismatches, tc.want) + } + }) + } +} + func TestInspectRestoreUsesDestinationNameForPodLookup(t *testing.T) { manifest := types.NewCheckpointManifest( "content-uid-123", @@ -91,6 +190,7 @@ func TestInspectRestoreUsesDestinationNameForPodLookup(t *testing.T) { types.CRIUDumpManifest{}, types.NewSourcePodManifest("source-id", 456, "node-1", "source-pod", "default", "10.0.0.11", nil), types.OverlayManifest{}, + types.NewHostManifest("6.17.0"), ) rt := &restoreFakeRuntime{} _, _, err := inspectRestore( @@ -133,6 +233,7 @@ func TestValidateRestoreManifest(t *testing.T) { types.CRIUDumpManifest{}, types.NewSourcePodManifest("source-id", 456, "node-1", "source-pod", "team-a", "10.0.0.11", nil), types.OverlayManifest{}, + types.HostManifest{}, ) for _, tc := range []struct { @@ -172,8 +273,11 @@ func TestRestoreInNamespaceRejectsMultiGPUCheckpointWithoutLaunchJobState(t *tes types.CRIUDumpManifest{}, types.NewSourcePodManifest("source-id", 456, "node-1", "source-pod", "default", "10.0.0.11", nil), types.OverlayManifest{}, + types.HostManifest{}, ) - manifest.CUDA = types.NewCUDAManifest([]int{42, 43}, []string{"GPU-aaa", "GPU-bbb"}) + manifest.CUDA = types.NewCUDAManifest([]int{42, 43}, compat.GPUFacts{ + Devices: []compat.GPUDevice{{UUID: "GPU-aaa"}, {UUID: "GPU-bbb"}}, + }) if err := types.WriteManifest(checkpointDir, manifest); err != nil { t.Fatalf("WriteManifest: %v", err) } @@ -193,3 +297,23 @@ func TestRemainingDuration(t *testing.T) { t.Fatal("remainingDuration should not go negative") } } + +func TestExistingMounts(t *testing.T) { + targetRoot := t.TempDir() + if err := os.MkdirAll(filepath.Join(targetRoot, "model-cache"), 0o755); err != nil { + t.Fatalf("MkdirAll: %v", err) + } + if err := os.WriteFile(filepath.Join(targetRoot, "etc-hostname"), nil, 0o600); err != nil { + t.Fatalf("WriteFile: %v", err) + } + + got := existingMounts(targetRoot, []string{"/model-cache", "/data", "/etc-hostname"}) + want := []string{"/model-cache", "/etc-hostname"} + if !reflect.DeepEqual(got, want) { + t.Errorf("existingMounts = %#v, want %#v", got, want) + } + + if got := existingMounts(targetRoot, nil); len(got) != 0 { + t.Errorf("existingMounts of nothing = %#v, want empty", got) + } +} diff --git a/agent/internal/runtime/image.go b/agent/internal/runtime/image.go new file mode 100644 index 00000000..4ad67c84 --- /dev/null +++ b/agent/internal/runtime/image.go @@ -0,0 +1,33 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package runtime + +import ( + "context" + "fmt" + "time" + + runtimeapi "k8s.io/cri-api/pkg/apis/runtime/v1" +) + +const containerImageIDTimeout = 10 * time.Second + +type containerStatusService interface { + ContainerStatus(context.Context, string, bool) (*runtimeapi.ContainerStatusResponse, error) +} + +func resolveContainerImageID(ctx context.Context, svc containerStatusService, containerID string) (string, error) { + ctx, cancel := context.WithTimeout(ctx, containerImageIDTimeout) + defer cancel() + + response, err := svc.ContainerStatus(ctx, containerID, false) + if err != nil { + return "", fmt.Errorf("failed to get container status for %s: %w", containerID, err) + } + imageID := response.GetStatus().GetImageId() + if imageID == "" { + return "", fmt.Errorf("container status for %s has no image ID", containerID) + } + return imageID, nil +} diff --git a/agent/internal/runtime/image_test.go b/agent/internal/runtime/image_test.go new file mode 100644 index 00000000..c31a230e --- /dev/null +++ b/agent/internal/runtime/image_test.go @@ -0,0 +1,72 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package runtime + +import ( + "context" + "errors" + "strings" + "testing" + + runtimeapi "k8s.io/cri-api/pkg/apis/runtime/v1" +) + +type fakeContainerStatusService struct { + response *runtimeapi.ContainerStatusResponse + err error + verbose bool +} + +func (f *fakeContainerStatusService) ContainerStatus(_ context.Context, _ string, verbose bool) (*runtimeapi.ContainerStatusResponse, error) { + f.verbose = verbose + return f.response, f.err +} + +func TestResolveContainerImageID(t *testing.T) { + tests := []struct { + name string + service *fakeContainerStatusService + want string + wantError string + }{ + { + name: "returns the runtime image ID", + service: &fakeContainerStatusService{response: &runtimeapi.ContainerStatusResponse{ + Status: &runtimeapi.ContainerStatus{ImageId: "sha256:config"}, + }}, + want: "sha256:config", + }, + { + name: "rejects a missing status", + service: &fakeContainerStatusService{response: &runtimeapi.ContainerStatusResponse{}}, + wantError: "has no image ID", + }, + { + name: "reports the runtime error", + service: &fakeContainerStatusService{err: errors.New("unavailable")}, + wantError: "unavailable", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got, err := resolveContainerImageID(context.Background(), tc.service, "container-id") + if tc.wantError != "" { + if err == nil || !strings.Contains(err.Error(), tc.wantError) { + t.Fatalf("error = %v, want containing %q", err, tc.wantError) + } + return + } + if err != nil { + t.Fatalf("resolveContainerImageID: %v", err) + } + if got != tc.want { + t.Fatalf("image ID = %q, want %q", got, tc.want) + } + if tc.service.verbose { + t.Fatal("ContainerStatus requested verbose output") + } + }) + } +} diff --git a/agent/internal/runtime/oci_containerd.go b/agent/internal/runtime/oci_containerd.go index 02350286..2459ee4e 100644 --- a/agent/internal/runtime/oci_containerd.go +++ b/agent/internal/runtime/oci_containerd.go @@ -5,11 +5,15 @@ package runtime import ( "context" + "errors" "fmt" + "time" containerd "github.com/containerd/containerd/v2/client" "github.com/containerd/containerd/v2/pkg/namespaces" specs "github.com/opencontainers/runtime-spec/specs-go" + internalapi "k8s.io/cri-api/pkg/apis" + remote "k8s.io/cri-client/pkg" ) // k8sNamespace is containerd's conventional namespace for kubelet-managed @@ -18,6 +22,7 @@ const k8sNamespace = "k8s.io" type ContainerdRuntime struct { client *containerd.Client + cri internalapi.RuntimeService } func NewContainerdRuntime(socket string) (*ContainerdRuntime, error) { @@ -25,11 +30,20 @@ func NewContainerdRuntime(socket string) (*ContainerdRuntime, error) { if err != nil { return nil, fmt.Errorf("failed to dial containerd at %s: %w", socket, err) } - return &ContainerdRuntime{client: client}, nil + cri, err := remote.NewRemoteRuntimeService(context.Background(), socket, 2*time.Second, nil, false) + if err != nil { + _ = client.Close() + return nil, fmt.Errorf("failed to dial containerd CRI at %s: %w", socket, err) + } + return &ContainerdRuntime{client: client, cri: cri}, nil } func (r *ContainerdRuntime) Close() error { - return r.client.Close() + return errors.Join(r.cri.Close(context.Background()), r.client.Close()) +} + +func (r *ContainerdRuntime) ResolveContainerImageID(ctx context.Context, containerID string) (string, error) { + return resolveContainerImageID(ctx, r.cri, containerID) } func (r *ContainerdRuntime) ResolveContainer(ctx context.Context, containerID string) (int, *specs.Spec, error) { diff --git a/agent/internal/runtime/oci_crio.go b/agent/internal/runtime/oci_crio.go index 797f0de8..80afb14f 100644 --- a/agent/internal/runtime/oci_crio.go +++ b/agent/internal/runtime/oci_crio.go @@ -34,9 +34,13 @@ func NewCRIORuntime(socket string) (*CRIORuntime, error) { return &CRIORuntime{svc: svc}, nil } -// Close is a no-op: k8s.io/cri-client's RuntimeService interface doesn't -// expose one. The gRPC connection is released at process exit. -func (r *CRIORuntime) Close() error { return nil } +func (r *CRIORuntime) Close() error { + return r.svc.Close(context.Background()) +} + +func (r *CRIORuntime) ResolveContainerImageID(ctx context.Context, containerID string) (string, error) { + return resolveContainerImageID(ctx, r.svc, containerID) +} func (r *CRIORuntime) ResolveContainer(ctx context.Context, id string) (int, *specs.Spec, error) { ctx, cancel := context.WithTimeout(ctx, crioCallTimeout) diff --git a/agent/internal/runtime/process.go b/agent/internal/runtime/process.go index ad7fb7ac..1c2da1ac 100644 --- a/agent/internal/runtime/process.go +++ b/agent/internal/runtime/process.go @@ -19,6 +19,21 @@ import ( // HostProcPath is the mount point for the host's /proc in DaemonSet pods. const HostProcPath = "/host/proc" +// ReadKernelVersion reads the host's /proc/sys/kernel/osrelease entry documented at +// https://man7.org/linux/man-pages/man5/proc_sys_kernel.5.html. +func ReadKernelVersion(hostProcPath string) (string, error) { + releasePath := filepath.Join(hostProcPath, "sys", "kernel", "osrelease") + content, err := os.ReadFile(releasePath) + if err != nil { + return "", fmt.Errorf("failed to read kernel version from %s: %w", releasePath, err) + } + release := strings.TrimSpace(string(content)) + if release == "" { + return "", fmt.Errorf("kernel version at %s is empty", releasePath) + } + return release, nil +} + // ProcessDetails captures the parent link plus the observed, outermost, and innermost // PID views for one proc entry. ObservedPID is relative to the proc root being read. type ProcessDetails struct { diff --git a/agent/internal/runtime/process_test.go b/agent/internal/runtime/process_test.go index 259478b3..1880a8f6 100644 --- a/agent/internal/runtime/process_test.go +++ b/agent/internal/runtime/process_test.go @@ -230,3 +230,42 @@ func TestResolveManifestPIDsToObservedPIDsFailsWhenNamespaceDepthIsNotTwo(t *tes t.Fatal("ResolveManifestPIDsToObservedPIDs(...) unexpectedly succeeded") } } + +func writeKernelRelease(t *testing.T, procPath, content string) { + t.Helper() + dir := filepath.Join(procPath, "sys", "kernel") + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatalf("MkdirAll: %v", err) + } + if err := os.WriteFile(filepath.Join(dir, "osrelease"), []byte(content), 0o600); err != nil { + t.Fatalf("WriteFile: %v", err) + } +} + +func TestReadKernelVersion(t *testing.T) { + procPath := t.TempDir() + writeKernelRelease(t, procPath, "5.15.0-119-generic\n") + + got, err := ReadKernelVersion(procPath) + if err != nil { + t.Fatalf("ReadKernelVersion: %v", err) + } + if got != "5.15.0-119-generic" { + t.Fatalf("ReadKernelVersion() = %q, want 5.15.0-119-generic", got) + } +} + +// An unreadable or blank release is reported rather than passed on: a fact +// recorded as the empty string would be indistinguishable from one this agent +// version never recorded at all. +func TestReadKernelVersionRejectsWhatItCannotRead(t *testing.T) { + if _, err := ReadKernelVersion(t.TempDir()); err == nil { + t.Fatal("expected an error when the host proc mount has no osrelease") + } + + procPath := t.TempDir() + writeKernelRelease(t, procPath, "\n") + if _, err := ReadKernelVersion(procPath); err == nil { + t.Fatal("expected an error when osrelease is blank") + } +} diff --git a/agent/internal/runtime/runtime.go b/agent/internal/runtime/runtime.go index 41545d7a..2b08bc47 100644 --- a/agent/internal/runtime/runtime.go +++ b/agent/internal/runtime/runtime.go @@ -29,6 +29,7 @@ type Runtime interface { ResolveContainer(ctx context.Context, id string) (int, *specs.Spec, error) ResolveContainerIDByPod(ctx context.Context, pod, ns, ctr string) (string, error) ResolveContainerByPod(ctx context.Context, pod, ns, ctr string) (int, *specs.Spec, error) + ResolveContainerImageID(ctx context.Context, id string) (string, error) Close() error } diff --git a/agent/internal/types/config.go b/agent/internal/types/config.go index 2d3a29ea..93fbd16d 100644 --- a/agent/internal/types/config.go +++ b/agent/internal/types/config.go @@ -15,14 +15,14 @@ import ( // helper independently enforces the same path. const CheckpointBasePath = "/checkpoints" -// AgentConfig holds the full agent configuration: static checkpoint settings -// from the ConfigMap YAML, plus runtime fields from environment variables. +// AgentConfig holds static checkpoint settings plus runtime fields populated at startup. type AgentConfig struct { - NodeName string `yaml:"-"` - Storage StorageSpec `yaml:"storage"` - Overlay OverlaySettings `yaml:"overlay"` - Restore RestoreSpec `yaml:"restore"` - CRIU CRIUSettings `yaml:"criu"` + NodeName string `yaml:"-"` + HostKernelVersion string `yaml:"-"` + Storage StorageSpec `yaml:"storage"` + Overlay OverlaySettings `yaml:"overlay"` + Restore RestoreSpec `yaml:"restore"` + CRIU CRIUSettings `yaml:"criu"` } func (c *AgentConfig) LoadEnvOverrides() { @@ -70,6 +70,11 @@ type StorageSpec struct { // RestoreSpec holds settings for the CRIU restore process. type RestoreSpec struct { RestoreTimeoutSeconds int `yaml:"restoreTimeoutSeconds"` + + // SkipCompatCheck turns the restore compatibility gate off for every + // restore this agent handles. It is the per-node escape hatch, for a + // cluster admin who would otherwise be stuck annotating pods one by one. + SkipCompatCheck bool `yaml:"skipCompatCheck"` } func (c *RestoreSpec) RestoreTimeout() time.Duration { diff --git a/agent/internal/types/config_test.go b/agent/internal/types/config_test.go index 27e1af81..d45a1128 100644 --- a/agent/internal/types/config_test.go +++ b/agent/internal/types/config_test.go @@ -3,7 +3,11 @@ package types -import "testing" +import ( + "testing" + + "gopkg.in/yaml.v3" +) func validAgentConfig() *AgentConfig { return &AgentConfig{ @@ -17,6 +21,25 @@ func validAgentConfig() *AgentConfig { } } +// The key an admin flips in the ConfigMap has to be the key the agent reads, +// and an agent whose ConfigMap predates it has to keep the gate on. +func TestRestoreSpecParsesSkipCompatCheck(t *testing.T) { + cases := map[string]bool{ + "restore:\n skipCompatCheck: true\n": true, + "restore:\n skipCompatCheck: false\n": false, + "restore:\n restoreTimeoutSeconds: 60\n": false, + } + for document, want := range cases { + cfg := &AgentConfig{} + if err := yaml.Unmarshal([]byte(document), cfg); err != nil { + t.Fatalf("unmarshal %q: %v", document, err) + } + if cfg.Restore.SkipCompatCheck != want { + t.Errorf("%q parsed skipCompatCheck = %v, want %v", document, cfg.Restore.SkipCompatCheck, want) + } + } +} + func TestAgentConfigValidateRequiresFixedStorageBasePath(t *testing.T) { for _, basePath := range []string{"checkpoints", " /checkpoints ", "/checkpoints/../other", "/other"} { cfg := validAgentConfig() diff --git a/agent/internal/types/inspect.go b/agent/internal/types/inspect.go index 740f3001..023b3ea9 100644 --- a/agent/internal/types/inspect.go +++ b/agent/internal/types/inspect.go @@ -5,6 +5,8 @@ package types import ( specs "github.com/opencontainers/runtime-spec/specs-go" + + "github.com/ai-dynamo/snapshot/api/compat" ) // MountInfo holds parsed mount information from /proc/pid/mountinfo. @@ -21,6 +23,7 @@ type MountInfo struct { // CheckpointContainerSnapshot holds runtime container state collected during checkpoint inspection. type CheckpointContainerSnapshot struct { PID int + ImageID string RootFS string UpperDir string OCISpec *specs.Spec @@ -30,7 +33,10 @@ type CheckpointContainerSnapshot struct { HostCgroupPath string // host filesystem path for CRIU's --freeze-cgroup CUDAHostPIDs []int // host-visible PIDs used for checkpoint-side CUDA actions CUDANSPIDs []int // namespace-relative PIDs stored in the checkpoint manifest - GPUUUIDs []string // source GPU UUIDs from kubelet PodResources API + + // GPUs holds the GPUs the checkpointed container could see, in allocation + // order, with the model and driver version where they could be read. + GPUs compat.GPUFacts } // RestoreContainerSnapshot holds inspected state for the restore target. diff --git a/agent/internal/types/manifest.go b/agent/internal/types/manifest.go index 75e35606..a4dd8a53 100644 --- a/agent/internal/types/manifest.go +++ b/agent/internal/types/manifest.go @@ -7,12 +7,16 @@ import ( "fmt" "os" "path/filepath" + "runtime" + "sort" "strings" "time" criurpc "github.com/checkpoint-restore/go-criu/v8/rpc" specs "github.com/opencontainers/runtime-spec/specs-go" "gopkg.in/yaml.v3" + + "github.com/ai-dynamo/snapshot/api/compat" ) const manifestFilename = "manifest.yaml" @@ -26,6 +30,7 @@ type CheckpointManifest struct { K8s SourcePodManifest `yaml:"k8s"` Overlay OverlayManifest `yaml:"overlay"` CUDA CUDAManifest `yaml:"cudaRestore,omitempty"` + Host HostManifest `yaml:"host,omitempty"` } // ArtifactManifest pins an on-disk checkpoint to the Kubernetes content object @@ -41,6 +46,7 @@ func NewCheckpointManifest( criuDump CRIUDumpManifest, k8s SourcePodManifest, overlay OverlayManifest, + host HostManifest, ) *CheckpointManifest { return &CheckpointManifest{ Artifact: ArtifactManifest{ @@ -51,6 +57,24 @@ func NewCheckpointManifest( CRIUDump: criuDump, K8s: k8s, Overlay: overlay, + Host: host, + } +} + +// HostManifest records the machine a checkpoint was captured on. A fact the +// agent could not read is left out rather than written empty, so it reads as +// unknown instead of as a value that happens to be blank. +type HostManifest struct { + KernelVersion string `yaml:"kernelVersion,omitempty"` + CPUArch string `yaml:"cpuArch,omitempty"` +} + +// NewHostManifest takes the architecture from the agent binary rather than +// asking the node: this binary is running on that node, so they agree. +func NewHostManifest(kernelVersion string) HostManifest { + return HostManifest{ + KernelVersion: kernelVersion, + CPUArch: runtime.GOARCH, } } @@ -94,6 +118,12 @@ type SourcePodManifest struct { // StdioFDs holds readlink targets for FDs 0, 1, 2 (e.g. "pipe:[12345]"). StdioFDs []string `yaml:"stdioFDs,omitempty"` + + // ImageID is CRI ContainerStatus.image_id, not kubelet's image_ref alias. + Image string `yaml:"image,omitempty"` + ImageID string `yaml:"imageId,omitempty"` + CPULimit string `yaml:"cpuLimit,omitempty"` + MemoryLimit string `yaml:"memoryLimit,omitempty"` } func NewSourcePodManifest(containerID string, pid int, sourceNode, podName, podNamespace, podIP string, stdioFDs []string) SourcePodManifest { @@ -142,13 +172,33 @@ func NewOverlayManifest(exclusions OverlaySettings, upperDir string, ociSpec *sp type CUDAManifest struct { PIDs []int `yaml:"pids"` SourceGPUUUIDs []string `yaml:"sourceGpuUuids"` + + // SourceGPUs describes the same GPUs as SourceGPUUUIDs, in the same order. + // The UUIDs stay where they are because the device map is built from them + // and artifacts written before this field exists still restore. + SourceGPUs []GPUManifest `yaml:"sourceGpus,omitempty"` + SourceDriverVersion string `yaml:"sourceDriverVersion,omitempty"` +} + +// GPUManifest is one GPU the checkpointed process could see. +type GPUManifest struct { + UUID string `yaml:"uuid"` + ProductName string `yaml:"productName,omitempty"` } -func NewCUDAManifest(pids []int, sourceGPUUUIDs []string) CUDAManifest { - return CUDAManifest{ - PIDs: append([]int(nil), pids...), - SourceGPUUUIDs: append([]string(nil), sourceGPUUUIDs...), +func NewCUDAManifest(pids []int, gpus compat.GPUFacts) CUDAManifest { + m := CUDAManifest{ + PIDs: append([]int(nil), pids...), + SourceDriverVersion: gpus.DriverVersion, + } + for _, device := range gpus.Devices { + m.SourceGPUUUIDs = append(m.SourceGPUUUIDs, device.UUID) + m.SourceGPUs = append(m.SourceGPUs, GPUManifest{ + UUID: device.UUID, + ProductName: device.ProductName, + }) } + return m } func (m CUDAManifest) IsEmpty() bool { @@ -206,3 +256,65 @@ func validateArtifactManifest(artifact ArtifactManifest) error { } return nil } + +// CompatFacts maps the manifest onto the fact model the compatibility gates +// compare. Both gates read it from here, so the two cannot disagree about what +// the checkpoint recorded. +func (m *CheckpointManifest) CompatFacts() compat.Facts { + gpus := m.gpuFacts() + return compat.Facts{ + KernelVersion: m.Host.KernelVersion, + CPUArch: m.Host.CPUArch, + Image: m.K8s.Image, + ImageID: m.K8s.ImageID, + CPULimit: m.K8s.CPULimit, + MemoryLimit: m.K8s.MemoryLimit, + DriverVersion: gpus.DriverVersion, + GPUDevices: gpus.Devices, + ExternalizedMounts: m.externalizedMounts(), + } +} + +// WithPodFacts records what the captured container ran as. It is the inverse of +// the pod half of CompatFacts, and sits next to it so the two field lists cannot +// drift apart. +func (m SourcePodManifest) WithPodFacts(facts compat.Facts) SourcePodManifest { + m.Image = facts.Image + m.ImageID = facts.ImageID + m.CPULimit = facts.CPULimit + m.MemoryLimit = facts.MemoryLimit + return m +} + +// gpuFacts prefers the described GPUs and falls back to the UUID list, so an +// artifact captured before the models were recorded still reports its GPU count. +func (m *CheckpointManifest) gpuFacts() compat.GPUFacts { + facts := compat.GPUFacts{DriverVersion: m.CUDA.SourceDriverVersion} + if len(m.CUDA.SourceGPUs) > 0 { + for _, gpu := range m.CUDA.SourceGPUs { + facts.Devices = append(facts.Devices, compat.GPUDevice{ + UUID: gpu.UUID, + ProductName: gpu.ProductName, + }) + } + return facts + } + for _, uuid := range m.CUDA.SourceGPUUUIDs { + facts.Devices = append(facts.Devices, compat.GPUDevice{UUID: uuid}) + } + return facts +} + +// externalizedMounts returns the destinations CRIU externalized at capture, in a +// stable order so a refusal always names them the same way. +func (m *CheckpointManifest) externalizedMounts() []string { + if len(m.CRIUDump.ExtMnt) == 0 { + return nil + } + destinations := make([]string, 0, len(m.CRIUDump.ExtMnt)) + for destination := range m.CRIUDump.ExtMnt { + destinations = append(destinations, destination) + } + sort.Strings(destinations) + return destinations +} diff --git a/agent/internal/types/manifest_test.go b/agent/internal/types/manifest_test.go index ac46a85b..e335ed5b 100644 --- a/agent/internal/types/manifest_test.go +++ b/agent/internal/types/manifest_test.go @@ -6,10 +6,15 @@ package types import ( "os" "path/filepath" + "reflect" + "runtime" + "strings" "testing" criurpc "github.com/checkpoint-restore/go-criu/v8/rpc" "google.golang.org/protobuf/proto" + + "github.com/ai-dynamo/snapshot/api/compat" ) func TestManifestRoundTrip(t *testing.T) { @@ -35,8 +40,15 @@ func TestManifestRoundTrip(t *testing.T) { ExternalPaths: []string{"/proc/acpi"}, BindMountDests: []string{"/data"}, }, + NewHostManifest("5.15.0-1071-aws"), ) - original.CUDA = NewCUDAManifest([]int{42, 43}, []string{"GPU-aaa", "GPU-bbb"}) + original.CUDA = NewCUDAManifest([]int{42, 43}, compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{ + {UUID: "GPU-aaa", ProductName: "NVIDIA A100-SXM4-40GB"}, + {UUID: "GPU-bbb", ProductName: "NVIDIA A100-SXM4-40GB"}, + }, + }) if err := WriteManifest(dir, original); err != nil { t.Fatalf("WriteManifest: %v", err) @@ -90,6 +102,184 @@ func TestManifestRoundTrip(t *testing.T) { if len(loaded.CUDA.SourceGPUUUIDs) != 2 || loaded.CUDA.SourceGPUUUIDs[0] != "GPU-aaa" { t.Errorf("CUDA.SourceGPUUUIDs = %v", loaded.CUDA.SourceGPUUUIDs) } + if loaded.CUDA.SourceDriverVersion != "580.65.06" { + t.Errorf("CUDA.SourceDriverVersion = %q", loaded.CUDA.SourceDriverVersion) + } + wantGPUs := []GPUManifest{ + {UUID: "GPU-aaa", ProductName: "NVIDIA A100-SXM4-40GB"}, + {UUID: "GPU-bbb", ProductName: "NVIDIA A100-SXM4-40GB"}, + } + if !reflect.DeepEqual(loaded.CUDA.SourceGPUs, wantGPUs) { + t.Errorf("CUDA.SourceGPUs = %#v, want %#v", loaded.CUDA.SourceGPUs, wantGPUs) + } + wantHost := HostManifest{ + KernelVersion: "5.15.0-1071-aws", + CPUArch: runtime.GOARCH, + } + if !reflect.DeepEqual(loaded.Host, wantHost) { + t.Errorf("Host = %#v, want %#v", loaded.Host, wantHost) + } +} + +func TestSourcePodManifestRecordsTheImageAndItsLimits(t *testing.T) { + dir := t.TempDir() + original := &CheckpointManifest{Artifact: ArtifactManifest{ContentUID: "content-uid-123", ContainerName: "main"}} + original.K8s = NewSourcePodManifest("ctr-abc", 42, "node-1", "my-pod", "default", "10.0.0.11", nil) + original.K8s.Image = "nvcr.io/nvidia/tritonserver:24.09-py3" + original.K8s.ImageID = "docker-pullable://nvcr.io/nvidia/tritonserver@sha256:deadbeef" + original.K8s.CPULimit = "4" + original.K8s.MemoryLimit = "16Gi" + + if err := WriteManifest(dir, original); err != nil { + t.Fatalf("WriteManifest: %v", err) + } + loaded, err := ReadManifest(dir) + if err != nil { + t.Fatalf("ReadManifest: %v", err) + } + if !reflect.DeepEqual(loaded.K8s, original.K8s) { + t.Errorf("K8s = %#v, want %#v", loaded.K8s, original.K8s) + } +} + +// Every checkpoint already on disk was written before any of these facts +// existed. Such an artifact has to keep parsing, and the facts it never +// recorded have to come back unknown - the manifest carries no schema version, +// so absent keys are the entire compatibility mechanism. +func TestReadManifestAcceptsAnArtifactWrittenBeforeTheseFacts(t *testing.T) { + dir := t.TempDir() + legacy := `artifact: + contentUID: content-uid-123 + containerName: main +createdAt: 2026-03-31T00:00:00Z +criuDump: + criu: + logLevel: 4 + extMnt: + /etc/hostname: /etc/hostname +k8s: + containerId: ctr-abc + pid: 42 + sourceNode: node-1 + podName: my-pod + podNamespace: default +overlay: + upperDir: /var/lib/containerd/upper +cudaRestore: + pids: [42] + sourceGpuUuids: [GPU-aaa] +` + if err := os.WriteFile(filepath.Join(dir, manifestFilename), []byte(legacy), 0o600); err != nil { + t.Fatalf("WriteFile: %v", err) + } + + manifest, err := ReadManifest(dir) + if err != nil { + t.Fatalf("ReadManifest: %v", err) + } + facts := manifest.CompatFacts() + if facts.Image != "" || facts.ImageID != "" || facts.CPULimit != "" || facts.MemoryLimit != "" { + t.Errorf("pod facts = %#v, want unknown", facts) + } + if facts.KernelVersion != "" || facts.CPUArch != "" { + t.Errorf("host facts = %#v, want unknown", facts) + } + + // What the artifact does record still has to arrive, or the older + // checkpoints would stop being compared at all. + wantGPUs := []compat.GPUDevice{{UUID: "GPU-aaa"}} + if !reflect.DeepEqual(facts.GPUDevices, wantGPUs) { + t.Errorf("GPU devices = %#v, want %#v", facts.GPUDevices, wantGPUs) + } + if !reflect.DeepEqual(facts.ExternalizedMounts, []string{"/etc/hostname"}) { + t.Errorf("externalized mounts = %#v", facts.ExternalizedMounts) + } +} + +// A host fact the agent could not read has to stay absent in the file, because a +// comparison treats absent as unknown and an empty string as a value. +func TestHostManifestOmitsWhatTheAgentCouldNotRead(t *testing.T) { + dir := t.TempDir() + manifest := &CheckpointManifest{Artifact: ArtifactManifest{ContentUID: "content-uid-123", ContainerName: "main"}} + if err := WriteManifest(dir, manifest); err != nil { + t.Fatalf("WriteManifest: %v", err) + } + + content, err := os.ReadFile(filepath.Join(dir, manifestFilename)) + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + for _, key := range []string{"kernelVersion", "cpuArch"} { + if strings.Contains(string(content), key) { + t.Errorf("manifest wrote an unknown %s:\n%s", key, content) + } + } +} + +// The facts are recorded to be compared, so what a manifest carries has to come +// back out as the source side of a comparison, one group at a time. +func TestManifestFactsSurviveIntoTheComparison(t *testing.T) { + tests := []struct { + name string + manifest *CheckpointManifest + want compat.Facts + }{ + { + name: "host facts", + manifest: &CheckpointManifest{ + Host: NewHostManifest("5.15.0-1071-aws"), + }, + want: compat.Facts{KernelVersion: "5.15.0-1071-aws", CPUArch: runtime.GOARCH}, + }, + { + name: "pod facts", + manifest: &CheckpointManifest{K8s: SourcePodManifest{ + Image: "nvcr.io/nvidia/tritonserver:24.09-py3", + ImageID: "docker-pullable://nvcr.io/nvidia/tritonserver@sha256:deadbeef", + CPULimit: "4", + MemoryLimit: "16Gi", + }}, + want: compat.Facts{ + Image: "nvcr.io/nvidia/tritonserver:24.09-py3", + ImageID: "docker-pullable://nvcr.io/nvidia/tritonserver@sha256:deadbeef", + CPULimit: "4", + MemoryLimit: "16Gi", + }, + }, + { + name: "GPU facts", + manifest: &CheckpointManifest{ + CUDA: NewCUDAManifest([]int{1}, compat.GPUFacts{ + DriverVersion: "580.65.06", + Devices: []compat.GPUDevice{{UUID: "GPU-aaa", ProductName: "NVIDIA L4"}}, + }), + }, + want: compat.Facts{ + DriverVersion: "580.65.06", + GPUDevices: []compat.GPUDevice{{UUID: "GPU-aaa", ProductName: "NVIDIA L4"}}, + }, + }, + { + // An artifact captured before the models were recorded still has to + // report how many GPUs it used, or the count rule would silently stop + // applying to every checkpoint taken so far. + name: "GPU facts recorded as UUIDs alone", + manifest: &CheckpointManifest{ + CUDA: CUDAManifest{PIDs: []int{1}, SourceGPUUUIDs: []string{"GPU-aaa", "GPU-bbb"}}, + }, + want: compat.Facts{ + GPUDevices: []compat.GPUDevice{{UUID: "GPU-aaa"}, {UUID: "GPU-bbb"}}, + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := tc.manifest.CompatFacts(); !reflect.DeepEqual(got, tc.want) { + t.Errorf("CompatFacts = %#v, want %#v", got, tc.want) + } + }) + } } func TestNewCRIUDumpManifest(t *testing.T) { diff --git a/api/compat/checks.go b/api/compat/checks.go new file mode 100644 index 00000000..8937c3d8 --- /dev/null +++ b/api/compat/checks.go @@ -0,0 +1,320 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package compat + +import ( + "sort" + "strconv" + "strings" + + "k8s.io/apimachinery/pkg/api/resource" +) + +// CheckCPUArch refuses a restore onto a different instruction set. A checkpoint +// holds register state, and CRIU has nowhere to put an x86 register file on an +// ARM core, so this one can never be waived by a bigger machine or a newer +// driver - it is the hardest of the rules. +const CheckCPUArch Check = "cpu-arch" + +var cpuArchCheck = check{ + name: CheckCPUArch, + gate: GatePreflight, + compare: func(source, target Facts) []Mismatch { return mustMatch(source.CPUArch, target.CPUArch) }, +} + +// CheckKernelVersion refuses a restore onto a kernel other than the captured +// one. Restores that had worked for over a year have broken on a kernel upgrade +// alone: criu#2636. +const CheckKernelVersion Check = "kernel-version" + +// CheckKernelMinimum refuses a kernel too old to restore a modern glibc at all. +// glibc uses rseq, which needs 5.13 (criu#2229), and glibc 2.35 and newer +// segfault on restore below it (criu#2552). +const CheckKernelMinimum Check = "kernel-minimum" + +const ( + minKernelMajor = 5 + minKernelMinor = 13 +) + +var kernelVersionCheck = check{ + name: CheckKernelVersion, + gate: GatePreflight, + compare: func(source, target Facts) []Mismatch { + return mustMatch(source.KernelVersion, target.KernelVersion) + }, +} + +var kernelMinimumCheck = check{ + name: CheckKernelMinimum, + gate: GatePreflight, + compare: func(_, target Facts) []Mismatch { + major, minor, ok := parseKernelVersion(target.KernelVersion) + if !ok || major > minKernelMajor || (major == minKernelMajor && minor >= minKernelMinor) { + return nil + } + return []Mismatch{{ + Source: strconv.Itoa(minKernelMajor) + "." + strconv.Itoa(minKernelMinor) + " or newer", + Target: target.KernelVersion, + }} + }, +} + +// parseKernelVersion reads the leading major.minor of a kernel release, which +// carries a distro suffix after it as in "5.15.0-1071-aws". A release it cannot +// read is unknown, which leaves the floor to the equality rule above. +func parseKernelVersion(version string) (major, minor int, ok bool) { + fields := strings.SplitN(strings.TrimSpace(version), ".", 3) + if len(fields) < 2 { + return 0, 0, false + } + major, majorOK := leadingNumber(fields[0]) + minor, minorOK := leadingNumber(fields[1]) + return major, minor, majorOK && minorOK +} + +// leadingNumber reads the digits a version field starts with, ignoring whatever +// a distro appended to them. +func leadingNumber(field string) (int, bool) { + end := 0 + for end < len(field) && field[end] >= '0' && field[end] <= '9' { + end++ + } + if end == 0 { + return 0, false + } + value, err := strconv.Atoi(field[:end]) + return value, err == nil +} + +// CheckImageDigest refuses a restore into image content other than what was +// captured. The reference is not compared: the same content is reachable under +// more than one, and a rebuilt or moved tag is one reference over two contents. +const CheckImageDigest Check = "image-digest" + +var imageDigestCheck = check{ + name: CheckImageDigest, + gate: GateInspect, + compare: func(source, target Facts) []Mismatch { + return mustMatch(imageDigest(source.ImageID), imageDigest(target.ImageID)) + }, +} + +// CheckMemoryLimit refuses a restore into less memory than the checkpoint was +// captured with. Restoring faults the whole recorded address space back in, so a +// lower ceiling is not a slower restore but an OOM kill partway through one. +const CheckMemoryLimit Check = "memory-limit" + +var memoryLimitCheck = check{ + name: CheckMemoryLimit, + gate: GatePreflight, + compare: func(source, target Facts) []Mismatch { + return atLeastSource(source.MemoryLimit, target.MemoryLimit) + }, +} + +// CheckCPULimit refuses a restore into less CPU than the checkpoint was captured +// with. Unlike memory, too little does not fail: the workload restores, reports +// success, and runs measurably slower forever, which is the worst outcome of the +// three because nothing says it happened. +const CheckCPULimit Check = "cpu-limit" + +var cpuLimitCheck = check{ + name: CheckCPULimit, + gate: GatePreflight, + compare: func(source, target Facts) []Mismatch { + return atLeastSource(source.CPULimit, target.CPULimit) + }, +} + +// atLeastSource reports a mismatch when the target is given less than the +// checkpoint was captured with. A quantity absent or unreadable on either side is +// unknown - which is also how an unlimited pod reads, since a pod with no limit +// records none. +// +// It is deliberately blunt: a deployment that was genuinely over-provisioned and +// is being trimmed on purpose is refused too, and the escape hatch is the way +// out of that. +func atLeastSource(source, target string) []Mismatch { + sourceQuantity, err := resource.ParseQuantity(source) + if err != nil { + return nil + } + targetQuantity, err := resource.ParseQuantity(target) + if err != nil { + return nil + } + if targetQuantity.Cmp(sourceQuantity) >= 0 { + return nil + } + return []Mismatch{{Source: source, Target: target}} +} + +// imageDigest reduces a container status image ID to the digest inside it. +// Runtimes disagree on the wrapping - containerd reports a bare "sha256:...", +// others a scheme and a repository around it - and the artifact keeps whichever +// form it was given, so the two are only comparable after this. +func imageDigest(imageID string) string { + digest := strings.TrimSpace(imageID) + if scheme := strings.Index(digest, "://"); scheme >= 0 { + digest = digest[scheme+len("://"):] + } + if at := strings.LastIndex(digest, "@"); at >= 0 { + digest = digest[at+1:] + } + return digest +} + +// CheckMount refuses a restore into a pod that is missing a path the checkpoint +// had mounted. CRIU was told to leave those mounts alone and expect them to be +// there; where one is absent, the restored process gets a working directory or a +// dataset that simply is not there, and finds out by failing later. +const CheckMount Check = "mount" + +// criuHandledMounts are recorded as externalized but reconstructed by CRIU +// itself, so their absence from the target pod is not a missing volume. +var criuHandledMounts = map[string]bool{ + "/": true, + "/dev/shm": true, +} + +var mountCheck = check{ + name: CheckMount, + gate: GateInspect, + compare: func(source, target Facts) []Mismatch { + existing := make(map[string]bool, len(target.ExistingMounts)) + for _, path := range target.ExistingMounts { + existing[path] = true + } + + var mismatches []Mismatch + for _, path := range source.ExternalizedMounts { + if existing[path] || criuHandledMounts[path] { + continue + } + mismatches = append(mismatches, Mismatch{Source: path, Target: "missing"}) + } + return mismatches + }, +} + +// CheckGPUModel refuses a restore onto a different GPU model. A CUDA checkpoint +// carries device state built for one architecture's memory layout and +// capabilities, and no amount of driver compatibility makes an A100 replay what +// an L4 was doing. +const CheckGPUModel Check = "gpu-model" + +var gpuModelCheck = check{ + name: CheckGPUModel, + gate: GateInspect, + compare: func(source, target Facts) []Mismatch { + sourceModels, sourceOK := gpuModels(source.GPUDevices) + targetModels, targetOK := gpuModels(target.GPUDevices) + if !sourceOK || !targetOK || sourceModels == targetModels { + return nil + } + return []Mismatch{{Source: sourceModels, Target: targetModels}} + }, +} + +// CheckGPUCount refuses a restore onto a different number of GPUs. A multi-GPU +// checkpoint holds one piece of device state per GPU with a rank each, and there +// is no meaning to be given to a rank that has nowhere to land - or to a GPU no +// rank was recorded for. +// +// A target with no GPUs at all is that same refusal and is reported as one. It +// reaches here only once discovery has run, so none found means none, and the +// alternative is the unnamed device-map error further in. +const CheckGPUCount Check = "gpu-count" + +var gpuCountCheck = check{ + name: CheckGPUCount, + gate: GateInspect, + compare: func(source, target Facts) []Mismatch { + sourceCount := len(source.GPUDevices) + targetCount := len(target.GPUDevices) + if sourceCount == 0 || sourceCount == targetCount { + return nil + } + return []Mismatch{{ + Source: strconv.Itoa(sourceCount), + Target: strconv.Itoa(targetCount), + }} + }, +} + +// CheckDriverVersion refuses a restore on a driver build other than the captured +// one. Build granularity is not caution for its own sake: upstream reproduces a +// restore failure between 560.35.03 and 560.35.05. +const CheckDriverVersion Check = "driver-version" + +// CheckDriverMinimum refuses a driver older than CUDA checkpoint and restore is +// supported on at all. +const CheckDriverMinimum Check = "driver-minimum" + +const minDriverMajor = 580 + +var driverVersionCheck = check{ + name: CheckDriverVersion, + gate: GateInspect, + compare: func(source, target Facts) []Mismatch { + return mustMatch(source.DriverVersion, target.DriverVersion) + }, +} + +var driverMinimumCheck = check{ + name: CheckDriverMinimum, + gate: GateInspect, + compare: func(_, target Facts) []Mismatch { + major, ok := leadingNumber(target.DriverVersion) + if !ok || major >= minDriverMajor { + return nil + } + return []Mismatch{{ + Source: strconv.Itoa(minDriverMajor) + " or newer", + Target: target.DriverVersion, + }} + }, +} + +// gpuModels builds a stable model summary: sorting ignores allocation order, +// while "xN" preserves how many GPUs have each name. ProductName comes from +// nvidia-smi --query-gpu=name, documented as the GPU's official product name: +// https://docs.nvidia.com/deploy/nvidia-smi/index.html#product-name +// +// It returns unknown if any GPU has no name because partial data cannot prove +// that the source and target models differ. +func gpuModels(devices []GPUDevice) (string, bool) { + if len(devices) == 0 { + return "", false + } + counts := make(map[string]int, len(devices)) + for _, device := range devices { + model := strings.TrimSpace(device.ProductName) + if model == "" { + return "", false + } + counts[model]++ + } + + models := make([]string, 0, len(counts)) + for model := range counts { + models = append(models, model) + } + sort.Strings(models) + for i, model := range models { + models[i] = model + " x" + strconv.Itoa(counts[model]) + } + return strings.Join(models, ", "), true +} + +// mustMatch reports a mismatch unless the two values are identical. A value +// absent on either side is unknown, and an unknown fact never refuses a restore: +// a checkpoint captured before it was ever recorded has to stay restorable. +func mustMatch(source, target string) []Mismatch { + if source == "" || target == "" || source == target { + return nil + } + return []Mismatch{{Source: source, Target: target}} +} diff --git a/api/compat/checks_test.go b/api/compat/checks_test.go new file mode 100644 index 00000000..4d0811e3 --- /dev/null +++ b/api/compat/checks_test.go @@ -0,0 +1,645 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package compat + +import ( + "reflect" + "testing" +) + +func TestCPUArchCheck(t *testing.T) { + arch := func(value string) Facts { + return Facts{CPUArch: value} + } + + tests := []struct { + name string + source Facts + target Facts + want []Mismatch + }{ + { + name: "same architecture", + source: arch("amd64"), + target: arch("amd64"), + }, + { + name: "different architecture", + source: arch("amd64"), + target: arch("arm64"), + want: []Mismatch{{Check: CheckCPUArch, Source: "amd64", Target: "arm64"}}, + }, + { + name: "checkpoint taken before the architecture was recorded", + target: arch("arm64"), + }, + { + name: "target architecture unknown", + source: arch("amd64"), + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := Compare(GatePreflight, tc.source, tc.target) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("Compare = %+v, want %+v", got, tc.want) + } + }) + } + + // The architecture is decidable from the manifest and the node alone, so + // waiting for the placeholder container would delay the refusal for nothing. + if got := Compare(GateInspect, arch("amd64"), arch("arm64")); len(got) != 0 { + t.Errorf("the second gate repeated the architecture check: %+v", got) + } +} + +func TestKernelVersionCheck(t *testing.T) { + kernel := func(value string) Facts { + return Facts{KernelVersion: value} + } + + tests := []struct { + name string + source Facts + target Facts + want []Mismatch + }{ + { + name: "same kernel release", + source: kernel("5.15.0-1071-aws"), + target: kernel("5.15.0-1071-aws"), + }, + { + // A kernel upgrade alone is enough to break a restore that worked, + // so the same major.minor on a different release is not good enough. + name: "same major and minor, different release", + source: kernel("5.15.0-1071-aws"), + target: kernel("5.15.0-1082-aws"), + want: []Mismatch{{ + Check: CheckKernelVersion, + Source: "5.15.0-1071-aws", + Target: "5.15.0-1082-aws", + }}, + }, + { + name: "checkpoint taken before the kernel was recorded", + target: kernel("5.15.0-1071-aws"), + }, + { + name: "target kernel unknown", + source: kernel("5.15.0-1071-aws"), + }, + { + // Both rules fire: the node runs a different kernel, and one no + // restore of a modern glibc can succeed on. + name: "target below the floor", + source: kernel("5.15.0-1071-aws"), + target: kernel("5.4.0-150-generic"), + want: []Mismatch{ + {Check: CheckKernelVersion, Source: "5.15.0-1071-aws", Target: "5.4.0-150-generic"}, + {Check: CheckKernelMinimum, Source: "5.13 or newer", Target: "5.4.0-150-generic"}, + }, + }, + { + name: "both sides below the floor", + source: kernel("4.19.0-25-amd64"), + target: kernel("4.19.0-25-amd64"), + want: []Mismatch{ + {Check: CheckKernelMinimum, Source: "5.13 or newer", Target: "4.19.0-25-amd64"}, + }, + }, + { + name: "exactly at the floor", + source: kernel("5.13.0-52-generic"), + target: kernel("5.13.0-52-generic"), + }, + { + name: "a newer major is above the floor", + source: kernel("6.8.0-45-generic"), + target: kernel("6.8.0-45-generic"), + }, + { + // A release the floor cannot read is unknown rather than old, so a + // kernel string in a form nobody anticipated does not refuse every + // restore on the node. + name: "unreadable release", + source: kernel("custom-build"), + target: kernel("custom-build"), + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := Compare(GatePreflight, tc.source, tc.target) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("Compare = %+v, want %+v", got, tc.want) + } + }) + } +} + +func TestImageDigestCheck(t *testing.T) { + const ( + captured = "sha256:1111111111111111111111111111111111111111111111111111111111111111" + rebuilt = "sha256:2222222222222222222222222222222222222222222222222222222222222222" + ) + imageID := func(id string) Facts { + return Facts{ImageID: id} + } + + tests := []struct { + name string + source Facts + target Facts + want []Mismatch + }{ + { + name: "same content", + source: imageID(captured), + target: imageID(captured), + }, + { + name: "different index aliases for the same runtime content", + source: Facts{Image: "registry.example/source@sha256:index-a", ImageID: captured}, + target: Facts{Image: "registry.example/target@sha256:index-b", ImageID: captured}, + }, + { + // The same reference resolved to different content, which is what a + // rebuilt or moved tag looks like from here. + name: "same reference, rebuilt content", + source: imageID(captured), + target: imageID(rebuilt), + want: []Mismatch{{Check: CheckImageDigest, Source: captured, Target: rebuilt}}, + }, + { + // Runtimes wrap the digest differently, and the artifact keeps + // whatever it was given. The same content must not read as a + // mismatch because one side spells it out and the other does not. + name: "the same digest wrapped differently", + source: imageID("docker-pullable://nvcr.io/nvidia/tritonserver@" + captured), + target: imageID(captured), + }, + { + name: "different content, wrapped differently", + source: imageID("docker-pullable://nvcr.io/nvidia/tritonserver@" + captured), + target: imageID(rebuilt), + want: []Mismatch{{Check: CheckImageDigest, Source: captured, Target: rebuilt}}, + }, + { + name: "checkpoint taken before the image ID was recorded", + target: imageID(captured), + }, + { + name: "runtime image ID is unavailable", + source: imageID(captured), + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := Compare(GateInspect, tc.source, tc.target) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("Compare = %+v, want %+v", got, tc.want) + } + }) + } + + if got := Compare(GatePreflight, imageID(captured), imageID(rebuilt)); len(got) != 0 { + t.Errorf("the first gate judged a runtime fact it cannot read: %+v", got) + } +} + +func TestMemoryLimitCheck(t *testing.T) { + memory := func(limit string) Facts { + return Facts{MemoryLimit: limit} + } + + tests := []struct { + name string + source Facts + target Facts + want []Mismatch + }{ + { + name: "the same limit", + source: memory("32Gi"), + target: memory("32Gi"), + }, + { + name: "a larger limit", + source: memory("32Gi"), + target: memory("64Gi"), + }, + { + name: "a smaller limit", + source: memory("32Gi"), + target: memory("1Gi"), + want: []Mismatch{{Check: CheckMemoryLimit, Source: "32Gi", Target: "1Gi"}}, + }, + { + // The same amount written another way is the same amount. + name: "the same limit in different units", + source: memory("32Gi"), + target: memory("34359738368"), + }, + { + // A pod with no limit records none, so this is also how a restore + // into an unlimited pod reads: nothing to compare. + name: "the target has no limit", + source: memory("32Gi"), + }, + { + name: "the checkpoint recorded no limit", + target: memory("1Gi"), + }, + { + name: "an unreadable quantity", + source: memory("32Gi"), + target: memory("plenty"), + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := Compare(GatePreflight, tc.source, tc.target) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("Compare = %+v, want %+v", got, tc.want) + } + }) + } +} + +func TestCPULimitCheck(t *testing.T) { + cpu := func(limit string) Facts { + return Facts{CPULimit: limit} + } + + tests := []struct { + name string + source Facts + target Facts + want []Mismatch + }{ + { + name: "the same limit", + source: cpu("4"), + target: cpu("4"), + }, + { + name: "a larger limit", + source: cpu("4"), + target: cpu("8"), + }, + { + name: "a smaller limit", + source: cpu("4"), + target: cpu("1"), + want: []Mismatch{{Check: CheckCPULimit, Source: "4", Target: "1"}}, + }, + { + // Millicores and whole cores are the same scale, so the comparison + // has to see through the notation. + name: "millicores below whole cores", + source: cpu("4"), + target: cpu("500m"), + want: []Mismatch{{Check: CheckCPULimit, Source: "4", Target: "500m"}}, + }, + { + name: "millicores above whole cores", + source: cpu("1"), + target: cpu("4500m"), + }, + { + name: "the same limit written as millicores", + source: cpu("4"), + target: cpu("4000m"), + }, + { + name: "the target has no limit", + source: cpu("4"), + }, + { + name: "the checkpoint recorded no limit", + target: cpu("1"), + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := Compare(GatePreflight, tc.source, tc.target) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("Compare = %+v, want %+v", got, tc.want) + } + }) + } +} + +func TestMountCheck(t *testing.T) { + mounts := func(externalized, existing []string) Facts { + return Facts{ExternalizedMounts: externalized, ExistingMounts: existing} + } + + tests := []struct { + name string + facts Facts + want []Mismatch + }{ + { + name: "every mount is there", + facts: mounts([]string{"/model-cache", "/data"}, []string{"/model-cache", "/data"}), + }, + { + name: "one mount is missing", + facts: mounts([]string{"/model-cache", "/data"}, []string{"/model-cache"}), + want: []Mismatch{{Check: CheckMount, Source: "/data", Target: "missing"}}, + }, + { + // Each missing volume is named, since a user fixing their pod needs + // to know about all of them and not one at a time. + name: "the pod has none of them", + facts: mounts([]string{"/model-cache", "/data"}, nil), + want: []Mismatch{ + {Check: CheckMount, Source: "/model-cache", Target: "missing"}, + {Check: CheckMount, Source: "/data", Target: "missing"}, + }, + }, + { + // CRIU reconstructs these itself, so their absence from the pod is + // not a volume anybody forgot to declare. + name: "the mounts CRIU restores itself", + facts: mounts([]string{"/", "/dev/shm", "/model-cache"}, []string{"/model-cache"}), + }, + { + name: "the checkpoint externalized nothing", + facts: mounts(nil, nil), + }, + { + // The target side is resolved from the recorded list, so a path the + // pod has and the checkpoint never used is not the gate's business. + name: "the pod has more than the checkpoint used", + facts: mounts([]string{"/model-cache"}, []string{"/model-cache", "/scratch"}), + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := Compare(GateInspect, tc.facts, tc.facts) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("Compare = %+v, want %+v", got, tc.want) + } + }) + } + + // Whether a path resolves is only knowable once the placeholder container + // exists, which is after the first gate has already run. + missing := mounts([]string{"/model-cache"}, nil) + if got := Compare(GatePreflight, missing, missing); len(got) != 0 { + t.Errorf("the first gate judged a mount it cannot see: %+v", got) + } +} + +func gpus(models ...string) Facts { + devices := make([]GPUDevice, 0, len(models)) + for i, model := range models { + devices = append(devices, GPUDevice{UUID: "GPU-" + string(rune('a'+i)), ProductName: model}) + } + return Facts{GPUDevices: devices} +} + +func TestGPUModelCheck(t *testing.T) { + tests := []struct { + name string + source Facts + target Facts + want []Mismatch + }{ + { + name: "same model", + source: gpus("NVIDIA L4"), + target: gpus("NVIDIA L4"), + }, + { + name: "different model", + source: gpus("NVIDIA L4"), + target: gpus("NVIDIA A100-SXM4-40GB"), + want: []Mismatch{{ + Check: CheckGPUModel, + Source: "NVIDIA L4 x1", + Target: "NVIDIA A100-SXM4-40GB x1", + }}, + }, + { + // Which GPU is allocated at which index is the device map's concern. + // For the model, two of the same is two of the same. + name: "the same models in another order", + source: gpus("NVIDIA L4", "Tesla T4"), + target: gpus("Tesla T4", "NVIDIA L4"), + }, + { + name: "one model replaced in a mixed set", + source: gpus("NVIDIA L4", "Tesla T4"), + target: gpus("NVIDIA L4", "NVIDIA L4"), + want: []Mismatch{{ + Check: CheckGPUModel, + Source: "NVIDIA L4 x1, Tesla T4 x1", + Target: "NVIDIA L4 x2", + }}, + }, + { + name: "checkpoint taken before the models were recorded", + source: Facts{GPUDevices: []GPUDevice{{UUID: "GPU-a"}}}, + target: gpus("NVIDIA L4"), + }, + { + // Surrounding whitespace is already insignificant to the blank + // guard, and a refusal printing two identical-looking names is one + // nobody can act on. + name: "the same model padded on one side", + source: gpus("NVIDIA L4"), + target: gpus(" NVIDIA L4 "), + }, + { + // The GPUs were found but could not be described, which is not the + // same as being different. + name: "target models could not be read", + source: gpus("NVIDIA L4"), + target: Facts{GPUDevices: []GPUDevice{{UUID: "GPU-a"}}}, + }, + { + // Nothing to name is not another name, so the model rule stays + // quiet and the count rule is the one that refuses. + name: "no GPUs on the target at all", + source: gpus("NVIDIA L4"), + want: []Mismatch{{Check: CheckGPUCount, Source: "1", Target: "0"}}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := Compare(GateInspect, tc.source, tc.target) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("Compare = %+v, want %+v", got, tc.want) + } + }) + } + + // The GPUs a container can see are only readable once it exists, so the + // first gate has nothing to compare. + if got := Compare(GatePreflight, gpus("NVIDIA L4"), gpus("Tesla T4")); len(got) != 0 { + t.Errorf("the first gate judged a GPU it cannot see: %+v", got) + } +} + +func TestGPUCountCheck(t *testing.T) { + tests := []struct { + name string + source Facts + target Facts + want []Mismatch + }{ + { + name: "the same number of GPUs", + source: gpus("NVIDIA L4", "NVIDIA L4"), + target: gpus("NVIDIA L4", "NVIDIA L4"), + }, + { + // Both rules speak: the set of models differs by a count, and so + // does the number of GPUs. Each says something the other does not, + // and a refusal names every rule it failed. + name: "fewer GPUs than were captured", + source: gpus("NVIDIA L4", "NVIDIA L4"), + target: gpus("NVIDIA L4"), + want: []Mismatch{ + {Check: CheckGPUModel, Source: "NVIDIA L4 x2", Target: "NVIDIA L4 x1"}, + {Check: CheckGPUCount, Source: "2", Target: "1"}, + }, + }, + { + // More is not better here: a checkpoint records one piece of device + // state per GPU, and a spare GPU has no rank to take. + name: "more GPUs than were captured", + source: gpus("NVIDIA L4"), + target: gpus("NVIDIA L4", "NVIDIA L4"), + want: []Mismatch{ + {Check: CheckGPUModel, Source: "NVIDIA L4 x1", Target: "NVIDIA L4 x2"}, + {Check: CheckGPUCount, Source: "1", Target: "2"}, + }, + }, + { + name: "the checkpoint recorded no GPUs", + target: gpus("NVIDIA L4"), + }, + { + // The device map pairs source and target UUIDs by position, so a + // target with none of its own fails there over a GPU it cannot + // name. Refuse it here, where it can be named. + name: "no GPUs were discovered on the target", + source: gpus("NVIDIA L4"), + want: []Mismatch{{Check: CheckGPUCount, Source: "1", Target: "0"}}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := Compare(GateInspect, tc.source, tc.target) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("Compare = %+v, want %+v", got, tc.want) + } + }) + } + + // A checkpoint from before the models were recorded still knows how many + // GPUs it used, so the count keeps applying where the model rule cannot. + unnamed := Facts{GPUDevices: []GPUDevice{{UUID: "GPU-a"}, {UUID: "GPU-b"}}} + want := []Mismatch{{Check: CheckGPUCount, Source: "2", Target: "1"}} + if got := Compare(GateInspect, unnamed, gpus("NVIDIA L4")); !reflect.DeepEqual(got, want) { + t.Errorf("Compare = %+v, want %+v", got, want) + } +} + +func TestDriverVersionCheck(t *testing.T) { + driver := func(version string) Facts { + return Facts{DriverVersion: version} + } + + tests := []struct { + name string + source Facts + target Facts + want []Mismatch + }{ + { + name: "the same driver build", + source: driver("580.65.06"), + target: driver("580.65.06"), + }, + { + // Two builds of the same release are not interchangeable: upstream + // reproduces a restore failure across exactly this distance. + name: "a different build of the same release", + source: driver("580.65.06"), + target: driver("580.65.08"), + want: []Mismatch{{ + Check: CheckDriverVersion, + Source: "580.65.06", + Target: "580.65.08", + }}, + }, + { + name: "a newer driver", + source: driver("580.65.06"), + target: driver("585.10.01"), + want: []Mismatch{{ + Check: CheckDriverVersion, + Source: "580.65.06", + Target: "585.10.01", + }}, + }, + { + // Both rules speak: a different driver, and one CUDA checkpoint and + // restore is not supported on at all. + name: "a driver below the floor", + source: driver("580.65.06"), + target: driver("560.35.03"), + want: []Mismatch{ + {Check: CheckDriverVersion, Source: "580.65.06", Target: "560.35.03"}, + {Check: CheckDriverMinimum, Source: "580 or newer", Target: "560.35.03"}, + }, + }, + { + name: "both sides below the floor", + source: driver("560.35.03"), + target: driver("560.35.03"), + want: []Mismatch{ + {Check: CheckDriverMinimum, Source: "580 or newer", Target: "560.35.03"}, + }, + }, + { + name: "checkpoint taken before the driver was recorded", + target: driver("580.65.06"), + }, + { + name: "the target driver could not be read", + source: driver("580.65.06"), + }, + { + // A version in a form the floor cannot read is unknown rather than + // old, so it does not refuse every restore on the node. + name: "an unreadable version", + source: driver("vendor-build"), + target: driver("vendor-build"), + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := Compare(GateInspect, tc.source, tc.target) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("Compare = %+v, want %+v", got, tc.want) + } + }) + } +} diff --git a/api/compat/compat.go b/api/compat/compat.go new file mode 100644 index 00000000..ed1d0bdb --- /dev/null +++ b/api/compat/compat.go @@ -0,0 +1,154 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Package compat compares what a checkpoint was captured on against what a +// restore target offers, so an incompatible restore is refused up front instead +// of failing deep inside CRIU with an unattributable error. +// +// The node agent is the only consumer today. It sits in the api module because +// the check names are protocol, published verbatim on pod conditions and events, +// and so the operator can surface the same recorded facts on the content status +// without a second vocabulary growing up beside this one. +package compat + +import "fmt" + +// Gate names the moment a comparison runs. The two gates see different facts: +// only the later one can read the node's GPUs and the target's rootfs. +type Gate string + +const ( + // GatePreflight runs before the agent claims a restore attempt, where the + // checkpoint manifest and the target pod are all that is readable. + GatePreflight Gate = "preflight" + + // GateInspect runs once the placeholder container is resolved, where the + // GPUs it sees and the mounts under its rootfs become readable. + GateInspect Gate = "inspect" +) + +// Check identifies one comparison rule. It is reported verbatim to users and to +// tooling that branches on it, so a name never changes once released. +type Check string + +// Facts is one side of a comparison: the machine, pod, GPU and mount state a +// checkpoint was captured on, or the state a restore target offers. +// +// Every field is optional. A fact missing on either side is unknown rather than +// mismatched, because a checkpoint captured before that fact was ever recorded +// has to stay restorable. A producer that reads only some of these returns a +// Facts with those set and leaves its caller to fill in the rest. +type Facts struct { + KernelVersion string + CPUArch string + + Image string + ImageID string + CPULimit string + MemoryLimit string + + DriverVersion string + GPUDevices []GPUDevice + + // ExternalizedMounts holds the mount destinations CRIU externalized at + // capture. + ExternalizedMounts []string + + // ExistingMounts holds the destinations that resolve on this machine. The + // agent resolves them before comparing, so a comparison never touches disk. + ExistingMounts []string +} + +// GPUFacts is what discovery reads off a node before it is folded into a Facts. +// Discovery has no business carrying kernel versions and image digests around, +// so it keeps a type of its own. +type GPUFacts struct { + DriverVersion string + Devices []GPUDevice +} + +// GPUDevice is one visible GPU. +type GPUDevice struct { + UUID string + ProductName string +} + +// Mismatch is one rule the target failed, carrying both compared values so the +// reported reason can name them. +type Mismatch struct { + Check Check + Source string + Target string +} + +// IncompatibleError reports a restore the target cannot run. It is terminal and +// distinct from every other restore error: no CRIU work was attempted, and +// retrying on this node cannot change the answer. Both gates raise it, so the +// caller reports one refusal whichever gate turned the restore down. +type IncompatibleError struct { + Gate Gate + Mismatches []Mismatch +} + +func NewIncompatibleError(gate Gate, mismatches []Mismatch) *IncompatibleError { + return &IncompatibleError{Gate: gate, Mismatches: append([]Mismatch(nil), mismatches...)} +} + +func (e *IncompatibleError) Error() string { + return "restore refused as incompatible: " + Reasons(e.Mismatches) +} + +// check is one row of the policy table. compare returns nil when the rule passes +// or when a fact it needs is unknown, and may report more than one mismatch when +// a rule covers several values. +type check struct { + name Check + gate Gate + compare func(source, target Facts) []Mismatch +} + +// checksByGate is the policy table: every compatibility rule, partitioned by the +// gate that can evaluate it and kept in the order they are reported. +var checksByGate = registerChecks( + cpuArchCheck, + kernelVersionCheck, + kernelMinimumCheck, + imageDigestCheck, + memoryLimitCheck, + cpuLimitCheck, + mountCheck, + gpuModelCheck, + gpuCountCheck, + driverVersionCheck, + driverMinimumCheck, +) + +// registerChecks partitions the policy table by the gate each rule runs at, so +// a comparison indexes its rules instead of walking past the ones belonging to +// the other gate. A rule pinned to a gate nothing calls would never run and +// nothing would say so, so the table refuses to be built at all. +func registerChecks(checks ...check) map[Gate][]check { + byGate := make(map[Gate][]check) + for _, c := range checks { + switch c.gate { + case GatePreflight, GateInspect: + byGate[c.gate] = append(byGate[c.gate], c) + default: + panic(fmt.Sprintf("compat: check %q runs at gate %q, which nothing calls", c.name, c.gate)) + } + } + return byGate +} + +// Compare reports every rule the target fails at the given gate. An empty result +// means the restore may proceed. +func Compare(gate Gate, source, target Facts) []Mismatch { + var mismatches []Mismatch + for _, c := range checksByGate[gate] { + for _, mismatch := range c.compare(source, target) { + mismatch.Check = c.name + mismatches = append(mismatches, mismatch) + } + } + return mismatches +} diff --git a/api/compat/compat_test.go b/api/compat/compat_test.go new file mode 100644 index 00000000..0f59be8b --- /dev/null +++ b/api/compat/compat_test.go @@ -0,0 +1,221 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package compat + +import ( + "errors" + "fmt" + "testing" +) + +func populatedFacts() Facts { + return Facts{ + KernelVersion: "6.8.0-45-generic", + CPUArch: "amd64", + Image: "nvcr.io/nvidia/tritonserver:24.09", + ImageID: "sha256:1111111111111111111111111111111111111111111111111111111111111111", + CPULimit: "8", + MemoryLimit: "32Gi", + DriverVersion: "580.82.07", + GPUDevices: []GPUDevice{ + {UUID: "GPU-1111", ProductName: "Tesla T4"}, + }, + ExternalizedMounts: []string{"/model-cache"}, + ExistingMounts: []string{"/model-cache"}, + } +} + +func differentFacts() Facts { + return Facts{ + KernelVersion: "5.15.0-89-generic", + CPUArch: "arm64", + Image: "nvcr.io/nvidia/tritonserver:24.01", + ImageID: "sha256:2222222222222222222222222222222222222222222222222222222222222222", + CPULimit: "1", + MemoryLimit: "1Gi", + DriverVersion: "560.35.03", + GPUDevices: []GPUDevice{ + {UUID: "GPU-2222", ProductName: "NVIDIA A100-SXM4-40GB"}, + }, + ExternalizedMounts: []string{"/model-cache"}, + ExistingMounts: nil, + } +} + +// deliberatelyNotSilent holds the rules that do refuse on a fact the target side +// does not carry, for the reasons given where each is defined: the mount rule is +// handed a target list resolved from the source list, and the GPU count is only +// ever compared after discovery has run. An absence in either is a thing looked +// for and not found rather than a thing nobody read. +var deliberatelyNotSilent = map[Check]bool{ + CheckMount: true, + CheckGPUCount: true, +} + +// Whatever rules are registered, a fact nobody recorded cannot refuse anything: +// every checkpoint captured before a fact existed has to stay restorable, and a +// target the agent could not read has to be given the benefit of the doubt. +func TestCompareIgnoresUnknownFacts(t *testing.T) { + tests := []struct { + name string + source Facts + target Facts + }{ + { + name: "neither side knows anything", + }, + { + name: "the checkpoint recorded facts the target cannot describe", + source: populatedFacts(), + }, + { + name: "the target describes facts the checkpoint never recorded", + target: populatedFacts(), + }, + { + name: "both sides agree", + source: populatedFacts(), + target: populatedFacts(), + }, + } + + for _, gate := range []Gate{GatePreflight, GateInspect} { + for _, tc := range tests { + t.Run(string(gate)+" "+tc.name, func(t *testing.T) { + for _, mismatch := range Compare(gate, tc.source, tc.target) { + if deliberatelyNotSilent[mismatch.Check] { + continue + } + t.Errorf("Compare(%q) reported %+v, want no mismatches", gate, mismatch) + } + }) + } + } +} + +// Every registered rule reports itself, or a refusal names an empty check and +// nobody can tell which rule turned the restore down. Registration has already +// rejected a rule at a gate nothing calls, so the name is what is left to check. +func TestEveryCheckIsNamedAndRegisteredOnce(t *testing.T) { + seen := map[Check]bool{} + for gate, registered := range checksByGate { + for _, c := range registered { + if c.name == "" { + t.Errorf("the %q gate holds a rule with no name", gate) + } + if seen[c.name] { + t.Errorf("check %q is registered twice", c.name) + } + seen[c.name] = true + } + } +} + +// Compare has to attribute a mismatch to the rule that found it, since the whole +// refusal vocabulary is built on the check name. +func TestCompareNamesTheFailingCheck(t *testing.T) { + mismatches := Compare(GatePreflight, populatedFacts(), differentFacts()) + + if len(mismatches) == 0 { + t.Fatal("Compare found nothing wrong between two entirely different machines") + } + for _, mismatch := range mismatches { + if mismatch.Check == "" { + t.Errorf("mismatch %+v does not name the check that reported it", mismatch) + } + } +} + +// withChecks swaps the policy table for the length of one test, so what the +// table itself does can be pinned with a rule made up here rather than with +// whichever real rules happen to be registered. +func withChecks(t *testing.T, checks ...check) { + t.Helper() + saved := checksByGate + checksByGate = registerChecks(checks...) + t.Cleanup(func() { checksByGate = saved }) +} + +// A rule runs at its own gate and at no other, and what it reports comes back +// named after it. The two gates read different facts, so a rule that ran at the +// wrong one would compare against facts nobody had gathered yet. +func TestCompareRunsTheRulesOfOneGate(t *testing.T) { + archCheck := check{ + name: "fixture", + gate: GatePreflight, + compare: func(source, target Facts) []Mismatch { + return []Mismatch{{Source: source.CPUArch, Target: target.CPUArch}} + }, + } + + t.Run("at its own gate", func(t *testing.T) { + withChecks(t, archCheck) + + want := Mismatch{Check: "fixture", Source: "amd64", Target: "arm64"} + got := Compare(GatePreflight, populatedFacts(), differentFacts()) + if len(got) != 1 || got[0] != want { + t.Fatalf("Compare at the preflight gate = %v, want exactly %v", got, want) + } + }) + + t.Run("and nowhere else", func(t *testing.T) { + withChecks(t, archCheck) + + if got := Compare(GateInspect, populatedFacts(), differentFacts()); len(got) != 0 { + t.Fatalf("Compare at the inspect gate = %v, want no mismatches", got) + } + }) + + t.Run("a table with no rules refuses nothing", func(t *testing.T) { + withChecks(t) + + for _, gate := range []Gate{GatePreflight, GateInspect} { + if got := Compare(gate, populatedFacts(), differentFacts()); len(got) != 0 { + t.Fatalf("Compare at the %q gate = %v, want no mismatches", gate, got) + } + } + }) +} + +// A rule pinned to a gate nothing calls would never run, and nothing would say +// so. Registration refuses to build such a table at all. +func TestRegisterChecksRejectsAGateNothingCalls(t *testing.T) { + defer func() { + if recover() == nil { + t.Fatal("registerChecks built a table holding a rule at a gate nothing calls") + } + }() + + registerChecks(check{name: "fixture", gate: "nowhere"}) +} + +// A refusal has to survive the trip back to the caller that reports it: it +// crosses the restore call chain and is wrapped on the way, and the reader still +// has to tell it apart from a CRIU failure and learn which gate produced it. +func TestNewIncompatibleError(t *testing.T) { + mismatches := []Mismatch{ + {Check: "cpu-arch", Source: "amd64", Target: "arm64"}, + {Check: "memory-limit", Source: "32Gi", Target: "1Gi"}, + } + err := NewIncompatibleError(GateInspect, mismatches) + + var incompatible *IncompatibleError + if !errors.As(fmt.Errorf("restore worker: %w", err), &incompatible) { + t.Fatal("wrapped incompatible error did not unwrap to *IncompatibleError") + } + if incompatible.Gate != GateInspect { + t.Fatalf("gate = %q, want %q", incompatible.Gate, GateInspect) + } + want := "restore refused as incompatible: cpu-arch: source amd64, target arm64; memory-limit: source 32Gi, target 1Gi" + if got := err.Error(); got != want { + t.Fatalf("error = %q, want %q", got, want) + } + + // The caller's slice keeps changing after the refusal is built, and the + // refusal is what gets reported. + mismatches[0].Target = "mutated" + if incompatible.Mismatches[0].Target != "arm64" { + t.Fatalf("error kept a reference to the caller's slice: %#v", incompatible.Mismatches) + } +} diff --git a/api/compat/reason.go b/api/compat/reason.go new file mode 100644 index 00000000..e5881e7e --- /dev/null +++ b/api/compat/reason.go @@ -0,0 +1,33 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package compat + +import ( + "fmt" + "strings" + + "github.com/ai-dynamo/snapshot/api/util" +) + +// reasonSeparator joins the reasons of a refusal that failed several rules. +const reasonSeparator = "; " + +// Reason renders a mismatch as the sentence a user reads. The log field, the pod +// event and the pod annotation all carry this exact string, so an operator can +// match on one and find the others. +func (m Mismatch) Reason() string { + return fmt.Sprintf("%s: source %s, target %s", m.Check, util.OrUnknown(m.Source), util.OrUnknown(m.Target)) +} + +// Reasons renders every mismatch of one refusal in report order. +func Reasons(mismatches []Mismatch) string { + if len(mismatches) == 0 { + return "" + } + reasons := make([]string, 0, len(mismatches)) + for _, mismatch := range mismatches { + reasons = append(reasons, mismatch.Reason()) + } + return strings.Join(reasons, reasonSeparator) +} diff --git a/api/compat/reason_test.go b/api/compat/reason_test.go new file mode 100644 index 00000000..74af25c9 --- /dev/null +++ b/api/compat/reason_test.go @@ -0,0 +1,79 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package compat + +import "testing" + +// Asserted literally: every rule added later inherits this shape rather than +// inventing its own, and the three report surfaces compare against it. +func TestMismatchReason(t *testing.T) { + tests := []struct { + name string + mismatch Mismatch + want string + }{ + { + name: "both values known", + mismatch: Mismatch{Check: "memory-limit", Source: "32Gi", Target: "1Gi"}, + want: "memory-limit: source 32Gi, target 1Gi", + }, + { + name: "source unrecorded", + mismatch: Mismatch{Check: CheckImageDigest, Target: "sha256:beef"}, + want: "image-digest: source unknown, target sha256:beef", + }, + { + name: "target unreadable", + mismatch: Mismatch{Check: CheckDriverVersion, Source: "580.82.07"}, + want: "driver-version: source 580.82.07, target unknown", + }, + { + name: "blank values are unknown", + mismatch: Mismatch{Check: "kernel-version", Source: " ", Target: ""}, + want: "kernel-version: source unknown, target unknown", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := tc.mismatch.Reason(); got != tc.want { + t.Fatalf("Reason() = %q, want %q", got, tc.want) + } + }) + } +} + +func TestReasons(t *testing.T) { + tests := []struct { + name string + mismatches []Mismatch + want string + }{ + { + name: "no mismatches", + want: "", + }, + { + name: "one mismatch", + mismatches: []Mismatch{{Check: "cpu-arch", Source: "amd64", Target: "arm64"}}, + want: "cpu-arch: source amd64, target arm64", + }, + { + name: "report order is preserved", + mismatches: []Mismatch{ + {Check: "cpu-arch", Source: "amd64", Target: "arm64"}, + {Check: "memory-limit", Source: "32Gi", Target: "1Gi"}, + }, + want: "cpu-arch: source amd64, target arm64; memory-limit: source 32Gi, target 1Gi", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := Reasons(tc.mismatches); got != tc.want { + t.Fatalf("Reasons() = %q, want %q", got, tc.want) + } + }) + } +} diff --git a/api/util/util.go b/api/util/util.go new file mode 100644 index 00000000..6f621ac6 --- /dev/null +++ b/api/util/util.go @@ -0,0 +1,22 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Package util holds the small helpers shared across the api module, so the +// components that import it describe the same thing the same way. +package util + +import "strings" + +// unknownValue stands in for a value nobody recorded. One word for the idea +// everywhere, so a reader never has to work out whether "unset" and "n/a" mean +// the same thing. +const unknownValue = "unknown" + +// OrUnknown renders a value that may be absent. Blank is absent rather than +// empty, so a reason never trails off into a dangling comma. +func OrUnknown(value string) string { + if strings.TrimSpace(value) == "" { + return unknownValue + } + return value +} diff --git a/api/util/util_test.go b/api/util/util_test.go new file mode 100644 index 00000000..431069d1 --- /dev/null +++ b/api/util/util_test.go @@ -0,0 +1,28 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package util + +import "testing" + +func TestOrUnknown(t *testing.T) { + tests := []struct { + name string + value string + want string + }{ + {name: "a recorded value is its own", value: "580.82.07", want: "580.82.07"}, + {name: "an empty value is unknown", value: "", want: "unknown"}, + {name: "a blank value is unknown", value: " ", want: "unknown"}, + {name: "a tab is unknown", value: "\t", want: "unknown"}, + {name: "surrounding space is kept", value: " 5.15.0 ", want: " 5.15.0 "}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := OrUnknown(tc.value); got != tc.want { + t.Errorf("OrUnknown(%q) = %q, want %q", tc.value, got, tc.want) + } + }) + } +} diff --git a/api/v1alpha1/constants.go b/api/v1alpha1/constants.go index 2d2c32ac..91d97d90 100644 --- a/api/v1alpha1/constants.go +++ b/api/v1alpha1/constants.go @@ -31,6 +31,12 @@ const ( // RestoredCondition is the Pod status condition owned by the node agent. RestoredCondition = "nvidia.com/Restored" + // SkipCompatCheckAnnotation turns the restore compatibility gate off for + // one restore pod, whether it is stamped on the pod template up front or + // added to a pod the gate already turned down. It is the per-restore escape + // hatch; restore.skipCompatCheck in the agent ConfigMap is the per-node one. + SkipCompatCheckAnnotation = "nvidia.com/snapshot-skip-compat-check" + CheckpointVolumeName = "checkpoint-storage" DefaultSeccompLocalhostProfile = "profiles/block-iouring.json" ) diff --git a/api/v1alpha1/protocol.go b/api/v1alpha1/protocol.go index 26c414bb..a846488a 100644 --- a/api/v1alpha1/protocol.go +++ b/api/v1alpha1/protocol.go @@ -5,6 +5,7 @@ package v1alpha1 import ( "fmt" + "strconv" "strings" "k8s.io/apimachinery/pkg/util/validation" @@ -84,3 +85,11 @@ func ValidateRestoreContainerMappings(mappings []RestoreContainerMapping, captur } return nil } + +// SkipCompatCheckFromAnnotations reports whether this pod asks for the restore +// compatibility gate to be skipped. Anything that is not a recognized true +// keeps the gate on, so a typo cannot silently disable it. +func SkipCompatCheckFromAnnotations(annotations map[string]string) bool { + skip, err := strconv.ParseBool(strings.TrimSpace(annotations[SkipCompatCheckAnnotation])) + return err == nil && skip +} diff --git a/api/v1alpha1/protocol_test.go b/api/v1alpha1/protocol_test.go index 81ac0161..d88bd8dd 100644 --- a/api/v1alpha1/protocol_test.go +++ b/api/v1alpha1/protocol_test.go @@ -107,3 +107,31 @@ func TestValidateRestoreContainerMappings(t *testing.T) { }) } } + +func TestSkipCompatCheckFromAnnotations(t *testing.T) { + cases := map[string]bool{ + "true": true, + "True": true, + "1": true, + "false": false, + "0": false, + " true": true, + // A value nobody parses as a boolean leaves the gate on. Turning it off + // by accident is the expensive direction: a restore that should have + // been refused instead fails somewhere inside CRIU. + "yes": false, + "": false, + "TRUE-ISH": false, + "true please": false, + } + for value, want := range cases { + annotations := map[string]string{SkipCompatCheckAnnotation: value} + if got := SkipCompatCheckFromAnnotations(annotations); got != want { + t.Errorf("SkipCompatCheckFromAnnotations(%q) = %v, want %v", value, got, want) + } + } + + if SkipCompatCheckFromAnnotations(nil) { + t.Error("an unannotated pod asked to skip the compatibility gate") + } +} diff --git a/charts/snapshot/templates/configmap.yaml b/charts/snapshot/templates/configmap.yaml index 5f5f9983..e104b6bb 100644 --- a/charts/snapshot/templates/configmap.yaml +++ b/charts/snapshot/templates/configmap.yaml @@ -27,6 +27,7 @@ data: restore: restoreTimeoutSeconds: {{ .Values.config.restore.restoreTimeoutSeconds }} + skipCompatCheck: {{ .Values.config.restore.skipCompatCheck }} criu: binaryPath: {{ .Values.config.criu.binaryPath | quote }} diff --git a/charts/snapshot/tests/config_test.yaml b/charts/snapshot/tests/config_test.yaml index 49e2603d..9b669878 100644 --- a/charts/snapshot/tests/config_test.yaml +++ b/charts/snapshot/tests/config_test.yaml @@ -11,3 +11,17 @@ tests: asserts: - failedTemplate: errorPattern: "storage.pvc.basePath is fixed at /checkpoints" + + - it: keeps the restore compatibility gate on by default + asserts: + - matchRegex: + path: data["config.yaml"] + pattern: "skipCompatCheck: false" + + - it: renders the node-wide compatibility gate opt-out + set: + config.restore.skipCompatCheck: true + asserts: + - matchRegex: + path: data["config.yaml"] + pattern: "skipCompatCheck: true" diff --git a/charts/snapshot/tests/role_test.yaml b/charts/snapshot/tests/role_test.yaml new file mode 100644 index 00000000..1ea59c18 --- /dev/null +++ b/charts/snapshot/tests/role_test.yaml @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +suite: snapshot agent permissions +templates: + - templates/role.yaml +tests: + - it: lets the agent publish a refused restore on the pod status + documentSelector: + path: metadata.name + value: RELEASE-NAME-snapshot-agent + asserts: + - contains: + path: rules + content: + apiGroups: [""] + resources: ["pods/status"] + verbs: ["patch"] diff --git a/charts/snapshot/values.yaml b/charts/snapshot/values.yaml index c1d89ca7..e3228969 100644 --- a/charts/snapshot/values.yaml +++ b/charts/snapshot/values.yaml @@ -177,6 +177,13 @@ config: restore: # Maximum seconds to allow a restore attempt before the agent marks it failed restoreTimeoutSeconds: 7200 + # Skip the restore compatibility check on every restore this agent handles. + # The gate refuses a restore whose checkpoint was captured somewhere the + # target node cannot reproduce; turning it off means those restores are + # attempted and fail inside CRIU instead. Per-node escape hatch; a single + # restore can opt out with the nvidia.com/snapshot-skip-compat-check + # annotation on the restore pod. + skipCompatCheck: false criu: # Path to the criu binary diff --git a/e2e/snapshot_e2e/k8s.py b/e2e/snapshot_e2e/k8s.py index aec6306e..0b1e99cd 100644 --- a/e2e/snapshot_e2e/k8s.py +++ b/e2e/snapshot_e2e/k8s.py @@ -60,6 +60,18 @@ def read_crd(name: str) -> client.V1CustomResourceDefinition: return client.ApiextensionsV1Api().read_custom_resource_definition(name) +def read_config_map(namespace: str, name: str) -> client.V1ConfigMap: + return client.CoreV1Api().read_namespaced_config_map(name, namespace) + + +def patch_config_map(namespace: str, name: str, data: dict[str, str]) -> client.V1ConfigMap: + return client.CoreV1Api().patch_namespaced_config_map(name, namespace, {"data": data}) + + +def read_node(name: str) -> client.V1Node: + return client.CoreV1Api().read_node(name) + + def list_events(namespace: str) -> list[client.CoreV1Event]: return client.CoreV1Api().list_namespaced_event(namespace).items @@ -167,6 +179,23 @@ def exec_command( ) +PAYLOAD_MARKER = "e2e-payload-follows" + + +def exec_payload(namespace: str, pod: str, command: str) -> str: + """Exec output with whatever the login shell printed first dropped. + + exec_command merges stderr into the stream, so a container whose profile + writes anything breaks every caller that parses the result rather than + matching a substring in it. + """ + output = exec_command(namespace, pod, f"echo {PAYLOAD_MARKER}; {command}") + _, marker, payload = output.partition(PAYLOAD_MARKER) + if not marker: + raise AssertionError(f"exec output carried no payload marker: {output!r}") + return payload.lstrip("\n") + + def snapshot_custom_resource_api_is_accessible(namespace: str) -> None: api = client.CustomObjectsApi() api.list_namespaced_custom_object( diff --git a/e2e/snapshot_e2e/lifecycle.py b/e2e/snapshot_e2e/lifecycle.py index 1603cf88..5113d2db 100644 --- a/e2e/snapshot_e2e/lifecycle.py +++ b/e2e/snapshot_e2e/lifecycle.py @@ -5,10 +5,13 @@ from __future__ import annotations +import json import shlex import time -from typing import Any, Callable +from contextlib import contextmanager +from typing import Any, Callable, Iterator +import yaml from kubernetes import client from kubernetes.client import ApiException @@ -28,6 +31,7 @@ GROUP = "nvidia.com" VERSION = "v1alpha1" +RESTORED_CONDITION = f"{GROUP}/Restored" PODSNAPSHOTS = "podsnapshots" PODSNAPSHOTCONTENTS = "podsnapshotcontents" SNAPSHOTJOBS = "snapshotjobs" @@ -151,21 +155,24 @@ def detail() -> str: return wait_for(f"pod {namespace}/{name} Ready", ready, timeout, detail=detail) +def file_present(namespace: str, pod: str, path: str) -> bool: + # Require a stdout marker because exec does not expose remote exit status, + # and look for it rather than match on it: exec returns stderr too, and a + # login shell is free to write to it. + marker = "__snapshot_e2e_file_present__" + command = f"[[ -f {shlex.quote(path)} ]] && printf '%s' {shlex.quote(marker)}" + return marker in k8s.exec_command(namespace, pod, command) + + def wait_for_file(namespace: str, pod: str, path: str, timeout: int = 180) -> None: last_error: str | None = None - marker = "__snapshot_e2e_file_present__" def exists() -> bool | None: nonlocal last_error try: - # Require a stdout marker because exec does not expose remote exit status. - command = ( - f"[[ -f {shlex.quote(path)} ]] && " - f"printf '%s' {shlex.quote(marker)}" - ) - response = k8s.exec_command(namespace, pod, command) + present = file_present(namespace, pod, path) last_error = None - return True if response == marker else None + return True if present else None except Exception as exc: last_error = f"{type(exc).__name__}: {exc}" return None @@ -387,7 +394,7 @@ def wait_for_restored_condition( ) -> client.V1Pod: def check() -> client.V1Pod | None: pod = k8s.read_pod(namespace, pod_name) - restored = pod_condition(pod, "nvidia.com/Restored") + restored = pod_condition(pod, RESTORED_CONDITION) if restored and restored.status == status and restored.reason == reason: return pod terminal_reasons = {"RestoreSucceeded", "RestorePartiallySucceeded", "RestoreFailed"} @@ -403,8 +410,8 @@ def detail() -> str: pod = k8s.read_pod(namespace, pod_name) except ApiException as exc: return f"api_error={k8s.api_error_detail(exc)}" - restored = pod_condition(pod, "nvidia.com/Restored") - return f"nvidia.com/Restored={restored or ''}" + restored = pod_condition(pod, RESTORED_CONDITION) + return f"{RESTORED_CONDITION}={restored or ''}" return wait_for( f"nvidia.com/Restored={status}/{reason} on {namespace}/{pod_name}", @@ -421,16 +428,89 @@ def pod_condition(pod: client.V1Pod, condition_type: str) -> client.V1PodConditi return None +def wait_for_restore_past_the_gate( + namespace: str, + pod_name: str, + timeout: int = 600, +) -> client.V1Pod: + """Wait for any reason the agent only records once the gate has let the restore through. + + RestoreInProgress is transient, so waiting for it alone is a race a fast + restore wins. A restore refused at the gate never reaches any of these. + """ + past = ("RestoreInProgress", "RestoreSucceeded", "RestoreFailed") + + def check() -> client.V1Pod | None: + pod = k8s.read_pod(namespace, pod_name) + restored = pod_condition(pod, RESTORED_CONDITION) + return pod if restored and restored.reason in past else None + + def detail() -> str: + try: + pod = k8s.read_pod(namespace, pod_name) + except ApiException as exc: + return f"api_error={k8s.api_error_detail(exc)}" + restored = pod_condition(pod, RESTORED_CONDITION) + return f"{RESTORED_CONDITION}={restored.reason if restored else ''}" + + return wait_for( + f"restore past the gate on {namespace}/{pod_name}", + check, + timeout, + detail=detail, + ) + + def checkpoint_artifact_manifest( config: k8s.E2EConfig, node: str, content_uid: str ) -> str: - return k8s.exec_command( + return k8s.exec_payload( config.namespace, checkpoint_agent_pod(config, node), f"cat {checkpoint_artifact_path(content_uid)}/manifest.yaml", ) +def checkpoint_manifest( + config: k8s.E2EConfig, node: str, content_uid: str +) -> dict[str, Any]: + """The manifest as the agent will read it back, rather than as text.""" + return yaml.safe_load(checkpoint_artifact_manifest(config, node, content_uid)) + + +def runtime_image_id(config: k8s.E2EConfig, node: str, container_id: str) -> str: + runtime_id = container_id.split("://", 1)[-1] + output = k8s.exec_payload( + config.namespace, + checkpoint_agent_pod(config, node), + f"nsenter -t 1 -m -- crictl inspect {shlex.quote(runtime_id)}", + ) + image_id = (json.loads(output).get("status") or {}).get("imageId") + if not image_id: + raise AssertionError(f"runtime reported no image ID for {container_id!r}") + return image_id + + +def visible_gpus(namespace: str, pod: str) -> list[dict[str, str]]: + """The GPUs a pod can see, as nvidia-smi inside that pod reports them. + + The same query the agent runs, so a test comparing the two is comparing what + the machine says against what the artifact recorded, not two spellings of it. + """ + output = k8s.exec_payload( + namespace, + pod, + "nvidia-smi --query-gpu=gpu_uuid,name,driver_version --format=csv,noheader", + ) + gpus = [] + for line in output.strip().splitlines(): + fields = [field.strip() for field in line.split(",")] + if len(fields) != 3: + raise AssertionError(f"unexpected nvidia-smi row {line!r}") + gpus.append({"uuid": fields[0], "name": fields[1], "driver": fields[2]}) + return gpus + + def checkpoint_artifact_listing( config: k8s.E2EConfig, node: str, content_uid: str ) -> str: @@ -479,6 +559,85 @@ def checkpoint_agent_pod(config: k8s.E2EConfig, node: str) -> str: return agents[0].metadata.name +AGENT_CONFIG_VOLUME = "config" +AGENT_CONFIG_KEY = "config.yaml" + + +def agent_config_source(config: k8s.E2EConfig) -> tuple[str, str]: + """The ConfigMap the agent reads from and the path it reads it at. + + Read off the DaemonSet so a test edits the file the agent is actually + mounting rather than one the chart happens to name the same way. + """ + daemonsets = k8s.list_snapshot_daemonsets( + config.namespace, config.release, "snapshot-agent" + ) + if len(daemonsets) != 1: + raise AssertionError(f"expected one snapshot agent DaemonSet, found {len(daemonsets)}") + + spec = daemonsets[0].spec.template.spec + volume = next(v for v in spec.volumes if v.name == AGENT_CONFIG_VOLUME) + mount = next( + m + for container in spec.containers + for m in container.volume_mounts or [] + if m.name == AGENT_CONFIG_VOLUME + ) + return volume.config_map.name, f"{mount.mount_path}/{AGENT_CONFIG_KEY}" + + +def wait_for_agent_config(config: k8s.E2EConfig, expected: str, timeout: int = 180) -> None: + """Wait for a ConfigMap edit to reach every agent in the release. + + The ConfigMap belongs to the release, so an edit reaches the whole + DaemonSet; waiting on one node would leave the others holding the old value. + The kubelet refreshes projected ConfigMaps on its own schedule, so the file + inside each container is the only honest signal that an edit has landed. + """ + _, path = agent_config_source(config) + command = f"cat {shlex.quote(path)}" + agents = [ + pod.metadata.name + for pod in k8s.list_snapshot_pods(config.namespace, config.release, "snapshot-agent") + ] + if not agents: + raise AssertionError("no snapshot agent pods to wait on") + + for agent in agents: + def projected(agent: str = agent) -> bool | None: + content = k8s.exec_command(config.namespace, agent, command) + return True if expected in content else None + + wait_for( + f"{expected!r} in {agent}:{path}", + projected, + timeout, + detail=lambda agent=agent: k8s.exec_command(config.namespace, agent, command), + ) + + +@contextmanager +def node_skip_compat_check(config: k8s.E2EConfig) -> Iterator[None]: + """Turn the node switch on for the body, and put it back afterwards. + + Waits for the projection on both edges: leaving the switch on would let a + later test's restore through the very gate it means to exercise. + """ + name, _ = agent_config_source(config) + original = k8s.read_config_map(config.namespace, name).data[AGENT_CONFIG_KEY] + off, on = "skipCompatCheck: false", "skipCompatCheck: true" + if off not in original: + raise AssertionError(f"{name}:{AGENT_CONFIG_KEY} does not carry {off!r}") + + k8s.patch_config_map(config.namespace, name, {AGENT_CONFIG_KEY: original.replace(off, on)}) + try: + wait_for_agent_config(config, on) + yield + finally: + k8s.patch_config_map(config.namespace, name, {AGENT_CONFIG_KEY: original}) + wait_for_agent_config(config, off) + + def assert_restored_state( namespace: str, pod: str, @@ -525,6 +684,15 @@ def debug_dump(config: k8s.E2EConfig, run: TestRun) -> None: for pod in pods: print(f"pod {pod.metadata.name} phase={pod.status.phase} node={pod.spec.node_name}") print(f"annotations={pod.metadata.annotations or {}}") + print( + "conditions=" + + str( + [ + (c.type, c.status, c.reason, c.message) + for c in pod.status.conditions or [] + ] + ) + ) print(k8s.pod_logs(config.namespace, pod.metadata.name, tail_lines=80)) print_custom_objects(config, run) print_snapshot_controller_logs(config) diff --git a/e2e/snapshot_e2e/workloads.py b/e2e/snapshot_e2e/workloads.py index 855fedd3..83279ae7 100644 --- a/e2e/snapshot_e2e/workloads.py +++ b/e2e/snapshot_e2e/workloads.py @@ -68,6 +68,7 @@ def source_pod( run: TestRun, gpu: bool, annotations: dict[str, str] | None = None, + memory_limit: str | None = None, ) -> dict[str, Any]: metadata = { "name": run.source_pod, @@ -75,7 +76,7 @@ def source_pod( "labels": run.labels, "annotations": annotations or {}, } - spec = base_pod_spec(config, run, source_command(run.image, gpu), gpu) + spec = base_pod_spec(config, run, source_command(run.image, gpu), gpu, memory_limit) spec["containers"][0]["env"] = [ {"name": SOURCE_TOKEN_ENV, "value": run.source_token}, ] @@ -132,11 +133,12 @@ def restore_pod( gpu: bool, source_node: str | None = None, snapshot_name: str | None = None, + memory_limit: str | None = None, ) -> dict[str, Any]: # snapshot_name defaults to the lifecycle flow's test-created PodSnapshot # (run.snapshot_name); a SnapshotJob-produced PodSnapshot is named after # the SnapshotJob instead, so those tests pass it explicitly. - spec = base_pod_spec(config, run, restore_command(run.image, gpu), gpu) + spec = base_pod_spec(config, run, restore_command(run.image, gpu), gpu, memory_limit) spec["securityContext"] = { "seccompProfile": { "type": "Localhost", @@ -238,6 +240,7 @@ def base_pod_spec( run: TestRun, command: str, gpu: bool, + memory_limit: str | None = None, *, control_volume: bool = True, checkpoint_pvc: bool = True, @@ -283,6 +286,9 @@ def base_pod_spec( if gpu: spec["runtimeClassName"] = "nvidia" container["resources"] = {"limits": {"nvidia.com/gpu": "1"}} + if memory_limit: + limits = container.setdefault("resources", {}).setdefault("limits", {}) + limits["memory"] = memory_limit return spec diff --git a/e2e/tests/test_snapshot_lifecycle.py b/e2e/tests/test_snapshot_lifecycle.py index a1070913..437cf848 100644 --- a/e2e/tests/test_snapshot_lifecycle.py +++ b/e2e/tests/test_snapshot_lifecycle.py @@ -3,6 +3,8 @@ from __future__ import annotations +import time + import pytest from snapshot_e2e import k8s @@ -72,6 +74,84 @@ def test_successful_snapshot_captures_cpu_gpu_and_fs( raise +# The value does not matter, only that one is set: a limit neither side records +# compares equal to itself and proves nothing. +RECORDED_MEMORY_LIMIT = "4Gi" + + +@pytest.mark.snapshot_success +@pytest.mark.gpu +def test_snapshot_records_the_facts_a_restore_is_checked_against( + config: k8s.E2EConfig, + run: snap.TestRun, +) -> None: + """The recorded facts have to be the machine's, not merely present. + + Everything the compatibility gates decide on is read at capture and can + never be recovered afterwards, so this compares each recorded fact against + the node object and against nvidia-smi inside the pod that was captured. + """ + try: + source, source_node = create_ready_source( + config, run, gpu=True, memory_limit=RECORDED_MEMORY_LIMIT + ) + snap.wait_for_state_observations( + config.namespace, + run.source_pod, + run.source_token, + gpu=True, + minimum=2, + ) + # Read before the capture, while the source container is still running. + visible_gpus = snap.visible_gpus(config.namespace, run.source_pod) + assert visible_gpus, "the GPU workload could not see a GPU" + source_status = next( + status + for status in source.status.container_statuses + if status.name == snap.CONTAINER + ) + source_image_id = snap.runtime_image_id( + config, source_node, source_status.container_id + ) + + snap.create_podsnapshot( + config.namespace, + run.snapshot_name, + run.source_pod, + source.metadata.uid, + ) + _, content = snap.wait_for_snapshot_ready(config.namespace, run.snapshot_name) + manifest = snap.checkpoint_manifest( + config, source_node, content["metadata"]["uid"] + ) + + node_info = k8s.read_node(source_node).status.node_info + host = manifest["host"] + assert host["kernelVersion"] == node_info.kernel_version + assert host["cpuArch"] == node_info.architecture + + pod = k8s.read_pod(config.namespace, run.source_pod) + container = next(c for c in pod.spec.containers if c.name == snap.CONTAINER) + limits = (container.resources.limits or {}) if container.resources else {} + recorded_pod = manifest["k8s"] + assert recorded_pod["image"] == container.image + assert recorded_pod["imageId"] == source_image_id + assert recorded_pod["memoryLimit"] == limits["memory"] + # This pod sets no CPU limit, and an absent fact is recorded as absent + # rather than invented, which is what makes it refuse nothing later. + assert "cpu" not in limits + assert "cpuLimit" not in recorded_pod + + cuda = manifest["cudaRestore"] + assert sorted( + (gpu["uuid"], gpu["productName"]) for gpu in cuda["sourceGpus"] + ) == sorted((gpu["uuid"], gpu["name"]) for gpu in visible_gpus) + assert cuda["sourceDriverVersion"] == visible_gpus[0]["driver"] + except Exception: + snap.debug_dump(config, run) + raise + + @pytest.mark.snapshot_success @pytest.mark.gpu def test_successful_restore_recovers_cpu_gpu_and_fs_from_snapshot( @@ -238,11 +318,154 @@ def test_failed_restore_gpu_checkpoint_into_non_gpu_target( raise +# A checkpoint captured with more memory than the target offers is the cheapest +# real mismatch to build: nothing about the node has to change for it. +CAPTURE_MEMORY_LIMIT = "4Gi" +SMALLER_MEMORY_LIMIT = "1Gi" + +# The restore informer resyncs every 30s, which is what would re-drive a refused +# pod. Waiting past two of them is how a retry loop would show itself. +RESTORE_RESYNC_SECONDS = 30 + + +@pytest.mark.snapshot_failure +@pytest.mark.gpu +def test_refused_restore_says_why_and_does_no_criu_work( + config: k8s.E2EConfig, + run: snap.TestRun, +) -> None: + try: + _, source_node, _ = create_valid_gpu_checkpoint( + config, run, memory_limit=CAPTURE_MEMORY_LIMIT + ) + k8s.delete_pod(config.namespace, run.source_pod) + snap.wait_for_pod_deleted(config.namespace, run.source_pod) + + k8s.create_pod( + snap.restore_pod( + config=config, + run=run, + gpu=True, + source_node=source_node, + memory_limit=SMALLER_MEMORY_LIMIT, + ) + ) + pod = snap.wait_for_restored_condition( + config.namespace, run.restore_pod, "False", "RestoreIncompatible" + ) + + refusal = snap.pod_condition(pod, snap.RESTORED_CONDITION) + assert "memory-limit" in refusal.message + assert CAPTURE_MEMORY_LIMIT in refusal.message + assert SMALLER_MEMORY_LIMIT in refusal.message + + assert_restore_events(config.namespace, run.restore_pod, {"RestoreIncompatible"}) + time.sleep(2 * RESTORE_RESYNC_SECONDS + 5) + assert restore_event_count(config.namespace, run.restore_pod, "RestoreIncompatible") == 1 + assert "RestoreFailed" not in restore_event_reasons(config.namespace, run.restore_pod) + + # The placeholder is still the placeholder: a refusal costs no CRIU work, + # so the workload never sees restore-complete. + assert not snap.file_present(config.namespace, run.restore_pod, snap.RESTORE_DONE) + except Exception: + snap.debug_dump(config, run) + raise + + +# Small enough to be refused, large enough to restore into once the checks are +# off: with a switch on, the restore these tests start actually runs. +SKIPPABLE_MEMORY_LIMIT = "3Gi" + + +@pytest.mark.snapshot_success +@pytest.mark.gpu +def test_skip_annotation_lets_a_refused_restore_through( + config: k8s.E2EConfig, + run: snap.TestRun, +) -> None: + try: + _, source_node, _ = create_valid_gpu_checkpoint( + config, run, memory_limit=CAPTURE_MEMORY_LIMIT + ) + k8s.delete_pod(config.namespace, run.source_pod) + snap.wait_for_pod_deleted(config.namespace, run.source_pod) + + body = snap.restore_pod( + config=config, + run=run, + gpu=True, + source_node=source_node, + memory_limit=SKIPPABLE_MEMORY_LIMIT, + ) + body["metadata"]["annotations"]["nvidia.com/snapshot-skip-compat-check"] = "true" + k8s.create_pod(body) + + pod = snap.wait_for_restore_past_the_gate(config.namespace, run.restore_pod) + assert snap.pod_condition(pod, snap.RESTORED_CONDITION).reason != "RestoreIncompatible" + assert "RestoreIncompatible" not in restore_event_reasons( + config.namespace, run.restore_pod + ) + except Exception: + snap.debug_dump(config, run) + raise + + +@pytest.mark.snapshot_success +@pytest.mark.gpu +def test_node_switch_lets_a_refused_restore_through_without_a_rollout( + config: k8s.E2EConfig, + run: snap.TestRun, +) -> None: + try: + _, source_node, _ = create_valid_gpu_checkpoint( + config, run, memory_limit=CAPTURE_MEMORY_LIMIT + ) + k8s.delete_pod(config.namespace, run.source_pod) + snap.wait_for_pod_deleted(config.namespace, run.source_pod) + + agent_before = k8s.read_pod( + config.namespace, snap.checkpoint_agent_pod(config, source_node) + ) + with snap.node_skip_compat_check(config): + k8s.create_pod( + snap.restore_pod( + config=config, + run=run, + gpu=True, + source_node=source_node, + memory_limit=SKIPPABLE_MEMORY_LIMIT, + ) + ) + pod = snap.wait_for_restore_past_the_gate(config.namespace, run.restore_pod) + assert snap.pod_condition(pod, snap.RESTORED_CONDITION).reason != "RestoreIncompatible" + assert "RestoreIncompatible" not in restore_event_reasons( + config.namespace, run.restore_pod + ) + + # The switch is worth having as a ConfigMap rather than an env var only + # if flipping it costs nothing, so the agent that honoured it has to be + # the same process that was running before. + agent_after = k8s.read_pod( + config.namespace, snap.checkpoint_agent_pod(config, source_node) + ) + assert agent_after.metadata.uid == agent_before.metadata.uid + assert agent_restarts(agent_after) == agent_restarts(agent_before) + except Exception: + snap.debug_dump(config, run) + raise + + +def agent_restarts(pod: object) -> int: + return sum(status.restart_count for status in pod.status.container_statuses or []) + + def create_valid_gpu_checkpoint( config: k8s.E2EConfig, run: snap.TestRun, + *, + memory_limit: str | None = None, ) -> tuple[object, str, int]: - source, source_node = create_ready_source(config, run, gpu=True) + source, source_node = create_ready_source(config, run, gpu=True, memory_limit=memory_limit) checkpoint_observations = snap.wait_for_state_observations( config.namespace, run.source_pod, @@ -264,6 +487,7 @@ def create_ready_source( *, gpu: bool, annotations: dict[str, str] | None = None, + memory_limit: str | None = None, ) -> tuple[object, str]: k8s.create_pod( snap.source_pod( @@ -271,6 +495,7 @@ def create_ready_source( run=run, gpu=gpu, annotations=annotations, + memory_limit=memory_limit, ) ) pod = snap.wait_for_pod_ready(config.namespace, run.source_pod) @@ -342,3 +567,18 @@ def restore_event_reasons(namespace: str, pod_name: str) -> set[str]: for event in events if event.involved_object and event.involved_object.name == pod_name } + + +def restore_event_count(namespace: str, pod_name: str, reason: str) -> int: + """How many times the pod was told this, not how many objects say it. + + Repeated events are aggregated into one object with a count, so counting + objects would report one no matter how often the agent repeated itself. + """ + return sum( + event.count or 1 + for event in k8s.list_events(namespace) + if event.involved_object + and event.involved_object.name == pod_name + and event.reason == reason + )