feat: Batch create job with custom status events sample (#12060)

* Add create_with_custom_status_events.py and update test_basics.py

- Added new script to create jobs with custom events to describe job's runnables.
- Modified test_basics.py to include tests for creating job with custom status events

* Minor fixes

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* Updated requirements.txt for the new version of cloud-batch.

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

batch/create/create_with_custom_status_events.py

@@ -0,0 +1,120 @@
+#  Copyright 2022 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+
+import google.auth
+
+# [START batch_custom_events]
+from google.cloud import batch_v1
+
+
+def create_job_with_status_events(
+    project_id: str, region: str, job_name: str
+) -> batch_v1.Job:
+    """
+    This method shows the creation of a Batch job with custom status events which describe runnables
+    Within the method, the state of a runnable is described by defining its display name.
+    The script text is modified to change the commands that are executed, and barriers are adjusted
+    to synchronize tasks at specific points.
+
+    Args:
+        project_id (str): project ID or project number of the Cloud project you want to use.
+        region (str): name of the region you want to use to run the job. Regions that are
+            available for Batch are listed on: https://cloud.google.com/batch/docs/locations
+        job_name (str): the name of the job that will be created.
+            It needs to be unique for each project and region pair.
+
+    Returns:
+        A job object representing the job created with additional runnables and custom events.
+    """
+    client = batch_v1.BatchServiceClient()
+
+    # Executes a simple script that prints a message.
+    runn1 = batch_v1.Runnable()
+    runn1.display_name = "Script 1"
+    runn1.script.text = "echo Hello world from Script 1 for task ${BATCH_TASK_INDEX}"
+
+    # Acts as a barrier to synchronize the execution of subsequent runnables.
+    runn2 = batch_v1.Runnable()
+    runn2.display_name = "Barrier 1"
+    runn2.barrier = batch_v1.Runnable.Barrier({"name": "hello-barrier"})
+
+    # Executes another script that prints a message, intended to run after the barrier.
+    runn3 = batch_v1.Runnable()
+    runn3.display_name = "Script 2"
+    runn3.script.text = "echo Hello world from Script 2 for task ${BATCH_TASK_INDEX}"
+
+    # Executes a script that imitates a delay and creates a custom event for monitoring purposes.
+    runn4 = batch_v1.Runnable()
+    runn4.script.text = (
+        'sleep 30; echo \'{"batch/custom/event": "EVENT_DESCRIPTION"}\'; sleep 30'
+    )
+
+    # Jobs can be divided into tasks. In this case, we have only one task.
+    task = batch_v1.TaskSpec()
+    # Assigning a list of runnables to the task.
+    task.runnables = [runn1, runn2, runn3, runn4]
+
+    # We can specify what resources are requested by each task.
+    resources = batch_v1.ComputeResource()
+    resources.cpu_milli = 2000  # in milliseconds per cpu-second. This means the task requires 2 whole CPUs.
+    resources.memory_mib = 16  # in MiB
+    task.compute_resource = resources
+
+    task.max_retry_count = 2
+    task.max_run_duration = "3600s"
+
+    # Tasks are grouped inside a job using TaskGroups.
+    # Currently, it's possible to have only one task group.
+    group = batch_v1.TaskGroup()
+
+    group.task_count = 4
+    group.task_spec = task
+
+    # Policies are used to define on what kind of virtual machines the tasks will run on.
+    # In this case, we tell the system to use "e2-standard-4" machine type.
+    # Read more about machine types here: https://cloud.google.com/compute/docs/machine-types
+    policy = batch_v1.AllocationPolicy.InstancePolicy()
+    policy.machine_type = "e2-standard-4"
+    instances = batch_v1.AllocationPolicy.InstancePolicyOrTemplate()
+    instances.policy = policy
+    allocation_policy = batch_v1.AllocationPolicy()
+    allocation_policy.instances = [instances]
+
+    job = batch_v1.Job()
+    job.task_groups = [group]
+    job.allocation_policy = allocation_policy
+    job.labels = {"env": "testing", "type": "container"}
+    # We use Cloud Logging as it's an out of the box available option
+    job.logs_policy = batch_v1.LogsPolicy()
+    job.logs_policy.destination = batch_v1.LogsPolicy.Destination.CLOUD_LOGGING
+
+    create_request = batch_v1.CreateJobRequest()
+    create_request.job = job
+    create_request.job_id = job_name
+    # The job's parent is the region in which the job will run
+    create_request.parent = f"projects/{project_id}/locations/{region}"
+
+    return client.create_job(create_request)
+
+
+# [END batch_custom_events]
+
+
+if __name__ == "__main__":
+    PROJECT_ID = google.auth.default()[1]
+    REGION = "europe-west4"
+    job_name = "test-job-name"
+    job = create_job_with_status_events(PROJECT_ID, REGION, job_name)
+    print(job)

batch/requirements.txt

@@ -1,2 +1,2 @@
-google-cloud-batch==0.11.0
+google-cloud-batch==0.17.22
 google-cloud-logging==3.5.0

batch/tests/test_basics.py

@@ -25,6 +25,7 @@
 import pytest
 
 from ..create.create_with_container_no_mounting import create_container_job
+from ..create.create_with_custom_status_events import create_job_with_status_events
 from ..create.create_with_gpu_no_mounting import create_gpu_job
 from ..create.create_with_persistent_disk import create_with_pd_job
 from ..create.create_with_script_no_mounting import create_script_job
@@ -137,6 +138,24 @@ def _check_secret_set(job: batch_v1.Job, secret_name: str):
     assert secret_name in job.task_groups[0].task_spec.environment.secret_variables
 
 
+def _check_custom_events(job: batch_v1.Job):
+    display_names = ["Script 1", "Barrier 1", "Script 2"]
+    custom_event_found = False
+    barrier_name_found = False
+
+    for runnable in job.task_groups[0].task_spec.runnables:
+        if runnable.display_name in display_names:
+            display_names.remove(runnable.display_name)
+        if runnable.barrier.name == "hello-barrier":
+            barrier_name_found = True
+        if '{"batch/custom/event": "EVENT_DESCRIPTION"}' in runnable.script.text:
+            custom_event_found = True
+
+    assert not display_names
+    assert custom_event_found
+    assert barrier_name_found
+
+
 @flaky(max_runs=3, min_passes=1)
 def test_script_job(job_name, capsys):
     job = create_script_job(PROJECT, REGION, job_name)
@@ -196,3 +215,9 @@ def test_pd_job(job_name, disk_name):
         additional_test=lambda: _check_policy(job, job_name, disk_names),
         region=region,
     )
+
+
+@flaky(max_runs=3, min_passes=1)
+def test_create_job_with_custom_events(job_name):
+    job = create_job_with_status_events(PROJECT, REGION, job_name)
+    _test_body(job, additional_test=lambda: _check_custom_events(job))