Repository: jcrist/skein
Pull request #119 (2018-12-19 20:01:13)

Driver works in YARN container
Two fixes to make the Skein driver/client work properly inside a YARN container:
- Properly set `io.netty.native.workdir` for the driver as well as the application master
- Load the current application's security configuration by default when run inside a container

This allows applications to be submitted from inside a YARN container just as easily as from the edge node.


Patch:
diff --git a/java/src/main/java/com/anaconda/skein/ApplicationMaster.java b/java/src/main/java/com/anaconda/skein/ApplicationMaster.java
index 3ff27e3..954be65 100644
--- a/java/src/main/java/com/anaconda/skein/ApplicationMaster.java
+++ b/java/src/main/java/com/anaconda/skein/ApplicationMaster.java
@@ -538,7 +538,7 @@ public class ApplicationMaster {
   public static void main(String[] args) {
     // Specify the netty native workdir. This is necessary for systems where
     // `/tmp` is not executable.
-    System.setProperty("io.netty.native.workdir", "./");
+    Utils.configureNettyNativeWorkDir();
 
     ApplicationMaster appMaster = new ApplicationMaster();
 
diff --git a/java/src/main/java/com/anaconda/skein/Driver.java b/java/src/main/java/com/anaconda/skein/Driver.java
index b8c27f1..c773141 100644
--- a/java/src/main/java/com/anaconda/skein/Driver.java
+++ b/java/src/main/java/com/anaconda/skein/Driver.java
@@ -159,7 +159,10 @@ public class Driver {
 
   /** Main Entry Point. **/
   public static void main(String[] args) {
-    boolean result = false;
+    // Maybe specify the netty native workdir. This is necessary for systems
+    // where `/tmp` is not executable.
+    Utils.configureNettyNativeWorkDir();
+
     try {
       Driver driver = new Driver();
       driver.init(args);
@@ -480,7 +483,7 @@ public class Driver {
 
     // Setup the LocalResources for the application master
     Map<String, LocalResource> lr = new HashMap<String, LocalResource>();
-    lr.put("skein.jar", newLocalResource(uploadCache, appDir, jarPath));
+    lr.put(".skein.jar", newLocalResource(uploadCache, appDir, jarPath));
     if (spec.getMaster().hasLogConfig()) {
       LocalResource logConfig = spec.getMaster().getLogConfig();
       finalizeLocalResource(uploadCache, appDir, logConfig, false);
diff --git a/java/src/main/java/com/anaconda/skein/Utils.java b/java/src/main/java/com/anaconda/skein/Utils.java
index 04d2577..d0462c8 100644
--- a/java/src/main/java/com/anaconda/skein/Utils.java
+++ b/java/src/main/java/com/anaconda/skein/Utils.java
@@ -86,6 +86,20 @@ public class Utils {
         new CustomThreadFactory(name, isDaemon));
   }
 
+  public static void configureNettyNativeWorkDir() {
+    final String property = "io.netty.native.workdir";
+
+    // Already set by user, do nothing
+    if (System.getProperty(property) != null) {
+      return;
+    }
+
+    // Set to current directory if we're in a container
+    if (System.getenv("CONTAINER_ID") != null) {
+      System.setProperty(property, "./");
+    }
+  }
+
   public static String formatAcl(List<String> users, List<String> groups) {
     // "*" in either groups or users -> "*"
     if (users.contains("*") || groups.contains("*")) {
diff --git a/skein/core.py b/skein/core.py
index d1ff0bf..97b95a9 100644
--- a/skein/core.py
+++ b/skein/core.py
@@ -14,9 +14,10 @@ import grpc
 
 from . import proto
 from .compatibility import PY2, makedirs, isidentifier, Mapping
-from .exceptions import (context, FileNotFoundError, ConnectionError,
-                         TimeoutError, ApplicationNotRunningError,
-                         ApplicationError, DriverNotRunningError, DriverError)
+
+from .exceptions import (context, ConnectionError, TimeoutError,
+                         ApplicationNotRunningError, ApplicationError,
+                         DriverNotRunningError, DriverError)
 from .kv import KeyValueStore
 from .ui import WebUI
 from .model import (Security, ApplicationSpec, ApplicationReport,
@@ -756,25 +757,12 @@ class ApplicationClient(_ClientBase):
         Useful for connecting to the application master from a running
         container in a application.
         """
-        address = properties.appmaster_address
-        app_id = properties.application_id
-        yarn_container_id = properties.yarn_container_id
-        local_dirs = os.environ.get('LOCAL_DIRS')
-
-        if any(p is None for p in [address, app_id, yarn_container_id,
-                                   local_dirs]):
+        if properties.application_id is None:
             raise context.ValueError("Not running inside a container")
 
-        for local_dir in local_dirs.split(','):
-            container_dir = os.path.join(local_dir, yarn_container_id)
-            crt_path = os.path.join(container_dir, '.skein.crt')
-            pem_path = os.path.join(container_dir, '.skein.pem')
-            if os.path.exists(crt_path) and os.path.exists(pem_path):
-                security = Security(cert_file=crt_path, key_file=pem_path)
-                return cls(address, app_id, security=security)
-
-        raise FileNotFoundError(
-            "Failed to resolve .skein.{crt,pem} in 'LOCAL_DIRS'")
+        return cls(properties.appmaster_address,
+                   properties.application_id,
+                   security=Security.from_default())
 
     def get_containers(self, services=None, states=None):
         """Get information on containers in this application.
diff --git a/skein/model.py b/skein/model.py
index 8cd457a..658a85e 100644
--- a/skein/model.py
+++ b/skein/model.py
@@ -309,8 +309,30 @@ class Security(Specification):
 
     @classmethod
     def from_default(cls):
-        """The default security configuration."""
+        """The default security configuration.
+
+        Usually this loads the credentials stored in the configuration
+        directory (``~/.skein`` by default). If these credentials don't already
+        exist, new ones will be created.
+
+        When run in a YARN container started by Skein, this loads the same
+        security credentials as used for the current application.
+        """
         from .core import properties
+
+        # Are we in a container started by skein?
+        if properties.application_id is not None:
+            # Find the container directory from all the options
+            for path in os.environ.get('LOCAL_DIRS', '').split(','):
+                container_dir = os.path.join(path, properties.yarn_container_id)
+                cert_path = os.path.join(container_dir, '.skein.crt')
+                key_path = os.path.join(container_dir, '.skein.pem')
+                if os.path.exists(cert_path) and os.path.exists(key_path):
+                    return Security(cert_file=cert_path, key_file=key_path)
+            raise context.FileNotFoundError(
+                "Failed to resolve .skein.{crt,pem} in 'LOCAL_DIRS'")
+
+        # Try to load from config_dir, and fallback to minting new credentials
         try:
             return cls.from_directory(properties.config_dir)
         except FileNotFoundError:


Test patch:
diff --git a/skein/test/test_core.py b/skein/test/test_core.py
index 37f4a33..322e50e 100644
--- a/skein/test/test_core.py
+++ b/skein/test/test_core.py
@@ -252,7 +252,7 @@ def test_client_login_from_keytab(security, not_logged_in):
         skein.Client(keytab=KEYTAB_PATH, security=security)
 
 
-def test_application_client_from_current(monkeypatch, tmpdir, security):
+def test_appclient_and_security_in_container(monkeypatch, tmpdir, security):
     # Not running in a container
     with pytest.raises(ValueError) as exc:
         skein.ApplicationClient.from_current()
@@ -277,6 +277,10 @@ def test_application_client_from_current(monkeypatch, tmpdir, security):
         skein.ApplicationClient.from_current()
     assert str(exc.value) == "Failed to resolve .skein.{crt,pem} in 'LOCAL_DIRS'"
 
+    with pytest.raises(FileNotFoundError) as exc:
+        skein.Security.from_default()
+    assert str(exc.value) == "Failed to resolve .skein.{crt,pem} in 'LOCAL_DIRS'"
+
     # Add proper LOCAL_DIRS environment
     good_dir = tmpdir.mkdir('good_dir')
     local_dir = good_dir.mkdir(container_id)
@@ -291,6 +295,10 @@ def test_application_client_from_current(monkeypatch, tmpdir, security):
     assert app.id == app_id
     assert app.address == address
 
+    security2 = skein.Security.from_default()
+    assert security._get_bytes('key') == security2._get_bytes('key')
+    assert security._get_bytes('cert') == security2._get_bytes('cert')
+
 
 def test_simple_app(client):
     with run_application(client) as app: