From 1f43d9550c05d835dbc78234a075f2b05e44af17 Mon Sep 17 00:00:00 2001 From: Symious Date: Mon, 28 Jun 2021 17:09:23 +0800 Subject: [PATCH] YARN-10835. Pass user defined variables of yarn.nodemanager.env-whitelist along jobs --- .../v2/app/job/impl/TaskAttemptImpl.java | 11 ++++ .../hadoop/yarn/conf/YarnConfiguration.java | 5 +- .../yarn/client/api/impl/YarnClientImpl.java | 30 +++++++++ .../client/api/impl/TestYarnClientImpl.java | 61 +++++++++++++++++++ 4 files changed, 106 insertions(+), 1 deletion(-) diff --git a/hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-app/src/main/java/org/apache/hadoop/mapreduce/v2/app/job/impl/TaskAttemptImpl.java b/hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-app/src/main/java/org/apache/hadoop/mapreduce/v2/app/job/impl/TaskAttemptImpl.java index 3943a3aa913f9..f0adfb11ebdbc 100644 --- a/hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-app/src/main/java/org/apache/hadoop/mapreduce/v2/app/job/impl/TaskAttemptImpl.java +++ b/hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-app/src/main/java/org/apache/hadoop/mapreduce/v2/app/job/impl/TaskAttemptImpl.java @@ -1137,6 +1137,17 @@ static ContainerLaunchContext createContainerLaunchContext( } MapReduceChildJVM.setVMEnv(myEnv, remoteTask); + if (conf.getBoolean(YarnConfiguration.NM_ENV_WHITELIST_EXPORT_ENABLED, + YarnConfiguration.DEFAULT_NM_ENV_WHITELIST_EXPORT_ENABLED)) { + String[] whitelistVars = conf.get(YarnConfiguration.NM_ENV_WHITELIST, + YarnConfiguration.DEFAULT_NM_ENV_WHITELIST).split(","); + for (String var : whitelistVars) { + if (System.getenv(var) != null && !System.getenv(var).isEmpty()) { + myEnv.put(var, System.getenv(var)); + } + } + } + // Set up the launch command List commands = MapReduceChildJVM.getVMCommand( taskAttemptListener.getAddress(), remoteTask, jvmID); diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/conf/YarnConfiguration.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/conf/YarnConfiguration.java index 2ebf79cfae3f3..0a05b096e8afe 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/conf/YarnConfiguration.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/conf/YarnConfiguration.java @@ -1265,7 +1265,10 @@ public static boolean isAclEnabled(Configuration conf) { ApplicationConstants.Environment.HADOOP_CONF_DIR.key(), ApplicationConstants.Environment.CLASSPATH_PREPEND_DISTCACHE.key(), ApplicationConstants.Environment.HADOOP_YARN_HOME.key())); - + public static final String NM_ENV_WHITELIST_EXPORT_ENABLED = NM_ENV_WHITELIST + + "export.enabled"; + public static final boolean DEFAULT_NM_ENV_WHITELIST_EXPORT_ENABLED = false; + /** address of node manager IPC.*/ public static final String NM_ADDRESS = NM_PREFIX + "address"; public static final int DEFAULT_NM_PORT = 0; diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/impl/YarnClientImpl.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/impl/YarnClientImpl.java index 9ae966f9a126e..ce4d83aa401c3 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/impl/YarnClientImpl.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/api/impl/YarnClientImpl.java @@ -310,6 +310,20 @@ public YarnClientApplication createApplication() throw new ApplicationIdNotProvidedException( "ApplicationId is not provided in ApplicationSubmissionContext"); } + + if (isNMEnvWhitelistExportEnabled()) { + String[] whitelistVars = getConfig().get( + YarnConfiguration.NM_ENV_WHITELIST, + YarnConfiguration.DEFAULT_NM_ENV_WHITELIST).split(","); + for(String var : whitelistVars) { + String value = getNMEnvWhitelistValue(var); + if (value != null) { + appContext.getAMContainerSpec().getEnvironment() + .put(var, value); + } + } + } + SubmitApplicationRequest request = Records.newRecord(SubmitApplicationRequest.class); request.setApplicationSubmissionContext(appContext); @@ -524,6 +538,12 @@ protected boolean isSecurityEnabled() { return UserGroupInformation.isSecurityEnabled(); } + protected boolean isNMEnvWhitelistExportEnabled() { + return getConfig().getBoolean( + YarnConfiguration.NM_ENV_WHITELIST_EXPORT_ENABLED, + YarnConfiguration.DEFAULT_NM_ENV_WHITELIST_EXPORT_ENABLED); + } + @Override public void failApplicationAttempt(ApplicationAttemptId attemptId) throws YarnException, IOException { @@ -1191,4 +1211,14 @@ public void shellToContainer(ContainerId containerId, LOG.error("Fail to shell to container: " + t.getMessage()); } } + + private String getNMEnvWhitelistValue(String key) { + if (System.getenv(key) != null && !System.getenv(key).isEmpty()) { + return System.getenv(key); + } + if (System.getProperty(key) != null && !System.getProperty(key).isEmpty()) { + return System.getProperty(key); + } + return null; + } } diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestYarnClientImpl.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestYarnClientImpl.java index cded253333ea5..a7b27c802d773 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestYarnClientImpl.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestYarnClientImpl.java @@ -59,6 +59,8 @@ import java.io.IOException; import java.nio.ByteBuffer; import java.util.Collection; +import java.util.HashMap; +import java.util.Map; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; @@ -381,4 +383,63 @@ public void testParseTimelineDelegationTokenRenewer() { client.stop(); } } + + @Test + public void testNMEnvWhitelistExportEnabled() { + Configuration conf = getConf(); + conf.set(YarnConfiguration.NM_ENV_WHITELIST_EXPORT_ENABLED, "true"); + conf.set(YarnConfiguration.NM_ENV_WHITELIST, "TEST_KEY_1"); + + // Set property for TEST_KEY_1 + System.setProperty("TEST_KEY_1", "TEST_VALUE_1"); + System.setProperty("TEST_KEY_2", "TEST_VALUE_2"); + + // Prepare a Mock YarnClientImpl + YarnClientImpl client = spy(new YarnClientImpl() { + + @Override + protected void serviceStart() { + rmClient = mock(ApplicationClientProtocol.class); + } + + @Override + protected void serviceStop() { + } + + @Override + public ApplicationReport getApplicationReport(ApplicationId appId) { + ApplicationReport report = mock(ApplicationReport.class); + when(report.getYarnApplicationState()) + .thenReturn(YarnApplicationState.RUNNING); + return report; + } + }); + conf.set( + YarnConfiguration.RM_ADDRESS, "localhost:8188"); + + // Prepare a ApplicationSubmissionContext and submit the app + ApplicationSubmissionContext context = + mock(ApplicationSubmissionContext.class); + ApplicationId applicationId = ApplicationId.newInstance(0, 1); + when(context.getApplicationId()).thenReturn(applicationId); + + Map env = new HashMap<>(); + ContainerLaunchContext clc = ContainerLaunchContext.newInstance( + null, env, null, null, null, null); + when(context.getAMContainerSpec()).thenReturn(clc); + + try { + client.init(conf); + client.start(); + client.submitApplication(context); + Assert.assertEquals("TEST_VALUE_1", + context.getAMContainerSpec().getEnvironment().get("TEST_KEY_1")); + Assert.assertNotEquals("TEST_VALUE_2", + context.getAMContainerSpec().getEnvironment().get("TEST_KEY_2")); + } catch (YarnException | IOException e) { + e.printStackTrace(); + } finally { + client.stop(); + } + } }