Safer Vulkan System Info

LLama/Extensions/ProcessExtensions.cs

@@ -0,0 +1,61 @@
+using System.Diagnostics;
+using System.Text;
+using System;
+
+namespace LLama.Extensions;
+
+internal static class ProcessExtensions
+{
+    public static void SafeKill(this Process process, bool entireProcessTree = true)
+    {
+        if (process.HasExited)
+            return;
+
+        // There's a race here! If the process closed between the above check
+        // and the below `Kill` call then an `InvalidOperationException` will
+        // be thrown! Catch it and move on.
+
+        try
+        {
+#if NET5_0_OR_GREATER
+            process.Kill(entireProcessTree);
+#else
+            process.Kill();
+#endif
+            process.WaitForExit(55);
+        }
+        catch (InvalidOperationException)
+        {
+        }
+    }
+
+    /// <summary>
+    /// Run a process for a certain amount of time and then terminate it
+    /// </summary>
+    /// <param name="process"></param>
+    /// <param name="timeout"></param>
+    /// <returns>return code, standard output, standard error, flag indicating if process exited or was terminated</returns>
+    public static (int ret, string stdOut, string stdErr, bool ok) SafeRun(this Process process, TimeSpan timeout)
+    {
+        var stdOut = new StringBuilder();
+        process.OutputDataReceived += (s, e) =>
+        {
+            stdOut.Append(e.Data);
+        };
+
+        var stdErr = new StringBuilder();
+        process.ErrorDataReceived += (s, e) =>
+        {
+            stdErr.Append(e.Data);
+        };
+
+        process.Start();
+        process.BeginOutputReadLine();
+        process.BeginErrorReadLine();
+
+        var ok = process.WaitForExit((int)timeout.TotalMilliseconds);
+        process.SafeKill();
+
+        return (process.ExitCode, stdOut.ToString(), stdErr.ToString(), ok);
+    }
+}

LLama/Native/Load/SystemInfo.cs

@@ -12,7 +12,7 @@ namespace LLama.Native
     /// <param name="OSPlatform"></param>
     /// <param name="CudaMajorVersion"></param>
     /// <param name="VulkanVersion"></param>
-    public record class SystemInfo(OSPlatform OSPlatform, int CudaMajorVersion, string? VulkanVersion)
+    public record SystemInfo(OSPlatform OSPlatform, int CudaMajorVersion, string? VulkanVersion)
     {
         /// <summary>
         /// Get the system information of the current machine.
@@ -71,36 +71,29 @@ public static SystemInfo Get()
             // Note: on Linux, this requires `vulkan-tools` to be installed. (`sudo apt install vulkan-tools`)
             try
             {
-                // Set up the process start info
-                ProcessStartInfo start = new()
-                {
-                    FileName = "vulkaninfo",
-                    Arguments = "--summary",
-                    RedirectStandardOutput = true,
-                    UseShellExecute = false,
-                    CreateNoWindow = true
-                };
-
-                // Start the process
+                // Start a process to read vulkan info
                 Process process = new()
                 {
-                    StartInfo = start
+                    StartInfo = new()
+                    {
+                        FileName = "vulkaninfo",
+                        Arguments = "--summary",
+                        RedirectStandardOutput = true,
+                        RedirectStandardError = true,
+                        UseShellExecute = false,
+                        CreateNoWindow = true
+                    }
                 };
-                process.Start();
-
-                // Read the output to a string
-                string output = process.StandardOutput.ReadToEnd();
-
-                // Wait for the process to exit
-                process.WaitForExit();
+                var (exitCode, output, error, ok) = process.SafeRun(TimeSpan.FromSeconds(1));
+
+                if (!ok)
+                    return null;
 
                 // Return the output
                 return output;
             }
-            catch (Exception e)
+            catch
             {
-                //Console.WriteLine(e);
-
                 // Return null if we failed to get the Vulkan version
                 return null;
             }