From dffb7bfe6cd2292162f08ec0bded736bc5194c3f Mon Sep 17 00:00:00 2001
From: Shane Kumpf <skumpf@apache.org>
Date: Wed, 5 Sep 2018 06:47:54 -0600
Subject: [PATCH] YARN-8638. Allow linux container runtimes to be pluggable.
 Contributed by Craig Condit

---
 .../hadoop/yarn/conf/YarnConfiguration.java   |   5 +-
 .../runtime/DefaultLinuxContainerRuntime.java |   8 ++
 .../DelegatingLinuxContainerRuntime.java      | 105 +++++++++++++-----
 .../runtime/DockerLinuxContainerRuntime.java  |   5 +
 .../JavaSandboxLinuxContainerRuntime.java     |   4 +-
 .../linux/runtime/LinuxContainerRuntime.java  |  11 ++
 .../runtime/MockLinuxContainerRuntime.java    |  61 ++++++++++
 .../TestDelegatingLinuxContainerRuntime.java  |  71 +++++++++---
 8 files changed, 225 insertions(+), 45 deletions(-)
 create mode 100644 hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/MockLinuxContainerRuntime.java

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 d525e4d2f4e..d681c03abbd 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
@@ -1865,7 +1865,7 @@ public class YarnConfiguration extends Configuration {
 
   /**
    * Comma separated list of runtimes that are allowed when using
-   * LinuxContainerExecutor. The allowed values are:
+   * LinuxContainerExecutor. The standard values are:
    * <ul>
    *   <li>default</li>
    *   <li>docker</li>
@@ -1875,6 +1875,9 @@ public class YarnConfiguration extends Configuration {
   public static final String LINUX_CONTAINER_RUNTIME_ALLOWED_RUNTIMES =
       LINUX_CONTAINER_RUNTIME_PREFIX + "allowed-runtimes";
 
+  public static final String LINUX_CONTAINER_RUNTIME_CLASS_FMT =
+      LINUX_CONTAINER_RUNTIME_PREFIX + "%s.class";
+
   /** The default list of allowed runtimes when using LinuxContainerExecutor. */
   public static final String[] DEFAULT_LINUX_CONTAINER_RUNTIME_ALLOWED_RUNTIMES
       = {"default"};
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/linux/runtime/DefaultLinuxContainerRuntime.java b/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
index b5c933aff23..2b8e8adfba4 100644
--- a/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
+++ b/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
@@ -32,11 +32,13 @@ import org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.privileg
 import org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.privileged.PrivilegedOperationExecutor;
 import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.ContainerExecutionException;
 import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.ContainerRuntime;
+import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.ContainerRuntimeConstants;
 import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.ContainerRuntimeContext;
 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.*;
 
@@ -66,6 +68,12 @@ public class DefaultLinuxContainerRuntime implements LinuxContainerRuntime {
     this.privilegedOperationExecutor = privilegedOperationExecutor;
   }
 
+  @Override
+  public boolean isRuntimeRequested(Map<String, String> env) {
+    String type = env.get(ContainerRuntimeConstants.ENV_CONTAINER_TYPE);
+    return type == null || type.equals("default");
+  }
+
   @Override
   public void initialize(Configuration conf, Context nmContext)
       throws ContainerExecutionException {
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/linux/runtime/DelegatingLinuxContainerRuntime.java b/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
index f936169b1c6..7199b54c9c0 100644
--- a/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
+++ b/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
@@ -24,6 +24,7 @@ import com.google.common.annotations.VisibleForTesting;
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceStability;
 import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.util.ReflectionUtils;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.server.nodemanager.Context;
 import org.apache.hadoop.yarn.server.nodemanager.containermanager.container.Container;
@@ -34,28 +35,32 @@ import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.Contai
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.EnumSet;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * This class is a {@link ContainerRuntime} implementation that delegates all
  * operations to a {@link DefaultLinuxContainerRuntime} instance, a
- * {@link DockerLinuxContainerRuntime} instance, or a
- * {@link JavaSandboxLinuxContainerRuntime} instance depending on whether
- * each instance believes the operation to be within its scope.
+ * {@link DockerLinuxContainerRuntime} instance, a
+ * {@link JavaSandboxLinuxContainerRuntime} instance, or a custom instance
+ * depending on whether each instance believes the operation to be within its
+ * scope.
  *
- * @see DockerLinuxContainerRuntime#isDockerContainerRequested
+ * @see LinuxContainerRuntime#isRuntimeRequested
  */
 @InterfaceAudience.Private
 @InterfaceStability.Unstable
 public class DelegatingLinuxContainerRuntime implements LinuxContainerRuntime {
   private static final Logger LOG =
-       LoggerFactory.getLogger(DelegatingLinuxContainerRuntime.class);
+      LoggerFactory.getLogger(DelegatingLinuxContainerRuntime.class);
   private DefaultLinuxContainerRuntime defaultLinuxContainerRuntime;
   private DockerLinuxContainerRuntime dockerLinuxContainerRuntime;
   private JavaSandboxLinuxContainerRuntime javaSandboxLinuxContainerRuntime;
-  private EnumSet<LinuxContainerRuntimeConstants.RuntimeType> allowedRuntimes =
-      EnumSet.noneOf(LinuxContainerRuntimeConstants.RuntimeType.class);
+  private Set<String> allowedRuntimes = new HashSet<>();
+  private List<LinuxContainerRuntime> pluggableRuntimes = new ArrayList<>();
 
   @Override
   public void initialize(Configuration conf, Context nmContext)
@@ -64,52 +69,62 @@ public class DelegatingLinuxContainerRuntime implements LinuxContainerRuntime {
         YarnConfiguration.LINUX_CONTAINER_RUNTIME_ALLOWED_RUNTIMES,
         YarnConfiguration.DEFAULT_LINUX_CONTAINER_RUNTIME_ALLOWED_RUNTIMES);
     for (String configuredRuntime : configuredRuntimes) {
-      try {
-        allowedRuntimes.add(
-            LinuxContainerRuntimeConstants.RuntimeType.valueOf(
-                configuredRuntime.toUpperCase()));
-      } catch (IllegalArgumentException e) {
-        throw new ContainerExecutionException("Invalid runtime set in "
-            + YarnConfiguration.LINUX_CONTAINER_RUNTIME_ALLOWED_RUNTIMES + " : "
-            + configuredRuntime);
+      String normRuntime = configuredRuntime.toUpperCase();
+      allowedRuntimes.add(normRuntime);
+      if (isPluggableRuntime(normRuntime)) {
+        LinuxContainerRuntime runtime = createPluggableRuntime(conf,
+            configuredRuntime);
+        runtime.initialize(conf, nmContext);
+        pluggableRuntimes.add(runtime);
       }
     }
     if (isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.JAVASANDBOX)) {
+        LinuxContainerRuntimeConstants.RuntimeType.JAVASANDBOX.name())) {
       javaSandboxLinuxContainerRuntime = new JavaSandboxLinuxContainerRuntime(
           PrivilegedOperationExecutor.getInstance(conf));
       javaSandboxLinuxContainerRuntime.initialize(conf, nmContext);
     }
     if (isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.DOCKER)) {
+        LinuxContainerRuntimeConstants.RuntimeType.DOCKER.name())) {
       dockerLinuxContainerRuntime = new DockerLinuxContainerRuntime(
           PrivilegedOperationExecutor.getInstance(conf));
       dockerLinuxContainerRuntime.initialize(conf, nmContext);
     }
     if (isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.DEFAULT)) {
+        LinuxContainerRuntimeConstants.RuntimeType.DEFAULT.name())) {
       defaultLinuxContainerRuntime = new DefaultLinuxContainerRuntime(
           PrivilegedOperationExecutor.getInstance(conf));
       defaultLinuxContainerRuntime.initialize(conf, nmContext);
     }
   }
 
