Add additional job attributes support to Titus Launcher

Netflix · Feb 23, 2021 · 8242dbe · 8242dbe
1 parent 33a9369
commit 8242dbe
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 6 deletions.
diff --git a/genie-docs/src/docs/asciidoc/_properties.adoc b/genie-docs/src/docs/asciidoc/_properties.adoc
@@ -172,6 +172,11 @@ https://docs.spring.io/spring-boot/docs/current/reference/html/spring-boot-featu
 |0
 |yes
 
+|genie.agent.launcher.titus.additional-job-attributes
+|Map of attributes to send with the Titus request in addition to whatever defaults there are
+|empty
+|yes
+
 |genie.agent.launcher.titus.additional-memory
 |An additional amount of memory that should be added to whatever the job originally requested
 |2GB

diff --git a/...-web/src/main/java/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImpl.java b/...-web/src/main/java/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImpl.java
@@ -267,6 +267,20 @@ private TitusBatchJobRequest createJobRequest(final ResolvedJob resolvedJob) {
         );
         final Duration runtimeLimit = this.titusAgentLauncherProperties.getRuntimeLimit();
 
+        final Map<String, String> jobAttributes = new HashMap<>();
+        jobAttributes.put(GENIE_USER_ATTR, resolvedJob.getJobMetadata().getUser());
+        jobAttributes.put(GENIE_SOURCE_HOST_ATTR, this.genieHostInfo.getHostname());
+        jobAttributes.put(GENIE_ENDPOINT_ATTR, this.titusAgentLauncherProperties.getGenieServerHost());
+        jobAttributes.put(GENIE_JOB_ID_ATTR, jobId);
+        jobAttributes.putAll(
+            this.binder
+                .bind(
+                    TitusAgentLauncherProperties.ADDITIONAL_JOB_ATTRIBUTES_PROPERTY,
+                    Bindable.mapOf(String.class, String.class)
+                )
+                .orElse(new HashMap<>())
+        );
+
         return new TitusBatchJobRequest(
             new TitusBatchJobRequest.Owner(this.titusAgentLauncherProperties.getOwnerEmail()),
             this.titusAgentLauncherProperties.getApplicationName(),
@@ -275,12 +289,7 @@ private TitusBatchJobRequest createJobRequest(final ResolvedJob resolvedJob) {
                 String.class,
                 this.titusAgentLauncherProperties.getCapacityGroup()
             ),
-            ImmutableMap.of(
-                GENIE_USER_ATTR, resolvedJob.getJobMetadata().getUser(),
-                GENIE_SOURCE_HOST_ATTR, this.genieHostInfo.getHostname(),
-                GENIE_ENDPOINT_ATTR, this.titusAgentLauncherProperties.getGenieServerHost(),
-                GENIE_JOB_ID_ATTR, jobId
-            ),
+            jobAttributes,
             new TitusBatchJobRequest.Container(
                 new TitusBatchJobRequest.Resources(
                     cpu,

diff --git a/genie-web/src/main/java/com/netflix/genie/web/properties/TitusAgentLauncherProperties.java b/genie-web/src/main/java/com/netflix/genie/web/properties/TitusAgentLauncherProperties.java
@@ -79,6 +79,11 @@ public class TitusAgentLauncherProperties {
      */
     public static final String ADDITIONAL_GPU_PROPERTY = PREFIX + ".additionalGPU";
 
+    /**
+     * Any additional job attributes that should be added to defaults.
+     */
+    public static final String ADDITIONAL_JOB_ATTRIBUTES_PROPERTY = PREFIX + ".additional-job-attributes";
+
     /**
      * An additional amount of memory that should be added to whatever the job originally requested.
      */
@@ -339,4 +344,10 @@ public class TitusAgentLauncherProperties {
      */
     @NotNull
     private Map<String, String> containerAttributes = new HashMap<>();
+
+    /**
+     * Additional job attributes.
+     */
+    @NotNull
+    private Map<String, String> additionalJobAttributes = new HashMap<>();
 }
diff --git a/.../src/test/groovy/com/netflix/genie/web/properties/TitusAgentLauncherPropertiesSpec.groovy b/.../src/test/groovy/com/netflix/genie/web/properties/TitusAgentLauncherPropertiesSpec.groovy
@@ -66,6 +66,7 @@ class TitusAgentLauncherPropertiesSpec extends Specification {
         p.getMinimumMemory().toGigabytes() == 4
         p.getMinimumGPU() == 0
         p.getContainerAttributes() == [:]
+        p.getAdditionalJobAttributes() == [:]
 
         when:
         p.setEnabled(true)
@@ -105,6 +106,7 @@ class TitusAgentLauncherPropertiesSpec extends Specification {
         p.setMinimumMemory(DataSize.ofGigabytes(8))
         p.setMinimumGPU(1)
         p.setContainerAttributes(["hi": "bye"])
+        p.setAdditionalJobAttributes(["new": "attribute"])
 
         then:
         p.isEnabled()
@@ -145,5 +147,6 @@ class TitusAgentLauncherPropertiesSpec extends Specification {
         p.getMinimumMemory().toGigabytes() == 8
         p.getMinimumGPU() == 1
         p.getContainerAttributes() == ["hi": "bye"]
+        p.getAdditionalJobAttributes() == ["new": "attribute"]
     }
 }