YARN-7677. Docker image cannot set HADOOP_CONF_DIR. Contributed by Jim Brennan

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/util/Apps.java

@@ -23,6 +23,7 @@ import static org.apache.hadoop.yarn.util.StringHelper.join;
 import static org.apache.hadoop.yarn.util.StringHelper.sjoin;
 
 import java.io.File;
+import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.Map;
 import java.util.regex.Matcher;
@@ -105,7 +106,26 @@ public class Apps {
       }
     }
   }
-  
+
+  /**
+   *
+   * @param envString String containing env variable definitions
+   * @param classPathSeparator String that separates the definitions
+   * @return ArrayList of environment variable names
+   */
+  public static ArrayList<String> getEnvVarsFromInputString(String envString,
+      String classPathSeparator) {
+    ArrayList<String> envList = new ArrayList<>();
+    if (envString != null && envString.length() > 0) {
+      Matcher varValMatcher = VARVAL_SPLITTER.matcher(envString);
+      while (varValMatcher.find()) {
+        String envVar = varValMatcher.group(1);
+        envList.add(envVar);
+      }
+    }
+    return envList;
+  }
+
   /**
    * This older version of this method is kept around for compatibility
    * because downstream frameworks like Spark and Tez have been using it.

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/util/AuxiliaryServiceHelper.java

@@ -45,7 +45,7 @@ public class AuxiliaryServiceHelper {
         Base64.encodeBase64String(byteData));
   }
 
-  private static String getPrefixServiceName(String serviceName) {
+  public static String getPrefixServiceName(String serviceName) {
     return NM_AUX_SERVICE + serviceName;
   }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/ContainerExecutor.java

@@ -27,6 +27,7 @@ import java.net.UnknownHostException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -316,14 +317,15 @@ public abstract class ContainerExecutor implements Configurable {
    * @param command the command that will be run
    * @param logDir the log dir to which to copy debugging information
    * @param user the username of the job owner
+   * @param nmVars the set of environment vars that are explicitly set by NM
    * @throws IOException if any errors happened writing to the OutputStream,
    * while creating symlinks
    */
   public void writeLaunchEnv(OutputStream out, Map<String, String> environment,
       Map<Path, List<String>> resources, List<String> command, Path logDir,
-      String user) throws IOException {
+      String user, LinkedHashSet<String> nmVars) throws IOException {
     this.writeLaunchEnv(out, environment, resources, command, logDir, user,
-        ContainerLaunch.CONTAINER_SCRIPT);
+        ContainerLaunch.CONTAINER_SCRIPT, nmVars);
   }
 
   /**
@@ -339,14 +341,15 @@ public abstract class ContainerExecutor implements Configurable {
    * @param logDir the log dir to which to copy debugging information
    * @param user the username of the job owner
    * @param outFilename the path to which to write the launch environment
+   * @param nmVars the set of environment vars that are explicitly set by NM
    * @throws IOException if any errors happened writing to the OutputStream,
    * while creating symlinks
    */
   @VisibleForTesting
   public void writeLaunchEnv(OutputStream out, Map<String, String> environment,
       Map<Path, List<String>> resources, List<String> command, Path logDir,
-      String user, String outFilename) throws IOException {
-    updateEnvForWhitelistVars(environment);
+      String user, String outFilename, LinkedHashSet<String> nmVars)
+      throws IOException {
 
     ContainerLaunch.ShellScriptBuilder sb =
         ContainerLaunch.ShellScriptBuilder.create();
@@ -361,8 +364,40 @@ public abstract class ContainerExecutor implements Configurable {
 
     if (environment != null) {
       sb.echo("Setting up env variables");
+      // Whitelist environment variables are treated specially.
+      // Only add them if they are not already defined in the environment.
+      // Add them using special syntax to prevent them from eclipsing
+      // variables that may be set explicitly in the container image (e.g,
+      // in a docker image).  Put these before the others to ensure the
+      // correct expansion is used.
+      for(String var : whitelistVars) {
+        if (!environment.containsKey(var)) {
+          String val = getNMEnvVar(var);
+          if (val != null) {
+            sb.whitelistedEnv(var, val);
+          }
+        }
+      }
+      // Now write vars that were set explicitly by nodemanager, preserving
+      // the order they were written in.
+      for (String nmEnvVar : nmVars) {
+        sb.env(nmEnvVar, environment.get(nmEnvVar));
+      }
+      // Now write the remaining environment variables.
       for (Map.Entry<String, String> env : environment.entrySet()) {
-        sb.env(env.getKey(), env.getValue());
+        if (!nmVars.contains(env.getKey())) {
+          sb.env(env.getKey(), env.getValue());
+        }
+      }
+      // Add the whitelist vars to the environment.  Do this after writing
+      // environment variables so they are not written twice.
+      for(String var : whitelistVars) {
+        if (!environment.containsKey(var)) {
+          String val = getNMEnvVar(var);
+          if (val != null) {
+            environment.put(var, val);
+          }
+        }
       }
     }
 
@@ -663,23 +698,6 @@ public abstract class ContainerExecutor implements Configurable {
     }
   }
 
-  /**
-   * Propagate variables from the nodemanager's environment into the
-   * container's environment if unspecified by the container.
-   * @param env the environment to update
-   * @see org.apache.hadoop.yarn.conf.YarnConfiguration#NM_ENV_WHITELIST
-   */
-  protected void updateEnvForWhitelistVars(Map<String, String> env) {
-    for(String var : whitelistVars) {
-      if (!env.containsKey(var)) {
-        String val = getNMEnvVar(var);
-        if (val != null) {
-          env.put(var, val);
-        }
-      }
-    }
-  }
-
   @VisibleForTesting
   protected String getNMEnvVar(String varname) {
     return System.getenv(varname);

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/LinuxContainerExecutor.java

@@ -66,7 +66,6 @@ import java.net.InetSocketAddress;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Map;
 import java.util.regex.Pattern;
 
 import static org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.runtime.LinuxContainerRuntimeConstants.*;
@@ -472,13 +471,6 @@ public class LinuxContainerExecutor extends ContainerExecutor {
     }
   }
 
-  @Override
-  protected void updateEnvForWhitelistVars(Map<String, String> env) {
-    if (linuxContainerRuntime.useWhitelistEnv(env)) {
-      super.updateEnvForWhitelistVars(env);
-    }
-  }
-
   @Override
   public int launchContainer(ContainerStartContext ctx)
       throws IOException, ConfigurationException {

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

@@ -33,7 +33,9 @@ import java.util.ArrayList;
 import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.List;
+import java.util.LinkedHashSet;
 import java.util.Map;
+import java.util.Set;
 import java.util.Map.Entry;
 import java.util.concurrent.Callable;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -217,6 +219,9 @@ public class ContainerLaunch implements Callable<Integer> {
           launchContext, containerLogDir);
       // /////////////////////////// End of variable expansion
 
+      // Use this to track variables that are added to the environment by nm.
+      LinkedHashSet<String> nmEnvVars = new LinkedHashSet<String>();
+
       FileContext lfs = FileContext.getLocalFSFileContext();
 
       Path nmPrivateContainerScriptPath = dirsHandler.getLocalPathForWrite(
@@ -261,7 +266,7 @@ public class ContainerLaunch implements Callable<Integer> {
       }
 
       // Set the token location too.
-      environment.put(
+      addToEnvMap(environment, nmEnvVars,
           ApplicationConstants.CONTAINER_TOKEN_FILE_ENV_NAME,
           new Path(containerWorkDir,
               FINAL_CONTAINER_TOKENS_FILE).toUri().getPath());
@@ -272,14 +277,15 @@ public class ContainerLaunch implements Callable<Integer> {
                    EnumSet.of(CREATE, OVERWRITE))) {
         // Sanitize the container's environment
         sanitizeEnv(environment, containerWorkDir, appDirs, userLocalDirs,
-            containerLogDirs, localResources, nmPrivateClasspathJarDir);
+            containerLogDirs, localResources, nmPrivateClasspathJarDir,
+            nmEnvVars);
 
         prepareContainer(localResources, containerLocalDirs);
 
         // Write out the environment
         exec.writeLaunchEnv(containerScriptOutStream, environment,
             localResources, launchContext.getCommands(),
-            containerLogDir, user);
+            containerLogDir, user, nmEnvVars);
       }
       // /////////// End of writing out container-script
 
@@ -1171,6 +1177,9 @@ public class ContainerLaunch implements Callable<Integer> {
 
     public abstract void env(String key, String value) throws IOException;
 
+    public abstract void whitelistedEnv(String key, String value)
+        throws IOException;
+
     public abstract void echo(String echoStr) throws IOException;
 
     public final void symlink(Path src, Path dst) throws IOException {
@@ -1290,6 +1299,11 @@ public class ContainerLaunch implements Callable<Integer> {
       line("export ", key, "=\"", value, "\"");
     }
 
+    @Override
+    public void whitelistedEnv(String key, String value) throws IOException {
+      line("export ", key, "=${", key, ":-", "\"", value, "\"}");
+    }
+
     @Override
     public void echo(final String echoStr) throws IOException {
       line("echo \"" + echoStr + "\"");
@@ -1380,6 +1394,11 @@ public class ContainerLaunch implements Callable<Integer> {
       errorCheck();
     }
 
+    @Override
+    public void whitelistedEnv(String key, String value) throws IOException {
+      env(key, value);
+    }
+
     @Override
     public void echo(final String echoStr) throws IOException {
       lineWithLenCheck("@echo \"", echoStr, "\"");
@@ -1435,60 +1454,70 @@ public class ContainerLaunch implements Callable<Integer> {
       putEnvIfNotNull(environment, variable, System.getenv(variable));
     }
   }
-  
+
+  private static void addToEnvMap(
+      Map<String, String> envMap, Set<String> envSet,
+      String envName, String envValue) {
+    envMap.put(envName, envValue);
+    envSet.add(envName);
+  }
+
   public void sanitizeEnv(Map<String, String> environment, Path pwd,
       List<Path> appDirs, List<String> userLocalDirs, List<String>
-      containerLogDirs,
-      Map<Path, List<String>> resources,
-      Path nmPrivateClasspathJarDir) throws IOException {
+      containerLogDirs, Map<Path, List<String>> resources,
+      Path nmPrivateClasspathJarDir,
+      Set<String> nmVars) throws IOException {
     /**
      * Non-modifiable environment variables
      */
 
-    environment.put(Environment.CONTAINER_ID.name(), container
-        .getContainerId().toString());
+    addToEnvMap(environment, nmVars, Environment.CONTAINER_ID.name(),
+        container.getContainerId().toString());
 
-    environment.put(Environment.NM_PORT.name(),
+    addToEnvMap(environment, nmVars, Environment.NM_PORT.name(),
       String.valueOf(this.context.getNodeId().getPort()));
 
-    environment.put(Environment.NM_HOST.name(), this.context.getNodeId()
-      .getHost());
+    addToEnvMap(environment, nmVars, Environment.NM_HOST.name(),
+        this.context.getNodeId().getHost());
 
-    environment.put(Environment.NM_HTTP_PORT.name(),
+    addToEnvMap(environment, nmVars, Environment.NM_HTTP_PORT.name(),
       String.valueOf(this.context.getHttpPort()));
 
-    environment.put(Environment.LOCAL_DIRS.name(),
+    addToEnvMap(environment, nmVars, Environment.LOCAL_DIRS.name(),
         StringUtils.join(",", appDirs));
 
-    environment.put(Environment.LOCAL_USER_DIRS.name(), StringUtils.join(",",
-        userLocalDirs));
+    addToEnvMap(environment, nmVars, Environment.LOCAL_USER_DIRS.name(),
+        StringUtils.join(",", userLocalDirs));
 
-    environment.put(Environment.LOG_DIRS.name(),
+    addToEnvMap(environment, nmVars, Environment.LOG_DIRS.name(),
       StringUtils.join(",", containerLogDirs));
 
-    environment.put(Environment.USER.name(), container.getUser());
-    
-    environment.put(Environment.LOGNAME.name(), container.getUser());
+    addToEnvMap(environment, nmVars, Environment.USER.name(),
+        container.getUser());
 
-    environment.put(Environment.HOME.name(),
+    addToEnvMap(environment, nmVars, Environment.LOGNAME.name(),
+        container.getUser());
+
+    addToEnvMap(environment, nmVars, Environment.HOME.name(),
         conf.get(
             YarnConfiguration.NM_USER_HOME_DIR, 
             YarnConfiguration.DEFAULT_NM_USER_HOME_DIR
             )
         );
-    
-    environment.put(Environment.PWD.name(), pwd.toString());
-    
-    putEnvIfAbsent(environment, Environment.HADOOP_CONF_DIR.name());
+
+    addToEnvMap(environment, nmVars, Environment.PWD.name(), pwd.toString());
 
     if (!Shell.WINDOWS) {
-      environment.put("JVM_PID", "$$");
+      addToEnvMap(environment, nmVars, "JVM_PID", "$$");
     }
 
     // variables here will be forced in, even if the container has specified them.
-    Apps.setEnvFromInputString(environment, conf.get(
-      YarnConfiguration.NM_ADMIN_USER_ENV,
-      YarnConfiguration.DEFAULT_NM_ADMIN_USER_ENV), File.pathSeparator);
+    String nmAdminUserEnv = conf.get(
+        YarnConfiguration.NM_ADMIN_USER_ENV,
+        YarnConfiguration.DEFAULT_NM_ADMIN_USER_ENV);
+    Apps.setEnvFromInputString(environment, nmAdminUserEnv, File.pathSeparator);
+    nmVars.addAll(Apps.getEnvVarsFromInputString(nmAdminUserEnv,
+        File.pathSeparator));
 
     // TODO: Remove Windows check and use this approach on all platforms after
     // additional testing.  See YARN-358.
@@ -1502,6 +1531,7 @@ public class ContainerLaunch implements Callable<Integer> {
         .getAuxServiceMetaData().entrySet()) {
       AuxiliaryServiceHelper.setServiceDataIntoEnv(
           meta.getKey(), meta.getValue(), environment);
+      nmVars.add(AuxiliaryServiceHelper.getPrefixServiceName(meta.getKey()));
     }
   }
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/DefaultLinuxContainerRuntime.java

@@ -37,7 +37,6 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.util.List;
-import java.util.Map;
 
 import static org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.runtime.LinuxContainerRuntimeConstants.*;
 
@@ -73,11 +72,6 @@ public class DefaultLinuxContainerRuntime implements LinuxContainerRuntime {
     this.conf = conf;
   }
 
-  @Override
-  public boolean useWhitelistEnv(Map<String, String> env) {
-    return true;
-  }
-
   @Override
   public void prepareContainer(ContainerRuntimeContext ctx)
       throws ContainerExecutionException {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/DelegatingLinuxContainerRuntime.java

@@ -94,17 +94,6 @@ public class DelegatingLinuxContainerRuntime implements LinuxContainerRuntime {
     }
   }
 
-  @Override
-  public boolean useWhitelistEnv(Map<String, String> env) {
-    try {
-      LinuxContainerRuntime runtime = pickContainerRuntime(env);
-      return runtime.useWhitelistEnv(env);
-    } catch (ContainerExecutionException e) {
-      LOG.debug("Unable to determine runtime");
-      return false;
-    }
-  }
-
   @VisibleForTesting
   LinuxContainerRuntime pickContainerRuntime(
       Map<String, String> environment) throws ContainerExecutionException {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/DockerLinuxContainerRuntime.java

@@ -371,13 +371,6 @@ public class DockerLinuxContainerRuntime implements LinuxContainerRuntime {
     return capabilities;
   }
 
-  @Override
-  public boolean useWhitelistEnv(Map<String, String> env) {
-    // Avoid propagating nodemanager environment variables into the container
-    // so those variables can be picked up from the Docker image instead.
-    return false;
-  }
-
   private String runDockerVolumeCommand(DockerVolumeCommand dockerVolumeCommand,
       Container container) throws ContainerExecutionException {
     try {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/runtime/ContainerRuntime.java

@@ -24,8 +24,6 @@ import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceStability;
 import org.apache.hadoop.yarn.server.nodemanager.containermanager.container.Container;
 
-import java.util.Map;
-
 /**
  * An abstraction for various container runtime implementations. Examples
  * include Process Tree, Docker, Appc runtimes etc. These implementations
@@ -85,13 +83,4 @@ public interface ContainerRuntime {
    * and hostname
    */
   String[] getIpAndHost(Container container) throws ContainerExecutionException;
-
-  /**
-   * Whether to propagate the whitelist of environment variables from the
-   * nodemanager environment into the container environment.
-   * @param env the container's environment variables
-   * @return true if whitelist variables should be propagated, false otherwise
-   * @see org.apache.hadoop.yarn.conf.YarnConfiguration#NM_ENV_WHITELIST
-   */
-  boolean useWhitelistEnv(Map<String, String> env);
 }

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

@@ -41,6 +41,7 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.StringTokenizer;
@@ -185,8 +186,10 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
       DefaultContainerExecutor defaultContainerExecutor =
           new DefaultContainerExecutor();
       defaultContainerExecutor.setConf(new YarnConfiguration());
+      LinkedHashSet<String> nmVars = new LinkedHashSet<>();
       defaultContainerExecutor.writeLaunchEnv(fos, env, resources, commands,
-          new Path(localLogDir.getAbsolutePath()), "user", tempFile.getName());
+          new Path(localLogDir.getAbsolutePath()), "user", tempFile.getName(),
+          nmVars);
       fos.flush();
       fos.close();
       FileUtil.setExecutable(tempFile, true);
@@ -260,8 +263,9 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
       DefaultContainerExecutor defaultContainerExecutor =
           new DefaultContainerExecutor();
       defaultContainerExecutor.setConf(new YarnConfiguration());
+      LinkedHashSet<String> nmVars = new LinkedHashSet<>();
       defaultContainerExecutor.writeLaunchEnv(fos, env, resources, commands,
-          new Path(localLogDir.getAbsolutePath()), "user");
+          new Path(localLogDir.getAbsolutePath()), "user", nmVars);
       fos.flush();
       fos.close();
       FileUtil.setExecutable(tempFile, true);
@@ -323,8 +327,9 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
     conf.set(YarnConfiguration.NM_ENV_WHITELIST,
         "HADOOP_MAPRED_HOME,HADOOP_YARN_HOME");
     defaultContainerExecutor.setConf(conf);
+    LinkedHashSet<String> nmVars = new LinkedHashSet<>();
     defaultContainerExecutor.writeLaunchEnv(fos, env, resources, commands,
-        new Path(localLogDir.getAbsolutePath()), "user");
+        new Path(localLogDir.getAbsolutePath()), "user", nmVars);
     String shellContent =
         new String(Files.readAllBytes(Paths.get(shellFile.getAbsolutePath())),
             StandardCharsets.UTF_8);
@@ -337,7 +342,8 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
     Assert.assertFalse(shellContent.contains("HADOOP_HDFS_HOME"));
     // Available in env and in whitelist
     Assert.assertTrue(shellContent.contains(
-        "export HADOOP_YARN_HOME=\"nodemanager_yarn_home\""));
+        "export HADOOP_YARN_HOME=${HADOOP_YARN_HOME:-\"nodemanager_yarn_home\"}"
+      ));
     fos.flush();
     fos.close();
   }
@@ -372,8 +378,9 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
     conf.set(YarnConfiguration.NM_ENV_WHITELIST,
         "HADOOP_MAPRED_HOME,HADOOP_YARN_HOME");
     lce.setConf(conf);
+    LinkedHashSet<String> nmVars = new LinkedHashSet<>();
     lce.writeLaunchEnv(fos, env, resources, commands,
-        new Path(localLogDir.getAbsolutePath()), "user");
+        new Path(localLogDir.getAbsolutePath()), "user", nmVars);
     String shellContent =
         new String(Files.readAllBytes(Paths.get(shellFile.getAbsolutePath())),
             StandardCharsets.UTF_8);
@@ -382,13 +389,106 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
     // Whitelisted variable overridden by container
     Assert.assertTrue(shellContent.contains(
         "export HADOOP_MAPRED_HOME=\"/opt/hadoopbuild\""));
-    // Verify no whitelisted variables inherited from NM env
+    // Available in env but not in whitelist
     Assert.assertFalse(shellContent.contains("HADOOP_HDFS_HOME"));
-    Assert.assertFalse(shellContent.contains("HADOOP_YARN_HOME"));
+    // Available in env and in whitelist
+    Assert.assertTrue(shellContent.contains(
+        "export HADOOP_YARN_HOME=${HADOOP_YARN_HOME:-\"nodemanager_yarn_home\"}"
+    ));
+    fos.flush();
+    fos.close();
+  }
+
+  @Test(timeout = 20000)
+  public void testWriteEnvOrder() throws Exception {
+    // Valid only for unix
+    assumeNotWindows();
+    List<String> commands = new ArrayList<String>();
+
+    // Setup user-defined environment
+    Map<String, String> env = new HashMap<String, String>();
+    env.put("USER_VAR_1", "1");
+    env.put("USER_VAR_2", "2");
+    env.put("NM_MODIFIED_VAR_1", "nm 1");
+    env.put("NM_MODIFIED_VAR_2", "nm 2");
+
+    // These represent vars explicitly set by NM
+    LinkedHashSet<String> trackedNmVars = new LinkedHashSet<>();
+    trackedNmVars.add("NM_MODIFIED_VAR_1");
+    trackedNmVars.add("NM_MODIFIED_VAR_2");
+
+    // Setup Nodemanager environment
+    final Map<String, String> nmEnv = new HashMap<>();
+    nmEnv.put("WHITELIST_VAR_1", "wl 1");
+    nmEnv.put("WHITELIST_VAR_2", "wl 2");
+    nmEnv.put("NON_WHITELIST_VAR_1", "nwl 1");
+    nmEnv.put("NON_WHITELIST_VAR_2", "nwl 2");
+    DefaultContainerExecutor defaultContainerExecutor =
+        new DefaultContainerExecutor() {
+          @Override
+          protected String getNMEnvVar(String varname) {
+            return nmEnv.get(varname);
+          }
+        };
+
+    // Setup conf with whitelisted variables
+    ArrayList<String> whitelistVars = new ArrayList<>();
+    whitelistVars.add("WHITELIST_VAR_1");
+    whitelistVars.add("WHITELIST_VAR_2");
+    YarnConfiguration conf = new YarnConfiguration();
+    conf.set(YarnConfiguration.NM_ENV_WHITELIST,
+        whitelistVars.get(0) + "," + whitelistVars.get(1));
+
+    // These are in the NM env, but not in the whitelist.
+    ArrayList<String> nonWhiteListEnv = new ArrayList<>();
+    nonWhiteListEnv.add("NON_WHITELIST_VAR_1");
+    nonWhiteListEnv.add("NON_WHITELIST_VAR_2");
+
+    // Write the launch script
+    File shellFile = Shell.appendScriptExtension(tmpDir, "hello");
+    Map<Path, List<String>> resources = new HashMap<Path, List<String>>();
+    FileOutputStream fos = new FileOutputStream(shellFile);
+    defaultContainerExecutor.setConf(conf);
+    defaultContainerExecutor.writeLaunchEnv(fos, env, resources, commands,
+        new Path(localLogDir.getAbsolutePath()), "user", trackedNmVars);
     fos.flush();
     fos.close();
+
+    // Examine the script
+    String shellContent =
+        new String(Files.readAllBytes(Paths.get(shellFile.getAbsolutePath())),
+            StandardCharsets.UTF_8);
+    // First make sure everything is there that's supposed to be
+    for (String envVar : env.keySet()) {
+      Assert.assertTrue(shellContent.contains(envVar + "="));
+    }
+    for (String wlVar : whitelistVars) {
+      Assert.assertTrue(shellContent.contains(wlVar + "="));
+    }
+    for (String nwlVar : nonWhiteListEnv) {
+      Assert.assertFalse(shellContent.contains(nwlVar + "="));
+    }
+    // Explicitly Set NM vars should be before user vars
+    for (String nmVar : trackedNmVars) {
+      for (String userVar : env.keySet()) {
+        // Need to skip nm vars and whitelist vars
+        if (!trackedNmVars.contains(userVar) &&
+            !whitelistVars.contains(userVar)) {
+          Assert.assertTrue(shellContent.indexOf(nmVar + "=") <
+              shellContent.indexOf(userVar + "="));
+        }
+      }
+    }
+    // Whitelisted vars should be before explicitly set NM vars
+    for (String wlVar : whitelistVars) {
+      for (String nmVar : trackedNmVars) {
+        Assert.assertTrue(shellContent.indexOf(wlVar + "=") <
+            shellContent.indexOf(nmVar + "="));
+      }
+    }
   }
 
+
   @Test (timeout = 20000)
   public void testInvalidEnvSyntaxDiagnostics() throws IOException  {
 
@@ -410,8 +510,9 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
       DefaultContainerExecutor defaultContainerExecutor =
           new DefaultContainerExecutor();
       defaultContainerExecutor.setConf(new YarnConfiguration());
+      LinkedHashSet<String> nmVars = new LinkedHashSet<>();
       defaultContainerExecutor.writeLaunchEnv(fos, env, resources, commands,
-          new Path(localLogDir.getAbsolutePath()), "user");
+          new Path(localLogDir.getAbsolutePath()), "user", nmVars);
       fos.flush();
       fos.close();
 
@@ -493,8 +594,9 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
       commands.add(command);
       ContainerExecutor exec = new DefaultContainerExecutor();
       exec.setConf(new YarnConfiguration());
+      LinkedHashSet<String> nmVars = new LinkedHashSet<>();
       exec.writeLaunchEnv(fos, env, resources, commands,
-          new Path(localLogDir.getAbsolutePath()), "user");
+          new Path(localLogDir.getAbsolutePath()), "user", nmVars);
       fos.flush();
       fos.close();
 
@@ -585,7 +687,7 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
     Path nmp = new Path(testDir);
 
     launch.sanitizeEnv(userSetEnv, pwd, appDirs, userLocalDirs, containerLogs,
-        resources, nmp);
+        resources, nmp, Collections.emptySet());
 
     List<String> result =
       getJarManifestClasspath(userSetEnv.get(Environment.CLASSPATH.name()));
@@ -604,7 +706,7 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
         dispatcher, exec, null, container, dirsHandler, containerManager);
 
     launch.sanitizeEnv(userSetEnv, pwd, appDirs, userLocalDirs, containerLogs,
-        resources, nmp);
+        resources, nmp, Collections.emptySet());
 
     result =
       getJarManifestClasspath(userSetEnv.get(Environment.CLASSPATH.name()));
