gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/test/gpu/sr_test.go (about)

     1  // Copyright 2023 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package sr_test runs checkpoint/restore tests for nvproxy.
    16  package sr_test
    17  
    18  import (
    19  	"testing"
    20  	"time"
    21  
    22  	"gvisor.dev/gvisor/pkg/context"
    23  	"gvisor.dev/gvisor/pkg/test/dockerutil"
    24  	"gvisor.dev/gvisor/pkg/test/testutil"
    25  )
    26  
    27  func TestGPUCheckpointRestore(t *testing.T) {
    28  	if !testutil.IsCheckpointSupported() {
    29  		t.Skip("Checkpoint is not supported.")
    30  	}
    31  	dockerutil.EnsureDockerExperimentalEnabled()
    32  
    33  	ctx := context.Background()
    34  	c := dockerutil.MakeContainer(ctx, t)
    35  	defer c.CleanUp(ctx)
    36  
    37  	opts := dockerutil.GPURunOpts()
    38  	opts.Image = "basic/cuda-vector-add"
    39  	if err := c.Spawn(ctx, opts, "sleep", "infinity"); err != nil {
    40  		t.Fatalf("could not run cuda-vector-add: %v", err)
    41  	}
    42  
    43  	// Run the vector add program.
    44  	vectorAddCmd := []string{"/bin/sh", "-c", "./vectorAdd"}
    45  	if _, err := c.Exec(ctx, dockerutil.ExecOpts{}, vectorAddCmd...); err != nil {
    46  		t.Fatalf("docker exec failed: %v", err)
    47  	}
    48  
    49  	// Create a snapshot.
    50  	if err := c.Checkpoint(ctx, "test"); err != nil {
    51  		t.Fatalf("docker checkpoint failed: %v", err)
    52  	}
    53  	if err := c.WaitTimeout(ctx, time.Minute); err != nil {
    54  		t.Fatalf("wait failed: %v", err)
    55  	}
    56  
    57  	// Restore the snapshot.
    58  	// TODO(b/143498576): Remove Poll after github.com/moby/moby/issues/38963 is fixed.
    59  	if err := testutil.Poll(func() error { return c.Restore(ctx, "test") }, time.Minute); err != nil {
    60  		t.Fatalf("docker restore failed: %v", err)
    61  	}
    62  
    63  	// Run the vector add program again to ensure GPUs are functional.
    64  	if _, err := c.Exec(ctx, dockerutil.ExecOpts{}, vectorAddCmd...); err != nil {
    65  		t.Fatalf("docker exec failed: %v", err)
    66  	}
    67  }