AC-1243 Added configurable resource constraints on argo message sending pod containers

woodheadio · woodheadio · commit d36b557fe58f · 2023-02-06T12:12:16.000Z
Signed-off-by: Alex Woodhead &lt;11213454+woodheadio@users.noreply.github.com&gt;
diff --git a/docs/changelog.md b/docs/changelog.md
@@ -22,3 +22,5 @@ Renamed the (Generated) Argo client to ArgoGeneratedClient, added a new ArgoClie
 Enhanced the ArgoClient -> Argo_Get_WorkflowLogsAsync method to decode the json better and make the logs extracted from Argo more readable.
 
 Added Mongo Migraions, to allow changes without breaking current stored entries
+
+Added resource constraints to the generated ARGO templates, so these can operate within a Kubernetes cluster with a Resource Quota in the ARGO namespace.
diff --git a/src/Shared/Configuration/TaskManagerConfiguration.cs b/src/Shared/Configuration/TaskManagerConfiguration.cs
@@ -50,5 +50,29 @@ public class ArgoPluginArguments
     {
         [ConfigurationKeyName("server_url")]
         public string ServerUrl { get; set; } = string.Empty;
+
+        [ConfigurationKeyName("initContainerCpuLimit")]
+        public string InitContainerCpuLimit { get; set; } = "1";
+
+        [ConfigurationKeyName("initContainerMemoryLimit")]
+        public string InitContainerMemoryLimit { get; set; } = "500Mi";
+
+        [ConfigurationKeyName("waitContainerCpuLimit")]
+        public string WaitContainerCpuLimit { get; set; } = "1";
+
+        [ConfigurationKeyName("waitContainerMemoryLimit")]
+        public string WaitContainerMemoryLimit { get; set; } = "500Mi";
+
+        [ConfigurationKeyName("messageGeneratorContainerCpuLimit")]
+        public string MessageGeneratorContainerCpuLimit { get; set; } = "1";
+
+        [ConfigurationKeyName("messageGeneratorContainerMemoryLimit")]
+        public string MessageGeneratorContainerMemoryLimit { get; set; } = "500Mi";
+
+        [ConfigurationKeyName("messageSenderContainerCpuLimit")]
+        public string MessageSenderContainerCpuLimit { get; set; } = "1";
+
+        [ConfigurationKeyName("messageSenderContainerMemoryLimit")]
+        public string MessageSenderContainerMemoryLimit { get; set; } = "500Mi";
     }
 }
diff --git a/src/Shared/Configuration/WorkflowManagerOptions.cs b/src/Shared/Configuration/WorkflowManagerOptions.cs
@@ -60,11 +60,11 @@ public class WorkflowManagerOptions
         /// </summary>
         public BackgroundServiceSettings BackgroundServiceSettings { get; set; }
 
-        [ConfigurationKeyName("argoTtlStatergySeconds")]
-        public int ArgoTtlStatergySeconds { get; set; } = 60 * 60 * 24 * 2; // 2 days before the pods get automatically cleaned up from argo
+        [ConfigurationKeyName("argoTtlStrategySeconds")]
+        public int ArgoTtlStrategySeconds { get; set; } = 60 * 60 * 24 * 2; // 2 days before the pods get automatically cleaned up from argo
 
