feat(snapshot): complete concurrent PageBroker staging

container/compliance/native_packages.yaml

@@ -185,3 +185,12 @@ packages:
     source: https://github.com/ai-dynamo/dynamo
     images:
       - snapshot-agent
+
+  - name: pagebroker
+    # First-party C++ daemon built from deploy/snapshot/pagebroker and copied
+    # into the snapshot-agent image.
+    version: 1.0
+    license: Apache-2.0
+    source: https://github.com/ai-dynamo/dynamo
+    images:
+      - snapshot-agent

deploy/helm/charts/snapshot/templates/daemonset.yaml

@@ -113,6 +113,8 @@ spec:
             {{- else }}
             - {{ printf "%s/pagebroker.sock" $pageBrokerControlPath | quote }}
             - {{ $pageBrokerStagingPath | quote }}
+            - --max-concurrency
+            - {{ default 16 .Values.pageBroker.maxConcurrency | quote }}
             {{- end }}
           volumeMounts:
             - name: pagebroker

deploy/helm/charts/snapshot/values.yaml

@@ -84,6 +84,8 @@ pageBroker:
   command: []
   # Optional arguments override. Defaults to the PageBroker socket and staging path.
   args: []
+  # Maximum simultaneous PageBroker socket handlers.
+  maxConcurrency: 16
   # Optional maximum size for tmpfs staging, for example "100Gi".
   stagingSizeLimit: ""
   image:

deploy/snapshot/internal/executor/restore.go

@@ -66,6 +66,10 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r
 	transactionID := ""
 	var broker pagebroker.Client
 	committed := false
+	var pageBrokerStageDuration time.Duration
+	var pageBrokerStagingMountDuration time.Duration
+	var pageBrokerStagingUnmountDuration time.Duration
+	var pageBrokerCommitDuration time.Duration
 	brokered := req.PageBrokerRequested && req.PageBrokerEnabled
 	if brokered {
 		transactionID = uuid.NewString()
@@ -77,7 +81,9 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r
 				_ = broker.Abort(abortCtx, transactionID)
 			}
 		}()
+		stageStart := time.Now()
 		staged, err := broker.StagedRestore(ctx, transactionID, req.CheckpointLocation)
+		pageBrokerStageDuration = time.Since(stageStart)
 		if err != nil {
 			return 0, fmt.Errorf("stage PageBroker restore: %w", err)
 		}
@@ -100,6 +106,7 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r
 	}
 	injectDuration := time.Since(injectStart)
 	defer func() {
+		unmountStart := time.Now()
 		// Pass a background context: mp.Unmount has its own internal timeout
 		// (nsmount.unmountTimeout) around the ns-bind-mount subprocess.
 		if cleanupErr := mp.Unmount(context.Background()); cleanupErr != nil {
@@ -109,10 +116,12 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r
 			// already restored successfully. Log it and let the pod continue.
 			log.Error(cleanupErr, "failed to unmount agent bundle from placeholder namespace")
 		}
+		log.Info("Agent bundle unmount timing", "duration", time.Since(unmountStart))
 	}()
 
 	var mountedStaging nsmount.MountPoint
 	if brokered {
+		stagingMountStart := time.Now()
 		stagingMounter, err := nsmount.New(req.CheckpointLocation, nsmount.PageBrokerDst, log)
 		if err != nil {
 			return 0, fmt.Errorf("create PageBroker staging mount: %w", err)
@@ -121,36 +130,51 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r
 		if err != nil {
 			return 0, fmt.Errorf("mount PageBroker staging: %w", err)
 		}
+		pageBrokerStagingMountDuration = time.Since(stagingMountStart)
 		req.ContainerCheckpointLocation = nsmount.PageBrokerDst
 	}
 
 	// Phase 3: Execute — nsrestore handles rootfs, CRIU restore, and CUDA restore inside namespace.
 	result, err := execNSRestore(ctx, log, req, snap, mp)
 	if mountedStaging != nil {
+		stagingUnmountStart := time.Now()
 		if cleanupErr := mountedStaging.Unmount(context.Background()); cleanupErr != nil {
 			log.Error(cleanupErr, "failed to unmount PageBroker staging from placeholder namespace")
 		}
+		pageBrokerStagingUnmountDuration = time.Since(stagingUnmountStart)
 	}
 	if err != nil {
 		return 0, fmt.Errorf("nsrestore failed: %w", err)
 	}
 	if brokered {
+		commitStart := time.Now()
 		if err := broker.Commit(ctx, transactionID); err != nil {
 			log.Error(err, "failed to commit PageBroker restore")
 		} else {
 			committed = true
 		}
+		pageBrokerCommitDuration = time.Since(commitStart)
 	}
-	restoreDuration := hostInspectDuration + injectDuration + result.TotalDuration()
+
+	validationStart := time.Now()
+	if err := validateRestoredProcess(snap.TargetRoot, result.RestoredPID, log); err != nil {
+		return 0, err
+	}
+	validationDuration := time.Since(validationStart)
 	log.Info("Restore timing summary",
 		"restore", map[string]any{
-			"duration": restoreDuration.String(),
+			"duration": time.Since(restoreStart).String(),
 			"phases": map[string]string{
-				"host_inspect_duration":    hostInspectDuration.String(),
-				"inject_duration":          injectDuration.String(),
-				"nsrestore_setup_duration": result.NSRestoreSetupDuration.String(),
-				"criu_restore_duration":    result.CRIURestoreDuration.String(),
-				"cuda_duration":            result.CUDADuration.String(),
+				"pagebroker_stage_duration":           pageBrokerStageDuration.String(),
+				"host_inspect_duration":               hostInspectDuration.String(),
+				"inject_duration":                     injectDuration.String(),
+				"pagebroker_staging_mount_duration":   pageBrokerStagingMountDuration.String(),
+				"nsrestore_setup_duration":            result.NSRestoreSetupDuration.String(),
+				"criu_restore_duration":               result.CRIURestoreDuration.String(),
+				"cuda_duration":                       result.CUDADuration.String(),
+				"pagebroker_staging_unmount_duration": pageBrokerStagingUnmountDuration.String(),
+				"pagebroker_commit_duration":          pageBrokerCommitDuration.String(),
+				"validation_duration":                 validationDuration.String(),
 			},
 		},
 	)
@@ -160,15 +184,10 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r
 		)
 	}
 
