diff --git a/agent/internal/cuda/prefetch.go b/agent/internal/cuda/prefetch.go deleted file mode 100644 index e00eadf4..00000000 --- a/agent/internal/cuda/prefetch.go +++ /dev/null @@ -1,129 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -package cuda - -import ( - "context" - "errors" - "fmt" - "io/fs" - "os" - "path/filepath" - "regexp" - "sort" - "time" - - "golang.org/x/sys/unix" -) - -const customStoragePrefetchBufferBytes = 8 * 1024 * 1024 - -var customStorageExtentFilePattern = regexp.MustCompile(`^device-[0-9]{4}\.bin(?:\.part-[0-9]{4})?$`) - -// CustomStoragePrefetchResult describes a completed best-effort page-cache -// preload. Duration is service time and may overlap CRIU restore. -type CustomStoragePrefetchResult struct { - Files int - Bytes int64 - Duration time.Duration -} - -// PrefetchCustomStorageArtifacts validates and reads every CUDA CustomStorage -// extent into the node page cache. Snapshot starts this before CRIU restore so -// durable storage I/O overlaps process restore; the CUDA helper still performs -// the authoritative read into registered host buffers after target PIDs exist. -func PrefetchCustomStorageArtifacts(ctx context.Context, checkpointDir string) (CustomStoragePrefetchResult, error) { - start := time.Now() - root := filepath.Join(checkpointDir, "cuda-custom-storage") - rootInfo, err := os.Lstat(root) - if err != nil { - return CustomStoragePrefetchResult{}, fmt.Errorf("inspect CUDA CustomStorage artifact directory: %w", err) - } - if !rootInfo.IsDir() { - return CustomStoragePrefetchResult{}, fmt.Errorf("CUDA CustomStorage artifact path is not a directory") - } - - var paths []string - err = filepath.WalkDir(root, func(path string, entry fs.DirEntry, walkErr error) error { - if walkErr != nil { - return walkErr - } - if entry.Type()&os.ModeSymlink != 0 { - if customStorageExtentFilePattern.MatchString(entry.Name()) { - return fmt.Errorf("CUDA CustomStorage extent %s is a symlink", path) - } - return nil - } - if entry.IsDir() || !customStorageExtentFilePattern.MatchString(entry.Name()) { - return nil - } - if !entry.Type().IsRegular() { - return fmt.Errorf("CUDA CustomStorage extent %s is not a regular file", path) - } - paths = append(paths, path) - return nil - }) - if err != nil { - return CustomStoragePrefetchResult{}, fmt.Errorf("discover CUDA CustomStorage extents: %w", err) - } - if len(paths) == 0 { - return CustomStoragePrefetchResult{}, fmt.Errorf("CUDA CustomStorage artifact contains no extent files") - } - sort.Strings(paths) - - buffer := make([]byte, customStoragePrefetchBufferBytes) - result := CustomStoragePrefetchResult{Files: len(paths)} - for _, path := range paths { - bytesRead, err := prefetchCustomStorageFile(ctx, path, buffer) - if err != nil { - return CustomStoragePrefetchResult{}, err - } - result.Bytes += bytesRead - } - result.Duration = time.Since(start) - return result, nil -} - -func prefetchCustomStorageFile(ctx context.Context, path string, buffer []byte) (int64, error) { - fd, err := unix.Open(path, unix.O_RDONLY|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0) - if err != nil { - return 0, fmt.Errorf("open CUDA CustomStorage extent %s: %w", path, err) - } - defer unix.Close(fd) - - var stat unix.Stat_t - if err := unix.Fstat(fd, &stat); err != nil { - return 0, fmt.Errorf("stat CUDA CustomStorage extent %s: %w", path, err) - } - if stat.Mode&unix.S_IFMT != unix.S_IFREG || stat.Size <= 0 { - return 0, fmt.Errorf("CUDA CustomStorage extent %s is not a nonempty regular file", path) - } - _ = unix.Fadvise(fd, 0, stat.Size, unix.FADV_WILLNEED) - - // FADV_WILLNEED is only an asynchronous hint and may return before the - // extent reaches the page cache. Reading the complete file makes this - // best-effort prefetch observable and ensures the later CUDA restore can - // consume cached pages when the filesystem honors normal buffered I/O. - var total int64 - for { - if err := ctx.Err(); err != nil { - return 0, fmt.Errorf("prefetch CUDA CustomStorage extent %s: %w", path, err) - } - read, err := unix.Read(fd, buffer) - if errors.Is(err, unix.EINTR) { - continue - } - if err != nil { - return 0, fmt.Errorf("read CUDA CustomStorage extent %s: %w", path, err) - } - total += int64(read) - if read == 0 { - break - } - } - if total != stat.Size { - return 0, fmt.Errorf("CUDA CustomStorage extent %s changed size while prefetching: read %d, expected %d", path, total, stat.Size) - } - return total, nil -} diff --git a/agent/internal/cuda/prefetch_test.go b/agent/internal/cuda/prefetch_test.go deleted file mode 100644 index 2448face..00000000 --- a/agent/internal/cuda/prefetch_test.go +++ /dev/null @@ -1,125 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -package cuda - -import ( - "context" - "os" - "path/filepath" - "strings" - "testing" -) - -func TestPrefetchCustomStorageArtifacts(t *testing.T) { - checkpointDir := t.TempDir() - processDir := filepath.Join(checkpointDir, "cuda-custom-storage", "process-nspid-42") - if err := os.MkdirAll(processDir, 0o700); err != nil { - t.Fatal(err) - } - files := map[string][]byte{ - "device-0000.bin.part-0000": []byte("first"), - "device-0000.bin.part-0001": []byte("second"), - "manifest.txt": []byte("ignored"), - } - for name, contents := range files { - if err := os.WriteFile(filepath.Join(processDir, name), contents, 0o600); err != nil { - t.Fatal(err) - } - } - - result, err := PrefetchCustomStorageArtifacts(context.Background(), checkpointDir) - if err != nil { - t.Fatal(err) - } - if result.Files != 2 || result.Bytes != int64(len("first")+len("second")) { - t.Fatalf("unexpected prefetch result: %+v", result) - } -} - -func TestPrefetchCustomStorageArtifactsRejectsExtentSymlink(t *testing.T) { - checkpointDir := t.TempDir() - processDir := filepath.Join(checkpointDir, "cuda-custom-storage", "process-nspid-42") - if err := os.MkdirAll(processDir, 0o700); err != nil { - t.Fatal(err) - } - target := filepath.Join(checkpointDir, "target") - if err := os.WriteFile(target, []byte("data"), 0o600); err != nil { - t.Fatal(err) - } - if err := os.Symlink(target, filepath.Join(processDir, "device-0000.bin")); err != nil { - t.Fatal(err) - } - - _, err := PrefetchCustomStorageArtifacts(context.Background(), checkpointDir) - if err == nil || !strings.Contains(err.Error(), "symlink") { - t.Fatalf("expected symlink rejection, got %v", err) - } -} - -func TestPrefetchCustomStorageArtifactsHonorsCancellation(t *testing.T) { - checkpointDir := t.TempDir() - processDir := filepath.Join(checkpointDir, "cuda-custom-storage", "process-nspid-42") - if err := os.MkdirAll(processDir, 0o700); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(processDir, "device-0000.bin"), []byte("data"), 0o600); err != nil { - t.Fatal(err) - } - ctx, cancel := context.WithCancel(context.Background()) - cancel() - - _, err := PrefetchCustomStorageArtifacts(ctx, checkpointDir) - if err == nil || !strings.Contains(err.Error(), "context canceled") { - t.Fatalf("expected cancellation, got %v", err) - } -} - -func TestPrefetchCustomStorageArtifactsRejectsIncompleteArtifacts(t *testing.T) { - tests := []struct { - name string - prepare func(t *testing.T, checkpointDir string) - wantError string - }{ - { - name: "missing artifact directory", - prepare: func(*testing.T, string) {}, - wantError: "inspect CUDA CustomStorage artifact directory", - }, - { - name: "no extent files", - prepare: func(t *testing.T, checkpointDir string) { - t.Helper() - if err := os.MkdirAll(filepath.Join(checkpointDir, "cuda-custom-storage"), 0o700); err != nil { - t.Fatal(err) - } - }, - wantError: "contains no extent files", - }, - { - name: "empty extent", - prepare: func(t *testing.T, checkpointDir string) { - t.Helper() - processDir := filepath.Join(checkpointDir, "cuda-custom-storage", "process-nspid-42") - if err := os.MkdirAll(processDir, 0o700); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(processDir, "device-0000.bin"), nil, 0o600); err != nil { - t.Fatal(err) - } - }, - wantError: "not a nonempty regular file", - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - checkpointDir := t.TempDir() - test.prepare(t, checkpointDir) - _, err := PrefetchCustomStorageArtifacts(context.Background(), checkpointDir) - if err == nil || !strings.Contains(err.Error(), test.wantError) { - t.Fatalf("expected error containing %q, got %v", test.wantError, err) - } - }) - } -} diff --git a/agent/internal/executor/restore.go b/agent/internal/executor/restore.go index f9b78f95..ef292fb8 100644 --- a/agent/internal/executor/restore.go +++ b/agent/internal/executor/restore.go @@ -51,34 +51,6 @@ type restoreMount struct { point nsmount.MountPoint } -type customStoragePrefetchOutcome struct { - result cuda.CustomStoragePrefetchResult - err error -} - -func waitForCustomStoragePrefetch( - ctx context.Context, - outcomes <-chan customStoragePrefetchOutcome, - cancel context.CancelFunc, - discard bool, -) (cuda.CustomStoragePrefetchResult, error) { - if outcomes == nil { - return cuda.CustomStoragePrefetchResult{}, nil - } - if discard { - cancel() - return cuda.CustomStoragePrefetchResult{}, nil - } - select { - case outcome := <-outcomes: - cancel() - return outcome.result, outcome.err - case <-ctx.Done(): - cancel() - return cuda.CustomStoragePrefetchResult{}, fmt.Errorf("wait for CUDA CustomStorage artifact prefetch: %w", ctx.Err()) - } -} - func cleanupRestoreMounts(ctx context.Context, mounts []restoreMount) error { var cleanupErr error cleanupCtx := context.WithoutCancel(ctx) @@ -184,59 +156,14 @@ func Restore(ctx context.Context, rt snapshotruntime.Runtime, log logr.Logger, r point: artifactMount, }) - var prefetch <-chan customStoragePrefetchOutcome - var cancelPrefetch context.CancelFunc - if snap.CUDAStorageMode == types.CUDAStorageModePOSIX { - prefetchCtx, cancel := context.WithCancel(ctx) - cancelPrefetch = cancel - outcomes := make(chan customStoragePrefetchOutcome, 1) - prefetch = outcomes - go func() { - result, err := cuda.PrefetchCustomStorageArtifacts(prefetchCtx, artifactPath) - outcomes <- customStoragePrefetchOutcome{result: result, err: err} - }() - } - awaitPrefetch := func(cancel bool) (cuda.CustomStoragePrefetchResult, error) { - if prefetch == nil { - return cuda.CustomStoragePrefetchResult{}, nil - } - result, err := waitForCustomStoragePrefetch(ctx, prefetch, cancelPrefetch, cancel) - prefetch = nil - cancelPrefetch = nil - return result, err - } - defer func() { - if prefetch != nil { - _, _ = awaitPrefetch(true) - } - }() - // NodeController.failRestore owns placeholder-wide termination for every // non-cleanup error returned after execution begins. Keeping cleanup in the // controller guarantees that RestoreFailed is not persisted until the // runtime-owned placeholder has actually been resolved and terminated. result, err := execNSRestore(ctx, log, req, snap, bundleMount, nsmount.CheckpointDst) if err != nil { - _, _ = awaitPrefetch(true) return 0, fmt.Errorf("nsrestore failed: %w", err) } - prefetchResult, err := awaitPrefetch(false) - if err != nil { - if ctx.Err() != nil { - return 0, fmt.Errorf("CUDA CustomStorage artifact prefetch interrupted after CRIU restore: %w", err) - } - // Prefetch only overlaps durable-storage reads with CRIU. The CUDA - // helper performs the authoritative read and validation, so an - // optimization failure must not strand a process after CRIU restore. - log.Error(err, "CUDA CustomStorage artifact prefetch failed; continuing with authoritative restore") - } else if prefetchResult.Files > 0 { - log.Info("CUDA CustomStorage artifact prefetch completed", - "files", prefetchResult.Files, - "bytes", prefetchResult.Bytes, - "service_duration", prefetchResult.Duration, - "overlapped_with_criu", true, - ) - } if result.CleanupError != nil { cleanupErr = errors.Join(cleanupErr, result.CleanupError) } diff --git a/agent/internal/executor/restore_test.go b/agent/internal/executor/restore_test.go index fe3105cc..3aeefdf1 100644 --- a/agent/internal/executor/restore_test.go +++ b/agent/internal/executor/restore_test.go @@ -177,47 +177,6 @@ func TestRemainingDuration(t *testing.T) { } } -func TestWaitForCustomStoragePrefetchDiscardDoesNotWait(t *testing.T) { - outcomes := make(chan customStoragePrefetchOutcome, 1) - ctx, cancel := context.WithCancel(context.Background()) - - done := make(chan struct{}) - go func() { - defer close(done) - if _, err := waitForCustomStoragePrefetch(ctx, outcomes, cancel, true); err != nil { - t.Errorf("waitForCustomStoragePrefetch(discard=true): %v", err) - } - }() - - select { - case <-done: - case <-time.After(time.Second): - t.Fatal("discarded prefetch waited for an outcome") - } - select { - case <-ctx.Done(): - case <-time.After(time.Second): - t.Fatal("discarded prefetch did not cancel its context") - } -} - -func TestWaitForCustomStoragePrefetchHonorsContextCancellation(t *testing.T) { - outcomes := make(chan customStoragePrefetchOutcome, 1) - ctx, cancelContext := context.WithCancel(context.Background()) - cancelContext() - prefetchCtx, cancelPrefetch := context.WithCancel(context.Background()) - - _, err := waitForCustomStoragePrefetch(ctx, outcomes, cancelPrefetch, false) - if !errors.Is(err, context.Canceled) { - t.Fatalf("waitForCustomStoragePrefetch() error = %v, want context.Canceled", err) - } - select { - case <-prefetchCtx.Done(): - case <-time.After(time.Second): - t.Fatal("canceled wait did not cancel prefetch") - } -} - func TestRestoreDeferredCUDAProcessesResolvesAndValidatesHostIdentity(t *testing.T) { namespaceProcess := snapshotruntime.ProcessDetails{ InnermostPID: 7,