-        [ConfigurationKeyName("minArgoTtlStatergySeconds")]
-        public int MinArgoTtlStatergySeconds { get; set; } = 30; // time to get logs before cleanup !
+        [ConfigurationKeyName("minArgoTtlStrategySeconds")]
+        public int MinArgoTtlStrategySeconds { get; set; } = 30; // time to get logs before cleanup !
 
         public WorkflowManagerOptions()
         {
diff --git a/src/TaskManager/Plug-ins/Argo/ArgoPlugin.cs b/src/TaskManager/Plug-ins/Argo/ArgoPlugin.cs
@@ -618,16 +618,16 @@ private void SetArgoTtlStratergy(Workflow workflow)
         {
             if (workflow.Spec.TtlStrategy is null)
             {
-                workflow.Spec.TtlStrategy = new TTLStrategy { SecondsAfterCompletion = _options.Value.ArgoTtlStatergySeconds };
+                workflow.Spec.TtlStrategy = new TTLStrategy { SecondsAfterCompletion = _options.Value.ArgoTtlStrategySeconds };
             }
             else
             {
                 if (workflow.Spec.TtlStrategy.SecondsAfterCompletion.HasValue)
-                    workflow.Spec.TtlStrategy.SecondsAfterCompletion = Math.Max(_options.Value.MinArgoTtlStatergySeconds, workflow.Spec.TtlStrategy.SecondsAfterCompletion.Value);
+                    workflow.Spec.TtlStrategy.SecondsAfterCompletion = Math.Max(_options.Value.MinArgoTtlStrategySeconds, workflow.Spec.TtlStrategy.SecondsAfterCompletion.Value);
                 if (workflow.Spec.TtlStrategy.SecondsAfterSuccess.HasValue)
-                    workflow.Spec.TtlStrategy.SecondsAfterSuccess = Math.Max(_options.Value.MinArgoTtlStatergySeconds, workflow.Spec.TtlStrategy.SecondsAfterSuccess.Value);
+                    workflow.Spec.TtlStrategy.SecondsAfterSuccess = Math.Max(_options.Value.MinArgoTtlStrategySeconds, workflow.Spec.TtlStrategy.SecondsAfterSuccess.Value);
                 if (workflow.Spec.TtlStrategy.SecondsAfterFailure.HasValue)
-                    workflow.Spec.TtlStrategy.SecondsAfterFailure = Math.Max(_options.Value.MinArgoTtlStatergySeconds, workflow.Spec.TtlStrategy.SecondsAfterFailure.Value);
+                    workflow.Spec.TtlStrategy.SecondsAfterFailure = Math.Max(_options.Value.MinArgoTtlStrategySeconds, workflow.Spec.TtlStrategy.SecondsAfterFailure.Value);
             }
             RemovePodGCStratergy(workflow);
         }
diff --git a/src/TaskManager/Plug-ins/Argo/ExitHookTemplate.cs b/src/TaskManager/Plug-ins/Argo/ExitHookTemplate.cs
@@ -74,9 +74,18 @@ public Template2 GenerateMessageTemplate(S3Artifact2 artifact)
                 Container = new Container2
                 {
                     Image = Strings.ExitHookGenerateMessageContainerImage,
+                    Resources = new ResourceRequirements
+                    {
+                        Limits = new Dictionary<string, string>
+                        {
+                            {"cpu", _options.TaskManager.ArgoPluginArguments.MessageGeneratorContainerCpuLimit},
+                            {"memory", _options.TaskManager.ArgoPluginArguments.MessageGeneratorContainerMemoryLimit}
+                        }
+                    },
                     Command = new List<string> { "/bin/sh", "-c" },
                     Args = new List<string> { $"echo \"{{{{inputs.parameters.message}}}}\" > {Strings.ExitHookOutputPath}{_messageFileName}; cat {Strings.ExitHookOutputPath}{_messageFileName};" }
                 },
+                PodSpecPatch = "{\"initContainers\":[{\"name\":\"init\",\"resources\":{\"limits\":{\"cpu\":\"" + _options.TaskManager.ArgoPluginArguments.InitContainerCpuLimit + "\",\"memory\": \"" + _options.TaskManager.ArgoPluginArguments.InitContainerMemoryLimit + "\"},\"requests\":{\"cpu\":\"0\",\"memory\":\"0Mi\"}}}],\"containers\":[{\"name\":\"wait\",\"resources\":{\"limits\":{\"cpu\":\"" + _options.TaskManager.ArgoPluginArguments.WaitContainerCpuLimit + "\",\"memory\":\"" + _options.TaskManager.ArgoPluginArguments.WaitContainerMemoryLimit + "\"},\"requests\":{\"cpu\":\"0\",\"memory\":\"0Mi\"}}}]}",
                 Outputs = new Outputs
                 {
                     Artifacts = new List<Artifact>()
@@ -155,6 +164,10 @@ public Template2 GenerateSendTemplate(S3Artifact2 artifact)
                 Container = new Container2
                 {
                     Image = _options.TaskManager.ArgoExitHookSendMessageContainerImage,
+                    Resources = new ResourceRequirements
+                    {
+                        Limits = new Dictionary<string, string>{ {"cpu", _options.TaskManager.ArgoPluginArguments.MessageSenderContainerCpuLimit}, {"memory", _options.TaskManager.ArgoPluginArguments.MessageSenderContainerMemoryLimit} }
+                    },
                     Command = new List<string> { "/rabtap" },
                     Args = new List<string> {
                         "pub",
@@ -164,7 +177,8 @@ public Template2 GenerateSendTemplate(S3Artifact2 artifact)
                         "--delay=0s",
                         "--confirms",
                         "--mandatory" }
-                }
+                },
+                PodSpecPatch = "{\"initContainers\":[{\"name\":\"init\",\"resources\":{\"limits\":{\"cpu\":\"" + _options.TaskManager.ArgoPluginArguments.InitContainerCpuLimit + "\",\"memory\": \"" + _options.TaskManager.ArgoPluginArguments.InitContainerMemoryLimit + "\"},\"requests\":{\"cpu\":\"0\",\"memory\":\"0Mi\"}}}],\"containers\":[{\"name\":\"wait\",\"resources\":{\"limits\":{\"cpu\":\"" + _options.TaskManager.ArgoPluginArguments.WaitContainerCpuLimit + "\",\"memory\":\"" + _options.TaskManager.ArgoPluginArguments.WaitContainerMemoryLimit + "\"},\"requests\":{\"cpu\":\"0\",\"memory\":\"0Mi\"}}}]}",
             };
         }
     }
