Exposed YARN node address in the ProtoBuf and Python API

Closes #38
jcrist · Jul 16, 2018 · 7055a8f · 7055a8f
1 parent b0da468
commit 7055a8f
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 6 deletions.
diff --git a/java/src/main/java/com/anaconda/skein/MsgUtils.java b/java/src/main/java/com/anaconda/skein/MsgUtils.java
@@ -7,6 +7,7 @@
 import org.apache.hadoop.yarn.api.records.LocalResource;
 import org.apache.hadoop.yarn.api.records.LocalResourceType;
 import org.apache.hadoop.yarn.api.records.LocalResourceVisibility;
+import org.apache.hadoop.yarn.api.records.NodeId;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.api.records.URL;
 import org.apache.hadoop.yarn.api.records.YarnApplicationState;
@@ -331,6 +332,10 @@ public static Msg.Container writeContainer(Model.Container container) {
     if (containerId != null) {
       builder.setYarnContainerId(containerId.toString());
     }
+    NodeId nodeId = container.getYarnNodeId();
+    if (nodeId != null) {
+      builder.setYarnNodeAddress(nodeId.getHost() + ":" + nodeId.getPort());
+    }
     return builder.build();
   }
 

diff --git a/java/src/main/proto/skein.proto b/java/src/main/proto/skein.proto
@@ -104,8 +104,9 @@ message Container {
   int32 instance = 2;
   State state = 3;
   string yarn_container_id = 4;
-  int64 start_time = 5;
-  int64 finish_time = 6;
+  string yarn_node_address = 5;
+  int64 start_time = 6;
+  int64 finish_time = 7;
 }
 
 

diff --git a/skein/model.py b/skein/model.py
@@ -1038,23 +1038,26 @@ class Container(Base):
         The current container state.
     yarn_container_id : str
         The YARN container id.
+    yarn_node_address : str
+        The YARN node address given as ``host:port``.
     start_time : datetime
         The start time, None if container has not started.
     finish_time : datetime
         The finish time, None if container has not finished.
     """
     __slots__ = ('service_name', 'instance', '_state', 'yarn_container_id',
-                 'start_time', 'finish_time')
+                 'yarn_node_address', 'start_time', 'finish_time')
     _params = ('service_name', 'instance', 'state', 'yarn_container_id',
-               'start_time', 'finish_time')
+               'yarn_node_address', 'start_time', 'finish_time')
     _protobuf_cls = _proto.Container
 
     def __init__(self, service_name, instance, state, yarn_container_id,
-                 start_time, finish_time):
+                 yarn_node_address, start_time, finish_time):
         self.service_name = service_name
         self.instance = instance
         self.state = state
         self.yarn_container_id = yarn_container_id
+        self.yarn_node_address = yarn_node_address
         self.start_time = start_time
         self.finish_time = finish_time
 
@@ -1077,6 +1080,7 @@ def _validate(self):
         self._check_is_type('instance', integer)
         self._check_is_type('state', ContainerState)
         self._check_is_type('yarn_container_id', string)
+        self._check_is_type('yarn_node_address', string)
         self._check_is_type('start_time', datetime, nullable=True)
         self._check_is_type('finish_time', datetime, nullable=True)
 
@@ -1098,6 +1102,7 @@ def from_dict(cls, obj):
                    instance=obj['instance'],
                    state=ContainerState(obj['state']),
                    yarn_container_id=obj['yarn_container_id'],
+                   yarn_node_address=obj['yarn_node_address'],
                    start_time=_datetime_from_millis(obj.get('start_time')),
                    finish_time=_datetime_from_millis(obj.get('finish_time')))
 
@@ -1108,5 +1113,6 @@ def from_protobuf(cls, obj):
                    instance=obj.instance,
                    state=ContainerState(_proto.Container.State.Name(obj.state)),
                    yarn_container_id=obj.yarn_container_id,
+                   yarn_node_address=obj.yarn_node_address,
                    start_time=_datetime_from_millis(obj.start_time),
                    finish_time=_datetime_from_millis(obj.finish_time))
diff --git a/skein/test/test_model.py b/skein/test/test_model.py
@@ -348,7 +348,8 @@ def test_container():
 
     kwargs = dict(service_name="foo",
                   instance=0,
-                  yarn_container_id='container_1528138529205_0038_01_000001')
+                  yarn_container_id='container_1528138529205_0038_01_000001',
+                  yarn_node_address='localhost:14420')
 
     c = Container(state='RUNNING',
                   start_time=start,