cos-extensions: skip postInstall when nvidia-gridd and persistenced unit exist

BUG=b/573970841
TEST=unit tests; use custom gpu-installer and cos-extensions to install gpu driver on a fractional VM, verify license and CUDA.
RELEASE_NOTE=None

#### IMPORTANT NOTE ####
Only merge this after merging go/cos-rev/227866, making a new cos-gpu-installer release and consuming that version on cos-extensions ebuild.

cos-patch: bug
Change-Id: I62b56ac4c23eb8f83821dd68fa1bb2ee45b0fc25
Reviewed-on: https://cos-review.googlesource.com/c/cos/cos-extensions/+/228667
Reviewed-by: Robert Kolchmeyer <rkolchmeyer@google.com>
Tested-by: Cusky Presubmit Bot <presubmit@cos-infra-prod.iam.gserviceaccount.com>

extensions/gpu/gpu.go

diff --git a/extensions/gpu/gpu.go b/extensions/gpu/gpu.go
index be4e07e..a78cbe2 100644
--- a/extensions/gpu/gpu.go
+++ b/extensions/gpu/gpu.go
@@ -32,11 +32,12 @@
 )
 
 var (
-	run                 = osutils.Run
-	parse               = osutils.Parse
-	artifactRegions     = []string{"us", "eu", "asia"}
-	apiDomainsPath      = "/etc/cloud-api-domains"
-	tpcGPUInstallerRepo = "docker.%s/%s/cos-cloud/cos-gpu-installer"
+	run                     = osutils.Run
+	parse                   = osutils.Parse
+	artifactRegions         = []string{"us", "eu", "asia"}
+	apiDomainsPath          = "/etc/cloud-api-domains"
+	tpcGPUInstallerRepo     = "docker.%s/%s/cos-cloud/cos-gpu-installer"
+	persistencedServicePath = "/etc/systemd/system/nvidia-persistenced.service"
 	// Note: Please update the app-admin/extensions-manager ebuild file.
 	DefaultInstaller = "gcr.io/cos-cloud/cos-gpu-installer:v2.4.4"
 )
@@ -125,6 +126,13 @@
 }
 
 func postInstall(hostDir string) {
+	if _, err := os.Stat(filepath.Join(hostDir, "bin/nvidia-gridd")); err == nil {
+		if _, err := os.Stat(persistencedServicePath); err == nil {
+			glog.Info("NVIDIA Grid driver detected; daemon setup is handled by cos-gpu-installer.")
+			return
+		}
+	}
+
 	glog.Info("Making the GPU driver installation path executable by re-mounting it.")
 	cmd := fmt.Sprintf(
 		`sudo -n -- sh -c "mount --bind %[1]s %[1]s; mount -o remount,exec %[1]s"`,

extensions/gpu/gpu_test.go

diff --git a/extensions/gpu/gpu_test.go b/extensions/gpu/gpu_test.go
index 1468e25..f7b6dd7 100644
--- a/extensions/gpu/gpu_test.go
+++ b/extensions/gpu/gpu_test.go
@@ -6,6 +6,7 @@
 	"io"
 	"net/http"
 	"os"
+	"path/filepath"
 	"strings"
 	"testing"
 
@@ -470,18 +471,76 @@
 func TestPostInstall(t *testing.T) {
 	f := FakeOSUtils{result: "x86_64"}
 	run = f.run
-	hostDir := "/tmp/nvidia"
+	hostDir := t.TempDir()
 	postInstall(hostDir)
 
 	wantCmds := []string{
 		"/bin/sh", "-c",
-		"sudo -n -- sh -c \"mount --bind /tmp/nvidia /tmp/nvidia; " +
-			"mount -o remount,exec /tmp/nvidia\"",
+		fmt.Sprintf("sudo -n -- sh -c \"mount --bind %[1]s %[1]s; "+
+			"mount -o remount,exec %[1]s\"", hostDir),
 		"/bin/sh", "-c",
-		"sudo /tmp/nvidia/bin/nvidia-persistenced",
+		fmt.Sprintf("sudo %s/bin/nvidia-persistenced", hostDir),
 	}
 
 	if !cmp.Equal(f.cmds, wantCmds) {
 		t.Errorf("TestPostInstall: returned unexpected difference in commands;\n want:%v\t got:%v", wantCmds, f.cmds)
 	}
 }
+
+func TestPostInstallWithGridd(t *testing.T) {
+	f := FakeOSUtils{result: "x86_64"}
+	run = f.run
+	hostDir := t.TempDir()
+	binDir := filepath.Join(hostDir, "bin")
+	if err := os.MkdirAll(binDir, 0755); err != nil {
+		t.Fatalf("failed to create bin dir: %v", err)
+	}
+	if err := os.WriteFile(filepath.Join(binDir, "nvidia-gridd"), []byte("fake"), 0755); err != nil {
+		t.Fatalf("failed to create fake nvidia-gridd: %v", err)
+	}
+
+	fakeService := filepath.Join(hostDir, "nvidia-persistenced.service")
+	if err := os.WriteFile(fakeService, []byte("[Unit]\n"), 0644); err != nil {
+		t.Fatalf("failed to create fake nvidia-persistenced.service: %v", err)
+	}
+	origServicePath := persistencedServicePath
+	persistencedServicePath = fakeService
+	defer func() { persistencedServicePath = origServicePath }()
+
+	postInstall(hostDir)
+
+	if len(f.cmds) != 0 {
+		t.Errorf("TestPostInstallWithGridd: expected no commands to be run when nvidia-gridd and nvidia-persistenced.service are present, got: %v", f.cmds)
+	}
+}
+
+func TestPostInstallWithGriddNoPersistencedService(t *testing.T) {
+	f := FakeOSUtils{result: "x86_64"}
+	run = f.run
+	hostDir := t.TempDir()
+	binDir := filepath.Join(hostDir, "bin")
+	if err := os.MkdirAll(binDir, 0755); err != nil {
+		t.Fatalf("failed to create bin dir: %v", err)
+	}
+	if err := os.WriteFile(filepath.Join(binDir, "nvidia-gridd"), []byte("fake"), 0755); err != nil {
+		t.Fatalf("failed to create fake nvidia-gridd: %v", err)
+	}
+
+	origServicePath := persistencedServicePath
+	persistencedServicePath = filepath.Join(hostDir, "nonexistent-nvidia-persistenced.service")
+	defer func() { persistencedServicePath = origServicePath }()
+
+	postInstall(hostDir)
+
+	wantCmds := []string{
+		"/bin/sh", "-c",
+		fmt.Sprintf("sudo -n -- sh -c \"mount --bind %[1]s %[1]s; "+
+			"mount -o remount,exec %[1]s\"", hostDir),
+		"/bin/sh", "-c",
+		fmt.Sprintf("sudo %s/bin/nvidia-persistenced", hostDir),
+	}
+
+	if !cmp.Equal(f.cmds, wantCmds) {
+		t.Errorf("TestPostInstallWithGriddNoPersistencedService: returned unexpected difference in commands;\n want:%v\t got:%v", wantCmds, f.cmds)
+	}
+}