in community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/conf.py [0:0]
def conflines(lkp: util.Lookup) -> str:
params = lkp.cfg.cloud_parameters
def get(key, default):
"""
Returns the value of the key in params if it exists and is not None,
otherwise returns supplied default.
We can't rely on the `dict.get` method because the value could be `None` as
well as empty NSDict, depending on type of the `cfg.cloud_parameters`.
TODO: Simplify once NSDict is removed from the codebase.
"""
if key not in params or params[key] is None:
return default
return params[key]
no_comma_params = get("no_comma_params", False)
any_gpus = any(
lkp.template_info(nodeset.instance_template).gpu
for nodeset in lkp.cfg.nodeset.values()
)
any_tpu = any(
tpu_nodeset is not None
for part in lkp.cfg.partitions.values()
for tpu_nodeset in part.partition_nodeset_tpu
)
any_dynamic = any(bool(p.partition_feature) for p in lkp.cfg.partitions.values())
comma_params = {
"LaunchParameters": [
"enable_nss_slurm",
"use_interactive_step",
],
"SlurmctldParameters": [
"cloud_reg_addrs" if any_dynamic or any_tpu else "cloud_dns",
"enable_configless",
"idle_on_node_suspend",
],
"GresTypes": [
"gpu" if any_gpus else None,
],
}
scripts_dir = lkp.cfg.install_dir or dirs.scripts
prolog_path = Path(dirs.custom_scripts / "prolog.d")
epilog_path = Path(dirs.custom_scripts / "epilog.d")
default_tree_width = 65533 if any_dynamic else 128
conf_options = {
**(comma_params if not no_comma_params else {}),
"Prolog": f"{prolog_path}/*" if lkp.cfg.prolog_scripts else None,
"Epilog": f"{epilog_path}/*" if lkp.cfg.epilog_scripts else None,
"PrivateData": get("private_data", []),
"SchedulerParameters": get("scheduler_parameters", [
"bf_continue",
"salloc_wait_nodes",
"ignore_prefer_validation",
]),
"ResumeProgram": f"{scripts_dir}/resume.py",
"ResumeFailProgram": f"{scripts_dir}/suspend.py",
"ResumeRate": get("resume_rate", 0),
"ResumeTimeout": get("resume_timeout", 300),
"SuspendProgram": f"{scripts_dir}/suspend.py",
"SuspendRate": get("suspend_rate", 0),
"SuspendTimeout": get("suspend_timeout", 300),
"TreeWidth": get("tree_width", default_tree_width),
"JobSubmitPlugins": "lua" if any_tpu else None,
"TopologyPlugin": topology_plugin(lkp),
"TopologyParam": get("topology_param", "SwitchAsNodeRank"),
}
return dict_to_conf(conf_options, delim="\n")