Merge pull request #753 from Project-MONAI/AC-2022

Ac 2022

Author: neildsouth

src/TaskManager/Plug-ins/Argo/ArgoPlugin.cs

@@ -434,7 +434,7 @@ private void ProcessTaskPluginArguments(Workflow workflow)
             Guard.Against.Null(workflow);
 
             var resources = Event.GetTaskPluginArgumentsParameter<Dictionary<string, string>>(Keys.ArgoResource);
-            var priorityClassName = Event.GetTaskPluginArgumentsParameter(Keys.TaskPriorityClassName);
+            var priorityClassName = Event.GetTaskPluginArgumentsParameter(Keys.TaskPriorityClassName) ?? "standard";
             var argoParameters = Event.GetTaskPluginArgumentsParameter<Dictionary<string, string>>(Keys.ArgoParameters);
 
             if (argoParameters is not null)
@@ -462,16 +462,9 @@ private void ProcessTaskPluginArguments(Workflow workflow)
                 AddRequest(resources, template, ResourcesKeys.CpuReservation);
                 AddRequest(resources, template, ResourcesKeys.MemoryReservation);
                 AddRequest(resources, template, ResourcesKeys.GpuLimit);
-
-                if (priorityClassName is not null)
-                {
-                    template.PriorityClassName = priorityClassName;
-                }
-            }
-            if (priorityClassName is not null)
-            {
-                workflow.Spec.PodPriorityClassName = priorityClassName;
+                template.PriorityClassName = priorityClassName;
             }
+            workflow.Spec.PodPriorityClassName = priorityClassName;
         }
 
         private static void AddLimit(Dictionary<string, string>? resources, Template2 template, ResourcesKey key)

src/WorkflowManager/WorkflowManager/Validators/WorkflowValidator.cs

@@ -41,6 +41,7 @@ public class WorkflowValidator
         public static readonly string Separator = ";";
         private const string Comma = ", ";
         private readonly ILogger<WorkflowValidator> _logger;
+        public static readonly string TaskPriorityClassName = "priority";
 
         /// <summary>
         /// Initializes a new instance of the <see cref="WorkflowValidator"/> class.
@@ -335,6 +336,18 @@ private void ValidateArgoTask(TaskObject currentTask)
             {
                 Errors.Add($"Task: '{currentTask.Id}' workflow_template_name must be specified{Comma}this corresponds to an Argo template name.");
             }
+
+            if (currentTask.Args.ContainsKey(TaskPriorityClassName))
+            {
+                switch (currentTask.Args[TaskPriorityClassName].ToLower())
+                {
+                    case "high" or "standard" or "low":
+                        break;
+                    default:
+                        Errors.Add($"Task: '{currentTask.Id}' TaskPriorityClassName must be one of \"high\"{Comma} \"standard\" or \"low\"");
+                        break;
+                }
+            }
         }
 
         private void ValidateClinicalReviewTask(TaskObject[] tasks, TaskObject currentTask)

tests/UnitTests/TaskManager.Argo.Tests/ArgoPluginTest.cs

@@ -806,6 +806,68 @@ public async Task ArgoPlugin_CreateArgoTemplate_Valid_json_Calls_Client()
             Times.Once);
     }
 
