Skip to content

perf(snapshot): remove duplicate CustomStorage restore prefetch #121

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hhzhang16 wants to merge 1 commit into from
Open

perf(snapshot): remove duplicate CustomStorage restore prefetch #121

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 0 additions & 129 deletions agent/internal/cuda/prefetch.go
View file
Open in desktop

This file was deleted.

125 changes: 0 additions & 125 deletions agent/internal/cuda/prefetch_test.go
View file
Open in desktop

This file was deleted.

73 changes: 0 additions & 73 deletions agent/internal/executor/restore.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
}
Expand Down
41 changes: 0 additions & 41 deletions agent/internal/executor/restore_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down