def createJobJSON()

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