diff --git a/README.md b/README.md index 5a794c9d..ec0383d5 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,10 @@ Snapshots are portable across compatible machines and can be restored on any nod | `nvidia.com/restore-from` | Namespaced | Added as a pod annotation to trigger restore from a named `PodSnapshot` in the same namespace. | | `nvidia.com/restore-container-map` | Namespaced | Optional comma-separated `source=destination` mappings used to clone the single captured container into one or more restore containers. | +Restore producers must also implement the versioned +[restore Pod contract](docs/restore-pod-contract.md). Go integrations should +use the public builder and validator from `github.com/ai-dynamo/snapshot/api/v1alpha1`. +   diff --git a/agent/internal/controller/controller.go b/agent/internal/controller/controller.go index ddd3dca1..fbe0d98c 100644 --- a/agent/internal/controller/controller.go +++ b/agent/internal/controller/controller.go @@ -600,64 +600,21 @@ func validateRestoreTarget(pod *corev1.Pod, snapshot *snapshotv1alpha1.PodSnapsh if err != nil { return nil, nil, err } - if err := validateRestoreMappingsForPod(&pod.Spec, mappings, containerName); err != nil { + if err := snapshotv1alpha1.ValidateRestoreContainerMappings(mappings, containerName); err != nil { return nil, nil, err } - return &restoreTarget{SnapshotName: snapshot.Name, ContentUID: string(content.UID), SourceContainerName: containerName}, mappings, nil -} - -// validateRestoreMappingsForPod validates the mapping contract and ensures -// every destination exists in the restore Pod spec. -func validateRestoreMappingsForPod(spec *corev1.PodSpec, mappings []snapshotv1alpha1.RestoreContainerMapping, capturedSource string) error { - if err := snapshotv1alpha1.ValidateRestoreContainerMappings(mappings, capturedSource); err != nil { - return err - } - if len(mappings) > 1 { - hasControlVolume := false - for _, volume := range spec.Volumes { - if volume.Name == snapshotv1alpha1.SnapshotControlVolumeName && volume.EmptyDir != nil { - hasControlVolume = true - break - } - } - if !hasControlVolume { - return fmt.Errorf("multi-container restore requires %s emptyDir volume", snapshotv1alpha1.SnapshotControlVolumeName) - } - } - for _, mapping := range mappings { - var destination *corev1.Container - for i := range spec.Containers { - if spec.Containers[i].Name == mapping.Destination { - destination = &spec.Containers[i] - break - } - } - if destination == nil { - return fmt.Errorf("restore pod has no destination container named %q", mapping.Destination) - } - if len(mappings) == 1 { - continue - } - validControlMount := false - for _, mount := range destination.VolumeMounts { - if mount.Name == snapshotv1alpha1.SnapshotControlVolumeName && - mount.MountPath == snapshotv1alpha1.SnapshotControlMountPath && - mount.SubPath == mapping.Destination { - validControlMount = true - break - } - } - if !validControlMount { - return fmt.Errorf( - "multi-container restore destination %q requires %s mounted at %s with subPath %q", - mapping.Destination, - snapshotv1alpha1.SnapshotControlVolumeName, - snapshotv1alpha1.SnapshotControlMountPath, - mapping.Destination, - ) - } + if err := snapshotv1alpha1.ValidateRestorePod( + pod, + snapshot.Name, + mappings, + // Seccomp injection is optional producer policy. The node agent validates + // the universal restore contract and leaves profile selection to the + // producer that shaped the Pod. + snapshotv1alpha1.RestorePodOptions{}, + ); err != nil { + return nil, nil, err } - return nil + return &restoreTarget{SnapshotName: snapshot.Name, ContentUID: string(content.UID), SourceContainerName: containerName}, mappings, nil } // resolveRestoreArtifact resolves the validated restore target to its physical diff --git a/agent/internal/controller/controller_test.go b/agent/internal/controller/controller_test.go index 06a23ba7..9a4a8894 100644 --- a/agent/internal/controller/controller_test.go +++ b/agent/internal/controller/controller_test.go @@ -206,7 +206,7 @@ func pendingRestoreReason(t *testing.T, err error) string { } func restorePod(annotations map[string]string) *corev1.Pod { - return &corev1.Pod{ + pod := &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{ Name: "restore-worker", Namespace: "inference", @@ -231,6 +231,21 @@ func restorePod(annotations map[string]string) *corev1.Pod { }}, }, } + snapshotName, restoreRequested := annotations[snapshotv1alpha1.RestoreFromAnnotation] + _, hasExplicitMapping := annotations[snapshotv1alpha1.RestoreContainerMapAnnotation] + if !restoreRequested || hasExplicitMapping { + return pod + } + shaped, err := snapshotv1alpha1.BuildRestorePod( + pod, + snapshotName, + []snapshotv1alpha1.RestoreContainerMapping{{Source: "main", Destination: "main"}}, + snapshotv1alpha1.RestorePodOptions{}, + ) + if err != nil { + panic(err) + } + return shaped } func multiRestorePod() *corev1.Pod { @@ -264,7 +279,19 @@ func multiRestorePod() *corev1.Pod { {Name: "engine-0", ContainerID: "containerd://engine-0-id", State: corev1.ContainerState{Running: &corev1.ContainerStateRunning{}}}, {Name: "engine-1", ContainerID: "containerd://engine-1-id", State: corev1.ContainerState{Running: &corev1.ContainerStateRunning{}}}, } - return pod + shaped, err := snapshotv1alpha1.BuildRestorePod( + pod, + "snapshot-a", + []snapshotv1alpha1.RestoreContainerMapping{ + {Source: "main", Destination: "engine-0"}, + {Source: "main", Destination: "engine-1"}, + }, + snapshotv1alpha1.RestorePodOptions{}, + ) + if err != nil { + panic(err) + } + return shaped } func processQueuedRestorePod(t *testing.T, w *NodeController, pod *corev1.Pod) { @@ -571,6 +598,51 @@ func TestPreflightRestoreRejectsSharedMultiDestinationControlMount(t *testing.T) assert.Contains(t, err.Error(), `subPath "engine-1"`) } +func TestPreflightRestoreRejectsInvalidRestorePodContract(t *testing.T) { + snapshot, content := readySnapshotObjects() + tests := map[string]func(*corev1.Pod){ + "control volume": func(pod *corev1.Pod) { + pod.Spec.Volumes = nil + }, + "control environment": func(pod *corev1.Pod) { + pod.Spec.Containers[0].Env = nil + }, + "startup gate": func(pod *corev1.Pod) { + pod.Spec.Containers[0].StartupProbe = nil + }, + } + for name, mutate := range tests { + t.Run(name, func(t *testing.T) { + pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + mutate(pod) + w := makeTestController(t, pod, snapshot, content) + + plan, err := w.preflightRestore(context.Background(), pod) + + assert.Nil(t, plan) + require.Error(t, err) + }) + } +} + +func TestValidateRestoreTargetAcceptsEquivalentRestoreStartupGate(t *testing.T) { + pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod.Spec.Containers[0].StartupProbe = &corev1.Probe{ + ProbeHandler: corev1.ProbeHandler{Exec: &corev1.ExecAction{Command: []string{ + "test", "-f", snapshotv1alpha1.SnapshotControlMountPath + "/" + snapshotv1alpha1.RestoreCompleteFile, + }}}, + PeriodSeconds: 7, + FailureThreshold: 42, + } + snapshot, content := readySnapshotObjects() + + target, mappings, err := validateRestoreTarget(pod, snapshot, content) + + require.NoError(t, err) + assert.Equal(t, "snapshot-a", target.SnapshotName) + assert.Equal(t, []snapshotv1alpha1.RestoreContainerMapping{{Source: "main", Destination: "main"}}, mappings) +} + func TestPreflightRestoreRetriesInProgressCondition(t *testing.T) { pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{ diff --git a/api/v1alpha1/constants.go b/api/v1alpha1/constants.go index 2d2c32ac..78f569b5 100644 --- a/api/v1alpha1/constants.go +++ b/api/v1alpha1/constants.go @@ -57,16 +57,28 @@ const ( // control mount path to the workload. SnapshotControlDirEnv = "SNAPSHOT_CONTROL_DIR" - // LegacySnapshotControlDirEnv is the environment variable exposing the - // control mount path to the workload. EnsureControlVolume injects both - // during the migration window so existing workload images (which read - // this name) keep working while new images can move to + // LegacySnapshotControlDirEnv is the deprecated environment variable + // exposing the control mount path to the workload. Snapshot Pod shaping + // injects both names during the migration window so existing workload + // images keep working while new images can move to // SnapshotControlDirEnv. // // Deprecated: use SnapshotControlDirEnv instead. Remove once no workload // image depends on this name. LegacySnapshotControlDirEnv = "DYN_SNAPSHOT_CONTROL_DIR" + // RestoreStandbyModeEnv asks standby-aware workload entrypoints to remain + // inert until Snapshot replaces them with restored processes. Snapshot's + // generic Pod builder does not inject this workload-specific setting. + RestoreStandbyModeEnv = "SNAPSHOT_RESTORE_STANDBY" + + // LegacyRestoreStandbyModeEnv is the deprecated Dynamo restore standby + // environment variable. Snapshot publishes the name for producers that + // support existing workload images but does not inject it. + // + // Deprecated: use RestoreStandbyModeEnv for new workload integrations. + LegacyRestoreStandbyModeEnv = "DYN_SNAPSHOT_RESTORE_STANDBY" + // SnapshotCompleteFile named the sentinel the agent used to release a // checkpointed workload when leave-running dumps existed. A checkpoint now // always terminates the source process, so the agent no longer writes it; diff --git a/api/v1alpha1/protocol.go b/api/v1alpha1/protocol.go index 26c414bb..9f70e7c0 100644 --- a/api/v1alpha1/protocol.go +++ b/api/v1alpha1/protocol.go @@ -21,7 +21,11 @@ type RestoreContainerMapping struct { // GetRestoreFromSnapshotName returns the same-namespace PodSnapshot named by // the restore-from annotation. func GetRestoreFromSnapshotName(annotations map[string]string) (string, error) { - snapshotName := strings.TrimSpace(annotations[RestoreFromAnnotation]) + return validateRestoreFromSnapshotName(annotations[RestoreFromAnnotation]) +} + +func validateRestoreFromSnapshotName(value string) (string, error) { + snapshotName := strings.TrimSpace(value) if snapshotName == "" { return "", fmt.Errorf("%s must name a PodSnapshot", RestoreFromAnnotation) } diff --git a/api/v1alpha1/restore_pod.go b/api/v1alpha1/restore_pod.go new file mode 100644 index 00000000..46c57465 --- /dev/null +++ b/api/v1alpha1/restore_pod.go @@ -0,0 +1,509 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package v1alpha1 + +import ( + "fmt" + "path" + "reflect" + "strings" + + corev1 "k8s.io/api/core/v1" +) + +const restoreStartupFailureThreshold int32 = 1800 // 30 minutes at 1s cadence. + +// RestorePodOptions controls generic Snapshot restore Pod shaping. An empty +// SeccompProfile leaves the Pod's seccomp configuration unchanged. +// +kubebuilder:object:generate=false +type RestorePodOptions struct { + // SeccompProfile is the kubelet-local profile applied at Pod scope. Empty + // leaves seccomp configuration entirely caller-owned. + SeccompProfile string +} + +// BuildRestorePod returns a restore-shaped deep copy of pod. The caller's Pod +// is never mutated, including when validation fails. Workload-specific standby +// behavior and container commands remain the caller's responsibility. Mapping +// sources must come from the referenced PodSnapshot; this pure builder performs +// no API read to derive them. +func BuildRestorePod( + pod *corev1.Pod, + snapshotName string, + mappings []RestoreContainerMapping, + options RestorePodOptions, +) (*corev1.Pod, error) { + if pod == nil { + return nil, fmt.Errorf("restore pod is nil") + } + snapshotName, source, err := validateRestorePodRequest(snapshotName, mappings) + if err != nil { + return nil, err + } + result := pod.DeepCopy() + if err := ensureRestoreAnnotations(result, snapshotName, source, mappings); err != nil { + return nil, err + } + if err := ensureRestorePodSpec(&result.Spec, mappings, options); err != nil { + return nil, err + } + if err := validateCanonicalRestorePod(result, snapshotName, mappings, options); err != nil { + return nil, fmt.Errorf("validate shaped restore pod: %w", err) + } + return result, nil +} + +// ValidateRestorePod verifies that pod implements the declarative Snapshot +// restore runtime contract for snapshotName and mappings. It accepts supported +// restore-completion gates independently of their probe timing so Pods remain +// valid across builder versions. It performs no API reads and never mutates the +// Pod. +func ValidateRestorePod( + pod *corev1.Pod, + snapshotName string, + mappings []RestoreContainerMapping, + options RestorePodOptions, +) error { + return validateRestorePod(pod, snapshotName, mappings, options, validateRestoreStartupProbe) +} + +func validateCanonicalRestorePod( + pod *corev1.Pod, + snapshotName string, + mappings []RestoreContainerMapping, + options RestorePodOptions, +) error { + return validateRestorePod(pod, snapshotName, mappings, options, validateCanonicalRestoreStartupProbe) +} + +func validateRestorePod( + pod *corev1.Pod, + snapshotName string, + mappings []RestoreContainerMapping, + options RestorePodOptions, + validateStartupProbe func(*corev1.Container) error, +) error { + if pod == nil { + return fmt.Errorf("restore pod is nil") + } + snapshotName, source, err := validateRestorePodRequest(snapshotName, mappings) + if err != nil { + return err + } + if err := validateRestoreAnnotations(pod.Annotations, snapshotName, source, mappings); err != nil { + return err + } + if err := validateControlVolume(&pod.Spec); err != nil { + return err + } + for _, mapping := range mappings { + container := findContainer(&pod.Spec, mapping.Destination) + if container == nil { + return fmt.Errorf("restore pod has no destination container named %q", mapping.Destination) + } + if err := validateControlMount(container); err != nil { + return err + } + if err := validateControlEnvironment(container); err != nil { + return err + } + if err := validateStartupProbe(container); err != nil { + return err + } + if err := validateContainerSeccompProfile(container, options.SeccompProfile); err != nil { + return err + } + } + return validatePodSeccompProfile(&pod.Spec, options.SeccompProfile) +} + +func validateRestorePodRequest(snapshotName string, mappings []RestoreContainerMapping) (string, string, error) { + snapshotName, err := validateRestoreFromSnapshotName(snapshotName) + if err != nil { + return "", "", err + } + if len(mappings) == 0 { + return "", "", fmt.Errorf("restore container mapping must contain at least one destination") + } + source := mappings[0].Source + if err := ValidateRestoreContainerMappings(mappings, source); err != nil { + return "", "", err + } + return snapshotName, source, nil +} + +func ensureRestoreAnnotations(pod *corev1.Pod, snapshotName, source string, mappings []RestoreContainerMapping) error { + if pod.Annotations == nil { + pod.Annotations = make(map[string]string, 2) + } + if existing, found := pod.Annotations[RestoreFromAnnotation]; found { + resolved, err := validateRestoreFromSnapshotName(existing) + if err != nil { + return err + } + if resolved != snapshotName { + return fmt.Errorf("%s names %q, conflicting with requested PodSnapshot %q", RestoreFromAnnotation, resolved, snapshotName) + } + } + pod.Annotations[RestoreFromAnnotation] = snapshotName + + formatted, needsMapping := formatRestoreContainerMappings(mappings) + if _, found := pod.Annotations[RestoreContainerMapAnnotation]; found { + existingMappings, err := RestoreContainerMappingsFromAnnotations(pod.Annotations, source) + if err != nil { + return err + } + if err := ValidateRestoreContainerMappings(existingMappings, source); err != nil { + return err + } + if !sameRestoreContainerMappings(existingMappings, mappings) { + return fmt.Errorf("%s conflicts with requested restore mappings", RestoreContainerMapAnnotation) + } + if needsMapping { + pod.Annotations[RestoreContainerMapAnnotation] = formatted + } + return nil + } + if needsMapping { + pod.Annotations[RestoreContainerMapAnnotation] = formatted + } + return nil +} + +func validateRestoreAnnotations(annotations map[string]string, snapshotName, source string, mappings []RestoreContainerMapping) error { + resolved, err := GetRestoreFromSnapshotName(annotations) + if err != nil { + return err + } + if resolved != snapshotName { + return fmt.Errorf("%s names %q, expected %q", RestoreFromAnnotation, resolved, snapshotName) + } + annotatedMappings, err := RestoreContainerMappingsFromAnnotations(annotations, source) + if err != nil { + return err + } + if err := ValidateRestoreContainerMappings(annotatedMappings, source); err != nil { + return err + } + if !sameRestoreContainerMappings(annotatedMappings, mappings) { + return fmt.Errorf("%s does not match the requested restore mappings", RestoreContainerMapAnnotation) + } + return nil +} + +func formatRestoreContainerMappings(mappings []RestoreContainerMapping) (string, bool) { + if len(mappings) == 1 && mappings[0].Source == mappings[0].Destination { + return "", false + } + formatted := make([]string, 0, len(mappings)) + for _, mapping := range mappings { + formatted = append(formatted, mapping.Source+"="+mapping.Destination) + } + return strings.Join(formatted, ","), true +} + +func sameRestoreContainerMappings(left, right []RestoreContainerMapping) bool { + if len(left) != len(right) { + return false + } + want := make(map[RestoreContainerMapping]struct{}, len(right)) + for _, mapping := range right { + want[mapping] = struct{}{} + } + for _, mapping := range left { + if _, found := want[mapping]; !found { + return false + } + } + return true +} + +func ensureRestorePodSpec(spec *corev1.PodSpec, mappings []RestoreContainerMapping, options RestorePodOptions) error { + if err := ensureControlVolume(spec); err != nil { + return err + } + if err := ensurePodSeccompProfile(spec, options.SeccompProfile); err != nil { + return err + } + for _, mapping := range mappings { + container := findContainer(spec, mapping.Destination) + if container == nil { + return fmt.Errorf("restore pod has no destination container named %q", mapping.Destination) + } + if err := validateContainerSeccompProfile(container, options.SeccompProfile); err != nil { + return err + } + if err := ensureControlMount(container); err != nil { + return err + } + if err := ensureControlEnvironment(container); err != nil { + return err + } + ensureRestoreStartupProbe(container) + } + return nil +} + +func findContainer(spec *corev1.PodSpec, name string) *corev1.Container { + for i := range spec.Containers { + if spec.Containers[i].Name == name { + return &spec.Containers[i] + } + } + return nil +} + +func ensureControlVolume(spec *corev1.PodSpec) error { + found, err := hasValidControlVolume(spec) + if err != nil { + return err + } + if !found { + spec.Volumes = append(spec.Volumes, corev1.Volume{ + Name: SnapshotControlVolumeName, + VolumeSource: corev1.VolumeSource{EmptyDir: &corev1.EmptyDirVolumeSource{}}, + }) + } + return nil +} + +func validateControlVolume(spec *corev1.PodSpec) error { + found, err := hasValidControlVolume(spec) + if err != nil { + return err + } + if !found { + return fmt.Errorf("missing %s emptyDir volume", SnapshotControlVolumeName) + } + return nil +} + +func hasValidControlVolume(spec *corev1.PodSpec) (bool, error) { + found := false + for i := range spec.Volumes { + volume := &spec.Volumes[i] + if volume.Name != SnapshotControlVolumeName { + continue + } + if found { + return false, fmt.Errorf("duplicate %s volume", SnapshotControlVolumeName) + } + found = true + if !isEmptyDirVolumeSource(volume.VolumeSource) { + return false, fmt.Errorf("volume %q must be an emptyDir", SnapshotControlVolumeName) + } + } + return found, nil +} + +func isEmptyDirVolumeSource(source corev1.VolumeSource) bool { + return source.EmptyDir != nil && reflect.DeepEqual(source, corev1.VolumeSource{EmptyDir: source.EmptyDir}) +} + +func ensureControlMount(container *corev1.Container) error { + found, err := hasValidControlMount(container) + if err != nil { + return err + } + if !found { + container.VolumeMounts = append(container.VolumeMounts, corev1.VolumeMount{ + Name: SnapshotControlVolumeName, + MountPath: SnapshotControlMountPath, + SubPath: container.Name, + }) + } + return nil +} + +func validateControlMount(container *corev1.Container) error { + found, err := hasValidControlMount(container) + if err != nil { + return err + } + if !found { + return fmt.Errorf("container %q is missing %s mounted at %s", container.Name, SnapshotControlVolumeName, SnapshotControlMountPath) + } + return nil +} + +func hasValidControlMount(container *corev1.Container) (bool, error) { + found := false + for i := range container.VolumeMounts { + mount := &container.VolumeMounts[i] + if mount.Name != SnapshotControlVolumeName && mount.MountPath != SnapshotControlMountPath { + continue + } + if found { + return false, fmt.Errorf("container %q has duplicate snapshot control mounts", container.Name) + } + found = true + if err := validateControlMountValue(container.Name, mount); err != nil { + return false, err + } + } + return found, nil +} + +func validateControlMountValue(containerName string, mount *corev1.VolumeMount) error { + if mount.Name != SnapshotControlVolumeName || mount.MountPath != SnapshotControlMountPath || mount.SubPath != containerName { + return fmt.Errorf("container %q requires volume %q mounted at %s with subPath %q", containerName, SnapshotControlVolumeName, SnapshotControlMountPath, containerName) + } + if mount.ReadOnly || mount.RecursiveReadOnly != nil || mount.SubPathExpr != "" || + (mount.MountPropagation != nil && *mount.MountPropagation != corev1.MountPropagationNone) { + return fmt.Errorf("container %q has conflicting options on the snapshot control mount", containerName) + } + return nil +} + +func ensureControlEnvironment(container *corev1.Container) error { + for _, name := range []string{SnapshotControlDirEnv, LegacySnapshotControlDirEnv} { + if err := ensureControlEnv(container, name); err != nil { + return err + } + } + return nil +} + +func ensureControlEnv(container *corev1.Container, name string) error { + found, err := hasValidControlEnv(container, name) + if err != nil { + return err + } + if !found { + container.Env = append(container.Env, corev1.EnvVar{Name: name, Value: SnapshotControlMountPath}) + } + return nil +} + +func validateControlEnvironment(container *corev1.Container) error { + found, err := hasValidControlEnv(container, SnapshotControlDirEnv) + if err != nil { + return err + } + if !found { + return fmt.Errorf("container %q is missing %s environment variable", container.Name, SnapshotControlDirEnv) + } + if _, err := hasValidControlEnv(container, LegacySnapshotControlDirEnv); err != nil { + return err + } + return nil +} + +func hasValidControlEnv(container *corev1.Container, name string) (bool, error) { + found := false + for i := range container.Env { + env := &container.Env[i] + if env.Name != name { + continue + } + if found { + return false, fmt.Errorf("container %q has duplicate %s environment variables", container.Name, name) + } + found = true + if env.Value != SnapshotControlMountPath || env.ValueFrom != nil { + return false, fmt.Errorf("container %q has conflicting %s environment variable", container.Name, name) + } + } + return found, nil +} + +func ensureRestoreStartupProbe(container *corev1.Container) { + container.StartupProbe = canonicalRestoreStartupProbe() +} + +func canonicalRestoreStartupProbe() *corev1.Probe { + return &corev1.Probe{ + ProbeHandler: corev1.ProbeHandler{Exec: &corev1.ExecAction{ + Command: []string{"cat", path.Join(SnapshotControlMountPath, RestoreCompleteFile)}, + }}, + TimeoutSeconds: 1, + PeriodSeconds: 1, + FailureThreshold: restoreStartupFailureThreshold, + SuccessThreshold: 1, + } +} + +func validateRestoreStartupProbe(container *corev1.Container) error { + probe := container.StartupProbe + if probe == nil { + return fmt.Errorf("container %q is missing the restore startup gate", container.Name) + } + if probe.Exec == nil || !isRestoreCompletionProbeCommand(probe.Exec.Command) || + probe.HTTPGet != nil || probe.TCPSocket != nil || probe.GRPC != nil { + return fmt.Errorf("container %q restore startup gate must check %s", container.Name, path.Join(SnapshotControlMountPath, RestoreCompleteFile)) + } + return nil +} + +func validateCanonicalRestoreStartupProbe(container *corev1.Container) error { + if !reflect.DeepEqual(container.StartupProbe, canonicalRestoreStartupProbe()) { + return fmt.Errorf("container %q has a conflicting restore startup gate", container.Name) + } + return nil +} + +func isRestoreCompletionProbeCommand(command []string) bool { + completionPath := path.Join(SnapshotControlMountPath, RestoreCompleteFile) + switch { + case len(command) == 2 && isSupportedProbeExecutable(command[0], "cat"): + return command[1] == completionPath + case len(command) == 3 && isSupportedProbeExecutable(command[0], "test"): + return command[1] == "-f" && command[2] == completionPath + default: + return false + } +} + +func isSupportedProbeExecutable(actual string, executable string) bool { + return actual == executable || actual == "/bin/"+executable || actual == "/usr/bin/"+executable +} + +func ensurePodSeccompProfile(spec *corev1.PodSpec, expected string) error { + if expected == "" { + return nil + } + if spec.SecurityContext == nil { + spec.SecurityContext = &corev1.PodSecurityContext{} + } + if spec.SecurityContext.SeccompProfile == nil { + spec.SecurityContext.SeccompProfile = localhostSeccompProfile(expected) + return nil + } + if !matchesLocalhostSeccompProfile(spec.SecurityContext.SeccompProfile, expected) { + return fmt.Errorf("pod has a conflicting seccomp profile; expected localhost profile %q", expected) + } + return nil +} + +func validatePodSeccompProfile(spec *corev1.PodSpec, expected string) error { + if expected == "" { + return nil + } + if spec.SecurityContext == nil || !matchesLocalhostSeccompProfile(spec.SecurityContext.SeccompProfile, expected) { + return fmt.Errorf("pod must use localhost seccomp profile %q", expected) + } + return nil +} + +func validateContainerSeccompProfile(container *corev1.Container, expected string) error { + if expected == "" || container.SecurityContext == nil || container.SecurityContext.SeccompProfile == nil { + return nil + } + if !matchesLocalhostSeccompProfile(container.SecurityContext.SeccompProfile, expected) { + return fmt.Errorf("container %q overrides required localhost seccomp profile %q", container.Name, expected) + } + return nil +} + +func localhostSeccompProfile(profile string) *corev1.SeccompProfile { + return &corev1.SeccompProfile{ + Type: corev1.SeccompProfileTypeLocalhost, + LocalhostProfile: &profile, + } +} + +func matchesLocalhostSeccompProfile(profile *corev1.SeccompProfile, expected string) bool { + return profile != nil && profile.Type == corev1.SeccompProfileTypeLocalhost && + profile.LocalhostProfile != nil && *profile.LocalhostProfile == expected +} diff --git a/api/v1alpha1/restore_pod_test.go b/api/v1alpha1/restore_pod_test.go new file mode 100644 index 00000000..a6169ea8 --- /dev/null +++ b/api/v1alpha1/restore_pod_test.go @@ -0,0 +1,412 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package v1alpha1 + +import ( + "reflect" + "testing" + + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" +) + +func restorePodFixture() *corev1.Pod { + return &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "restore-worker", + Namespace: "inference", + Annotations: map[string]string{"example.com/team": "inference"}, + }, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{ + { + Name: "main", + Image: "worker:latest", + Command: []string{"python3"}, + Args: []string{"serve.py"}, + }, + {Name: "sidecar", Image: "sidecar:latest"}, + }, + }, + } +} + +func singleRestoreMapping() []RestoreContainerMapping { + return []RestoreContainerMapping{{Source: "main", Destination: "main"}} +} + +func TestBuildRestorePodShapesSingleDestination(t *testing.T) { + original := restorePodFixture() + before := original.DeepCopy() + options := RestorePodOptions{SeccompProfile: DefaultSeccompLocalhostProfile} + + shaped, err := BuildRestorePod(original, "snapshot-a", singleRestoreMapping(), options) + if err != nil { + t.Fatalf("BuildRestorePod() failed: %v", err) + } + if !reflect.DeepEqual(original, before) { + t.Fatal("BuildRestorePod() mutated its input") + } + if shaped.Annotations[RestoreFromAnnotation] != "snapshot-a" { + t.Fatalf("%s = %q", RestoreFromAnnotation, shaped.Annotations[RestoreFromAnnotation]) + } + if _, found := shaped.Annotations[RestoreContainerMapAnnotation]; found { + t.Fatalf("same-name restore unexpectedly set %s", RestoreContainerMapAnnotation) + } + if shaped.Annotations["example.com/team"] != "inference" { + t.Fatal("unrelated annotation was not preserved") + } + if len(shaped.Spec.Volumes) != 1 || shaped.Spec.Volumes[0].EmptyDir == nil { + t.Fatalf("expected one snapshot control emptyDir, got %#v", shaped.Spec.Volumes) + } + + main := &shaped.Spec.Containers[0] + if !reflect.DeepEqual(main.Command, []string{"python3"}) || !reflect.DeepEqual(main.Args, []string{"serve.py"}) { + t.Fatalf("workload command changed: command=%v args=%v", main.Command, main.Args) + } + if len(main.VolumeMounts) != 1 || main.VolumeMounts[0].SubPath != "main" { + t.Fatalf("unexpected snapshot control mount: %#v", main.VolumeMounts) + } + for _, name := range []string{SnapshotControlDirEnv, LegacySnapshotControlDirEnv} { + if got := restoreEnvValue(main.Env, name); got != SnapshotControlMountPath { + t.Fatalf("%s = %q, want %q", name, got, SnapshotControlMountPath) + } + } + for _, name := range []string{RestoreStandbyModeEnv, LegacyRestoreStandbyModeEnv} { + if got := restoreEnvValue(main.Env, name); got != "" { + t.Fatalf("generic builder injected workload-specific %s value %q", name, got) + } + } + if main.StartupProbe == nil || main.StartupProbe.Exec == nil || + !reflect.DeepEqual(main.StartupProbe.Exec.Command, []string{"cat", SnapshotControlMountPath + "/" + RestoreCompleteFile}) { + t.Fatalf("unexpected restore startup gate: %#v", main.StartupProbe) + } + if shaped.Spec.SecurityContext == nil || + !matchesLocalhostSeccompProfile(shaped.Spec.SecurityContext.SeccompProfile, DefaultSeccompLocalhostProfile) { + t.Fatalf("missing expected seccomp profile: %#v", shaped.Spec.SecurityContext) + } + if !reflect.DeepEqual(shaped.Spec.Containers[1], before.Spec.Containers[1]) { + t.Fatal("non-destination sidecar was modified") + } + if err := ValidateRestorePod(shaped, "snapshot-a", singleRestoreMapping(), options); err != nil { + t.Fatalf("ValidateRestorePod() rejected builder output: %v", err) + } +} + +func TestBuildRestorePodShapesFanoutIdempotently(t *testing.T) { + pod := restorePodFixture() + pod.Spec.Containers = []corev1.Container{{Name: "engine-0"}, {Name: "engine-1"}} + mappings := []RestoreContainerMapping{ + {Source: "main", Destination: "engine-0"}, + {Source: "main", Destination: "engine-1"}, + } + + first, err := BuildRestorePod(pod, "snapshot-a", mappings, RestorePodOptions{}) + if err != nil { + t.Fatalf("first BuildRestorePod() failed: %v", err) + } + second, err := BuildRestorePod(first, "snapshot-a", mappings, RestorePodOptions{}) + if err != nil { + t.Fatalf("second BuildRestorePod() failed: %v", err) + } + if !reflect.DeepEqual(first, second) { + t.Fatal("BuildRestorePod() is not idempotent") + } + if got := first.Annotations[RestoreContainerMapAnnotation]; got != "main=engine-0,main=engine-1" { + t.Fatalf("%s = %q", RestoreContainerMapAnnotation, got) + } + if len(first.Spec.Volumes) != 1 { + t.Fatalf("expected one shared volume, got %d", len(first.Spec.Volumes)) + } + for i, name := range []string{"engine-0", "engine-1"} { + container := &first.Spec.Containers[i] + if len(container.VolumeMounts) != 1 || container.VolumeMounts[0].SubPath != name { + t.Fatalf("container %q mount = %#v", name, container.VolumeMounts) + } + } +} + +func TestBuildRestorePodUsesAuthoritativeRestoreGate(t *testing.T) { + pod := restorePodFixture() + pod.Spec.Containers[0].LivenessProbe = &corev1.Probe{ + ProbeHandler: corev1.ProbeHandler{HTTPGet: &corev1.HTTPGetAction{Path: "/live"}}, + PeriodSeconds: 10, + } + pod.Spec.Containers[0].ReadinessProbe = &corev1.Probe{ + ProbeHandler: corev1.ProbeHandler{HTTPGet: &corev1.HTTPGetAction{Path: "/ready"}}, + PeriodSeconds: 5, + } + pod.Spec.Containers[0].StartupProbe = &corev1.Probe{ + ProbeHandler: corev1.ProbeHandler{HTTPGet: &corev1.HTTPGetAction{Path: "/live"}}, + PeriodSeconds: 3, + } + beforeLiveness := pod.Spec.Containers[0].LivenessProbe.DeepCopy() + beforeReadiness := pod.Spec.Containers[0].ReadinessProbe.DeepCopy() + beforeStartup := pod.Spec.Containers[0].StartupProbe.DeepCopy() + + shaped, err := BuildRestorePod(pod, "snapshot-a", singleRestoreMapping(), RestorePodOptions{}) + if err != nil { + t.Fatalf("BuildRestorePod() failed: %v", err) + } + main := &shaped.Spec.Containers[0] + if !reflect.DeepEqual(main.LivenessProbe, beforeLiveness) || !reflect.DeepEqual(main.ReadinessProbe, beforeReadiness) { + t.Fatal("workload liveness or readiness probe was modified") + } + if reflect.DeepEqual(main.StartupProbe, beforeStartup) { + t.Fatal("workload startup probe was not replaced by the restore gate") + } + if !reflect.DeepEqual(pod.Spec.Containers[0].StartupProbe, beforeStartup) { + t.Fatal("BuildRestorePod() mutated the input startup probe") + } + expectedCommand := []string{"cat", SnapshotControlMountPath + "/" + RestoreCompleteFile} + if main.StartupProbe == nil || main.StartupProbe.Exec == nil || + !reflect.DeepEqual(main.StartupProbe.Exec.Command, expectedCommand) { + t.Fatalf("restore startup gate does not check the completion sentinel: %#v", main.StartupProbe) + } + if main.StartupProbe.InitialDelaySeconds != 0 || main.StartupProbe.TimeoutSeconds != 1 || main.StartupProbe.PeriodSeconds != 1 || + main.StartupProbe.FailureThreshold != restoreStartupFailureThreshold || main.StartupProbe.SuccessThreshold != 1 { + t.Fatalf("restore startup gate timing is incorrect: %#v", main.StartupProbe) + } +} + +func TestBuildRestorePodCanonicalizesEquivalentMapping(t *testing.T) { + pod := restorePodFixture() + pod.Spec.Containers = []corev1.Container{{Name: "engine-0"}, {Name: "engine-1"}} + pod.Annotations[RestoreContainerMapAnnotation] = "main=engine-1, main=engine-0" + mappings := []RestoreContainerMapping{ + {Source: "main", Destination: "engine-0"}, + {Source: "main", Destination: "engine-1"}, + } + + shaped, err := BuildRestorePod(pod, "snapshot-a", mappings, RestorePodOptions{}) + if err != nil { + t.Fatalf("BuildRestorePod() failed: %v", err) + } + if got := shaped.Annotations[RestoreContainerMapAnnotation]; got != "main=engine-0,main=engine-1" { + t.Fatalf("canonical mapping = %q", got) + } +} + +func TestBuildRestorePodRejectsConflictsAtomically(t *testing.T) { + runtimeDefault := corev1.SeccompProfile{Type: corev1.SeccompProfileTypeRuntimeDefault} + tests := []struct { + name string + mutate func(*corev1.Pod) + mappings []RestoreContainerMapping + }{ + { + name: "restore annotation", + mutate: func(pod *corev1.Pod) { + pod.Annotations[RestoreFromAnnotation] = "snapshot-b" + }, + }, + { + name: "mapping annotation", + mutate: func(pod *corev1.Pod) { + pod.Annotations[RestoreContainerMapAnnotation] = "main=sidecar" + }, + }, + { + name: "control volume", + mutate: func(pod *corev1.Pod) { + pod.Spec.Volumes = []corev1.Volume{{ + Name: SnapshotControlVolumeName, + VolumeSource: corev1.VolumeSource{PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ + ClaimName: "wrong", + }}, + }} + }, + }, + { + name: "control mount name", + mutate: func(pod *corev1.Pod) { + pod.Spec.Containers[0].VolumeMounts = []corev1.VolumeMount{{ + Name: SnapshotControlVolumeName, MountPath: "/wrong", SubPath: "main", + }} + }, + }, + { + name: "control mount path", + mutate: func(pod *corev1.Pod) { + pod.Spec.Containers[0].VolumeMounts = []corev1.VolumeMount{{ + Name: "other", MountPath: SnapshotControlMountPath, + }} + }, + }, + { + name: "control environment", + mutate: func(pod *corev1.Pod) { + pod.Spec.Containers[0].Env = []corev1.EnvVar{{Name: SnapshotControlDirEnv, Value: "/wrong"}} + }, + }, + { + name: "pod seccomp", + mutate: func(pod *corev1.Pod) { + pod.Spec.SecurityContext = &corev1.PodSecurityContext{SeccompProfile: runtimeDefault.DeepCopy()} + }, + }, + { + name: "container seccomp", + mutate: func(pod *corev1.Pod) { + pod.Spec.Containers[0].SecurityContext = &corev1.SecurityContext{SeccompProfile: runtimeDefault.DeepCopy()} + }, + }, + { + name: "missing destination", + mappings: []RestoreContainerMapping{{Source: "main", Destination: "missing"}}, + }, + { + name: "duplicate destination", + mappings: []RestoreContainerMapping{ + {Source: "main", Destination: "main"}, + {Source: "main", Destination: "main"}, + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + pod := restorePodFixture() + if test.mutate != nil { + test.mutate(pod) + } + before := pod.DeepCopy() + mappings := test.mappings + if mappings == nil { + mappings = singleRestoreMapping() + } + + if _, err := BuildRestorePod( + pod, + "snapshot-a", + mappings, + RestorePodOptions{SeccompProfile: DefaultSeccompLocalhostProfile}, + ); err == nil { + t.Fatal("BuildRestorePod() unexpectedly succeeded") + } + if !reflect.DeepEqual(pod, before) { + t.Fatal("BuildRestorePod() mutated input after failure") + } + }) + } +} + +func TestValidateRestorePodAllowsMissingLegacyControlEnvironment(t *testing.T) { + pod, err := BuildRestorePod(restorePodFixture(), "snapshot-a", singleRestoreMapping(), RestorePodOptions{}) + if err != nil { + t.Fatalf("BuildRestorePod() failed: %v", err) + } + pod.Spec.Containers[0].Env = removeRestoreEnv(pod.Spec.Containers[0].Env, LegacySnapshotControlDirEnv) + + if err := ValidateRestorePod(pod, "snapshot-a", singleRestoreMapping(), RestorePodOptions{}); err != nil { + t.Fatalf("ValidateRestorePod() rejected Pod without deprecated environment alias: %v", err) + } + pod.Spec.Containers[0].Env = append(pod.Spec.Containers[0].Env, corev1.EnvVar{ + Name: LegacySnapshotControlDirEnv, + Value: "/wrong", + }) + if err := ValidateRestorePod(pod, "snapshot-a", singleRestoreMapping(), RestorePodOptions{}); err == nil { + t.Fatal("ValidateRestorePod() accepted conflicting deprecated environment alias") + } +} + +func TestValidateRestorePodAcceptsEquivalentRestoreCompletionGates(t *testing.T) { + tests := []struct { + name string + command []string + }{ + { + name: "cat through an absolute path", + command: []string{"/usr/bin/cat", SnapshotControlMountPath + "/" + RestoreCompleteFile}, + }, + { + name: "test file existence", + command: []string{"test", "-f", SnapshotControlMountPath + "/" + RestoreCompleteFile}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + pod, err := BuildRestorePod(restorePodFixture(), "snapshot-a", singleRestoreMapping(), RestorePodOptions{}) + if err != nil { + t.Fatalf("BuildRestorePod() failed: %v", err) + } + pod.Spec.Containers[0].StartupProbe = &corev1.Probe{ + ProbeHandler: corev1.ProbeHandler{Exec: &corev1.ExecAction{Command: test.command}}, + PeriodSeconds: 7, + FailureThreshold: 42, + } + + if err := ValidateRestorePod(pod, "snapshot-a", singleRestoreMapping(), RestorePodOptions{}); err != nil { + t.Fatalf("ValidateRestorePod() rejected equivalent restore gate: %v", err) + } + }) + } +} + +func TestValidateRestorePodRejectsContractDrift(t *testing.T) { + options := RestorePodOptions{SeccompProfile: DefaultSeccompLocalhostProfile} + valid, err := BuildRestorePod(restorePodFixture(), "snapshot-a", singleRestoreMapping(), options) + if err != nil { + t.Fatalf("BuildRestorePod() failed: %v", err) + } + tests := map[string]func(*corev1.Pod){ + "annotation": func(pod *corev1.Pod) { + delete(pod.Annotations, RestoreFromAnnotation) + }, + "volume": func(pod *corev1.Pod) { + pod.Spec.Volumes = nil + }, + "mount": func(pod *corev1.Pod) { + pod.Spec.Containers[0].VolumeMounts[0].SubPath = "other" + }, + "environment": func(pod *corev1.Pod) { + pod.Spec.Containers[0].Env = removeRestoreEnv(pod.Spec.Containers[0].Env, SnapshotControlDirEnv) + }, + "startup gate": func(pod *corev1.Pod) { + pod.Spec.Containers[0].StartupProbe.Exec.Command = []string{"true"} + }, + "seccomp": func(pod *corev1.Pod) { + pod.Spec.SecurityContext.SeccompProfile = nil + }, + } + for name, mutate := range tests { + t.Run(name, func(t *testing.T) { + pod := valid.DeepCopy() + mutate(pod) + if err := ValidateRestorePod(pod, "snapshot-a", singleRestoreMapping(), options); err == nil { + t.Fatal("ValidateRestorePod() unexpectedly succeeded") + } + }) + } +} + +func TestBuildRestorePodAllowsUnmanagedSeccomp(t *testing.T) { + pod := restorePodFixture() + shaped, err := BuildRestorePod(pod, "snapshot-a", singleRestoreMapping(), RestorePodOptions{}) + if err != nil { + t.Fatalf("BuildRestorePod() failed: %v", err) + } + if shaped.Spec.SecurityContext != nil { + t.Fatalf("empty option unexpectedly changed security context: %#v", shaped.Spec.SecurityContext) + } +} + +func restoreEnvValue(env []corev1.EnvVar, name string) string { + for _, item := range env { + if item.Name == name { + return item.Value + } + } + return "" +} + +func removeRestoreEnv(env []corev1.EnvVar, name string) []corev1.EnvVar { + result := make([]corev1.EnvVar, 0, len(env)) + for _, item := range env { + if item.Name != name { + result = append(result, item) + } + } + return result +} diff --git a/docs/restore-pod-contract.md b/docs/restore-pod-contract.md new file mode 100644 index 00000000..36bae0e7 --- /dev/null +++ b/docs/restore-pod-contract.md @@ -0,0 +1,120 @@ + + +# Restore Pod contract + +A Pod requests restore by naming a ready `PodSnapshot` in its namespace. The +Snapshot node agent restores only Pods that implement this contract. + +Go integrations should use: + +```go +restored, err := snapshotv1alpha1.BuildRestorePod( + pod, + snapshotName, + mappings, + snapshotv1alpha1.RestorePodOptions{ + SeccompProfile: snapshotv1alpha1.DefaultSeccompLocalhostProfile, + }, +) +``` + +`BuildRestorePod` is a pure, atomic transformation: it returns a deep copy or +an error and never mutates its input. Reapplying it with the same arguments is +idempotent. It emits one canonical representation and validates that exact +output. `ValidateRestorePod` checks the stable runtime contract without +mutation or Kubernetes API reads, accepting supported equivalent completion +gates and caller-selected probe timing. Conflicting annotations, volumes, +mounts, environment, and security settings are rejected instead of overwritten. + +The producer derives each typed mapping source from the referenced +`PodSnapshot`. The builder validates one-source-to-many-destination consistency +but deliberately performs no API read to discover the captured source. + +Non-Go producers can implement the declarative form directly. This fan-out +example restores the captured `main` process into two destination containers: + +```yaml +apiVersion: v1 +kind: Pod +metadata: + name: restored-worker + annotations: + nvidia.com/restore-from: worker-snapshot + nvidia.com/restore-container-map: main=engine-0,main=engine-1 +spec: + securityContext: + seccompProfile: + type: Localhost + localhostProfile: profiles/block-iouring.json + volumes: + - name: snapshot-control + emptyDir: {} + containers: + - name: engine-0 + image: worker:latest + # The caller supplies an inert, long-running entrypoint. + command: ["/bin/sh", "-c", "exec sleep infinity"] + env: + - name: SNAPSHOT_CONTROL_DIR + value: /snapshot-control + volumeMounts: + - name: snapshot-control + mountPath: /snapshot-control + subPath: engine-0 + startupProbe: + exec: + command: ["cat", "/snapshot-control/restore-complete"] + timeoutSeconds: 1 + periodSeconds: 1 + failureThreshold: 1800 + successThreshold: 1 + - name: engine-1 + image: worker:latest + command: ["/bin/sh", "-c", "exec sleep infinity"] + env: + - name: SNAPSHOT_CONTROL_DIR + value: /snapshot-control + volumeMounts: + - name: snapshot-control + mountPath: /snapshot-control + subPath: engine-1 + startupProbe: + exec: + command: ["cat", "/snapshot-control/restore-complete"] + timeoutSeconds: 1 + periodSeconds: 1 + failureThreshold: 1800 + successThreshold: 1 +``` + +The container mapping is optional for a single same-name restore. Every +destination mounts the one shared `snapshot-control` `emptyDir` at +`/snapshot-control` with `subPath` equal to its container name. The canonical +`SNAPSHOT_CONTROL_DIR` environment variable points to that mount. The builder +also injects the deprecated `DYN_SNAPSHOT_CONTROL_DIR` alias during the +migration window; hand-authored Pods may omit the alias. + +The `restore-complete` sentinel probe shown above is the authoritative startup +gate. Kubernetes supports only one startup probe, so the builder replaces any +existing startup probe with the canonical restore gate while keeping workload +liveness and readiness probes unchanged. The runtime validator also accepts a +direct `test -f /snapshot-control/restore-complete` gate and does not pin probe +timing to a particular builder release. The canonical builder's extended +failure threshold allows 1,800 consecutive one-second startup-probe failures. +Kubernetes pauses liveness and readiness probes until the startup gate +succeeds; if restoration exceeds that failure budget, kubelet restarts the +placeholder according to the Pod's restart policy. + +`RestorePodOptions.SeccompProfile` controls Snapshot's pod-level localhost +profile. An empty value leaves seccomp unmanaged. A destination container must +not override a requested profile with a conflicting container-level profile. + +Snapshot does not modify container commands and does not inject +`SNAPSHOT_RESTORE_STANDBY`, its deprecated +`DYN_SNAPSHOT_RESTORE_STANDBY` alias, or any other workload-specific standby +setting. Both names are exported by the Go API so application owners can set +the convention they support. The producer must ensure each destination process +remains alive and inert until the agent replaces it with the restored process. diff --git a/e2e/snapshot_e2e/workloads.py b/e2e/snapshot_e2e/workloads.py index 855fedd3..40289494 100644 --- a/e2e/snapshot_e2e/workloads.py +++ b/e2e/snapshot_e2e/workloads.py @@ -146,12 +146,13 @@ def restore_pod( spec["containers"][0]["env"] = [ {"name": "DYN_SNAPSHOT_RESTORE_STANDBY", "value": "1"}, {"name": "SNAPSHOT_CONTROL_DIR", "value": CONTROL_DIR}, + {"name": "DYN_SNAPSHOT_CONTROL_DIR", "value": CONTROL_DIR}, {"name": RESTORE_TOKEN_ENV, "value": run.restore_token}, ] spec["containers"][0]["startupProbe"] = { - "exec": {"command": ["/bin/bash", "-lc", f"test -f {RESTORE_DONE}"]}, + "exec": {"command": ["cat", RESTORE_DONE]}, "periodSeconds": 1, - "failureThreshold": 1200, + "failureThreshold": 1800, } if source_node: spec["affinity"] = same_node_affinity(source_node) @@ -206,12 +207,13 @@ def multi_restore_pod( "env": [ {"name": "DYN_SNAPSHOT_RESTORE_STANDBY", "value": "1"}, {"name": "SNAPSHOT_CONTROL_DIR", "value": CONTROL_DIR}, + {"name": "DYN_SNAPSHOT_CONTROL_DIR", "value": CONTROL_DIR}, {"name": RESTORE_TOKEN_ENV, "value": restore_tokens[destination]}, ], "startupProbe": { - "exec": {"command": ["/bin/bash", "-lc", f"test -f {RESTORE_DONE}"]}, + "exec": {"command": ["cat", RESTORE_DONE]}, "periodSeconds": 1, - "failureThreshold": 1200, + "failureThreshold": 1800, }, } containers.append(container) @@ -269,7 +271,10 @@ def base_pod_spec( } ) if control_volume: - container["volumeMounts"].insert(0, {"name": "snapshot-control", "mountPath": CONTROL_DIR}) + container["volumeMounts"].insert( + 0, + {"name": "snapshot-control", "mountPath": CONTROL_DIR, "subPath": CONTAINER}, + ) volumes.insert(0, {"name": "snapshot-control", "emptyDir": {}}) spec: dict[str, Any] = { "restartPolicy": "Never", diff --git a/e2e/tests/test_workload_scripts.py b/e2e/tests/test_workload_scripts.py index 30b3a03b..963a4956 100644 --- a/e2e/tests/test_workload_scripts.py +++ b/e2e/tests/test_workload_scripts.py @@ -25,6 +25,7 @@ import pytest +from snapshot_e2e import k8s from snapshot_e2e import workloads @@ -111,3 +112,41 @@ def test_snapshotjob_exit_template_never_signals_ready() -> None: timeout=10, ) assert proc.returncode == exit_code + + +@pytest.mark.workload +def test_restore_manifests_use_canonical_control_mount_and_startup_gate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SNAPSHOT_E2E_WORKLOAD_IMAGE", "snapshot-workload:test") + config = k8s.E2EConfig( + namespace="snapshot-e2e", + release="snapshot", + pvc_name="snapshot-pvc", + kubeconfig=None, + ) + run = workloads.TestRun.new("manifest") + + single = workloads.restore_pod(config=config, run=run, gpu=False) + single_container = single["spec"]["containers"][0] + assert single_container["volumeMounts"][0] == { + "name": "snapshot-control", + "mountPath": workloads.CONTROL_DIR, + "subPath": workloads.CONTAINER, + } + assert single_container["startupProbe"]["exec"]["command"] == [ + "cat", + workloads.RESTORE_DONE, + ] + + multi, _ = workloads.multi_restore_pod( + config=config, + run=run, + source_node="source-node", + ) + for container in multi["spec"]["containers"]: + assert container["volumeMounts"][0]["subPath"] == container["name"] + assert container["startupProbe"]["exec"]["command"] == [ + "cat", + workloads.RESTORE_DONE, + ] diff --git a/operator/cmd/snapshotctl/README.md b/operator/cmd/snapshotctl/README.md index 5d2e315a..c6933de1 100644 --- a/operator/cmd/snapshotctl/README.md +++ b/operator/cmd/snapshotctl/README.md @@ -55,6 +55,14 @@ That pod manifest must: - describe the worker pod you want to checkpoint or restore - use the placeholder image for checkpoint-aware flows - match the runtime-relevant worker settings you care about preserving +- provide an inert, long-running entrypoint for each restore destination + +`snapshotctl restore` applies Snapshot's generic +[restore Pod contract](../../../docs/restore-pod-contract.md), including the +annotations, control volume, environment, startup gate, and seccomp profile. It +does not replace container commands or inject a workload-specific standby +environment variable. A platform-specific entrypoint convention remains the +platform's responsibility. In practice, start from the real worker pod spec you would normally run, then keep only the pod-level fields needed to recreate that worker accurately. diff --git a/operator/cmd/snapshotctl/restore.go b/operator/cmd/snapshotctl/restore.go index 651d91af..7c3ccac1 100644 --- a/operator/cmd/snapshotctl/restore.go +++ b/operator/cmd/snapshotctl/restore.go @@ -14,7 +14,6 @@ import ( "sigs.k8s.io/controller-runtime/pkg/client" snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" - snapshotprotocol "github.com/ai-dynamo/snapshot/operator/internal/protocol" ) type restoreOptions struct { @@ -41,22 +40,29 @@ func runRestoreFlow(ctx context.Context, opts restoreOptions) (*result, error) { if len(containers) != 1 || strings.TrimSpace(containers[0]) == "" { return nil, fmt.Errorf("PodSnapshot %s/%s must capture exactly one container", namespace, snapshotName) } + mappings, err := snapshotv1alpha1.RestoreContainerMappingsFromAnnotations(pod.Annotations, containers[0]) + if err != nil { + return nil, err + } - restorePod, err := snapshotprotocol.NewRestorePod(&corev1.Pod{ + candidate := &corev1.Pod{ TypeMeta: metav1.TypeMeta{APIVersion: "v1", Kind: "Pod"}, ObjectMeta: metav1.ObjectMeta{ Name: pod.Name, GenerateName: pod.GenerateName, + Namespace: namespace, Labels: pod.Labels, Annotations: pod.Annotations, }, Spec: *pod.Spec.DeepCopy(), - }, snapshotprotocol.PodOptions{ - Namespace: namespace, - SnapshotName: snapshotName, - SourceContainer: containers[0], - SeccompProfile: snapshotv1alpha1.DefaultSeccompLocalhostProfile, - }) + } + candidate.Spec.RestartPolicy = corev1.RestartPolicyNever + restorePod, err := snapshotv1alpha1.BuildRestorePod( + candidate, + snapshotName, + mappings, + snapshotv1alpha1.RestorePodOptions{SeccompProfile: snapshotv1alpha1.DefaultSeccompLocalhostProfile}, + ) if err != nil { return nil, err } diff --git a/operator/internal/protocol/restore.go b/operator/internal/protocol/restore.go deleted file mode 100644 index fa917158..00000000 --- a/operator/internal/protocol/restore.go +++ /dev/null @@ -1,223 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -package protocol - -import ( - "fmt" - "path/filepath" - - corev1 "k8s.io/api/core/v1" - - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" -) - -type PodOptions struct { - Namespace string - SnapshotName string - SourceContainer string - SeccompProfile string -} - -const ( - // RestoreStandbyModeEnv asks standby-aware workload entrypoints to capture - // restore context and sleep instead of cold-starting the workload. Generic - // images that do not honor this env must still provide their own inert - // restore command. - RestoreStandbyModeEnv = "DYN_SNAPSHOT_RESTORE_STANDBY" - restoreStartupFailureThreshold = 1800 // 30 minutes at 1s cadence. -) - -// NewRestorePod shapes every annotated target container for restore. -func NewRestorePod(pod *corev1.Pod, opts PodOptions) (*corev1.Pod, error) { - pod = pod.DeepCopy() - if pod.Annotations == nil { - pod.Annotations = map[string]string{} - } - pod.Annotations[snapshotv1alpha1.RestoreFromAnnotation] = opts.SnapshotName - if _, err := snapshotv1alpha1.GetRestoreFromSnapshotName(pod.Annotations); err != nil { - return nil, err - } - mappings, err := snapshotv1alpha1.RestoreContainerMappingsFromAnnotations(pod.Annotations, opts.SourceContainer) - if err != nil { - return nil, err - } - if err := snapshotv1alpha1.ValidateRestoreContainerMappings(mappings, opts.SourceContainer); err != nil { - return nil, err - } - if err := PrepareRestorePodSpec(&pod.Spec, mappings, opts.SeccompProfile, true); err != nil { - return nil, err - } - pod.Namespace = opts.Namespace - pod.Spec.RestartPolicy = corev1.RestartPolicyNever - return pod, nil -} - -// PrepareRestorePodSpec applies restore shaping to annotated target containers. -// It does not change container command/args. Once the checkpoint is ready, it -// sets DYN_SNAPSHOT_RESTORE_STANDBY=1 so standby-aware workload entrypoints -// sleep before CRIU restore; generic images that do not honor the env must -// still provide their own inert restore command. -func PrepareRestorePodSpec( - podSpec *corev1.PodSpec, - mappings []snapshotv1alpha1.RestoreContainerMapping, - seccompProfile string, - isCheckpointReady bool, -) error { - if podSpec == nil { - return fmt.Errorf("pod spec is nil") - } - if len(mappings) == 0 { - return fmt.Errorf("restore target container is required") - } - containers := make([]*corev1.Container, 0, len(mappings)) - for _, mapping := range mappings { - var container *corev1.Container - for i := range podSpec.Containers { - if podSpec.Containers[i].Name == mapping.Destination { - container = &podSpec.Containers[i] - break - } - } - if container == nil { - return fmt.Errorf("restore destination container %q not found in pod spec", mapping.Destination) - } - containers = append(containers, container) - } - - EnsureLocalhostSeccompProfile(podSpec, seccompProfile) - for _, container := range containers { - EnsureControlVolume(podSpec, container) - if isCheckpointReady { - // Standby-aware entrypoints honor this env by writing restore - // context and sleeping. Keep command/args intact so generic images - // can provide their own inert restore entrypoint when needed. - foundRestoreStandbyModeEnv := false - for i := range container.Env { - if container.Env[i].Name == RestoreStandbyModeEnv { - container.Env[i].Value = "1" - container.Env[i].ValueFrom = nil - foundRestoreStandbyModeEnv = true - break - } - } - if !foundRestoreStandbyModeEnv { - container.Env = append(container.Env, corev1.EnvVar{ - Name: RestoreStandbyModeEnv, - Value: "1", - }) - } - ensureRestoreStartupProbe(container) - } - } - return nil -} - -// ensureRestoreStartupProbe installs a StartupProbe that gates Ready until -// CRIU restore completes. It prefers the workload's existing Startup/Liveness/ -// Readiness probe (deep-copied with tightened cadence and infinite retries), -// and falls back to a sentinel-file exec probe when none is defined. -func ensureRestoreStartupProbe(container *corev1.Container) { - startup := container.StartupProbe - if startup == nil { - startup = container.LivenessProbe - if startup == nil { - startup = container.ReadinessProbe - } - } - if startup == nil { - container.StartupProbe = &corev1.Probe{ - ProbeHandler: corev1.ProbeHandler{ - Exec: &corev1.ExecAction{ - Command: []string{"cat", filepath.Join(snapshotv1alpha1.SnapshotControlMountPath, snapshotv1alpha1.RestoreCompleteFile)}, - }, - }, - TimeoutSeconds: 1, - PeriodSeconds: 1, - FailureThreshold: restoreStartupFailureThreshold, - SuccessThreshold: 1, - } - return - } - - startup = startup.DeepCopy() - startup.InitialDelaySeconds = 0 - startup.PeriodSeconds = 1 - startup.FailureThreshold = restoreStartupFailureThreshold - startup.SuccessThreshold = 1 - container.StartupProbe = startup -} - -// ValidateRestorePodSpec verifies the target containers are restore-shaped. -func ValidateRestorePodSpec( - podSpec *corev1.PodSpec, - mappings []snapshotv1alpha1.RestoreContainerMapping, - seccompProfile string, -) error { - if podSpec == nil { - return fmt.Errorf("pod spec is nil") - } - if len(mappings) == 0 { - return fmt.Errorf("restore target container is required") - } - hasControlVolume := false - for _, volume := range podSpec.Volumes { - if volume.Name == snapshotv1alpha1.SnapshotControlVolumeName && volume.EmptyDir != nil { - hasControlVolume = true - break - } - } - if !hasControlVolume { - return fmt.Errorf("missing %s emptyDir volume; add it via snapshotprotocol.EnsureControlVolume", snapshotv1alpha1.SnapshotControlVolumeName) - } - for _, mapping := range mappings { - name := mapping.Destination - var container *corev1.Container - for i := range podSpec.Containers { - if podSpec.Containers[i].Name == name { - container = &podSpec.Containers[i] - break - } - } - if container == nil { - return fmt.Errorf("restore target container %q not found in pod spec", name) - } - hasControlMount := false - for _, mount := range container.VolumeMounts { - if mount.Name == snapshotv1alpha1.SnapshotControlVolumeName && mount.MountPath == snapshotv1alpha1.SnapshotControlMountPath { - hasControlMount = true - if mount.SubPath != name { - return fmt.Errorf("expected SubPath %q for %s at %s on container %q, got %q", name, snapshotv1alpha1.SnapshotControlVolumeName, snapshotv1alpha1.SnapshotControlMountPath, name, mount.SubPath) - } - break - } - } - if !hasControlMount { - return fmt.Errorf("missing %s mount at %s on container %q", snapshotv1alpha1.SnapshotControlVolumeName, snapshotv1alpha1.SnapshotControlMountPath, name) - } - hasControlEnv := false - for _, env := range container.Env { - if env.Name == snapshotv1alpha1.SnapshotControlDirEnv { - hasControlEnv = true - break - } - } - if !hasControlEnv { - return fmt.Errorf("missing %s env var on container %q", snapshotv1alpha1.SnapshotControlDirEnv, name) - } - if container.StartupProbe == nil { - return fmt.Errorf("missing restore-complete startup probe on container %q", name) - } - } - if seccompProfile == "" { - return nil - } - if podSpec.SecurityContext == nil || podSpec.SecurityContext.SeccompProfile == nil { - return fmt.Errorf("missing localhost seccomp profile") - } - profile := podSpec.SecurityContext.SeccompProfile - if profile.Type != corev1.SeccompProfileTypeLocalhost || profile.LocalhostProfile == nil || *profile.LocalhostProfile != seccompProfile { - return fmt.Errorf("expected localhost seccomp profile %q", seccompProfile) - } - return nil -} diff --git a/operator/internal/protocol/restore_test.go b/operator/internal/protocol/restore_test.go deleted file mode 100644 index bf1db580..00000000 --- a/operator/internal/protocol/restore_test.go +++ /dev/null @@ -1,172 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -package protocol - -import ( - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - corev1 "k8s.io/api/core/v1" - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" -) - -func restorePodFixture() *corev1.Pod { - return &corev1.Pod{ - ObjectMeta: metav1.ObjectMeta{ - Name: "restore-worker", - Annotations: map[string]string{"example.com/team": "inference"}, - }, - Spec: corev1.PodSpec{ - Containers: []corev1.Container{ - { - Name: "main", - Image: "worker:latest", - Command: []string{"python3"}, - Args: []string{"serve.py"}, - }, - {Name: "sidecar", Image: "sidecar:latest"}, - }, - }, - } -} - -func TestNewRestorePodSetsRestoreFromAnnotation(t *testing.T) { - pod, err := NewRestorePod(restorePodFixture(), PodOptions{ - Namespace: "inference", - SnapshotName: "snapshot-a", - SourceContainer: "main", - SeccompProfile: snapshotv1alpha1.DefaultSeccompLocalhostProfile, - }) - require.NoError(t, err) - - assert.Equal(t, "inference", pod.Namespace) - assert.Equal(t, corev1.RestartPolicyNever, pod.Spec.RestartPolicy) - assert.Equal(t, "snapshot-a", pod.Annotations[snapshotv1alpha1.RestoreFromAnnotation]) - assert.NotContains(t, pod.Annotations, snapshotv1alpha1.RestoreContainerMapAnnotation) - assert.Equal(t, "inference", pod.Annotations["example.com/team"]) - - main := &pod.Spec.Containers[0] - assert.Equal(t, []string{"python3"}, main.Command) - assert.Equal(t, []string{"serve.py"}, main.Args) - assert.Equal(t, "1", envValue(main.Env, RestoreStandbyModeEnv)) - assert.Equal(t, snapshotv1alpha1.SnapshotControlMountPath, main.VolumeMounts[0].MountPath) - assert.Equal(t, "main", main.VolumeMounts[0].SubPath) - require.NotNil(t, main.StartupProbe) - - sidecar := &pod.Spec.Containers[1] - assert.Empty(t, sidecar.VolumeMounts) - assert.Empty(t, sidecar.Env) - assert.Nil(t, sidecar.StartupProbe) - require.NoError(t, ValidateRestorePodSpec(&pod.Spec, restoreMappings("main"), snapshotv1alpha1.DefaultSeccompLocalhostProfile)) -} - -func TestPrepareRestorePodSpecRequiresCapturedContainer(t *testing.T) { - spec := restorePodFixture().Spec - require.Error(t, PrepareRestorePodSpec(&spec, nil, "", true)) - err := PrepareRestorePodSpec(&spec, restoreMappings("missing"), "", true) - require.Error(t, err) - assert.Contains(t, err.Error(), `container "missing"`) -} - -func TestPrepareRestorePodSpecIsIdempotent(t *testing.T) { - spec := restorePodFixture().Spec - require.NoError(t, PrepareRestorePodSpec(&spec, restoreMappings("main"), "", true)) - require.NoError(t, PrepareRestorePodSpec(&spec, restoreMappings("main"), "", true)) - - main := &spec.Containers[0] - assert.Len(t, spec.Volumes, 1) - assert.Len(t, main.VolumeMounts, 1) - assert.Equal(t, "1", envValue(main.Env, RestoreStandbyModeEnv)) -} - -func TestPrepareRestorePodSpecReusesExistingProbe(t *testing.T) { - spec := restorePodFixture().Spec - spec.Containers[0].ReadinessProbe = &corev1.Probe{ - ProbeHandler: corev1.ProbeHandler{HTTPGet: &corev1.HTTPGetAction{Path: "/ready"}}, - InitialDelaySeconds: 30, - PeriodSeconds: 10, - FailureThreshold: 3, - } - require.NoError(t, PrepareRestorePodSpec(&spec, restoreMappings("main"), "", true)) - startup := spec.Containers[0].StartupProbe - require.NotNil(t, startup) - require.NotNil(t, startup.HTTPGet) - assert.Equal(t, "/ready", startup.HTTPGet.Path) - assert.Zero(t, startup.InitialDelaySeconds) - assert.Equal(t, int32(1), startup.PeriodSeconds) - assert.Equal(t, int32(restoreStartupFailureThreshold), startup.FailureThreshold) -} - -func TestValidateRestorePodSpecFailures(t *testing.T) { - spec := restorePodFixture().Spec - require.NoError(t, PrepareRestorePodSpec(&spec, restoreMappings("main"), snapshotv1alpha1.DefaultSeccompLocalhostProfile, true)) - - tests := map[string]func(*corev1.PodSpec){ - "volume": func(s *corev1.PodSpec) { s.Volumes = nil }, - "mount": func(s *corev1.PodSpec) { s.Containers[0].VolumeMounts = nil }, - "env": func(s *corev1.PodSpec) { - s.Containers[0].Env = nil - }, - "probe": func(s *corev1.PodSpec) { s.Containers[0].StartupProbe = nil }, - "seccomp": func(s *corev1.PodSpec) { - s.SecurityContext = nil - }, - } - for name, mutate := range tests { - t.Run(name, func(t *testing.T) { - bad := spec.DeepCopy() - mutate(bad) - require.Error(t, ValidateRestorePodSpec(bad, restoreMappings("main"), snapshotv1alpha1.DefaultSeccompLocalhostProfile)) - }) - } -} - -func TestNewRestorePodShapesMappedDestinations(t *testing.T) { - pod := restorePodFixture() - pod.Annotations[snapshotv1alpha1.RestoreContainerMapAnnotation] = "main=main,main=sidecar" - restored, err := NewRestorePod(pod, PodOptions{ - Namespace: "inference", - SnapshotName: "snapshot-a", - SourceContainer: "main", - }) - require.NoError(t, err) - - for i, name := range []string{"main", "sidecar"} { - container := &restored.Spec.Containers[i] - require.Equal(t, name, container.Name) - require.Len(t, container.VolumeMounts, 1) - assert.Equal(t, name, container.VolumeMounts[0].SubPath) - assert.Equal(t, "1", envValue(container.Env, RestoreStandbyModeEnv)) - require.NotNil(t, container.StartupProbe) - } -} - -func TestPrepareRestorePodSpecValidatesBeforeMutation(t *testing.T) { - spec := restorePodFixture().Spec - err := PrepareRestorePodSpec(&spec, restoreMappings("main", "missing"), snapshotv1alpha1.DefaultSeccompLocalhostProfile, true) - require.Error(t, err) - assert.Nil(t, spec.SecurityContext) - assert.Empty(t, spec.Volumes) - assert.Empty(t, spec.Containers[0].VolumeMounts) -} - -func restoreMappings(destinations ...string) []snapshotv1alpha1.RestoreContainerMapping { - mappings := make([]snapshotv1alpha1.RestoreContainerMapping, 0, len(destinations)) - for _, destination := range destinations { - mappings = append(mappings, snapshotv1alpha1.RestoreContainerMapping{Source: "main", Destination: destination}) - } - return mappings -} - -func envValue(env []corev1.EnvVar, name string) string { - for _, item := range env { - if item.Name == name { - return item.Value - } - } - return "" -}