migrate code from googleapis/python-batch

compute/batch/__init__.py

@@ -0,0 +1,13 @@
+#  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.

compute/batch/noxfile_config.py

@@ -0,0 +1,17 @@
+#  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.
+
+TEST_CONFIG_OVERRIDE = {
+    "gcloud_project_env": "BUILD_SPECIFIC_GCLOUD_PROJECT",
+}

compute/batch/requirements-test.txt

@@ -0,0 +1,4 @@
+pytest==7.2.0
+google-cloud-compute==1.6.1
+google-cloud-resource-manager==1.6.3
+google-cloud-storage==2.5.0

compute/batch/requirements.txt

@@ -0,0 +1,4 @@
+isort==5.10.1
+black==22.10.0
+google-cloud-batch==0.4.1
+google-cloud-logging==3.2.5

compute/batch/snippets/create/create_with_container_no_mounting.py

@@ -0,0 +1,87 @@
+#  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.
+
+# [START batch_create_container_job]
+from google.cloud import batch_v1
+
+
+def create_container_job(project_id: str, region: str, job_name: str) -> batch_v1.Job:
+    """
+    This method shows how to create a sample Batch Job that will run
+    a simple command inside a container on Cloud Compute instances.
+
+    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.
+
+    Returns:
+        A job object representing the job created.
+    """
+    client = batch_v1.BatchServiceClient()
+
+    # Define what will be done as part of the job.
+    runnable = batch_v1.Runnable()
+    runnable.container = batch_v1.Runnable.Container()
+    runnable.container.image_uri = "gcr.io/google-containers/busybox"
+    runnable.container.entrypoint = "/bin/sh"
+    runnable.container.commands = ["-c", "echo Hello world! This is task ${BATCH_TASK_INDEX}. This job has a total of ${BATCH_TASK_COUNT} tasks."]
+
+    # Jobs can be divided into tasks. In this case, we have only one task.
+    task = batch_v1.TaskSpec()
+    task.runnables = [runnable]
+
+    # 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_create_container_job]

compute/batch/snippets/create/create_with_mounted_bucket.py

@@ -0,0 +1,91 @@
+#  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.
+
+# [START batch_create_script_job_with_bucket]
+from google.cloud import batch_v1
+
+
+def create_script_job_with_bucket(project_id: str, region: str, job_name: str, bucket_name: str) -> batch_v1.Job:
+    """
+    This method shows how to create a sample Batch Job that will run
+    a simple command on Cloud Compute instances.
+
+    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.
+        bucket_name: name of the bucket to be mounted for your 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/share/output_task_${BATCH_TASK_INDEX}.txt"
+    task.runnables = [runnable]
+
+    gcs_bucket = batch_v1.GCS()
+    gcs_bucket.remote_path = bucket_name
+    gcs_volume = batch_v1.Volume()
+    gcs_volume.gcs = gcs_bucket
+    gcs_volume.mount_path = '/mnt/share'
+    task.volumes = [gcs_volume]
+
+    # We can specify what resources are requested by each task.
+    resources = batch_v1.ComputeResource()
+    resources.cpu_milli = 500  # in milliseconds per cpu-second. This means the task requires 50% of a single CPUs.
+    resources.memory_mib = 16
+    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
+    allocation_policy = batch_v1.AllocationPolicy()
+    policy = batch_v1.AllocationPolicy.InstancePolicy()
+    policy.machine_type = "e2-standard-4"
+    instances = batch_v1.AllocationPolicy.InstancePolicyOrTemplate()
+    instances.policy = policy
+    allocation_policy.instances = [instances]
+
+    job = batch_v1.Job()
+    job.task_groups = [group]
+    job.allocation_policy = allocation_policy
+    job.labels = {"env": "testing", "type": "script", "mount": "bucket"}
+    # 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_create_script_job_with_bucket]

compute/batch/snippets/create/create_with_script_no_mounting.py

@@ -0,0 +1,87 @@
+#  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.
+
+# [START batch_create_script_job]
+from google.cloud import batch_v1
+
+
+def create_script_job(project_id: str, region: str, job_name: str) -> batch_v1.Job:
+    """
+    This method shows how to create a sample Batch Job that will run
+    a simple command on Cloud Compute instances.
+
+    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.
+
+    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! This is task ${BATCH_TASK_INDEX}. This job has a total of ${BATCH_TASK_COUNT} tasks."
+    # You can also run a script from a file. Just remember, that needs to be a script that's
+    # already on the VM that will be running the job. Using runnable.script.text and runnable.script.path is mutually
+    # exclusive.
+    # runnable.script.path = '/tmp/test.sh'
+    task.runnables = [runnable]
+
+    # 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
+    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
+    allocation_policy = batch_v1.AllocationPolicy()
+    policy = batch_v1.AllocationPolicy.InstancePolicy()
+    policy.machine_type = "e2-standard-4"
+    instances = batch_v1.AllocationPolicy.InstancePolicyOrTemplate()
+    instances.policy = policy
+    allocation_policy.instances = [instances]
+
+    job = batch_v1.Job()
+    job.task_groups = [group]
+    job.allocation_policy = allocation_policy
+    job.labels = {"env": "testing", "type": "script"}
+    # 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_create_script_job]

compute/batch/snippets/create/create_with_template.py

@@ -0,0 +1,88 @@
+#  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.
+
+# [START batch_create_job_with_template]
+from google.cloud import batch_v1
+
+
+def create_script_job_with_template(project_id: str, region: str, job_name: str, template_link: str) -> batch_v1.Job:
+    """
+    This method shows how to create a sample Batch Job that will run
+    a simple command on Cloud Compute instances created using a provided Template.
+
+    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.
+        template_link: a link to an existing Instance Template. Acceptable formats:
+            * "projects/{project_id}/global/instanceTemplates/{template_name}"
+            * "{template_name}" - if the template is defined in the same project as used to create the 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! This is task ${BATCH_TASK_INDEX}. This job has a total of ${BATCH_TASK_COUNT} tasks."
+    # You can also run a script from a file. Just remember, that needs to be a script that's
+    # already on the VM that will be running the job. Using runnable.script.text and runnable.script.path is mutually
+    # exclusive.
+    # runnable.script.path = '/tmp/test.sh'
+    task.runnables = [runnable]
+
+    # 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
+    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 an instance template that defines all the
+    # required parameters.
+    allocation_policy = batch_v1.AllocationPolicy()
+    instances = batch_v1.AllocationPolicy.InstancePolicyOrTemplate()
+    instances.instance_template = template_link
+    allocation_policy.instances = [instances]
+
+    job = batch_v1.Job()
+    job.task_groups = [group]
+    job.allocation_policy = allocation_policy
+    job.labels = {"env": "testing", "type": "script"}
+    # 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_create_job_with_template]