Modifying pcluster_fleet_config_generator.py to make --output-file and --total-min-count as mutually exclusive arguments

Author: Himani Anil Deshpande

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

@@ -196,6 +196,25 @@ def _load_cluster_config(input_file_path):
     with open(input_file_path, encoding="utf-8") as input_file:
         return yaml.load(input_file, Loader=yaml.SafeLoader)
 
+def get_total_min_count(input_file: str):
+    "Gives total count of Static Nodes in a cluster."
+    cluster_config = _load_cluster_config(input_file)
+    queue_name, compute_resource_name = None, None
+    try:
+        total_min_count = 0
+        for queue_config in cluster_config["Scheduling"]["SlurmQueues"]:
+            for compute_resource_config in queue_config["ComputeResources"]:
+                total_min_count = total_min_count + compute_resource_config["MinCount"]
+    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}" if queue_name else ""
+        log.error(message)
+        raise CriticalError(message)
+    return total_min_count
+
 
 def main():
     try:
@@ -204,14 +223,19 @@ def main():
         )
         log.info("Running ParallelCluster Fleet Config Generator")
         parser = argparse.ArgumentParser(description="Take in fleet configuration generator related parameters")
-        parser.add_argument("--output-file", help="The output file for generated json fleet config", required=True)
         parser.add_argument(
             "--input-file",
             help="Yaml file containing pcluster CLI configuration file with default values",
             required=True,
         )
+        exclusive_group = parser.add_mutually_exclusive_group(required=True)
+        exclusive_group.add_argument("--output-file", help="The output file for generated json fleet config")
+        exclusive_group.add_argument("--total-min-count", help="Gets the total min count of Scheduler", action='store_true')
         args = parser.parse_args()
-        generate_fleet_config_file(args.output_file, args.input_file)
+        if args.total_min_count:
+            log.info(f"The total MinCount of cluster is {get_total_min_count(args.input_file)}")
+        else:
+            generate_fleet_config_file(args.output_file, args.input_file)
     except Exception as e:
         log.exception("Failed to generate Fleet configuration, exception: %s", e)
         raise

cookbooks/aws-parallelcluster-slurm/libraries/helpers.rb

@@ -179,7 +179,7 @@ def wait_cluster_ready
 end
 
 def get_static_node_count
-  cmd = Mixlib::ShellOut.new("cat #{node['cluster']['slurm']['install_dir']}/etc/slurm_parallelcluster.conf | grep -o '#TOTAL_MIN_COUNT=\([0-9]*\)' | cut -d'=' -f2")
+  cmd = Mixlib::ShellOut.new("#{cookbook_virtualenv_path}/bin/python #{node['cluster']['scripts_dir']}/slurm/pcluster_fleet_config_generator.py  --input-file #{node['cluster']['cluster_config_path']} --total-min-count | grep -o 'The total MinCount of cluster is =\([0-9]*\)' | cut -d'=' -f2")
   cmd.run_command.stdout.strip
 end