diff --git a/StabilityMatrix.Core/Python/PyVenvRunner.cs b/StabilityMatrix.Core/Python/PyVenvRunner.cs index eeb5d94f..93c4fb76 100644 --- a/StabilityMatrix.Core/Python/PyVenvRunner.cs +++ b/StabilityMatrix.Core/Python/PyVenvRunner.cs @@ -20,20 +20,20 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable private static readonly Logger Logger = LogManager.GetCurrentClassLogger(); public const string TorchPipInstallArgsCuda = - "torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118"; - public const string TorchPipInstallArgsCpu = - "torch torchvision torchaudio"; - public const string TorchPipInstallArgsDirectML = - "torch-directml"; + "torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118"; + public const string TorchPipInstallArgsCpu = "torch torchvision torchaudio"; + public const string TorchPipInstallArgsDirectML = "torch-directml"; /// /// Relative path to the site-packages folder from the venv root. /// This is platform specific. /// - public static string RelativeSitePackagesPath => Compat.Switch( - (PlatformKind.Windows, "Lib/site-packages"), - (PlatformKind.Unix, "lib/python3.10/site-packages")); - + public static string RelativeSitePackagesPath => + Compat.Switch( + (PlatformKind.Windows, "Lib/site-packages"), + (PlatformKind.Unix, "lib/python3.10/site-packages") + ); + /// /// The process running the python executable. /// @@ -48,51 +48,54 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable /// Optional working directory for the python process. /// public DirectoryPath? WorkingDirectory { get; set; } - + /// /// Optional environment variables for the python process. /// public IReadOnlyDictionary? EnvironmentVariables { get; set; } - + /// /// Name of the python binary folder. /// 'Scripts' on Windows, 'bin' on Unix. /// - public static string RelativeBinPath => Compat.Switch( - (PlatformKind.Windows, "Scripts"), - (PlatformKind.Unix, "bin")); - + public static string RelativeBinPath => + Compat.Switch((PlatformKind.Windows, "Scripts"), (PlatformKind.Unix, "bin")); + /// /// The relative path to the python executable. /// - public static string RelativePythonPath => Compat.Switch( - (PlatformKind.Windows, Path.Combine("Scripts", "python.exe")), - (PlatformKind.Unix, Path.Combine("bin", "python3"))); + public static string RelativePythonPath => + Compat.Switch( + (PlatformKind.Windows, Path.Combine("Scripts", "python.exe")), + (PlatformKind.Unix, Path.Combine("bin", "python3")) + ); /// /// The full path to the python executable. /// public FilePath PythonPath => RootPath.JoinFile(RelativePythonPath); - + /// /// The relative path to the pip executable. /// - public static string RelativePipPath => Compat.Switch( - (PlatformKind.Windows, Path.Combine("Scripts", "pip.exe")), - (PlatformKind.Unix, Path.Combine("bin", "pip3"))); + public static string RelativePipPath => + Compat.Switch( + (PlatformKind.Windows, Path.Combine("Scripts", "pip.exe")), + (PlatformKind.Unix, Path.Combine("bin", "pip3")) + ); /// /// The full path to the pip executable. /// public FilePath PipPath => RootPath.JoinFile(RelativePipPath); - + /// /// List of substrings to suppress from the output. /// When a line contains any of these substrings, it will not be forwarded to callbacks. /// A corresponding Info log will be written instead. /// public List SuppressOutput { get; } = new() { "fatal: not a git repository" }; - + public PyVenvRunner(DirectoryPath path) { RootPath = path; @@ -115,8 +118,13 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable RootPath.Create(); // Create venv (copy mode if windows) - var args = new string[] { "-m", "virtualenv", - Compat.IsWindows ? "--always-copy" : "", RootPath }; + var args = new string[] + { + "-m", + "virtualenv", + Compat.IsWindows ? "--always-copy" : "", + RootPath + }; var venvProc = ProcessRunner.StartAnsiProcess(PyRunner.PythonExePath, args); await venvProc.WaitForExitAsync().ConfigureAwait(false); @@ -126,7 +134,9 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable { var output = await venvProc.StandardOutput.ReadToEndAsync().ConfigureAwait(false); output += await venvProc.StandardError.ReadToEndAsync().ConfigureAwait(false); - throw new InvalidOperationException($"Venv creation failed with code {returnCode}: {output}"); + throw new InvalidOperationException( + $"Venv creation failed with code {returnCode}: {output}" + ); } } @@ -137,7 +147,8 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable private void SetPyvenvCfg(string pythonDirectory) { // Skip if we are not created yet - if (!Exists()) return; + if (!Exists()) + return; // Path to pyvenv.cfg var cfgPath = Path.Combine(RootPath, "pyvenv.cfg"); @@ -145,22 +156,25 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable { throw new FileNotFoundException("pyvenv.cfg not found", cfgPath); } - + Logger.Info("Updating pyvenv.cfg with embedded Python directory {PyDir}", pythonDirectory); - + // Insert a top section var topSection = "[top]" + Environment.NewLine; var cfg = new ConfigParser(topSection + File.ReadAllText(cfgPath)); - + // Need to set all path keys - home, base-prefix, base-exec-prefix, base-executable cfg.SetValue("top", "home", pythonDirectory); cfg.SetValue("top", "base-prefix", pythonDirectory); - + cfg.SetValue("top", "base-exec-prefix", pythonDirectory); - - cfg.SetValue("top", "base-executable", - Path.Combine(pythonDirectory, Compat.IsWindows ? "python.exe" : RelativePythonPath)); - + + cfg.SetValue( + "top", + "base-executable", + Path.Combine(pythonDirectory, Compat.IsWindows ? "python.exe" : RelativePythonPath) + ); + // Convert to string for writing, strip the top section var cfgString = cfg.ToString()!.Replace(topSection, ""); File.WriteAllText(cfgPath, cfgString); @@ -176,31 +190,35 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable { throw new FileNotFoundException("pip not found", PipPath); } - + // Record output for errors var output = new StringBuilder(); - - var outputAction = outputDataReceived == null ? null : new Action(s => - { - Logger.Debug($"Pip output: {s.Text}"); - // Record to output - output.Append(s.Text); - // Forward to callback - outputDataReceived(s); - }); - + + var outputAction = + outputDataReceived == null + ? null + : new Action(s => + { + Logger.Debug($"Pip output: {s.Text}"); + // Record to output + output.Append(s.Text); + // Forward to callback + outputDataReceived(s); + }); + SetPyvenvCfg(PyRunner.PythonDir); RunDetached($"-m pip install {args}", outputAction); await Process.WaitForExitAsync().ConfigureAwait(false); - + // Check return code if (Process.ExitCode != 0) { throw new ProcessException( - $"pip install failed with code {Process.ExitCode}: {output.ToString().ToRepr()}"); + $"pip install failed with code {Process.ExitCode}: {output.ToString().ToRepr()}" + ); } } - + /// /// Run a command using the venv Python executable and return the result. /// @@ -209,17 +227,23 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable { // Record output for errors var output = new StringBuilder(); - + var outputAction = new Action(s => { - if (s == null) return; + if (s == null) + return; Logger.Debug("Pip output: {Text}", s); output.Append(s); }); - + SetPyvenvCfg(PyRunner.PythonDir); - using var process = ProcessRunner.StartProcess(PythonPath, arguments, - WorkingDirectory?.FullPath, outputAction, EnvironmentVariables); + using var process = ProcessRunner.StartProcess( + PythonPath, + arguments, + WorkingDirectory?.FullPath, + outputAction, + EnvironmentVariables + ); await process.WaitForExitAsync().ConfigureAwait(false); return new ProcessResult @@ -231,45 +255,55 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable [MemberNotNull(nameof(Process))] public void RunDetached( - string arguments, + string arguments, Action? outputDataReceived, Action? onExit = null, - bool unbuffered = true) + bool unbuffered = true + ) { - if (!Exists()) + if (!PythonPath.Exists) { - throw new InvalidOperationException("Venv python process does not exist"); + throw new FileNotFoundException("Venv python not found", PythonPath); } SetPyvenvCfg(PyRunner.PythonDir); - - Logger.Debug($"Launching RunDetached at {PythonPath} with args {arguments}"); - var filteredOutput = outputDataReceived == null ? null : new Action(s => - { - if (SuppressOutput.Any(s.Text.Contains)) - { - Logger.Info("Filtered output: {S}", s); - return; - } - outputDataReceived.Invoke(s); - }); + Logger.Info( + "Launching venv process [{PythonPath}] " + + "in working directory [{WorkingDirectory}] with args {arguments.ToRepr()}", + PythonPath, + WorkingDirectory, + arguments.ToRepr() + ); + + var filteredOutput = + outputDataReceived == null + ? null + : new Action(s => + { + if (SuppressOutput.Any(s.Text.Contains)) + { + Logger.Info("Filtered output: {S}", s); + return; + } + outputDataReceived.Invoke(s); + }); var env = new Dictionary(); if (EnvironmentVariables != null) { env.Update(EnvironmentVariables); } - + // Disable pip caching - uses significant memory for large packages like torch env["PIP_NO_CACHE_DIR"] = "true"; - // On windows, add portable git + // On windows, add portable git if (Compat.IsWindows) { var portableGit = GlobalConfig.LibraryDir.JoinDir("PortableGit", "bin"); env["PATH"] = Compat.GetEnvPathWithExtensions(portableGit); } - + if (unbuffered) { env["PYTHONUNBUFFERED"] = "1"; @@ -286,10 +320,13 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable } } - Process = ProcessRunner.StartAnsiProcess(PythonPath, arguments, + Process = ProcessRunner.StartAnsiProcess( + PythonPath, + arguments, workingDirectory: WorkingDirectory?.FullPath, outputDataReceived: filteredOutput, - environmentVariables: env); + environmentVariables: env + ); if (onExit != null) { @@ -300,7 +337,7 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable }; } } - + /// /// Get entry points for a package. /// https://packaging.python.org/en/latest/specifications/entry-points/#entry-points @@ -314,7 +351,7 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable results = entry_points(group='console_scripts', name='{entryPointName}') print(tuple(results)[0].value, end='') """; - + var result = await Run($"-c \"{code}\"").ConfigureAwait(false); if (result.ExitCode == 0 && !string.IsNullOrWhiteSpace(result.StandardOutput)) { @@ -323,7 +360,7 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable return null; } - + /// /// Kills the running process and cancels stream readers, does not wait for exit. /// @@ -335,11 +372,11 @@ public class PyVenvRunner : IDisposable, IAsyncDisposable Process.Kill(); Process.Dispose(); } - + Process = null; GC.SuppressFinalize(this); } - + /// /// Kills the running process, waits for exit. ///