+  @Override
+  public boolean isRuntimeRequested(Map<String, String> env) {
+    return true;
+  }
+
   @VisibleForTesting
   LinuxContainerRuntime pickContainerRuntime(
       Map<String, String> environment) throws ContainerExecutionException {
     LinuxContainerRuntime runtime;
     //Sandbox checked first to ensure DockerRuntime doesn't circumvent controls
     if (javaSandboxLinuxContainerRuntime != null &&
-        javaSandboxLinuxContainerRuntime.isSandboxContainerRequested()){
+        javaSandboxLinuxContainerRuntime.isRuntimeRequested(environment)){
       runtime = javaSandboxLinuxContainerRuntime;
     } else if (dockerLinuxContainerRuntime != null &&
-        DockerLinuxContainerRuntime.isDockerContainerRequested(environment)){
+        dockerLinuxContainerRuntime.isRuntimeRequested(environment)) {
       runtime = dockerLinuxContainerRuntime;
-    } else if (defaultLinuxContainerRuntime != null &&
-        !DockerLinuxContainerRuntime.isDockerContainerRequested(environment)) {
-      runtime = defaultLinuxContainerRuntime;
     } else {
-      throw new ContainerExecutionException("Requested runtime not allowed.");
+      LinuxContainerRuntime pluggableRuntime = pickPluggableRuntime(
+          environment);
+      if (pluggableRuntime != null) {
+        runtime = pluggableRuntime;
+      } else if (defaultLinuxContainerRuntime != null &&
+          defaultLinuxContainerRuntime.isRuntimeRequested(environment)) {
+        runtime = defaultLinuxContainerRuntime;
+      } else {
+        throw new ContainerExecutionException("Requested runtime not allowed.");
+      }
     }
 
     if (LOG.isDebugEnabled()) {
@@ -120,6 +135,16 @@ public class DelegatingLinuxContainerRuntime implements LinuxContainerRuntime {
     return runtime;
   }
 
+  private LinuxContainerRuntime pickPluggableRuntime(
+      Map<String, String> environment) {
+    for (LinuxContainerRuntime runtime : pluggableRuntimes) {
+      if (runtime.isRuntimeRequested(environment)) {
+        return runtime;
+      }
+    }
+    return null;
+  }
+
   private LinuxContainerRuntime pickContainerRuntime(Container container)
       throws ContainerExecutionException {
     return pickContainerRuntime(container.getLaunchContext().getEnvironment());
@@ -175,9 +200,33 @@ public class DelegatingLinuxContainerRuntime implements LinuxContainerRuntime {
     return runtime.getIpAndHost(container);
   }
 
+  private boolean isPluggableRuntime(String runtimeType) {
+    for (LinuxContainerRuntimeConstants.RuntimeType type :
+        LinuxContainerRuntimeConstants.RuntimeType.values()) {
+      if (type.name().equalsIgnoreCase(runtimeType)) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  private LinuxContainerRuntime createPluggableRuntime(Configuration conf,
+      String runtimeType) throws ContainerExecutionException {
+    String confKey = String.format(
+        YarnConfiguration.LINUX_CONTAINER_RUNTIME_CLASS_FMT, runtimeType);
+    Class<? extends LinuxContainerRuntime> clazz = conf.getClass(
+        confKey, null, LinuxContainerRuntime.class);
+    if (clazz == null) {
+      throw new ContainerExecutionException("Invalid runtime set in "
+          + YarnConfiguration.LINUX_CONTAINER_RUNTIME_ALLOWED_RUNTIMES + " : "
+          + runtimeType + " : Missing configuration " + confKey);
+    }
+    return ReflectionUtils.newInstance(clazz, conf);
+  }
+
   @VisibleForTesting
-  boolean isRuntimeAllowed(
-      LinuxContainerRuntimeConstants.RuntimeType runtimeType) {
-    return allowedRuntimes.contains(runtimeType);
+  boolean isRuntimeAllowed(String runtimeType) {
+    return runtimeType != null && allowedRuntimes.contains(
+        runtimeType.toUpperCase());
   }
 }
\ No newline at end of file
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/linux/runtime/DockerLinuxContainerRuntime.java b/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
index 0ae3d0fa755..0665269339b 100644
--- a/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
+++ b/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
@@ -373,6 +373,11 @@ public class DockerLinuxContainerRuntime implements LinuxContainerRuntime {
         YarnConfiguration.NM_DOCKER_DEFAULT_TMPFS_MOUNTS)));
   }
 
+  @Override
+  public boolean isRuntimeRequested(Map<String, String> env) {
+    return isDockerContainerRequested(env);
+  }
+
   private Set<String> getDockerCapabilitiesFromConf() throws
       ContainerExecutionException {
     Set<String> caps = new HashSet<>(Arrays.asList(
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/linux/runtime/JavaSandboxLinuxContainerRuntime.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/JavaSandboxLinuxContainerRuntime.java
index 947c4bd8be8..dcc86b43bcc 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/JavaSandboxLinuxContainerRuntime.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/JavaSandboxLinuxContainerRuntime.java
@@ -282,9 +282,11 @@ public class JavaSandboxLinuxContainerRuntime
    * Determine if JVMSandboxLinuxContainerRuntime should be used.  This is
    * decided based on the value of
    * {@value YarnConfiguration#YARN_CONTAINER_SANDBOX}
+   * @param env the environment variable settings for the operation
    * @return true if Sandbox is requested, false otherwise
    */
-  boolean isSandboxContainerRequested() {
+  @Override
+  public boolean isRuntimeRequested(Map<String, String> env) {
     return sandboxMode != SandboxMode.disabled;
   }
 
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/linux/runtime/LinuxContainerRuntime.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/LinuxContainerRuntime.java
index 1203ce54f04..7381f7a2e2a 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/LinuxContainerRuntime.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/LinuxContainerRuntime.java
@@ -27,6 +27,8 @@ import org.apache.hadoop.yarn.server.nodemanager.Context;
 import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.ContainerExecutionException;
 import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.ContainerRuntime;
 
+import java.util.Map;
+
 /**
  * Linux-specific container runtime implementations must implement this
  * interface.
@@ -44,5 +46,14 @@ public interface LinuxContainerRuntime extends ContainerRuntime {
    * the runtime
    */
   void initialize(Configuration conf, Context nmContext) throws ContainerExecutionException;
+
+  /**
+   * Return whether the given environment variables indicate that the operation
+   * is requesting this runtime.
+   *
+   * @param env the environment variable settings for the operation
+   * @return whether this runtime is requested
+   */
+  boolean isRuntimeRequested(Map<String, String> env);
 }
 
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/linux/runtime/MockLinuxContainerRuntime.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/MockLinuxContainerRuntime.java
new file mode 100644
index 00000000000..6bfaed10f3f
--- /dev/null
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/MockLinuxContainerRuntime.java
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.runtime;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.yarn.server.nodemanager.Context;
+import org.apache.hadoop.yarn.server.nodemanager.containermanager.container.Container;
+import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.ContainerRuntimeConstants;
+import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.ContainerRuntimeContext;
+
+import java.util.Map;
+
+public class MockLinuxContainerRuntime implements LinuxContainerRuntime {
+  @Override
+  public void initialize(Configuration conf, Context nmContext) {}
+
+  @Override
+  public boolean isRuntimeRequested(Map<String, String> env) {
+    if (env == null) {
+      return false;
+    }
+
+    String type = env.get(ContainerRuntimeConstants.ENV_CONTAINER_TYPE);
+
+    return type != null && type.equals("mock");
+  }
+
+  @Override
+  public void prepareContainer(ContainerRuntimeContext ctx) {}
+
+  @Override
+  public void launchContainer(ContainerRuntimeContext ctx) {}
+
+  @Override
+  public void relaunchContainer(ContainerRuntimeContext ctx) {}
+
+  @Override
+  public void signalContainer(ContainerRuntimeContext ctx) {}
+
+  @Override
+  public void reapContainer(ContainerRuntimeContext ctx) {}
+
+  @Override
+  public String[] getIpAndHost(Container container) {
+    return new String[0];
+  }
+}
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/linux/runtime/TestDelegatingLinuxContainerRuntime.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/TestDelegatingLinuxContainerRuntime.java
index 907b1222ed6..fc34d9cb786 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/TestDelegatingLinuxContainerRuntime.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/TestDelegatingLinuxContainerRuntime.java
@@ -18,6 +18,7 @@ package org.apache.hadoop.yarn.server.nodemanager.containermanager.linux.runtime
 
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
+import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.ContainerExecutionException;
 import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.ContainerRuntime;
 import org.apache.hadoop.yarn.server.nodemanager.containermanager.runtime.ContainerRuntimeConstants;
 import org.junit.Before;
@@ -52,11 +53,11 @@ public class TestDelegatingLinuxContainerRuntime {
         YarnConfiguration.LINUX_CONTAINER_RUNTIME_ALLOWED_RUNTIMES));
     delegatingLinuxContainerRuntime.initialize(conf, null);
     assertTrue(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.DEFAULT));
+        LinuxContainerRuntimeConstants.RuntimeType.DEFAULT.name()));
     assertFalse(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.DOCKER));
+        LinuxContainerRuntimeConstants.RuntimeType.DOCKER.name()));
     assertFalse(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.JAVASANDBOX));