diff --git a/src/TaskManager/TaskManager/appsettings.json b/src/TaskManager/TaskManager/appsettings.json
@@ -5,8 +5,8 @@
   },
   "WorkflowManager": {
     "taskTimeoutMinutes": 60,
-    "argoTtlStatergySeconds": 172800, // (60*60*24*2) 2 days before the pods get automatically cleaned up from argo,
-    "minArgoTtlStatergySeconds": 30,
+    "argoTtlStrategySeconds": 172800, // (60*60*24*2) 2 days before the pods get automatically cleaned up from argo,
+    "minArgoTtlStrategySeconds": 30,
     "endpointSettings": {
       "defaultPageSize": 10,
       "maxPageSize": 1000000
@@ -25,7 +25,15 @@
         "test": "Monai.Deploy.WorkflowManager.TaskManager.TestPlugin.Repositories.TestPluginRepository, Monai.Deploy.WorkflowManager.TaskManager.TestPlugin"
       },
       "argoPluginArguments": {
-        "server_url": "http://argo-argo-workflows-server.argo:2746"
+        "server_url": "http://argo-argo-workflows-server.argo:2746",
+        "initContainerCpuLimit": "1",
+        "initContainerMemoryLimit": "500Mi",
+        "waitContainerCpuLimit": "1",
+        "waitContainerMemoryLimit": "500Mi",
+        "messageGeneratorContainerCpuLimit": "1",
+        "messageGeneratorContainerMemoryLimit": "500Mi",
+        "messageSenderContainerCpuLimit": "1",
+        "messageSenderContainerMemoryLimit": "500Mi"
       },
       "argoExitHookSendMessageContainerImage": "ghcr.io/jandelgado/rabtap:latest"
     },
diff --git a/tests/UnitTests/TaskManager.Argo.Tests/ArgoPluginTest.cs b/tests/UnitTests/TaskManager.Argo.Tests/ArgoPluginTest.cs
@@ -60,6 +60,14 @@ public class ArgoPluginTest
     private Workflow? _submittedArgoTemplate;
     private readonly int _argoTtlStatergySeconds = 360;
     private readonly int _minAgoTtlStatergySeconds = 30;
