From abe38d82ac3634264608259d1c60003cdd53658a Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Wed, 16 Oct 2024 16:45:22 -0700
Subject: [PATCH] Docker GPU tests: Use the sniffer on `exec`'d commands too.

Without this, tests that `exec` commands into existing GPU containers (e.g.
the CUDA samples test) would not actually get ioctl compatibility
enforcement.

Updates issue #10885.

PiperOrigin-RevId: 686686310
---
 pkg/test/dockerutil/container.go | 22 +++++++++++++++++++---
 pkg/test/dockerutil/exec.go      |  8 ++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/pkg/test/dockerutil/container.go b/pkg/test/dockerutil/container.go
index a2f4e9d57..3b20e7524 100644
--- a/pkg/test/dockerutil/container.go
+++ b/pkg/test/dockerutil/container.go
@@ -58,6 +58,9 @@ type Container struct {
 
 	// profile is the profiling hook associated with this container.
 	profile *profile
+
+	// sniffGPUOpts, if set, sets the rules for GPU sniffing for this container.
+	sniffGPUOpts *SniffGPUOpts
 }
 
 // RunOpts are options for running a container.
@@ -207,6 +210,7 @@ func (c *Container) SpawnProcess(ctx context.Context, r RunOpts, args ...string)
 	}
 
 	c.cleanups = append(c.cleanups, func() { stream.Close() })
+	c.sniffGPUOpts = r.sniffGPUOpts
 
 	if err := c.Start(ctx); err != nil {
 		return Process{}, err
@@ -310,15 +314,19 @@ func (c *Container) config(ctx context.Context, r RunOpts, args []string) (*cont
 		c.cleanups = append(c.cleanups, func() {
 			r.sniffGPUOpts.cleanup()
 		})
-		if len(entrypoint) == 0 && len(args) == 0 {
+		if len(entrypoint) == 0 || len(args) == 0 {
 			// Need to look up the image's default entrypoint/args so we can prepend to them.
 			// If we don't, then we will end up overwriting them.
 			imageInfo, _, err := c.client.ImageInspectWithRaw(ctx, image)
 			if err != nil {
 				return nil, fmt.Errorf("cannot inspect image %q: %w", image, err)
 			}
-			entrypoint = []string(imageInfo.Config.Entrypoint)
-			args = []string(imageInfo.Config.Cmd)
+			if len(entrypoint) == 0 {
+				entrypoint = []string(imageInfo.Config.Entrypoint)
+			}
+			if len(args) == 0 {
+				args = []string(imageInfo.Config.Cmd)
+			}
 		}
 		if len(entrypoint) != 0 {
 			entrypoint = r.sniffGPUOpts.prepend(entrypoint)
@@ -413,6 +421,14 @@ func (c *Container) Logs(ctx context.Context) (string, error) {
 	return out.String(), err
 }
 
+// OutputStreams gets the container's stdout and stderr streams separately.
+func (c *Container) OutputStreams(ctx context.Context) (string, string, error) {
+	var stdout bytes.Buffer
+	var stderr bytes.Buffer
+	err := c.logs(ctx, &stdout, &stderr)
+	return stdout.String(), stderr.String(), err
+}
+
 func (c *Container) logs(ctx context.Context, stdout, stderr *bytes.Buffer) error {
 	opts := types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true}
 	writer, err := c.client.ContainerLogs(ctx, c.id, opts)
diff --git a/pkg/test/dockerutil/exec.go b/pkg/test/dockerutil/exec.go
index 1ae4fd990..f4ee4cd6c 100644
--- a/pkg/test/dockerutil/exec.go
+++ b/pkg/test/dockerutil/exec.go
@@ -40,6 +40,11 @@ type ExecOpts struct {
 
 	// WorkDir is the working directory of the process.
 	WorkDir string
+
+	// NoWrap, if true, indicates that no command-line wrapping may be performed
+	// on the command line being exec'd. Otherwise, settings are inherited from
+	// the container.
+	NoWrap bool
 }
 
 // ExecError is returned when a process terminated with a non-zero exit status.
@@ -111,6 +116,9 @@ func (c *Container) doExec(ctx context.Context, r ExecOpts, args []string) (Proc
 
 func (c *Container) execConfig(r ExecOpts, cmd []string) types.ExecConfig {
 	env := append(r.Env, fmt.Sprintf("RUNSC_TEST_NAME=%s", c.Name))
+	if !r.NoWrap && c.sniffGPUOpts != nil {
+		cmd = c.sniffGPUOpts.prepend(cmd)
+	}
 	return types.ExecConfig{
 		AttachStdin:  r.UseTTY,
 		AttachStderr: true,