Simplify generate_fleet_config_file my moving compute resource logic into a separate function

Signed-off-by: Enrico Usai <[email protected]>

Author: enrico-usai

cookbooks/aws-parallelcluster-slurm/files/default/head_node_slurm/slurm/pcluster_fleet_config_generator.py

@@ -14,12 +14,20 @@
 import json
 import logging
 import traceback
+from typing import List
 
 import yaml
 
 log = logging.getLogger()
 
 
+CAPACITY_TYPE_MAP = {
+    "ONDEMAND": "on-demand",
+    "SPOT": "spot",
+    "CAPACITY_BLOCK": "capacity-block",
+}
+
+
 class CriticalError(Exception):
     """Critical error for the script."""
 
@@ -32,7 +40,7 @@ class ConfigurationFieldNotFoundError(Exception):
     pass
 
 
-def generate_fleet_config_file(output_file, input_file):
+def generate_fleet_config_file(output_file: str, input_file: str):
     """
     Generate configuration file used by Fleet Manager in node daemon package.
 
@@ -64,21 +72,15 @@ def generate_fleet_config_file(output_file, input_file):
         }
     }
     """
-    capacity_type_map = {
-        "ONDEMAND": "on-demand",
-        "SPOT": "spot",
-        "CAPACITY_BLOCK": "capacity-block",
-    }
-
     cluster_config = _load_cluster_config(input_file)
-    queue, compute_resource = None, None
+    queue_name, compute_resource_name = None, None
     try:
         fleet_config = {}
         for queue_config in cluster_config["Scheduling"]["SlurmQueues"]:
-            queue = queue_config["Name"]
+            queue_name = queue_config["Name"]
 
-            # Retrieve capacity info from the queue, if there
-            queue_capacity_type = capacity_type_map.get(queue_config.get("CapacityType", "ONDEMAND"))
+            # Retrieve capacity info from the queue_name, if there
+            queue_capacity_type = CAPACITY_TYPE_MAP.get(queue_config.get("CapacityType", "ONDEMAND"))
             queue_allocation_strategy = queue_config.get("AllocationStrategy")
             queue_capacity_reservation_target = queue_config.get("CapacityReservationTarget", {})
             queue_capacity_reservation = (
@@ -87,66 +89,25 @@ def generate_fleet_config_file(output_file, input_file):
                 else None
             )
 
-            fleet_config[queue] = {}
+            fleet_config[queue_name] = {}
 
             for compute_resource_config in queue_config["ComputeResources"]:
