Merge pull request #1222 from Crelex/crelex-nvidia-cuda-container-version-get-fix

Fix for getting CUDA Version inside nvidia/cuda containers

Author: martindevans

LLama/Native/Load/SystemInfo.cs

@@ -41,31 +41,26 @@ public static SystemInfo Get()
 
             return new SystemInfo(platform, GetCudaMajorVersion(), GetVulkanVersion());
         }
-        
         #region Vulkan version
         private static string? GetVulkanVersion()
         {
             // Get Vulkan Summary
             string? vulkanSummary = GetVulkanSummary();
-            
             // If we have a Vulkan summary
             if (vulkanSummary != null)
             {
                 // Extract Vulkan version from summary
                 string? vulkanVersion = ExtractVulkanVersionFromSummary(vulkanSummary);
-                
                 // If we have a Vulkan version
                 if (vulkanVersion != null)
                 {
                     // Return the Vulkan version
                     return vulkanVersion;
                 }
             }
-            
             // Return null if we failed to get the Vulkan version
             return null;
         }
-        
         private static string? GetVulkanSummary()
         {
             // Note: on Linux, this requires `vulkan-tools` to be installed. (`sudo apt install vulkan-tools`)
@@ -85,7 +80,6 @@ public static SystemInfo Get()
                     }
                 };
                 var (exitCode, output, error, ok) = process.SafeRun(TimeSpan.FromSeconds(12));
-
                 // If ok return the output else return null 
                 return ok ? output :
                     null;
@@ -96,38 +90,30 @@ public static SystemInfo Get()
                 return null;
             }
         }
-
         static string? ExtractVulkanVersionFromSummary(string vulkanSummary)
         {
             // We have three ways of parsing the Vulkan version from the summary (output is a different between Windows and Linux)
             // For now, I have decided to go with the full version number, and leave it up to the user to parse it further if needed
             // I have left the other patterns in, in case we need them in the future
-            
             // Output on linux : 4206847 (1.3.255)
             // Output on windows : 1.3.255
             string pattern = @"apiVersion\s*=\s*([^\r\n]+)";
-            
             // Output on linux : 4206847
             // Output on windows : 1.3.255
             //string pattern = @"apiVersion\s*=\s*([\d\.]+)";
-            
             // Output on linux : 1.3.255
             // Output on windows : 1.3.255
             //string pattern = @"apiVersion\s*=\s*(?:\d+\s*)?(?:\(\s*)?([\d]+\.[\d]+\.[\d]+)(?:\s*\))?";
-            
             // Create a Regex object to match the pattern
             Regex regex = new Regex(pattern);
-
             // Match the pattern in the input string
             Match match = regex.Match(vulkanSummary);
-
             // If a match is found
             if (match.Success)
             {
                 // Return the version number
                 return match.Groups[1].Value;
             }
-
             // Return null if no match is found
             return null;
         }
@@ -145,37 +131,37 @@ private static int GetCudaMajorVersion()
                 {
                     return -1;
                 }
-
                 //Ensuring cuda bin path is reachable. Especially for MAUI environment.
                 string cudaBinPath = Path.Combine(cudaPath, "bin");
-
                 if (Directory.Exists(cudaBinPath))
                 {
                     AddDllDirectory(cudaBinPath);
                 }
-
                 version = GetCudaVersionFromPath(cudaPath);
             }
             else if (RuntimeInformation.IsOSPlatform(OSPlatform.Linux))
             {
+                string? env_version = Environment.GetEnvironmentVariable("CUDA_VERSION");
+                if (env_version is not null)
+                {
+                    return ExtractMajorVersion(ref env_version);
+                }
                 // List of default cuda paths
                 string[] defaultCudaPaths =
                 [
                     "/usr/local/bin/cuda",
                     "/usr/local/cuda",
                 ];
-                
                 // Loop through every default path to find the version
                 foreach (var path in defaultCudaPaths)
                 {
                     // Attempt to get the version from the path
                     version = GetCudaVersionFromPath(path);
-                    
+
                     // If a CUDA version is found, break the loop
                     if (!string.IsNullOrEmpty(version))
                         break;
                 }
-                
                 if (string.IsNullOrEmpty(version))
                 {
                     cudaPath = Environment.GetEnvironmentVariable("LD_LIBRARY_PATH");
@@ -193,17 +179,18 @@ private static int GetCudaMajorVersion()
                     }
                 }
             }
-
             if (string.IsNullOrEmpty(version))
                 return -1;
 
+            return ExtractMajorVersion(ref version);
+        }
+        private static int ExtractMajorVersion(ref string version)
+        {
             version = version.Split('.')[0];
             if (int.TryParse(version, out var majorVersion))
                 return majorVersion;
-
             return -1;
         }
-
         private static string GetCudaVersionFromPath(string cudaPath)
         {
             try
@@ -226,12 +213,10 @@ private static string GetCudaVersionFromPath(string cudaPath)
                 return string.Empty;
             }
         }
-
         // Put it here to avoid calling NativeApi when getting the cuda version.
         [DllImport("kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
         internal static extern int AddDllDirectory(string NewDirectory);
-
         private const string cudaVersionFile = "version.json";
         #endregion
     }
-}
+}