Configure task communication using an MPI library

This document explains how to configure a Batch job with tightly coupled tasks that communicate with each other across different VMs by using a Message Passing Interface (MPI) library.

For Batch jobs, coupling describes the interdependency tasks; this influences how you configure the number of tasks that can run in parallel (instead of sequentially) using a job's parallelism field. Tasks can be described using the following types of coupling:

Loosely coupled tasks: tasks that can run independently.
Tightly coupled tasks: tasks that are dependent on each other to run.

Optionally, you can create a job that uses an MPI library to let tightly coupled tasks communicate with each other across different VM instances. A common use case for MPI is tightly coupled high-performance computing (HPC) workloads.

Before you begin

If you haven't used Batch before, review Get started with Batch and enable Batch by completing the prerequisites for projects and users.
To get the permissions that you need to create a job, ask your administrator to grant you the following IAM roles:
- Batch Job Editor (roles/batch.jobsEditor) on the project
- Service Account User (roles/iam.serviceAccountUser) on the job's service account, which by default is the default Compute Engine service account
For more information about granting roles, see Manage access to projects, folders, and organizations.

You might also be able to get the required permissions through custom roles or other predefined roles.
If you specify the network for this job, make sure the network has a firewall rule that allows connections between the job's VMs. Learn how to Configure VPC firewall rules for common use cases.

Create and run a job that uses MPI for tightly coupled tasks

This section provides examples for how to create a job that can use MPI. Notably, the example job has 3 runnables:

The first runnable is a script that prepares the job for MPI by disabling simultaneous multithreading and installing Intel MPI.
The second runnable is an empty barrier runnable (formatted as { "barrier": {} }), which ensures that all tasks finish setting up MPI before continuing to future runnables.
The third runnable (and any subsequent runnables) is available for the job's workload.

You can create a job that uses MPI for tightly coupled tasks using the gcloud CLI or Batch API.

gcloud

To create a script job that uses MPI for tightly coupled tasks by using the gcloud CLI, do the following:

Create a JSON configuration file with the following contents:

{
    "taskGroups": [
        {
            "taskSpec": {
                "runnables": [
                    {
                        "script": {
                            "text": "google_mpi_tuning --nosmt; google_install_mpi --intel_mpi;"
                        }
                    },
                    { "barrier": {} },
                    {
                        "script": {
                            SCRIPT
                        }
                    }
                ]
            },
            "taskCount": TASK_COUNT,
            "taskCountPerNode": TASK_COUNT_PER_NODE,
            "requireHostsFile": REQUIRE_HOSTS_FILE,
            "permissiveSsh": PERMISSIVE_SSH
        }
    ]
}

Replace the following:

SCRIPT: a script runnable for a workload that uses MPI.
TASK_COUNT: the number of tasks for the job. The value must be a whole number between 1 and the tasks per task group limit. To use the MPI libraries provided by Batch, this field is required and must be set to 2 or higher.
TASK_COUNT_PER_NODE: the number of tasks that a job can run concurrently on a VM instance. To use the MPI libraries provided by Batch, this field is required and must be set to 1, which equals to running one VM instance per task.
REQUIRE_HOSTS_FILE: when set to true, the job creates a file listing the VM instances running in a task group. The file path is stored in the BATCH_HOSTS_FILE environment variable. To use the MPI libraries provided by Batch, this field must be set to true.
PERMISSIVE_SSH: when set to true, Batch configures SSH to allow passwordless communication among the VM instances running in a task group. To use the MPI libraries provided by Batch, this field must be set to true.

To create the job, use the gcloud batch jobs submit command.
```
gcloud batch jobs submit JOB_NAME \
  --location LOCATION \
  --config JSON_CONFIGURATION_FILE
```
Replace the following:
- JOB_NAME: the name of the job.
- LOCATION: the location of the job.
- JSON_CONFIGURATION_FILE: the path for a JSON file with the job's configuration details.

Optionally, you can increase the performance of the MPI libraries provided by Batch by doing the following:

Create and run a job using a VM instance template and specify an HPC-ready VM instance template in your job's resources.
Specify Cloud Logging as the destination for your job's logging information in your job's resources.

For example, to create a script job from an instance template that uses MPI and makes 1 task output the hostname of the 3 tasks in the task group:

Create a JSON file in the current directory named example-job-uses-mpi.json with the following contents:

