Repository: jcrist/skein
Pull request #44 (2018-07-26 18:08:50)

Added support for YARN node label expressions
The node label expression can be defined either for the whole application

```python
skein.ApplicationSpec(node_label="gpu")
```

or for each service separately

```python
skein.Service(..., node_label="gpu")
```

superbobry: IIRC node labelling was introduced in Hadoop 2.7.X which is not available in CDH5. We have backported this feature to 2.6.X internally, but for the docker-compose, I guess we'd have to drop CDH to have the labelling tested. Are you OK with that?
jcrist: Hmmm. The methods appear to be available in 2.6.0 (see e.g. http://hadoop.apache.org/docs/r2.6.0/api/org/apache/hadoop/yarn/api/records/ApplicationSubmissionContext.html#setNodeLabelExpression(java.lang.String), and are found with `javap` on the test cdh5 install.
jcrist: [The IBM docs](https://developer.ibm.com/hadoop/2017/03/10/yarn-node-labels/) say it was introduced in 2.6.0, but not great to use until 2.8. I suspect that's mostly from the scheduler perspective, and the user api would be stable.
superbobry: Hmm, OK, thanks for the references. I will try to make it work in the hadoop-test-cluster. Stay tuned for updates.
superbobry: @jcrist do you think we can merge this PR given that the tests pass on a two-worker cluster? I can rebase on master.
jcrist: Yeah, apologies for the delay here. If you handle fixing the merge conflicts, I'll handle whatever else I see is needed to get this in. Thanks for being patient here.
superbobry: @jcrist I have merged with master (rebase was a bit more painful). I hope it's OK.
jcrist: Thanks for your patience @superbobry, I've pushed up some small changes:

- Removed the runtime tests. I wasn't happy with the hacks to get this to work with hadoop 2.6 (larger investigation showed these were just bugs in how hadoop 2.6 handled node labels), and the tests were *mostly* smoketests anyway. Since this is just forwarding parameters from the spec to the the container requests, I'm fine with testing by inspection only (though it's not ideal).
- Removed propagating the application level `node_label` ourselves. When I mentioned you should do that above I wasn't aware that YARN already does this internally. No need to duplicate work.
- Added docs

Merging, thanks again for all your contributions.

Patch:
diff --git a/docs/source/specification.rst b/docs/source/specification.rst
index 7a76d90..5ea0b50 100644
--- a/docs/source/specification.rst
+++ b/docs/source/specification.rst
@@ -31,6 +31,14 @@ The name of the application. Optional, defaults to ``skein``.
 
 The queue to submit the application to. Optional, defaults to ``default``.
 
+``node_label``
+~~~~~~~~~~~~~~
+
+The `node_label
+<https://hadoop.apache.org/docs/stable/hadoop-yarn/hadoop-yarn-site/NodeLabel.html>`__
+to request for all containers in this application. Services can override this
+value by setting ``node_label`` on the service directly. Default is no label.
+
 ``max_attempts``
 ~~~~~~~~~~~~~~~~
 
@@ -189,6 +197,28 @@ Optional, default is 1.
       my_service:
         instances: 4  # Start 4 instances
 
+``node_label``
+~~~~~~~~~~~~~~
+
+The `node_label
+<https://hadoop.apache.org/docs/stable/hadoop-yarn/hadoop-yarn-site/NodeLabel.html>`__
+to request for all containers in this service. If not set, defaults to the
+application-level ``node_label`` (if set).
+
+**Example**
+
+.. code-block:: none
+
+    node_label: mylabel
+
+    services:
+      my_service1:
+        node_label: GPU  # This service will be allocated on "GPU" nodes only.
+        ...
+
+      my_service2:
+        # node_label is not set, the application label "mylabel" will be used.
+        ...
 
 ``max_restarts``
 ~~~~~~~~~~~~~~~~
@@ -197,7 +227,7 @@ The maximum number of restarts allowed for this service. Must be >= -1. On
 failure, a container will be restarted if the total number of restarts for its
 service is < ``max_restarts``. Once this limit is exceeded, the service is
 marked as failed and the application will be terminated. Set to -1 to always
-restart, or 0 to never restart. Optinoal, default is 0.
+restart, or 0 to never restart. Optional, default is 0.
 
 **Example**
 
diff --git a/java/src/main/java/com/anaconda/skein/ApplicationMaster.java b/java/src/main/java/com/anaconda/skein/ApplicationMaster.java
index ebf8959..e9fcbfb 100644
--- a/java/src/main/java/com/anaconda/skein/ApplicationMaster.java
+++ b/java/src/main/java/com/anaconda/skein/ApplicationMaster.java
@@ -1,6 +1,7 @@
 package com.anaconda.skein;
 
 import com.google.common.base.Function;
+import com.google.common.base.Strings;
 import com.google.common.collect.Lists;
 import com.google.protobuf.ByteString;
 
@@ -837,8 +838,13 @@ public class ApplicationMaster {
 
     private synchronized void requestContainer(Model.Container container) {
       Priority priority = newPriority(this);
-      ContainerRequest req = new ContainerRequest(service.getResources(),
-                                                  null, null, priority);
+      ContainerRequest req = new ContainerRequest(
+          service.getResources(),
+          null,
+          null,
+          priority,
+          true,
+          Strings.emptyToNull(service.getNodeLabel()));
       container.setContainerRequest(req);
       rmClient.addContainerRequest(req);
       requested.put(priority, container);
diff --git a/java/src/main/java/com/anaconda/skein/Daemon.java b/java/src/main/java/com/anaconda/skein/Daemon.java
index 28521e9..75f16f5 100644
--- a/java/src/main/java/com/anaconda/skein/Daemon.java
+++ b/java/src/main/java/com/anaconda/skein/Daemon.java
@@ -1,5 +1,6 @@
 package com.anaconda.skein;
 
+import com.google.common.base.Strings;
 import com.google.common.collect.ObjectArrays;
 
 import io.grpc.Server;
@@ -301,6 +302,7 @@ public class Daemon {
     appContext.setPriority(Priority.newInstance(0));
     appContext.setQueue(spec.getQueue());
     appContext.setMaxAppAttempts(spec.getMaxAttempts());
+    appContext.setNodeLabelExpression(Strings.emptyToNull(spec.getNodeLabel()));
     appContext.setApplicationTags(spec.getTags());
 
     LOG.info("Submitting application...");
diff --git a/java/src/main/java/com/anaconda/skein/Model.java b/java/src/main/java/com/anaconda/skein/Model.java
index 0b7d36c..5b46c28 100644
--- a/java/src/main/java/com/anaconda/skein/Model.java
+++ b/java/src/main/java/com/anaconda/skein/Model.java
@@ -34,6 +34,7 @@ public class Model {
 
   public static class Service {
     private int instances;
+    private String nodeLabel;
     private int maxRestarts;
     private Resource resources;
     private Map<String, LocalResource> localResources;
@@ -44,6 +45,7 @@ public class Model {
     public Service() {}
 
     public Service(int instances,
+                   String nodeLabel,
                    int maxRestarts,
                    Resource resources,
                    Map<String, LocalResource> localResources,
@@ -51,6 +53,7 @@ public class Model {
                    List<String> commands,
                    Set<String> depends) {
       this.instances = instances;
+      this.nodeLabel = nodeLabel;
       this.maxRestarts = maxRestarts;
       this.resources = resources;
       this.localResources = localResources;
@@ -62,6 +65,7 @@ public class Model {
     public String toString() {
       return ("Service:\n"
               + "instances: " + instances + "\n"
+              + "nodeLabel: " + nodeLabel + "\n"
               + "maxRestarts: " + maxRestarts + "\n"
               + "resources: " + resources + "\n"
               + "localResources: " + localResources + "\n"
@@ -73,6 +77,9 @@ public class Model {
     public void setInstances(int instances) { this.instances = instances; }
     public int getInstances() { return instances; }
 
+    public void setNodeLabel(String nodeLabel) { this.nodeLabel = nodeLabel; }
+    public String getNodeLabel() { return nodeLabel; }
+
     public void setMaxRestarts(int maxRestarts) { this.maxRestarts = maxRestarts; }
     public int getMaxRestarts() { return maxRestarts; }
 
@@ -177,6 +184,7 @@ public class Model {
   public static class ApplicationSpec {
     private String name;
     private String queue;
+    private String nodeLabel;
     private int maxAttempts;
     private Set<String> tags;
     private List<Path> fileSystems;
@@ -186,12 +194,13 @@ public class Model {
 
     public ApplicationSpec() {}
 
-    public ApplicationSpec(String name, String queue, int maxAttempts,
+    public ApplicationSpec(String name, String queue, String nodeLabel, int maxAttempts,
                            Set<String> tags, List<Path> fileSystems,
                            Acls acls, Master master,
                            Map<String, Service> services) {
       this.name = name;
       this.queue = queue;
+      this.nodeLabel = nodeLabel;
       this.maxAttempts = maxAttempts;
       this.tags = tags;
       this.fileSystems = fileSystems;
@@ -204,6 +213,7 @@ public class Model {
       return ("ApplicationSpec<"
               + "name: " + name + ", "
               + "queue: " + queue + ", "
+              + "nodeLabel: " + nodeLabel + ", "
               + "maxAttempts: " + maxAttempts + ", "
               + "tags: " + tags + ", "
               + "fileSystems" + fileSystems + ", "
@@ -216,6 +226,9 @@ public class Model {
     public void setQueue(String queue) { this.queue = queue; }
     public String getQueue() { return queue; }
 
+    public void setNodeLabel(String nodeLabel) { this.nodeLabel = nodeLabel; }
+    public String getNodeLabel() { return nodeLabel; }
+
     public void setMaxAttempts(int maxAttempts) { this.maxAttempts = maxAttempts; }
     public int getMaxAttempts() { return maxAttempts; }
 
diff --git a/java/src/main/java/com/anaconda/skein/MsgUtils.java b/java/src/main/java/com/anaconda/skein/MsgUtils.java
index 63c04f4..06b2cf7 100644
--- a/java/src/main/java/com/anaconda/skein/MsgUtils.java
+++ b/java/src/main/java/com/anaconda/skein/MsgUtils.java
@@ -271,6 +271,7 @@ public class MsgUtils {
   public static Msg.Service writeService(Model.Service service) {
     Msg.Service.Builder builder = Msg.Service.newBuilder()
         .setInstances(service.getInstances())
+        .setNodeLabel(service.getNodeLabel())
         .setMaxRestarts(service.getMaxRestarts())
         .setResources(writeResources(service.getResources()))
         .putAllEnv(service.getEnv())
@@ -290,6 +291,7 @@ public class MsgUtils {
     }
     return new Model.Service(
         service.getInstances(),
+        service.getNodeLabel(),
         service.getMaxRestarts(),
         readResources(service.getResources()),
         localResources,
@@ -387,6 +389,7 @@ public class MsgUtils {
     Msg.ApplicationSpec.Builder builder = Msg.ApplicationSpec.newBuilder()
         .setName(spec.getName())
         .setQueue(spec.getQueue())
+        .setNodeLabel(spec.getNodeLabel())
         .setMaxAttempts(spec.getMaxAttempts())
         .addAllTags(spec.getTags())
         .setAcls(writeAcls(spec.getAcls()))
@@ -412,6 +415,7 @@ public class MsgUtils {
 
     return new Model.ApplicationSpec(spec.getName(),
                                      spec.getQueue(),
+                                     spec.getNodeLabel(),
                                      spec.getMaxAttempts(),
                                      new HashSet<String>(spec.getTagsList()),
                                      fileSystems,
diff --git a/java/src/main/proto/skein.proto b/java/src/main/proto/skein.proto
index 2139c9f..0b679f2 100644
--- a/java/src/main/proto/skein.proto
+++ b/java/src/main/proto/skein.proto
@@ -73,12 +73,13 @@ message File {
 
 message Service {
   int32 instances = 1;
-  int32 max_restarts = 2;
-  Resources resources = 3;
-  map<string, File> files = 4;
-  map<string, string> env = 5;
-  repeated string commands = 6;
-  repeated string depends = 7;
+  string node_label = 2;
+  int32 max_restarts = 3;
+  Resources resources = 4;
+  map<string, File> files = 5;
+  map<string, string> env = 6;
+  repeated string commands = 7;
+  repeated string depends = 8;
 }
 
 
@@ -115,12 +116,13 @@ message Master {
 message ApplicationSpec {
   string name = 1;
   string queue = 2;
-  int32 max_attempts = 3;
-  repeated string tags = 4;
-  repeated string file_systems = 5;
-  Acls acls = 6;
-  Master master = 7;
-  map<string, Service> services = 8;
+  string node_label = 3;
+  int32 max_attempts = 4;
+  repeated string tags = 5;
+  repeated string file_systems = 6;
+  Acls acls = 7;
+  Master master = 8;
+  map<string, Service> services = 9;
 }
 
 
diff --git a/skein/model.py b/skein/model.py
index 3119707..be63112 100644
--- a/skein/model.py
+++ b/skein/model.py
@@ -366,6 +366,10 @@ class Service(Specification):
         Describes the resources needed to run the service.
     instances : int, optional
         The number of instances to create on startup. Default is 1.
+    node_label : str, optional
+        The node label expression to use when requesting containers for this
+        service. If not set, defaults to the application-level ``node_label``
+        (if set).
     max_restarts : int, optional
         The maximum number of restarts to allow for this service. Containers
         are only restarted on failure, and the cap is set for all containers in
@@ -382,15 +386,17 @@ class Service(Specification):
         A set of service names that this service depends on. The service will
         only be started after all its dependencies have been started.
     """
-    __slots__ = ('commands', 'resources', 'instances', 'max_restarts', 'files',
-                 'env', 'depends')
+    __slots__ = ('commands', 'resources', 'instances', 'node_label',
+                 'max_restarts', 'files', 'env', 'depends')
     _protobuf_cls = _proto.Service
 
     def __init__(self, commands=required, resources=required, instances=1,
-                 max_restarts=0, files=None, env=None, depends=None):
+                 node_label='', max_restarts=0, files=None, env=None,
+                 depends=None):
         self._assign_required('commands', commands)
         self._assign_required('resources', resources)
         self.instances = instances
+        self.node_label = node_label
         self.max_restarts = max_restarts
         if files is not None:
             files = {k: v if isinstance(v, File) else File(v)
@@ -407,6 +413,7 @@ class Service(Specification):
 
     def _validate(self):
         self._check_is_bounded_int('instances', min=0)
+        self._check_is_type('node_label', string)
         self._check_is_bounded_int('max_restarts', min=-1)
 
         self._check_is_type('resources', Resources)
@@ -451,6 +458,7 @@ class Service(Specification):
         resources = Resources.from_protobuf(obj.resources)
         files = {k: File.from_protobuf(v) for k, v in obj.files.items()}
         kwargs = {'instances': obj.instances,
+                  'node_label': obj.node_label,
                   'max_restarts': obj.max_restarts,
                   'resources': resources,
                   'files': files,
@@ -677,6 +685,10 @@ class ApplicationSpec(Specification):
         The name of the application, defaults to 'skein'.
     queue : string, optional
         The queue to submit to. Defaults to the default queue.
+    node_label : string, optional
+        The node label expression to use when requesting containers for this
+        application. Services can override this setting by specifying
+        ``node_label`` on the service directly. Default is no label.
     tags : set, optional
         A set of strings to use as tags for this application.
     file_systems : list, optional
@@ -693,16 +705,17 @@ class ApplicationSpec(Specification):
         application as failed. Note that this only considers failures of the
         application master during startup. Default is 1.
     """
-    __slots__ = ('services', 'name', 'queue', 'tags', 'file_systems', 'acls',
-                 'master', 'max_attempts')
+    __slots__ = ('services', 'name', 'queue', 'node_label', 'tags',
+                 'file_systems', 'acls', 'master', 'max_attempts')
     _protobuf_cls = _proto.ApplicationSpec
 
     def __init__(self, services=required, name='skein', queue='default',
-                 tags=None, file_systems=None, acls=None, master=None,
+                 node_label='', tags=None, file_systems=None, acls=None, master=None,
                  max_attempts=1):
         self._assign_required('services', services)
         self.name = name
         self.queue = queue
+        self.node_label = node_label
         self.tags = set() if tags is None else set(tags)
         self.file_systems = [] if file_systems is None else file_systems
         self.acls = ACLs() if acls is None else acls
@@ -717,6 +730,7 @@ class ApplicationSpec(Specification):
     def _validate(self):
         self._check_is_type('name', string)
         self._check_is_type('queue', string)
+        self._check_is_type('node_label', string)
         self._check_is_set_of('tags', string)
         self._check_is_list_of('file_systems', string)
         self._check_is_bounded_int('max_attempts', min=1)
@@ -770,6 +784,7 @@ class ApplicationSpec(Specification):
                     for k, v in obj.services.items()}
         return cls(name=obj.name,
                    queue=obj.queue,
+                   node_label=obj.node_label,
                    tags=set(obj.tags),
                    file_systems=list(obj.file_systems),
                    max_attempts=min(1, obj.max_attempts),


Test patch:
diff --git a/skein/test/test_model.py b/skein/test/test_model.py
index f53ebe4..0462282 100644
--- a/skein/test/test_model.py
+++ b/skein/test/test_model.py
@@ -20,6 +20,7 @@ def indent(s, n):
 
 
 service_spec = """\
+node_label: gpu
 resources:
     vcores: 1
     memory: 1024
@@ -40,6 +41,7 @@ commands:
 app_spec = """\
 name: test
 queue: default
+node_label: cpu
 max_attempts: 2
 tags:
     - tag1
@@ -201,6 +203,7 @@ def test_service():
     r = Resources(memory=1024, vcores=1)
     c = ['commands']
     s1 = Service(resources=r, commands=c,
+                 node_label="testlabel",
                  files={'file': File(source='/test/path')})
     s2 = Service(resources=r, commands=c,
                  files={'file': File(source='/test/path', size=1024)})
@@ -254,7 +257,10 @@ def test_application_spec():
                  files={'file': File(source='/test/path')})
     s2 = Service(resources=r, commands=c,
                  files={'file': File(source='/test/path', size=1024)})
-    spec1 = ApplicationSpec(services={'service': s1})
+    spec1 = ApplicationSpec(name='test',
+                            queue='testqueue',
+                            node_label='testlabel',
+                            services={'service': s1})
     spec2 = ApplicationSpec(services={'service': s2})
     check_specification_methods(spec1, spec2)
 
@@ -299,6 +305,8 @@ def test_service_from_yaml():
     s = Service.from_yaml(service_spec)
     assert isinstance(s, Service)
 
+    assert s.node_label == 'gpu'
+
     assert s.resources.vcores == 1
     assert s.resources.memory == 1024
 
@@ -330,6 +338,7 @@ def test_application_spec_from_yaml():
 
     assert spec.name == 'test'
     assert spec.queue == 'default'
+    assert spec.node_label == 'cpu'
     assert spec.tags == {'tag1', 'tag2'}
     assert spec.file_systems == ['hdfs://preprod']
     assert spec.max_attempts == 2