feat: refuse incompatible restores before CRIU

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)

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")
+		}
+	})
+}

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")
 	}

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
+}