Adding some retry logic to the lifecycle hook

The hook is run before the container image is downloaded (and ollama started) when you run from a completely clean session (like GitHub Actions), and this poses a problem because it would timeout trying to check if the model exists. Added a really simple bit of retry logic in to counter that

Author: aaronpowell

src/CommunityToolkit.Aspire.Hosting.Ollama/OllamaResourceLifecycleHook.cs

@@ -1,17 +1,28 @@
-using Aspire.Hosting.ApplicationModel;
+using Aspire.Hosting;
+using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.Lifecycle;
+using Microsoft.Extensions.Logging;
 using OllamaSharp;
 using OllamaSharp.Models;
+using System.Globalization;
 
 namespace CommunityToolkit.Aspire.Hosting.Ollama;
-internal class OllamaResourceLifecycleHook(ResourceNotificationService notificationService) : IDistributedApplicationLifecycleHook, IAsyncDisposable
+internal class OllamaResourceLifecycleHook(
+    ResourceLoggerService loggerService,
+    ResourceNotificationService notificationService,
+    DistributedApplicationExecutionContext context) : IDistributedApplicationLifecycleHook, IAsyncDisposable
 {
     private readonly ResourceNotificationService _notificationService = notificationService;
 
     private readonly CancellationTokenSource _tokenSource = new();
 
     public Task AfterResourcesCreatedAsync(DistributedApplicationModel appModel, CancellationToken cancellationToken = default)
     {
+        if (context.IsPublishMode)
+        {
+            return Task.CompletedTask;
+        }
+
         foreach (var resource in appModel.Resources.OfType<OllamaResource>())
         {
             DownloadModel(resource, _tokenSource.Token);
@@ -27,11 +38,13 @@ private void DownloadModel(OllamaResource resource, CancellationToken cancellati
             return;
         }
 
+        var logger = loggerService.GetLogger(resource);
+
         _ = Task.Run(async () =>
         {
             try
             {
-                var connectionString = await resource.ConnectionStringExpression.GetValueAsync(cancellationToken);
+                var connectionString = await resource.ConnectionStringExpression.GetValueAsync(cancellationToken).ConfigureAwait(false);
 
                 if (string.IsNullOrWhiteSpace(connectionString))
                 {
@@ -47,7 +60,18 @@ private void DownloadModel(OllamaResource resource, CancellationToken cancellati
 
                 if (!hasModel)
                 {
-                    await PullModel(resource, ollamaClient, model, cancellationToken);
+                    logger.LogInformation("{TimeStamp}: [{Model}] needs to be downloaded for {ResourceName}",
+                        DateTime.UtcNow.ToString("yyyy-MM-ddTHH:mm:ss.fffZ", CultureInfo.InvariantCulture),
+                        resource.ModelName,
+                        resource.Name);
+                    await PullModel(resource, ollamaClient, model, logger, cancellationToken);
+                }
+                else
+                {
+                    logger.LogInformation("{TimeStamp}: [{Model}] already exists for {ResourceName}",
+                        DateTime.UtcNow.ToString("yyyy-MM-ddTHH:mm:ss.fffZ", CultureInfo.InvariantCulture),
+                        resource.ModelName,
+                        resource.Name);
                 }
 
                 await _notificationService.PublishUpdateAsync(resource, state => state with { State = new ResourceStateSnapshot("Running", KnownResourceStateStyles.Success) });
@@ -57,24 +81,42 @@ private void DownloadModel(OllamaResource resource, CancellationToken cancellati
                 await _notificationService.PublishUpdateAsync(resource, state => state with { State = new ResourceStateSnapshot(ex.Message, KnownResourceStateStyles.Error) });
             }
 
-        }, cancellationToken);
+        }, cancellationToken).ConfigureAwait(false);
     }
 
     private static async Task<bool> HasModelAsync(OllamaApiClient ollamaClient, string model, CancellationToken cancellationToken)
     {
-        var localModels = await ollamaClient.ListLocalModels(cancellationToken);
-        return localModels.Any(m => m.Name.StartsWith(model));
+        int retryCount = 0;
+        while (retryCount < 5)
+        {
+            try
+            {
+                var localModels = await ollamaClient.ListLocalModels(cancellationToken);
+                return localModels.Any(m => m.Name.StartsWith(model));
+            }
+            catch (TaskCanceledException)
+            {
+                // wait 30 seconds before retrying
+                await Task.Delay(TimeSpan.FromSeconds(30), cancellationToken);
+                retryCount++;
+            }
+        }
+
+        throw new TimeoutException("Failed to list local models after 5 retries. Likely that the container image was not pulled in time, or the container is not running.");
     }
 
-    private async Task PullModel(OllamaResource resource, OllamaApiClient ollamaClient, string model, CancellationToken cancellationToken)
+    private async Task PullModel(OllamaResource resource, OllamaApiClient ollamaClient, string model, ILogger logger, CancellationToken cancellationToken)
     {
+        logger.LogInformation("{TimeStamp}: Pulling ollama model {Model}...",
+            DateTime.UtcNow.ToString("yyyy-MM-ddTHH:mm:ss.fffZ", CultureInfo.InvariantCulture),
+            model);
         await _notificationService.PublishUpdateAsync(resource, state => state with { State = new ResourceStateSnapshot("Downloading model", KnownResourceStateStyles.Info) });
 
         long percentage = 0;
 
         await foreach (PullModelResponse? status in ollamaClient.PullModel(model, cancellationToken))
         {
-            if (status == null)
+            if (status is null)
             {
                 continue;
             }
@@ -95,6 +137,10 @@ await _notificationService.PublishUpdateAsync(resource,
                 }
             }
         }
+
+        logger.LogInformation("{TimeStamp}: Finished pulling ollama model {Model}",
+            DateTime.UtcNow.ToString("yyyy-MM-ddTHH:mm:ss.fffZ", CultureInfo.InvariantCulture),
+            model);
     }
 
     public ValueTask DisposeAsync()