+    private readonly string _initContainerCpuLimit = "100m";
+    private readonly string _initContainerMemoryLimit = "200Mi";
+    private readonly string _waitContainerCpuLimit = "200m";
+    private readonly string _waitContainerMemoryLimit = "300Mi";
+    private readonly string _messageGeneratorContainerCpuLimit = "300m";
+    private readonly string _messageGeneratorContainerMemoryLimit = "400Mi";
+    private readonly string _messageSenderContainerCpuLimit = "400m";
+    private readonly string _messageSenderContainerMemoryLimit = "500Mi";
 
     public ArgoPluginTest()
     {
@@ -80,8 +88,16 @@ public ArgoPluginTest()
         _options.Value.Messaging.PublisherSettings.Add("exchange", "exchange");
         _options.Value.Messaging.PublisherSettings.Add("virtualHost", "vhost");
         _options.Value.Messaging.Topics.TaskCallbackRequest = "md.tasks.callback";
-        _options.Value.ArgoTtlStatergySeconds = _argoTtlStatergySeconds;
-        _options.Value.MinArgoTtlStatergySeconds = _minAgoTtlStatergySeconds;
+        _options.Value.ArgoTtlStrategySeconds = _argoTtlStatergySeconds;
+        _options.Value.MinArgoTtlStrategySeconds = _minAgoTtlStatergySeconds;
+        _options.Value.TaskManager.ArgoPluginArguments.InitContainerCpuLimit = _initContainerCpuLimit;
+        _options.Value.TaskManager.ArgoPluginArguments.InitContainerMemoryLimit = _initContainerMemoryLimit;
+        _options.Value.TaskManager.ArgoPluginArguments.WaitContainerCpuLimit = _waitContainerCpuLimit;
+        _options.Value.TaskManager.ArgoPluginArguments.WaitContainerMemoryLimit = _waitContainerMemoryLimit;
+        _options.Value.TaskManager.ArgoPluginArguments.MessageGeneratorContainerCpuLimit = _messageGeneratorContainerCpuLimit;
+        _options.Value.TaskManager.ArgoPluginArguments.MessageGeneratorContainerMemoryLimit = _messageGeneratorContainerMemoryLimit;
+        _options.Value.TaskManager.ArgoPluginArguments.MessageSenderContainerCpuLimit = _messageSenderContainerCpuLimit;
+        _options.Value.TaskManager.ArgoPluginArguments.MessageSenderContainerMemoryLimit = _messageSenderContainerMemoryLimit;
 
         _serviceScopeFactory.Setup(p => p.CreateScope()).Returns(_serviceScope.Object);
 
@@ -630,8 +646,8 @@ public async Task ArgoPlugin_Copies_ImagePullSecrets()
         Assert.Equal(secret, _submittedArgoTemplate?.Spec.ImagePullSecrets.First());
     }
 
-    [Fact(DisplayName = "TTL gets added if not pressent")]
-    public async Task ArgoPlugin_Ensures_TTL_Added_If_Not_pressent()
+    [Fact(DisplayName = "TTL gets added if not present")]
+    public async Task ArgoPlugin_Ensures_TTL_Added_If_Not_present()
     {
         var argoTemplate = LoadArgoTemplate("SimpleTemplate.yml");
         Assert.NotNull(argoTemplate);
@@ -647,6 +663,34 @@ public async Task ArgoPlugin_Ensures_TTL_Added_If_Not_pressent()
         Assert.Equal(_argoTtlStatergySeconds, _submittedArgoTemplate?.Spec.TtlStrategy?.SecondsAfterCompletion);
     }
 
+    [Fact(DisplayName = "Argo Plugin adds required resource limits")]
+    public async Task ArgoPlugin_Adds_Container_Resource_Restrictions_Based_On_Configured_Values()
+    {
+        var argoTemplate = LoadArgoTemplate("SimpleTemplate.yml");
+        Assert.NotNull(argoTemplate);
+
+        SetUpSimpleArgoWorkFlow(argoTemplate);
+
+        var message = GenerateTaskDispatchEventWithValidArguments();
+
+        var expectedPodSpecPatch = "{\"initContainers\":[{\"name\":\"init\",\"resources\":{\"limits\":{\"cpu\":\"" + _initContainerCpuLimit + "\",\"memory\": \"" +
+                                                   _initContainerMemoryLimit +
+                                                   "\"},\"requests\":{\"cpu\":\"0\",\"memory\":\"0Mi\"}}}],\"containers\":[{\"name\":\"wait\",\"resources\":{\"limits\":{\"cpu\":\"" +
+                                                   _waitContainerCpuLimit + "\",\"memory\":\"" + _waitContainerMemoryLimit +
+                                                   "\"},\"requests\":{\"cpu\":\"0\",\"memory\":\"0Mi\"}}}]}";
+
+        var runner = new ArgoPlugin(_serviceScopeFactory.Object, _logger.Object, _options, message);
+        var result = await runner.ExecuteTask(CancellationToken.None).ConfigureAwait(false);
+
+        Assert.Equal(TaskExecutionStatus.Accepted, result.Status);
+        Assert.Equal(_messageGeneratorContainerCpuLimit, _submittedArgoTemplate?.Spec.Templates.FirstOrDefault(p => p.Name == Strings.ExitHookTemplateGenerateTemplateName).Container.Resources.Limits["cpu"]);
+        Assert.Equal(_messageGeneratorContainerMemoryLimit, _submittedArgoTemplate?.Spec.Templates.FirstOrDefault(p => p.Name == Strings.ExitHookTemplateGenerateTemplateName).Container.Resources.Limits["memory"]);
+        Assert.Equal(expectedPodSpecPatch, _submittedArgoTemplate?.Spec.Templates.FirstOrDefault(p => p.Name == Strings.ExitHookTemplateGenerateTemplateName).PodSpecPatch);
+        Assert.Equal(_messageSenderContainerCpuLimit, _submittedArgoTemplate?.Spec.Templates.FirstOrDefault(p => p.Name == Strings.ExitHookTemplateSendTemplateName).Container.Resources.Limits["cpu"]);
+        Assert.Equal(_messageSenderContainerMemoryLimit, _submittedArgoTemplate?.Spec.Templates.FirstOrDefault(p => p.Name == Strings.ExitHookTemplateSendTemplateName).Container.Resources.Limits["memory"]);
+        Assert.Equal(expectedPodSpecPatch, _submittedArgoTemplate?.Spec.Templates.FirstOrDefault(p => p.Name == Strings.ExitHookTemplateSendTemplateName).PodSpecPatch);
+    }
+
     [Theory(DisplayName = "TTL gets extended if too short")]
     [InlineData(31, 31, 29)]
     [InlineData(1, null, null)]