+        LinuxContainerRuntimeConstants.RuntimeType.JAVASANDBOX.name()));
   }
 
   @Test
@@ -65,11 +66,11 @@ public class TestDelegatingLinuxContainerRuntime {
         "docker");
     delegatingLinuxContainerRuntime.initialize(conf, null);
     assertTrue(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.DOCKER));
+        LinuxContainerRuntimeConstants.RuntimeType.DOCKER.name()));
     assertFalse(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.DEFAULT));
+        LinuxContainerRuntimeConstants.RuntimeType.DEFAULT.name()));
     assertFalse(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.JAVASANDBOX));
+        LinuxContainerRuntimeConstants.RuntimeType.JAVASANDBOX.name()));
   }
 
   @Test
@@ -78,11 +79,11 @@ public class TestDelegatingLinuxContainerRuntime {
         "javasandbox");
     delegatingLinuxContainerRuntime.initialize(conf, null);
     assertTrue(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.JAVASANDBOX));
+        LinuxContainerRuntimeConstants.RuntimeType.JAVASANDBOX.name()));
     assertFalse(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.DEFAULT));
+        LinuxContainerRuntimeConstants.RuntimeType.DEFAULT.name()));
     assertFalse(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.DOCKER));
+        LinuxContainerRuntimeConstants.RuntimeType.DOCKER.name()));
   }
 
   @Test
