Move model saving to engine definition

Fix model cleaning hanfire storage
sillsdev · Feb 5, 2024 · 54c2a35 · 54c2a35
1 parent 86910b3
commit 54c2a35
Show file tree

Hide file tree

Showing 19 changed files with 134 additions and 81 deletions.
diff --git a/src/SIL.Machine.AspNetCore/Configuration/IMachineBuilderExtensions.cs b/src/SIL.Machine.AspNetCore/Configuration/IMachineBuilderExtensions.cs
@@ -151,6 +151,21 @@ private static IMachineBuilder AddHangfireBuildJobRunner(this IMachineBuilder bu
         return builder;
     }
 
+    private static MongoStorageOptions GetMongoStorageOptions()
+    {
+        var mongoStorageOptions = new MongoStorageOptions
+        {
+            MigrationOptions = new MongoMigrationOptions
+            {
+                MigrationStrategy = new MigrateMongoMigrationStrategy(),
+                BackupStrategy = new CollectionMongoBackupStrategy()
+            },
+            CheckConnection = true,
+            CheckQueuedJobsStrategy = CheckQueuedJobsStrategy.TailNotificationsCollection,
+        };
+        return mongoStorageOptions;
+    }
+
     public static IMachineBuilder AddMongoHangfireJobClient(
         this IMachineBuilder builder,
         string? connectionString = null
@@ -164,19 +179,7 @@ public static IMachineBuilder AddMongoHangfireJobClient(
             c.SetDataCompatibilityLevel(CompatibilityLevel.Version_170)
                 .UseSimpleAssemblyNameTypeSerializer()
                 .UseRecommendedSerializerSettings()
-                .UseMongoStorage(
-                    connectionString,
-                    new MongoStorageOptions
-                    {
-                        MigrationOptions = new MongoMigrationOptions
-                        {
-                            MigrationStrategy = new MigrateMongoMigrationStrategy(),
-                            BackupStrategy = new CollectionMongoBackupStrategy()
-                        },
-                        CheckConnection = true,
-                        CheckQueuedJobsStrategy = CheckQueuedJobsStrategy.TailNotificationsCollection,
-                    }
-                )
+                .UseMongoStorage(connectionString, GetMongoStorageOptions())
                 .UseFilter(new AutomaticRetryAttribute { Attempts = 0 })
         );
         builder.Services.AddHealthChecks().AddCheck<HangfireHealthCheck>(name: "Hangfire");
@@ -208,7 +211,7 @@ public static IMachineBuilder AddHangfireJobServer(
 
         builder.Services.AddHangfireServer(o =>
         {
-            o.Queues = queues.ToArray();
+            o.Queues = [.. queues];
         });
         return builder;
     }
@@ -402,6 +405,20 @@ public static IMachineBuilder AddBuildJobService(this IMachineBuilder builder)
         return builder;
     }
 
+    public static IMachineBuilder AddModelCleanupJob(this IMachineBuilder builder, string? connectionString = null)
+    {
+        connectionString ??= builder.Configuration?.GetConnectionString("Hangfire");
+        if (connectionString is null)
+            throw new InvalidOperationException("Hangfire connection string is required");
+
+        var mongoClientSettings = MongoClientSettings.FromUrl(new MongoUrl(connectionString));
+        JobStorage.Current = new MongoStorage(mongoClientSettings, "recurring_job", GetMongoStorageOptions());
+        builder.Services.AddSingleton<ICleanupOldModelsJob, CleanupOldModelsJob>();
+        RecurringJobOptions options = new() { TimeZone = TimeZoneInfo.Utc };
+        RecurringJob.AddOrUpdate<ICleanupOldModelsJob>("Cleanup-job", x => x.RunAsync(), Cron.Daily, options);
+        return builder;
+    }
+
     private static IMachineBuilder AddBuildJobService(this IMachineBuilder builder, BuildJobOptions options)
     {
         builder.Services.AddScoped<IBuildJobService, BuildJobService>();

diff --git a/src/SIL.Machine.AspNetCore/Models/TranslationEngine.cs b/src/SIL.Machine.AspNetCore/Models/TranslationEngine.cs
@@ -7,6 +7,7 @@ public class TranslationEngine : IEntity
     public string EngineId { get; set; } = default!;
     public string SourceLanguage { get; set; } = default!;
     public string TargetLanguage { get; set; } = default!;
+    public bool IsModelRetrievable { get; set; } = false;
     public int BuildRevision { get; set; }
     public Build? CurrentBuild { get; set; }
 }
