7478: Unify arvados node size access across different drivers

[arvados.git] / services / nodemanager / arvnodeman / computenode / driver / gce.py

diff --git a/services/nodemanager/arvnodeman/computenode/driver/gce.py b/services/nodemanager/arvnodeman/computenode/driver/gce.py
index 419557fe288ded9c4c4706bcd47a58b035ce2e65..f8d81cab041d974ac253e3d11f68e95dcfd1abd3 100644 (file)
--- a/services/nodemanager/arvnodeman/computenode/driver/gce.py
+++ b/services/nodemanager/arvnodeman/computenode/driver/gce.py
@@ -38,7 +38,7 @@ class ComputeNodeDriver(BaseComputeNodeDriver):
         super(ComputeNodeDriver, self).__init__(
             auth_kwargs, list_kwargs, create_kwargs,
             driver_class)
-        self._sizes_by_name = {sz.name: sz for sz in self.sizes.itervalues()}
+        self._sizes_by_id = {sz.id: sz for sz in self.sizes.itervalues()}
         self._disktype_links = {dt.name: self._object_link(dt)
                                 for dt in self.real.ex_list_disktypes()}
 
@@ -109,6 +109,11 @@ class ComputeNodeDriver(BaseComputeNodeDriver):
         return result
 
 
+    def create_node(self, size, arvados_node):
+        # Set up tag indicating the Arvados assigned Cloud Size id.
+        self.create_kwargs['ex_metadata'].update({'arvados_node_size': size.id})
+        return super(ComputeNodeDriver, self).create_node(size, arvados_node)
+
     def list_nodes(self):
         # The GCE libcloud driver only supports filtering node lists by zone.
         # Do our own filtering based on tag list.
@@ -120,7 +125,8 @@ class ComputeNodeDriver(BaseComputeNodeDriver):
         # and monkeypatch the results when that's the case.
         if nodelist and not hasattr(nodelist[0].size, 'id'):
             for node in nodelist:
-                node.size = self._sizes_by_name[node.size]
+                node.size = self._sizes_by_id[node.size]
+                node.extra['arvados_node_size'] = node.extra.get('metadata', {}).get('arvados_node_size')
         return nodelist
 
     @classmethod