From f30d1092bc5e7275328b113a37e52f65785c9697 Mon Sep 17 00:00:00 2001 From: Eric Badger Date: Mon, 8 Mar 2021 20:20:52 +0000 Subject: [PATCH] YARN-10664. Allow parameter expansion in NM_ADMIN_USER_ENV. Contributed by Jim Brennan. --- .../launcher/ContainerLaunch.java | 20 ++-- .../launcher/TestContainerLaunch.java | 112 +++++++++++++++++- 2 files changed, 117 insertions(+), 15 deletions(-) diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/ContainerLaunch.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/ContainerLaunch.java index 84cc79e3d3e..15dddc0337b 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/ContainerLaunch.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/ContainerLaunch.java @@ -175,15 +175,13 @@ public class ContainerLaunch implements Callable { return var; } - private Map expandAllEnvironmentVars( - ContainerLaunchContext launchContext, Path containerLogDir) { - Map environment = launchContext.getEnvironment(); + private void expandAllEnvironmentVars( + Map environment, Path containerLogDir) { for (Entry entry : environment.entrySet()) { String value = entry.getValue(); value = expandEnvironment(value, containerLogDir); entry.setValue(value); } - return environment; } @Override @@ -218,8 +216,10 @@ public class ContainerLaunch implements Callable { } launchContext.setCommands(newCmds); - Map environment = expandAllEnvironmentVars( - launchContext, containerLogDir); + // The actual expansion of environment variables happens after calling + // sanitizeEnv. This allows variables specified in NM_ADMIN_USER_ENV + // to reference user or container-defined variables. + Map environment = launchContext.getEnvironment(); // /////////////////////////// End of variable expansion // Use this to track variables that are added to the environment by nm. @@ -283,6 +283,8 @@ public class ContainerLaunch implements Callable { containerLogDirs, localResources, nmPrivateClasspathJarDir, nmEnvVars); + expandAllEnvironmentVars(environment, containerLogDir); + prepareContainer(localResources, containerLocalDirs); // Write out the environment @@ -1564,13 +1566,13 @@ public class ContainerLaunch implements Callable { } // variables here will be forced in, even if the container has - // specified them. + // specified them. Note: we do not track these in nmVars, to + // allow them to be ordered properly if they reference variables + // defined by the user. String defEnvStr = conf.get(YarnConfiguration.DEFAULT_NM_ADMIN_USER_ENV); Apps.setEnvFromInputProperty(environment, YarnConfiguration.NM_ADMIN_USER_ENV, defEnvStr, conf, File.pathSeparator); - nmVars.addAll(Apps.getEnvVarsFromInputProperty( - YarnConfiguration.NM_ADMIN_USER_ENV, defEnvStr, conf)); if (!Shell.WINDOWS) { // maybe force path components diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/TestContainerLaunch.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/TestContainerLaunch.java index d700850fd28..0f856da8634 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/TestContainerLaunch.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/launcher/TestContainerLaunch.java @@ -25,6 +25,7 @@ import static org.mockito.Matchers.any; import static org.mockito.Mockito.*; import java.io.BufferedReader; +import java.io.DataOutputStream; import java.io.File; import java.io.FileOutputStream; import java.io.FileReader; @@ -55,9 +56,8 @@ import java.util.function.Supplier; import com.google.common.collect.Lists; import org.apache.commons.codec.binary.Base64; import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.FileUtil; -import org.apache.hadoop.fs.Path; -import org.apache.hadoop.fs.UnsupportedFileSystemException; +import org.apache.hadoop.fs.*; +import org.apache.hadoop.io.IOUtils; import org.apache.hadoop.security.Credentials; import org.apache.hadoop.security.token.SecretManager.InvalidToken; import org.apache.hadoop.test.GenericTestUtils; @@ -124,6 +124,8 @@ import org.junit.Assume; import org.junit.Before; import org.junit.Test; import org.mockito.ArgumentCaptor; +import org.mockito.invocation.InvocationOnMock; +import org.mockito.stubbing.Answer; public class TestContainerLaunch extends BaseContainerManagerTest { @@ -664,7 +666,7 @@ public class TestContainerLaunch extends BaseContainerManagerTest { Container container = mock(Container.class); when(container.getContainerId()).thenReturn(cId); when(container.getLaunchContext()).thenReturn(containerLaunchContext); - when(container.getLocalizedResources()).thenReturn(null); + when(container.localizationCountersAsString()).thenReturn("1,2,3,4,5"); Dispatcher dispatcher = mock(Dispatcher.class); EventHandler eventHandler = new EventHandler() { public void handle(Event event) { @@ -805,8 +807,6 @@ public class TestContainerLaunch extends BaseContainerManagerTest { Assert.assertTrue(userSetEnv.containsKey(testKey1)); Assert.assertTrue(userSetEnv.containsKey(testKey2)); Assert.assertTrue(userSetEnv.containsKey(testKey3)); - Assert.assertTrue(nmEnvTrack.contains("MALLOC_ARENA_MAX")); - Assert.assertTrue(nmEnvTrack.contains("MOUNT_LIST")); Assert.assertEquals(userMallocArenaMaxVal + File.pathSeparator + mallocArenaMaxVal, userSetEnv.get("MALLOC_ARENA_MAX")); Assert.assertEquals(testVal1, userSetEnv.get(testKey1)); @@ -1848,6 +1848,7 @@ public class TestContainerLaunch extends BaseContainerManagerTest { when(id.toString()).thenReturn("1"); when(container.getContainerId()).thenReturn(id); when(container.getUser()).thenReturn("user"); + when(container.localizationCountersAsString()).thenReturn("1,2,3,4,5"); ContainerLaunchContext clc = mock(ContainerLaunchContext.class); when(clc.getCommands()).thenReturn(Lists.newArrayList()); when(container.getLaunchContext()).thenReturn(clc); @@ -2444,6 +2445,7 @@ public class TestContainerLaunch extends BaseContainerManagerTest { .newContainerId(ApplicationAttemptId.newInstance(appId, 1), 1); when(container.getContainerId()).thenReturn(containerId); when(container.getUser()).thenReturn("test"); + when(container.localizationCountersAsString()).thenReturn("1,2,3,4,5"); when(container.getLocalizedResources()) .thenReturn(Collections.> emptyMap()); @@ -2512,4 +2514,102 @@ public class TestContainerLaunch extends BaseContainerManagerTest { launch.getUserFilecacheDirs(localDirsForRead)), StringUtils.join(",", ctx.getUserFilecacheDirs())); } + + private String readStringFromPath(Path p) throws IOException { + FileSystem fs = FileSystem.get(conf); + try (FSDataInputStream is = fs.open(p)) { + byte[] bytes = IOUtils.readFullyToByteArray(is); + return new String(bytes); + } + } + + @Test(timeout = 20000) + public void testExpandNmAdmEnv() throws Exception { + // setup mocks + Dispatcher dispatcher = mock(Dispatcher.class); + EventHandler handler = mock(EventHandler.class); + when(dispatcher.getEventHandler()).thenReturn(handler); + ContainerExecutor containerExecutor = mock(ContainerExecutor.class); + doAnswer(new Answer() { + @Override + public Void answer(InvocationOnMock invocation) throws Throwable { + Object[] args = invocation.getArguments(); + DataOutputStream dos = (DataOutputStream) args[0]; + dos.writeBytes("script"); + return null; + } + }).when(containerExecutor).writeLaunchEnv( + any(), any(), any(), any(), any(), any(), any()); + Application app = mock(Application.class); + ApplicationId appId = mock(ApplicationId.class); + when(appId.toString()).thenReturn("1"); + when(app.getAppId()).thenReturn(appId); + Container container = mock(Container.class); + ContainerId id = mock(ContainerId.class); + when(id.toString()).thenReturn("1"); + when(container.getContainerId()).thenReturn(id); + when(container.getUser()).thenReturn("user"); + ContainerLaunchContext clc = mock(ContainerLaunchContext.class); + when(clc.getCommands()).thenReturn(Lists.newArrayList()); + when(container.getLaunchContext()).thenReturn(clc); + Credentials credentials = mock(Credentials.class); + when(container.getCredentials()).thenReturn(credentials); + when(container.localizationCountersAsString()).thenReturn("1,2,3,4,5"); + + // Define user environment variables. + Map userSetEnv = new HashMap(); + String userVar = "USER_VAR"; + String userVarVal = "user-var-value"; + userSetEnv.put(userVar, userVarVal); + when(clc.getEnvironment()).thenReturn(userSetEnv); + + YarnConfiguration localConf = new YarnConfiguration(conf); + + // Admin Env var that depends on USER_VAR1 + String testKey1 = "TEST_KEY1"; + String testVal1 = "relies on {{USER_VAR}}"; + localConf.set( + YarnConfiguration.NM_ADMIN_USER_ENV + "." + testKey1, testVal1); + String testVal1Expanded; // this is what we expect after {{}} expansion + if (Shell.WINDOWS) { + testVal1Expanded = "relies on %USER_VAR%"; + } else { + testVal1Expanded = "relies on $USER_VAR"; + } + // Another Admin Env var that depends on the first one + String testKey2 = "TEST_KEY2"; + String testVal2 = "relies on {{TEST_KEY1}}"; + localConf.set( + YarnConfiguration.NM_ADMIN_USER_ENV + "." + testKey2, testVal2); + String testVal2Expanded; // this is what we expect after {{}} expansion + if (Shell.WINDOWS) { + testVal2Expanded = "relies on %TEST_KEY1%"; + } else { + testVal2Expanded = "relies on $TEST_KEY1"; + } + + // call containerLaunch + ContainerLaunch containerLaunch = new ContainerLaunch( + distContext, localConf, dispatcher, + containerExecutor, app, container, dirsHandler, containerManager); + containerLaunch.call(); + + // verify the nmPrivate paths and files + ArgumentCaptor cscArgument = + ArgumentCaptor.forClass(ContainerStartContext.class); + verify(containerExecutor, times(1)).launchContainer(cscArgument.capture()); + ContainerStartContext csc = cscArgument.getValue(); + Assert.assertEquals("script", + readStringFromPath(csc.getNmPrivateContainerScriptPath())); + + // verify env + ArgumentCaptor envArgument = ArgumentCaptor.forClass(Map.class); + verify(containerExecutor, times(1)).writeLaunchEnv(any(), + envArgument.capture(), any(), any(), any(), any(), any()); + Map env = envArgument.getValue(); + Assert.assertEquals(userVarVal, env.get(userVar)); + Assert.assertEquals(testVal1Expanded, env.get(testKey1)); + Assert.assertEquals(testVal2Expanded, env.get(testKey2)); + } + }