From 1403a9662eaaa35b5da2ecca47514dfbc7ca0c66 Mon Sep 17 00:00:00 2001
From: Dan Feigin <dfeigin@nvidia.com>
Date: Thu, 13 Aug 2026 10:59:38 +0300
Subject: [PATCH] feat(snapshot): deploy opted-in tmpfs PageBroker

---
 .../charts/snapshot/templates/_helpers.tpl    |  2 ++
 .../charts/snapshot/templates/configmap.yaml  |  9 +++++++
 .../charts/snapshot/templates/daemonset.yaml  | 27 +++++++++++++++++++
 deploy/helm/charts/snapshot/values.yaml       |  9 +++++++
 4 files changed, 47 insertions(+)

diff --git a/deploy/helm/charts/snapshot/templates/_helpers.tpl b/deploy/helm/charts/snapshot/templates/_helpers.tpl
index c62950700e1f..45cb70b12dfa 100644
--- a/deploy/helm/charts/snapshot/templates/_helpers.tpl
+++ b/deploy/helm/charts/snapshot/templates/_helpers.tpl
@@ -105,3 +105,5 @@ reads for rootfs-diff capture, and CRI-O config.json fallback).
 {{- if eq .Values.runtime.type "crio" -}}/var/lib/containers{{- else -}}/var/lib/containerd{{- end -}}
 {{- end }}
 
+{{- define "snapshot.pageBrokerControlPath" -}}/pagebroker/control{{- end -}}
+{{- define "snapshot.pageBrokerStagingPath" -}}/pagebroker/staging{{- end -}}
diff --git a/deploy/helm/charts/snapshot/templates/configmap.yaml b/deploy/helm/charts/snapshot/templates/configmap.yaml
index 9b323cd98292..c29efb9ebafd 100644
--- a/deploy/helm/charts/snapshot/templates/configmap.yaml
+++ b/deploy/helm/charts/snapshot/templates/configmap.yaml
@@ -9,6 +9,9 @@ kind: ConfigMap
 {{- if not (has .Values.storage.accessMode (list "agentMount" "podMount")) }}
 {{- fail (printf "snapshot.storage.accessMode=%q is not supported; expected agentMount or podMount" .Values.storage.accessMode) }}
 {{- end }}
+{{- if and .Values.pageBroker.enabled (ne .Values.storage.accessMode "agentMount") }}
+{{- fail "snapshot.pageBroker requires storage.accessMode=agentMount" }}
+{{- end }}
 metadata:
   name: {{ include "snapshot.fullname" . }}-config
   namespace: {{ .Release.Namespace }}
@@ -26,6 +29,12 @@ data:
     overlay:
       exclusions: {{ toYaml .Values.config.overlay.exclusions | nindent 8 }}
 
+    pageBroker:
+      enabled: {{ .Values.pageBroker.enabled }}
+      {{- if .Values.pageBroker.enabled }}
+      controlSocketPath: {{ printf "%s/pagebroker.sock" (include "snapshot.pageBrokerControlPath" .) | quote }}
+      {{- end }}
+
     restore:
       restoreTimeoutSeconds: {{ .Values.config.restore.restoreTimeoutSeconds }}
 
diff --git a/deploy/helm/charts/snapshot/templates/daemonset.yaml b/deploy/helm/charts/snapshot/templates/daemonset.yaml
index b780784be487..26ca3da2dff8 100644
--- a/deploy/helm/charts/snapshot/templates/daemonset.yaml
+++ b/deploy/helm/charts/snapshot/templates/daemonset.yaml
@@ -2,6 +2,8 @@
 {{- $socket := include "snapshot.runtimeSocket" . -}}
 {{- $socketDir := $socket | dir -}}
 {{- $storageDir := include "snapshot.runtimeStorageDir" . -}}
+{{- $pageBrokerControlPath := include "snapshot.pageBrokerControlPath" . -}}
+{{- $pageBrokerStagingPath := include "snapshot.pageBrokerStagingPath" . -}}
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
@@ -95,6 +97,19 @@ spec:
               mountPath: /host-seccomp
       {{- end }}
       containers:
+        {{- if .Values.pageBroker.enabled }}
+        - name: pagebroker
+          image: "{{ .Values.pageBroker.image.repository }}:{{ .Values.pageBroker.image.tag }}"
+          imagePullPolicy: {{ .Values.pageBroker.image.pullPolicy }}
+          args:
+            - {{ printf "%s/pagebroker.sock" $pageBrokerControlPath | quote }}
+            - {{ $pageBrokerStagingPath | quote }}
+          volumeMounts:
+            - name: pagebroker
+              mountPath: /pagebroker
+            - name: checkpoints
+              mountPath: {{ .Values.storage.pvc.basePath }}
+        {{- end }}
         - name: agent
           image: "{{ .Values.daemonset.image.repository }}:{{ .Values.daemonset.image.tag }}"
           imagePullPolicy: {{ .Values.daemonset.image.pullPolicy }}
@@ -125,6 +140,10 @@ spec:
             - name: config
               mountPath: /etc/snapshot
               readOnly: true
+            {{- if .Values.pageBroker.enabled }}
+            - name: pagebroker
+              mountPath: /pagebroker
+            {{- end }}
             {{- if and (eq .Values.storage.type "pvc") (eq .Values.storage.accessMode "agentMount") }}
             # Mount the checkpoint PVC directly in namespace-local agent mode
             - name: checkpoints
@@ -171,6 +190,14 @@ spec:
         - name: config
           configMap:
             name: {{ include "snapshot.fullname" . }}-config
+        {{- if .Values.pageBroker.enabled }}
+        - name: pagebroker
+          emptyDir:
+            medium: Memory
+            {{- with .Values.pageBroker.stagingSizeLimit }}
+            sizeLimit: {{ . | quote }}
+            {{- end }}
+        {{- end }}
         {{- if .Values.seccomp.deploy }}
         # Seccomp profile ConfigMap (used by initContainer)
         - name: seccomp-profiles
diff --git a/deploy/helm/charts/snapshot/values.yaml b/deploy/helm/charts/snapshot/values.yaml
index 551bf90f9261..c8d7f3600aae 100644
--- a/deploy/helm/charts/snapshot/values.yaml
+++ b/deploy/helm/charts/snapshot/values.yaml
@@ -75,6 +75,15 @@ storage:
   oci:
     uri: ""
 
+pageBroker:
+  enabled: false
+  # Optional maximum size for tmpfs staging, for example "100Gi".
+  stagingSizeLimit: ""
+  image:
+    repository: nvcr.io/nvidia/ai-dynamo/pagebroker
+    tag: 1.4.0
+    pullPolicy: Always
+
 # DaemonSet configuration for snapshot (checkpoint/restore) agent
 daemonset:
   # Container image