From 62a9bb4dc9807f72f7532b008ea6d3a99547a05f Mon Sep 17 00:00:00 2001 From: Julien Mancuso Date: Mon, 31 Aug 2026 11:53:56 -0600 Subject: [PATCH] feat(api): publish restore Pod contract Signed-off-by: Julien Mancuso --- README.md | 5 + agent/internal/controller/controller.go | 110 +--- agent/internal/controller/controller_test.go | 217 +++++--- agent/internal/cuda/job.go | 16 +- agent/internal/cuda/job_test.go | 34 +- .../cuda/shim_restore_job_file_test.go | 6 +- agent/internal/executor/nsrestore.go | 4 +- agent/internal/runtime/control.go | 6 +- api/podcontract/doc.go | 6 + api/podcontract/protocol.go | 165 ++++++ .../protocol_test.go | 56 +- api/podcontract/restore_pod.go | 494 ++++++++++++++++++ api/podcontract/restore_pod_test.go | 436 ++++++++++++++++ .../restore_status.go | 16 +- .../restore_status_test.go | 18 +- api/v1alpha1/constants.go | 85 +-- api/v1alpha1/protocol.go | 86 --- docs/restore-pod-contract.md | 152 ++++++ e2e/snapshot_e2e/workloads.py | 15 +- e2e/tests/test_workload_scripts.py | 39 ++ operator/cmd/snapshotctl/README.md | 8 + operator/cmd/snapshotctl/checkpoint.go | 4 +- operator/cmd/snapshotctl/restore.go | 29 +- .../internal/controller/snapshotjob_job.go | 3 +- .../controller/snapshotjob_job_test.go | 7 +- operator/internal/protocol/control_volume.go | 17 +- .../internal/protocol/control_volume_test.go | 25 +- operator/internal/protocol/restore.go | 223 -------- operator/internal/protocol/restore_test.go | 172 ------ operator/internal/protocol/source_job.go | 10 +- .../protocol/source_job_identity_test.go | 4 +- operator/internal/protocol/source_job_test.go | 28 +- 32 files changed, 1659 insertions(+), 837 deletions(-) create mode 100644 api/podcontract/doc.go create mode 100644 api/podcontract/protocol.go rename api/{v1alpha1 => podcontract}/protocol_test.go (52%) create mode 100644 api/podcontract/restore_pod.go create mode 100644 api/podcontract/restore_pod_test.go rename api/{v1alpha1 => podcontract}/restore_status.go (86%) rename api/{v1alpha1 => podcontract}/restore_status_test.go (82%) delete mode 100644 api/v1alpha1/protocol.go create mode 100644 docs/restore-pod-contract.md delete mode 100644 operator/internal/protocol/restore.go delete mode 100644 operator/internal/protocol/restore_test.go diff --git a/README.md b/README.md index f277b8de..ba95e405 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,11 @@ 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 self-contained constants, builder, validator, and restore-outcome API +from `github.com/ai-dynamo/snapshot/api/podcontract`. +   diff --git a/agent/internal/controller/controller.go b/agent/internal/controller/controller.go index eeeec4ce..459ade95 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/podcontract" snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" ) @@ -99,7 +100,7 @@ type restoreTarget struct { type restorePlan struct { artifact *restoreArtifact - mappings []snapshotv1alpha1.RestoreContainerMapping + mappings []podcontract.ContainerMapping } type restoreResultState int @@ -231,7 +232,7 @@ func (w *NodeController) Run(ctx context.Context) error { ctx = logr.NewContext(ctx, w.log) w.log.Info("Starting snapshot node controller", "node", w.config.NodeName, - "restore_from_annotation", snapshotv1alpha1.RestoreFromAnnotation, + "restore_from_annotation", podcontract.RestoreFromAnnotation, ) w.log.Info("Watching pods cluster-wide (all namespaces)") @@ -384,7 +385,7 @@ func (w *NodeController) restorePodRequested(pod *corev1.Pod) bool { if pod.Spec.NodeName != w.config.NodeName { return false } - _, requested := pod.Annotations[snapshotv1alpha1.RestoreFromAnnotation] + _, requested := pod.Annotations[podcontract.RestoreFromAnnotation] return requested } @@ -514,7 +515,7 @@ func (w *NodeController) preflightRestore(ctx context.Context, pod *corev1.Pod) } func (w *NodeController) getPodSnapshotFromPod(ctx context.Context, pod *corev1.Pod) (*snapshotv1alpha1.PodSnapshot, error) { - snapshotName, err := snapshotv1alpha1.GetRestoreFromSnapshotName(pod.Annotations) + snapshotName, err := podcontract.GetRestoreFromSnapshotName(pod.Annotations) if err != nil { return nil, err } @@ -593,75 +594,25 @@ func validatePodSnapshotContentForRestore(content *snapshotv1alpha1.PodSnapshotC // validateRestoreTarget resolves the captured source and validates every // requested destination against the restore Pod. -func validateRestoreTarget(pod *corev1.Pod, snapshot *snapshotv1alpha1.PodSnapshot, content *snapshotv1alpha1.PodSnapshotContent) (*restoreTarget, []snapshotv1alpha1.RestoreContainerMapping, error) { +func validateRestoreTarget(pod *corev1.Pod, snapshot *snapshotv1alpha1.PodSnapshot, content *snapshotv1alpha1.PodSnapshotContent) (*restoreTarget, []podcontract.ContainerMapping, error) { containerName, err := singleTargetContainer(content) if err != nil { return nil, nil, err } - mappings, err := snapshotv1alpha1.RestoreContainerMappingsFromAnnotations(pod.Annotations, containerName) + mappings, err := podcontract.ContainerMappingsFromAnnotations(pod.Annotations, containerName) if err != nil { return nil, nil, err } - if err := validateRestoreMappingsForPod(&pod.Spec, mappings, containerName); err != nil { + if err := podcontract.Validate(pod, podcontract.Request{ + SnapshotName: snapshot.Name, + SourceContainer: containerName, + Mappings: mappings, + }); 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, - ) - } - } - return nil -} - // resolveRestoreArtifact resolves the validated restore target to its physical // checkpoint directory. A nil error always returns a complete artifact. func (w *NodeController) resolveRestoreArtifact(podKey string, target *restoreTarget) (*restoreArtifact, error) { @@ -698,7 +649,7 @@ func (w *NodeController) restorePodContainers(ctx context.Context, pod *corev1.P // considering a CRIU replay. recovering := restoreInProgress(pod) message := fmt.Sprintf("Restoring %d destination container(s) from PodSnapshot %s", len(plan.mappings), plan.artifact.SnapshotName) - if err := w.applyRestoredCondition(ctx, pod, corev1.ConditionFalse, snapshotv1alpha1.RestoreReasonInProgress, message); err != nil { + if err := w.applyRestoredCondition(ctx, pod, corev1.ConditionFalse, podcontract.RestoreReasonInProgress, message); err != nil { emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, corev1.EventTypeWarning, restoreStatusUpdateFailedReason, err.Error()) return true } @@ -732,7 +683,7 @@ func (w *NodeController) recordRestoreResults(ctx context.Context, pod *corev1.P "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 err := w.applyRestoredCondition(ctx, pod, corev1.ConditionFalse, snapshotv1alpha1.RestoreReasonInProgress, message); err != nil { + if err := w.applyRestoredCondition(ctx, pod, corev1.ConditionFalse, podcontract.RestoreReasonInProgress, message); err != nil { emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, corev1.EventTypeWarning, restoreStatusUpdateFailedReason, err.Error()) } return true @@ -740,14 +691,14 @@ func (w *NodeController) recordRestoreResults(ctx context.Context, pod *corev1.P if len(failed) == 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, snapshotv1alpha1.RestoreReasonSucceeded, message) != nil + return w.finishRestore(ctx, pod, corev1.ConditionTrue, podcontract.RestoreReasonSucceeded, 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, ", ")) - return w.finishRestore(ctx, pod, corev1.ConditionFalse, snapshotv1alpha1.RestoreReasonPartiallySucceeded, message) != nil + return w.finishRestore(ctx, pod, corev1.ConditionFalse, podcontract.RestoreReasonPartiallySucceeded, 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, snapshotv1alpha1.RestoreReasonFailed, message) != nil + return w.finishRestore(ctx, pod, corev1.ConditionFalse, podcontract.RestoreReasonFailed, message) != nil } // restoreDestination resolves and restores one destination independently of @@ -870,7 +821,7 @@ func (w *NodeController) runRestore(ctx context.Context, pod *corev1.Pod, artifa // completion sentinel proves the operation already finished. func (op *restoreOperation) recoverCompletedRestore(ctx context.Context) (bool, error) { condition := findRestoredCondition(op.pod) - if condition == nil || condition.Status != corev1.ConditionFalse || condition.Reason != snapshotv1alpha1.RestoreReasonInProgress { + if condition == nil || condition.Status != corev1.ConditionFalse || condition.Reason != podcontract.RestoreReasonInProgress { return false, nil } @@ -878,7 +829,7 @@ func (op *restoreOperation) recoverCompletedRestore(ctx context.Context) (bool, if err != nil { return false, fmt.Errorf("resolve restore container before checking completion sentinel: %w", err) } - exists, err := op.controller.controlSentinelExistsFn(hostPID, snapshotv1alpha1.RestoreCompleteFile) + exists, err := op.controller.controlSentinelExistsFn(hostPID, podcontract.RestoreCompleteFile) if err != nil { return false, fmt.Errorf("check restore completion sentinel: %w", err) } @@ -942,7 +893,7 @@ func (op *restoreOperation) completeRestore(ctx context.Context, placeholderHost w := op.controller // Any PID inside the container mount namespace reaches the control // volume through /host/proc//root. - if err := w.writeControlSentinelFn(placeholderHostPID, snapshotv1alpha1.RestoreCompleteFile); err != nil { + if err := w.writeControlSentinelFn(placeholderHostPID, podcontract.RestoreCompleteFile); err != nil { op.log.Error(err, "Failed to write restore-complete sentinel") if killErr := w.sendSignalFn(op.log, placeholderHostPID, syscall.SIGKILL, "restore sentinel failed"); killErr != nil { return errors.Join(fmt.Errorf("failed to write restore-complete sentinel: %w", err), fmt.Errorf("placeholder could not be killed: %w", killErr)) @@ -957,7 +908,7 @@ func (op *restoreOperation) completeRestore(ctx context.Context, placeholderHost // owns only nvidia.com/Restored and does not replace kubelet-owned conditions. func (w *NodeController) applyRestoredCondition(ctx context.Context, pod *corev1.Pod, status corev1.ConditionStatus, reason, message string) error { setPodCondition(&pod.Status, corev1.PodCondition{ - Type: corev1.PodConditionType(snapshotv1alpha1.RestoredCondition), + Type: corev1.PodConditionType(podcontract.RestoredCondition), Status: status, Reason: reason, Message: message, @@ -1068,7 +1019,7 @@ func (w *NodeController) failRestorePod(ctx context.Context, pod *corev1.Pod, ca ctx, pod, corev1.ConditionFalse, - snapshotv1alpha1.RestoreReasonFailed, + podcontract.RestoreReasonFailed, cause.Error(), ) return err != nil @@ -1217,7 +1168,7 @@ func restoreContainerResolveAttemptContext(ctx context.Context, deadlineAt time. func findRestoredCondition(pod *corev1.Pod) *corev1.PodCondition { for i := range pod.Status.Conditions { condition := &pod.Status.Conditions[i] - if condition.Type == corev1.PodConditionType(snapshotv1alpha1.RestoredCondition) { + if condition.Type == corev1.PodConditionType(podcontract.RestoredCondition) { return condition } } @@ -1226,27 +1177,20 @@ func findRestoredCondition(pod *corev1.Pod) *corev1.PodCondition { func restoreInProgress(pod *corev1.Pod) bool { condition := findRestoredCondition(pod) - return condition != nil && condition.Status == corev1.ConditionFalse && condition.Reason == snapshotv1alpha1.RestoreReasonInProgress + return condition != nil && condition.Status == corev1.ConditionFalse && condition.Reason == podcontract.RestoreReasonInProgress } func isRestoreSucceeded(pod *corev1.Pod) bool { - return snapshotv1alpha1.ClassifyRestoreOutcome(pod.Status.Conditions) == snapshotv1alpha1.RestoreOutcomeSucceeded + return podcontract.ClassifyRestoreOutcome(pod.Status.Conditions) == podcontract.RestoreOutcomeSucceeded } // isRestorePartiallySucceeded reports the terminal mixed worker outcome. func isRestorePartiallySucceeded(pod *corev1.Pod) bool { - return snapshotv1alpha1.ClassifyRestoreOutcome(pod.Status.Conditions) == snapshotv1alpha1.RestoreOutcomePartiallySucceeded + return podcontract.ClassifyRestoreOutcome(pod.Status.Conditions) == podcontract.RestoreOutcomePartiallySucceeded } func isRestoreTerminal(pod *corev1.Pod) bool { - switch snapshotv1alpha1.ClassifyRestoreOutcome(pod.Status.Conditions) { - case snapshotv1alpha1.RestoreOutcomeSucceeded, - snapshotv1alpha1.RestoreOutcomeFailed, - snapshotv1alpha1.RestoreOutcomePartiallySucceeded: - return true - default: - return false - } + return podcontract.ClassifyRestoreOutcome(pod.Status.Conditions).Terminal() } func isRestorePodActive(pod *corev1.Pod) bool { diff --git a/agent/internal/controller/controller_test.go b/agent/internal/controller/controller_test.go index 8dd6d5c5..dd29265f 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/podcontract" snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" ) @@ -206,7 +207,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,31 +232,48 @@ func restorePod(annotations map[string]string) *corev1.Pod { }}, }, } + snapshotName, restoreRequested := annotations[podcontract.RestoreFromAnnotation] + _, hasExplicitMapping := annotations[podcontract.RestoreContainerMapAnnotation] + if !restoreRequested || hasExplicitMapping { + return pod + } + shaped, err := podcontract.Build( + pod, + podcontract.Request{ + SnapshotName: snapshotName, + SourceContainer: "main", + }, + podcontract.Options{}, + ) + if err != nil { + panic(err) + } + return shaped } func multiRestorePod() *corev1.Pod { pod := restorePod(map[string]string{ - snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a", - snapshotv1alpha1.RestoreContainerMapAnnotation: "main=engine-0,main=engine-1", + podcontract.RestoreFromAnnotation: "snapshot-a", + podcontract.RestoreContainerMapAnnotation: "main=engine-0,main=engine-1", }) pod.Spec.Volumes = []corev1.Volume{{ - Name: snapshotv1alpha1.SnapshotControlVolumeName, + Name: podcontract.SnapshotControlVolumeName, VolumeSource: corev1.VolumeSource{EmptyDir: &corev1.EmptyDirVolumeSource{}}, }} pod.Spec.Containers = []corev1.Container{ { Name: "engine-0", VolumeMounts: []corev1.VolumeMount{{ - Name: snapshotv1alpha1.SnapshotControlVolumeName, - MountPath: snapshotv1alpha1.SnapshotControlMountPath, + Name: podcontract.SnapshotControlVolumeName, + MountPath: podcontract.SnapshotControlMountPath, SubPath: "engine-0", }}, }, { Name: "engine-1", VolumeMounts: []corev1.VolumeMount{{ - Name: snapshotv1alpha1.SnapshotControlVolumeName, - MountPath: snapshotv1alpha1.SnapshotControlMountPath, + Name: podcontract.SnapshotControlVolumeName, + MountPath: podcontract.SnapshotControlMountPath, SubPath: "engine-1", }}, }, @@ -264,7 +282,22 @@ 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 := podcontract.Build( + pod, + podcontract.Request{ + SnapshotName: "snapshot-a", + SourceContainer: "main", + Mappings: []podcontract.ContainerMapping{ + {Source: "main", Destination: "engine-0"}, + {Source: "main", Destination: "engine-1"}, + }, + }, + podcontract.Options{}, + ) + if err != nil { + panic(err) + } + return shaped } func processQueuedRestorePod(t *testing.T, w *NodeController, pod *corev1.Pod) { @@ -280,7 +313,7 @@ func processQueuedRestorePod(t *testing.T, w *NodeController, pod *corev1.Pod) { func restoredPodCondition(pod *corev1.Pod) *corev1.PodCondition { for i := range pod.Status.Conditions { - if pod.Status.Conditions[i].Type == corev1.PodConditionType(snapshotv1alpha1.RestoredCondition) { + if pod.Status.Conditions[i].Type == corev1.PodConditionType(podcontract.RestoredCondition) { return &pod.Status.Conditions[i] } } @@ -331,7 +364,7 @@ func TestTweakNodePodListOptions(t *testing.T) { } func TestEnqueueRestorePodFiltersAndDeduplicatesEvents(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) w := makeTestController(t, pod) t.Cleanup(w.restoreQueue.ShutDown) @@ -370,7 +403,7 @@ func TestRestoreQueueRequeuesDirtyKeyAfterDone(t *testing.T) { } func TestProcessRestoreQueueItemUsesCachedPod(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) w := makeTestController(t, pod) w.restoreQueue.ShutDown() testClock := clocktesting.NewFakeClock(time.Now()) @@ -393,7 +426,7 @@ func TestProcessRestoreQueueItemUsesCachedPod(t *testing.T) { } func TestPreflightRestore(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) snapshot, content := readySnapshotObjects() w := makeTestController(t, pod, snapshot, content) path, err := nsmount.ResolveArtifactPath(w.config.Storage.BasePath, string(content.UID), "main") @@ -410,7 +443,7 @@ func TestPreflightRestore(t *testing.T) { assert.Equal(t, string(content.UID), plan.artifact.ContentUID) assert.Equal(t, "main", plan.artifact.SourceContainerName) assert.Equal(t, path, plan.artifact.Path) - assert.Equal(t, []snapshotv1alpha1.RestoreContainerMapping{{Source: "main", Destination: "main"}}, plan.mappings) + assert.Equal(t, []podcontract.ContainerMapping{{Source: "main", Destination: "main"}}, plan.mappings) } func TestReconcileRestorePodRunsMappedDestinationsConcurrently(t *testing.T) { @@ -433,7 +466,7 @@ func TestReconcileRestorePodRunsMappedDestinationsConcurrently(t *testing.T) { } sentinels := make(chan int, 2) w.writeControlSentinelFn = func(pid int, name string) error { - assert.Equal(t, snapshotv1alpha1.RestoreCompleteFile, name) + assert.Equal(t, podcontract.RestoreCompleteFile, name) sentinels <- pid return nil } @@ -516,19 +549,19 @@ func TestIsRestoreTerminalRequiresKnownTerminalOutcome(t *testing.T) { reason string want bool }{ - {name: "succeeded", status: corev1.ConditionTrue, reason: snapshotv1alpha1.RestoreReasonSucceeded, want: true}, - {name: "failed", status: corev1.ConditionFalse, reason: snapshotv1alpha1.RestoreReasonFailed, want: true}, - {name: "partially succeeded", status: corev1.ConditionFalse, reason: snapshotv1alpha1.RestoreReasonPartiallySucceeded, want: true}, - {name: "in progress", status: corev1.ConditionFalse, reason: snapshotv1alpha1.RestoreReasonInProgress}, + {name: "succeeded", status: corev1.ConditionTrue, reason: podcontract.RestoreReasonSucceeded, want: true}, + {name: "failed", status: corev1.ConditionFalse, reason: podcontract.RestoreReasonFailed, want: true}, + {name: "partially succeeded", status: corev1.ConditionFalse, reason: podcontract.RestoreReasonPartiallySucceeded, want: true}, + {name: "in progress", status: corev1.ConditionFalse, reason: podcontract.RestoreReasonInProgress}, {name: "unrecognized reason", status: corev1.ConditionFalse, reason: "RestoreIncompatible"}, - {name: "unknown status", status: corev1.ConditionUnknown, reason: snapshotv1alpha1.RestoreReasonSucceeded}, + {name: "unknown status", status: corev1.ConditionUnknown, reason: podcontract.RestoreReasonSucceeded}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { pod := restorePod(nil) pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{ - Type: corev1.PodConditionType(snapshotv1alpha1.RestoredCondition), + Type: corev1.PodConditionType(podcontract.RestoredCondition), Status: test.status, Reason: test.reason, }) @@ -550,7 +583,7 @@ func TestRestorePodContainersKeepsAggregateInProgressWhileDestinationIsPending(t } plan := &restorePlan{ artifact: &restoreArtifact{SnapshotName: "snapshot-a", ContentUID: "content-uid", SourceContainerName: "main"}, - mappings: []snapshotv1alpha1.RestoreContainerMapping{ + mappings: []podcontract.ContainerMapping{ {Source: "main", Destination: "engine-0"}, {Source: "main", Destination: "engine-1"}, }, @@ -566,7 +599,7 @@ func TestRestorePodContainersKeepsAggregateInProgressWhileDestinationIsPending(t func TestPreflightRestoreRejectsInvalidMappingBeforeExecution(t *testing.T) { pod := multiRestorePod() - pod.Annotations[snapshotv1alpha1.RestoreContainerMapAnnotation] = "worker=engine-0" + pod.Annotations[podcontract.RestoreContainerMapAnnotation] = "worker=engine-0" snapshot, content := readySnapshotObjects() w := makeTestController(t, pod, snapshot, content) restoreCalls := 0 @@ -600,10 +633,50 @@ 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 + }, + } + for name, mutate := range tests { + t.Run(name, func(t *testing.T) { + pod := restorePod(map[string]string{podcontract.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 TestValidateRestoreTargetIgnoresWorkloadStartupProbe(t *testing.T) { + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) + pod.Spec.Containers[0].StartupProbe = &corev1.Probe{ + ProbeHandler: corev1.ProbeHandler{Exec: &corev1.ExecAction{Command: []string{"true"}}}, + 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, []podcontract.ContainerMapping{{Source: "main", Destination: "main"}}, mappings) +} + func TestPreflightRestoreRetriesInProgressCondition(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{ - Type: corev1.PodConditionType(snapshotv1alpha1.RestoredCondition), Status: corev1.ConditionFalse, Reason: "RestoreInProgress", + Type: corev1.PodConditionType(podcontract.RestoredCondition), Status: corev1.ConditionFalse, Reason: "RestoreInProgress", }) snapshot, content := readySnapshotObjects() w := makeTestController(t, pod, snapshot, content) @@ -619,7 +692,7 @@ func TestPreflightRestoreRetriesInProgressCondition(t *testing.T) { } func TestPreflightRestorePendingStates(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) t.Run("missing snapshot", func(t *testing.T) { w := makeTestController(t, pod) artifact, err := w.preflightRestore(context.Background(), pod) @@ -651,9 +724,9 @@ func TestPreflightRestorePendingStates(t *testing.T) { } func TestPreflightRestoreFailsWhenInProgressSnapshotDisappears(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{ - Type: corev1.PodConditionType(snapshotv1alpha1.RestoredCondition), Status: corev1.ConditionFalse, Reason: snapshotv1alpha1.RestoreReasonInProgress, + Type: corev1.PodConditionType(podcontract.RestoredCondition), Status: corev1.ConditionFalse, Reason: podcontract.RestoreReasonInProgress, }) w := makeTestController(t, pod) @@ -667,7 +740,7 @@ func TestPreflightRestoreFailsWhenInProgressSnapshotDisappears(t *testing.T) { } func TestRestorePreflightAPIErrorsUseStablePendingMessages(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) snapshot, _ := readySnapshotObjects() tests := []struct { @@ -719,7 +792,7 @@ func TestRestorePreflightAPIErrorsUseStablePendingMessages(t *testing.T) { } func TestDeletingSnapshotDependencyFailsRestoreBeforeLaunch(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) snapshot, content := readySnapshotObjects() now := metav1.Now() snapshot.DeletionTimestamp = &now @@ -741,7 +814,7 @@ func TestDeletingSnapshotDependencyFailsRestoreBeforeLaunch(t *testing.T) { } func TestDeletingContentDependencyFailsRestoreBeforeLaunch(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) snapshot, content := readySnapshotObjects() now := metav1.Now() content.DeletionTimestamp = &now @@ -763,7 +836,7 @@ func TestDeletingContentDependencyFailsRestoreBeforeLaunch(t *testing.T) { } func TestPreflightRestoreWaitsForPodIPBeforeTCPRestore(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) snapshot, content := readySnapshotObjects() w := makeTestController(t, pod, snapshot, content) w.config.CRIU.TcpEstablished = true @@ -778,7 +851,7 @@ func TestPreflightRestoreWaitsForPodIPBeforeTCPRestore(t *testing.T) { } func TestReconcileRestorePodReportsPendingPreflightConditionAndEvent(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) w := makeTestController(t, pod) requeue := w.reconcileRestorePod(context.Background(), pod) @@ -791,9 +864,9 @@ func TestReconcileRestorePodReportsPendingPreflightConditionAndEvent(t *testing. } func TestPendingDependencyDoesNotOverwriteRestoreInProgress(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{ - Type: corev1.PodConditionType(snapshotv1alpha1.RestoredCondition), Status: corev1.ConditionFalse, Reason: snapshotv1alpha1.RestoreReasonInProgress, + Type: corev1.PodConditionType(podcontract.RestoredCondition), Status: corev1.ConditionFalse, Reason: podcontract.RestoreReasonInProgress, }) w := makeTestController(t, pod) @@ -801,12 +874,12 @@ func TestPendingDependencyDoesNotOverwriteRestoreInProgress(t *testing.T) { assert.True(t, requeue) assert.False(t, hasPodStatusApply(w)) - assert.Equal(t, snapshotv1alpha1.RestoreReasonInProgress, restoredPodCondition(pod).Reason) + assert.Equal(t, podcontract.RestoreReasonInProgress, restoredPodCondition(pod).Reason) assert.True(t, sawEventReason(w.clientset.(*fake.Clientset), "ArtifactPending")) } func TestProcessRestoreQueueItemReportsNonRunningPhaseAsFailed(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) pod.Status.Phase = corev1.PodFailed w := makeTestController(t, pod) @@ -815,11 +888,11 @@ func TestProcessRestoreQueueItemReportsNonRunningPhaseAsFailed(t *testing.T) { payload := string(lastPodStatusApply(t, w).GetPatch()) assert.Contains(t, payload, `"reason":"RestoreFailed"`) assert.Contains(t, payload, "phase Failed") - assert.True(t, sawEventReason(w.clientset.(*fake.Clientset), snapshotv1alpha1.RestoreReasonFailed)) + assert.True(t, sawEventReason(w.clientset.(*fake.Clientset), podcontract.RestoreReasonFailed)) } func TestContainerPollingDoesNotSetInProgressBeforeExecution(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) pod.Status.ContainerStatuses[0].ContainerID = "" snapshot, content := readySnapshotObjects() w := makeTestController(t, pod, snapshot, content) @@ -839,7 +912,7 @@ func TestContainerPollingDoesNotSetInProgressBeforeExecution(t *testing.T) { } func TestPreflightRestoreValidatesContentBacklink(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) tests := map[string]func(*snapshotv1alpha1.PodSnapshotContent){ "namespace": func(content *snapshotv1alpha1.PodSnapshotContent) { content.Spec.PodSnapshotRef.Namespace = "other" @@ -867,7 +940,7 @@ func TestPreflightRestoreValidatesContentBacklink(t *testing.T) { } func TestPreflightRestoreTerminalStates(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) tests := map[string]struct { mutateSnapshot func(*snapshotv1alpha1.PodSnapshot) mutateContent func(*snapshotv1alpha1.PodSnapshotContent) @@ -920,7 +993,7 @@ func TestPreflightRestoreTerminalStates(t *testing.T) { } func TestReconcileRestorePodReportsInvalidBacklinkAsFailedCondition(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) snapshot, content := readySnapshotObjects() content.Spec.PodSnapshotRef.UID = "other-uid" w := makeTestController(t, pod, snapshot, content) @@ -933,10 +1006,10 @@ func TestReconcileRestorePodReportsInvalidBacklinkAsFailedCondition(t *testing.T } func TestProcessRestoreQueueItemEmitsEventWhenRestoreAlreadyCompleted(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) pod.Finalizers = []string{restorePodFinalizer} pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{ - Type: corev1.PodConditionType(snapshotv1alpha1.RestoredCondition), Status: corev1.ConditionTrue, Reason: "RestoreSucceeded", + Type: corev1.PodConditionType(podcontract.RestoredCondition), Status: corev1.ConditionTrue, Reason: "RestoreSucceeded", }) w := makeTestController(t, pod) @@ -949,13 +1022,13 @@ func TestProcessRestoreQueueItemEmitsEventWhenRestoreAlreadyCompleted(t *testing } func TestProcessRestoreQueueItemIgnoresFailedRestoreDuringPreflight(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) pod.Finalizers = []string{restorePodFinalizer} pod.Status.Phase = corev1.PodFailed pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{ - Type: corev1.PodConditionType(snapshotv1alpha1.RestoredCondition), + Type: corev1.PodConditionType(podcontract.RestoredCondition), Status: corev1.ConditionFalse, - Reason: snapshotv1alpha1.RestoreReasonFailed, + Reason: podcontract.RestoreReasonFailed, Message: "original restore failure", }) w := makeTestController(t, pod) @@ -977,7 +1050,7 @@ func TestProcessRestoreQueueItemIgnoresFailedRestoreDuringPreflight(t *testing.T assert.Zero(t, getCalls, "failed restore preflight must not read PodSnapshot or PodSnapshotContent") condition := restoredPodCondition(pod) require.NotNil(t, condition) - assert.Equal(t, snapshotv1alpha1.RestoreReasonFailed, condition.Reason) + assert.Equal(t, podcontract.RestoreReasonFailed, condition.Reason) assert.Equal(t, "original restore failure", condition.Message) event := eventForReason(w.clientset.(*fake.Clientset), "RestoreAlreadyFailed") require.NotNil(t, event) @@ -989,12 +1062,12 @@ func TestProcessRestoreQueueItemIgnoresFailedRestoreDuringPreflight(t *testing.T } func TestDeletingInProgressRestoreRemovesFinalizer(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) pod.Finalizers = []string{restorePodFinalizer} now := metav1.Now() pod.DeletionTimestamp = &now pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{ - Type: corev1.PodConditionType(snapshotv1alpha1.RestoredCondition), Status: corev1.ConditionFalse, Reason: snapshotv1alpha1.RestoreReasonInProgress, + Type: corev1.PodConditionType(podcontract.RestoredCondition), Status: corev1.ConditionFalse, Reason: podcontract.RestoreReasonInProgress, }) w := makeTestController(t, pod) @@ -1007,7 +1080,7 @@ func TestDeletingInProgressRestoreRemovesFinalizer(t *testing.T) { } func TestReconcileRestorePodEmitsEventWhenStatusUpdateFails(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) w := makeTestController(t, pod) w.clientset.(*fake.Clientset).PrependReactor("patch", "pods", func(clientgotesting.Action) (bool, runtime.Object, error) { return true, nil, errors.New("status patch failed") @@ -1021,8 +1094,8 @@ func TestReconcileRestorePodEmitsEventWhenStatusUpdateFails(t *testing.T) { func TestPreflightRestoreAllowsUnrelatedAnnotations(t *testing.T) { pod := restorePod(map[string]string{ - snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a", - "example.com/team": "inference", + podcontract.RestoreFromAnnotation: "snapshot-a", + "example.com/team": "inference", }) snapshot, content := readySnapshotObjects() w := makeTestController(t, pod, snapshot, content) @@ -1038,7 +1111,7 @@ func TestPreflightRestoreAllowsUnrelatedAnnotations(t *testing.T) { } func TestReconcileRestorePodRunsPreflightOnce(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) snapshot, content := readySnapshotObjects() w := makeTestController(t, pod, snapshot, content) path, err := nsmount.ResolveArtifactPath(w.config.Storage.BasePath, string(content.UID), "main") @@ -1067,7 +1140,7 @@ func TestReconcileRestorePodRunsPreflightOnce(t *testing.T) { } func TestRestoreFinalizerIsNotAddedWhilePreflightIsPending(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "missing-snapshot"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "missing-snapshot"}) w := makeTestController(t, pod) processQueuedRestorePod(t, w, pod) @@ -1087,7 +1160,7 @@ func TestRestoreFinalizerIsNotAddedWhilePreflightIsPending(t *testing.T) { } func TestRestoreFinalizerProtectsExecutionAndIsRemovedAfterSuccess(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) snapshot, content := readySnapshotObjects() w := makeTestController(t, pod, snapshot, content) path, err := nsmount.ResolveArtifactPath(w.config.Storage.BasePath, string(content.UID), "main") @@ -1116,7 +1189,7 @@ func TestRestoreFinalizerProtectsExecutionAndIsRemovedAfterSuccess(t *testing.T) } func TestRestoreStatusRetryUsesCompletionSentinelWithoutReplayingRestore(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) snapshot, content := readySnapshotObjects() w := makeTestController(t, pod, snapshot, content) w.runtime = &fakeRuntime{resolveContainerPID: 4242} @@ -1132,13 +1205,13 @@ func TestRestoreStatusRetryUsesCompletionSentinelWithoutReplayingRestore(t *test sentinelWritten := false w.writeControlSentinelFn = func(pid int, name string) error { assert.Equal(t, 4242, pid) - assert.Equal(t, snapshotv1alpha1.RestoreCompleteFile, name) + assert.Equal(t, podcontract.RestoreCompleteFile, name) sentinelWritten = true return nil } w.controlSentinelExistsFn = func(pid int, name string) (bool, error) { assert.Equal(t, 4242, pid) - assert.Equal(t, snapshotv1alpha1.RestoreCompleteFile, name) + assert.Equal(t, podcontract.RestoreCompleteFile, name) return sentinelWritten, nil } statusPatches := 0 @@ -1164,7 +1237,7 @@ func TestRestoreStatusRetryUsesCompletionSentinelWithoutReplayingRestore(t *test assert.True(t, hasFinalizer(live, restorePodFinalizer)) condition := restoredPodCondition(live) require.NotNil(t, condition) - assert.Equal(t, snapshotv1alpha1.RestoreReasonInProgress, condition.Reason) + assert.Equal(t, podcontract.RestoreReasonInProgress, condition.Reason) processQueuedRestorePod(t, w, live) assert.Equal(t, 1, restoreCalls, "completion sentinel must prevent CRIU replay") @@ -1176,11 +1249,11 @@ func TestRestoreStatusRetryUsesCompletionSentinelWithoutReplayingRestore(t *test condition = restoredPodCondition(live) require.NotNil(t, condition) assert.Equal(t, corev1.ConditionTrue, condition.Status) - assert.Equal(t, snapshotv1alpha1.RestoreReasonSucceeded, condition.Reason) + assert.Equal(t, podcontract.RestoreReasonSucceeded, condition.Reason) } func TestRestoreFinalizerRemovalRetriesWithoutReplayingRestore(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) snapshot, content := readySnapshotObjects() w := makeTestController(t, pod, snapshot, content) path, err := nsmount.ResolveArtifactPath(w.config.Storage.BasePath, string(content.UID), "main") @@ -1218,7 +1291,7 @@ func TestRestoreFinalizerRemovalRetriesWithoutReplayingRestore(t *testing.T) { } func TestApplyRestoredConditionUsesServerSideApply(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) w := makeTestController(t, pod) err := w.applyRestoredCondition(context.Background(), pod, corev1.ConditionFalse, "SnapshotPending", "waiting") require.NoError(t, err) @@ -1235,9 +1308,9 @@ func TestApplyRestoredConditionUsesServerSideApply(t *testing.T) { func TestApplyRestoredConditionPreservesTransitionTimeForSameStatus(t *testing.T) { transition := metav1.NewTime(time.Unix(123, 0)) - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{ - Type: corev1.PodConditionType(snapshotv1alpha1.RestoredCondition), Status: corev1.ConditionFalse, LastTransitionTime: transition, + Type: corev1.PodConditionType(podcontract.RestoredCondition), Status: corev1.ConditionFalse, LastTransitionTime: transition, }) w := makeTestController(t, pod) err := w.applyRestoredCondition(context.Background(), pod, corev1.ConditionFalse, "ArtifactPending", "waiting") @@ -1258,7 +1331,7 @@ func TestInFlightKeyIsDeduplicatedForCapture(t *testing.T) { } func TestRunRestoreCleanupFailureStillCompletesRestore(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) w := makeTestController(t, pod) artifactPath := t.TempDir() artifact := &restoreArtifact{ @@ -1290,7 +1363,7 @@ func TestRunRestoreCleanupFailureStillCompletesRestore(t *testing.T) { } func TestRunRestoreRetriesFullRestoreUntilFailureCleanupSucceeds(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) w := makeTestController(t, pod) w.runtime = &fakeRuntime{resolveContainerPID: 4242} artifact := &restoreArtifact{SnapshotName: "snapshot-a", ContentUID: "content-uid", SourceContainerName: "main"} @@ -1318,7 +1391,7 @@ func TestRunRestoreRetriesFullRestoreUntilFailureCleanupSucceeds(t *testing.T) { } func TestRunRestoreFailureKillsPlaceholder(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) pod.Finalizers = []string{restorePodFinalizer} w := makeTestController(t, pod) w.runtime = &fakeRuntime{resolveContainerPID: 4242} @@ -1340,17 +1413,17 @@ func TestRunRestoreFailureKillsPlaceholder(t *testing.T) { } func TestRunRestoreFinalizesExistingCompletionSentinelWithoutReplay(t *testing.T) { - pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"}) + pod := restorePod(map[string]string{podcontract.RestoreFromAnnotation: "snapshot-a"}) pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{ - Type: corev1.PodConditionType(snapshotv1alpha1.RestoredCondition), + Type: corev1.PodConditionType(podcontract.RestoredCondition), Status: corev1.ConditionFalse, - Reason: snapshotv1alpha1.RestoreReasonInProgress, + Reason: podcontract.RestoreReasonInProgress, }) w := makeTestController(t, pod) w.runtime = &fakeRuntime{resolveContainerPID: 4242} w.controlSentinelExistsFn = func(pid int, name string) (bool, error) { assert.Equal(t, 4242, pid) - assert.Equal(t, snapshotv1alpha1.RestoreCompleteFile, name) + assert.Equal(t, podcontract.RestoreCompleteFile, name) return true, nil } w.restoreFn = func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (int, error) { diff --git a/agent/internal/cuda/job.go b/agent/internal/cuda/job.go index 25ef350a..855d5250 100644 --- a/agent/internal/cuda/job.go +++ b/agent/internal/cuda/job.go @@ -10,7 +10,7 @@ import ( "path/filepath" "strings" - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" + "github.com/ai-dynamo/snapshot/api/podcontract" "golang.org/x/sys/unix" ) @@ -24,12 +24,12 @@ const JobFileEnv = "CUDA_CHECKPOINT_JOB_FILE" // launch wrapper persists the driver-created file at a fixed path before // starting the workload. func StageJobFile(sourceRootPath, checkpointDir string, sourceGPUCount int) (string, error) { - sourcePath := filepath.Join(sourceRootPath, strings.TrimPrefix(snapshotv1alpha1.CUDAJobFilePath, string(os.PathSeparator))) - destinationPath := filepath.Join(checkpointDir, snapshotv1alpha1.CUDAJobFileName) + sourcePath := filepath.Join(sourceRootPath, strings.TrimPrefix(podcontract.CUDAJobFilePath, string(os.PathSeparator))) + destinationPath := filepath.Join(checkpointDir, podcontract.CUDAJobFileName) if err := copyJobFile(sourcePath, destinationPath); err != nil { if os.IsNotExist(err) { if sourceGPUCount > 1 { - return "", fmt.Errorf("multi-GPU CUDA source is missing %s; source must be launched under cuda-checkpoint --launch-job", snapshotv1alpha1.CUDAJobFilePath) + return "", fmt.Errorf("multi-GPU CUDA source is missing %s; source must be launched under cuda-checkpoint --launch-job", podcontract.CUDAJobFilePath) } return "", nil } @@ -45,7 +45,7 @@ func refreshJobFileArtifact(liveJobFile, checkpointDir string) error { if liveJobFile == "" { return nil } - destinationPath := filepath.Join(checkpointDir, snapshotv1alpha1.CUDAJobFileName) + destinationPath := filepath.Join(checkpointDir, podcontract.CUDAJobFileName) if err := prepareLiveJobFile(liveJobFile, destinationPath); err != nil { return fmt.Errorf("refresh CUDA checkpoint job file: %w", err) } @@ -58,15 +58,15 @@ func refreshJobFileArtifact(liveJobFile, checkpointDir string) error { // The returned path is the per-restore working copy that CUDA helpers must use; // the staged artifact remains immutable so it can seed later restores. func PrepareLiveJobFile(stagedJobFile string) (string, error) { - if err := prepareLiveJobFile(stagedJobFile, snapshotv1alpha1.CUDAJobFilePath); err != nil { + if err := prepareLiveJobFile(stagedJobFile, podcontract.CUDAJobFilePath); err != nil { return "", err } - return snapshotv1alpha1.CUDAJobFilePath, nil + return podcontract.CUDAJobFilePath, nil } // JobFileFromCheckpoint returns the staged job file when an artifact contains one. func JobFileFromCheckpoint(checkpointDir string) (string, error) { - jobFile := filepath.Join(checkpointDir, snapshotv1alpha1.CUDAJobFileName) + jobFile := filepath.Join(checkpointDir, podcontract.CUDAJobFileName) info, err := os.Lstat(jobFile) if os.IsNotExist(err) { return "", nil diff --git a/agent/internal/cuda/job_test.go b/agent/internal/cuda/job_test.go index 9c1b210d..3255aaad 100644 --- a/agent/internal/cuda/job_test.go +++ b/agent/internal/cuda/job_test.go @@ -13,13 +13,13 @@ import ( "github.com/go-logr/logr" "golang.org/x/sys/unix" - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" + "github.com/ai-dynamo/snapshot/api/podcontract" ) func TestStageJobFile(t *testing.T) { sourceRoot := t.TempDir() checkpointDir := t.TempDir() - jobFile := snapshotv1alpha1.CUDAJobFilePath + jobFile := podcontract.CUDAJobFilePath if err := os.MkdirAll(filepath.Join(sourceRoot, filepath.Dir(jobFile)), 0700); err != nil { t.Fatal(err) } @@ -35,7 +35,7 @@ func TestStageJobFile(t *testing.T) { if helperJobFile != wantHelperJobFile { t.Fatalf("StageJobFile() = %q, want %q", helperJobFile, wantHelperJobFile) } - artifact := filepath.Join(checkpointDir, snapshotv1alpha1.CUDAJobFileName) + artifact := filepath.Join(checkpointDir, podcontract.CUDAJobFileName) content, err := os.ReadFile(artifact) if err != nil { t.Fatal(err) @@ -48,7 +48,7 @@ func TestStageJobFile(t *testing.T) { func TestStageJobFileRejectsSymlink(t *testing.T) { sourceRoot := t.TempDir() checkpointDir := t.TempDir() - jobFile := snapshotv1alpha1.CUDAJobFilePath + jobFile := podcontract.CUDAJobFilePath if err := os.MkdirAll(filepath.Join(sourceRoot, filepath.Dir(jobFile)), 0700); err != nil { t.Fatal(err) } @@ -65,15 +65,15 @@ func TestStageJobFileRejectsSymlink(t *testing.T) { if err == nil { t.Fatal("expected symlink source to be rejected") } - if _, statErr := os.Stat(filepath.Join(checkpointDir, snapshotv1alpha1.CUDAJobFileName)); !os.IsNotExist(statErr) { + if _, statErr := os.Stat(filepath.Join(checkpointDir, podcontract.CUDAJobFileName)); !os.IsNotExist(statErr) { t.Fatalf("staged file exists after rejected symlink: %v", statErr) } } func TestRefreshJobFileArtifactCapturesPostCheckpointState(t *testing.T) { checkpointDir := t.TempDir() - live := filepath.Join(t.TempDir(), snapshotv1alpha1.CUDAJobFileName) - artifact := filepath.Join(checkpointDir, snapshotv1alpha1.CUDAJobFileName) + live := filepath.Join(t.TempDir(), podcontract.CUDAJobFileName) + artifact := filepath.Join(checkpointDir, podcontract.CUDAJobFileName) if err := os.WriteFile(live, []byte("pre-checkpoint-state"), 0600); err != nil { t.Fatal(err) } @@ -147,7 +147,7 @@ if [ "$action" = checkpoint ]; then printf '|%s' "$pid" >> "$job_file"; fi if err := os.WriteFile(liveJobFile, []byte("initial"), 0600); err != nil { t.Fatal(err) } - if err := os.WriteFile(filepath.Join(checkpointDir, snapshotv1alpha1.CUDAJobFileName), []byte("validation-copy"), 0600); err != nil { + if err := os.WriteFile(filepath.Join(checkpointDir, podcontract.CUDAJobFileName), []byte("validation-copy"), 0600); err != nil { t.Fatal(err) } t.Setenv("DYNAMO_TEST_TRACE", trace) @@ -163,7 +163,7 @@ if [ "$action" = checkpoint ]; then printf '|%s' "$pid" >> "$job_file"; fi if got, want := string(traceContent), "lock 101\nlock 202\ncheckpoint 101\ncheckpoint 202\n"; got != want { t.Fatalf("helper call order = %q, want %q", got, want) } - artifact, err := os.ReadFile(filepath.Join(checkpointDir, snapshotv1alpha1.CUDAJobFileName)) + artifact, err := os.ReadFile(filepath.Join(checkpointDir, podcontract.CUDAJobFileName)) if err != nil { t.Fatal(err) } @@ -217,7 +217,7 @@ func TestJobFileFromCheckpointRejectsSymlink(t *testing.T) { if err := os.WriteFile(target, []byte("job-state"), 0600); err != nil { t.Fatal(err) } - if err := os.Symlink(target, filepath.Join(checkpointDir, snapshotv1alpha1.CUDAJobFileName)); err != nil { + if err := os.Symlink(target, filepath.Join(checkpointDir, podcontract.CUDAJobFileName)); err != nil { t.Fatal(err) } @@ -228,8 +228,8 @@ func TestJobFileFromCheckpointRejectsSymlink(t *testing.T) { } func TestPrepareLiveJobFileReplacesMutatedState(t *testing.T) { - staged := filepath.Join(t.TempDir(), snapshotv1alpha1.CUDAJobFileName) - live := filepath.Join(t.TempDir(), snapshotv1alpha1.CUDAJobFileName) + staged := filepath.Join(t.TempDir(), podcontract.CUDAJobFileName) + live := filepath.Join(t.TempDir(), podcontract.CUDAJobFileName) if err := os.WriteFile(staged, []byte("capture-time-state"), 0600); err != nil { t.Fatal(err) } @@ -257,12 +257,12 @@ func TestPrepareLiveJobFileReplacesMutatedState(t *testing.T) { } func TestPrepareLiveJobFileKeepsRestoreTargetsIsolated(t *testing.T) { - staged := filepath.Join(t.TempDir(), snapshotv1alpha1.CUDAJobFileName) + staged := filepath.Join(t.TempDir(), podcontract.CUDAJobFileName) if err := os.WriteFile(staged, []byte("capture-time-state"), 0600); err != nil { t.Fatal(err) } - first := filepath.Join(t.TempDir(), snapshotv1alpha1.CUDAJobFileName) - second := filepath.Join(t.TempDir(), snapshotv1alpha1.CUDAJobFileName) + first := filepath.Join(t.TempDir(), podcontract.CUDAJobFileName) + second := filepath.Join(t.TempDir(), podcontract.CUDAJobFileName) if err := prepareLiveJobFile(staged, first); err != nil { t.Fatal(err) } @@ -282,7 +282,7 @@ func TestPrepareLiveJobFileKeepsRestoreTargetsIsolated(t *testing.T) { } func TestPrepareLiveJobFileRejectsSymlinkDestination(t *testing.T) { - staged := filepath.Join(t.TempDir(), snapshotv1alpha1.CUDAJobFileName) + staged := filepath.Join(t.TempDir(), podcontract.CUDAJobFileName) if err := os.WriteFile(staged, []byte("capture-time-state"), 0600); err != nil { t.Fatal(err) } @@ -291,7 +291,7 @@ func TestPrepareLiveJobFileRejectsSymlinkDestination(t *testing.T) { if err := os.WriteFile(target, []byte("must-not-change"), 0600); err != nil { t.Fatal(err) } - live := filepath.Join(destinationDir, snapshotv1alpha1.CUDAJobFileName) + live := filepath.Join(destinationDir, podcontract.CUDAJobFileName) if err := os.Symlink(target, live); err != nil { t.Fatal(err) } diff --git a/agent/internal/cuda/shim_restore_job_file_test.go b/agent/internal/cuda/shim_restore_job_file_test.go index e7d0d469..1a149db5 100644 --- a/agent/internal/cuda/shim_restore_job_file_test.go +++ b/agent/internal/cuda/shim_restore_job_file_test.go @@ -11,13 +11,13 @@ import ( "github.com/go-logr/logr" - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" + "github.com/ai-dynamo/snapshot/api/podcontract" ) func TestRunActionInheritsJobFileEnvironment(t *testing.T) { trace := filepath.Join(t.TempDir(), "trace") installFakeCUDAHelper(t, "printf '%s' \"$CUDA_CHECKPOINT_JOB_FILE\" > \""+trace+"\"\n") - t.Setenv(JobFileEnv, snapshotv1alpha1.CUDAJobFilePath) + t.Setenv(JobFileEnv, podcontract.CUDAJobFilePath) if err := runAction(context.Background(), 11, actionRestore, "", cudaCheckpointHelperBinary, logr.Discard()); err != nil { t.Fatalf("runAction() error = %v", err) @@ -26,7 +26,7 @@ func TestRunActionInheritsJobFileEnvironment(t *testing.T) { if err != nil { t.Fatal(err) } - if got, want := string(content), snapshotv1alpha1.CUDAJobFilePath; got != want { + if got, want := string(content), podcontract.CUDAJobFilePath; got != want { t.Fatalf("helper environment = %q, want %q", got, want) } } diff --git a/agent/internal/executor/nsrestore.go b/agent/internal/executor/nsrestore.go index 7ab89d3f..1aa97f23 100644 --- a/agent/internal/executor/nsrestore.go +++ b/agent/internal/executor/nsrestore.go @@ -18,7 +18,7 @@ import ( "github.com/ai-dynamo/snapshot/agent/internal/cuda" snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime" "github.com/ai-dynamo/snapshot/agent/internal/types" - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" + "github.com/ai-dynamo/snapshot/api/podcontract" ) // RestoreOptions holds configuration for an in-namespace restore. @@ -183,7 +183,7 @@ func executeRestore( // leftover from an earlier incarnation cannot release the restored process // before this CRIU/CUDA attempt finishes. A missing file is already gone; // a missing mount is a hard error. - if err := snapshotruntime.RemoveControlSentinel(snapshotv1alpha1.SnapshotControlMountPath, snapshotv1alpha1.RestoreCompleteFile); err != nil { + if err := snapshotruntime.RemoveControlSentinel(podcontract.SnapshotControlMountPath, podcontract.RestoreCompleteFile); err != nil { return nil, 0, nil, fmt.Errorf("remove stale restore-complete sentinel: %w", err) } diff --git a/agent/internal/runtime/control.go b/agent/internal/runtime/control.go index 48565df0..27a6ab09 100644 --- a/agent/internal/runtime/control.go +++ b/agent/internal/runtime/control.go @@ -9,7 +9,7 @@ import ( "path/filepath" "strconv" - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" + "github.com/ai-dynamo/snapshot/api/podcontract" ) // WriteControlSentinel writes a sentinel file into the workload container's @@ -28,7 +28,7 @@ func WriteControlSentinel(hostPID int, name string) error { if hostPID <= 0 { return fmt.Errorf("invalid host PID %d for control sentinel %q", hostPID, name) } - dir := filepath.Join(HostProcPath, strconv.Itoa(hostPID), "root", snapshotv1alpha1.SnapshotControlMountPath) + dir := filepath.Join(HostProcPath, strconv.Itoa(hostPID), "root", podcontract.SnapshotControlMountPath) return writeSentinelInDir(dir, name) } @@ -40,7 +40,7 @@ func ControlSentinelExists(hostPID int, name string) (bool, error) { if hostPID <= 0 { return false, fmt.Errorf("invalid host PID %d for control sentinel %q", hostPID, name) } - dir := filepath.Join(HostProcPath, strconv.Itoa(hostPID), "root", snapshotv1alpha1.SnapshotControlMountPath) + dir := filepath.Join(HostProcPath, strconv.Itoa(hostPID), "root", podcontract.SnapshotControlMountPath) return controlSentinelExistsInDir(dir, name) } diff --git a/api/podcontract/doc.go b/api/podcontract/doc.go new file mode 100644 index 00000000..a115a7c4 --- /dev/null +++ b/api/podcontract/doc.go @@ -0,0 +1,6 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Package podcontract defines Snapshot's workload Pod contract and provides +// helpers to build, validate, and interpret restore Pods. +package podcontract diff --git a/api/podcontract/protocol.go b/api/podcontract/protocol.go new file mode 100644 index 00000000..31bb60c8 --- /dev/null +++ b/api/podcontract/protocol.go @@ -0,0 +1,165 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package podcontract + +import ( + "fmt" + "strings" + + "k8s.io/apimachinery/pkg/util/validation" +) + +// Snapshot workload Pod annotations and control protocol. +const ( + // RestoreFromAnnotation names the PodSnapshot to restore in the Pod's + // namespace. + RestoreFromAnnotation = "nvidia.com/restore-from" + + // RestoreContainerMapAnnotation optionally maps the single captured source + // container to one or more restore destinations. Its value is a + // comma-separated list of source=destination pairs. When absent, restore + // uses the captured container name as the destination. + RestoreContainerMapAnnotation = "nvidia.com/restore-container-map" + + // DefaultSeccompLocalhostProfile is the kubelet-local profile installed by + // the Snapshot Helm chart to block io_uring for CRIU. + DefaultSeccompLocalhostProfile = "profiles/block-iouring.json" + + // SnapshotControlVolumeName is the per-Pod emptyDir used to carry capture and + // restore lifecycle files between Snapshot and the workload. + SnapshotControlVolumeName = "snapshot-control" + + // SnapshotControlMountPath is where the control volume is mounted inside a + // Snapshot-managed workload container. + SnapshotControlMountPath = "/snapshot-control" + + // SnapshotControlDirEnv is the canonical environment variable exposing the + // control mount path to the workload. + SnapshotControlDirEnv = "SNAPSHOT_CONTROL_DIR" + + // LegacySnapshotControlDirEnv is the deprecated environment variable + // exposing the control mount path to the workload. + // + // Deprecated: use SnapshotControlDirEnv instead. + LegacySnapshotControlDirEnv = "DYN_SNAPSHOT_CONTROL_DIR" + + // ReadyForSnapshotFile is written by the workload when the model is loaded + // and the workload is ready for capture. The source Pod's kubelet readiness + // probe observes it through the control volume. + ReadyForSnapshotFile = "ready-for-snapshot" + + // CUDAJobFileName is the stable name under which the + // cuda-checkpoint-helper launch-job wrapper persists the CUDA checkpoint job + // file inside the control volume. + CUDAJobFileName = "cuda-checkpoint-job" + + // CUDAJobFilePath is the full stable path to the persisted CUDA checkpoint + // job file. + CUDAJobFilePath = SnapshotControlMountPath + "/" + CUDAJobFileName + + // RestoreStandbyModeEnv asks standby-aware workload entrypoints to remain + // inert until Snapshot replaces them with restored processes. Snapshot does + // not inject this workload-specific setting. + RestoreStandbyModeEnv = "SNAPSHOT_RESTORE_STANDBY" + + // LegacyRestoreStandbyModeEnv is the deprecated Dynamo restore standby + // environment variable. Snapshot publishes the name but does not inject it. + // + // Deprecated: use RestoreStandbyModeEnv for new workload integrations. + LegacyRestoreStandbyModeEnv = "DYN_SNAPSHOT_RESTORE_STANDBY" + + // RestoreCompleteFile is written by the Snapshot agent when restore has + // completed and the workload may resume. + RestoreCompleteFile = "restore-complete" +) + +// ContainerMapping maps the one captured source container to a restore +// destination in the target Pod. +type ContainerMapping struct { + Source string + Destination string +} + +// GetRestoreFromSnapshotName returns the same-namespace PodSnapshot named by +// the restore-from annotation. +func GetRestoreFromSnapshotName(annotations map[string]string) (string, error) { + 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) + } + if errs := validation.IsDNS1123Subdomain(snapshotName); len(errs) != 0 { + return "", fmt.Errorf( + "%s value %q is not a valid PodSnapshot name: %s", + RestoreFromAnnotation, + snapshotName, + strings.Join(errs, "; "), + ) + } + return snapshotName, nil +} + +// ContainerMappingsFromAnnotations parses the optional flat restore mapping. +// Absence keeps the existing same-name restore behavior. +func ContainerMappingsFromAnnotations( + annotations map[string]string, + capturedSource string, +) ([]ContainerMapping, error) { + raw, ok := annotations[RestoreContainerMapAnnotation] + if !ok { + capturedSource = strings.TrimSpace(capturedSource) + return []ContainerMapping{{Source: capturedSource, Destination: capturedSource}}, nil + } + parts := strings.Split(strings.TrimSpace(raw), ",") + mappings := make([]ContainerMapping, 0, len(parts)) + for _, part := range parts { + pair := strings.Split(part, "=") + if len(pair) != 2 { + return nil, fmt.Errorf( + "invalid %s entry %q: expected source=destination", + RestoreContainerMapAnnotation, + strings.TrimSpace(part), + ) + } + mappings = append(mappings, ContainerMapping{ + Source: strings.TrimSpace(pair[0]), + Destination: strings.TrimSpace(pair[1]), + }) + } + return mappings, nil +} + +// ValidateContainerMappings enforces the one-source-to-many-destinations +// contract after parsing. +func ValidateContainerMappings(mappings []ContainerMapping, capturedSource string) error { + capturedSource = strings.TrimSpace(capturedSource) + if errs := validation.IsDNS1123Label(capturedSource); len(errs) != 0 { + return fmt.Errorf("captured source container %q is invalid: %s", capturedSource, strings.Join(errs, "; ")) + } + if len(mappings) == 0 { + return fmt.Errorf("restore container mapping must contain at least one destination") + } + destinations := make(map[string]struct{}, len(mappings)) + for _, mapping := range mappings { + source := mapping.Source + destination := mapping.Destination + if errs := validation.IsDNS1123Label(source); len(errs) != 0 { + return fmt.Errorf("invalid restore source container %q: %s", source, strings.Join(errs, "; ")) + } + if errs := validation.IsDNS1123Label(destination); len(errs) != 0 { + return fmt.Errorf("invalid restore destination container %q: %s", destination, strings.Join(errs, "; ")) + } + if source != capturedSource { + return fmt.Errorf("restore source container %q does not match captured container %q", source, capturedSource) + } + if _, duplicate := destinations[destination]; duplicate { + return fmt.Errorf("duplicate restore destination container %q", destination) + } + destinations[destination] = struct{}{} + } + return nil +} diff --git a/api/v1alpha1/protocol_test.go b/api/podcontract/protocol_test.go similarity index 52% rename from api/v1alpha1/protocol_test.go rename to api/podcontract/protocol_test.go index 81ac0161..a5c4c25b 100644 --- a/api/v1alpha1/protocol_test.go +++ b/api/podcontract/protocol_test.go @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -package v1alpha1 +package podcontract import ( "reflect" @@ -33,26 +33,26 @@ func TestGetRestoreFromSnapshotName(t *testing.T) { } } -func TestRestoreContainerMappingsFromAnnotations(t *testing.T) { +func TestContainerMappingsFromAnnotations(t *testing.T) { t.Run("defaults to captured container", func(t *testing.T) { - got, err := RestoreContainerMappingsFromAnnotations(nil, "main") + got, err := ContainerMappingsFromAnnotations(nil, "main") if err != nil { t.Fatal(err) } - want := []RestoreContainerMapping{{Source: "main", Destination: "main"}} + want := []ContainerMapping{{Source: "main", Destination: "main"}} if !reflect.DeepEqual(got, want) { t.Fatalf("mappings = %#v, want %#v", got, want) } }) t.Run("parses one source to many destinations", func(t *testing.T) { - got, err := RestoreContainerMappingsFromAnnotations(map[string]string{ + got, err := ContainerMappingsFromAnnotations(map[string]string{ RestoreContainerMapAnnotation: " main = engine-0,main=engine-1 ", }, "main") if err != nil { t.Fatal(err) } - want := []RestoreContainerMapping{ + want := []ContainerMapping{ {Source: "main", Destination: "engine-0"}, {Source: "main", Destination: "engine-1"}, } @@ -67,42 +67,54 @@ func TestRestoreContainerMappingsFromAnnotations(t *testing.T) { "too many equals": "main=engine=0", } { t.Run(name, func(t *testing.T) { - _, err := RestoreContainerMappingsFromAnnotations(map[string]string{ + _, err := ContainerMappingsFromAnnotations(map[string]string{ RestoreContainerMapAnnotation: value, }, "main") if err == nil { - t.Fatal("RestoreContainerMappingsFromAnnotations() unexpectedly succeeded") + t.Fatal("ContainerMappingsFromAnnotations() unexpectedly succeeded") } }) } } -func TestValidateRestoreContainerMappings(t *testing.T) { - valid := []RestoreContainerMapping{ +func TestValidateContainerMappings(t *testing.T) { + valid := []ContainerMapping{ {Source: "main", Destination: "engine-0"}, {Source: "main", Destination: "engine-1"}, } - if err := ValidateRestoreContainerMappings(valid, "main"); err != nil { + if err := ValidateContainerMappings(valid, "main"); err != nil { t.Fatal(err) } tests := map[string]struct { - mappings []RestoreContainerMapping + mappings []ContainerMapping source string }{ - "empty": {source: "main"}, - "empty source": {mappings: []RestoreContainerMapping{{Source: "", Destination: "engine-0"}}, source: "main"}, - "empty destination": {mappings: []RestoreContainerMapping{{Source: "main", Destination: ""}}, source: "main"}, - "invalid source": {mappings: []RestoreContainerMapping{{Source: "Main", Destination: "engine-0"}}, source: "main"}, - "invalid destination": {mappings: []RestoreContainerMapping{{Source: "main", Destination: "Engine_0"}}, source: "main"}, - "source mismatch": {mappings: []RestoreContainerMapping{{Source: "worker", Destination: "engine-0"}}, source: "main"}, - "multiple sources": {mappings: []RestoreContainerMapping{{Source: "main", Destination: "engine-0"}, {Source: "worker", Destination: "engine-1"}}, source: "main"}, - "duplicate destination": {mappings: []RestoreContainerMapping{{Source: "main", Destination: "engine-0"}, {Source: "main", Destination: "engine-0"}}, source: "main"}, + "empty": {source: "main"}, + "empty source": {mappings: []ContainerMapping{{Source: "", Destination: "engine-0"}}, source: "main"}, + "empty destination": {mappings: []ContainerMapping{{Source: "main", Destination: ""}}, source: "main"}, + "invalid source": {mappings: []ContainerMapping{{Source: "Main", Destination: "engine-0"}}, source: "main"}, + "invalid destination": {mappings: []ContainerMapping{{Source: "main", Destination: "Engine_0"}}, source: "main"}, + "source mismatch": {mappings: []ContainerMapping{{Source: "worker", Destination: "engine-0"}}, source: "main"}, + "multiple sources": { + mappings: []ContainerMapping{ + {Source: "main", Destination: "engine-0"}, + {Source: "worker", Destination: "engine-1"}, + }, + source: "main", + }, + "duplicate destination": { + mappings: []ContainerMapping{ + {Source: "main", Destination: "engine-0"}, + {Source: "main", Destination: "engine-0"}, + }, + source: "main", + }, } for name, test := range tests { t.Run(name, func(t *testing.T) { - if err := ValidateRestoreContainerMappings(test.mappings, test.source); err == nil { - t.Fatal("ValidateRestoreContainerMappings() unexpectedly succeeded") + if err := ValidateContainerMappings(test.mappings, test.source); err == nil { + t.Fatal("ValidateContainerMappings() unexpectedly succeeded") } }) } diff --git a/api/podcontract/restore_pod.go b/api/podcontract/restore_pod.go new file mode 100644 index 00000000..bf38b6a2 --- /dev/null +++ b/api/podcontract/restore_pod.go @@ -0,0 +1,494 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package podcontract + +import ( + "fmt" + "path" + "reflect" + "sort" + "strings" + + corev1 "k8s.io/api/core/v1" +) + +const restoreStartupFailureThreshold int32 = 1800 // 30 minutes at 1s cadence. + +// Request identifies the PodSnapshot source and its restore destinations. +// Empty Mappings default to restoring SourceContainer into a same-named +// destination container. +type Request struct { + SnapshotName string + SourceContainer string + Mappings []ContainerMapping +} + +// Options controls optional restore Pod shaping. Its zero value applies only +// the mechanics required by Snapshot's minimum restore protocol. +type Options struct { + // SeccompProfile is the kubelet-local profile applied to each restore + // destination container. Empty leaves seccomp configuration entirely + // caller-owned. + SeccompProfile string + + // EnableStartupGate replaces each destination container's startup probe + // with a restore-completion gate. Workload owners opt in when Kubernetes + // must withhold readiness and liveness until restore completes. + EnableStartupGate bool +} + +// Build 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 Build(pod *corev1.Pod, request Request, options Options) (*corev1.Pod, error) { + if pod == nil { + return nil, fmt.Errorf("restore pod is nil") + } + request, err := normalizeRequest(request) + if err != nil { + return nil, err + } + result := pod.DeepCopy() + if err := ensureRestoreAnnotations(result, request); err != nil { + return nil, err + } + if err := ensureRestorePodSpec(&result.Spec, request.Mappings, options); err != nil { + return nil, err + } + if err := validateCanonicalRestorePod(result, request, options); err != nil { + return nil, fmt.Errorf("validate shaped restore pod: %w", err) + } + return result, nil +} + +// Validate verifies that pod implements Snapshot's minimum restore protocol. +// Optional workload policy such as startup gating and seccomp selection is not +// part of this agent-facing validation. Validate performs no API reads and +// never mutates the Pod. +func Validate(pod *corev1.Pod, request Request) error { + request, err := normalizeRequest(request) + if err != nil { + return err + } + return validateMinimumRestorePod(pod, request) +} + +func validateMinimumRestorePod(pod *corev1.Pod, request Request) error { + if pod == nil { + return fmt.Errorf("restore pod is nil") + } + if err := validateRestoreAnnotations(pod.Annotations, request); err != nil { + return err + } + if err := validateControlVolume(&pod.Spec); err != nil { + return err + } + for _, mapping := range request.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 + } + } + return nil +} + +func validateCanonicalRestorePod(pod *corev1.Pod, request Request, options Options) error { + if err := validateMinimumRestorePod(pod, request); err != nil { + return err + } + for _, mapping := range request.Mappings { + container := findContainer(&pod.Spec, mapping.Destination) + if options.EnableStartupGate { + if err := validateCanonicalRestoreStartupProbe(container); err != nil { + return err + } + } + if err := validateContainerSeccompProfile(container, options.SeccompProfile); err != nil { + return err + } + } + return nil +} + +func normalizeRequest(request Request) (Request, error) { + snapshotName, err := validateRestoreFromSnapshotName(request.SnapshotName) + if err != nil { + return Request{}, err + } + source := strings.TrimSpace(request.SourceContainer) + mappings := request.Mappings + if len(mappings) == 0 { + mappings = []ContainerMapping{{Source: source, Destination: source}} + } + if err := ValidateContainerMappings(mappings, source); err != nil { + return Request{}, err + } + request.SnapshotName = snapshotName + request.SourceContainer = source + request.Mappings = append([]ContainerMapping(nil), mappings...) + return request, nil +} + +func ensureRestoreAnnotations(pod *corev1.Pod, request Request) 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 != request.SnapshotName { + return fmt.Errorf( + "%s names %q, conflicting with requested PodSnapshot %q", + RestoreFromAnnotation, + resolved, + request.SnapshotName, + ) + } + } + pod.Annotations[RestoreFromAnnotation] = request.SnapshotName + + formatted, needsMapping := formatContainerMappings(request.Mappings) + if _, found := pod.Annotations[RestoreContainerMapAnnotation]; found { + existingMappings, err := ContainerMappingsFromAnnotations(pod.Annotations, request.SourceContainer) + if err != nil { + return err + } + if err := ValidateContainerMappings(existingMappings, request.SourceContainer); err != nil { + return err + } + if !sameContainerMappings(existingMappings, request.Mappings) { + return fmt.Errorf("%s conflicts with requested restore mappings", RestoreContainerMapAnnotation) + } + if needsMapping { + pod.Annotations[RestoreContainerMapAnnotation] = formatted + } else { + delete(pod.Annotations, RestoreContainerMapAnnotation) + } + return nil + } + if needsMapping { + pod.Annotations[RestoreContainerMapAnnotation] = formatted + } + return nil +} + +func validateRestoreAnnotations(annotations map[string]string, request Request) error { + resolved, err := GetRestoreFromSnapshotName(annotations) + if err != nil { + return err + } + if resolved != request.SnapshotName { + return fmt.Errorf("%s names %q, expected %q", RestoreFromAnnotation, resolved, request.SnapshotName) + } + annotatedMappings, err := ContainerMappingsFromAnnotations(annotations, request.SourceContainer) + if err != nil { + return err + } + if err := ValidateContainerMappings(annotatedMappings, request.SourceContainer); err != nil { + return err + } + if !sameContainerMappings(annotatedMappings, request.Mappings) { + return fmt.Errorf("%s does not match the requested restore mappings", RestoreContainerMapAnnotation) + } + return nil +} + +func formatContainerMappings(mappings []ContainerMapping) (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) + } + sort.Strings(formatted) + return strings.Join(formatted, ","), true +} + +func sameContainerMappings(left, right []ContainerMapping) bool { + if len(left) != len(right) { + return false + } + want := make(map[ContainerMapping]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 []ContainerMapping, options Options) error { + if err := ensureControlVolume(spec); 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 := ensureContainerSeccompProfile(container, options.SeccompProfile); err != nil { + return err + } + if err := ensureControlMount(container); err != nil { + return err + } + if err := ensureControlEnvironment(container); err != nil { + return err + } + if options.EnableStartupGate { + 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 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 ensureContainerSeccompProfile(container *corev1.Container, expected string) error { + if expected == "" { + return nil + } + if container.SecurityContext == nil { + container.SecurityContext = &corev1.SecurityContext{} + } + if container.SecurityContext.SeccompProfile == nil { + container.SecurityContext.SeccompProfile = localhostSeccompProfile(expected) + return nil + } + return validateContainerSeccompProfile(container, expected) +} + +func validateContainerSeccompProfile(container *corev1.Container, expected string) error { + if expected == "" { + return nil + } + if container.SecurityContext == nil || + !matchesLocalhostSeccompProfile(container.SecurityContext.SeccompProfile, expected) { + return fmt.Errorf("container %q must use 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/podcontract/restore_pod_test.go b/api/podcontract/restore_pod_test.go new file mode 100644 index 00000000..c17488a6 --- /dev/null +++ b/api/podcontract/restore_pod_test.go @@ -0,0 +1,436 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package podcontract + +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 restoreRequest(mappings ...ContainerMapping) Request { + return Request{ + SnapshotName: "snapshot-a", + SourceContainer: "main", + Mappings: mappings, + } +} + +func TestBuildShapesSingleDestination(t *testing.T) { + original := restorePodFixture() + before := original.DeepCopy() + options := Options{ + SeccompProfile: DefaultSeccompLocalhostProfile, + EnableStartupGate: true, + } + + shaped, err := Build(original, restoreRequest(), options) + if err != nil { + t.Fatalf("Build() failed: %v", err) + } + if !reflect.DeepEqual(original, before) { + t.Fatal("Build() 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 { + t.Fatalf("restore profile was applied at Pod scope: %#v", shaped.Spec.SecurityContext) + } + if main.SecurityContext == nil || + !matchesLocalhostSeccompProfile(main.SecurityContext.SeccompProfile, DefaultSeccompLocalhostProfile) { + t.Fatalf("destination is missing expected seccomp profile: %#v", main.SecurityContext) + } + if !reflect.DeepEqual(shaped.Spec.Containers[1], before.Spec.Containers[1]) { + t.Fatal("non-destination sidecar was modified") + } + if err := Validate(shaped, restoreRequest()); err != nil { + t.Fatalf("Validate() rejected builder output: %v", err) + } +} + +func TestBuildShapesFanoutIdempotently(t *testing.T) { + pod := restorePodFixture() + pod.Spec.Containers = []corev1.Container{{Name: "engine-0"}, {Name: "engine-1"}} + request := restoreRequest( + ContainerMapping{Source: "main", Destination: "engine-1"}, + ContainerMapping{Source: "main", Destination: "engine-0"}, + ) + options := Options{SeccompProfile: DefaultSeccompLocalhostProfile} + + first, err := Build(pod, request, options) + if err != nil { + t.Fatalf("first Build() failed: %v", err) + } + second, err := Build(first, request, options) + if err != nil { + t.Fatalf("second Build() failed: %v", err) + } + if !reflect.DeepEqual(first, second) { + t.Fatal("Build() 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) + } + if container.SecurityContext == nil || + !matchesLocalhostSeccompProfile(container.SecurityContext.SeccompProfile, DefaultSeccompLocalhostProfile) { + t.Fatalf("container %q is missing expected seccomp profile: %#v", name, container.SecurityContext) + } + } +} + +func TestBuildStartupGateIsOptIn(t *testing.T) { + pod := restorePodFixture() + pod.Spec.Containers[0].LivenessProbe = &corev1.Probe{ + ProbeHandler: corev1.ProbeHandler{HTTPGet: &corev1.HTTPGetAction{Path: "/live"}}, + } + pod.Spec.Containers[0].ReadinessProbe = &corev1.Probe{ + ProbeHandler: corev1.ProbeHandler{HTTPGet: &corev1.HTTPGetAction{Path: "/ready"}}, + } + pod.Spec.Containers[0].StartupProbe = &corev1.Probe{ + ProbeHandler: corev1.ProbeHandler{HTTPGet: &corev1.HTTPGetAction{Path: "/startup"}}, + PeriodSeconds: 3, + } + beforeLiveness := pod.Spec.Containers[0].LivenessProbe.DeepCopy() + beforeReadiness := pod.Spec.Containers[0].ReadinessProbe.DeepCopy() + beforeStartup := pod.Spec.Containers[0].StartupProbe.DeepCopy() + + ungated, err := Build(pod, restoreRequest(), Options{}) + if err != nil { + t.Fatalf("ungated Build() failed: %v", err) + } + if !reflect.DeepEqual(ungated.Spec.Containers[0].StartupProbe, beforeStartup) { + t.Fatal("zero-value options changed the workload startup probe") + } + + gated, err := Build(pod, restoreRequest(), Options{EnableStartupGate: true}) + if err != nil { + t.Fatalf("gated Build() failed: %v", err) + } + main := &gated.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("Build() mutated the input startup probe") + } + if !reflect.DeepEqual(main.StartupProbe, canonicalRestoreStartupProbe()) { + t.Fatalf("unexpected restore startup gate: %#v", main.StartupProbe) + } +} + +func TestBuildCanonicalizesEquivalentMappings(t *testing.T) { + t.Run("fanout", func(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" + request := restoreRequest( + ContainerMapping{Source: "main", Destination: "engine-0"}, + ContainerMapping{Source: "main", Destination: "engine-1"}, + ) + + shaped, err := Build(pod, request, Options{}) + if err != nil { + t.Fatalf("Build() failed: %v", err) + } + if got := shaped.Annotations[RestoreContainerMapAnnotation]; got != "main=engine-0,main=engine-1" { + t.Fatalf("canonical mapping = %q", got) + } + }) + + t.Run("same name", func(t *testing.T) { + pod := restorePodFixture() + pod.Annotations[RestoreContainerMapAnnotation] = "main=main" + shaped, err := Build(pod, restoreRequest(), Options{}) + if err != nil { + t.Fatalf("Build() failed: %v", err) + } + if _, found := shaped.Annotations[RestoreContainerMapAnnotation]; found { + t.Fatalf("canonical same-name restore retained %s", RestoreContainerMapAnnotation) + } + }) +} + +func TestBuildRejectsConflictsAtomically(t *testing.T) { + runtimeDefault := corev1.SeccompProfile{Type: corev1.SeccompProfileTypeRuntimeDefault} + tests := []struct { + name string + mutate func(*corev1.Pod) + mappings []ContainerMapping + }{ + { + 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: "container seccomp", + mutate: func(pod *corev1.Pod) { + pod.Spec.Containers[0].SecurityContext = &corev1.SecurityContext{SeccompProfile: runtimeDefault.DeepCopy()} + }, + }, + { + name: "missing destination", + mappings: []ContainerMapping{{Source: "main", Destination: "missing"}}, + }, + { + name: "duplicate destination", + mappings: []ContainerMapping{ + {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() + + if _, err := Build( + pod, + restoreRequest(test.mappings...), + Options{SeccompProfile: DefaultSeccompLocalhostProfile}, + ); err == nil { + t.Fatal("Build() unexpectedly succeeded") + } + if !reflect.DeepEqual(pod, before) { + t.Fatal("Build() mutated input after failure") + } + }) + } +} + +func TestBuildScopesSeccompToRestoreDestinations(t *testing.T) { + runtimeDefault := &corev1.SeccompProfile{Type: corev1.SeccompProfileTypeRuntimeDefault} + pod := restorePodFixture() + pod.Spec.SecurityContext = &corev1.PodSecurityContext{SeccompProfile: runtimeDefault.DeepCopy()} + beforeSidecar := pod.Spec.Containers[1].DeepCopy() + + shaped, err := Build( + pod, + restoreRequest(), + Options{SeccompProfile: DefaultSeccompLocalhostProfile}, + ) + if err != nil { + t.Fatalf("Build() failed: %v", err) + } + if !reflect.DeepEqual(shaped.Spec.SecurityContext.SeccompProfile, runtimeDefault) { + t.Fatalf("Pod seccomp profile changed: %#v", shaped.Spec.SecurityContext.SeccompProfile) + } + main := &shaped.Spec.Containers[0] + if main.SecurityContext == nil || + !matchesLocalhostSeccompProfile(main.SecurityContext.SeccompProfile, DefaultSeccompLocalhostProfile) { + t.Fatalf("destination is missing expected seccomp profile: %#v", main.SecurityContext) + } + if !reflect.DeepEqual(&shaped.Spec.Containers[1], beforeSidecar) { + t.Fatalf("non-destination sidecar changed: %#v", shaped.Spec.Containers[1]) + } +} + +func TestValidateMinimumContract(t *testing.T) { + pod, err := Build(restorePodFixture(), restoreRequest(), Options{ + SeccompProfile: DefaultSeccompLocalhostProfile, + EnableStartupGate: true, + }) + if err != nil { + t.Fatalf("Build() failed: %v", err) + } + + // Legacy environment compatibility, startup gating, and seccomp selection + // are producer policy rather than agent-facing minimum protocol. + pod.Spec.Containers[0].Env = removeRestoreEnv(pod.Spec.Containers[0].Env, LegacySnapshotControlDirEnv) + pod.Spec.Containers[0].StartupProbe = &corev1.Probe{ + ProbeHandler: corev1.ProbeHandler{Exec: &corev1.ExecAction{Command: []string{"true"}}}, + } + pod.Spec.Containers[0].SecurityContext.SeccompProfile = nil + if err := Validate(pod, restoreRequest()); err != nil { + t.Fatalf("Validate() rejected minimum restore protocol: %v", err) + } + + pod.Spec.Containers[0].Env = append(pod.Spec.Containers[0].Env, corev1.EnvVar{ + Name: LegacySnapshotControlDirEnv, + Value: "/wrong", + }) + if err := Validate(pod, restoreRequest()); err == nil { + t.Fatal("Validate() accepted conflicting deprecated environment alias") + } +} + +func TestValidateRejectsMinimumContractDrift(t *testing.T) { + valid, err := Build(restorePodFixture(), restoreRequest(), Options{}) + if err != nil { + t.Fatalf("Build() 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) + }, + } + for name, mutate := range tests { + t.Run(name, func(t *testing.T) { + pod := valid.DeepCopy() + mutate(pod) + if err := Validate(pod, restoreRequest()); err == nil { + t.Fatal("Validate() unexpectedly succeeded") + } + }) + } +} + +func TestBuildAllowsUnmanagedSeccomp(t *testing.T) { + pod := restorePodFixture() + shaped, err := Build(pod, restoreRequest(), Options{}) + if err != nil { + t.Fatalf("Build() failed: %v", err) + } + if shaped.Spec.SecurityContext != nil { + t.Fatalf("empty option unexpectedly changed security context: %#v", shaped.Spec.SecurityContext) + } + if shaped.Spec.Containers[0].SecurityContext != nil { + t.Fatalf( + "empty option unexpectedly changed destination security context: %#v", + shaped.Spec.Containers[0].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/api/v1alpha1/restore_status.go b/api/podcontract/restore_status.go similarity index 86% rename from api/v1alpha1/restore_status.go rename to api/podcontract/restore_status.go index 2bf114da..598dd1ad 100644 --- a/api/v1alpha1/restore_status.go +++ b/api/podcontract/restore_status.go @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -package v1alpha1 +package podcontract import corev1 "k8s.io/api/core/v1" @@ -24,9 +24,23 @@ const ( RestoreOutcomePartiallySucceeded RestoreOutcome = "PartiallySucceeded" ) +// Terminal reports whether the restore outcome is final. +func (o RestoreOutcome) Terminal() bool { + switch o { + case RestoreOutcomeSucceeded, RestoreOutcomeFailed, RestoreOutcomePartiallySucceeded: + return true + default: + return false + } +} + // Stable reasons used on the nvidia.com/Restored Pod condition. Dependency-wait // reasons remain agent-internal; consumers should use ClassifyRestoreOutcome. const ( + // RestoredCondition is the Pod status condition owned by the Snapshot node + // agent. + RestoredCondition = "nvidia.com/Restored" + // RestoreReasonInProgress marks active restore execution. RestoreReasonInProgress = "RestoreInProgress" // RestoreReasonSucceeded marks a terminal all-destinations success. diff --git a/api/v1alpha1/restore_status_test.go b/api/podcontract/restore_status_test.go similarity index 82% rename from api/v1alpha1/restore_status_test.go rename to api/podcontract/restore_status_test.go index 658e9caf..871bce71 100644 --- a/api/v1alpha1/restore_status_test.go +++ b/api/podcontract/restore_status_test.go @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -package v1alpha1 +package podcontract import ( "testing" @@ -79,3 +79,19 @@ func TestClassifyRestoreOutcome(t *testing.T) { }) } } + +func TestRestoreOutcomeTerminal(t *testing.T) { + tests := map[RestoreOutcome]bool{ + RestoreOutcomeUnknown: false, + RestoreOutcomePending: false, + RestoreOutcomeSucceeded: true, + RestoreOutcomeFailed: true, + RestoreOutcomePartiallySucceeded: true, + RestoreOutcome("future-outcome"): false, + } + for outcome, want := range tests { + if got := outcome.Terminal(); got != want { + t.Errorf("RestoreOutcome(%q).Terminal() = %t, want %t", outcome, got, want) + } + } +} diff --git a/api/v1alpha1/constants.go b/api/v1alpha1/constants.go index 1c55ecb0..297d674f 100644 --- a/api/v1alpha1/constants.go +++ b/api/v1alpha1/constants.go @@ -3,11 +3,8 @@ package v1alpha1 -// Snapshot control-plane contract: the labels, annotations, and control-volume -// vocabulary the operator stamps and the node agent reads (and that external -// consumers such as the restore webhook and the workload depend on). This is the -// single versioned home for these constants so both sides pin them from the api -// module. +// Snapshot control-plane labels and storage vocabulary shared by the operator +// and node agent. Workload Pod contract constants live in api/podcontract. const ( // CaptureEligibleLabel is the gate-applied promotion label: the node agent's pre-bind gate // adds it only after the source pod passes validation. The agent's source-pod capture @@ -18,81 +15,5 @@ const ( // object so the per-node agent's cache can label-select work for its node. SnapshotNodeLabel = "nvidia.com/snapshot-node" - // RestoreFromAnnotation names the PodSnapshot to restore in the pod's - // namespace. - RestoreFromAnnotation = "nvidia.com/restore-from" - - // RestoreContainerMapAnnotation optionally maps the single captured source - // container to one or more restore destinations. Its value is a comma-separated - // list of source=destination pairs. When absent, restore uses the captured - // container name as the destination. - RestoreContainerMapAnnotation = "nvidia.com/restore-container-map" - - // RestoredCondition is the Pod status condition owned by the node agent. - RestoredCondition = "nvidia.com/Restored" - - CheckpointVolumeName = "checkpoint-storage" - DefaultSeccompLocalhostProfile = "profiles/block-iouring.json" -) - -// Control-volume contract: the per-pod emptyDir carrying checkpoint/restore -// lifecycle sentinels written by the snapshot agent and observed by the workload. -const ( - // SnapshotControlVolumeName is the per-pod emptyDir used to carry - // checkpoint/restore lifecycle sentinels written by the snapshot agent - // and observed by the workload. It replaces the SIGUSR1/SIGCONT signals - // that previously required the workload to run as PID 1. - // - // When a pod targets multiple containers (e.g. failover engine-0 + - // engine-1), each container mounts the emptyDir with - // subPath=, so sentinels are isolated per-container on - // disk while each container still sees them at SnapshotControlMountPath. - SnapshotControlVolumeName = "snapshot-control" - - // SnapshotControlMountPath is where the control volume is mounted inside - // the workload container. - SnapshotControlMountPath = "/snapshot-control" - - // SnapshotControlDirEnv is the canonical environment variable exposing the - // 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 - // SnapshotControlDirEnv. - // - // Deprecated: use SnapshotControlDirEnv instead. Remove once no workload - // image depends on this name. - LegacySnapshotControlDirEnv = "DYN_SNAPSHOT_CONTROL_DIR" - - // 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; - // a workload polling for this file is killed instead of released. - // - // Deprecated: no longer written. Remove once no workload image waits on it. - SnapshotCompleteFile = "snapshot-complete" - - // RestoreCompleteFile is written by the snapshot agent inside the - // control volume when a restore has completed and the workload may - // resume. - RestoreCompleteFile = "restore-complete" - - // ReadyForSnapshotFile is written by the workload inside the control - // volume when the model is loaded and the workload is ready for a - // checkpoint. Observed by the source job's kubelet readiness probe - // on the worker container. - ReadyForSnapshotFile = "ready-for-snapshot" - - // CUDAJobFileName is the stable name the cuda-checkpoint-helper launch-job - // wrapper persists the CUDA checkpoint job file under, inside the control - // volume. The stable location survives past the transient path the CUDA - // driver reports via CUDA_CHECKPOINT_JOB_FILE. - CUDAJobFileName = "cuda-checkpoint-job" - - // CUDAJobFilePath is the full stable path to the persisted CUDA checkpoint - // job file. - CUDAJobFilePath = SnapshotControlMountPath + "/" + CUDAJobFileName + CheckpointVolumeName = "checkpoint-storage" ) diff --git a/api/v1alpha1/protocol.go b/api/v1alpha1/protocol.go deleted file mode 100644 index 26c414bb..00000000 --- a/api/v1alpha1/protocol.go +++ /dev/null @@ -1,86 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -package v1alpha1 - -import ( - "fmt" - "strings" - - "k8s.io/apimachinery/pkg/util/validation" -) - -// RestoreContainerMapping maps the one captured source container to a restore -// destination in the target Pod. -// +kubebuilder:object:generate=false -type RestoreContainerMapping struct { - Source string - Destination string -} - -// 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]) - if snapshotName == "" { - return "", fmt.Errorf("%s must name a PodSnapshot", RestoreFromAnnotation) - } - if errs := validation.IsDNS1123Subdomain(snapshotName); len(errs) != 0 { - return "", fmt.Errorf("%s value %q is not a valid PodSnapshot name: %s", RestoreFromAnnotation, snapshotName, strings.Join(errs, "; ")) - } - return snapshotName, nil -} - -// RestoreContainerMappingsFromAnnotations parses the optional flat restore -// mapping. Absence keeps the existing same-name restore behavior. -func RestoreContainerMappingsFromAnnotations(annotations map[string]string, capturedSource string) ([]RestoreContainerMapping, error) { - raw, ok := annotations[RestoreContainerMapAnnotation] - if !ok { - capturedSource = strings.TrimSpace(capturedSource) - return []RestoreContainerMapping{{Source: capturedSource, Destination: capturedSource}}, nil - } - parts := strings.Split(strings.TrimSpace(raw), ",") - mappings := make([]RestoreContainerMapping, 0, len(parts)) - for _, part := range parts { - pair := strings.Split(part, "=") - if len(pair) != 2 { - return nil, fmt.Errorf("invalid %s entry %q: expected source=destination", RestoreContainerMapAnnotation, strings.TrimSpace(part)) - } - mappings = append(mappings, RestoreContainerMapping{ - Source: strings.TrimSpace(pair[0]), - Destination: strings.TrimSpace(pair[1]), - }) - } - return mappings, nil -} - -// ValidateRestoreContainerMappings enforces the one-source-to-many-destinations -// contract after parsing. -func ValidateRestoreContainerMappings(mappings []RestoreContainerMapping, capturedSource string) error { - capturedSource = strings.TrimSpace(capturedSource) - if errs := validation.IsDNS1123Label(capturedSource); len(errs) != 0 { - return fmt.Errorf("captured source container %q is invalid: %s", capturedSource, strings.Join(errs, "; ")) - } - if len(mappings) == 0 { - return fmt.Errorf("restore container mapping must contain at least one destination") - } - destinations := make(map[string]struct{}, len(mappings)) - for _, mapping := range mappings { - source := mapping.Source - destination := mapping.Destination - if errs := validation.IsDNS1123Label(source); len(errs) != 0 { - return fmt.Errorf("invalid restore source container %q: %s", source, strings.Join(errs, "; ")) - } - if errs := validation.IsDNS1123Label(destination); len(errs) != 0 { - return fmt.Errorf("invalid restore destination container %q: %s", destination, strings.Join(errs, "; ")) - } - if source != capturedSource { - return fmt.Errorf("restore source container %q does not match captured container %q", source, capturedSource) - } - if _, duplicate := destinations[destination]; duplicate { - return fmt.Errorf("duplicate restore destination container %q", destination) - } - destinations[destination] = struct{}{} - } - return nil -} diff --git a/docs/restore-pod-contract.md b/docs/restore-pod-contract.md new file mode 100644 index 00000000..8ccc5341 --- /dev/null +++ b/docs/restore-pod-contract.md @@ -0,0 +1,152 @@ + + +# 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 := podcontract.Build( + pod, + podcontract.Request{ + SnapshotName: snapshotName, + SourceContainer: capturedContainer, + Mappings: mappings, + }, + podcontract.Options{ + SeccompProfile: podcontract.DefaultSeccompLocalhostProfile, + EnableStartupGate: true, + }, +) +``` + +This release intentionally moves the Go Pod-contract API from +`api/restorepod` and `api/v1alpha1` to `api/podcontract` without compatibility +aliases. Consumers must update restore annotations, mappings, status reasons, +`RestoreOutcome`, and `ClassifyRestoreOutcome` references to the new package. +Custom-resource types remain in `api/v1alpha1`. + +`podcontract.Build` 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. `podcontract.Validate` checks only Snapshot's minimum restore protocol, +without mutation or Kubernetes API reads. Optional workload policy such as +startup gating and seccomp selection remains producer-owned. Conflicting +annotations, volumes, mounts, environment, and requested security settings are +rejected instead of overwritten. + +Consumers can interpret the agent-owned `nvidia.com/Restored` Pod condition +without importing Snapshot's custom-resource API: + +```go +outcome := podcontract.ClassifyRestoreOutcome(pod.Status.Conditions) +``` + +The package also owns the shared control-volume, environment, capture-sentinel, +and CUDA job-file names used by Snapshot-managed source and restore Pods. + +The producer supplies `SourceContainer` from the referenced `PodSnapshot`. +Empty mappings restore that source into the same-named destination only. 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: + volumes: + - name: snapshot-control + emptyDir: {} + containers: + - name: engine-0 + image: worker:latest + securityContext: + seccompProfile: + type: Localhost + localhostProfile: profiles/block-iouring.json + # 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 + securityContext: + seccompProfile: + type: Localhost + localhostProfile: profiles/block-iouring.json + 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. It is a +required part of the minimum contract so every Snapshot-managed workload and +tool can discover the control directory without hard-coding its location. 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` startup probe shown above is optional workload lifecycle +policy. Set `Options.EnableStartupGate` when Kubernetes must withhold readiness +and liveness until restore completes. Because Kubernetes supports only one +startup probe, the builder then replaces any existing startup probe with the +canonical restore-completion gate while preserving workload liveness and +readiness probes. Its failure threshold allows 1,800 consecutive one-second +failures; if restoration exceeds that budget, kubelet restarts the placeholder +according to the Pod's restart policy. The node agent always writes the +sentinel but does not require or inspect the optional probe. + +`Options.SeccompProfile` controls the localhost profile on restore destination +containers only. The standard Snapshot Helm installation deploys +`DefaultSeccompLocalhostProfile` to block io_uring for CRIU. Applying it at +container scope preserves the Pod-level policy inherited by unrelated +sidecars. An empty value leaves seccomp unmanaged for environments that provide +the restriction elsewhere. A destination container must not override a +requested profile with a conflicting 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 `api/podcontract` 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/checkpoint.go b/operator/cmd/snapshotctl/checkpoint.go index f2819d3e..734bcdc0 100644 --- a/operator/cmd/snapshotctl/checkpoint.go +++ b/operator/cmd/snapshotctl/checkpoint.go @@ -15,7 +15,7 @@ import ( "k8s.io/apimachinery/pkg/util/rand" "k8s.io/apimachinery/pkg/util/validation" - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" + "github.com/ai-dynamo/snapshot/api/podcontract" snapshotprotocol "github.com/ai-dynamo/snapshot/operator/internal/protocol" ) @@ -73,7 +73,7 @@ func runCheckpointFlow(ctx context.Context, opts checkpointOptions) (_ *result, }, snapshotprotocol.SourceJobOptions{ Namespace: namespace, TargetContainer: containerName, - SeccompProfile: snapshotv1alpha1.DefaultSeccompLocalhostProfile, + SeccompProfile: podcontract.DefaultSeccompLocalhostProfile, Name: checkpointJobName, WrapLaunchJob: opts.CudaCheckpointWrap, }) diff --git a/operator/cmd/snapshotctl/restore.go b/operator/cmd/snapshotctl/restore.go index 651d91af..e463b675 100644 --- a/operator/cmd/snapshotctl/restore.go +++ b/operator/cmd/snapshotctl/restore.go @@ -13,8 +13,8 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "sigs.k8s.io/controller-runtime/pkg/client" + "github.com/ai-dynamo/snapshot/api/podcontract" snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" - snapshotprotocol "github.com/ai-dynamo/snapshot/operator/internal/protocol" ) type restoreOptions struct { @@ -41,22 +41,35 @@ 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 := podcontract.ContainerMappingsFromAnnotations(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 := podcontract.Build( + candidate, + podcontract.Request{ + SnapshotName: snapshotName, + SourceContainer: containers[0], + Mappings: mappings, + }, + podcontract.Options{ + SeccompProfile: podcontract.DefaultSeccompLocalhostProfile, + EnableStartupGate: true, + }, + ) if err != nil { return nil, err } diff --git a/operator/internal/controller/snapshotjob_job.go b/operator/internal/controller/snapshotjob_job.go index df2e58c7..ee75c68d 100644 --- a/operator/internal/controller/snapshotjob_job.go +++ b/operator/internal/controller/snapshotjob_job.go @@ -10,6 +10,7 @@ import ( batchv1 "k8s.io/api/batch/v1" contentvalidation "k8s.io/apimachinery/pkg/api/validate/content" + "github.com/ai-dynamo/snapshot/api/podcontract" snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" "github.com/ai-dynamo/snapshot/operator/internal/protocol" ) @@ -51,7 +52,7 @@ func buildSourceJob(sj *snapshotv1alpha1.SnapshotJob) (*batchv1.Job, error) { Namespace: sj.Namespace, Name: sj.Name, TargetContainer: targetContainer, - SeccompProfile: snapshotv1alpha1.DefaultSeccompLocalhostProfile, + SeccompProfile: podcontract.DefaultSeccompLocalhostProfile, ActiveDeadlineSeconds: sj.Spec.ActiveDeadlineSeconds, TTLSecondsAfterFinish: nil, WrapLaunchJob: false, diff --git a/operator/internal/controller/snapshotjob_job_test.go b/operator/internal/controller/snapshotjob_job_test.go index a06e867b..d2c09a04 100644 --- a/operator/internal/controller/snapshotjob_job_test.go +++ b/operator/internal/controller/snapshotjob_job_test.go @@ -14,6 +14,7 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/utils/ptr" + "github.com/ai-dynamo/snapshot/api/podcontract" snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" ) @@ -53,7 +54,7 @@ func TestBuildSourceJob(t *testing.T) { require.NotNil(t, job.Spec.Template.Spec.SecurityContext.SeccompProfile) profile := job.Spec.Template.Spec.SecurityContext.SeccompProfile assert.Equal(t, corev1.SeccompProfileTypeLocalhost, profile.Type) - assert.Equal(t, ptr.To(snapshotv1alpha1.DefaultSeccompLocalhostProfile), profile.LocalhostProfile) + assert.Equal(t, ptr.To(podcontract.DefaultSeccompLocalhostProfile), profile.LocalhostProfile) }) t.Run("stamps the owner label without clobbering existing pod template labels", func(t *testing.T) { @@ -111,14 +112,14 @@ func TestBuildSourceJob(t *testing.T) { require.NotNil(t, main.ReadinessProbe) require.NotNil(t, main.ReadinessProbe.Exec) assert.Equal(t, - []string{"cat", snapshotv1alpha1.SnapshotControlMountPath + "/" + snapshotv1alpha1.ReadyForSnapshotFile}, + []string{"cat", podcontract.SnapshotControlMountPath + "/" + podcontract.ReadyForSnapshotFile}, main.ReadinessProbe.Exec.Command) var mountPaths []string for _, m := range main.VolumeMounts { mountPaths = append(mountPaths, m.MountPath) } - assert.Contains(t, mountPaths, snapshotv1alpha1.SnapshotControlMountPath, + assert.Contains(t, mountPaths, podcontract.SnapshotControlMountPath, "the target container must mount the control volume the probe and sentinel live in") }) diff --git a/operator/internal/protocol/control_volume.go b/operator/internal/protocol/control_volume.go index 5578192c..83f86082 100644 --- a/operator/internal/protocol/control_volume.go +++ b/operator/internal/protocol/control_volume.go @@ -4,9 +4,8 @@ package protocol import ( + "github.com/ai-dynamo/snapshot/api/podcontract" corev1 "k8s.io/api/core/v1" - - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" ) // EnsureControlVolume adds the snapshot-control emptyDir to the pod spec, @@ -30,14 +29,14 @@ func EnsureControlVolume(podSpec *corev1.PodSpec, container *corev1.Container) { hasVolume := false for _, v := range podSpec.Volumes { - if v.Name == snapshotv1alpha1.SnapshotControlVolumeName { + if v.Name == podcontract.SnapshotControlVolumeName { hasVolume = true break } } if !hasVolume { podSpec.Volumes = append(podSpec.Volumes, corev1.Volume{ - Name: snapshotv1alpha1.SnapshotControlVolumeName, + Name: podcontract.SnapshotControlVolumeName, VolumeSource: corev1.VolumeSource{EmptyDir: &corev1.EmptyDirVolumeSource{}}, }) } @@ -50,21 +49,21 @@ func EnsureControlVolume(podSpec *corev1.PodSpec, container *corev1.Container) { hasMount := false for _, m := range container.VolumeMounts { - if m.Name == snapshotv1alpha1.SnapshotControlVolumeName { + if m.Name == podcontract.SnapshotControlVolumeName { hasMount = true break } } if !hasMount { container.VolumeMounts = append(container.VolumeMounts, corev1.VolumeMount{ - Name: snapshotv1alpha1.SnapshotControlVolumeName, - MountPath: snapshotv1alpha1.SnapshotControlMountPath, + Name: podcontract.SnapshotControlVolumeName, + MountPath: podcontract.SnapshotControlMountPath, SubPath: subPath, }) } - ensureEnv(container, snapshotv1alpha1.SnapshotControlDirEnv, snapshotv1alpha1.SnapshotControlMountPath) - ensureEnv(container, snapshotv1alpha1.LegacySnapshotControlDirEnv, snapshotv1alpha1.SnapshotControlMountPath) + ensureEnv(container, podcontract.SnapshotControlDirEnv, podcontract.SnapshotControlMountPath) + ensureEnv(container, podcontract.LegacySnapshotControlDirEnv, podcontract.SnapshotControlMountPath) } // ensureEnv sets name=value on the container if name is not already present, diff --git a/operator/internal/protocol/control_volume_test.go b/operator/internal/protocol/control_volume_test.go index 13ee8c4a..10a76bd3 100644 --- a/operator/internal/protocol/control_volume_test.go +++ b/operator/internal/protocol/control_volume_test.go @@ -6,9 +6,8 @@ package protocol import ( "testing" + "github.com/ai-dynamo/snapshot/api/podcontract" corev1 "k8s.io/api/core/v1" - - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" ) func TestEnsureControlVolume(t *testing.T) { @@ -16,12 +15,12 @@ func TestEnsureControlVolume(t *testing.T) { ps := &corev1.PodSpec{Containers: []corev1.Container{{Name: "main"}}} EnsureControlVolume(ps, &ps.Containers[0]) - if len(ps.Volumes) != 1 || ps.Volumes[0].Name != snapshotv1alpha1.SnapshotControlVolumeName || ps.Volumes[0].EmptyDir == nil { - t.Fatalf("expected one %s emptyDir volume, got %#v", snapshotv1alpha1.SnapshotControlVolumeName, ps.Volumes) + if len(ps.Volumes) != 1 || ps.Volumes[0].Name != podcontract.SnapshotControlVolumeName || ps.Volumes[0].EmptyDir == nil { + t.Fatalf("expected one %s emptyDir volume, got %#v", podcontract.SnapshotControlVolumeName, ps.Volumes) } c := ps.Containers[0] - if len(c.VolumeMounts) != 1 || c.VolumeMounts[0].Name != snapshotv1alpha1.SnapshotControlVolumeName || c.VolumeMounts[0].MountPath != snapshotv1alpha1.SnapshotControlMountPath { - t.Fatalf("expected one %s mount at %s, got %#v", snapshotv1alpha1.SnapshotControlVolumeName, snapshotv1alpha1.SnapshotControlMountPath, c.VolumeMounts) + if len(c.VolumeMounts) != 1 || c.VolumeMounts[0].Name != podcontract.SnapshotControlVolumeName || c.VolumeMounts[0].MountPath != podcontract.SnapshotControlMountPath { + t.Fatalf("expected one %s mount at %s, got %#v", podcontract.SnapshotControlVolumeName, podcontract.SnapshotControlMountPath, c.VolumeMounts) } if c.VolumeMounts[0].SubPath != "main" { t.Fatalf("expected subPath=%q, got %q", "main", c.VolumeMounts[0].SubPath) @@ -29,13 +28,13 @@ func TestEnsureControlVolume(t *testing.T) { if len(c.Env) != 2 { t.Fatalf("expected two control-dir env vars, got %#v", c.Env) } - for _, name := range []string{snapshotv1alpha1.SnapshotControlDirEnv, snapshotv1alpha1.LegacySnapshotControlDirEnv} { + for _, name := range []string{podcontract.SnapshotControlDirEnv, podcontract.LegacySnapshotControlDirEnv} { found := false for _, e := range c.Env { if e.Name == name { found = true - if e.Value != snapshotv1alpha1.SnapshotControlMountPath { - t.Fatalf("expected env %s=%s, got %#v", name, snapshotv1alpha1.SnapshotControlMountPath, e) + if e.Value != podcontract.SnapshotControlMountPath { + t.Fatalf("expected env %s=%s, got %#v", name, podcontract.SnapshotControlMountPath, e) } } } @@ -77,7 +76,7 @@ func TestEnsureControlVolume(t *testing.T) { t.Run("legacy env pre-set backfills canonical", func(t *testing.T) { ps := &corev1.PodSpec{Containers: []corev1.Container{{ Name: "main", - Env: []corev1.EnvVar{{Name: snapshotv1alpha1.LegacySnapshotControlDirEnv, Value: snapshotv1alpha1.SnapshotControlMountPath}}, + Env: []corev1.EnvVar{{Name: podcontract.LegacySnapshotControlDirEnv, Value: podcontract.SnapshotControlMountPath}}, }}} EnsureControlVolume(ps, &ps.Containers[0]) c := ps.Containers[0] @@ -86,7 +85,7 @@ func TestEnsureControlVolume(t *testing.T) { } legacyCount := 0 for _, e := range c.Env { - if e.Name == snapshotv1alpha1.LegacySnapshotControlDirEnv { + if e.Name == podcontract.LegacySnapshotControlDirEnv { legacyCount++ } } @@ -98,7 +97,7 @@ func TestEnsureControlVolume(t *testing.T) { t.Run("canonical env pre-set backfills legacy", func(t *testing.T) { ps := &corev1.PodSpec{Containers: []corev1.Container{{ Name: "main", - Env: []corev1.EnvVar{{Name: snapshotv1alpha1.SnapshotControlDirEnv, Value: snapshotv1alpha1.SnapshotControlMountPath}}, + Env: []corev1.EnvVar{{Name: podcontract.SnapshotControlDirEnv, Value: podcontract.SnapshotControlMountPath}}, }}} EnsureControlVolume(ps, &ps.Containers[0]) c := ps.Containers[0] @@ -107,7 +106,7 @@ func TestEnsureControlVolume(t *testing.T) { } canonicalCount := 0 for _, e := range c.Env { - if e.Name == snapshotv1alpha1.SnapshotControlDirEnv { + if e.Name == podcontract.SnapshotControlDirEnv { canonicalCount++ } } 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 "" -} diff --git a/operator/internal/protocol/source_job.go b/operator/internal/protocol/source_job.go index 1f68df98..7a10cb70 100644 --- a/operator/internal/protocol/source_job.go +++ b/operator/internal/protocol/source_job.go @@ -12,7 +12,7 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/utils/ptr" - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" + "github.com/ai-dynamo/snapshot/api/podcontract" ) type SourceJobOptions struct { @@ -28,8 +28,8 @@ type SourceJobOptions struct { func NewSourceJob(podTemplate *corev1.PodTemplateSpec, opts SourceJobOptions) (*batchv1.Job, error) { podTemplate = podTemplate.DeepCopy() for _, annotation := range []string{ - snapshotv1alpha1.RestoreFromAnnotation, - snapshotv1alpha1.RestoreContainerMapAnnotation, + podcontract.RestoreFromAnnotation, + podcontract.RestoreContainerMapAnnotation, } { if _, restoreRequested := podTemplate.Annotations[annotation]; restoreRequested { return nil, fmt.Errorf("source job pod template must not set %s", annotation) @@ -78,7 +78,7 @@ func NewSourceJob(podTemplate *corev1.PodTemplateSpec, opts SourceJobOptions) (* targetContainer.ReadinessProbe = &corev1.Probe{ ProbeHandler: corev1.ProbeHandler{ Exec: &corev1.ExecAction{ - Command: []string{"cat", filepath.Join(snapshotv1alpha1.SnapshotControlMountPath, snapshotv1alpha1.ReadyForSnapshotFile)}, + Command: []string{"cat", filepath.Join(podcontract.SnapshotControlMountPath, podcontract.ReadyForSnapshotFile)}, }, }, PeriodSeconds: 1, @@ -165,7 +165,7 @@ export CUDA_CHECKPOINT_JOB_FILE="$job_file" exec "$@"` wrappedArgs := make([]string, 0, len(command)+len(args)+7) - wrappedArgs = append(wrappedArgs, "--launch-job", "/bin/sh", "-c", persistJobFileScript, "dynamo-cuda-checkpoint", snapshotv1alpha1.CUDAJobFilePath) + wrappedArgs = append(wrappedArgs, "--launch-job", "/bin/sh", "-c", persistJobFileScript, "dynamo-cuda-checkpoint", podcontract.CUDAJobFilePath) wrappedArgs = append(wrappedArgs, command...) wrappedArgs = append(wrappedArgs, args...) return []string{"cuda-checkpoint"}, wrappedArgs diff --git a/operator/internal/protocol/source_job_identity_test.go b/operator/internal/protocol/source_job_identity_test.go index 2ff19d15..b5151721 100644 --- a/operator/internal/protocol/source_job_identity_test.go +++ b/operator/internal/protocol/source_job_identity_test.go @@ -10,13 +10,13 @@ import ( "strings" "testing" - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" + "github.com/ai-dynamo/snapshot/api/podcontract" ) func TestCudaCheckpointLaunchJobWrapperPersistsJobFile(t *testing.T) { tempDir := t.TempDir() transientJobFile := filepath.Join(tempDir, "transient-job") - stableJobFile := filepath.Join(tempDir, "checkpoint", snapshotv1alpha1.CUDAJobFileName) + stableJobFile := filepath.Join(tempDir, "checkpoint", podcontract.CUDAJobFileName) observedEnvironment := filepath.Join(tempDir, "observed-environment") if err := os.WriteFile(transientJobFile, []byte("job-state"), 0600); err != nil { t.Fatal(err) diff --git a/operator/internal/protocol/source_job_test.go b/operator/internal/protocol/source_job_test.go index a08ea931..49330db8 100644 --- a/operator/internal/protocol/source_job_test.go +++ b/operator/internal/protocol/source_job_test.go @@ -11,7 +11,7 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/utils/ptr" - snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1" + "github.com/ai-dynamo/snapshot/api/podcontract" ) func requireCheckpointContainer(t *testing.T, containers []corev1.Container, name string) *corev1.Container { @@ -36,8 +36,8 @@ func requireStableLaunchJobWrapper(t *testing.T, container *corev1.Container, or if container.Args[0] != "--launch-job" || container.Args[1] != "/bin/sh" || container.Args[2] != "-c" || container.Args[4] != "dynamo-cuda-checkpoint" { t.Fatalf("unexpected launch-job wrapper prefix: %#v", container.Args[:6]) } - if container.Args[5] != snapshotv1alpha1.CUDAJobFilePath { - t.Fatalf("stable job file = %q, want %q", container.Args[5], snapshotv1alpha1.CUDAJobFilePath) + if container.Args[5] != podcontract.CUDAJobFilePath { + t.Fatalf("stable job file = %q, want %q", container.Args[5], podcontract.CUDAJobFilePath) } if got := container.Args[6:]; strings.Join(got, "|") != strings.Join(original, "|") { t.Fatalf("original command = %#v, want %#v", got, original) @@ -64,7 +64,7 @@ func TestNewSourceJob(t *testing.T) { }, SourceJobOptions{ Namespace: "test-ns", TargetContainer: "main", - SeccompProfile: snapshotv1alpha1.DefaultSeccompLocalhostProfile, + SeccompProfile: podcontract.DefaultSeccompLocalhostProfile, Name: "test-job", ActiveDeadlineSeconds: ptr.To(int64(60)), TTLSecondsAfterFinish: ptr.To(int32(300)), @@ -77,12 +77,12 @@ func TestNewSourceJob(t *testing.T) { if job.Name != "test-job" || job.Namespace != "test-ns" { t.Fatalf("unexpected job identity: %#v", job.ObjectMeta) } - if len(job.Spec.Template.Spec.Volumes) != 1 || job.Spec.Template.Spec.Volumes[0].Name != snapshotv1alpha1.SnapshotControlVolumeName { - t.Fatalf("expected only %s volume, got %#v", snapshotv1alpha1.SnapshotControlVolumeName, job.Spec.Template.Spec.Volumes) + if len(job.Spec.Template.Spec.Volumes) != 1 || job.Spec.Template.Spec.Volumes[0].Name != podcontract.SnapshotControlVolumeName { + t.Fatalf("expected only %s volume, got %#v", podcontract.SnapshotControlVolumeName, job.Spec.Template.Spec.Volumes) } main := &job.Spec.Template.Spec.Containers[0] - if len(main.VolumeMounts) != 1 || main.VolumeMounts[0].MountPath != snapshotv1alpha1.SnapshotControlMountPath { - t.Fatalf("expected only %s mount at %s, got %#v", snapshotv1alpha1.SnapshotControlVolumeName, snapshotv1alpha1.SnapshotControlMountPath, main.VolumeMounts) + if len(main.VolumeMounts) != 1 || main.VolumeMounts[0].MountPath != podcontract.SnapshotControlMountPath { + t.Fatalf("expected only %s mount at %s, got %#v", podcontract.SnapshotControlVolumeName, podcontract.SnapshotControlMountPath, main.VolumeMounts) } if main.VolumeMounts[0].SubPath != "main" { t.Fatalf("expected control mount subPath=main, got %#v", main.VolumeMounts[0]) @@ -90,7 +90,7 @@ func TestNewSourceJob(t *testing.T) { if main.ReadinessProbe == nil || main.ReadinessProbe.Exec == nil { t.Fatalf("expected ready-file readiness probe, got %#v", main.ReadinessProbe) } - expectedProbe := []string{"cat", snapshotv1alpha1.SnapshotControlMountPath + "/" + snapshotv1alpha1.ReadyForSnapshotFile} + expectedProbe := []string{"cat", podcontract.SnapshotControlMountPath + "/" + podcontract.ReadyForSnapshotFile} if strings.Join(main.ReadinessProbe.Exec.Command, " ") != strings.Join(expectedProbe, " ") { t.Fatalf("expected readiness probe %#v, got %#v", expectedProbe, main.ReadinessProbe.Exec.Command) } @@ -147,12 +147,12 @@ func TestNewSourceJobWrapsTargetContainer(t *testing.T) { } // Sidecar does not get a control volume mount, snapshot env, or ready probe. for _, mount := range sidecar.VolumeMounts { - if mount.Name == snapshotv1alpha1.SnapshotControlVolumeName { + if mount.Name == podcontract.SnapshotControlVolumeName { t.Fatalf("sidecar should not have control volume mount: %#v", sidecar.VolumeMounts) } } for _, env := range sidecar.Env { - if env.Name == snapshotv1alpha1.SnapshotControlDirEnv { + if env.Name == podcontract.SnapshotControlDirEnv { t.Fatalf("sidecar should not have control env: %#v", sidecar.Env) } } @@ -262,12 +262,12 @@ func TestNewSourceJobRequiresTarget(t *testing.T) { func TestNewSourceJobRejectsRestoreAnnotations(t *testing.T) { for _, annotation := range []string{ - snapshotv1alpha1.RestoreFromAnnotation, - snapshotv1alpha1.RestoreContainerMapAnnotation, + podcontract.RestoreFromAnnotation, + podcontract.RestoreContainerMapAnnotation, } { t.Run(annotation, func(t *testing.T) { value := "" - if annotation == snapshotv1alpha1.RestoreFromAnnotation { + if annotation == podcontract.RestoreFromAnnotation { value = "snapshot-a" } _, err := NewSourceJob(&corev1.PodTemplateSpec{