Passed shared file folder to ClearML NMT job (#157)

sillsdev · Jan 16, 2024 · c843ccd · c843ccd
1 parent d941cd2
commit c843ccd
Show file tree

Hide file tree

Showing 2 changed files with 99 additions and 1 deletion.
diff --git a/src/SIL.Machine.AspNetCore/Services/NmtClearMLBuildJobFactory.cs b/src/SIL.Machine.AspNetCore/Services/NmtClearMLBuildJobFactory.cs
@@ -37,14 +37,18 @@ public async Task<string> CreateJobScriptAsync(
             if (engine is null)
                 throw new InvalidOperationException("The engine does not exist.");
 
+            Uri sharedFileUri = _sharedFileService.GetBaseUri();
+            string baseUri = sharedFileUri.GetComponents(UriComponents.SchemeAndServer, UriFormat.Unescaped);
+            string folder = sharedFileUri.GetComponents(UriComponents.Path, UriFormat.Unescaped);
             return "from machine.jobs.build_nmt_engine import run\n"
                 + "args = {\n"
                 + $"    'model_type': '{_options.CurrentValue.ModelType}',\n"
                 + $"    'engine_id': '{engineId}',\n"
                 + $"    'build_id': '{buildId}',\n"
                 + $"    'src_lang': '{_languageTagService.ConvertToFlores200Code(engine.SourceLanguage)}',\n"
                 + $"    'trg_lang': '{_languageTagService.ConvertToFlores200Code(engine.TargetLanguage)}',\n"
-                + $"    'shared_file_uri': '{_sharedFileService.GetBaseUri()}',\n"
+                + $"    'shared_file_uri': '{baseUri}',\n"
+                + $"    'shared_file_folder': '{folder}',\n"
                 + (buildOptions is not null ? $"    'build_options': '''{buildOptions}''',\n" : "")
                 + $"    'clearml': True\n"
                 + "}\n"

diff --git a/tests/SIL.Machine.AspNetCore.Tests/Services/NmtClearMLBuildJobFactoryTests.cs b/tests/SIL.Machine.AspNetCore.Tests/Services/NmtClearMLBuildJobFactoryTests.cs
@@ -0,0 +1,94 @@
+namespace SIL.Machine.AspNetCore.Services;
+
+[TestFixture]
+public class NmtClearMLBuildJobFactoryTests
+{
+    [Test]
+    public async Task CreateJobScriptAsync_BuildOptions()
+    {
+        var env = new TestEnvironment();
+        string script = await env.BuildJobFactory.CreateJobScriptAsync(
+            "engine1",
+            "build1",
+            NmtBuildStages.Train,
+            buildOptions: "{ \"max_steps\": \"10\" }"
+        );
+        Assert.That(
+            script,
+            Is.EqualTo(
+                @"from machine.jobs.build_nmt_engine import run
+args = {
+    'model_type': 'test_model',
+    'engine_id': 'engine1',
+    'build_id': 'build1',
+    'src_lang': 'spa_Latn',
+    'trg_lang': 'eng_Latn',
+    'shared_file_uri': 's3://bucket',
+    'shared_file_folder': 'folder1/folder2',
+    'build_options': '''{ ""max_steps"": ""10"" }''',
+    'clearml': True
+}
+run(args)
+".ReplaceLineEndings("\n")
+            )
+        );
+    }
+
+    [Test]
+    public async Task CreateJobScriptAsync_NoBuildOptions()
+    {
+        var env = new TestEnvironment();
+        string script = await env.BuildJobFactory.CreateJobScriptAsync("engine1", "build1", NmtBuildStages.Train);
+        Assert.That(
+            script,
+            Is.EqualTo(
+                @"from machine.jobs.build_nmt_engine import run
+args = {
+    'model_type': 'test_model',
+    'engine_id': 'engine1',
+    'build_id': 'build1',
+    'src_lang': 'spa_Latn',
+    'trg_lang': 'eng_Latn',
+    'shared_file_uri': 's3://bucket',
+    'shared_file_folder': 'folder1/folder2',
+    'clearml': True
+}
+run(args)
+".ReplaceLineEndings("\n")
+            )
+        );
+    }
+
+    private class TestEnvironment
+    {
+        public ISharedFileService SharedFileService { get; }
+        public MemoryRepository<TranslationEngine> Engines { get; }
+        public IOptionsMonitor<ClearMLOptions> Options { get; }
+        public ILanguageTagService LanguageTagService { get; }
+        public NmtClearMLBuildJobFactory BuildJobFactory { get; }
+
+        public TestEnvironment()
+        {
+            Engines = new MemoryRepository<TranslationEngine>();
+            Engines.Add(
+                new TranslationEngine
+                {
+                    Id = "engine1",
+                    EngineId = "engine1",
+                    SourceLanguage = "es",
+                    TargetLanguage = "en",
+                    BuildRevision = 1,
+                    CurrentBuild = new Build { BuildId = "build1", JobState = BuildJobState.Pending }
+                }
+            );
+            Options = Substitute.For<IOptionsMonitor<ClearMLOptions>>();
+            Options.CurrentValue.Returns(new ClearMLOptions { ModelType = "test_model" });
+            SharedFileService = Substitute.For<ISharedFileService>();
+            SharedFileService.GetBaseUri().Returns(new Uri("s3://bucket/folder1/folder2"));
+            LanguageTagService = Substitute.For<ILanguageTagService>();
+            LanguageTagService.ConvertToFlores200Code("es").Returns("spa_Latn");
+            LanguageTagService.ConvertToFlores200Code("en").Returns("eng_Latn");
+            BuildJobFactory = new NmtClearMLBuildJobFactory(SharedFileService, LanguageTagService, Engines, Options);
+        }
+    }
+}