Support development instance_types_data configuration

This commit integrates the support for the new `instance_types_data` configuration parameter, allowing custom instance types data to be used from ParallelCluster daemons.

Signed-off-by: ddeidda <[email protected]>

Author: ddeidda

src/common/utils.py

@@ -11,6 +11,7 @@
 # See the License for the specific language governing permissions and limitations under the License.
 import collections
 import itertools
+import json
 import logging
 import os
 import pwd
@@ -274,20 +275,29 @@ def _fetch_instance_info(region, proxy_config, instance_type):
         raise CriticalError(emsg)
 
 
-def _get_instance_info(region, proxy_config, instance_type):
+def _get_instance_info(region, proxy_config, instance_type, additional_instance_types_data=None):
     """
     Call the DescribeInstanceTypes to get number of vcpus and gpus for the given instance type.
 
     :return: (the number of vcpus or -1 if the instance type cannot be found,
                 number of gpus or None if the instance does not have gpu)
     """
-    log.debug("Fetching info for instance_type {0}".format(instance_type))
-    instance_info = _fetch_instance_info(region, proxy_config, instance_type)
-    log.debug("Received the following information for instance type {0}: {1}".format(instance_type, instance_info))
+    instance_info = None
+
+    # First attempt to describe the instance is from configuration data, if present
+    if additional_instance_types_data:
+        instance_info = additional_instance_types_data.get(instance_type, None)
+
+    # If no data is provided from configuration we retrieve it from ec2
+    if not instance_info:
+        log.debug("Fetching info for instance_type {0}".format(instance_type))
+        instance_info = _fetch_instance_info(region, proxy_config, instance_type)
+        log.debug("Received the following information for instance type {0}: {1}".format(instance_type, instance_info))
+
     return _get_vcpus_from_instance_info(instance_info), _get_gpus_from_instance_info(instance_info)
 
 
