google_cloud_batch_script_job_creation_with_task_groups.py

1from google.cloud import batch_v1
2
3def create_script_job(project_id: str, region: str, job_name: str) -> batch_v1.Job:
4    """
5    This method shows how to create a sample Batch Job that will run a simple HTTP GET request.
6
7    Args:
8        project_id: project ID or project number of the Cloud project you want to use.
9        region: name of the region you want to use to run the job. Regions that are
10            available for Batch are listed at: https://cloud.google.com/batch/docs/get-started#locations
11        job_name: the name of the job that will be created.
12
13    Returns:
14        A job object representing the job created.
15    """
16    client = batch_v1.BatchServiceClient()
17
18    # Define what will be done as part of the job.
19    task = batch_v1.TaskSpec()
20    runnable = batch_v1.Runnable()
21    runnable.script = batch_v1.Runnable.Script()
22    runnable.script.text = "echo Hello world! This is task ${BATCH_TASK_INDEX}. This job has ${BATCH_TASK_COUNT} tasks."
23    # You can also run a container:
24    # runnable.container = batch_v1.Runnable.Container()
25    # runnable.container.image_uri = "gcr.io/google-containers/busybox"
26    # runnable.container.entrypoint = "/bin/sh"
27    # runnable.container.commands = ["-c", "echo Hello world! This is task ${BATCH_TASK_INDEX}. This job has ${BATCH_TASK_COUNT} tasks."]
28    task.runnables = [runnable]
29
30    # We can specify what resources are requested by each task.
31    resources = batch_v1.ComputeResource()
32    resources.cpu_milli = 2000  # in milliseconds per cpu-second. This means the task requires 2 whole CPUs.
33    resources.memory_mib = 16  # in MiB
34    task.compute_resource = resources
35
36    task.max_retry_count = 2
37    task.max_run_duration = "3600s"
38
39    # Tasks are grouped into a job using TaskGroups.
40    # Currently, a Batch job can have only one TaskGroup.
41    group = batch_v1.TaskGroup()
42    group.task_count = 4
43    group.task_spec = task
44
45    # Policies are used to define on what kind of virtual machines the tasks will run on.
46    # In this case, we tell the system to use "e2-standard-4" machine type.
47    # Read more about machine types here: https://cloud.google.com/compute/docs/machine-types
48    policy = batch_v1.AllocationPolicy()
49    instance_policy = batch_v1.AllocationPolicy.InstancePolicy()
50    instance_policy.machine_type = "e2-standard-4"
51    instances = batch_v1.AllocationPolicy.InstancePolicyOrTemplate()
52    instances.policy = instance_policy
53    policy.instances = [instances]
54
55    job = batch_v1.Job()
56    job.task_groups = [group]
57    job.allocation_policy = policy
58    job.labels = {"env": "testing", "type": "script"}
59
60    request = batch_v1.CreateJobRequest()
61    request.job = job
62    request.job_id = job_name
63    # The job's parent is the region in which the job will run
64    request.parent = f"projects/{project_id}/locations/{region}"
65
66    return client.create_job(request=request)