+    [Fact(DisplayName = "podPriorityClassName gets set if not given in workflow")]
+    public async Task ArgoPlugin_Ensures_podPriorityClassName_is_set()
+    {
+        var argoTemplate = LoadArgoTemplate("SimpleTemplate.yml");
+        Assert.NotNull(argoTemplate);
+
+        SetUpSimpleArgoWorkFlow(argoTemplate);
+
+        var message = GenerateTaskDispatchEventWithValidArguments();
+
+        WorkflowCreateRequest? requestMade = default;
+        ArgoClient.Setup(a => a.Argo_CreateWorkflowAsync(It.IsAny<string>(), It.IsAny<WorkflowCreateRequest>(), It.IsAny<CancellationToken>()))
+         .Callback<string, WorkflowCreateRequest, CancellationToken>((name, request, token) =>
+             {
+                 requestMade = request;
+             });
+
+        var defaultClassName = "standard";
+        var runner = new ArgoPlugin(ServiceScopeFactory.Object, _logger.Object, Options, message);
+        var result = await runner.ExecuteTask(CancellationToken.None).ConfigureAwait(false);
+
+        Assert.NotNull(requestMade);
+        Assert.Equal(defaultClassName, requestMade.Workflow.Spec.PodPriorityClassName);
+
+        foreach (var template in requestMade.Workflow.Spec.Templates)
+        {
+            Assert.Equal(defaultClassName, template.PriorityClassName);
+        }
+    }
+
+    [Fact(DisplayName = "podPriorityClassName gets set if given in workflow")]
+    public async Task ArgoPlugin_Ensures_podPriorityClassName_is_set_as_given()
+    {
+        var argoTemplate = LoadArgoTemplate("SimpleTemplate.yml");
+        Assert.NotNull(argoTemplate);
+
+        SetUpSimpleArgoWorkFlow(argoTemplate);
+
+        var message = GenerateTaskDispatchEventWithValidArguments();
+
+        var givenClassName = "fred";
+        message.TaskPluginArguments.Add("priority", givenClassName);
+
+        WorkflowCreateRequest? requestMade = default;
+        ArgoClient.Setup(a => a.Argo_CreateWorkflowAsync(It.IsAny<string>(), It.IsAny<WorkflowCreateRequest>(), It.IsAny<CancellationToken>()))
+         .Callback<string, WorkflowCreateRequest, CancellationToken>((name, request, token) =>
+         {
+             requestMade = request;
+         });
+
+        var runner = new ArgoPlugin(ServiceScopeFactory.Object, _logger.Object, Options, message);
+        var result = await runner.ExecuteTask(CancellationToken.None).ConfigureAwait(false);
+
+        Assert.NotNull(requestMade);
+        Assert.Equal(givenClassName, requestMade.Workflow.Spec.PodPriorityClassName);
+
+        foreach (var template in requestMade.Workflow.Spec.Templates)
+        {
+            Assert.Equal(givenClassName, template.PriorityClassName);
+        }
+    }
+
     private void SetUpSimpleArgoWorkFlow(WorkflowTemplate argoTemplate)
     {
         Assert.NotNull(argoTemplate);

tests/UnitTests/WorkflowManager.Tests/Validators/WorkflowValidatorTests.cs

@@ -17,6 +17,7 @@
 using System;
 using System.Security.Cryptography.Xml;
 using System.Threading.Tasks;
+using Amazon.Runtime.Internal.Transform;
 using Microsoft.Extensions.Logging;
 using Monai.Deploy.WorkflowManager.Common.Interfaces;
 using Monai.Deploy.WorkflowManager.Contracts.Models;
@@ -535,5 +536,97 @@ public async Task ValidateWorkflow_ValidateWorkflow_ReturnsNoErrors()
                 Assert.True(errors.Count == 0);
             }
         }
+
+        [Fact]
+        public async Task ValidateWorkflow_Incorrect_podPriorityClassName_ReturnsErrors()
+        {
+            var workflow = new Workflow
+            {
+                Name = "Workflowname1",
+                Description = "Workflowdesc1",
+                Version = "1",
+                InformaticsGateway = new InformaticsGateway
+                {
+                    AeTitle = "aetitle",
+                    ExportDestinations = new string[] { "oneDestination", "twoDestination", "threeDestination" }
+                },
+                Tasks = new TaskObject[]
+                {
+                    new TaskObject
+                    {
+                        Args = new System.Collections.Generic.Dictionary<string, string>{
+                            { "priority" ,"god" },
+                            { "workflow_template_name" ,"spot"}
+                        },
+                        Id = "rootTask",
+                        Type = "argo",
+                        Description = "TestDesc",
+                        Artifacts = new ArtifactMap
+                        {
+                            Input = new Artifact[]{
+                                new Artifact
+                                {
+                                    Name = "non_unique_artifact",
+                                    Value = "Example Value"
+                                }
+                            }
+                        }
+                    },
+                }
+            };
+
+            _workflowService.Setup(w => w.GetByNameAsync(It.IsAny<string>()))
+                .ReturnsAsync(null, TimeSpan.FromSeconds(.1));
+
+            var errors = await _workflowValidator.ValidateWorkflow(workflow);
+
+            Assert.Single(errors);
+        }
+
+        [Fact]
+        public async Task ValidateWorkflow_correct_podPriorityClassName_ReturnsNoErrors()
+        {
+            var workflow = new Workflow
+            {
+                Name = "Workflowname1",
+                Description = "Workflowdesc1",
+                Version = "1",
+                InformaticsGateway = new InformaticsGateway
+                {
+                    AeTitle = "aetitle",
+                    ExportDestinations = new string[] { "oneDestination", "twoDestination", "threeDestination" }
+                },
+                Tasks = new TaskObject[]
+                {
+                    new TaskObject
+                    {
+                        Args = new System.Collections.Generic.Dictionary<string, string>{
+                            { "priority" ,"high" },
+                            { "workflow_template_name" ,"spot"}
+                        },
+                        Id = "rootTask",
+                        Type = "argo",
+                        Description = "TestDesc",
+                        Artifacts = new ArtifactMap
+                        {
+                            Input = new Artifact[]{
+                                new Artifact
+                                {
+                                    Name = "non_unique_artifact",
+                                    Value = "Example Value"
+                                }
+                            }
+                        }
+                    },
+                }
+            };
+
+            _workflowService.Setup(w => w.GetByNameAsync(It.IsAny<string>()))
+                .ReturnsAsync(null, TimeSpan.FromSeconds(.1));
+
+            var errors = await _workflowValidator.ValidateWorkflow(workflow);
+
+            Assert.Empty(errors);
+        }
     }
 }