@@ -1528,9 +1630,10 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
         FileOutputStream fos = new FileOutputStream(tempFile);
         ContainerExecutor exec = new DefaultContainerExecutor();
         exec.setConf(conf);
+        LinkedHashSet<String> nmVars = new LinkedHashSet<>();
         exec.writeLaunchEnv(fos, env, resources, commands,
             new Path(localLogDir.getAbsolutePath()), "user",
-            tempFile.getName());
+            tempFile.getName(), nmVars);
         fos.flush();
         fos.close();
         FileUtil.setExecutable(tempFile, true);
@@ -1753,8 +1856,9 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
       List<String> commands = new ArrayList<String>();
       DefaultContainerExecutor executor = new DefaultContainerExecutor();
       executor.setConf(new Configuration());
+      LinkedHashSet<String> nmVars = new LinkedHashSet<>();
       executor.writeLaunchEnv(fos, env, resources, commands,
-          new Path(localLogDir.getAbsolutePath()), user);
+          new Path(localLogDir.getAbsolutePath()), user, nmVars);
       fos.flush();
       fos.close();
 
@@ -1798,8 +1902,9 @@ public class TestContainerLaunch extends BaseContainerManagerTest {
       Configuration execConf = new Configuration();
       execConf.setBoolean(YarnConfiguration.NM_LOG_CONTAINER_DEBUG_INFO, false);
       executor.setConf(execConf);
+      LinkedHashSet<String> nmVars = new LinkedHashSet<>();
       executor.writeLaunchEnv(fos, env, resources, commands,
-          new Path(localLogDir.getAbsolutePath()), user);
+          new Path(localLogDir.getAbsolutePath()), user, nmVars);
       fos.flush();
       fos.close();