diff --git a/src/SIL.Machine.AspNetCore/Services/CleanupJob.cs b/src/SIL.Machine.AspNetCore/Services/CleanupJob.cs
@@ -6,11 +6,9 @@ public class CleanupOldModelsJob(ISharedFileService sharedFileService) : ICleanu
     private List<string> _filesPreviouslyMarkedForDeletion = [];
     private readonly List<string> _filesNewlyMarkedForDeletion = [];
 
-    private static readonly string MODEL_DIRECTORY = "models/";
-
     public async Task RunAsync()
     {
-        var files = await SharedFileService.ListFilesAsync(MODEL_DIRECTORY);
+        var files = await SharedFileService.ListFilesAsync(ISharedFileService.ModelDirectory);
         // split name by underscore into engineID and buildRevision
         Dictionary<string, int> modelsByEngineId = [];
         foreach (string file in files)

diff --git a/src/SIL.Machine.AspNetCore/Services/HangfireHealthCheck.cs b/src/SIL.Machine.AspNetCore/Services/HangfireHealthCheck.cs
@@ -1,15 +1,9 @@
 namespace SIL.Machine.AspNetCore.Services;
 
-public class HangfireHealthCheck : IHealthCheck
+public class HangfireHealthCheck(JobStorage jobStorage, IOptions<BackgroundJobServerOptions> options) : IHealthCheck
 {
-    private readonly JobStorage _jobStorage;
-    private readonly IOptions<BackgroundJobServerOptions> _options;
-
-    public HangfireHealthCheck(JobStorage jobStorage, IOptions<BackgroundJobServerOptions> options)
-    {
-        _jobStorage = jobStorage;
-        _options = options;
-    }
+    private readonly JobStorage _jobStorage = jobStorage;
+    private readonly IOptions<BackgroundJobServerOptions> _options = options;
 
     public Task<HealthCheckResult> CheckHealthAsync(
         HealthCheckContext context,

diff --git a/src/SIL.Machine.AspNetCore/Services/IFileStorage.cs b/src/SIL.Machine.AspNetCore/Services/IFileStorage.cs
@@ -14,7 +14,7 @@ Task<IReadOnlyCollection<string>> ListFilesAsync(
 
     Task<Stream> OpenWriteAsync(string path, CancellationToken cancellationToken = default);
 
-    Task<string> GetPresignedUrlAsync(string path, CancellationToken cancellationToken = default);
+    Task<string> GetPresignedUrlAsync(string path, int minutesToExpire, CancellationToken cancellationToken = default);
 
     Task DeleteAsync(string path, bool recurse = false, CancellationToken cancellationToken = default);
 }
diff --git a/src/SIL.Machine.AspNetCore/Services/ISharedFileService.cs b/src/SIL.Machine.AspNetCore/Services/ISharedFileService.cs
@@ -2,11 +2,13 @@
 
 public interface ISharedFileService
 {
+    public const string ModelDirectory = "models/";
+
     Uri GetBaseUri();
 
     Uri GetResolvedUri(string path);
 
-    Task<Uri> GetPresignedUrlAsync(string path);
+    Task<Uri> GetPresignedUrlAsync(string path, int minutesToExpire);
 
     Task<IReadOnlyCollection<string>> ListFilesAsync(
         string path,

diff --git a/src/SIL.Machine.AspNetCore/Services/ITranslationEngineService.cs b/src/SIL.Machine.AspNetCore/Services/ITranslationEngineService.cs
@@ -9,6 +9,7 @@ Task CreateAsync(
         string? engineName,
         string sourceLanguage,
         string targetLanguage,
+        bool isModelRetrievable = false,
         CancellationToken cancellationToken = default
     );
     Task DeleteAsync(string engineId, CancellationToken cancellationToken = default);

diff --git a/src/SIL.Machine.AspNetCore/Services/InMemoryStorage.cs b/src/SIL.Machine.AspNetCore/Services/InMemoryStorage.cs
@@ -96,7 +96,11 @@ public Task<IReadOnlyCollection<string>> ListFilesAsync(
         );
     }
 
-    public Task<string> GetPresignedUrlAsync(string path, CancellationToken cancellationToken = default)
+    public Task<string> GetPresignedUrlAsync(
+        string path,
+        int minutesToExpire,
+        CancellationToken cancellationToken = default
+    )
     {
         return Task.FromResult(path);
     }

diff --git a/src/SIL.Machine.AspNetCore/Services/LocalStorage.cs b/src/SIL.Machine.AspNetCore/Services/LocalStorage.cs
@@ -36,7 +36,11 @@ public Task<IReadOnlyCollection<string>> ListFilesAsync(
         );
     }
 
-    public Task<string> GetPresignedUrlAsync(string path, CancellationToken cancellationToken = default)
+    public Task<string> GetPresignedUrlAsync(
+        string path,
+        int minutesToExpire,
+        CancellationToken cancellationToken = default
+    )
     {
         return Task.FromResult(path);
     }

diff --git a/src/SIL.Machine.AspNetCore/Services/NmtClearMLBuildJobFactory.cs b/src/SIL.Machine.AspNetCore/Services/NmtClearMLBuildJobFactory.cs
@@ -52,6 +52,7 @@ public async Task<string> CreateJobScriptAsync(
                 + $"    'shared_file_uri': '{baseUri}',\n"
                 + $"    'shared_file_folder': '{folder}',\n"
                 + (buildOptions is not null ? $"    'build_options': '''{buildOptions}''',\n" : "")
+                + (engine.IsModelRetrievable ? $"    'save_model': '{engineId}_{engine.BuildRevision + 1}',\n" : "")
                 + $"    'clearml': True\n"
                 + "}\n"
                 + "run(args)\n";

diff --git a/src/SIL.Machine.AspNetCore/Services/NmtEngineService.cs b/src/SIL.Machine.AspNetCore/Services/NmtEngineService.cs
@@ -29,11 +29,14 @@ ISharedFileService sharedFileService
 
     public TranslationEngineType Type => TranslationEngineType.Nmt;
 
+    private const int MinutesToExpire = 60;
+
     public async Task CreateAsync(
         string engineId,
         string? engineName,
         string sourceLanguage,
         string targetLanguage,
+        bool isModelRetrievable = false,
         CancellationToken cancellationToken = default
     )
     {
@@ -43,12 +46,13 @@ await _engines.InsertAsync(
             {
                 EngineId = engineId,
                 SourceLanguage = sourceLanguage,
-                TargetLanguage = targetLanguage
+                TargetLanguage = targetLanguage,
+                IsModelRetrievable = isModelRetrievable
             },
             cancellationToken
         );
         await _buildJobService.CreateEngineAsync(
-            new[] { BuildJobType.Cpu, BuildJobType.Gpu },
+            [BuildJobType.Cpu, BuildJobType.Gpu],
             engineId,
             engineName,
             cancellationToken
@@ -115,16 +119,35 @@ public async Task<ModelPresignedUrl> GetModelPresignedUrlAsync(
         CancellationToken cancellationToken = default
     )
     {
-        var files = await _sharedFileService.ListFilesAsync($"models/", cancellationToken: cancellationToken);
-        // find latest file that start with the engineId
-        var latestFile = files.Where(f => f.StartsWith(engineId)).OrderByDescending(f => f).FirstOrDefault();
-        if (latestFile is null)
-            throw new FileNotFoundException("No built, saved model found for engine.", engineId);
-        string buildRevision = latestFile.Split('_').Last();
+        TranslationEngine engine = await GetEngineAsync(engineId, cancellationToken);
+        if (!engine.IsModelRetrievable)
+            throw new InvalidOperationException(
+                "The model cannot be downloaded. "
+                    + "To enable downloading the model, recreate the engine with IsModelRetrievable property to true."
+            );
+        if (engine.BuildRevision == 0)
+            throw new InvalidOperationException("The engine has not been built yet.");
+        string filename = $"{engineId}_{engine.BuildRevision}.tar.gz";
+        bool fileExists = await _sharedFileService.ExistsAsync(
+            ISharedFileService.ModelDirectory + filename,
+            cancellationToken
+        );
+        if (!fileExists)
+            throw new FileNotFoundException(
+                $"The model should exist to be downloaded but is not there for BuildRevision {engine.BuildRevision}."
+            );
         var modelInfo = new ModelPresignedUrl
         {
-            PresignedUrl = (await _sharedFileService.GetPresignedUrlAsync($"models/{latestFile}")).ToString(),
-            BuildRevision = buildRevision
+            PresignedUrl = (
+                await _sharedFileService.GetPresignedUrlAsync(
+                    ISharedFileService.ModelDirectory + filename,
+                    MinutesToExpire
+                )
+            ).ToString(),
+            BuildRevision = engine.BuildRevision,
+            UrlExpirationTime = DateTime
+                .UtcNow.AddMinutes(MinutesToExpire)
+                .ToString("yyyy-MM-ddTHH\\:mm\\:ss.fffffffzzz", CultureInfo.InvariantCulture)
         };
         return modelInfo;
     }
@@ -178,4 +201,12 @@ private async Task CancelBuildJobAsync(string engineId, CancellationToken cancel
         if (buildId is not null && jobState is BuildJobState.None)
             await _platformService.BuildCanceledAsync(buildId, CancellationToken.None);
     }
+
+    private async Task<TranslationEngine> GetEngineAsync(string engineId, CancellationToken cancellationToken)
+    {
+        TranslationEngine? engine = await _engines.GetAsync(e => e.EngineId == engineId, cancellationToken);
+        if (engine is null)
+            throw new InvalidOperationException($"The engine {engineId} does not exist.");
+        return engine;
+    }
 }
diff --git a/src/SIL.Machine.AspNetCore/Services/NmtTrainBuildJob.cs b/src/SIL.Machine.AspNetCore/Services/NmtTrainBuildJob.cs
@@ -68,6 +68,11 @@ await PipInstallModuleAsync(
                         + $"    'trg_lang': '{ConvertLanguageTag(engine.TargetLanguage)}',\n"
                         + $"    'shared_file_uri': '{_sharedFileService.GetBaseUri()}',\n"
                         + (buildOptions is not null ? $"    'build_options': '''{buildOptions}''',\n" : "")
+                        + (
+                            engine.IsModelRetrievable
+                                ? $"    'save_model': '{engine.Id}_{engine.BuildRevision + 1}',\n"
+                                : ""
+                        )
                         + $"    'clearml': False\n"
                         + "}\n"
                         + "run(args)\n"

diff --git a/src/SIL.Machine.AspNetCore/Services/S3FileStorage.cs b/src/SIL.Machine.AspNetCore/Services/S3FileStorage.cs
@@ -65,15 +65,19 @@ public async Task<IReadOnlyCollection<string>> ListFilesAsync(
         return response.S3Objects.Select(s3Obj => s3Obj.Key[_basePath.Length..]).ToList();
     }
 
-    public Task<string> GetPresignedUrlAsync(string path, CancellationToken cancellationToken = default)
+    public Task<string> GetPresignedUrlAsync(
+        string path,
+        int minutesToExpire,
+        CancellationToken cancellationToken = default
+    )
     {
         return Task.FromResult(
             _client.GetPreSignedURL(
                 new GetPreSignedUrlRequest
                 {
                     BucketName = _bucketName,
                     Key = _basePath + Normalize(path),
-                    Expires = DateTime.UtcNow.AddMinutes(60)
+                    Expires = DateTime.UtcNow.AddMinutes(minutesToExpire)
                 }
             )
         );

diff --git a/src/SIL.Machine.AspNetCore/Services/ServalTranslationEngineServiceV1.cs b/src/SIL.Machine.AspNetCore/Services/ServalTranslationEngineServiceV1.cs
@@ -23,6 +23,7 @@ await engineService.CreateAsync(
             request.HasEngineName ? request.EngineName : null,
             request.SourceLanguage,
             request.TargetLanguage,
+            request.IsModelRetrievable,
             context.CancellationToken
         );
         return Empty;
@@ -132,7 +133,8 @@ ServerCallContext context
         return new GetModelPresignedUrlResponse
         {
             PresignedUrl = modelPresignedUrl.PresignedUrl,
-            BuildRevision = modelPresignedUrl.BuildRevision
+            BuildRevision = modelPresignedUrl.BuildRevision,
+            UrlExpirationTime = modelPresignedUrl.UrlExpirationTime
         };
     }
 

diff --git a/src/SIL.Machine.AspNetCore/Services/SharedFileService.cs b/src/SIL.Machine.AspNetCore/Services/SharedFileService.cs
@@ -55,12 +55,12 @@ public Uri GetResolvedUri(string path)
         return new Uri(_baseUri, path);
     }
 
-    public async Task<Uri> GetPresignedUrlAsync(string path)
+    public async Task<Uri> GetPresignedUrlAsync(string path, int minutesToExpire)
     {
         string presignedUrl = path;
         if (_baseUri is not null)
             if (_baseUri.Scheme == "s3")
-                presignedUrl = await _fileStorage.GetPresignedUrlAsync(path);
+                presignedUrl = await _fileStorage.GetPresignedUrlAsync(path, minutesToExpire);
         var url = GetResolvedUri(presignedUrl);
         return url;
     }