def main()

in cookbooks/aws-parallelcluster-slurm/files/default/head_node_slurm/slurm/pcluster_slurm_config_generator.py [0:0]
69 lines of code
5 McCabe index (conditional complexity)

def main():
    def memory_ratio_float(arg):
        """Type function for realmemory_to_ec2memory_ratio with custom lower and upper bounds."""
        # We cannot allow 0 as minimum value because `RealMemory=0` is not valid in Slurm.
        # We put a minimum value= 0.1. It doesn't make sense to put such a low value anyway.
        min_value = 0.1
        max_value = 1.0
        try:
            f = float(arg)
        except ValueError:
            raise argparse.ArgumentTypeError("The argument must be a floating point number")
        if f < min_value or f > max_value:
            raise argparse.ArgumentTypeError(
                f"The argument must be greater or equal than {str(min_value)} and less "
                f"or equal than {str(max_value)}"
            )
        return f

    try:
        _setup_logger()
        log.info("Running ParallelCluster Slurm Config Generator")
        parser = argparse.ArgumentParser(description="Take in slurm configuration generator related parameters")
        parser.add_argument(
            "--output-directory", help="The output directory for generated slurm configs", required=True
        )
        parser.add_argument("--template-directory", help="The directory storing slurm config templates", required=True)
        parser.add_argument("--input-file", help="Yaml file containing pcluster configuration file", required=True)
        parser.add_argument("--instance-types-data", help="JSON file containing info about instance types")
        parser.add_argument(
            "--dryrun",
            action="store_true",
            help="dryrun",
            required=False,
            default=False,
        )
        parser.add_argument(
            "--no-gpu",
            action="store_true",
            help="no gpu configuration",
            required=False,
            default=False,
        )
        parser.add_argument(
            "--compute-node-bootstrap-timeout",
            type=int,
            help="Configure ResumeTimeout",
            required=False,
            default=1800,
        )
        parser.add_argument(
            "--realmemory-to-ec2memory-ratio",
            type=memory_ratio_float,
            help="Configure ratio between RealMemory and memory advertised by EC2",
            required=True,
        )
        parser.add_argument("--slurmdbd-user", help="User for the slurmdbd service.", required=True)
        parser.add_argument("--cluster-name", help="Name of the cluster.", required=True)
        args = parser.parse_args()
        generate_slurm_config_files(
            args.output_directory,
            args.template_directory,
            args.input_file,
            args.instance_types_data,
            args.dryrun,
            args.no_gpu,
            args.compute_node_bootstrap_timeout,
            args.realmemory_to_ec2memory_ratio,
            args.slurmdbd_user,
            args.cluster_name,
        )
    except Exception as e:
        log.exception("Failed to generate slurm configurations, exception: %s", e)
        raise