Skip to content

Commit 7444528

Browse files
Thoughtseize1Thoughtseize1
authored
feat: Create Batch job with custom network sample (GoogleCloudPlatform#12280)
* Create Batch job with custom network
1 parent ef7da8d commit 7444528

2 files changed

Lines changed: 139 additions & 0 deletions

File tree

‎batch/create/create_with_specific_network.py‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import google.auth
16+
17+
# [START batch_create_custom_network]
18+
from google.cloud import batch_v1
19+
20+
21+
def create_with_custom_network(
22+
project_id: str,
23+
region: str,
24+
network_name: str,
25+
subnet_name: str,
26+
job_name: str,
27+
) -> batch_v1.Job:
28+
"""Create a Batch job that runs on a specific network and subnet.
29+
30+
Args:
31+
project_id: project ID or project number of the Cloud project you want to use.
32+
region: name of the region you want to use to run the job. Regions that are
33+
available for Batch are listed on: https://cloud.google.com/batch/docs/locations
34+
network_name: The name of a VPC network in the current project or a Shared VPC network.
35+
subnet_name: Name of the subnetwork to be used within the specified region.
36+
job_name: the name of the job that will be created.
37+
It needs to be unique for each project and region pair.
38+
Returns:
39+
A job object representing the job created.
40+
"""
41+
42+
client = batch_v1.BatchServiceClient()
43+
44+
# Define what will be done as part of the job.
45+
runnable = batch_v1.Runnable()
46+
runnable.script = batch_v1.Runnable.Script()
47+
runnable.script.text = "echo Hello world! This is task ${BATCH_TASK_INDEX}. This job has a total of ${BATCH_TASK_COUNT} tasks."
48+
49+
# Jobs can be divided into tasks. In this case, we have only one task.
50+
task = batch_v1.TaskSpec()
51+
task.runnables = [runnable]
52+
53+
# We can specify what resources are requested by each task.
54+
resources = batch_v1.ComputeResource()
55+
resources.cpu_milli = 2000 # in milliseconds per cpu-second. This means the task requires 2 whole CPUs.
56+
resources.memory_mib = 16 # in MiB
57+
task.compute_resource = resources
58+
59+
task.max_retry_count = 2
60+
task.max_run_duration = "3600s"
61+
62+
# Tasks are grouped inside a job using TaskGroups.
63+
# Currently, it's possible to have only one task group.
64+
group = batch_v1.TaskGroup()
65+
group.task_count = 3
66+
group.task_spec = task
67+
68+
# Policies are used to define on what kind of virtual machines the tasks will run on.
69+
# In this case, we tell the system to use "e2-standard-4" machine type.
70+
# Read more about machine types here: https://cloud.google.com/compute/docs/machine-types
71+
policy = batch_v1.AllocationPolicy.InstancePolicy()
72+
policy.machine_type = "e2-standard-4"
73+
instances = batch_v1.AllocationPolicy.InstancePolicyOrTemplate()
74+
instances.policy = policy
75+
allocation_policy = batch_v1.AllocationPolicy()
76+
allocation_policy.instances = [instances]
77+
78+
# Create a NetworkInterface object to specify network settings for the job
79+
network_interface = batch_v1.AllocationPolicy.NetworkInterface()
80+
# Set the network to the specified network name within the project
81+
network_interface.network = f"projects/{project_id}/global/networks/{network_name}"
82+
# Set the subnetwork to the specified subnetwork within the region
83+
network_interface.subnetwork = (
84+
f"projects/{project_id}/regions/{region}/subnetworks/{subnet_name}"
85+
)
86+
allocation_policy.network.network_interfaces = [network_interface]
87+
88+
job = batch_v1.Job()
89+
job.task_groups = [group]
90+
job.allocation_policy = allocation_policy
91+
# We use Cloud Logging as it's an out of the box available option
92+
job.logs_policy = batch_v1.LogsPolicy()
93+
job.logs_policy.destination = batch_v1.LogsPolicy.Destination.CLOUD_LOGGING
94+
95+
create_request = batch_v1.CreateJobRequest()
96+
create_request.job = job
97+
create_request.job_id = job_name
98+
# The job's parent is the region in which the job will run
99+
create_request.parent = f"projects/{project_id}/locations/{region}"
100+
return client.create_job(create_request)
101+
102+
103+
# [END batch_create_custom_network]
104+
105+
if __name__ == "__main__":
106+
PROJECT_ID = google.auth.default()[1]
107+
REGION = "europe-west"
108+
vpc_network = "VPC-network-name"
109+
subnet = "subnet-name"
110+
job_name = "test-job-name"
111+
job = create_with_custom_network(
112+
PROJECT_ID,
113+
REGION,
114+
vpc_network,
115+
subnet,
116+
job_name,
117+
)

‎batch/tests/test_basics.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
from ..create.create_with_script_no_mounting import create_script_job
3636
from ..create.create_with_secret_manager import create_with_secret_manager
3737
from ..create.create_with_service_account import create_with_custom_service_account_job
38+
from ..create.create_with_specific_network import create_with_custom_network
3839
from ..create.create_with_ssd import create_local_ssd_job
3940

4041
from ..delete.delete_job import delete_job
@@ -189,6 +190,17 @@ def _check_nfs_mounting(
189190
assert expected_script_text in job.task_groups[0].task_spec.runnables[0].script.text
190191

191192

193+
def _check_custom_networks(job, network_name, subnet_name):
194+
assert (
195+
f"/networks/{network_name}"
196+
in job.allocation_policy.network.network_interfaces[0].network
197+
)
198+
assert (
199+
f"/subnetworks/{subnet_name}"
200+
in job.allocation_policy.network.network_interfaces[0].subnetwork
201+
)
202+
203+
192204
@flaky(max_runs=3, min_passes=1)
193205
def test_script_job(job_name, capsys):
194206
job = create_script_job(PROJECT, REGION, job_name)
@@ -285,3 +297,13 @@ def test_check_nfs_job(job_name):
285297
region="us-central1",
286298
project=project_with_nfs_filestore,
287299
)
300+
301+
302+
@flaky(max_runs=3, min_passes=1)
303+
def test_job_with_custom_network(job_name):
304+
network_name = "default"
305+
subnet = "default"
306+
job = create_with_custom_network(PROJECT, REGION, network_name, subnet, job_name)
307+
_test_body(
308+
job, additional_test=lambda: _check_custom_networks(job, network_name, subnet)
309+
)

0 commit comments

Comments
 (0)