@@ -725,7 +769,7 @@ public async Task ArgoPlugin_Ensures_TTL_Remains(int? secondsAfterCompletion, in
         Assert.Equal(secondsAfterFailure, _submittedArgoTemplate?.Spec.TtlStrategy.SecondsAfterFailure);
     }
 
-    [Fact(DisplayName = "pocGC gets removed if pressent")]
+    [Fact(DisplayName = "pocGC gets removed if present")]
     public async Task ArgoPlugin_Ensures_podGC_is_removed()
     {
         var argoTemplate = LoadArgoTemplate("SimpleTemplate.yml");

Original file line number	Diff line number	Diff line change
`@@ -618,16 +618,16 @@ private void SetArgoTtlStratergy(Workflow workflow)`
`618`	`618`	`{`
`619`	`619`	`if (workflow.Spec.TtlStrategy is null)`
`620`	`620`	`{`
`621`		`- workflow.Spec.TtlStrategy = new TTLStrategy { SecondsAfterCompletion = _options.Value.ArgoTtlStatergySeconds };`
	`621`	`+ workflow.Spec.TtlStrategy = new TTLStrategy { SecondsAfterCompletion = _options.Value.ArgoTtlStrategySeconds };`
`622`	`622`	`}`
`623`	`623`	`else`
`624`	`624`	`{`
`625`	`625`	`if (workflow.Spec.TtlStrategy.SecondsAfterCompletion.HasValue)`
`626`		`- workflow.Spec.TtlStrategy.SecondsAfterCompletion = Math.Max(_options.Value.MinArgoTtlStatergySeconds, workflow.Spec.TtlStrategy.SecondsAfterCompletion.Value);`
	`626`	`+ workflow.Spec.TtlStrategy.SecondsAfterCompletion = Math.Max(_options.Value.MinArgoTtlStrategySeconds, workflow.Spec.TtlStrategy.SecondsAfterCompletion.Value);`
`627`	`627`	`if (workflow.Spec.TtlStrategy.SecondsAfterSuccess.HasValue)`
`628`		`- workflow.Spec.TtlStrategy.SecondsAfterSuccess = Math.Max(_options.Value.MinArgoTtlStatergySeconds, workflow.Spec.TtlStrategy.SecondsAfterSuccess.Value);`
	`628`	`+ workflow.Spec.TtlStrategy.SecondsAfterSuccess = Math.Max(_options.Value.MinArgoTtlStrategySeconds, workflow.Spec.TtlStrategy.SecondsAfterSuccess.Value);`
`629`	`629`	`if (workflow.Spec.TtlStrategy.SecondsAfterFailure.HasValue)`
`630`		`- workflow.Spec.TtlStrategy.SecondsAfterFailure = Math.Max(_options.Value.MinArgoTtlStatergySeconds, workflow.Spec.TtlStrategy.SecondsAfterFailure.Value);`
	`630`	`+ workflow.Spec.TtlStrategy.SecondsAfterFailure = Math.Max(_options.Value.MinArgoTtlStrategySeconds, workflow.Spec.TtlStrategy.SecondsAfterFailure.Value);`
`631`	`631`	`}`
`632`	`632`	`RemovePodGCStratergy(workflow);`
`633`	`633`	`}`