feat(api): publish restore Pod contract

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`.
+
  
 
 

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

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,33 @@ 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 TestPreflightRestoreRetriesInProgressCondition(t *testing.T) {
 	pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"})
 	pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{

api/v1alpha1/constants.go

@@ -57,10 +57,10 @@ 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

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