@@ -91,11 +92,11 @@ public class TestDelegatingLinuxContainerRuntime {
         "docker,javasandbox");
     delegatingLinuxContainerRuntime.initialize(conf, null);
     assertTrue(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.DOCKER));
+        LinuxContainerRuntimeConstants.RuntimeType.DOCKER.name()));
     assertTrue(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.JAVASANDBOX));
+        LinuxContainerRuntimeConstants.RuntimeType.JAVASANDBOX.name()));
     assertFalse(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.DEFAULT));
+        LinuxContainerRuntimeConstants.RuntimeType.DEFAULT.name()));
   }
 
   @Test
@@ -104,11 +105,38 @@ public class TestDelegatingLinuxContainerRuntime {
         "default,docker,javasandbox");
     delegatingLinuxContainerRuntime.initialize(conf, null);
     assertTrue(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.DEFAULT));
+        LinuxContainerRuntimeConstants.RuntimeType.DEFAULT.name()));
     assertTrue(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.DOCKER));
+        LinuxContainerRuntimeConstants.RuntimeType.DOCKER.name()));
     assertTrue(delegatingLinuxContainerRuntime.isRuntimeAllowed(
-        LinuxContainerRuntimeConstants.RuntimeType.JAVASANDBOX));
+        LinuxContainerRuntimeConstants.RuntimeType.JAVASANDBOX.name()));
+  }
+
+  @Test
+  public void testInitializeMissingRuntimeClass() throws Exception {
+    conf.set(YarnConfiguration.LINUX_CONTAINER_RUNTIME_ALLOWED_RUNTIMES,
+        "mock");
+    try {
+      delegatingLinuxContainerRuntime.initialize(conf, null);
+      fail("initialize should fail");
+    } catch (ContainerExecutionException e) {
+      assert(e.getMessage().contains("Invalid runtime set"));
+    }
+  }
+  @Test
+  public void testIsRuntimeAllowedMock() throws Exception {
+    conf.set(YarnConfiguration.LINUX_CONTAINER_RUNTIME_ALLOWED_RUNTIMES,
+        "mock");
+    conf.set(String.format(YarnConfiguration.LINUX_CONTAINER_RUNTIME_CLASS_FMT,
+        "mock"), MockLinuxContainerRuntime.class.getName());
+    delegatingLinuxContainerRuntime.initialize(conf, null);
+    assertFalse(delegatingLinuxContainerRuntime.isRuntimeAllowed(
+        LinuxContainerRuntimeConstants.RuntimeType.DEFAULT.name()));
+    assertFalse(delegatingLinuxContainerRuntime.isRuntimeAllowed(
+        LinuxContainerRuntimeConstants.RuntimeType.DOCKER.name()));
+    assertFalse(delegatingLinuxContainerRuntime.isRuntimeAllowed(
+        LinuxContainerRuntimeConstants.RuntimeType.JAVASANDBOX.name()));
+    assertTrue(delegatingLinuxContainerRuntime.isRuntimeAllowed("mock"));
   }
 
   @Test
@@ -134,4 +162,17 @@ public class TestDelegatingLinuxContainerRuntime {
         delegatingLinuxContainerRuntime.pickContainerRuntime(env);
     assertTrue(runtime instanceof DockerLinuxContainerRuntime);
   }
+
+  @Test
+  public void testMockRuntimeSelected() throws Exception {
+    env.put(ContainerRuntimeConstants.ENV_CONTAINER_TYPE, "mock");
+    conf.set(String.format(YarnConfiguration.LINUX_CONTAINER_RUNTIME_CLASS_FMT,
+        "mock"), MockLinuxContainerRuntime.class.getName());
+    conf.set(YarnConfiguration.LINUX_CONTAINER_RUNTIME_ALLOWED_RUNTIMES,
+        "mock");
+    delegatingLinuxContainerRuntime.initialize(conf, null);
+    ContainerRuntime runtime =
+        delegatingLinuxContainerRuntime.pickContainerRuntime(env);
+    assertTrue(runtime instanceof MockLinuxContainerRuntime);
+  }
 }
\ No newline at end of file