-                compute_resource = compute_resource_config["Name"]
-
-                # Override capacity info from the compute resource.
-                # CapacityReservationTarget can be specified on both queue and compute resource level.
-                # CapacityType and AllocationStrategy are not yet supported at compute resource level from the CLI,
-                # but this code is ready to use them.
-                capacity_type = capacity_type_map.get(compute_resource_config.get("CapacityType"), queue_capacity_type)
-                allocation_strategy = compute_resource_config.get("AllocationStrategy", queue_allocation_strategy)
-                capacity_reservation_target = compute_resource_config.get("CapacityReservationTarget", {})
-                capacity_reservation = (
-                    capacity_reservation_target.get("CapacityReservationId", queue_capacity_reservation)
-                    if capacity_reservation_target
-                    else queue_capacity_reservation
+                compute_resource_name, config_for_fleet = _generate_compute_resource_fleet_config(
+                    compute_resource_config=compute_resource_config,
+                    queue_name=queue_name,
+                    queue_allocation_strategy=queue_allocation_strategy,
+                    queue_capacity_reservation=queue_capacity_reservation,
+                    queue_capacity_type=queue_capacity_type,
+                    queue_subnets=queue_config["Networking"]["SubnetIds"],
                 )
-
-                config_for_fleet = {"CapacityType": capacity_type}
-                if capacity_reservation:
-                    config_for_fleet.update({"CapacityReservationId": capacity_reservation})
-
-                if compute_resource_config.get("Instances"):
-                    # multiple instance types, create-fleet api
-                    config_for_fleet.update(
-                        {
-                            "Api": "create-fleet",
-                            "Instances": copy.deepcopy(compute_resource_config["Instances"]),
-                            "Networking": {"SubnetIds": queue_config["Networking"]["SubnetIds"]},
-                        }
-                    )
-                    if allocation_strategy:
-                        config_for_fleet.update({"AllocationStrategy": allocation_strategy})
-                    if capacity_type == "spot" and compute_resource_config["SpotPrice"]:
-                        config_for_fleet.update({"MaxPrice": compute_resource_config["SpotPrice"]})
-
-                elif compute_resource_config.get("InstanceType"):
-                    # single instance type, run-instances api
-                    config_for_fleet.update(
-                        {
-                            "Api": "run-instances",
-                            "Instances": [{"InstanceType": compute_resource_config["InstanceType"]}],
-                        }
-                    )
-
-                else:
-                    raise ConfigurationFieldNotFoundError(
-                        "Instances or InstanceType field not found "
-                        f"in queue: {queue}, compute resource: {compute_resource} configuration"
-                    )
-
-                fleet_config[queue][compute_resource] = config_for_fleet
+                fleet_config[queue_name][compute_resource_name] = config_for_fleet
 
     except (KeyError, AttributeError) as e:
         if isinstance(e, KeyError):
             message = f"Unable to find key {e} in the configuration file."
         else:
             message = f"Error parsing configuration file. {e}. {traceback.format_exc()}."
-        message += f" Queue: {queue}" if queue else ""
-        message += f" Compute resource: {compute_resource}" if compute_resource else ""
+        message += f" Queue: {queue_name}" if queue_name else ""
         log.error(message)
         raise CriticalError(message)
 
@@ -157,6 +118,79 @@ def generate_fleet_config_file(output_file, input_file):
     log.info("Finished.")
 
 
+def _generate_compute_resource_fleet_config(
+    compute_resource_config: dict,
+    queue_name: str,
+    queue_allocation_strategy: str,
+    queue_capacity_reservation: str,
+    queue_capacity_type: str,
+    queue_subnets: List,
+):
+    """
+    Generate compute resource config to add in the fleet-config.json, overriding values from the queue.
+
+    CapacityReservationTarget can be specified on both queue and compute resource level.
+    CapacityType and AllocationStrategy are not yet supported at compute resource level from the CLI,
+    but this code is ready to use them.
+
+    Returns compute_resource name and fleet-config section for the given compute resource.
+    """
+    compute_resource_name = compute_resource_config["Name"]
+
+    try:
+        capacity_type = CAPACITY_TYPE_MAP.get(compute_resource_config.get("CapacityType"), queue_capacity_type)
+        config_for_fleet = {"CapacityType": capacity_type}
+
+        capacity_reservation_target = compute_resource_config.get("CapacityReservationTarget", {})
+        capacity_reservation = (
+            capacity_reservation_target.get("CapacityReservationId", queue_capacity_reservation)
+            if capacity_reservation_target
+            else queue_capacity_reservation
+        )
+        if capacity_reservation:
+            config_for_fleet.update({"CapacityReservationId": capacity_reservation})
+
+        if compute_resource_config.get("Instances"):
+            # multiple instance types, create-fleet api
+            config_for_fleet.update(
+                {
+                    "Api": "create-fleet",
+                    "Instances": copy.deepcopy(compute_resource_config["Instances"]),
+                    "Networking": {"SubnetIds": queue_subnets},
+                }
+            )
+            allocation_strategy = compute_resource_config.get("AllocationStrategy", queue_allocation_strategy)
+            if allocation_strategy:
+                config_for_fleet.update({"AllocationStrategy": allocation_strategy})
+            if capacity_type == "spot" and compute_resource_config["SpotPrice"]:
+                config_for_fleet.update({"MaxPrice": compute_resource_config["SpotPrice"]})
+
+        elif compute_resource_config.get("InstanceType"):
+            # single instance type, run-instances api
+            config_for_fleet.update(
+                {
+                    "Api": "run-instances",
+                    "Instances": [{"InstanceType": compute_resource_config["InstanceType"]}],
+                }
+            )
+
+        else:
+            raise ConfigurationFieldNotFoundError(
+                "Instances or InstanceType field not found "
+                f"in queue: {queue_name}, compute resource: {compute_resource_name} configuration"
+            )
+    except (KeyError, AttributeError) as e:
+        if isinstance(e, KeyError):
+            message = f"Unable to find key {e} in the configuration file."
+        else:
+            message = f"Error parsing configuration file. {e}. {traceback.format_exc()}."
+        message += f" Queue: {queue_name}, Compute resource: {compute_resource_name}"
+        log.error(message)
+        raise CriticalError(message)
+
+    return compute_resource_name, config_for_fleet
+
+
 def _load_cluster_config(input_file_path):
     """Load cluster config file."""
     with open(input_file_path, encoding="utf-8") as input_file:

test/unit/slurm/test_fleet_config_generator.py

@@ -44,6 +44,22 @@
                 }
             },
             CriticalError,
+            "Unable to find key 'Networking' in the configuration file. Queue: q1",
+        ),
+        (
+            {
+                "Scheduling": {
+                    "SlurmQueues": [
+                        {
+                            "Name": "q1",
+                            "CapacityType": "SPOT",
+                            "ComputeResources": [{"Instances": []}],
+                            "Networking": {"SubnetIds": ["123"]},
+                        }
+                    ]
+                }
+            },
+            CriticalError,
             "Unable to find key 'Name' in the configuration file. Queue: q1",
         ),
         (
@@ -137,7 +153,7 @@
                 }
             },
             CriticalError,
-            "Unable to find key 'SpotPrice' in the configuration file. Queue: q1 Compute resource: cr1",
+            "Unable to find key 'SpotPrice' in the configuration file. Queue: q1, Compute resource: cr1",
         ),
         (
             {
@@ -172,7 +188,7 @@
                 }
             },
             CriticalError,
-            "Unable to find key 'Networking' in the configuration file. Queue: q1 Compute resource: cr1",
+            "Unable to find key 'Networking' in the configuration file. Queue: q1",
         ),
         (
             {
@@ -190,7 +206,7 @@
                 }
             },
             CriticalError,
-            "Unable to find key 'SubnetIds' in the configuration file. Queue: q1 Compute resource: cr1",
+            "Unable to find key 'SubnetIds' in the configuration file. Queue: q1",
         ),
         (
             {