From f6205460ffe8508a7978695e0ca6efeb080d666d Mon Sep 17 00:00:00 2001 From: Marco Primi Date: Wed, 4 Oct 2017 14:52:24 -0700 Subject: [PATCH] Make command and cluster criteria tags available to job as environment variables --- .../netflix/genie/core/jobs/JobConstants.java | 10 ++++ .../jobs/workflow/impl/InitialSetupTask.java | 60 +++++++++++++++++++ .../impl/InitialSetupTaskUnitTest.java | 37 ++++++++++++ .../linux-runsh.txt | 5 ++ .../non-linux-runsh.txt | 5 ++ 5 files changed, 117 insertions(+) diff --git a/genie-core/src/main/java/com/netflix/genie/core/jobs/JobConstants.java b/genie-core/src/main/java/com/netflix/genie/core/jobs/JobConstants.java index 9f81db5c05c..31488be935e 100644 --- a/genie-core/src/main/java/com/netflix/genie/core/jobs/JobConstants.java +++ b/genie-core/src/main/java/com/netflix/genie/core/jobs/JobConstants.java @@ -245,6 +245,16 @@ public final class JobConstants { */ public static final String GENIE_JOB_MEMORY_ENV_VAR = "GENIE_JOB_MEMORY"; + /** + * Environment variable for the Genie command tags in the job request. + */ + public static final String GENIE_REQUESTED_COMMAND_TAGS_ENV_VAR = "GENIE_REQUESTED_COMMAND_TAGS"; + + /** + * Environment variable for the Genie cluster criteria tags in the job request. + */ + public static final String GENIE_REQUESTED_CLUSTER_TAGS_ENV_VAR = "GENIE_REQUESTED_CLUSTER_TAGS"; + /** * Process ID. **/ diff --git a/genie-core/src/main/java/com/netflix/genie/core/jobs/workflow/impl/InitialSetupTask.java b/genie-core/src/main/java/com/netflix/genie/core/jobs/workflow/impl/InitialSetupTask.java index 8fa89b42d59..2b068232f42 100644 --- a/genie-core/src/main/java/com/netflix/genie/core/jobs/workflow/impl/InitialSetupTask.java +++ b/genie-core/src/main/java/com/netflix/genie/core/jobs/workflow/impl/InitialSetupTask.java @@ -19,7 +19,9 @@ import com.google.common.annotations.VisibleForTesting; import com.netflix.genie.common.dto.Cluster; +import com.netflix.genie.common.dto.ClusterCriteria; import com.netflix.genie.common.dto.Command; +import com.netflix.genie.common.dto.JobRequest; import com.netflix.genie.common.exceptions.GenieException; import com.netflix.genie.common.exceptions.GeniePreconditionException; import com.netflix.genie.common.exceptions.GenieServerException; @@ -38,6 +40,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.Comparator; +import java.util.List; import java.util.Map; import java.util.Set; import java.util.concurrent.TimeUnit; @@ -106,6 +109,9 @@ public void executeTask(@NotNull final Map context) throws Genie jobExecEnv.getMemory() ); + // create environment variables for the job request + this.createJobRequestEnvironmentVariables(writer, jobExecEnv.getJobRequest()); + //Export the Genie Version writer.write(GENIE_VERSION_EXPORT); writer.write(LINE_SEPARATOR); @@ -376,6 +382,60 @@ void createJobEnvironmentVariables( writer.write(LINE_SEPARATOR); } + @VisibleForTesting + void createJobRequestEnvironmentVariables( + final Writer writer, + final JobRequest jobRequest + ) throws IOException { + + // create environment variable for the command tags/criteria in the job request + writer.write(JobConstants.EXPORT + + JobConstants.GENIE_REQUESTED_COMMAND_TAGS_ENV_VAR + + JobConstants.EQUALS_SYMBOL + + JobConstants.DOUBLE_QUOTE_SYMBOL + + tagsToString(jobRequest.getCommandCriteria()) + + JobConstants.DOUBLE_QUOTE_SYMBOL + + LINE_SEPARATOR); + + final List clusterCriterias = jobRequest.getClusterCriterias(); + + final List clusterCriteriasStrings = new ArrayList<>(clusterCriterias.size()); + + for (ClusterCriteria clusterCriteria : clusterCriterias) { + clusterCriteriasStrings.add("[" + tagsToString(clusterCriteria.getTags()) + "]"); + } + + // Append new line + writer.write(System.lineSeparator()); + + // create environment variable for the list of cluster tags/criteria in the job request as a single + // value in the form "[[x,y,z],[a,b,c]]" + writer.write(JobConstants.EXPORT + + JobConstants.GENIE_REQUESTED_CLUSTER_TAGS_ENV_VAR + + JobConstants.EQUALS_SYMBOL + + JobConstants.DOUBLE_QUOTE_SYMBOL + + "[" + StringUtils.join(clusterCriteriasStrings, ',') + "]" + + JobConstants.DOUBLE_QUOTE_SYMBOL + + LINE_SEPARATOR); + + // create environment variables for individual tags/criteria in the job request + for (int i = 0; i < clusterCriterias.size(); i++) { + final ClusterCriteria clusterCriteria = clusterCriterias.get(i); + + // create environment variable for the job name + writer.write(JobConstants.EXPORT + + JobConstants.GENIE_REQUESTED_CLUSTER_TAGS_ENV_VAR + "_" + i + + JobConstants.EQUALS_SYMBOL + + JobConstants.DOUBLE_QUOTE_SYMBOL + + tagsToString(clusterCriteria.getTags()) + + JobConstants.DOUBLE_QUOTE_SYMBOL + + LINE_SEPARATOR); + } + + // Append new line + writer.write(LINE_SEPARATOR); + } + /** * Helper to convert a set of tags into a string that is a suitable value for a shell environment variable. * Adds double quotes as necessary (i.e. in case of spaces, newlines), performs escaping of in-tag quotes. diff --git a/genie-core/src/test/java/com/netflix/genie/core/jobs/workflow/impl/InitialSetupTaskUnitTest.java b/genie-core/src/test/java/com/netflix/genie/core/jobs/workflow/impl/InitialSetupTaskUnitTest.java index 68a002b4c1e..b64cc9bff95 100644 --- a/genie-core/src/test/java/com/netflix/genie/core/jobs/workflow/impl/InitialSetupTaskUnitTest.java +++ b/genie-core/src/test/java/com/netflix/genie/core/jobs/workflow/impl/InitialSetupTaskUnitTest.java @@ -17,8 +17,11 @@ */ package com.netflix.genie.core.jobs.workflow.impl; +import com.google.common.collect.Sets; import com.netflix.genie.common.dto.Cluster; +import com.netflix.genie.common.dto.ClusterCriteria; import com.netflix.genie.common.dto.Command; +import com.netflix.genie.common.dto.JobRequest; import com.netflix.genie.core.jobs.JobConstants; import com.netflix.genie.test.categories.UnitTest; import com.netflix.spectator.api.Registry; @@ -117,6 +120,12 @@ public void testCreateEnvironmentVariables() throws Exception { final String commandTag2 = "cmd-foo"; final String jobName = "The Job"; final int memory = 1000; + final String cmdCritTag1 = "tagX"; + final String cmdCritTag2 = "tagY"; + final String cmdCritTag3 = "tagZ"; + final String cltCritTag1 = "bar"; + final String cltCritTag2 = "baz"; + final String cltCritTag3 = "foo"; final Cluster mockCluster = Mockito.mock(Cluster.class); Mockito.when(mockCluster.getId()).thenReturn(java.util.Optional.of(clusterId)); @@ -129,12 +138,29 @@ public void testCreateEnvironmentVariables() throws Exception { Mockito.when(mockCommand.getId()).thenReturn(java.util.Optional.of(commandName)); Mockito.when(mockCommand.getTags()).thenReturn(new HashSet<>(Arrays.asList(commandTag2, commandTag1))); + final JobRequest mockJobRequest = Mockito.mock(JobRequest.class); + Mockito + .when(mockJobRequest.getCommandCriteria()) + .thenReturn( + Sets.newHashSet(Arrays.asList(cmdCritTag3, cmdCritTag2, cmdCritTag1)) + ); + Mockito + .when(mockJobRequest.getClusterCriterias()) + .thenReturn( + Arrays.asList( + new ClusterCriteria(Sets.newHashSet(cltCritTag3, cltCritTag1, cltCritTag2)), + new ClusterCriteria(Sets.newHashSet(cltCritTag3, cltCritTag1)), + new ClusterCriteria(Sets.newHashSet(cltCritTag3)) + ) + ); + final StringWriter mockWriter = new StringWriter(); this.initialSetupTask.createJobDirEnvironmentVariables(mockWriter, tempDir.getRoot().getAbsolutePath()); this.initialSetupTask.createApplicationEnvironmentVariables(mockWriter); this.initialSetupTask.createCommandEnvironmentVariables(mockWriter, mockCommand); this.initialSetupTask.createClusterEnvironmentVariables(mockWriter, mockCluster); this.initialSetupTask.createJobEnvironmentVariables(mockWriter, jobId, jobName, memory); + this.initialSetupTask.createJobRequestEnvironmentVariables(mockWriter, mockJobRequest); final String expextedOutput = "" + "export GENIE_JOB_DIR=\"" + tempDir.getRoot().getAbsolutePath() + "\"\n" @@ -162,6 +188,17 @@ public void testCreateEnvironmentVariables() throws Exception { + "export GENIE_JOB_NAME=\"" + jobName + "\"\n" + "\n" + "export GENIE_JOB_MEMORY=" + memory + "\n" + + "\n" + + "export GENIE_REQUESTED_COMMAND_TAGS=\"" + cmdCritTag1 + "," + cmdCritTag2 + "," + cmdCritTag3 + "\"\n" + + "\n" + + "export GENIE_REQUESTED_CLUSTER_TAGS=\"[" + + "[" + cltCritTag1 + "," + cltCritTag2 + "," + cltCritTag3 + "]," + + "[" + cltCritTag1 + "," + cltCritTag3 + "]," + + "[" + cltCritTag3 + "]" + + "]\"" + "\n" + + "export GENIE_REQUESTED_CLUSTER_TAGS_0=\"" + cltCritTag1 + "," + cltCritTag2 + "," + cltCritTag3 + "\"\n" + + "export GENIE_REQUESTED_CLUSTER_TAGS_1=\"" + cltCritTag1 + "," + cltCritTag3 + "\"\n" + + "export GENIE_REQUESTED_CLUSTER_TAGS_2=\"" + cltCritTag3 + "\"\n" + "\n"; Assert.assertEquals(expextedOutput, mockWriter.getString()); diff --git a/genie-web/src/test/resources/com/netflix/genie/web/controllers/JobRestControllerIntegrationTests/linux-runsh.txt b/genie-web/src/test/resources/com/netflix/genie/web/controllers/JobRestControllerIntegrationTests/linux-runsh.txt index 39a41ee64da..8a1e8bba76e 100644 --- a/genie-web/src/test/resources/com/netflix/genie/web/controllers/JobRestControllerIntegrationTests/linux-runsh.txt +++ b/genie-web/src/test/resources/com/netflix/genie/web/controllers/JobRestControllerIntegrationTests/linux-runsh.txt @@ -89,6 +89,11 @@ export GENIE_JOB_NAME="List * ... Directories bash job" export GENIE_JOB_MEMORY=1024 +export GENIE_REQUESTED_COMMAND_TAGS="bash" + +export GENIE_REQUESTED_CLUSTER_TAGS="[[localhost]]" +export GENIE_REQUESTED_CLUSTER_TAGS_0="localhost" + export GENIE_VERSION=3 # Sourcing setup file from Cluster: cluster1 diff --git a/genie-web/src/test/resources/com/netflix/genie/web/controllers/JobRestControllerIntegrationTests/non-linux-runsh.txt b/genie-web/src/test/resources/com/netflix/genie/web/controllers/JobRestControllerIntegrationTests/non-linux-runsh.txt index 39a41ee64da..8a1e8bba76e 100644 --- a/genie-web/src/test/resources/com/netflix/genie/web/controllers/JobRestControllerIntegrationTests/non-linux-runsh.txt +++ b/genie-web/src/test/resources/com/netflix/genie/web/controllers/JobRestControllerIntegrationTests/non-linux-runsh.txt @@ -89,6 +89,11 @@ export GENIE_JOB_NAME="List * ... Directories bash job" export GENIE_JOB_MEMORY=1024 +export GENIE_REQUESTED_COMMAND_TAGS="bash" + +export GENIE_REQUESTED_CLUSTER_TAGS="[[localhost]]" +export GENIE_REQUESTED_CLUSTER_TAGS_0="localhost" + export GENIE_VERSION=3 # Sourcing setup file from Cluster: cluster1