feat: create job with persistent disk sample (GoogleCloudPlatform#11898)

Stepan Rasputny · gcf-owl-bot[bot] · web-flow · commit cb84e9d58b5b · 2024-07-03T12:53:46.000+02:00
* feat: create job with persistent disk sample * Fix lint issues * fix: added existing disk attachment * fix: added existing disk as a volume * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * fix: added test assert for second disk --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
diff --git a/batch/create/create_with_persistent_disk.py b/batch/create/create_with_persistent_disk.py
@@ -0,0 +1,132 @@
+#  Copyright 2024 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.
+
+# [START batch_create_persistent_disk_job]
+from google.cloud import batch_v1
+
+
+def create_with_pd_job(
+    project_id: str,
+    region: str,
+    job_name: str,
+    disk_name: str,
+    zone: str,
+    existing_disk_name=None,
+) -> batch_v1.Job:
+    """
+    This method shows how to create a sample Batch Job that will run
+    a simple command on Cloud Compute instances with mounted persistent disk.
+
+    Args:
+        project_id: project ID or project number of the Cloud project you want to use.
+        region: 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/get-started#locations
+        job_name: the name of the job that will be created.
+            It needs to be unique for each project and region pair.
+        disk_name: name of the disk to be mounted for your Job.
+        existing_disk_name(optional): existing disk name, which you want to attach to a job
+
+    Returns:
+        A job object representing the job created.
+    """
+    client = batch_v1.BatchServiceClient()
+
+    # Define what will be done as part of the job.
+    task = batch_v1.TaskSpec()
+    runnable = batch_v1.Runnable()
+    runnable.script = batch_v1.Runnable.Script()
+    runnable.script.text = (
+        "echo Hello world from task ${BATCH_TASK_INDEX}. >> /mnt/disks/"
+        + disk_name
+        + "/output_task_${BATCH_TASK_INDEX}.txt"
+    )
+    task.runnables = [runnable]
+    task.max_retry_count = 2
+    task.max_run_duration = "3600s"
+
+    volume = batch_v1.Volume()
+    volume.device_name = disk_name
+    volume.mount_path = f"/mnt/disks/{disk_name}"
+    task.volumes = [volume]
+
+    if existing_disk_name:
+        volume2 = batch_v1.Volume()
+        volume2.device_name = existing_disk_name
+        volume2.mount_path = f"/mnt/disks/{existing_disk_name}"
+        task.volumes.append(volume2)
+
+    # 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
+
+    disk = batch_v1.AllocationPolicy.Disk()
+    # The disk type of the new persistent disk, either pd-standard,
+    # pd-balanced, pd-ssd, or pd-extreme. For Batch jobs, the default is pd-balanced
+    disk.type_ = "pd-balanced"
+    disk.size_gb = 10
+
+    # Policies are used to define on what kind of virtual machines the tasks will run on.
+    # Read more about local disks here: https://cloud.google.com/compute/docs/disks/persistent-disks
+    policy = batch_v1.AllocationPolicy.InstancePolicy()
+    policy.machine_type = "n1-standard-1"
+
+    attached_disk = batch_v1.AllocationPolicy.AttachedDisk()
+    attached_disk.new_disk = disk
+    attached_disk.device_name = disk_name
+    policy.disks = [attached_disk]
+
+    if existing_disk_name:
+        attached_disk2 = batch_v1.AllocationPolicy.AttachedDisk()
+        attached_disk2.existing_disk = (
+            f"projects/{project_id}/zones/{zone}/disks/{existing_disk_name}"
+        )
+        attached_disk2.device_name = existing_disk_name
+        policy.disks.append(attached_disk2)
+
+    instances = batch_v1.AllocationPolicy.InstancePolicyOrTemplate()
+    instances.policy = policy
+
+    allocation_policy = batch_v1.AllocationPolicy()
+    allocation_policy.instances = [instances]
+
+    location = batch_v1.AllocationPolicy.LocationPolicy()
+    location.allowed_locations = [f"zones/{zone}"]
+    allocation_policy.location = location
+
+    job = batch_v1.Job()
+    job.task_groups = [group]
+    job.allocation_policy = allocation_policy
+    job.labels = {"env": "testing", "type": "script"}
+
+    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_create_persistent_disk_job]
+
+if __name__ == "__main__":
+    import google.auth
+
+    PROJECT = google.auth.default()[1]
+    REGION = "europe-west4"
+    ZONE = "europe-west4-c"
+    job = create_with_pd_job(PROJECT, REGION, "pd-job-batch", "pd-1", ZONE)
+    print(job)
diff --git a/batch/tests/test_basics.py b/batch/tests/test_basics.py
@@ -15,6 +15,7 @@
 
 from collections.abc import Callable
 import time
+from typing import Tuple
 import uuid
 
 from flaky import flaky
@@ -26,6 +27,7 @@
 
 from ..create.create_with_container_no_mounting import create_container_job
 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
 from ..create.create_with_service_account import create_with_custom_service_account_job
 from ..create.create_with_ssd import create_local_ssd_job
@@ -40,6 +42,7 @@
 PROJECT = google.auth.default()[1]
 REGION = "europe-central2"
 ZONE = "europe-central2-b"
+
 TIMEOUT = 600  # 10 minutes
 
 WAIT_STATES = {
@@ -68,7 +71,7 @@ def service_account() -> str:
 
 @pytest.fixture
 def disk_name():
-    return f"test-ssd-{uuid.uuid4().hex[:10]}"
+    return f"test-disk-{uuid.uuid4().hex[:10]}"
 
 
 def _test_body(test_job: batch_v1.Job, additional_test: Callable = None, region=REGION):
@@ -108,6 +111,12 @@ def _check_tasks(job_name):
     print("Tasks tested")
 
 
+def _check_policy(job: batch_v1.Job, job_name: str, disk_names: Tuple[str]):
+    assert job_name in job.name
+    assert job.allocation_policy.instances[0].policy.disks[0].device_name in disk_names
+    assert job.allocation_policy.instances[0].policy.disks[1].device_name in disk_names
+
+
 def _check_logs(job, capsys):
     print_job_logs(PROJECT, job)
     output = [
@@ -155,3 +164,19 @@ def test_service_account_job(job_name, service_account):
 def test_ssd_job(job_name: str, disk_name: str, capsys: "pytest.CaptureFixture[str]"):
     job = create_local_ssd_job(PROJECT, REGION, job_name, disk_name)
     _test_body(job, additional_test=lambda: _check_logs(job, capsys))
+
+
+@flaky(max_runs=3, min_passes=1)
+def test_pd_job(job_name, disk_name):
+    region = "europe-north1"
+    zone = "europe-north1-c"
+    existing_disk_name = "permanent-batch-testing"
+    job = create_with_pd_job(
+        PROJECT, region, job_name, disk_name, zone, existing_disk_name
+    )
+    disk_names = (disk_name, existing_disk_name)
+    _test_body(
+        job,
+        additional_test=lambda: _check_policy(job, job_name, disk_names),
+        region=region,
+    )
