Merge branch 'master' into 13804-no-shutdown-wanted-nodes

[arvados.git] / services / nodemanager / arvnodeman / config.py

diff --git a/services/nodemanager/arvnodeman/config.py b/services/nodemanager/arvnodeman/config.py
index 617b5626b335cccb92334363acb52ec1a694386e..4fda7e76d69ed87aa1b5b032f1e3b1b8e608f5b5 100644 (file)
--- a/services/nodemanager/arvnodeman/config.py
+++ b/services/nodemanager/arvnodeman/config.py
@@ -17,6 +17,7 @@ from apiclient import errors as apierror
 
 from .baseactor import BaseNodeManagerActor
 
+from functools import partial
 from libcloud.common.types import LibcloudError
 from libcloud.common.exceptions import BaseHTTPError
 
@@ -70,12 +71,23 @@ class NodeManagerConfig(ConfigParser.SafeConfigParser):
                 if not self.has_option(sec_name, opt_name):
                     self.set(sec_name, opt_name, value)
 
-    def get_section(self, section, transformer=None):
+    def get_section(self, section, transformers={}, default_transformer=None):
+        transformer_map = {
+            str: self.get,
+            int: self.getint,
+            bool: self.getboolean,
+            float: self.getfloat,
+        }
         result = self._dict()
         for key, value in self.items(section):
+            transformer = None
+            if transformers.get(key) in transformer_map:
+                transformer = partial(transformer_map[transformers[key]], section)
+            elif default_transformer in transformer_map:
+                transformer = partial(transformer_map[default_transformer], section)
             if transformer is not None:
                 try:
-                    value = transformer(value)
+                    value = transformer(key)
                 except (TypeError, ValueError):
                     pass
             result[key] = value
@@ -129,31 +141,41 @@ class NodeManagerConfig(ConfigParser.SafeConfigParser):
                                         self.get_section('Cloud Create'),
                                         driver_class=driver_class)
 
-    def node_sizes(self, all_sizes):
+    def node_sizes(self):
         """Finds all acceptable NodeSizes for our installation.
 
         Returns a list of (NodeSize, kwargs) pairs for each NodeSize object
         returned by libcloud that matches a size listed in our config file.
         """
-
+        all_sizes = self.new_cloud_client().list_sizes()
         size_kwargs = {}
+        section_types = {
+            'instance_type': str,
+            'price': float,
+            'preemptible': bool,
+        }
         for sec_name in self.sections():
             sec_words = sec_name.split(None, 2)
             if sec_words[0] != 'Size':
                 continue
-            size_spec = self.get_section(sec_name, int)
-            if 'price' in size_spec:
-                size_spec['price'] = float(size_spec['price'])
+            size_spec = self.get_section(sec_name, section_types, int)
+            if 'preemptible' not in size_spec:
+                size_spec['preemptible'] = False
+            if 'instance_type' not in size_spec:
+                # Assume instance type is Size name if missing
+                size_spec['instance_type'] = sec_words[1]
+            size_spec['id'] = sec_words[1]
             size_kwargs[sec_words[1]] = size_spec
         # EC2 node sizes are identified by id. GCE sizes are identified by name.
         matching_sizes = []
         for size in all_sizes:
-            if size.id in size_kwargs:
-                matching_sizes.append((size, size_kwargs[size.id]))
-            elif size.name in size_kwargs:
-                matching_sizes.append((size, size_kwargs[size.name]))
+            matching_sizes += [
+                (size, size_kwargs[s]) for s in size_kwargs
+                if size_kwargs[s]['instance_type'] == size.id
+                or size_kwargs[s]['instance_type'] == size.name
+            ]
         return matching_sizes
 
     def shutdown_windows(self):
-        return [int(n)
+        return [float(n)
                 for n in self.get('Cloud', 'shutdown_windows').split(',')]