feat(agent): route opted-in restores through PageBroker

agent/internal/controller/controller.go

@@ -563,16 +563,19 @@ func (w *NodeController) runRestore(ctx context.Context, pod *corev1.Pod, contai
 	}
 
 	req := executor.RestoreRequest{
-		CheckpointID:    checkpointID,
-		ArtifactVersion: artifactVersionFromPod(pod),
-		BasePath:        w.config.Storage.BasePath,
-		ContainerID:     containerID,
-		StartedAt:       startedAt,
-		PodName:         pod.Name,
-		PodNamespace:    pod.Namespace,
-		TargetPodIP:     pod.Status.PodIP,
-		ContainerName:   containerName,
-		Clientset:       w.clientset,
+		CheckpointID:                checkpointID,
+		ArtifactVersion:             artifactVersionFromPod(pod),
+		BasePath:                    w.config.Storage.BasePath,
+		ContainerID:                 containerID,
+		StartedAt:                   startedAt,
+		PodName:                     pod.Name,
+		PodNamespace:                pod.Namespace,
+		TargetPodIP:                 pod.Status.PodIP,
+		ContainerName:               containerName,
+		Clientset:                   w.clientset,
+		PageBrokerRequested:         pod.Annotations[snapshotv1alpha1.PageBrokerAnnotation] == snapshotv1alpha1.PageBrokerAnnotationEnabled,
+		PageBrokerEnabled:           w.config.PageBroker.Enabled,
+		PageBrokerControlSocketPath: w.config.PageBroker.ControlSocketPath,
 	}
 	placeholderHostPID, err := w.restoreFn(restoreCtx, w.runtime, log, req, w.injector)
 	if err != nil {

agent/internal/controller/controller_test.go

@@ -76,6 +76,10 @@ func (noopInjector) MountArtifact(_ context.Context, _ nsmount.MountPoint, _ str
 	return noopMountPoint{}, nil
 }
 
+func (noopInjector) MountPageBroker(_ context.Context, _ nsmount.MountPoint, _ string) (nsmount.MountPoint, error) {
+	return noopMountPoint{}, nil
+}
+
 type noopMountPoint struct{}
 
 func (noopMountPoint) Unmount(context.Context) error { return nil }
@@ -92,6 +96,10 @@ func (e errorInjector) MountArtifact(_ context.Context, _ nsmount.MountPoint, _
 	return nil, e.err
 }
 
+func (e errorInjector) MountPageBroker(_ context.Context, _ nsmount.MountPoint, _ string) (nsmount.MountPoint, error) {
+	return nil, e.err
+}
+
 type recordedMountCall struct {
 	src            string
 	dst            string
@@ -121,6 +129,14 @@ func (r recordingInjector) MountArtifact(_ context.Context, namespaceMount nsmou
 	return noopMountPoint{}, nil
 }
 
+func (r recordingInjector) MountPageBroker(_ context.Context, namespaceMount nsmount.MountPoint, src string) (nsmount.MountPoint, error) {
+	*r.calls = append(*r.calls, recordedMountCall{src: src, dst: nsmount.PageBrokerDst, namespaceMount: namespaceMount})
+	if r.artifactErr != nil {
+		return nil, r.artifactErr
+	}
+	return noopMountPoint{}, nil
+}
+
 var _ executor.RestoreMounter = noopInjector{}
 var _ executor.RestoreMounter = errorInjector{}
 var _ executor.RestoreMounter = recordingInjector{}

agent/internal/executor/restore.go

@@ -16,12 +16,14 @@ import (
 	"time"
 
 	"github.com/go-logr/logr"
+	"github.com/google/uuid"
 	"k8s.io/client-go/kubernetes"
 
 	"github.com/ai-dynamo/snapshot/agent/internal/criu"
 	"github.com/ai-dynamo/snapshot/agent/internal/cuda"
 	"github.com/ai-dynamo/snapshot/agent/internal/logging"
 	"github.com/ai-dynamo/snapshot/agent/internal/nsmount"
+	"github.com/ai-dynamo/snapshot/agent/internal/pagebroker"
 	snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime"
 	"github.com/ai-dynamo/snapshot/agent/internal/types"
 )
@@ -31,6 +33,7 @@ import (
 type RestoreMounter interface {
 	MountBundle(ctx context.Context, pid int) (nsmount.MountPoint, error)
 	MountArtifact(ctx context.Context, namespaceMount nsmount.MountPoint, artifactPath string) (nsmount.MountPoint, error)
+	MountPageBroker(ctx context.Context, namespaceMount nsmount.MountPoint, stagingPath string) (nsmount.MountPoint, error)
 }
 
 // RestoreCleanupError reports a successful restore whose cleanup did not fully
@@ -64,16 +67,19 @@ func cleanupRestoreMounts(ctx context.Context, mounts []restoreMount) error {
 
 // RestoreRequest holds the parameters for a restore operation.
 type RestoreRequest struct {
-	CheckpointID    string
-	ArtifactVersion string
-	BasePath        string
-	ContainerID     string
-	StartedAt       time.Time
-	PodName         string
-	PodNamespace    string
-	TargetPodIP     string
-	ContainerName   string
-	Clientset       kubernetes.Interface
+	CheckpointID                string
+	ArtifactVersion             string
+	BasePath                    string
+	ContainerID                 string
+	StartedAt                   time.Time
+	PodName                     string
+	PodNamespace                string
+	TargetPodIP                 string
+	ContainerName               string
+	Clientset                   kubernetes.Interface
+	PageBrokerRequested         bool
+	PageBrokerEnabled           bool
+	PageBrokerControlSocketPath string
 }
 
 // Restore performs external restore for the given request.
@@ -140,19 +146,56 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r
 		point:  bundleMount,
 	})
 
-	artifactMount, err := mounts.MountArtifact(ctx, bundleMount, artifactPath)
-	if err != nil {
-		return 0, fmt.Errorf("mount checkpoint artifact into placeholder: %w", err)
+	containerCheckpointPath := nsmount.CheckpointDst
+	brokered := req.PageBrokerRequested && req.PageBrokerEnabled
+	transactionID := ""
+	var broker pagebroker.Client
+	committed := false
+	if brokered {
+		transactionID = uuid.NewString()
+		broker = pagebroker.Client{ControlSocketPath: req.PageBrokerControlSocketPath}
+		defer func() {
+			if !committed {
+				abortCtx, cancel := context.WithTimeout(context.Background(), pageBrokerAbortTimeout)
+				defer cancel()
+				_ = broker.Abort(abortCtx, transactionID)
+			}
+		}()
+		staged, err := broker.StagedRestore(ctx, transactionID, artifactPath)
+		if err != nil {
+			return 0, fmt.Errorf("stage PageBroker restore: %w", err)
+		}
+		stagingMount, err := mounts.MountPageBroker(ctx, bundleMount, staged)
+		if err != nil {
+			return 0, fmt.Errorf("mount PageBroker staging: %w", err)
+		}
+		activeMounts = append(activeMounts, restoreMount{
+			action: "unmount PageBroker staging from placeholder",
+			point:  stagingMount,
+		})
+		containerCheckpointPath = nsmount.PageBrokerDst
+	} else {
+		artifactMount, err := mounts.MountArtifact(ctx, bundleMount, artifactPath)
+		if err != nil {
+			return 0, fmt.Errorf("mount checkpoint artifact into placeholder: %w", err)
+		}
+		activeMounts = append(activeMounts, restoreMount{
+			action: "unmount checkpoint artifact from placeholder",
+			point:  artifactMount,
+		})
 	}
-	activeMounts = append(activeMounts, restoreMount{
-		action: "unmount checkpoint artifact from placeholder",
-		point:  artifactMount,
-	})
 
-	result, err := execNSRestore(ctx, log, req, snap, bundleMount, nsmount.CheckpointDst)
+	result, err := execNSRestore(ctx, log, req, snap, bundleMount, containerCheckpointPath)
 	if err != nil {
 		return 0, fmt.Errorf("nsrestore failed: %w", err)
 	}
+	if brokered {
+		if err := broker.Commit(ctx, transactionID); err != nil {
+			log.Error(err, "failed to commit PageBroker restore")
+		} else {
+			committed = true
+		}
+	}
 	if result.CleanupError != nil {
 		cleanupErr = errors.Join(cleanupErr, result.CleanupError)
 	}

agent/internal/pagebroker/client.go

@@ -29,6 +29,16 @@ type Client struct {
 	ControlSocketPath string
 }
 
+func (c Client) StagedRestore(ctx context.Context, transactionID, source string) (string, error) {
+	response, err := c.request(ctx, transactionID, &Request_StagedRestore{
+		StagedRestore: &StagedRestoreRequest{Source: filesystem(source), IoEngine: posixCopy()},
+	})
+	if err != nil {
+		return "", err
+	}
+	return imageDirectory(response.GetStagedRestoreDirectory().GetImageDirectory())
+}
+
 func (c Client) PrepareCheckpoint(ctx context.Context, transactionID, destination string) (string, error) {
 	response, err := c.request(ctx, transactionID, &Request_PrepareStagedCheckpoint{
 		PrepareStagedCheckpoint: &PrepareStagedCheckpointRequest{Destination: filesystem(destination), IoEngine: posixCopy()},

agent/internal/pagebroker/client_test.go

@@ -173,6 +173,17 @@ func TestStagingRequestsRejectEmptyDirectory(t *testing.T) {
 		call  func(Client, context.Context) error
 		reply func(*Request) *Response
 	}{
+		{
+			name: "restore",
+			call: func(client Client, ctx context.Context) error {
+				_, err := client.StagedRestore(ctx, "transaction", "/checkpoints/source")
+				return err
+			},
+			reply: func(request *Request) *Response {
+				return &Response{RequestId: request.RequestId, TransactionId: request.TransactionId,
+					Result: &Response_StagedRestoreDirectory{StagedRestoreDirectory: &StagedRestoreDirectory{}}}
+			},
+		},
 		{
 			name: "checkpoint",
 			call: func(client Client, ctx context.Context) error {