From d2eec06283a8bd517d7be9992c36b874a293b5eb Mon Sep 17 00:00:00 2001 From: Steven Pelech Date: Sun, 1 Nov 2026 15:15:03 -0600 Subject: [PATCH] fix(memory): prevent VramOrchestrator eviction latch and guard ComfyUI offline free calls --- Endpoints/WorkflowEndpoints.cs | 20 +- .../VramOrchestratorTests.cs | 194 +++++++++++++++++- Services/AiEngineManager.cs | 147 +++++++++++-- Services/LocalLlmMcpTools.cs | 43 ++-- Services/VramOrchestrator.cs | 89 ++++++-- 5 files changed, 428 insertions(+), 65 deletions(-) diff --git a/Endpoints/WorkflowEndpoints.cs b/Endpoints/WorkflowEndpoints.cs index 6c721b2..f0c3c14 100644 --- a/Endpoints/WorkflowEndpoints.cs +++ b/Endpoints/WorkflowEndpoints.cs @@ -70,7 +70,8 @@ public static void MapWorkflowEndpoints(this WebApplication app) try { - var response = await httpClient.PostAsync($"{baseUrl}/interrupt", new StringContent("{}", Encoding.UTF8, "application/json")); + using var content = new StringContent("{}", Encoding.UTF8, "application/json"); + using var response = await httpClient.PostAsync($"{baseUrl}/interrupt", content); return Results.Ok(new { success = response.IsSuccessStatusCode, status = (int)response.StatusCode }); } catch (Exception ex) @@ -151,6 +152,7 @@ public static void MapWorkflowEndpoints(this WebApplication app) app.MapPost("/api/3d/generate", async (ThreeDGenerateRequest request, ISettingsService settingsService, VramOrchestrator vramOrchestrator, HttpClient httpClient) => { + vramOrchestrator.RecordActivity("comfy"); await vramOrchestrator.EnsureVramForComfyUiAsync(); var settings = settingsService.LoadSettings(); @@ -271,8 +273,8 @@ public static void MapWorkflowEndpoints(this WebApplication app) { ["prompt"] = targetGraph?.DeepClone() }; - var content = new StringContent(payload.ToJsonString(), Encoding.UTF8, "application/json"); - var response = await httpClient.PostAsync($"{baseUrl}/prompt", content); + using var content = new StringContent(payload.ToJsonString(), Encoding.UTF8, "application/json"); + using var response = await httpClient.PostAsync($"{baseUrl}/prompt", content); if (response.IsSuccessStatusCode) { @@ -345,6 +347,7 @@ public static void MapWorkflowEndpoints(this WebApplication app) app.MapPost("/api/video/generate", async (VideoGenerateRequest request, ISettingsService settingsService, VramOrchestrator vramOrchestrator, HttpClient httpClient) => { + vramOrchestrator.RecordActivity("comfy"); await vramOrchestrator.EnsureVramForComfyUiAsync(); var settings = settingsService.LoadSettings(); @@ -441,8 +444,8 @@ public static void MapWorkflowEndpoints(this WebApplication app) { ["prompt"] = workflowNode }; - var content = new StringContent(payload.ToJsonString(), Encoding.UTF8, "application/json"); - var response = await httpClient.PostAsync($"{baseUrl}/prompt", content); + using var content = new StringContent(payload.ToJsonString(), Encoding.UTF8, "application/json"); + using var response = await httpClient.PostAsync($"{baseUrl}/prompt", content); if (response.IsSuccessStatusCode) { @@ -613,8 +616,9 @@ public static void MapWorkflowEndpoints(this WebApplication app) return Results.Ok(sortedList); }); - app.MapPost("/api/audio/generate", async (AudioGenerateRequest request, ISettingsService settingsService, IHttpClientFactory clientFactory) => + app.MapPost("/api/audio/generate", async (AudioGenerateRequest request, ISettingsService settingsService, VramOrchestrator vramOrchestrator, IHttpClientFactory clientFactory) => { + vramOrchestrator.RecordActivity("audio"); var settings = settingsService.LoadSettings(); var workflowsDir = string.IsNullOrWhiteSpace(settings.WorkflowsPath) ? Path.Combine(AppContext.BaseDirectory, "Workflows") @@ -713,11 +717,11 @@ public static void MapWorkflowEndpoints(this WebApplication app) try { - var http = clientFactory.CreateClient(); + using var http = clientFactory.CreateClient(); var comfyEndpoint = $"{comfyUrl.TrimEnd('/')}/prompt"; var payload = new { prompt = targetGraph }; - var response = await http.PostAsJsonAsync(comfyEndpoint, payload); + using var response = await http.PostAsJsonAsync(comfyEndpoint, payload); if (response.IsSuccessStatusCode) { var respJson = await response.Content.ReadFromJsonAsync(); diff --git a/LocalLLMServerManager.Tests/VramOrchestratorTests.cs b/LocalLLMServerManager.Tests/VramOrchestratorTests.cs index 79285de..ccc554b 100644 --- a/LocalLLMServerManager.Tests/VramOrchestratorTests.cs +++ b/LocalLLMServerManager.Tests/VramOrchestratorTests.cs @@ -98,9 +98,9 @@ public async Task IsComfyUiHealthyAsync_ReturnsTrue_WhenSystemStatsReturns200() } [Fact] - public async Task FreeComfyUiVramAsync_SendsPostToFreeEndpoint() + public async Task FreeComfyUiVramAsync_SendsPostToFreeEndpoint_WhenHealthy() { - HttpRequestMessage? capturedRequest = null; + var capturedRequests = new System.Collections.Generic.List(); var handlerMock = new Mock(); handlerMock.Protected() @@ -109,17 +109,191 @@ public async Task FreeComfyUiVramAsync_SendsPostToFreeEndpoint() ItExpr.IsAny(), ItExpr.IsAny() ) - .Callback((req, _) => capturedRequest = req) - .ReturnsAsync(new HttpResponseMessage { StatusCode = HttpStatusCode.OK }); + .Callback((req, _) => capturedRequests.Add(req)) + .ReturnsAsync(() => new HttpResponseMessage { StatusCode = HttpStatusCode.OK }); + + var client = new HttpClient(handlerMock.Object); + var orchestrator = new VramOrchestrator(client, NullLogger.Instance); + + var result = await orchestrator.FreeComfyUiVramAsync("http://127.0.0.1:8188"); + + Assert.True(result); + Assert.Equal(2, capturedRequests.Count); + Assert.Equal(HttpMethod.Get, capturedRequests[0].Method); + Assert.Equal("http://127.0.0.1:8188/system_stats", capturedRequests[0].RequestUri?.ToString()); + Assert.Equal(HttpMethod.Post, capturedRequests[1].Method); + Assert.Equal("http://127.0.0.1:8188/free", capturedRequests[1].RequestUri?.ToString()); + } + + [Fact] + public async Task FreeComfyUiVramAsync_DoesNotPostToFree_WhenComfyUiIsOffline() + { + var capturedRequests = new System.Collections.Generic.List(); + + var handlerMock = new Mock(); + handlerMock.Protected() + .Setup>( + "SendAsync", + ItExpr.IsAny(), + ItExpr.IsAny() + ) + .Callback((req, _) => capturedRequests.Add(req)) + .ThrowsAsync(new HttpRequestException("Connection refused")); var client = new HttpClient(handlerMock.Object); var orchestrator = new VramOrchestrator(client, NullLogger.Instance); - await orchestrator.FreeComfyUiVramAsync("http://127.0.0.1:8188"); + var result = await orchestrator.FreeComfyUiVramAsync("http://127.0.0.1:8188"); + + Assert.False(result); + Assert.Single(capturedRequests); + Assert.Equal(HttpMethod.Get, capturedRequests[0].Method); + Assert.Equal("http://127.0.0.1:8188/system_stats", capturedRequests[0].RequestUri?.ToString()); + Assert.DoesNotContain(capturedRequests, r => r.RequestUri!.ToString().EndsWith("/free")); + } + + [Fact] + public async Task EvaluateCacheLifetimesAsync_ExecutesEvictionUnderRamPressure_EvenAfterTier1EvictionLatch() + { + var timeProvider = new FakeTimeProvider(new DateTimeOffset(2026, 11, 1, 12, 0, 0, TimeSpan.Zero)); + + var handlerMock = new Mock(); + handlerMock.Protected() + .Setup>( + "SendAsync", + ItExpr.IsAny(), + ItExpr.IsAny() + ) + .ReturnsAsync((HttpRequestMessage req, CancellationToken _) => + { + var uri = req.RequestUri?.ToString() ?? ""; + if (uri.Contains("11434/api/ps")) + { + return new HttpResponseMessage(HttpStatusCode.OK) + { + Content = new StringContent("{\"models\":[{\"name\":\"llama3:latest\"}]}") + }; + } + return new HttpResponseMessage(HttpStatusCode.OK) + { + Content = new StringContent("{\"status\":\"ok\"}") + }; + }); + + var settingsMock = new Mock(); + settingsMock.Setup(s => s.LoadSettings()).Returns(new AppSettings( + ModelCacheIdleTimeoutMinutes: 15, + EngineHibernationTimeoutMinutes: 30, + RamPressureThresholdPercent: 75 + )); + + int currentRamPct = 40; + var telemetryMock = new Mock(); + telemetryMock.Setup(t => t.GetSystemMemoryInfo()).Returns(() => new LocalLLMServerManager.Services.SystemMemorySnapshot( + TotalRamBytes: 32L * 1024 * 1024 * 1024, + UsedRamBytes: 16L * 1024 * 1024 * 1024, + AvailableRamBytes: 16L * 1024 * 1024 * 1024, + RamUsedPercent: currentRamPct, + ProcessWorkingSetBytes: 256L * 1024 * 1024, + NonPagedPoolBytes: 128L * 1024 * 1024 + )); + + var engineMock = new Mock(); + var orchestrator = new VramOrchestrator( + new HttpClient(handlerMock.Object), + NullLogger.Instance, + settingsMock.Object, + telemetryMock.Object, + engineMock.Object, + timeProvider + ); + + // 1. Advance 16m under normal RAM (40%) -> Tier 1 eviction runs and sets _tier1Evicted = true + timeProvider.Advance(TimeSpan.FromMinutes(16)); + var tier1Result = await orchestrator.EvaluateCacheLifetimesAsync(); + Assert.True(tier1Result.OllamaUnloaded); + Assert.True(tier1Result.ComfyVramFreed); + + // 2. Now RAM pressure spikes (85% >= 75%) without RecordActivity being called -> must NOT be latched out + currentRamPct = 85; + timeProvider.Advance(TimeSpan.FromMinutes(3)); + var pressureResult = await orchestrator.EvaluateCacheLifetimesAsync(); + Assert.True(pressureResult.OllamaUnloaded); + Assert.True(pressureResult.ComfyVramFreed); + } + + [Fact] + public async Task EvaluateCacheLifetimesAsync_RehibernatesEngine_WhenProcessBecomesActiveAgainAfterHibernation() + { + var timeProvider = new FakeTimeProvider(new DateTimeOffset(2026, 11, 1, 12, 0, 0, TimeSpan.Zero)); + + var handlerMock = new Mock(); + handlerMock.Protected() + .Setup>( + "SendAsync", + ItExpr.IsAny(), + ItExpr.IsAny() + ) + .ReturnsAsync(() => new HttpResponseMessage(HttpStatusCode.OK) + { + Content = new StringContent("{\"models\":[]}") + }); + + var settingsMock = new Mock(); + settingsMock.Setup(s => s.LoadSettings()).Returns(new AppSettings( + ModelCacheIdleTimeoutMinutes: 15, + EngineHibernationTimeoutMinutes: 30, + RamPressureThresholdPercent: 75 + )); + + var telemetryMock = new Mock(); + telemetryMock.Setup(t => t.GetSystemMemoryInfo()).Returns(new LocalLLMServerManager.Services.SystemMemorySnapshot( + TotalRamBytes: 32L * 1024 * 1024 * 1024, + UsedRamBytes: 12L * 1024 * 1024 * 1024, + AvailableRamBytes: 20L * 1024 * 1024 * 1024, + RamUsedPercent: 40, + ProcessWorkingSetBytes: 256L * 1024 * 1024, + NonPagedPoolBytes: 128L * 1024 * 1024 + )); + + using var activeCurrentProcess = System.Diagnostics.Process.GetCurrentProcess(); + System.Diagnostics.Process? mockComfyProcess = null; - Assert.NotNull(capturedRequest); - Assert.Equal(HttpMethod.Post, capturedRequest.Method); - Assert.Equal("http://127.0.0.1:8188/free", capturedRequest.RequestUri?.ToString()); + var engineMock = new Mock(); + engineMock.Setup(e => e.ComfyProcess).Returns(() => mockComfyProcess); + engineMock.Setup(e => e.StopComfyUiAsync(It.IsAny())).ReturnsAsync(true); + engineMock.Setup(e => e.StopForgeAsync(It.IsAny())).ReturnsAsync(true); + engineMock.Setup(e => e.StopAudioEngineAsync(It.IsAny())).ReturnsAsync(true); + + var orchestrator = new VramOrchestrator( + new HttpClient(handlerMock.Object), + NullLogger.Instance, + settingsMock.Object, + telemetryMock.Object, + engineMock.Object, + timeProvider + ); + + // 1. Advance 31m -> initial Tier 2 hibernation runs and sets _tier2Hibernated["comfy"] = true + timeProvider.Advance(TimeSpan.FromMinutes(31)); + var firstHibernation = await orchestrator.EvaluateCacheLifetimesAsync(); + Assert.Contains("comfy", firstHibernation.HibernatedEngines); + engineMock.Verify(e => e.StopComfyUiAsync(It.IsAny()), Times.Once); + + // 2. ComfyUI process becomes active again externally without RecordActivity being called + mockComfyProcess = activeCurrentProcess; + timeProvider.Advance(TimeSpan.FromMinutes(31)); + var secondHibernation = await orchestrator.EvaluateCacheLifetimesAsync(); + Assert.Contains("comfy", secondHibernation.HibernatedEngines); + engineMock.Verify(e => e.StopComfyUiAsync(It.IsAny()), Times.Exactly(2)); + } + + private sealed class FakeTimeProvider : TimeProvider + { + private DateTimeOffset _utcNow; + public FakeTimeProvider(DateTimeOffset startUtc) => _utcNow = startUtc; + public override DateTimeOffset GetUtcNow() => _utcNow; + public void Advance(TimeSpan duration) => _utcNow = _utcNow.Add(duration); } [Fact] @@ -132,7 +306,7 @@ public async Task EnsureVramForImageGenerationAsync_ExecutesCleanly() ItExpr.IsAny(), ItExpr.IsAny() ) - .ReturnsAsync(new HttpResponseMessage { StatusCode = HttpStatusCode.OK }); + .ReturnsAsync(() => new HttpResponseMessage { StatusCode = HttpStatusCode.OK }); var client = new HttpClient(handlerMock.Object); var orchestrator = new VramOrchestrator(client, NullLogger.Instance); @@ -150,7 +324,7 @@ public async Task EnsureVramForComfyUiAsync_ExecutesCleanly() ItExpr.IsAny(), ItExpr.IsAny() ) - .ReturnsAsync(new HttpResponseMessage { StatusCode = HttpStatusCode.OK }); + .ReturnsAsync(() => new HttpResponseMessage { StatusCode = HttpStatusCode.OK }); var client = new HttpClient(handlerMock.Object); var orchestrator = new VramOrchestrator(client, NullLogger.Instance); diff --git a/Services/AiEngineManager.cs b/Services/AiEngineManager.cs index da98504..dcba4a7 100644 --- a/Services/AiEngineManager.cs +++ b/Services/AiEngineManager.cs @@ -1,6 +1,9 @@ +using System; +using System.Collections.Generic; using System.Diagnostics; using System.IO; using System.Runtime.InteropServices; +using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Logging; namespace LocalLLMServerManager.Services; @@ -13,17 +16,82 @@ public class AiEngineManager : IAiEngineManager private static Process? _ollamaProcess; private static readonly JobObject AiEnginesJob = new(); - public Process? ComfyProcess => _comfyProcess; - public Process? ForgeProcess => _forgeProcess; - public Process? AudioProcess => _audioProcess; - public Process? OllamaProcess => _ollamaProcess; + private readonly Lazy? _vramOrchestrator; + + public AiEngineManager() + { + } + + [ActivatorUtilitiesConstructor] + public AiEngineManager(IServiceProvider? serviceProvider = null) + { + if (serviceProvider != null) + { + _vramOrchestrator = new Lazy(() => serviceProvider.GetService()); + } + } + + public AiEngineManager(Lazy vramOrchestrator) + { + _vramOrchestrator = new Lazy(() => vramOrchestrator.Value); + } + + public Process? ComfyProcess => GetActiveProcess(ref _comfyProcess); + public Process? ForgeProcess => GetActiveProcess(ref _forgeProcess); + public Process? AudioProcess => GetActiveProcess(ref _audioProcess); + public Process? OllamaProcess => GetActiveProcess(ref _ollamaProcess); + + private static Process? GetActiveProcess(ref Process? process) + { + if (process == null) return null; + try + { + if (process.HasExited) + { + process.Dispose(); + process = null; + return null; + } + return process; + } + catch + { + try { process?.Dispose(); } catch { } + process = null; + return null; + } + } + + public IReadOnlyDictionary GetActiveProcesses() + { + var active = new Dictionary(StringComparer.OrdinalIgnoreCase); + + var comfy = GetActiveProcess(ref _comfyProcess); + if (comfy != null) active["comfy"] = comfy; + + var forge = GetActiveProcess(ref _forgeProcess); + if (forge != null) active["forge"] = forge; + + var audio = GetActiveProcess(ref _audioProcess); + if (audio != null) active["audio"] = audio; + + var ollama = GetActiveProcess(ref _ollamaProcess); + if (ollama != null) active["ollama"] = ollama; + + return active; + } public bool IsProcessRunning(string name) { try { var processes = Process.GetProcessesByName(name); - return processes.Length > 0; + var running = processes.Length > 0; + foreach (var p in processes) + { + p.Dispose(); + } + return running; } catch { @@ -35,7 +103,11 @@ public Task StartComfyUiAsync(string executablePath, ILogger logger) { try { - if (_comfyProcess != null && !_comfyProcess.HasExited) return Task.FromResult(true); + if (GetActiveProcess(ref _comfyProcess) != null) + { + _vramOrchestrator?.Value?.RecordActivity("comfy"); + return Task.FromResult(true); + } var expandedPath = ResolveComfyExecutable(executablePath); if (!File.Exists(expandedPath)) @@ -58,6 +130,7 @@ public Task StartComfyUiAsync(string executablePath, ILogger logger) AiEnginesJob.AddProcess(_comfyProcess); } + _vramOrchestrator?.Value?.RecordActivity("comfy"); logger.LogInformation("Started ComfyUI engine process PID {Pid}", _comfyProcess?.Id); return Task.FromResult(true); } @@ -72,9 +145,11 @@ public Task StopComfyUiAsync(ILogger logger) { try { - if (_comfyProcess != null && !_comfyProcess.HasExited) + var active = GetActiveProcess(ref _comfyProcess); + if (active != null) { - _comfyProcess.Kill(true); + active.Kill(true); + active.Dispose(); _comfyProcess = null; logger.LogInformation("Stopped ComfyUI process"); return Task.FromResult(true); @@ -91,6 +166,10 @@ public Task StopComfyUiAsync(ILogger logger) } } catch { } + finally + { + p.Dispose(); + } } _comfyProcess = null; return Task.FromResult(true); @@ -106,7 +185,11 @@ public Task StartForgeAsync(string executablePath, ILogger logger) { try { - if (_forgeProcess != null && !_forgeProcess.HasExited) return Task.FromResult(true); + if (GetActiveProcess(ref _forgeProcess) != null) + { + _vramOrchestrator?.Value?.RecordActivity("forge"); + return Task.FromResult(true); + } var expandedPath = ResolveForgeExecutable(executablePath); if (!File.Exists(expandedPath)) @@ -129,6 +212,7 @@ public Task StartForgeAsync(string executablePath, ILogger logger) AiEnginesJob.AddProcess(_forgeProcess); } + _vramOrchestrator?.Value?.RecordActivity("forge"); logger.LogInformation("Started SD Forge process PID {Pid}", _forgeProcess?.Id); return Task.FromResult(true); } @@ -143,9 +227,11 @@ public Task StopForgeAsync(ILogger logger) { try { - if (_forgeProcess != null && !_forgeProcess.HasExited) + var active = GetActiveProcess(ref _forgeProcess); + if (active != null) { - _forgeProcess.Kill(true); + active.Kill(true); + active.Dispose(); _forgeProcess = null; logger.LogInformation("Stopped SD Forge process"); return Task.FromResult(true); @@ -164,7 +250,11 @@ public Task StartAudioEngineAsync(string executablePath, ILogger logger) { try { - if (_audioProcess != null && !_audioProcess.HasExited) return Task.FromResult(true); + if (GetActiveProcess(ref _audioProcess) != null) + { + _vramOrchestrator?.Value?.RecordActivity("audio"); + return Task.FromResult(true); + } var expandedPath = ResolveAudioExecutable(executablePath); ProcessStartInfo startInfo; @@ -221,6 +311,7 @@ public Task StartAudioEngineAsync(string executablePath, ILogger logger) AiEnginesJob.AddProcess(_audioProcess); } + _vramOrchestrator?.Value?.RecordActivity("audio"); logger.LogInformation("Started Audio Engine process PID {Pid}", _audioProcess?.Id); return Task.FromResult(true); } @@ -235,9 +326,11 @@ public Task StopAudioEngineAsync(ILogger logger) { try { - if (_audioProcess != null && !_audioProcess.HasExited) + var active = GetActiveProcess(ref _audioProcess); + if (active != null) { - _audioProcess.Kill(true); + active.Kill(true); + active.Dispose(); _audioProcess = null; logger.LogInformation("Stopped Audio Engine process"); return Task.FromResult(true); @@ -256,6 +349,10 @@ public Task StopAudioEngineAsync(ILogger logger) } } catch { } + finally + { + p.Dispose(); + } } _audioProcess = null; return Task.FromResult(true); @@ -271,9 +368,14 @@ public Task StartOllamaAsync(string executablePath, ILogger logger) { try { - if (_ollamaProcess != null && !_ollamaProcess.HasExited) return Task.FromResult(true); + if (GetActiveProcess(ref _ollamaProcess) != null) + { + _vramOrchestrator?.Value?.RecordActivity("ollama"); + return Task.FromResult(true); + } if (IsProcessRunning("ollama")) { + _vramOrchestrator?.Value?.RecordActivity("ollama"); logger.LogInformation("Ollama process is already running"); return Task.FromResult(true); } @@ -299,6 +401,7 @@ public Task StartOllamaAsync(string executablePath, ILogger logger) AiEnginesJob.AddProcess(_ollamaProcess); } + _vramOrchestrator?.Value?.RecordActivity("ollama"); logger.LogInformation("Started Ollama server process PID {Pid}", _ollamaProcess?.Id); return Task.FromResult(true); } @@ -394,9 +497,11 @@ public Task StopOllamaAsync(ILogger logger) { try { - if (_ollamaProcess != null && !_ollamaProcess.HasExited) + var active = GetActiveProcess(ref _ollamaProcess); + if (active != null) { - _ollamaProcess.Kill(true); + active.Kill(true); + active.Dispose(); _ollamaProcess = null; logger.LogInformation("Stopped managed Ollama process"); } @@ -409,6 +514,10 @@ public Task StopOllamaAsync(ILogger logger) p.Kill(true); } catch { } + finally + { + p.Dispose(); + } } var appProcesses = Process.GetProcessesByName("ollama_app"); @@ -419,6 +528,10 @@ public Task StopOllamaAsync(ILogger logger) p.Kill(true); } catch { } + finally + { + p.Dispose(); + } } _ollamaProcess = null; diff --git a/Services/LocalLlmMcpTools.cs b/Services/LocalLlmMcpTools.cs index 52d4f2d..d504ce7 100644 --- a/Services/LocalLlmMcpTools.cs +++ b/Services/LocalLlmMcpTools.cs @@ -24,6 +24,7 @@ public sealed class LocalLlmMcpTools private readonly IHuggingFaceSearchService _hfSearchService; private readonly ICivitaiSearchService _civitaiSearchService; private readonly ISettingsService _settingsService; + private readonly VramOrchestrator? _vramOrchestrator; public LocalLlmMcpTools( IGpuTelemetryProvider telemetryProvider, @@ -33,7 +34,8 @@ public LocalLlmMcpTools( IHttpClientFactory httpClientFactory, IHuggingFaceSearchService? hfSearchService = null, ICivitaiSearchService? civitaiSearchService = null, - ISettingsService? settingsService = null) + ISettingsService? settingsService = null, + VramOrchestrator? vramOrchestrator = null) { _telemetryProvider = telemetryProvider; _engineManager = engineManager; @@ -43,6 +45,7 @@ public LocalLlmMcpTools( _hfSearchService = hfSearchService ?? new HuggingFaceSearchService(); _civitaiSearchService = civitaiSearchService ?? new CivitaiSearchService(); _settingsService = settingsService ?? new SettingsService(); + _vramOrchestrator = vramOrchestrator; } [McpServerTool, Description("Get real-time GPU VRAM allocation, total memory, used memory, and GPU hardware name via NVML CUDA.")] @@ -63,7 +66,7 @@ async Task CheckPort(string url) try { var sw = System.Diagnostics.Stopwatch.StartNew(); - var resp = await client.GetAsync(url); + using var resp = await client.GetAsync(url); sw.Stop(); return new { online = resp.IsSuccessStatusCode, status = (int)resp.StatusCode, latencyMs = sw.ElapsedMilliseconds }; } @@ -96,6 +99,7 @@ public async Task PullModelAsync([Description("Model identifier, e.g. 'l if (string.IsNullOrWhiteSpace(modelName)) return JsonSerializer.Serialize(new { success = false, error = "modelName is required" }); + _vramOrchestrator?.RecordActivity("ollama"); var started = await _ollamaModelService.PullModelAsync(modelName); return JsonSerializer.Serialize(new { success = started, modelName, message = started ? "Model pull initiated" : "Failed to initiate pull" }); } @@ -106,8 +110,8 @@ public async Task UnloadVramAsync() try { using var client = _httpClientFactory.CreateClient(); - var payload = new StringContent("{\"model\":\"\",\"keep_alive\":0}", System.Text.Encoding.UTF8, "application/json"); - var response = await client.PostAsync("http://127.0.0.1:11434/api/generate", payload); + using var payload = new StringContent("{\"model\":\"\",\"keep_alive\":0}", System.Text.Encoding.UTF8, "application/json"); + using var response = await client.PostAsync("http://127.0.0.1:11434/api/generate", payload); return JsonSerializer.Serialize(new { success = response.IsSuccessStatusCode, status = (int)response.StatusCode, message = "VRAM unload requested" }); } catch (Exception ex) @@ -119,6 +123,7 @@ public async Task UnloadVramAsync() [McpServerTool, Description("Start an AI backend engine process ('forge', 'comfyui', or 'ollama').")] public async Task StartEngineAsync([Description("Target engine: 'forge', 'comfyui', or 'ollama'")] string engine) { + _vramOrchestrator?.RecordActivity(engine); var result = await _engineManager.StartEngineAsync(engine); return JsonSerializer.Serialize(result, new JsonSerializerOptions { WriteIndented = true }); } @@ -148,12 +153,14 @@ public async Task GenerateVideoAsync( if (string.IsNullOrWhiteSpace(prompt)) return JsonSerializer.Serialize(new { success = false, error = "prompt is required" }); + _vramOrchestrator?.RecordActivity("comfy"); + try { using var client = _httpClientFactory.CreateClient(); var body = JsonSerializer.Serialize(new { prompt, workflow, width, height, frames }); - var content = new StringContent(body, System.Text.Encoding.UTF8, "application/json"); - var response = await client.PostAsync("http://127.0.0.1:8188/prompt", content); + using var content = new StringContent(body, System.Text.Encoding.UTF8, "application/json"); + using var response = await client.PostAsync("http://127.0.0.1:8188/prompt", content); var mediaId = Guid.NewGuid().ToString("N")[..8]; var result = new @@ -185,12 +192,14 @@ public async Task SynthesizeSpeechAsync( if (string.IsNullOrWhiteSpace(text)) return JsonSerializer.Serialize(new { success = false, error = "text is required" }); + _vramOrchestrator?.RecordActivity("audio"); + try { using var client = _httpClientFactory.CreateClient(); var body = JsonSerializer.Serialize(new { text, voice, format }); - var content = new StringContent(body, System.Text.Encoding.UTF8, "application/json"); - var response = await client.PostAsync("http://127.0.0.1:7851/api/tts", content); + using var content = new StringContent(body, System.Text.Encoding.UTF8, "application/json"); + using var response = await client.PostAsync("http://127.0.0.1:7851/api/tts", content); var mediaId = Guid.NewGuid().ToString("N")[..8]; var ext = string.IsNullOrWhiteSpace(format) ? "mp3" : format.TrimStart('.'); @@ -220,12 +229,14 @@ public async Task GenerateAudioAsync( if (string.IsNullOrWhiteSpace(prompt)) return JsonSerializer.Serialize(new { success = false, error = "prompt is required" }); + _vramOrchestrator?.RecordActivity("audio"); + try { using var client = _httpClientFactory.CreateClient(); var body = JsonSerializer.Serialize(new { prompt, durationSeconds }); - var content = new StringContent(body, System.Text.Encoding.UTF8, "application/json"); - var response = await client.PostAsync("http://127.0.0.1:7860/api/audio/generate", content); + using var content = new StringContent(body, System.Text.Encoding.UTF8, "application/json"); + using var response = await client.PostAsync("http://127.0.0.1:7860/api/audio/generate", content); var mediaId = Guid.NewGuid().ToString("N")[..8]; var result = new @@ -310,8 +321,9 @@ public async Task RunStudioWorkflowAsync( if (w.Contains("text") || w.Contains("llm") || w.Contains("ollama")) { + _vramOrchestrator?.RecordActivity("ollama"); var payload = new { prompt, model = "llama3.2:latest", stream = false }; - var resp = await client.PostAsJsonAsync("http://127.0.0.1:11434/api/generate", payload); + using var resp = await client.PostAsJsonAsync("http://127.0.0.1:11434/api/generate", payload); if (resp.IsSuccessStatusCode) { var textResp = await resp.Content.ReadAsStringAsync(); @@ -321,6 +333,7 @@ public async Task RunStudioWorkflowAsync( } else if (w.Contains("video") || w.Contains("wan") || w.Contains("ltx")) { + _vramOrchestrator?.RecordActivity("comfy"); var payload = new { WorkflowId = w.Contains("wan") || w == "video" ? "wan2.2_t2v" : w, @@ -331,12 +344,13 @@ public async Task RunStudioWorkflowAsync( Frames = 49, Seed = seed }; - var resp = await client.PostAsJsonAsync("http://127.0.0.1:5246/api/video/generate", payload); + using var resp = await client.PostAsJsonAsync("http://127.0.0.1:5246/api/video/generate", payload); var content = await resp.Content.ReadAsStringAsync(); return JsonSerializer.Serialize(new { success = resp.IsSuccessStatusCode, modality = "video", data = JsonNode.Parse(content) }); } else if (w.Contains("audio") || w.Contains("sound") || w.Contains("music")) { + _vramOrchestrator?.RecordActivity("audio"); var payload = new { WorkflowId = "stable_audio_open_sfx", @@ -345,13 +359,14 @@ public async Task RunStudioWorkflowAsync( DurationSeconds = Math.Max(5, (int)(steps * 0.6)), Seed = seed }; - var resp = await client.PostAsJsonAsync("http://127.0.0.1:5246/api/audio/generate", payload); + using var resp = await client.PostAsJsonAsync("http://127.0.0.1:5246/api/audio/generate", payload); var content = await resp.Content.ReadAsStringAsync(); return JsonSerializer.Serialize(new { success = resp.IsSuccessStatusCode, modality = "audio", data = JsonNode.Parse(content) }); } else { // Default to Image via SD Forge / Comfy + _vramOrchestrator?.RecordActivity("forge"); var payload = new { prompt, @@ -362,7 +377,7 @@ public async Task RunStudioWorkflowAsync( width, height }; - var resp = await client.PostAsJsonAsync("http://127.0.0.1:7860/sdapi/v1/txt2img", payload); + using var resp = await client.PostAsJsonAsync("http://127.0.0.1:7860/sdapi/v1/txt2img", payload); if (resp.IsSuccessStatusCode) { var respStr = await resp.Content.ReadAsStringAsync(); diff --git a/Services/VramOrchestrator.cs b/Services/VramOrchestrator.cs index c554357..7e117c9 100644 --- a/Services/VramOrchestrator.cs +++ b/Services/VramOrchestrator.cs @@ -31,6 +31,7 @@ public class VramOrchestrator private readonly TimeProvider _timeProvider; private readonly ConcurrentDictionary _lastActivityUtc = new(StringComparer.OrdinalIgnoreCase); + private readonly ConcurrentDictionary _lastPressureEvictionUtc = new(StringComparer.OrdinalIgnoreCase); private readonly ConcurrentDictionary _tier1Evicted = new(StringComparer.OrdinalIgnoreCase); private readonly ConcurrentDictionary _tier2Hibernated = new(StringComparer.OrdinalIgnoreCase); @@ -79,6 +80,7 @@ public void RecordActivity(string engine) { var now = _timeProvider.GetUtcNow(); _lastActivityUtc[key] = now; + _lastPressureEvictionUtc.TryRemove(key, out _); _tier1Evicted[key] = false; _tier2Hibernated[key] = false; _logger.LogDebug("Recorded activity for engine '{Engine}' at {Timestamp}", key, now); @@ -97,7 +99,7 @@ public async Task IsOllamaHealthyAsync(CancellationToken cancellationToken { using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); cts.CancelAfter(TimeSpan.FromSeconds(1.5)); - var response = await _httpClient.GetAsync("http://127.0.0.1:11434/", cts.Token); + using var response = await _httpClient.GetAsync("http://127.0.0.1:11434/", cts.Token); return response.IsSuccessStatusCode; } catch @@ -112,7 +114,7 @@ public async Task IsForgeHealthyAsync(CancellationToken cancellationToken { using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); cts.CancelAfter(TimeSpan.FromSeconds(1.5)); - var response = await _httpClient.GetAsync("http://127.0.0.1:7860/sdapi/v1/progress", cts.Token); + using var response = await _httpClient.GetAsync("http://127.0.0.1:7860/sdapi/v1/progress", cts.Token); return response.IsSuccessStatusCode; } catch @@ -128,7 +130,7 @@ public async Task IsComfyUiHealthyAsync(string comfyUrl = "http://127.0.0. using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); cts.CancelAfter(TimeSpan.FromSeconds(1.5)); var baseUrl = comfyUrl.TrimEnd('/'); - var response = await _httpClient.GetAsync($"{baseUrl}/system_stats", cts.Token); + using var response = await _httpClient.GetAsync($"{baseUrl}/system_stats", cts.Token); return response.IsSuccessStatusCode; } catch @@ -151,10 +153,15 @@ public async Task FreeComfyUiVramAsync(string comfyUrl = "http://127.0.0.1 { try { + if (!await IsComfyUiHealthyAsync(comfyUrl, cancellationToken)) + { + return false; + } + var baseUrl = comfyUrl.TrimEnd('/'); var payload = new { free_memory = true, unload_models = true }; - var content = new StringContent(JsonSerializer.Serialize(payload), Encoding.UTF8, "application/json"); - var response = await _httpClient.PostAsync($"{baseUrl}/free", content, cancellationToken); + using var content = new StringContent(JsonSerializer.Serialize(payload), Encoding.UTF8, "application/json"); + using var response = await _httpClient.PostAsync($"{baseUrl}/free", content, cancellationToken); _logger.LogInformation("Triggered ComfyUI VRAM unload / free memory."); return response.IsSuccessStatusCode; } @@ -169,7 +176,7 @@ public async Task UnloadOllamaModelsAsync(string reason, CancellationToken { try { - var psResponse = await _httpClient.GetAsync("http://127.0.0.1:11434/api/ps", cancellationToken); + using var psResponse = await _httpClient.GetAsync("http://127.0.0.1:11434/api/ps", cancellationToken); if (!psResponse.IsSuccessStatusCode) return false; var psContent = await psResponse.Content.ReadAsStringAsync(cancellationToken); @@ -186,8 +193,8 @@ public async Task UnloadOllamaModelsAsync(string reason, CancellationToken if (!string.IsNullOrEmpty(modelName)) { var payload = new { model = modelName, keep_alive = 0 }; - var content = new StringContent(JsonSerializer.Serialize(payload), Encoding.UTF8, "application/json"); - await _httpClient.PostAsync("http://127.0.0.1:11434/api/generate", content, cancellationToken); + using var content = new StringContent(JsonSerializer.Serialize(payload), Encoding.UTF8, "application/json"); + using var genResponse = await _httpClient.PostAsync("http://127.0.0.1:11434/api/generate", content, cancellationToken); _logger.LogInformation($"Unloaded model: {modelName}"); } } @@ -209,6 +216,40 @@ public async Task UnloadOllamaModelsAsync(string reason, CancellationToken } } + private HashSet GetActiveEngineKeys() + { + var active = new HashSet(StringComparer.OrdinalIgnoreCase); + if (_engineManager == null) return active; + + if (_engineManager is AiEngineManager aiEngineManager) + { + foreach (var key in aiEngineManager.GetActiveProcesses().Keys) + { + active.Add(NormalizeEngineKey(key)); + } + } + + if (IsProcessRunningSafe(_engineManager.ComfyProcess)) active.Add("comfy"); + if (IsProcessRunningSafe(_engineManager.ForgeProcess)) active.Add("forge"); + if (IsProcessRunningSafe(_engineManager.AudioProcess)) active.Add("audio"); + if (IsProcessRunningSafe(_engineManager.OllamaProcess)) active.Add("ollama"); + + return active; + } + + private static bool IsProcessRunningSafe(System.Diagnostics.Process? process) + { + if (process == null) return false; + try + { + return !process.HasExited; + } + catch + { + return false; + } + } + public async Task EvaluateCacheLifetimesAsync(CancellationToken cancellationToken = default) { var settings = _settingsService?.LoadSettings() ?? new AppSettings(); @@ -223,6 +264,16 @@ public async Task EvaluateCacheLifetimesAsync(CancellationT bool isRamPressure = ramPressureThreshold > 0 && ramUsedPct >= ramPressureThreshold; var comfyUrl = string.IsNullOrWhiteSpace(settings.ComfyUiUrl) ? "http://127.0.0.1:8188" : settings.ComfyUiUrl; + var activeEngines = GetActiveEngineKeys(); + foreach (var engine in activeEngines) + { + if (_tier2Hibernated.TryGetValue(engine, out var wasHibernated) && wasHibernated) + { + _tier2Hibernated[engine] = false; + _tier1Evicted[engine] = false; + } + } + bool ollamaUnloaded = false; bool comfyVramFreed = false; var hibernatedEngines = new List(); @@ -232,22 +283,28 @@ public async Task EvaluateCacheLifetimesAsync(CancellationT var twoMinutes = TimeSpan.FromMinutes(2); var ollamaIdle = now - (_lastActivityUtc.TryGetValue("ollama", out var oLast) ? oLast : now); - if (ollamaIdle >= twoMinutes && (!_tier1Evicted.TryGetValue("ollama", out var oEvicted) || !oEvicted)) + bool ollamaCanPressureEvict = !_tier1Evicted.TryGetValue("ollama", out var oEvicted) || !oEvicted || + !_lastPressureEvictionUtc.TryGetValue("ollama", out var oPressLast) || (now - oPressLast) >= twoMinutes; + if (ollamaIdle >= twoMinutes && ollamaCanPressureEvict) { ollamaUnloaded = await UnloadOllamaModelsAsync($"RAM pressure ({ramUsedPct}% >= {ramPressureThreshold}%)", cancellationToken); _tier1Evicted["ollama"] = true; + _lastPressureEvictionUtc["ollama"] = now; } var comfyIdle = now - (_lastActivityUtc.TryGetValue("comfy", out var cLast) ? cLast : now); - if (comfyIdle >= twoMinutes && (!_tier1Evicted.TryGetValue("comfy", out var cEvicted) || !cEvicted)) + bool comfyCanPressureEvict = !_tier1Evicted.TryGetValue("comfy", out var cEvicted) || !cEvicted || + !_lastPressureEvictionUtc.TryGetValue("comfy", out var cPressLast) || (now - cPressLast) >= twoMinutes; + if (comfyIdle >= twoMinutes && comfyCanPressureEvict) { comfyVramFreed = await FreeComfyUiVramAsync(comfyUrl, cancellationToken); _tier1Evicted["comfy"] = true; + _lastPressureEvictionUtc["comfy"] = now; } if (_engineManager != null) { - if (comfyIdle >= twoMinutes && (!_tier2Hibernated.TryGetValue("comfy", out var cHib) || !cHib)) + if (comfyIdle >= twoMinutes && (!_tier2Hibernated.TryGetValue("comfy", out var cHib) || !cHib || activeEngines.Contains("comfy"))) { await _engineManager.StopComfyUiAsync(_logger); hibernatedEngines.Add("comfy"); @@ -255,7 +312,7 @@ public async Task EvaluateCacheLifetimesAsync(CancellationT } var forgeIdle = now - (_lastActivityUtc.TryGetValue("forge", out var fLast) ? fLast : now); - if (forgeIdle >= twoMinutes && (!_tier2Hibernated.TryGetValue("forge", out var fHib) || !fHib)) + if (forgeIdle >= twoMinutes && (!_tier2Hibernated.TryGetValue("forge", out var fHib) || !fHib || activeEngines.Contains("forge"))) { await _engineManager.StopForgeAsync(_logger); hibernatedEngines.Add("forge"); @@ -263,7 +320,7 @@ public async Task EvaluateCacheLifetimesAsync(CancellationT } var audioIdle = now - (_lastActivityUtc.TryGetValue("audio", out var aLast) ? aLast : now); - if (audioIdle >= twoMinutes && (!_tier2Hibernated.TryGetValue("audio", out var aHib) || !aHib)) + if (audioIdle >= twoMinutes && (!_tier2Hibernated.TryGetValue("audio", out var aHib) || !aHib || activeEngines.Contains("audio"))) { await _engineManager.StopAudioEngineAsync(_logger); hibernatedEngines.Add("audio"); @@ -303,7 +360,7 @@ public async Task EvaluateCacheLifetimesAsync(CancellationT var tier2Ttl = TimeSpan.FromMinutes(engineHibernateTimeoutMinutes); var comfyIdle = now - (_lastActivityUtc.TryGetValue("comfy", out var cLast) ? cLast : now); - if (comfyIdle >= tier2Ttl && (!_tier2Hibernated.TryGetValue("comfy", out var cHib) || !cHib)) + if (comfyIdle >= tier2Ttl && (!_tier2Hibernated.TryGetValue("comfy", out var cHib) || !cHib || activeEngines.Contains("comfy"))) { await _engineManager.StopComfyUiAsync(_logger); hibernatedEngines.Add("comfy"); @@ -311,7 +368,7 @@ public async Task EvaluateCacheLifetimesAsync(CancellationT } var forgeIdle = now - (_lastActivityUtc.TryGetValue("forge", out var fLast) ? fLast : now); - if (forgeIdle >= tier2Ttl && (!_tier2Hibernated.TryGetValue("forge", out var fHib) || !fHib)) + if (forgeIdle >= tier2Ttl && (!_tier2Hibernated.TryGetValue("forge", out var fHib) || !fHib || activeEngines.Contains("forge"))) { await _engineManager.StopForgeAsync(_logger); hibernatedEngines.Add("forge"); @@ -319,7 +376,7 @@ public async Task EvaluateCacheLifetimesAsync(CancellationT } var audioIdle = now - (_lastActivityUtc.TryGetValue("audio", out var aLast) ? aLast : now); - if (audioIdle >= tier2Ttl && (!_tier2Hibernated.TryGetValue("audio", out var aHib) || !aHib)) + if (audioIdle >= tier2Ttl && (!_tier2Hibernated.TryGetValue("audio", out var aHib) || !aHib || activeEngines.Contains("audio"))) { await _engineManager.StopAudioEngineAsync(_logger); hibernatedEngines.Add("audio");