Apache
/
hadoop
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

YARN-10664. Allow parameter expansion in NM_ADMIN_USER_ENV. Contributed by Jim 

Brennan.

(cherry picked from commit 9d088639bf)
This commit is contained in:
Eric Badger 2021-03-05 19:49:49 +00:00
parent e607d03995
commit 1a33878279
2 changed files with 136 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

66
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
View File

@ -179,15 +179,36 @@ public class ContainerLaunch implements Callable<Integer> {
return var;
}
private Map<String, String> expandAllEnvironmentVars(
ContainerLaunchContext launchContext, Path containerLogDir) {
Map<String, String> environment = launchContext.getEnvironment();
private void expandAllEnvironmentVars(
Map<String, String> environment, Path containerLogDir) {
for (Entry<String, String> entry : environment.entrySet()) {
String value = entry.getValue();
value = expandEnvironment(value, containerLogDir);
entry.setValue(value);
}
return environment;
}
private void addKeystoreVars(Map<String, String> environment,
Path containerWorkDir) {
environment.put(ApplicationConstants.KEYSTORE_FILE_LOCATION_ENV_NAME,
new Path(containerWorkDir,
ContainerLaunch.KEYSTORE_FILE).toUri().getPath());
environment.put(ApplicationConstants.KEYSTORE_PASSWORD_ENV_NAME,
new String(container.getCredentials().getSecretKey(
AMSecretKeys.YARN_APPLICATION_AM_KEYSTORE_PASSWORD),
StandardCharsets.UTF_8));
}
private void addTruststoreVars(Map<String, String> environment,
Path containerWorkDir) {
environment.put(
ApplicationConstants.TRUSTSTORE_FILE_LOCATION_ENV_NAME,
new Path(containerWorkDir,
ContainerLaunch.TRUSTSTORE_FILE).toUri().getPath());
environment.put(ApplicationConstants.TRUSTSTORE_PASSWORD_ENV_NAME,
new String(container.getCredentials().getSecretKey(
AMSecretKeys.YARN_APPLICATION_AM_TRUSTSTORE_PASSWORD),
StandardCharsets.UTF_8));
}
@Override
@ -222,8 +243,10 @@ public class ContainerLaunch implements Callable<Integer> {
}
launchContext.setCommands(newCmds);
Map<String, String> 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<String, String> environment = launchContext.getEnvironment();
// /////////////////////////// End of variable expansion
// Use this to track variables that are added to the environment by nm.
@ -289,13 +312,6 @@ public class ContainerLaunch implements Callable<Integer> {
lfs.create(nmPrivateKeystorePath,
EnumSet.of(CREATE, OVERWRITE))) {
keystoreOutStream.write(keystore);
environment.put(ApplicationConstants.KEYSTORE_FILE_LOCATION_ENV_NAME,
new Path(containerWorkDir,
ContainerLaunch.KEYSTORE_FILE).toUri().getPath());
environment.put(ApplicationConstants.KEYSTORE_PASSWORD_ENV_NAME,
new String(container.getCredentials().getSecretKey(
AMSecretKeys.YARN_APPLICATION_AM_KEYSTORE_PASSWORD),
StandardCharsets.UTF_8));
}
} else {
nmPrivateKeystorePath = null;
@ -307,14 +323,6 @@ public class ContainerLaunch implements Callable<Integer> {
lfs.create(nmPrivateTruststorePath,
EnumSet.of(CREATE, OVERWRITE))) {
truststoreOutStream.write(truststore);
environment.put(
ApplicationConstants.TRUSTSTORE_FILE_LOCATION_ENV_NAME,
new Path(containerWorkDir,
ContainerLaunch.TRUSTSTORE_FILE).toUri().getPath());
environment.put(ApplicationConstants.TRUSTSTORE_PASSWORD_ENV_NAME,
new String(container.getCredentials().getSecretKey(
AMSecretKeys.YARN_APPLICATION_AM_TRUSTSTORE_PASSWORD),
StandardCharsets.UTF_8));
}
} else {
nmPrivateTruststorePath = null;
@ -335,6 +343,16 @@ public class ContainerLaunch implements Callable<Integer> {
containerLogDirs, localResources, nmPrivateClasspathJarDir,
nmEnvVars);
expandAllEnvironmentVars(environment, containerLogDir);
// Add these if needed after expanding so we don't expand key values.
if (keystore != null) {
addKeystoreVars(environment, containerWorkDir);
}
if (truststore != null) {
addTruststoreVars(environment, containerWorkDir);
}
prepareContainer(localResources, containerLocalDirs);
// Write out the environment
@ -1628,13 +1646,13 @@ public class ContainerLaunch implements Callable<Integer> {
}
// 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

97
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
View File

@ -673,7 +673,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<Event> eventHandler = new EventHandler<Event>() {
public void handle(Event event) {
@ -814,8 +814,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));
@ -1857,6 +1855,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);
@ -2453,6 +2452,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.<Path, List<String>> emptyMap());
@ -2562,6 +2562,7 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
when(container.getLaunchContext()).thenReturn(clc);
Credentials credentials = mock(Credentials.class);
when(container.getCredentials()).thenReturn(credentials);
when(container.localizationCountersAsString()).thenReturn("1,2,3,4,5");
doAnswer(new Answer<Void>() {
@Override
public Void answer(InvocationOnMock invocation) throws Throwable {
@ -2662,4 +2663,94 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
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<Void>() {
@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<String, String> userSetEnv = new HashMap<String, String>();
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<ContainerStartContext> 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<Map> 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));
}
}