From 6199fc8395aa5be3a4bfec9ce3f14b34d5449032 Mon Sep 17 00:00:00 2001
From: Anthony Cui <anthonycui@google.com>
Date: Fri, 16 Aug 2024 16:30:38 -0700
Subject: [PATCH] Fix sniffer_test to work.

Previously, run_sniffer was not correctly reporting when unsupported ioctls
were found with the compatibility flag set. At the same time, the sniffer test
was not correctly testing a supported CUDA program, since it was using
run_sample which is currently broken with the sniffer.

This also adds the sniffer test to the list of gpu tests.

PiperOrigin-RevId: 663911778
---
 Makefile                           | 2 ++
 test/gpu/sniffer_test.go           | 8 ++++----
 tools/ioctl_sniffer/run_sniffer.go | 4 ++++
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/Makefile b/Makefile
index 86cd5f896..1f35f7317 100644
--- a/Makefile
+++ b/Makefile
@@ -308,6 +308,7 @@ gpu-all-tests: gpu-images gpu-smoke-tests $(RUNTIME_BIN)
 	@$(call sudo,test/gpu:imagegen_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 	@$(call sudo,test/gpu:sr_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 	@$(call sudo,test/gpu:nccl_test,--runtime=$(RUNTIME) -test.v $(ARGS))
+	@$(call test,--test_env=RUNTIME=$(RUNTIME) test/gpu:sniffer_test)
 .PHONY: gpu-all-tests
 
 cos-gpu-all-tests: gpu-images cos-gpu-smoke-tests $(RUNTIME_BIN)
@@ -317,6 +318,7 @@ cos-gpu-all-tests: gpu-images cos-gpu-smoke-tests $(RUNTIME_BIN)
 	@$(call sudo,test/gpu:imagegen_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 	@$(call sudo,test/gpu:sr_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 	@$(call sudo,test/gpu:nccl_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
+	@$(call test,--test_env=RUNTIME=$(RUNTIME) --test_arg="--cos-gpu" test/gpu:sniffer_test)
 .PHONY: cos-gpu-all-tests
 
 portforward-tests: load-basic_redis load-basic_nginx $(RUNTIME_BIN)
diff --git a/test/gpu/sniffer_test.go b/test/gpu/sniffer_test.go
index b79fa9605..f10debdbc 100644
--- a/test/gpu/sniffer_test.go
+++ b/test/gpu/sniffer_test.go
@@ -32,7 +32,7 @@ const maxDuration = 1 * time.Minute
 // RunCommand runs the given command via the sniffer, with the -enforce_compatibility flag.
 //
 //	It's run in a docker container, with the cuda-tests image.
-func runCommand(t *testing.T, cmd ...string) (string, error) {
+func runCUDATestsCommand(t *testing.T, cmd ...string) (string, error) {
 	// Find the sniffer binary
 	cliPath, err := testutil.FindFile("tools/ioctl_sniffer/run_sniffer")
 	if err != nil {
@@ -62,19 +62,19 @@ func runCommand(t *testing.T, cmd ...string) (string, error) {
 }
 
 func TestSupportedCUDAProgram(t *testing.T) {
-	output, err := runCommand(t, "/run_sample", "0_Introduction/vectorAdd")
+	output, err := runCUDATestsCommand(t, "/run_smoke.sh")
 	t.Logf("%s", output)
 	if err != nil {
 		t.Logf("Error: %v", err)
 		if strings.Contains(output, "unsupported ioctls found") {
 			t.Fatalf("'unsupported ioctls found' found in output")
 		}
-		t.Fatalf("Failed to run vectorAdd")
+		t.Fatalf("Failed to run run_smoke.sh")
 	}
 }
 
 func TestUnsupportedCUDAProgram(t *testing.T) {
-	output, err := runCommand(t, "/unsupported_ioctl")
+	output, err := runCUDATestsCommand(t, "/unsupported_ioctl")
 	t.Logf("%s", output)
 	if err == nil {
 		t.Fatalf("Expected run_sniffer to fail")
diff --git a/tools/ioctl_sniffer/run_sniffer.go b/tools/ioctl_sniffer/run_sniffer.go
index 88ff9ad42..0d5d8dd52 100644
--- a/tools/ioctl_sniffer/run_sniffer.go
+++ b/tools/ioctl_sniffer/run_sniffer.go
@@ -115,6 +115,10 @@ func Main(ctx context.Context) error {
 
 	// Merge results from each connection.
 	finalResults := server.AllResults()
+	if *enforceCompatability && finalResults.HasUnsupportedIoctl() {
+		return fmt.Errorf("unsupported ioctls found: %v", finalResults)
+	}
+
 	log.Infof("============== Unsupported ioctls ==============")
 	log.Infof("%v", finalResults)