diff --git a/src/Ngino.Client/LlamaCppManager.cs b/src/Ngino.Client/LlamaCppManager.cs index 2aea578..eccdbe3 100644 --- a/src/Ngino.Client/LlamaCppManager.cs +++ b/src/Ngino.Client/LlamaCppManager.cs @@ -1,5 +1,7 @@ using System.Collections.Concurrent; using System.Diagnostics; +using System.Net.Http; +using System.Net.Sockets; using System.Text.Json; using System.Text.RegularExpressions; using Microsoft.Extensions.Logging; @@ -13,6 +15,7 @@ internal sealed partial class LlamaCppManager : IAsyncDisposable private const int DefaultBasePort = 8081; private const string NginoContainerLabel = "ngino-llamacpp"; private static readonly TimeSpan DockerTimeout = TimeSpan.FromSeconds(60); + private static readonly TimeSpan ContainerStartTimeout = TimeSpan.FromMinutes(5); private readonly string _blobsPath; private readonly string _manifestsPath; @@ -75,6 +78,43 @@ internal sealed partial class LlamaCppManager : IAsyncDisposable return _modelPorts.ContainsKey(ollamaModelName); } + public async Task IsContainerRunningAsync(string ollamaModelName) + { + if (!_modelPorts.ContainsKey(ollamaModelName)) + { + return false; + } + + var containerName = SanitizeContainerName($"ngino-llamacpp-{ollamaModelName}"); + var (exitCode, output) = await RunDockerWithOutputAsync( + ["ps", "--filter", $"name=^{containerName}$", "--format", "{{.ID}}"], + CancellationToken.None); + + var running = exitCode == 0 && !string.IsNullOrWhiteSpace(output); + if (!running) + { + _logger.LogWarning( + "llama.cpp container {ContainerName} is no longer running. Invalidating cached port for {Model}.", + containerName, ollamaModelName); + _modelPorts.TryRemove(ollamaModelName, out _); + } + + return running; + } + + public bool RemoveModelMapping(string ollamaModelName) + { + if (_modelPorts.TryRemove(ollamaModelName, out var port)) + { + _logger.LogWarning( + "Removed stale llama.cpp port mapping for {Model} (port {Port}).", + ollamaModelName, port); + return true; + } + + return false; + } + public Uri? GetUpstream(string ollamaModelName) { if (_modelPorts.TryGetValue(ollamaModelName, out var port)) @@ -136,8 +176,32 @@ internal sealed partial class LlamaCppManager : IAsyncDisposable return false; } + _logger.LogInformation( + "llama.cpp container for {Model} started on port {Port}. Waiting for it to become ready...", + ollamaName, port); + + var ready = await WaitForServerReadyAsync("localhost", port, cancellationToken); + if (!ready) + { + _logger.LogError( + "llama.cpp container for {Model} did not become ready on port {Port} within {Timeout}. Stopping it.", + ollamaName, port, ContainerStartTimeout); + + try + { + await RunDockerAsync(["stop", "--time", "10", containerName], CancellationToken.None); + await RunDockerAsync(["rm", "-f", containerName], CancellationToken.None); + } + catch (Exception cleanupException) + { + _logger.LogWarning(cleanupException, "Failed to clean up container {ContainerName}", containerName); + } + + return false; + } + _modelPorts[ollamaName] = port; - _logger.LogInformation("llama.cpp container for {Model} started on port {Port}", ollamaName, port); + _logger.LogInformation("llama.cpp container for {Model} is ready on port {Port}.", ollamaName, port); return true; } catch (Exception ex) @@ -423,6 +487,77 @@ internal sealed partial class LlamaCppManager : IAsyncDisposable return port; } + private static async Task WaitForServerReadyAsync( + string host, int port, CancellationToken cancellationToken) + { + var deadline = DateTime.UtcNow + ContainerStartTimeout; + + if (!await WaitForTcpPortAsync(host, port, deadline, cancellationToken)) + { + return false; + } + + using var handler = new SocketsHttpHandler + { + ConnectTimeout = TimeSpan.FromSeconds(3), + UseProxy = false + }; + using var httpClient = new HttpClient(handler) { Timeout = TimeSpan.FromSeconds(10) }; + + while (DateTime.UtcNow < deadline) + { + cancellationToken.ThrowIfCancellationRequested(); + + try + { + using var response = await httpClient.GetAsync( + $"http://{host}:{port}/health", cancellationToken); + if (response.IsSuccessStatusCode) + { + return true; + } + } + catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested) + { + } + catch (Exception) + { + return false; + } + + await Task.Delay(TimeSpan.FromSeconds(2), cancellationToken); + } + + return false; + } + + private static async Task WaitForTcpPortAsync( + string host, int port, DateTime deadline, CancellationToken cancellationToken) + { + while (DateTime.UtcNow < deadline) + { + cancellationToken.ThrowIfCancellationRequested(); + + try + { + using var client = new TcpClient(); + await client.ConnectAsync(host, port, cancellationToken); + return true; + } + catch (OperationCanceledException) + { + throw; + } + catch (SocketException) + { + } + + await Task.Delay(TimeSpan.FromSeconds(2), cancellationToken); + } + + return false; + } + private static string SanitizeContainerName(string name) { var sanitized = InvalidContainerNameChars().Replace(name, "_"); diff --git a/src/Ngino.Client/TunnelClient.cs b/src/Ngino.Client/TunnelClient.cs index 528295e..498401d 100644 --- a/src/Ngino.Client/TunnelClient.cs +++ b/src/Ngino.Client/TunnelClient.cs @@ -482,7 +482,7 @@ internal sealed class TunnelClient Type = TunnelMessageTypes.ModelCommandResult, RequestId = message.RequestId, StatusCode = 500, - Error = $"Failed to start llama.cpp container for model '{modelName}'." + Error = $"Unable to load model '{modelName}'." }; } @@ -694,28 +694,44 @@ internal sealed class TunnelClient if (modelName is not null) { effectiveUpstream = _llamaCppManager.GetUpstream(modelName); + if (effectiveUpstream is not null) + { + var running = await _llamaCppManager.IsContainerRunningAsync(modelName); + if (!running) + { + _logger.LogWarning( + "Cached llama.cpp container for model '{Model}' is not running anymore. Starting a fresh one on demand...", + modelName); + effectiveUpstream = null; + } + } + if (effectiveUpstream is null) { - _logger.LogInformation( - "Request for model '{Model}' but no llama.cpp container is running. Starting one on demand...", - modelName); - var model = _llamaCppManager.DiscoverModelsWithBlob() .FirstOrDefault(m => string.Equals(m.OllamaName, modelName, StringComparison.OrdinalIgnoreCase)); if (model is not null) { + _logger.LogInformation( + "Request for model '{Model}' but no llama.cpp container is running. Starting one on demand...", + modelName); + var started = await _llamaCppManager.StartModelContainerAsync(model, cancellationToken); if (started) { effectiveUpstream = _llamaCppManager.GetUpstream(modelName); } + else + { + await SendModelLoadErrorAsync(socket, message, modelName, cancellationToken); + return; + } } - - if (effectiveUpstream is null) + else { _logger.LogWarning( - "Failed to start llama.cpp container for model '{Model}'. Falling back to default upstream.", + "Model '{Model}' was not found in the Ollama models path. Falling back to default upstream.", modelName); } } @@ -766,7 +782,14 @@ internal sealed class TunnelClient effectiveUpstream: effectiveUpstream, responseHandler: responseHandler, pathTransform: pathTransform, - bodyTransform: bodyTransform); + bodyTransform: bodyTransform, + onConnectionRefused: () => + { + if (_llamaCppManager is not null && modelName is not null) + { + _llamaCppManager.RemoveModelMapping(modelName); + } + }); if (!_activeRequests.TryAdd(message.RequestId, request)) { @@ -807,6 +830,50 @@ internal sealed class TunnelClient } } + private async Task SendModelLoadErrorAsync( + ClientWebSocket socket, TunnelMessage message, string modelName, CancellationToken cancellationToken) + { + _logger.LogWarning( + "Unable to load model '{Model}' via llama.cpp. Notifying caller.", modelName); + + var body = JsonSerializer.SerializeToUtf8Bytes( + new { error = $"Unable to load model '{modelName}'" }, + JsonOptions); + + _pendingRequestBodies.TryRemove(message.RequestId, out _); + + await SendAsync( + socket, + new TunnelMessage + { + Type = TunnelMessageTypes.HttpResponseHeaders, + RequestId = message.RequestId, + StatusCode = 500, + ReasonPhrase = "Internal Server Error", + Headers = [new HeaderPair("Content-Type", "application/json")] + }, + CancellationToken.None); + + await SendAsync( + socket, + new TunnelMessage + { + Type = TunnelMessageTypes.HttpResponseBody, + RequestId = message.RequestId, + Body = body + }, + CancellationToken.None); + + await SendAsync( + socket, + new TunnelMessage + { + Type = TunnelMessageTypes.HttpResponseComplete, + RequestId = message.RequestId + }, + CancellationToken.None); + } + private void CancelAllActiveRequests() { foreach (var pair in _activeRequests.ToArray()) diff --git a/src/Ngino.Client/UpstreamRequest.cs b/src/Ngino.Client/UpstreamRequest.cs index e3bd13d..2f5f38c 100644 --- a/src/Ngino.Client/UpstreamRequest.cs +++ b/src/Ngino.Client/UpstreamRequest.cs @@ -1,4 +1,5 @@ using System.Net.Http.Headers; +using System.Net.Sockets; using System.Threading.Channels; using Ngino.Protocol; @@ -39,6 +40,7 @@ internal sealed class UpstreamRequest private readonly Func? _responseHandler; private readonly Func? _pathTransform; private readonly Func? _bodyTransform; + private readonly Action? _onConnectionRefused; private readonly List _bufferedBody = []; public UpstreamRequest( @@ -51,7 +53,8 @@ internal sealed class UpstreamRequest Uri? effectiveUpstream = null, Func? responseHandler = null, Func? pathTransform = null, - Func? bodyTransform = null) + Func? bodyTransform = null, + Action? onConnectionRefused = null) { _httpClient = httpClient; _initialMessage = initialMessage; @@ -61,6 +64,7 @@ internal sealed class UpstreamRequest _responseHandler = responseHandler; _pathTransform = pathTransform; _bodyTransform = bodyTransform; + _onConnectionRefused = onConnectionRefused; _cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); if (!initialMessage.HasBody) @@ -143,6 +147,17 @@ internal sealed class UpstreamRequest } catch (Exception exception) { + if (IsConnectionRefused(exception)) + { + try + { + _onConnectionRefused?.Invoke(); + } + catch + { + } + } + await SendErrorAsync(exception); } finally @@ -296,6 +311,25 @@ internal sealed class UpstreamRequest } } + internal static bool IsConnectionRefused(Exception exception) + { + for (var current = exception; current is not null; current = current.InnerException) + { + if (current is SocketException socketException + && socketException.SocketErrorCode == SocketError.ConnectionRefused) + { + return true; + } + + if (current.Message.Contains("Connection refused", StringComparison.OrdinalIgnoreCase)) + { + return true; + } + } + + return false; + } + private static List CollectResponseHeaders(HttpResponseMessage response) { var headers = new List(); diff --git a/src/Ngino.Server/ReverseProxyEndpoint.cs b/src/Ngino.Server/ReverseProxyEndpoint.cs index 5349194..6ddc300 100644 --- a/src/Ngino.Server/ReverseProxyEndpoint.cs +++ b/src/Ngino.Server/ReverseProxyEndpoint.cs @@ -10,6 +10,8 @@ internal static class ReverseProxyEndpoint { private const string UnauthorizedMessage = "Missing or invalid Ngino token."; + private const string OllamaVersion = "0.32.5"; + private static readonly HashSet HopByHopHeaders = new(StringComparer.OrdinalIgnoreCase) { "Connection", @@ -72,6 +74,12 @@ internal static class ReverseProxyEndpoint return; } + if (IsVersionRequest(context.Request, proxyPath)) + { + await WriteVersionResponseAsync(context); + return; + } + if (TryGetClientAddress(proxyPath, out var pathClientId, out var clientPath)) { if (!groupAccess.IsClientAllowed(pathClientId)) @@ -317,6 +325,15 @@ internal static class ReverseProxyEndpoint HttpMethods.IsGet(request.Method) && string.Equals(proxyPath.Value, "/api/tags", StringComparison.OrdinalIgnoreCase); + private static bool IsVersionRequest(HttpRequest request, PathString proxyPath) => + HttpMethods.IsGet(request.Method) + && string.Equals(proxyPath.Value, "/api/version", StringComparison.OrdinalIgnoreCase); + + private static Task WriteVersionResponseAsync(HttpContext context) => + context.Response.WriteAsJsonAsync( + new { version = OllamaVersion }, + context.RequestAborted); + private static async Task HandleTagsAsync( HttpContext context, TunnelHub hub, diff --git a/tests/Ngino.Client.Tests/UpstreamRequestTests.cs b/tests/Ngino.Client.Tests/UpstreamRequestTests.cs index cffe339..668a28a 100644 --- a/tests/Ngino.Client.Tests/UpstreamRequestTests.cs +++ b/tests/Ngino.Client.Tests/UpstreamRequestTests.cs @@ -1,3 +1,4 @@ +using System.Net.Sockets; using Ngino.Client; using Xunit; @@ -32,4 +33,29 @@ public sealed class UpstreamRequestTests Assert.Contains("origin-form path", exception.Message); } + + [Fact] + public void IsConnectionRefused_DetectsSocketConnectionRefused() + { + var socketException = new SocketException((int)SocketError.ConnectionRefused); + var exception = new HttpRequestException("Connection refused", socketException); + + Assert.True(UpstreamRequest.IsConnectionRefused(exception)); + } + + [Fact] + public void IsConnectionRefused_DetectsConnectionRefusedByMessage() + { + var exception = new HttpRequestException("Connection refused (localhost:8081)"); + + Assert.True(UpstreamRequest.IsConnectionRefused(exception)); + } + + [Fact] + public void IsConnectionRefused_IgnoresUnrelatedFailures() + { + var exception = new HttpRequestException("Connection reset by peer"); + + Assert.False(UpstreamRequest.IsConnectionRefused(exception)); + } }