diff --git a/deploy/snapshot/internal/controller/podsnapshotcontent.go b/deploy/snapshot/internal/controller/podsnapshotcontent.go index f2901681622f..61fd087aba37 100644 --- a/deploy/snapshot/internal/controller/podsnapshotcontent.go +++ b/deploy/snapshot/internal/controller/podsnapshotcontent.go @@ -441,16 +441,17 @@ func (w *NodeController) executorCheckpoint(ctx context.Context, params Checkpoi log := logr.FromContextOrDiscard(ctx) req := executor.CheckpointRequest{ - ContainerID: params.ContainerID, - ContainerName: params.ContainerName, - CheckpointID: params.CheckpointID, - CheckpointLocation: params.HostPath, - StartedAt: params.StartedAt, - NodeName: w.config.NodeName, - PodName: params.Pod.Name, - PodNamespace: params.Pod.Namespace, - PodIP: params.Pod.Status.PodIP, - Clientset: w.clientset, + ContainerID: params.ContainerID, + ContainerName: params.ContainerName, + CheckpointID: params.CheckpointID, + CheckpointLocation: params.HostPath, + StartedAt: params.StartedAt, + NodeName: w.config.NodeName, + PodName: params.Pod.Name, + PodNamespace: params.Pod.Namespace, + PodIP: params.Pod.Status.PodIP, + Clientset: w.clientset, + PageBrokerRequested: params.Pod.Annotations[snapshotprotocol.PageBrokerAnnotation] == snapshotprotocol.PageBrokerAnnotationEnabled, } if err := executor.Checkpoint(ctx, w.runtime, log, req, w.config); err != nil { w.killCheckpointProcess(log, params.ContainerPID, "checkpoint failed") diff --git a/deploy/snapshot/internal/executor/checkpoint.go b/deploy/snapshot/internal/executor/checkpoint.go index 7c4ad2abda72..da204f2244c7 100644 --- a/deploy/snapshot/internal/executor/checkpoint.go +++ b/deploy/snapshot/internal/executor/checkpoint.go @@ -17,22 +17,26 @@ import ( "github.com/ai-dynamo/dynamo/deploy/snapshot/internal/criu" "github.com/ai-dynamo/dynamo/deploy/snapshot/internal/cuda" + "github.com/ai-dynamo/dynamo/deploy/snapshot/internal/pagebroker" snapshotruntime "github.com/ai-dynamo/dynamo/deploy/snapshot/internal/runtime" "github.com/ai-dynamo/dynamo/deploy/snapshot/internal/types" ) +const pageBrokerAbortTimeout = 5 * time.Second + // CheckpointRequest holds per-checkpoint identifiers for a checkpoint operation. type CheckpointRequest struct { - ContainerID string - ContainerName string - CheckpointID string - CheckpointLocation string - StartedAt time.Time - NodeName string - PodName string - PodNamespace string - PodIP string - Clientset kubernetes.Interface + ContainerID string + ContainerName string + CheckpointID string + CheckpointLocation string + StartedAt time.Time + NodeName string + PodName string + PodNamespace string + PodIP string + Clientset kubernetes.Interface + PageBrokerRequested bool } type checkpointPhaseTimings struct { @@ -46,9 +50,8 @@ type checkpointPhaseTimings struct { // Checkpoint performs a CRIU dump of a container. // The operation has three phases: inspect, configure, capture. // -// The checkpoint directory is staged under tmp/ during the operation. -// On success, the previous checkpoint is removed and the staged directory is -// renamed into place at the base path root. +// PageBroker uses its tmpfs staging directory when both the Pod and deployment enable it; +// otherwise the existing tmp/ staging and rename path is unchanged. func Checkpoint(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, req CheckpointRequest, cfg *types.AgentConfig) error { checkpointStart := time.Now() phaseTimings := checkpointPhaseTimings{} @@ -63,15 +66,36 @@ func Checkpoint(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger } finalDir := req.CheckpointLocation - tmpRoot := filepath.Join(filepath.Dir(finalDir), "tmp") - if err := os.MkdirAll(tmpRoot, 0700); err != nil { - return fmt.Errorf("failed to create checkpoint staging root: %w", err) - } - tmpDir := filepath.Join(tmpRoot, uuid.NewString()) - if err := os.Mkdir(tmpDir, 0700); err != nil { - return fmt.Errorf("failed to create checkpoint staging directory: %w", err) + tmpDir := "" + brokered := req.PageBrokerRequested && cfg.PageBroker.Enabled + transactionID := uuid.NewString() + var broker pagebroker.Client + committed := false + if brokered { + broker = pagebroker.Client{ControlSocketPath: cfg.PageBroker.ControlSocketPath} + defer func() { + if !committed { + abortCtx, cancel := context.WithTimeout(context.Background(), pageBrokerAbortTimeout) + defer cancel() + _ = broker.Abort(abortCtx, transactionID) + } + }() + var err error + tmpDir, err = broker.PrepareCheckpoint(ctx, transactionID, finalDir) + if err != nil { + return fmt.Errorf("prepare PageBroker checkpoint: %w", err) + } + } else { + tmpRoot := filepath.Join(filepath.Dir(finalDir), "tmp") + if err := os.MkdirAll(tmpRoot, 0700); err != nil { + return fmt.Errorf("failed to create checkpoint staging root: %w", err) + } + tmpDir = filepath.Join(tmpRoot, transactionID) + if err := os.Mkdir(tmpDir, 0700); err != nil { + return fmt.Errorf("failed to create checkpoint staging directory: %w", err) + } + defer os.RemoveAll(tmpDir) } - defer os.RemoveAll(tmpDir) // Phase 1: Inspect container state state, err := inspectContainer(ctx, rt, log, req) @@ -95,14 +119,21 @@ func Checkpoint(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger phaseTimings.CRIUDumpDuration = captureTimings.CRIUDumpDuration phaseTimings.OverlayCaptureDuration = captureTimings.OverlayCaptureDuration - // Remove any previous checkpoint with the same identity hash, then - // promote the staged checkpoint directory into place. finalizeStart := time.Now() - if err := os.RemoveAll(finalDir); err != nil { - return fmt.Errorf("failed to remove previous checkpoint directory: %w", err) - } - if err := os.Rename(tmpDir, finalDir); err != nil { - return fmt.Errorf("failed to finalize checkpoint directory: %w", err) + if brokered { + if err := broker.Commit(ctx, transactionID); err != nil { + return fmt.Errorf("commit PageBroker checkpoint: %w", err) + } + committed = true + } else { + // Remove any previous checkpoint with the same identity hash, then + // promote the staged checkpoint directory into place. + if err := os.RemoveAll(finalDir); err != nil { + return fmt.Errorf("failed to remove previous checkpoint directory: %w", err) + } + if err := os.Rename(tmpDir, finalDir); err != nil { + return fmt.Errorf("failed to finalize checkpoint directory: %w", err) + } } phaseTimings.FinalizeDuration = time.Since(finalizeStart) diff --git a/deploy/snapshot/internal/pagebroker/client.go b/deploy/snapshot/internal/pagebroker/client.go new file mode 100644 index 000000000000..20b0a9d203b4 --- /dev/null +++ b/deploy/snapshot/internal/pagebroker/client.go @@ -0,0 +1,181 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package pagebroker + +import ( + "context" + "encoding/binary" + "errors" + "fmt" + "io" + "net" + "time" + + "github.com/google/uuid" + "google.golang.org/protobuf/proto" +) + +const ( + // PageBroker control requests and responses are limited to 64 KiB. + maxMessageSize = 64 << 10 + commitRetryDelay = 100 * time.Millisecond +) + +var errMessageTooLarge = fmt.Errorf("message exceeds %d bytes", maxMessageSize) + +// Client uses the deployment-wide filesystem/POSIX PageBroker plan. +type Client struct { + ControlSocketPath string +} + +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()}, + }) + if err != nil { + return "", err + } + if response.GetStagedCheckpointDirectory() == nil { + return "", fmt.Errorf("unexpected PageBroker checkpoint response") + } + return response.GetStagedCheckpointDirectory().GetImageDirectory(), nil +} + +func (c Client) Commit(ctx context.Context, transactionID string) error { + for { + response, err := c.request(ctx, transactionID, &Request_Commit{Commit: &CommitRequest{}}) + if err == nil { + if response.GetCommitComplete() != nil { + return nil + } + return fmt.Errorf("unexpected PageBroker commit response") + } + var transport transportError + if !errors.As(err, &transport) { + return err + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(commitRetryDelay): + } + } +} + +func (c Client) Abort(ctx context.Context, transactionID string) error { + response, err := c.request(ctx, transactionID, &Request_Abort{Abort: &AbortRequest{}}) + if err != nil { + return err + } + if response.GetAbortComplete() == nil { + return fmt.Errorf("unexpected PageBroker abort response") + } + return nil +} + +func (c Client) request(ctx context.Context, transactionID string, command isRequest_Command) (*Response, error) { + connection, err := (&net.Dialer{}).DialContext(ctx, "unix", c.ControlSocketPath) + if err != nil { + return nil, transportError{cause: fmt.Errorf("dial PageBroker: %w", err)} + } + defer connection.Close() + stopCancel := context.AfterFunc(ctx, func() { _ = connection.Close() }) + defer stopCancel() + + requestID := uuid.NewString() + request := &Request{RequestId: &requestID, TransactionId: &transactionID, Command: command} + message, err := proto.Marshal(request) + if err != nil { + return nil, fmt.Errorf("marshal PageBroker request: %w", err) + } + if err := writeMessage(connection, message); err != nil { + return nil, transportError{cause: fmt.Errorf("write PageBroker request: %w", err)} + } + message, err = readMessage(connection) + if err != nil { + if errors.Is(err, errMessageTooLarge) { + return nil, err + } + return nil, transportError{cause: fmt.Errorf("read PageBroker response: %w", err)} + } + response := new(Response) + if err := proto.Unmarshal(message, response); err != nil { + return nil, fmt.Errorf("unmarshal PageBroker response: %w", err) + } + if response.GetRequestId() != requestID || response.GetTransactionId() != transactionID { + return nil, fmt.Errorf("PageBroker response identifiers do not match request") + } + if failure := response.GetFailure(); failure != nil { + return nil, failureError{code: failureCode(failure.GetCode()), message: failure.GetMessage()} + } + return response, nil +} + +type failureError struct { + code Failure_Code + message string +} + +func failureCode(code Failure_Code) Failure_Code { + switch code { + case Failure_UNSPECIFIED, Failure_INVALID_REQUEST, Failure_TRANSACTION_NOT_FOUND, Failure_TRANSACTION_CONFLICT, + Failure_INSUFFICIENT_STORAGE, Failure_STORAGE_ERROR, Failure_INTERNAL_ERROR: + return code + default: + return Failure_UNSPECIFIED + } +} + +type transportError struct { + cause error +} + +func (e transportError) Error() string { return e.cause.Error() } + +func (e transportError) Unwrap() error { return e.cause } + +func (e failureError) Error() string { + return fmt.Sprintf("PageBroker %s: %s", e.code, e.message) +} + +func filesystem(directory string) *StorageBackend { + return &StorageBackend{Kind: &StorageBackend_Filesystem{Filesystem: &FilesystemStorage{Directory: &directory}}} +} + +func posixCopy() *IOEngine { + return &IOEngine{Kind: &IOEngine_PosixCopy{PosixCopy: &PosixCopyIOEngine{}}} +} + +func writeMessage(writer io.Writer, message []byte) error { + if len(message) > maxMessageSize { + return fmt.Errorf("message exceeds %d bytes", maxMessageSize) + } + if err := binary.Write(writer, binary.BigEndian, uint32(len(message))); err != nil { + return err + } + for len(message) > 0 { + written, err := writer.Write(message) + if err != nil { + return err + } + if written == 0 { + return io.ErrShortWrite + } + message = message[written:] + } + return nil +} + +func readMessage(reader io.Reader) ([]byte, error) { + var size uint32 + if err := binary.Read(reader, binary.BigEndian, &size); err != nil { + return nil, err + } + if size > maxMessageSize { + return nil, errMessageTooLarge + } + message := make([]byte, size) + _, err := io.ReadFull(reader, message) + return message, err +} diff --git a/deploy/snapshot/internal/pagebroker/client_test.go b/deploy/snapshot/internal/pagebroker/client_test.go new file mode 100644 index 000000000000..394f35b30d61 --- /dev/null +++ b/deploy/snapshot/internal/pagebroker/client_test.go @@ -0,0 +1,201 @@ +package pagebroker + +import ( + "context" + "encoding/binary" + "errors" + "net" + "path/filepath" + "testing" + "time" + + "google.golang.org/protobuf/proto" +) + +func TestFailureCodeMapsUnknownValuesToUnspecified(t *testing.T) { + if got := failureCode(Failure_Code(99)); got != Failure_UNSPECIFIED { + t.Fatalf("failureCode(99) = %v, want %v", got, Failure_UNSPECIFIED) + } +} + +func TestRequestStopsWhenContextIsCanceled(t *testing.T) { + listener, err := net.Listen("unix", filepath.Join(t.TempDir(), "pagebroker.sock")) + if err != nil { + t.Fatal(err) + } + defer listener.Close() + + accepted := make(chan net.Conn, 1) + go func() { + connection, err := listener.Accept() + if err == nil { + accepted <- connection + } + }() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + result := make(chan error, 1) + go func() { + result <- (Client{ControlSocketPath: listener.Addr().String()}).Abort(ctx, "transaction") + }() + + connection := <-accepted + defer connection.Close() + if _, err := readMessage(connection); err != nil { + t.Fatal(err) + } + + cancel() + select { + case err := <-result: + if err == nil { + t.Fatal("request succeeded after its context was canceled") + } + case <-time.After(time.Second): + t.Fatal("request did not stop after its context was canceled") + } +} + +func TestCommitRetriesLostResponses(t *testing.T) { + listener, err := net.Listen("unix", filepath.Join(t.TempDir(), "pagebroker.sock")) + if err != nil { + t.Fatal(err) + } + defer listener.Close() + + requests := make(chan *Request, 3) + server := make(chan error, 1) + go func() { + for attempt := 0; attempt < 3; attempt++ { + connection, err := listener.Accept() + if err != nil { + server <- err + return + } + message, err := readMessage(connection) + if err != nil { + _ = connection.Close() + server <- err + return + } + request := new(Request) + if err := proto.Unmarshal(message, request); err != nil { + _ = connection.Close() + server <- err + return + } + requests <- request + if attempt == 2 { + response := &Response{ + RequestId: request.RequestId, + TransactionId: request.TransactionId, + Result: &Response_CommitComplete{CommitComplete: &CommitComplete{}}, + } + message, err = proto.Marshal(response) + if err == nil { + err = writeMessage(connection, message) + } + if err != nil { + _ = connection.Close() + server <- err + return + } + } + _ = connection.Close() + } + server <- nil + }() + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := (Client{ControlSocketPath: listener.Addr().String()}).Commit(ctx, "transaction"); err != nil { + t.Fatal(err) + } + if err := <-server; err != nil { + t.Fatal(err) + } + for range 3 { + request := <-requests + if request.GetTransactionId() != "transaction" || request.GetCommit() == nil { + t.Fatalf("unexpected retry request: %v", request) + } + } +} + +func TestAbortRequiresAbortComplete(t *testing.T) { + listener, err := net.Listen("unix", filepath.Join(t.TempDir(), "pagebroker.sock")) + if err != nil { + t.Fatal(err) + } + defer listener.Close() + + server := make(chan error, 1) + go func() { + connection, err := listener.Accept() + if err != nil { + server <- err + return + } + defer connection.Close() + message, err := readMessage(connection) + if err != nil { + server <- err + return + } + request := new(Request) + if err := proto.Unmarshal(message, request); err != nil { + server <- err + return + } + message, err = proto.Marshal(&Response{ + RequestId: request.RequestId, + TransactionId: request.TransactionId, + Result: &Response_CommitComplete{CommitComplete: &CommitComplete{}}, + }) + if err == nil { + err = writeMessage(connection, message) + } + server <- err + }() + + if err := (Client{ControlSocketPath: listener.Addr().String()}).Abort(context.Background(), "transaction"); err == nil { + t.Fatal("abort accepted a commit response") + } + if err := <-server; err != nil { + t.Fatal(err) + } +} + +func TestCommitDoesNotRetryInvalidFrame(t *testing.T) { + listener, err := net.Listen("unix", filepath.Join(t.TempDir(), "pagebroker.sock")) + if err != nil { + t.Fatal(err) + } + defer listener.Close() + + server := make(chan error, 1) + go func() { + connection, err := listener.Accept() + if err != nil { + server <- err + return + } + defer connection.Close() + if _, err := readMessage(connection); err != nil { + server <- err + return + } + server <- binary.Write(connection, binary.BigEndian, uint32(maxMessageSize+1)) + }() + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + err = (Client{ControlSocketPath: listener.Addr().String()}).Commit(ctx, "transaction") + if !errors.Is(err, errMessageTooLarge) { + t.Fatalf("Commit() error = %v, want invalid frame error", err) + } + if err := <-server; err != nil { + t.Fatal(err) + } +} diff --git a/deploy/snapshot/internal/types/config.go b/deploy/snapshot/internal/types/config.go index 26e4797f213e..f76dccbf80c6 100644 --- a/deploy/snapshot/internal/types/config.go +++ b/deploy/snapshot/internal/types/config.go @@ -15,6 +15,7 @@ type AgentConfig struct { RestrictedNamespace string `yaml:"-"` Storage StorageSpec `yaml:"storage"` Overlay OverlaySettings `yaml:"overlay"` + PageBroker PageBrokerSpec `yaml:"pageBroker"` Restore RestoreSpec `yaml:"restore"` CRIU CRIUSettings `yaml:"criu"` } @@ -90,6 +91,11 @@ type StorageSpec struct { AccessMode string `yaml:"accessMode"` } +type PageBrokerSpec struct { + Enabled bool `yaml:"enabled"` + ControlSocketPath string `yaml:"controlSocketPath"` +} + // RestoreSpec holds settings for the CRIU restore process. type RestoreSpec struct { RestoreTimeoutSeconds int `yaml:"restoreTimeoutSeconds"` diff --git a/deploy/snapshot/pagebroker/Makefile b/deploy/snapshot/pagebroker/Makefile index da90583081f4..ba4b42f55e28 100644 --- a/deploy/snapshot/pagebroker/Makefile +++ b/deploy/snapshot/pagebroker/Makefile @@ -1,13 +1,12 @@ 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 .PHONY: generate test generate: protoc --proto_path=. --cpp_out=. $(PROTO) -test: generate $(BROKER_SOURCES) daemon_test.cpp +test: generate broker.cpp daemon_test.cpp mkdir -p build - c++ -I. -std=c++20 -Wall -Werror $(BROKER_SOURCES) daemon_test.cpp v1/pagebroker.pb.cc -lprotobuf $(GTEST_FLAGS) -o build/pagebroker-test + c++ -I. -std=c++20 -Wall -Werror broker.cpp daemon_test.cpp v1/pagebroker.pb.cc -lprotobuf $(GTEST_FLAGS) -o build/pagebroker-test ./build/pagebroker-test diff --git a/deploy/snapshot/protocol/common.go b/deploy/snapshot/protocol/common.go index f869d9dd5f0f..995f73628fd9 100644 --- a/deploy/snapshot/protocol/common.go +++ b/deploy/snapshot/protocol/common.go @@ -24,6 +24,8 @@ const ( RestoreTargetLabel = "nvidia.com/snapshot-is-restore-target" CheckpointArtifactVersionAnnotation = "nvidia.com/snapshot-artifact-version" + PageBrokerAnnotation = "nvidia.com/pagebroker" + PageBrokerAnnotationEnabled = "enabled" // SnapshotNodeLabel mirrors PodSnapshotContent.spec.source.nodeName onto the // object so the per-node agent's cache can label-select work for its node.