-	validationStart := time.Now()
-	if err := validateRestoredProcess(snap.TargetRoot, result.RestoredPID, log); err != nil {
-		return 0, err
-	}
-
 	log.Info("=== External restore completed ===",
 		"restored_pid", result.RestoredPID,
 		"placeholder_host_pid", snap.PlaceholderPID,
-		"validation_duration", time.Since(validationStart),
+		"validation_duration", validationDuration,
 		"total_duration", time.Since(restoreStart),
 	)
 

deploy/snapshot/internal/pagebroker/client.go

@@ -20,6 +20,7 @@ const (
 	// PageBroker control requests and responses are limited to 64 KiB.
 	maxMessageSize   = 64 << 10
 	commitRetryDelay = 100 * time.Millisecond
+	commitRetryLimit = 30 * time.Second
 )
 
 var errMessageTooLarge = fmt.Errorf("message exceeds %d bytes", maxMessageSize)
@@ -56,6 +57,7 @@ func (c Client) PrepareCheckpoint(ctx context.Context, transactionID, destinatio
 }
 
 func (c Client) Commit(ctx context.Context, transactionID string) error {
+	var retryDeadline time.Time
 	for {
 		response, err := c.request(ctx, transactionID, &Request_Commit{Commit: &CommitRequest{}})
 		if err == nil {
@@ -68,10 +70,17 @@ func (c Client) Commit(ctx context.Context, transactionID string) error {
 		if !errors.As(err, &transport) {
 			return err
 		}
+		if retryDeadline.IsZero() {
+			retryDeadline = time.Now().Add(commitRetryLimit)
+		}
+		delay := min(commitRetryDelay, time.Until(retryDeadline))
+		if delay <= 0 {
+			return err
+		}
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
-		case <-time.After(commitRetryDelay):
+		case <-time.After(delay):
 		}
 	}
 }

deploy/snapshot/internal/types/config.go

@@ -67,6 +67,9 @@ func (c *AgentConfig) Validate() error {
 		}
 	}
 	c.Storage.AccessMode = accessMode
+	if c.PageBroker.Enabled && strings.TrimSpace(c.PageBroker.ControlSocketPath) == "" {
+		return &ConfigError{Field: "pageBroker.controlSocketPath", Message: "pageBroker.controlSocketPath is required when PageBroker is enabled"}
+	}
 	if c.CRIU.TcpClose && c.CRIU.TcpEstablished {
 		return &ConfigError{
 			Field:   "criu",

deploy/snapshot/internal/types/config_test.go

@@ -50,3 +50,12 @@ func TestAgentConfigValidateDefaultsStorageAccessMode(t *testing.T) {
 		t.Fatalf("Storage.AccessMode = %q, want %q", cfg.Storage.AccessMode, StorageAccessModeAgentMount)
 	}
 }
+
+func TestAgentConfigValidateRequiresPageBrokerControlSocket(t *testing.T) {
+	cfg := validAgentConfig()
+	cfg.PageBroker.Enabled = true
+
+	if err := cfg.Validate(); err == nil {
+		t.Fatal("expected error for missing PageBroker control socket")
+	}
+}

deploy/snapshot/pagebroker/Makefile

@@ -1,7 +1,7 @@
 PROTO := v1/pagebroker.proto
-GTEST_FLAGS := $(shell pkg-config --cflags --libs gtest_main)
-BROKER_SOURCES := broker.cpp checkpoint_transaction_descriptor.cpp posix_copy_engine.cpp restore_transaction_descriptor.cpp transfer_engine.cpp
-DAEMON_SOURCES := $(BROKER_SOURCES) daemon.cpp file_descriptor.cpp
+GTEST_FLAGS = $(shell pkg-config --cflags --libs gtest_main)
+BROKER_SOURCES := broker.cpp checkpoint_transaction_descriptor.cpp posix_copy_engine.cpp restore_transaction_descriptor.cpp transaction.cpp transfer_engine.cpp
+DAEMON_SOURCES := $(BROKER_SOURCES) daemon.cpp main.cpp file_descriptor.cpp
 
 .PHONY: daemon generate test