diff --git a/agent/Dockerfile b/agent/Dockerfile index ab4e310c..2a72e819 100644 --- a/agent/Dockerfile +++ b/agent/Dockerfile @@ -81,6 +81,19 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ RUN go test ./... -v && make -C pagebroker test +# ============================================================================= +# Stage: PageBroker Builder +# ============================================================================= +FROM go-base AS pagebroker-builder + +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + libprotobuf-dev \ + protobuf-compiler \ + && rm -rf /var/lib/apt/lists/* + +RUN make -C pagebroker daemon + # ============================================================================= # Stage: Builder - Build Go binaries # ============================================================================= @@ -189,6 +202,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ libnet1 \ libnl-3-200 \ libnl-route-3-200 \ + libprotobuf32t64 \ libprotobuf-c1 \ libgnutls30t64 \ libnftables1 \ @@ -242,6 +256,10 @@ RUN chmod +x /usr/local/sbin/cuda-checkpoint /usr/local/bin/cuda-checkpoint-help # Copy the built binaries COPY --from=builder /snapshot-agent /usr/local/bin/snapshot-agent +COPY --from=pagebroker-builder /workspace/pagebroker/pagebroker /usr/local/bin/pagebroker +RUN output=$(/usr/local/bin/pagebroker 2>&1); status=$?; \ + [ "$status" -eq 2 ]; \ + [ "$output" = "usage: pagebroker socket_path staging_directory storage_root --max-concurrent-requests max_concurrent_requests" ] # Assemble the injection bundle: bind-mounted read-only into the placeholder's # mount namespace at /tmp/snapshot-binaries during restore (internal/nsmount). diff --git a/agent/pagebroker/Dockerfile b/agent/pagebroker/Dockerfile index 5ed43faa..a7c801b5 100644 --- a/agent/pagebroker/Dockerfile +++ b/agent/pagebroker/Dockerfile @@ -15,7 +15,7 @@ RUN make daemon FROM ubuntu:24.04 -RUN apt-get update && apt-get install -y --no-install-recommends libprotobuf32 \ +RUN apt-get update && apt-get install -y --no-install-recommends libprotobuf32t64 \ && rm -rf /var/lib/apt/lists/* COPY --from=build /src/pagebroker /usr/local/bin/pagebroker diff --git a/charts/snapshot/templates/daemonset.yaml b/charts/snapshot/templates/daemonset.yaml index 5463b3ed..d3e8d583 100644 --- a/charts/snapshot/templates/daemonset.yaml +++ b/charts/snapshot/templates/daemonset.yaml @@ -110,12 +110,22 @@ spec: - name: pagebroker image: "{{ .Values.pageBroker.image.repository }}:{{ .Values.pageBroker.image.tag }}" imagePullPolicy: {{ .Values.pageBroker.image.pullPolicy }} + command: + {{- with .Values.pageBroker.command }} + {{- toYaml . | nindent 12 }} + {{- else }} + - /usr/local/bin/pagebroker + {{- end }} args: + {{- with .Values.pageBroker.args }} + {{- toYaml . | nindent 12 }} + {{- else }} - {{ printf "%s/pagebroker.sock" $pageBrokerControlPath | quote }} - {{ $pageBrokerStagingPath | quote }} - {{ .Values.storage.pvc.basePath | quote }} - --max-concurrent-requests - {{ .Values.pageBroker.maxConcurrentRequests | quote }} + {{- end }} resources: {{- toYaml .Values.pageBroker.resources | nindent 12 }} volumeMounts: diff --git a/charts/snapshot/values.yaml b/charts/snapshot/values.yaml index f5c67051..85015f5e 100644 --- a/charts/snapshot/values.yaml +++ b/charts/snapshot/values.yaml @@ -94,6 +94,8 @@ storage: pageBroker: enabled: false + command: [] + args: [] maxConcurrentRequests: 16 stagingSizeLimit: 64Gi resources: @@ -104,7 +106,7 @@ pageBroker: cpu: 32 memory: 72Gi image: - repository: nvcr.io/nvidia/ai-dynamo/pagebroker + repository: ghcr.io/ai-dynamo/snapshot/agent tag: 1.4.0 pullPolicy: Always