-def get_instance_properties(region, proxy_config, instance_type):
+def get_instance_properties(region, proxy_config, instance_type, additional_instance_types_data=None):
     """
     Get instance properties for the given instance type, according to the cfn_scheduler_slots configuration parameter.
 
@@ -299,7 +309,7 @@ def get_instance_properties(region, proxy_config, instance_type):
 
     if instance_type not in get_instance_properties.cache:
         # get vcpus and gpus from the pricing file, gpus = 0 if instance does not have GPU
-        vcpus, gpus = _get_instance_info(region, proxy_config, instance_type)
+        vcpus, gpus = _get_instance_info(region, proxy_config, instance_type, additional_instance_types_data)
 
         try:
             cfnconfig_params = _read_cfnconfig()
@@ -421,3 +431,30 @@ def grouper(iterable, n):
         if not chunk:
             return
         yield chunk
+
+
+def load_additional_instance_types_data(config, section):
+    """Load instance types data from configuration, if set; an empty dict is returned otherwise."""
+    instance_types_data = {}
+    if config.has_option(section, "instance_types_data"):
+        instance_types_data_str = config.get(section, "instance_types_data")
+        if instance_types_data_str:
+            try:
+                instance_types_data_str = str(instance_types_data_str).strip()
+
+                # Load json value if not empty
+                if instance_types_data_str:
+                    instance_types_data = json.loads(instance_types_data_str)
+
+                # Fallback to empty dict if value is None
+                if not instance_types_data:
+                    instance_types_data = {}
+
+                log.info(
+                    "Additional instance types data loaded for instance types '{0}': {1}".format(
+                        instance_types_data.keys(), instance_types_data
+                    )
+                )
+            except Exception as e:
+                raise CriticalError("Error loading instance types data from configuration: {0}".format(e))
+    return instance_types_data

src/jobwatcher/jobwatcher.py

@@ -25,6 +25,7 @@
     get_asg_settings,
     get_compute_instance_type,
     get_instance_properties,
+    load_additional_instance_types_data,
     load_module,
     sleep_remaining_loop_time,
 )
@@ -36,7 +37,7 @@
 
 
 JobwatcherConfig = collections.namedtuple(
-    "JobwatcherConfig", ["region", "scheduler", "stack_name", "pcluster_dir", "proxy_config"]
+    "JobwatcherConfig", ["region", "scheduler", "stack_name", "pcluster_dir", "proxy_config", "instance_types_data"]
 )
 
 
@@ -59,21 +60,23 @@ def _get_config():
     scheduler = config.get("jobwatcher", "scheduler")
     stack_name = config.get("jobwatcher", "stack_name")
     pcluster_dir = config.get("jobwatcher", "cfncluster_dir")
+    instance_types_data = load_additional_instance_types_data(config, "jobwatcher")
 
     _proxy = config.get("jobwatcher", "proxy")
     proxy_config = Config()
     if _proxy != "NONE":
         proxy_config = Config(proxies={"https": _proxy})
 
     log.info(
-        "Configured parameters: region=%s scheduler=%s stack_name=%s pcluster_dir=%s proxy=%s",
+        "Configured parameters: region=%s scheduler=%s stack_name=%s pcluster_dir=%s proxy=%s instance_types_data=%s",
         region,
         scheduler,
         stack_name,
         pcluster_dir,
         _proxy,
+        instance_types_data,
     )
-    return JobwatcherConfig(region, scheduler, stack_name, pcluster_dir, proxy_config)
+    return JobwatcherConfig(region, scheduler, stack_name, pcluster_dir, proxy_config, instance_types_data)
 
 
 def _poll_scheduler_status(config, asg_name, scheduler_module):
@@ -99,7 +102,9 @@ def _poll_scheduler_status(config, asg_name, scheduler_module):
             )
             if new_instance_type != instance_type:
                 instance_type = new_instance_type
-                instance_properties = get_instance_properties(config.region, config.proxy_config, instance_type)
+                instance_properties = get_instance_properties(
+                    config.region, config.proxy_config, instance_type, config.instance_types_data
+                )
         update_instance_properties_timer += LOOP_TIME
 
         # get current limits

src/nodewatcher/nodewatcher.py

@@ -31,6 +31,7 @@
     get_asg_name,
     get_asg_settings,
     get_instance_properties,
+    load_additional_instance_types_data,
     load_module,
     retrieve_max_cluster_size,
     sleep_remaining_loop_time,
@@ -48,7 +49,8 @@
 CLUSTER_PROPERTIES_REFRESH_INTERVAL = 300
 
 NodewatcherConfig = collections.namedtuple(
-    "NodewatcherConfig", ["region", "scheduler", "stack_name", "scaledown_idletime", "proxy_config"]
+    "NodewatcherConfig",
+    ["region", "scheduler", "stack_name", "scaledown_idletime", "proxy_config", "instance_types_data"],
 )
 
 
@@ -71,21 +73,24 @@ def _get_config():
     scheduler = config.get("nodewatcher", "scheduler")
     stack_name = config.get("nodewatcher", "stack_name")
     scaledown_idletime = int(config.get("nodewatcher", "scaledown_idletime"))
+    instance_types_data = load_additional_instance_types_data(config, "nodewatcher")
 
     _proxy = config.get("nodewatcher", "proxy")
     proxy_config = Config()
     if _proxy != "NONE":
         proxy_config = Config(proxies={"https": _proxy})
 
     log.info(
-        "Configured parameters: region=%s scheduler=%s stack_name=%s scaledown_idletime=%s proxy=%s",
+        "Configured parameters: region=%s scheduler=%s stack_name=%s scaledown_idletime=%s proxy=%s "
+        "instance_types_data=%s",
         region,
         scheduler,
         stack_name,
         scaledown_idletime,
         _proxy,
+        instance_types_data,
     )
-    return NodewatcherConfig(region, scheduler, stack_name, scaledown_idletime, proxy_config)
+    return NodewatcherConfig(region, scheduler, stack_name, scaledown_idletime, proxy_config, instance_types_data)
 
 
 def _get_metadata(metadata_path):
@@ -328,7 +333,9 @@ def _poll_instance_status(config, scheduler_module, asg_name, hostname, instance
     _terminate_if_down(scheduler_module, config, asg_name, instance_id, INITIAL_TERMINATE_TIMEOUT)
 
     idletime = _init_idletime()
-    instance_properties = get_instance_properties(config.region, config.proxy_config, instance_type)
+    instance_properties = get_instance_properties(
+        config.region, config.proxy_config, instance_type, config.instance_types_data
+    )
     start_time = None
     while True:
         sleep_remaining_loop_time(LOOP_TIME, start_time)

src/sqswatcher/sqswatcher.py

@@ -31,6 +31,7 @@
     get_asg_name,
     get_compute_instance_type,
     get_instance_properties,
+    load_additional_instance_types_data,
     load_module,
     retrieve_max_cluster_size,
     sleep_remaining_loop_time,
@@ -59,6 +60,7 @@ class QueryConfigError(Exception):
         "proxy_config",
         "stack_name",
         "max_processed_messages",
+        "instance_types_data",
     ],
 )
 
@@ -84,6 +86,7 @@ def _get_config():
     table_name = config.get("sqswatcher", "table_name")
     cluster_user = config.get("sqswatcher", "cluster_user")
     stack_name = config.get("sqswatcher", "stack_name")
+    instance_types_data = load_additional_instance_types_data(config, "sqswatcher")
     max_processed_messages = int(
         config.get("sqswatcher", "max_processed_messages", fallback=DEFAULT_MAX_PROCESSED_MESSAGES)
     )
@@ -95,7 +98,7 @@ def _get_config():
 
     log.info(
         "Configured parameters: region=%s scheduler=%s sqsqueue=%s table_name=%s cluster_user=%s "
-        "proxy=%s stack_name=%s max_processed_messages=%s",
+        "proxy=%s stack_name=%s max_processed_messages=%s instance_types_data=%s",
         region,
         scheduler,
         sqsqueue,
@@ -104,9 +107,18 @@ def _get_config():
         _proxy,
         stack_name,
         max_processed_messages,
+        instance_types_data,
     )
     return SQSWatcherConfig(
-        region, scheduler, sqsqueue, table_name, cluster_user, proxy_config, stack_name, max_processed_messages
+        region,
+        scheduler,
+        sqsqueue,
+        table_name,
+        cluster_user,
+        proxy_config,
+        stack_name,
+        max_processed_messages,
+        instance_types_data,
     )
 
 
@@ -226,7 +238,7 @@ def _retrieve_all_sqs_messages(queue, max_processed_messages):
     return messages
 
 
-def _parse_sqs_messages(sqs_config_region, sqs_config_proxy, messages, table, queue):
+def _parse_sqs_messages(sqs_config_region, sqs_config_proxy, messages, table, queue, additional_instance_types_data):
     add_events = []
     remove_events = []
     for message in messages:
@@ -240,7 +252,9 @@ def _parse_sqs_messages(sqs_config_region, sqs_config_proxy, messages, table, qu
             continue
 
         if event_type == "parallelcluster:COMPUTE_READY":
-            add_event = _process_compute_ready_event(sqs_config_region, sqs_config_proxy, message_attrs, message, table)
+            add_event = _process_compute_ready_event(
+                sqs_config_region, sqs_config_proxy, message_attrs, message, table, additional_instance_types_data
+            )
             if add_event:
                 add_events.append(add_event)
         elif event_type == "autoscaling:EC2_INSTANCE_TERMINATE":
@@ -266,12 +280,16 @@ def _parse_sqs_messages(sqs_config_region, sqs_config_proxy, messages, table, qu
     return update_events.values()
 
 
-def _process_compute_ready_event(sqs_config_region, sqs_config_proxy, message_attrs, message, table):
+def _process_compute_ready_event(
+    sqs_config_region, sqs_config_proxy, message_attrs, message, table, additional_instance_types_data
+):
     instance_id = message_attrs.get("EC2InstanceId")
     instance_type = message_attrs.get("EC2InstanceType")
     # Get instances properties for each event because instance types
     # from instance and CloudFormation could be out-of-sync
-    instance_properties = get_instance_properties(sqs_config_region, sqs_config_proxy, instance_type)
+    instance_properties = get_instance_properties(
+        sqs_config_region, sqs_config_proxy, instance_type, additional_instance_types_data
+    )
     gpus = instance_properties["gpus"]
     slots = message_attrs.get("Slots")
     hostname = message_attrs.get("LocalHostname").split(".")[0]
@@ -392,12 +410,16 @@ def _poll_queue(sqs_config, queue, table, asg_name):
             force_cluster_update = new_max_cluster_size != max_cluster_size or new_instance_type != instance_type
             if new_instance_type != instance_type:
                 instance_type = new_instance_type
-                instance_properties = get_instance_properties(sqs_config.region, sqs_config.proxy_config, instance_type)
+                instance_properties = get_instance_properties(
+                    sqs_config.region, sqs_config.proxy_config, instance_type, sqs_config.instance_types_data
+                )
             max_cluster_size = new_max_cluster_size
         cluster_properties_refresh_timer += LOOP_TIME
 
         messages = _retrieve_all_sqs_messages(queue, sqs_config.max_processed_messages)
-        update_events = _parse_sqs_messages(sqs_config.region, sqs_config.proxy_config, messages, table, queue)
+        update_events = _parse_sqs_messages(
+            sqs_config.region, sqs_config.proxy_config, messages, table, queue, sqs_config.instance_types_data
+        )
         _process_sqs_messages(
             update_events,
             scheduler_module,