From c38750169329821cc4fee34dc59db1b3d0283e4e Mon Sep 17 00:00:00 2001
From: Valeriano Manassero
 <14011549+valeriano-manassero@users.noreply.github.com>
Date: Thu, 2 Feb 2023 11:57:33 +0100
Subject: [PATCH] Added: affinity parameter

---
 .../templates/agentk8sglue-configmap.yaml         | 15 +++++++++++++++
 .../templates/agentk8sglue-deployment.yaml        |  8 ++++++++
 charts/clearml-agent/values.yaml                  | 12 +++++++++---
 3 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/charts/clearml-agent/templates/agentk8sglue-configmap.yaml b/charts/clearml-agent/templates/agentk8sglue-configmap.yaml
index 62ae17a..48ec4b6 100644
--- a/charts/clearml-agent/templates/agentk8sglue-configmap.yaml
+++ b/charts/clearml-agent/templates/agentk8sglue-configmap.yaml
@@ -172,6 +172,17 @@ data:
             {{- toYaml . | nindent 10 }}
           {{- end }}
         {{- end }}
+        {{- if $value.templateOverrides.affinity }}
+          {{- with $value.templateOverrides.affinity }}
+        affinity:
+            {{- toYaml . | nindent 10 }}
+          {{- end }}
+        {{- else if $.Values.agentk8sglue.basePodTemplate.affinity }}
+          {{- with $.Values.agentk8sglue.basePodTemplate.affinity }}
+        affinity:
+            {{- toYaml . | nindent 10 }}
+          {{- end }}
+        {{- end }}
     {{- end }}
   secrets.yaml: |
     {{- range $key, $value := $.Values.enterpriseFeatures.queues }}
@@ -250,6 +261,10 @@ data:
       tolerations:
           {{- toYaml . | nindent 8 }}
         {{- end }}
+        {{- with .Values.agentk8sglue.basePodTemplate.affinity }}
+      affinity:
+          {{- toYaml . | nindent 8 }}
+        {{- end }}
 {{- end }}
 {{- if .Values.sessions.portModeEnabled }}
 {{- range untilStep 1 ( ( add .Values.sessions.maxServices 1 ) | int ) 1 }}
diff --git a/charts/clearml-agent/templates/agentk8sglue-deployment.yaml b/charts/clearml-agent/templates/agentk8sglue-deployment.yaml
index 6ecae0d..a434871 100644
--- a/charts/clearml-agent/templates/agentk8sglue-deployment.yaml
+++ b/charts/clearml-agent/templates/agentk8sglue-deployment.yaml
@@ -177,6 +177,14 @@ spec:
       nodeSelector:
         {{- toYaml . | nindent 8 }}
       {{- end }}
+      {{- with .Values.agentk8sglue.affinity }}
+      affinity:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with .Values.agentk8sglue.tolerations }}
+      tolerations:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
       volumes:
         - name: {{ include "clearml.name" . }}-pt
           configMap:
diff --git a/charts/clearml-agent/values.yaml b/charts/clearml-agent/values.yaml
index cf33136..4b5c9b9 100644
--- a/charts/clearml-agent/values.yaml
+++ b/charts/clearml-agent/values.yaml
@@ -82,6 +82,10 @@ agentk8sglue:
   # -- nodeSelector setup for Agent pod (example in values.yaml comments)
   nodeSelector: {}
     # fleet: agent-nodes
+  # -- tolerations setup for Agent pod (example in values.yaml comments)
+  tolerations: []
+  # -- affinity setup for Agent pod (example in values.yaml comments)
+  affinity: {}
   # -- volumes definition for Glue Agent (example in values.yaml comments)
   volumes: []
     # - name: "yourvolume"
@@ -162,14 +166,16 @@ agentk8sglue:
     resources: {}
       # limits:
       #   nvidia.com/gpu: 1
+    # -- nodeSelector setup for pods spawned to consume ClearML Task (example in values.yaml comments)
+    nodeSelector: {}
+      # fleet: gpu-nodes
     # -- tolerations setup for pods spawned to consume ClearML Task (example in values.yaml comments)
     tolerations: []
       # - key: "nvidia.com/gpu"
       #   operator: Exists
       #   effect: "NoSchedule"
-    # -- nodeSelector setup for pods spawned to consume ClearML Task (example in values.yaml comments)
-    nodeSelector: {}
-      # fleet: gpu-nodes
+    # -- affinity setup for pods spawned to consume ClearML Task
+    affinity: {}
     # -- securityContext setup for pods spawned to consume ClearML Task (example in values.yaml comments)
     securityContext: {}
       #  runAsUser: 1001