{
    "taskGroups": [
        {
            "taskSpec": {
                "runnables": [
                    {
                        "script": {
                            "text": "google_mpi_tuning --nosmt; google_install_mpi --intel_mpi;"
                        }
                    },
                    { "barrier": {} },
                    {
                        "script": {
                            "text":
                                "if [ $BATCH_TASK_INDEX = 0 ]; then
                                mpirun -hostfile $BATCH_HOSTS_FILE -np 3 hostname;
                                fi"
                        }
                    },
                    { "barrier": {} }
                ]
            },
            "taskCount": 3,
            "taskCountPerNode": 1,
            "requireHostsFile": true,
            "permissiveSsh": true
        }
    ],
    "allocationPolicy": {
        "instances": [
            {
                "instanceTemplate": "example-template-job-uses-mpi"
            }
        ]
    },
    "logsPolicy": {
        "destination": "CLOUD_LOGGING"
    }
}

Run the following command:

gcloud batch jobs submit example-template-job-uses-mpi \
  --location us-central1 \
  --config example-job-uses-mpi.json

API

To create a script job that uses a MPI for tightly coupled tasks by using the Batch API, use the jobs.create method and specify the permissiveSsh, requireHostsFile, taskCount, and taskCountPerNode fields.

POST https://batch.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/jobs?job_id=JOB_ID

{
    "taskGroups": [
        {
            "taskSpec": {
                "runnables": [
                    {
                        "script": {
                            "text": "google_mpi_tuning --nosmt; google_install_mpi --intel_mpi;"
                        }
                    },
                    { "barrier": {} },
                    {
                        "script": {
                            SCRIPT
                        }
                    }
                ]
            },
            "taskCount": TASK_COUNT,
            "taskCountPerNode": TASK_COUNT_PER_NODE,
            "requireHostsFile": REQUIRE_HOSTS_FILE,
            "permissiveSsh": PERMISSIVE_SSH
        }
    ]
}

Replace the following:

PROJECT_ID: the project ID of your project.
LOCATION: the location of the job.
JOB_NAME: the name of the job.
SCRIPT: the script runnable for a workload that uses MPI.
TASK_COUNT: the number of tasks for the job. The value must be a whole number between 1 and the tasks per task group limit. To use the MPI libraries provided by Batch, this field is required and must be set to 2 or higher.
TASK_COUNT_PER_NODE: the number of tasks that a job can run concurrently on a VM instance. To use the MPI libraries provided by Batch, this field is required and must be set to 1, which equals to running one VM instance per task.
REQUIRE_HOSTS_FILE: when set to true, the job creates a file listing the VM instances running in a task group. The file path is stored in the BATCH_HOSTS_FILE environment variable. To use the MPI libraries provided by Batch, this field must be set to true.
PERMISSIVE_SSH: when set to true, Batch configures SSH to allow passwordless communication among the VM instances running in a task group. To use the MPI libraries provided by Batch, this field must be set to true.

Optionally, you can increase the performance of the MPI libraries provided by Batch by doing the following:

Create and run a job using a VM instance template and specify an HPC-ready VM instance template in your job's resources.
Specify Cloud Logging as the destination for your job's logging information in your job's resources.

For example, to create a script job from an instance template that uses MPI and makes 1 task output the hostname of the 3 tasks in the task group, use the following request:

POST https://batch.googleapis.com/v1/projects/PROJECT_ID/locations/us-central1/jobs?job_id=example-template-job-uses-mpi

{
    "taskGroups": [
        {
            "taskSpec": {
                "runnables": [
                    {
                        "script": {
                            "text": "google_mpi_tuning --nosmt; google_install_mpi --intel_mpi;"
                        }
                    },
                    { "barrier": {} },
                    {
                        "script": {
                            "text":
                                "if [ $BATCH_TASK_INDEX = 0 ]; then
                                mpirun -hostfile $BATCH_HOSTS_FILE -np 3 hostname;
                                fi"
                        }
                    },
                    { "barrier": {} }
                ]
            },
            "taskCount": 3,
            "taskCountPerNode": 1,
            "requireHostsFile": true,
            "permissiveSsh": true
        }
    ],
    "allocationPolicy": {
        "instances": [
            {
                "instanceTemplate": "example-template-job-uses-mpi"
            }
        ]
    },
    "logsPolicy": {
        "destination": "CLOUD_LOGGING"
    }
}

where PROJECT_ID is the project ID of your project.

What's next

For a more detailed example of a job that uses MPI for tightly coupled tasks, see Running the Weather Research and Forecasting model with Batch.
Learn how to reduce latency for jobs with tightly coupled tasks.
If you have issues creating or running a job, see Troubleshooting.
View jobs and tasks.
Learn about more job creation options.