Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,10 @@ Snapshots are portable across compatible machines and can be restored on any nod
| `nvidia.com/restore-from` | Namespaced | Added as a pod annotation to trigger restore from a named `PodSnapshot` in the same namespace. |
| `nvidia.com/restore-container-map` | Namespaced | Optional comma-separated `source=destination` mappings used to clone the single captured container into one or more restore containers. |

Restore producers must also implement the versioned
[restore Pod contract](docs/restore-pod-contract.md). Go integrations should
use the public builder and validator from `github.com/ai-dynamo/snapshot/api/v1alpha1`.

 


Expand Down
67 changes: 12 additions & 55 deletions agent/internal/controller/controller.go
Original file line number Diff line number Diff line change
Expand Up @@ -600,64 +600,21 @@ func validateRestoreTarget(pod *corev1.Pod, snapshot *snapshotv1alpha1.PodSnapsh
if err != nil {
return nil, nil, err
}
if err := validateRestoreMappingsForPod(&pod.Spec, mappings, containerName); err != nil {
if err := snapshotv1alpha1.ValidateRestoreContainerMappings(mappings, containerName); err != nil {
return nil, nil, err
}
return &restoreTarget{SnapshotName: snapshot.Name, ContentUID: string(content.UID), SourceContainerName: containerName}, mappings, nil
}

// validateRestoreMappingsForPod validates the mapping contract and ensures
// every destination exists in the restore Pod spec.
func validateRestoreMappingsForPod(spec *corev1.PodSpec, mappings []snapshotv1alpha1.RestoreContainerMapping, capturedSource string) error {
if err := snapshotv1alpha1.ValidateRestoreContainerMappings(mappings, capturedSource); err != nil {
return err
}
if len(mappings) > 1 {
hasControlVolume := false
for _, volume := range spec.Volumes {
if volume.Name == snapshotv1alpha1.SnapshotControlVolumeName && volume.EmptyDir != nil {
hasControlVolume = true
break
}
}
if !hasControlVolume {
return fmt.Errorf("multi-container restore requires %s emptyDir volume", snapshotv1alpha1.SnapshotControlVolumeName)
}
}
for _, mapping := range mappings {
var destination *corev1.Container
for i := range spec.Containers {
if spec.Containers[i].Name == mapping.Destination {
destination = &spec.Containers[i]
break
}
}
if destination == nil {
return fmt.Errorf("restore pod has no destination container named %q", mapping.Destination)
}
if len(mappings) == 1 {
continue
}
validControlMount := false
for _, mount := range destination.VolumeMounts {
if mount.Name == snapshotv1alpha1.SnapshotControlVolumeName &&
mount.MountPath == snapshotv1alpha1.SnapshotControlMountPath &&
mount.SubPath == mapping.Destination {
validControlMount = true
break
}
}
if !validControlMount {
return fmt.Errorf(
"multi-container restore destination %q requires %s mounted at %s with subPath %q",
mapping.Destination,
snapshotv1alpha1.SnapshotControlVolumeName,
snapshotv1alpha1.SnapshotControlMountPath,
mapping.Destination,
)
}
if err := snapshotv1alpha1.ValidateRestorePod(
pod,
snapshot.Name,
mappings,
// Seccomp injection is optional producer policy. The node agent validates
// the universal restore contract and leaves profile selection to the
// producer that shaped the Pod.
snapshotv1alpha1.RestorePodOptions{},
); err != nil {
return nil, nil, err
}
return nil
return &restoreTarget{SnapshotName: snapshot.Name, ContentUID: string(content.UID), SourceContainerName: containerName}, mappings, nil
}

// resolveRestoreArtifact resolves the validated restore target to its physical
Expand Down
76 changes: 74 additions & 2 deletions agent/internal/controller/controller_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -206,7 +206,7 @@ func pendingRestoreReason(t *testing.T, err error) string {
}

func restorePod(annotations map[string]string) *corev1.Pod {
return &corev1.Pod{
pod := &corev1.Pod{
ObjectMeta: metav1.ObjectMeta{
Name: "restore-worker",
Namespace: "inference",
Expand All @@ -231,6 +231,21 @@ func restorePod(annotations map[string]string) *corev1.Pod {
}},
},
}
snapshotName, restoreRequested := annotations[snapshotv1alpha1.RestoreFromAnnotation]
_, hasExplicitMapping := annotations[snapshotv1alpha1.RestoreContainerMapAnnotation]
if !restoreRequested || hasExplicitMapping {
return pod
}
shaped, err := snapshotv1alpha1.BuildRestorePod(
pod,
snapshotName,
[]snapshotv1alpha1.RestoreContainerMapping{{Source: "main", Destination: "main"}},
snapshotv1alpha1.RestorePodOptions{},
)
if err != nil {
panic(err)
}
return shaped
}

func multiRestorePod() *corev1.Pod {
Expand Down Expand Up @@ -264,7 +279,19 @@ func multiRestorePod() *corev1.Pod {
{Name: "engine-0", ContainerID: "containerd://engine-0-id", State: corev1.ContainerState{Running: &corev1.ContainerStateRunning{}}},
{Name: "engine-1", ContainerID: "containerd://engine-1-id", State: corev1.ContainerState{Running: &corev1.ContainerStateRunning{}}},
}
return pod
shaped, err := snapshotv1alpha1.BuildRestorePod(
pod,
"snapshot-a",
[]snapshotv1alpha1.RestoreContainerMapping{
{Source: "main", Destination: "engine-0"},
{Source: "main", Destination: "engine-1"},
},
snapshotv1alpha1.RestorePodOptions{},
)
if err != nil {
panic(err)
}
return shaped
}

func processQueuedRestorePod(t *testing.T, w *NodeController, pod *corev1.Pod) {
Expand Down Expand Up @@ -571,6 +598,51 @@ func TestPreflightRestoreRejectsSharedMultiDestinationControlMount(t *testing.T)
assert.Contains(t, err.Error(), `subPath "engine-1"`)
}

func TestPreflightRestoreRejectsInvalidRestorePodContract(t *testing.T) {
snapshot, content := readySnapshotObjects()
tests := map[string]func(*corev1.Pod){
"control volume": func(pod *corev1.Pod) {
pod.Spec.Volumes = nil
},
"control environment": func(pod *corev1.Pod) {
pod.Spec.Containers[0].Env = nil
},
"startup gate": func(pod *corev1.Pod) {
pod.Spec.Containers[0].StartupProbe = nil
},
}
for name, mutate := range tests {
t.Run(name, func(t *testing.T) {
pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"})
mutate(pod)
w := makeTestController(t, pod, snapshot, content)

plan, err := w.preflightRestore(context.Background(), pod)

assert.Nil(t, plan)
require.Error(t, err)
})
}
}

func TestValidateRestoreTargetAcceptsEquivalentRestoreStartupGate(t *testing.T) {
pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"})
pod.Spec.Containers[0].StartupProbe = &corev1.Probe{
ProbeHandler: corev1.ProbeHandler{Exec: &corev1.ExecAction{Command: []string{
"test", "-f", snapshotv1alpha1.SnapshotControlMountPath + "/" + snapshotv1alpha1.RestoreCompleteFile,
}}},
PeriodSeconds: 7,
FailureThreshold: 42,
}
snapshot, content := readySnapshotObjects()

target, mappings, err := validateRestoreTarget(pod, snapshot, content)

require.NoError(t, err)
assert.Equal(t, "snapshot-a", target.SnapshotName)
assert.Equal(t, []snapshotv1alpha1.RestoreContainerMapping{{Source: "main", Destination: "main"}}, mappings)
}

func TestPreflightRestoreRetriesInProgressCondition(t *testing.T) {
pod := restorePod(map[string]string{snapshotv1alpha1.RestoreFromAnnotation: "snapshot-a"})
pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{
Expand Down
20 changes: 16 additions & 4 deletions api/v1alpha1/constants.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,16 +57,28 @@ const (
// control mount path to the workload.
SnapshotControlDirEnv = "SNAPSHOT_CONTROL_DIR"

// LegacySnapshotControlDirEnv is the environment variable exposing the
// control mount path to the workload. EnsureControlVolume injects both
// during the migration window so existing workload images (which read
// this name) keep working while new images can move to
// LegacySnapshotControlDirEnv is the deprecated environment variable
// exposing the control mount path to the workload. Snapshot Pod shaping
// injects both names during the migration window so existing workload
// images keep working while new images can move to
// SnapshotControlDirEnv.
//
// Deprecated: use SnapshotControlDirEnv instead. Remove once no workload
// image depends on this name.
LegacySnapshotControlDirEnv = "DYN_SNAPSHOT_CONTROL_DIR"

// RestoreStandbyModeEnv asks standby-aware workload entrypoints to remain
// inert until Snapshot replaces them with restored processes. Snapshot's
// generic Pod builder does not inject this workload-specific setting.
RestoreStandbyModeEnv = "SNAPSHOT_RESTORE_STANDBY"

// LegacyRestoreStandbyModeEnv is the deprecated Dynamo restore standby
// environment variable. Snapshot publishes the name for producers that
// support existing workload images but does not inject it.
//
// Deprecated: use RestoreStandbyModeEnv for new workload integrations.
LegacyRestoreStandbyModeEnv = "DYN_SNAPSHOT_RESTORE_STANDBY"

// SnapshotCompleteFile named the sentinel the agent used to release a
// checkpointed workload when leave-running dumps existed. A checkpoint now
// always terminates the source process, so the agent no longer writes it;
Expand Down
6 changes: 5 additions & 1 deletion api/v1alpha1/protocol.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,11 @@ type RestoreContainerMapping struct {
// GetRestoreFromSnapshotName returns the same-namespace PodSnapshot named by
// the restore-from annotation.
func GetRestoreFromSnapshotName(annotations map[string]string) (string, error) {
snapshotName := strings.TrimSpace(annotations[RestoreFromAnnotation])
return validateRestoreFromSnapshotName(annotations[RestoreFromAnnotation])
}

func validateRestoreFromSnapshotName(value string) (string, error) {
snapshotName := strings.TrimSpace(value)
if snapshotName == "" {
return "", fmt.Errorf("%s must name a PodSnapshot", RestoreFromAnnotation)
}
Expand Down
Loading