feat: Allow tolerations for ephemeral agents (#1579)

* feat: Allow tolerations for ephemeral agents
AzBuilder · Nov 27, 2024 · 60a7f83 · 60a7f83
1 parent 234cdf2
commit 60a7f83
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/...g/terrakube/api/plugin/scheduler/job/tcl/executor/ephemeral/EphemeralExecutorService.java b/...g/terrakube/api/plugin/scheduler/job/tcl/executor/ephemeral/EphemeralExecutorService.java
@@ -18,6 +18,7 @@
 public class EphemeralExecutorService {
 
     private static final String NODE_SELECTOR = "EPHEMERAL_CONFIG_NODE_SELECTOR_TAGS";
+    private static final String TOLERATIONS = "EPHEMERAL_CONFIG_TOLERATIONS";
     private static final String SERVICE_ACCOUNT = "EPHEMERAL_CONFIG_SERVICE_ACCOUNT";
     private static final String ANNOTATIONS = "EPHEMERAL_CONFIG_ANNOTATIONS";
     private static final String CONFIG_MAP_NAME = "EPHEMERAL_CONFIG_MAP_NAME";
@@ -67,6 +68,23 @@ public ExecutorContext sendToEphemeralExecutor(Job job, ExecutorContext executor
             nodeSelectorInfo = ephemeralConfiguration.getNodeSelector();
         }
 
+        Optional<String> tolerationsInfo = Optional.ofNullable(
+                executorContext.getEnvironmentVariables().getOrDefault(TOLERATIONS, null));
+        List<Toleration> tolerations = new ArrayList<>();
+
+        if (tolerationsInfo.isPresent()) {
+            for (String tolerationData : tolerationsInfo.get().split(";")) {
+                String[] info = tolerationData.split(":");
+                Toleration toleration = new Toleration();
+
+                toleration.setKey(info[0]);
+                toleration.setOperator(info.length > 1 ? info[1] : "Exists");
+                toleration.setEffect(info.length > 2 ? info[2] : null);
+
+                tolerations.add(toleration);
+            }
+        }
+
         Optional<String> annotationsInfo = Optional.ofNullable(executorContext.getEnvironmentVariables().getOrDefault(ANNOTATIONS, null));
         Map<String, String> annotations = new HashMap();
         log.info("Custom Annotations: {}", annotationsInfo.isPresent());
@@ -121,6 +139,7 @@ public ExecutorContext sendToEphemeralExecutor(Job job, ExecutorContext executor
                 .withNewSpec()
                 .withNodeSelector(nodeSelectorInfo)
                 .withServiceAccountName(serviceAccount)
+                .withTolerations(tolerations)
                 .withVolumes(volumes)
                 .addNewContainer()
                 .withName(jobName)