Skip to content

fix(memory): prevent VramOrchestrator eviction latch and guard ComfyUI offline free calls #106

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
spelech merged 1 commit into from
Nov 1, 2026
Merged

fix(memory): prevent VramOrchestrator eviction latch and guard ComfyUI offline free calls #106

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 12 additions & 8 deletions Endpoints/WorkflowEndpoints.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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)
{
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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)
{
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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<JsonElement>();
Expand Down
194 changes: 184 additions & 10 deletions LocalLLMServerManager.Tests/VramOrchestratorTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<HttpRequestMessage>();

var handlerMock = new Mock<HttpMessageHandler>();
handlerMock.Protected()
Expand All @@ -109,17 +109,191 @@ public async Task FreeComfyUiVramAsync_SendsPostToFreeEndpoint()
ItExpr.IsAny<HttpRequestMessage>(),
ItExpr.IsAny<CancellationToken>()
)
.Callback<HttpRequestMessage, CancellationToken>((req, _) => capturedRequest = req)
.ReturnsAsync(new HttpResponseMessage { StatusCode = HttpStatusCode.OK });
.Callback<HttpRequestMessage, CancellationToken>((req, _) => capturedRequests.Add(req))
.ReturnsAsync(() => new HttpResponseMessage { StatusCode = HttpStatusCode.OK });

var client = new HttpClient(handlerMock.Object);
var orchestrator = new VramOrchestrator(client, NullLogger<VramOrchestrator>.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<HttpRequestMessage>();

var handlerMock = new Mock<HttpMessageHandler>();
handlerMock.Protected()
.Setup<Task<HttpResponseMessage>>(
"SendAsync",
ItExpr.IsAny<HttpRequestMessage>(),
ItExpr.IsAny<CancellationToken>()
)
.Callback<HttpRequestMessage, CancellationToken>((req, _) => capturedRequests.Add(req))
.ThrowsAsync(new HttpRequestException("Connection refused"));

var client = new HttpClient(handlerMock.Object);
var orchestrator = new VramOrchestrator(client, NullLogger<VramOrchestrator>.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<HttpMessageHandler>();
handlerMock.Protected()
.Setup<Task<HttpResponseMessage>>(
"SendAsync",
ItExpr.IsAny<HttpRequestMessage>(),
ItExpr.IsAny<CancellationToken>()
)
.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<LocalLLMServerManager.Shared.Interfaces.ISettingsService>();
settingsMock.Setup(s => s.LoadSettings()).Returns(new AppSettings(
ModelCacheIdleTimeoutMinutes: 15,
EngineHibernationTimeoutMinutes: 30,
RamPressureThresholdPercent: 75
));

int currentRamPct = 40;
var telemetryMock = new Mock<LocalLLMServerManager.Services.IGpuTelemetryProvider>();
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<LocalLLMServerManager.Services.IAiEngineManager>();
var orchestrator = new VramOrchestrator(
new HttpClient(handlerMock.Object),
NullLogger<VramOrchestrator>.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<HttpMessageHandler>();
handlerMock.Protected()
.Setup<Task<HttpResponseMessage>>(
"SendAsync",
ItExpr.IsAny<HttpRequestMessage>(),
ItExpr.IsAny<CancellationToken>()
)
.ReturnsAsync(() => new HttpResponseMessage(HttpStatusCode.OK)
{
Content = new StringContent("{\"models\":[]}")
});

var settingsMock = new Mock<LocalLLMServerManager.Shared.Interfaces.ISettingsService>();
settingsMock.Setup(s => s.LoadSettings()).Returns(new AppSettings(
ModelCacheIdleTimeoutMinutes: 15,
EngineHibernationTimeoutMinutes: 30,
RamPressureThresholdPercent: 75
));

var telemetryMock = new Mock<LocalLLMServerManager.Services.IGpuTelemetryProvider>();
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<LocalLLMServerManager.Services.IAiEngineManager>();
engineMock.Setup(e => e.ComfyProcess).Returns(() => mockComfyProcess);
engineMock.Setup(e => e.StopComfyUiAsync(It.IsAny<Microsoft.Extensions.Logging.ILogger>())).ReturnsAsync(true);
engineMock.Setup(e => e.StopForgeAsync(It.IsAny<Microsoft.Extensions.Logging.ILogger>())).ReturnsAsync(true);
engineMock.Setup(e => e.StopAudioEngineAsync(It.IsAny<Microsoft.Extensions.Logging.ILogger>())).ReturnsAsync(true);

var orchestrator = new VramOrchestrator(
new HttpClient(handlerMock.Object),
NullLogger<VramOrchestrator>.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<Microsoft.Extensions.Logging.ILogger>()), 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<Microsoft.Extensions.Logging.ILogger>()), 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]
Expand All @@ -132,7 +306,7 @@ public async Task EnsureVramForImageGenerationAsync_ExecutesCleanly()
ItExpr.IsAny<HttpRequestMessage>(),
ItExpr.IsAny<CancellationToken>()
)
.ReturnsAsync(new HttpResponseMessage { StatusCode = HttpStatusCode.OK });
.ReturnsAsync(() => new HttpResponseMessage { StatusCode = HttpStatusCode.OK });

var client = new HttpClient(handlerMock.Object);
var orchestrator = new VramOrchestrator(client, NullLogger<VramOrchestrator>.Instance);
Expand All @@ -150,7 +324,7 @@ public async Task EnsureVramForComfyUiAsync_ExecutesCleanly()
ItExpr.IsAny<HttpRequestMessage>(),
ItExpr.IsAny<CancellationToken>()
)
.ReturnsAsync(new HttpResponseMessage { StatusCode = HttpStatusCode.OK });
.ReturnsAsync(() => new HttpResponseMessage { StatusCode = HttpStatusCode.OK });

var client = new HttpClient(handlerMock.Object);
var orchestrator = new VramOrchestrator(client, NullLogger<VramOrchestrator>.Instance);
Expand Down
Loading
Loading