Determine driver version from folders

Signed-off-by: Evan Lezar <[email protected]>

Author: elezar

cmd/nvidia-cdi-hook/compat-libs/compat-libs.go

@@ -33,6 +33,7 @@ type command struct {
 }
 
 type options struct {
+	folders       cli.StringSlice
 	driverVersion string
 	containerSpec string
 }
@@ -66,6 +67,11 @@ func (m command) build() *cli.Command {
 			Usage:       "Specify the host driver version",
 			Destination: &cfg.driverVersion,
 		},
+		&cli.StringSliceFlag{
+			Name:        "folder",
+			Usage:       "Specify the folders that are added to the container. These are used to locate libcuda.so.RM_VERSION if not specified.",
+			Destination: &cfg.folders,
+		},
 		&cli.StringFlag{
 			Name:        "container-spec",
 			Usage:       "Specify the path to the OCI container spec. If empty or '-' the spec will be read from STDIN",
@@ -120,7 +126,22 @@ func (m command) run(c *cli.Context, cfg *options) error {
 
 	compatVersion := strings.TrimPrefix(filepath.Base(libs[0]), "libcuda.so.")
 	compatMajor := strings.SplitN(compatVersion, ".", 2)[0]
-	driverMajor := strings.SplitN(cfg.driverVersion, ".", 2)[0]
+
+	driverVersion := cfg.driverVersion
+	if driverVersion == "" {
+		for _, folder := range cfg.folders.Value() {
+			libs, err := root(containerRoot).glob(filepath.Join(folder, "libcuda.so.*.*"))
+			if err != nil || len(libs) == 0 {
+				continue
+			}
+			if len(libs) != 1 {
+				m.logger.Warningf("Unexpected number of CUDA compat libraries: %v", libs)
+			}
+
+			driverVersion = strings.TrimPrefix(filepath.Base(libs[0]), "libcuda.so.")
+		}
+	}
+	driverMajor := strings.SplitN(driverVersion, ".", 2)[0]
 
 	if driverMajor > compatMajor {
 		return nil