in slurm-to-batch/convert_to_batch_job.py [0:0]
def createJobJSON(slurm_conf: SlurmJobConfig) -> dict:
slurm_setup = generate_gres_conf_script(slurm_conf) + generate_slurm_conf_script(slurm_conf) + start_slurmctld()
job_definition = {
"taskGroups": [
{
"task_spec": {
"runnables": [
{
"script": {
"text": slurm_setup,
},
},
{
"barrier": {
"name": "slurmctld-started"
}
},
{
"script": {
"text": start_slurmd(slurm_conf)
}
},
{
"barrier": {
"name": "slurmd-started-all-vms"
}
},
{
"script": {
"text": work_load()
}
},
{
"barrier": {
"name": "slurm-job-finished-all-vms"
}
},
],
},
"task_count_per_node": 1,
"task_count": slurm_conf.node_count,
"require_hosts_file": True
}
],
"allocation_policy": {
"location": {
"allowed_locations": ["zones/<SELECTED_ZONE>"]
},
"instances": {
"policy": {
"accelerators": {
"type": "<CUSTOM_GPU_TYPE>",
"count": slurm_conf.gpu_per_node
},
"boot_disk": {
"image": "<CUSTOM_BOOT_IMAGE>",
"size_gb": "<CUSTOM_BOOT_DISK_SIZE>",
}
},
"install_gpu_drivers": True
}
},
"labels": {
"goog-batch-dynamic-workload-scheduler": "true"
},
"logs_policy": {
"destination": "CLOUD_LOGGING"
}
}
return job_definition