diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index b8ca6cb4be0d..7f85efce1a2e 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -20,6 +20,7 @@ /appengine/flexible/django_cloudsql/**/* @glasnt @GoogleCloudPlatform/aap-dpes @GoogleCloudPlatform/python-samples-reviewers /appengine/standard_python3/spanner/* @GoogleCloudPlatform/api-spanner-python @GoogleCloudPlatform/python-samples-reviewers /auth/**/* @arithmetic1728 @GoogleCloudPlatform/python-samples-reviewers +/batch/**/* @m-strzelczyk @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/python-samples-reviewers /bigquery/**/* @chalmerlowe @GoogleCloudPlatform/python-samples-reviewers /bigquery/remote_function/**/* @autoerr @GoogleCloudPlatform/python-samples-reviewers /billing/**/* @GoogleCloudPlatform/billing-samples-maintainers @GoogleCloudPlatform/python-samples-reviewers @@ -29,7 +30,7 @@ /cloud-sql/**/* @GoogleCloudPlatform/infra-db-dpes @GoogleCloudPlatform/python-samples-reviewers /codelabs/**/* @GoogleCloudPlatform/python-samples-reviewers /composer/**/* @leahecole @rachael-ds @rafalbiegacz @GoogleCloudPlatform/python-samples-reviewers -/compute/**/* @m-strzelczyk @GoogleCloudPlatform/python-samples-reviewers +/compute/**/* @m-strzelczyk @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/python-samples-reviewers /container/**/* @GoogleCloudPlatform/dee-platform-ops @GoogleCloudPlatform/python-samples-reviewers /data-science-onramp/ @leahecole @bradmiro @GoogleCloudPlatform/python-samples-reviewers /dataflow/**/* @davidcavazos @GoogleCloudPlatform/python-samples-reviewers diff --git a/.github/blunderbuss.yml b/.github/blunderbuss.yml index 5a26676163b6..b6e017af18fc 100644 --- a/.github/blunderbuss.yml +++ b/.github/blunderbuss.yml @@ -25,6 +25,10 @@ assign_issues_by: - 'api: auth' to: - arithmetic1728 +- labels: + - 'api: batch' + to: + - m-strzelczyk - labels: - 'api: bigquery' to: diff --git a/batch/__init__.py b/batch/__init__.py new file mode 100644 index 000000000000..4bbe0ffdb061 --- /dev/null +++ b/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. diff --git a/batch/create/__init__.py b/batch/create/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/batch/create/create_with_container_no_mounting.py b/batch/create/create_with_container_no_mounting.py new file mode 100644 index 000000000000..fd7f655e478b --- /dev/null +++ b/batch/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] diff --git a/batch/create/create_with_mounted_bucket.py b/batch/create/create_with_mounted_bucket.py new file mode 100644 index 000000000000..d64821cfa632 --- /dev/null +++ b/batch/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] diff --git a/batch/create/create_with_script_no_mounting.py b/batch/create/create_with_script_no_mounting.py new file mode 100644 index 000000000000..87a1dfe680ad --- /dev/null +++ b/batch/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] diff --git a/batch/create/create_with_template.py b/batch/create/create_with_template.py new file mode 100644 index 000000000000..c37a94f9e352 --- /dev/null +++ b/batch/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] diff --git a/batch/delete/delete_job.py b/batch/delete/delete_job.py new file mode 100644 index 000000000000..f4ebdcdea0ab --- /dev/null +++ b/batch/delete/delete_job.py @@ -0,0 +1,37 @@ +# 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_delete_job] +from google.api_core.operation import Operation + +from google.cloud import batch_v1 + + +def delete_job(project_id: str, region: str, job_name: str) -> Operation: + """ + Triggers the deletion of a Job. + + Args: + project_id: project ID or project number of the Cloud project you want to use. + region: name of the region hosts the job. + job_name: the name of the job that you want to delete. + + Returns: + An operation object related to the deletion. You can call `.result()` + on it to wait for its completion. + """ + client = batch_v1.BatchServiceClient() + + return client.delete_job(name=f"projects/{project_id}/locations/{region}/jobs/{job_name}") +# [END batch_delete_job] diff --git a/batch/get/get_job.py b/batch/get/get_job.py new file mode 100644 index 000000000000..8dff7376ecf7 --- /dev/null +++ b/batch/get/get_job.py @@ -0,0 +1,35 @@ +# 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_get_job] + +from google.cloud import batch_v1 + + +def get_job(project_id: str, region: str, job_name: str) -> batch_v1.Job: + """ + Retrieve information about a Batch Job. + + Args: + project_id: project ID or project number of the Cloud project you want to use. + region: name of the region hosts the job. + job_name: the name of the job you want to retrieve information about. + + Returns: + A Job object representing the specified job. + """ + client = batch_v1.BatchServiceClient() + + return client.get_job(name=f"projects/{project_id}/locations/{region}/jobs/{job_name}") +# [END batch_get_job] diff --git a/batch/get/get_task.py b/batch/get/get_task.py new file mode 100644 index 000000000000..ea5fc5c9c22f --- /dev/null +++ b/batch/get/get_task.py @@ -0,0 +1,38 @@ +# 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_get_task] + +from google.cloud import batch_v1 + + +def get_task(project_id: str, region: str, job_name: str, group_name: str, task_number: int) -> batch_v1.Task: + """ + Retrieve information about a Task. + + Args: + project_id: project ID or project number of the Cloud project you want to use. + region: name of the region hosts the job. + job_name: the name of the job you want to retrieve information about. + group_name: the name of the group that owns the task you want to check. Usually it's `group0`. + task_number: number of the task you want to look up. + + Returns: + A Task object representing the specified task. + """ + client = batch_v1.BatchServiceClient() + + return client.get_task(name=f"projects/{project_id}/locations/{region}/jobs/{job_name}" + f"/taskGroups/{group_name}/tasks/{task_number}") +# [END batch_get_task] diff --git a/batch/list/list_jobs.py b/batch/list/list_jobs.py new file mode 100644 index 000000000000..e52f4defe876 --- /dev/null +++ b/batch/list/list_jobs.py @@ -0,0 +1,35 @@ +# 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_list_jobs] +from typing import Iterable + +from google.cloud import batch_v1 + + +def list_jobs(project_id: str, region: str) -> Iterable[batch_v1.Job]: + """ + Get a list of all jobs defined in given region. + + Args: + project_id: project ID or project number of the Cloud project you want to use. + region: name of the region hosting the jobs. + + Returns: + An iterable collection of Job object. + """ + client = batch_v1.BatchServiceClient() + + return client.list_jobs(parent=f"projects/{project_id}/locations/{region}") +# [END batch_list_jobs] diff --git a/batch/list/list_tasks.py b/batch/list/list_tasks.py new file mode 100644 index 000000000000..9ef6674ec8d1 --- /dev/null +++ b/batch/list/list_tasks.py @@ -0,0 +1,37 @@ +# 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_list_tasks] +from typing import Iterable + +from google.cloud import batch_v1 + + +def list_tasks(project_id: str, region: str, job_name: str, group_name: str) -> Iterable[batch_v1.Task]: + """ + Get a list of all jobs defined in given region. + + Args: + project_id: project ID or project number of the Cloud project you want to use. + region: name of the region hosting the jobs. + job_name: name of the job which tasks you want to list. + group_name: name of the group of tasks. Usually it's `group0`. + + Returns: + An iterable collection of Task objects. + """ + client = batch_v1.BatchServiceClient() + + return client.list_tasks(parent=f"projects/{project_id}/locations/{region}/jobs/{job_name}/taskGroups/{group_name}") +# [END batch_list_tasks] diff --git a/batch/logs/read_job_logs.py b/batch/logs/read_job_logs.py new file mode 100644 index 000000000000..d9c227a1082b --- /dev/null +++ b/batch/logs/read_job_logs.py @@ -0,0 +1,39 @@ +# 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_job_logs] +from typing import NoReturn + +from google.cloud import batch_v1 +from google.cloud import logging + + +def print_job_logs(project_id: str, job: batch_v1.Job) -> NoReturn: + """ + Prints the log messages created by given job. + + Args: + project_id: name of the project hosting the job. + job: the job which logs you want to print. + """ + # Initialize client that will be used to send requests across threads. This + # client only needs to be created once, and can be reused for multiple requests. + log_client = logging.Client(project=project_id) + logger = log_client.logger("batch_task_logs") + + for log_entry in logger.list_entries(filter_=f"labels.job_uid={job.uid}"): + print(log_entry.payload) + +# [END batch_job_logs] diff --git a/batch/noxfile_config.py b/batch/noxfile_config.py new file mode 100644 index 000000000000..d27b90b5086b --- /dev/null +++ b/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", +} diff --git a/batch/requirements-test.txt b/batch/requirements-test.txt new file mode 100644 index 000000000000..586e344619cf --- /dev/null +++ b/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 diff --git a/batch/requirements.txt b/batch/requirements.txt new file mode 100644 index 000000000000..7edc9015390f --- /dev/null +++ b/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 diff --git a/batch/tests/__init__.py b/batch/tests/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/batch/tests/test_basics.py b/batch/tests/test_basics.py new file mode 100644 index 000000000000..829a92729062 --- /dev/null +++ b/batch/tests/test_basics.py @@ -0,0 +1,100 @@ +# 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 time +from typing import Callable +import uuid + +import google.auth +from google.cloud import batch_v1 +import pytest + +from ..create.create_with_container_no_mounting import create_container_job +from ..create.create_with_script_no_mounting import create_script_job + +from ..delete.delete_job import delete_job +from ..get.get_job import get_job +from ..get.get_task import get_task +from ..list.list_jobs import list_jobs +from ..list.list_tasks import list_tasks +from ..logs.read_job_logs import print_job_logs + +PROJECT = google.auth.default()[1] +REGION = 'europe-north1' + +TIMEOUT = 600 # 10 minutes + +WAIT_STATES = { + batch_v1.JobStatus.State.STATE_UNSPECIFIED, + batch_v1.JobStatus.State.QUEUED, + batch_v1.JobStatus.State.RUNNING, + batch_v1.JobStatus.State.SCHEDULED, + batch_v1.JobStatus.State.DELETION_IN_PROGRESS +} + + +@pytest.fixture +def job_name(): + return f"test-job-{uuid.uuid4().hex[:10]}" + + +def _test_body(test_job: batch_v1.Job, additional_test: Callable = None): + start_time = time.time() + try: + while test_job.status.state in WAIT_STATES: + if time.time() - start_time > TIMEOUT: + pytest.fail("Timed out while waiting for job to complete!") + test_job = get_job(PROJECT, REGION, test_job.name.rsplit('/', maxsplit=1)[1]) + time.sleep(5) + + assert test_job.status.state == batch_v1.JobStatus.State.SUCCEEDED + + for job in list_jobs(PROJECT, REGION): + if test_job.uid == job.uid: + break + else: + pytest.fail(f"Couldn't find job {test_job.uid} on the list of jobs.") + + if additional_test: + additional_test() + finally: + delete_job(PROJECT, REGION, test_job.name.rsplit('/', maxsplit=1)[1]).result() + + for job in list_jobs(PROJECT, REGION): + if job.uid == test_job.uid: + pytest.fail("The test job should be deleted at this point!") + + +def _check_tasks(job_name): + tasks = list_tasks(PROJECT, REGION, job_name, 'group0') + assert len(list(tasks)) == 4 + for i in range(4): + assert get_task(PROJECT, REGION, job_name, 'group0', i) is not None + print('Tasks tested') + + +def _check_logs(job, capsys): + print_job_logs(PROJECT, job) + output = [line for line in capsys.readouterr().out.splitlines(keepends=False) if line != ""] + assert len(output) == 4 + assert all("Hello world!" in log_msg for log_msg in output) + + +def test_script_job(job_name, capsys): + job = create_script_job(PROJECT, REGION, job_name) + _test_body(job, additional_test=lambda: _check_logs(job, capsys)) + + +def test_container_job(job_name): + job = create_container_job(PROJECT, REGION, job_name) + _test_body(job, additional_test=lambda: _check_tasks(job_name)) diff --git a/batch/tests/test_bucket.py b/batch/tests/test_bucket.py new file mode 100644 index 000000000000..ad8a347fbba0 --- /dev/null +++ b/batch/tests/test_bucket.py @@ -0,0 +1,70 @@ +# 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 uuid + + +import google.auth +from google.cloud import batch_v1 +from google.cloud import storage +import pytest + +from .test_basics import _test_body +from ..create.create_with_mounted_bucket import create_script_job_with_bucket + +PROJECT = google.auth.default()[1] +REGION = 'europe-north1' + +TIMEOUT = 600 # 10 minutes + +WAIT_STATES = { + batch_v1.JobStatus.State.STATE_UNSPECIFIED, + batch_v1.JobStatus.State.QUEUED, + batch_v1.JobStatus.State.RUNNING, + batch_v1.JobStatus.State.SCHEDULED, +} + + +@pytest.fixture +def job_name(): + return f"test-job-{uuid.uuid4().hex[:10]}" + + +@pytest.fixture() +def test_bucket(): + bucket_name = f"test-bucket-{uuid.uuid4().hex[:8]}" + client = storage.Client() + client.create_bucket(bucket_name, location="eu") + + yield bucket_name + + bucket = client.get_bucket(bucket_name) + bucket.delete(force=True) + + +def _test_bucket_content(test_bucket): + client = storage.Client() + bucket = client.get_bucket(test_bucket) + + file_name_template = "output_task_{task_number}.txt" + file_content_template = "Hello world from task {task_number}.\n" + + for i in range(4): + blob = bucket.blob(file_name_template.format(task_number=i)) + content = blob.download_as_bytes().decode() + assert content == file_content_template.format(task_number=i) + + +def test_bucket_job(job_name, test_bucket): + job = create_script_job_with_bucket(PROJECT, REGION, job_name, test_bucket) + _test_body(job, lambda: _test_bucket_content(test_bucket)) diff --git a/batch/tests/test_template.py b/batch/tests/test_template.py new file mode 100644 index 000000000000..572811731e0a --- /dev/null +++ b/batch/tests/test_template.py @@ -0,0 +1,114 @@ +# 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 uuid + +import google.auth +from google.cloud import batch_v1 +from google.cloud import compute_v1 +from google.cloud import resourcemanager_v3 +import pytest + + +from .test_basics import _test_body + +from ..create.create_with_template import create_script_job_with_template + +PROJECT = google.auth.default()[1] + +PROJECT_NUMBER = resourcemanager_v3.ProjectsClient().get_project(name=f"projects/{PROJECT}").name.split("/")[1] + +REGION = 'europe-north1' + +TIMEOUT = 600 # 10 minutes + +WAIT_STATES = { + batch_v1.JobStatus.State.STATE_UNSPECIFIED, + batch_v1.JobStatus.State.QUEUED, + batch_v1.JobStatus.State.RUNNING, + batch_v1.JobStatus.State.SCHEDULED, +} + + +@pytest.fixture +def job_name(): + return f"test-job-{uuid.uuid4().hex[:10]}" + + +@pytest.fixture +def instance_template(): + disk = compute_v1.AttachedDisk() + initialize_params = compute_v1.AttachedDiskInitializeParams() + initialize_params.source_image = ( + "projects/ubuntu-os-cloud/global/images/family/ubuntu-2204-lts" + ) + initialize_params.disk_size_gb = 25 + initialize_params.disk_type = 'pd-balanced' + disk.initialize_params = initialize_params + disk.auto_delete = True + disk.boot = True + + network_interface = compute_v1.NetworkInterface() + network_interface.name = "global/networks/default" + + access = compute_v1.AccessConfig() + access.type_ = compute_v1.AccessConfig.Type.ONE_TO_ONE_NAT.name + access.name = "External NAT" + access.network_tier = access.NetworkTier.PREMIUM.name + network_interface.access_configs = [access] + + template = compute_v1.InstanceTemplate() + template.name = "test-template-" + uuid.uuid4().hex[:10] + template.properties = compute_v1.InstanceProperties() + template.properties.disks = [disk] + template.properties.machine_type = "e2-standard-16" + template.properties.network_interfaces = [network_interface] + + template.properties.scheduling = compute_v1.Scheduling() + template.properties.scheduling.on_host_maintenance = compute_v1.Scheduling.OnHostMaintenance.MIGRATE.name + template.properties.scheduling.provisioning_model = compute_v1.Scheduling.ProvisioningModel.STANDARD.name + template.properties.scheduling.automatic_restart = True + + template.properties.service_accounts = [ + { + "email": f"{PROJECT_NUMBER}-compute@developer.gserviceaccount.com", + "scopes": [ + "https://www.googleapis.com/auth/devstorage.read_only", + "https://www.googleapis.com/auth/logging.write", + "https://www.googleapis.com/auth/monitoring.write", + "https://www.googleapis.com/auth/servicecontrol", + "https://www.googleapis.com/auth/service.management.readonly", + "https://www.googleapis.com/auth/trace.append" + ] + } + ] + + template_client = compute_v1.InstanceTemplatesClient() + operation_client = compute_v1.GlobalOperationsClient() + op = template_client.insert_unary( + project=PROJECT, instance_template_resource=template + ) + operation_client.wait(project=PROJECT, operation=op.name) + + template = template_client.get(project=PROJECT, instance_template=template.name) + + yield template + + op = template_client.delete_unary(project=PROJECT, instance_template=template.name) + operation_client.wait(project=PROJECT, operation=op.name) + + +def test_template_job(job_name, instance_template): + job = create_script_job_with_template(PROJECT, REGION, job_name, instance_template.self_link) + _test_body(job)