Skip to content

Commit c4cc554

Browse files
authored
refactor(batch): create Batch samples (#9376)
* implemented : batch_create_gpu_job, batch_create_local_ssd_job, batch_create_persistent_disk_job * implemented : batch_create_custom_serive_account, batch_create_using_secret_manager * created tests * created tests * fixes * fixes * fix comments * fixes * fixes
1 parent b836905 commit c4cc554

File tree

7 files changed

+1172
-0
lines changed

7 files changed

+1172
-0
lines changed
Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
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+
package com.example.batch;
16+
17+
// [START batch_notifications]
18+
19+
import com.google.cloud.batch.v1.BatchServiceClient;
20+
import com.google.cloud.batch.v1.CreateJobRequest;
21+
import com.google.cloud.batch.v1.Job;
22+
import com.google.cloud.batch.v1.JobNotification;
23+
import com.google.cloud.batch.v1.JobNotification.Message;
24+
import com.google.cloud.batch.v1.JobNotification.Type;
25+
import com.google.cloud.batch.v1.LogsPolicy;
26+
import com.google.cloud.batch.v1.LogsPolicy.Destination;
27+
import com.google.cloud.batch.v1.Runnable;
28+
import com.google.cloud.batch.v1.Runnable.Script;
29+
import com.google.cloud.batch.v1.TaskGroup;
30+
import com.google.cloud.batch.v1.TaskSpec;
31+
import com.google.cloud.batch.v1.TaskStatus.State;
32+
import com.google.common.collect.Lists;
33+
import com.google.protobuf.Duration;
34+
import java.io.IOException;
35+
import java.util.concurrent.ExecutionException;
36+
import java.util.concurrent.TimeUnit;
37+
import java.util.concurrent.TimeoutException;
38+
39+
public class CreateBatchNotification {
40+
41+
public static void main(String[] args)
42+
throws IOException, ExecutionException, InterruptedException, TimeoutException {
43+
// TODO(developer): Replace these variables before running the sample.
44+
// Project ID or project number of the Google Cloud project you want to use.
45+
String projectId = "YOUR_PROJECT_ID";
46+
// Name of the region you want to use to run the job. Regions that are
47+
// available for Batch are listed on: https://cloud.google.com/batch/docs/get-started#locations
48+
String region = "europe-central2";
49+
// The name of the job that will be created.
50+
// It needs to be unique for each project and region pair.
51+
String jobName = "JOB_NAME";
52+
// The Pub/Sub topic ID to send the notifications to.
53+
String topicId = "TOPIC_ID";
54+
55+
createBatchNotification(projectId, region, jobName, topicId);
56+
}
57+
58+
// Create a Batch job that sends notifications to Pub/Sub
59+
public static Job createBatchNotification(String projectId, String region, String jobName,
60+
String topicId)
61+
throws IOException, ExecutionException, InterruptedException, TimeoutException {
62+
// Initialize client that will be used to send requests. This client only needs to be created
63+
// once, and can be reused for multiple requests.
64+
try (BatchServiceClient batchServiceClient = BatchServiceClient.create()) {
65+
// Define what will be done as part of the job.
66+
Runnable runnable =
67+
Runnable.newBuilder()
68+
.setScript(
69+
Script.newBuilder()
70+
.setText(
71+
"echo Hello world! This is task ${BATCH_TASK_INDEX}. "
72+
+ "This job has a total of ${BATCH_TASK_COUNT} tasks.")
73+
// You can also run a script from a file. Just remember, that needs to be a
74+
// script that's already on the VM that will be running the job.
75+
// Using setText() and setPath() is mutually exclusive.
76+
// .setPath("/tmp/test.sh")
77+
.build())
78+
.build();
79+
80+
TaskSpec task = TaskSpec.newBuilder()
81+
// Jobs can be divided into tasks. In this case, we have only one task.
82+
.addRunnables(runnable)
83+
.setMaxRetryCount(2)
84+
.setMaxRunDuration(Duration.newBuilder().setSeconds(3600).build())
85+
.build();
86+
87+
// Tasks are grouped inside a job using TaskGroups.
88+
// Currently, it's possible to have only one task group.
89+
TaskGroup taskGroup = TaskGroup.newBuilder()
90+
.setTaskCount(3)
91+
.setParallelism(1)
92+
.setTaskSpec(task)
93+
.build();
94+
95+
Job job =
96+
Job.newBuilder()
97+
.addTaskGroups(taskGroup)
98+
.addAllNotifications(buildNotifications(projectId, topicId))
99+
.putLabels("env", "testing")
100+
.putLabels("type", "script")
101+
// We use Cloud Logging as it's an out of the box available option.
102+
.setLogsPolicy(
103+
LogsPolicy.newBuilder().setDestination(Destination.CLOUD_LOGGING))
104+
.build();
105+
106+
CreateJobRequest createJobRequest =
107+
CreateJobRequest.newBuilder()
108+
// The job's parent is the region in which the job will run.
109+
.setParent(String.format("projects/%s/locations/%s", projectId, region))
110+
.setJob(job)
111+
.setJobId(jobName)
112+
.build();
113+
114+
Job result =
115+
batchServiceClient
116+
.createJobCallable()
117+
.futureCall(createJobRequest)
118+
.get(5, TimeUnit.MINUTES);
119+
120+
System.out.printf("Successfully created the job: %s", result.getName());
121+
122+
return result;
123+
}
124+
}
125+
126+
// Creates notification configurations to send messages to Pub/Sub when the state is changed
127+
private static Iterable<JobNotification> buildNotifications(String projectId, String topicId) {
128+
String pubsubTopic = String.format("projects/%s/topics/%s", projectId, topicId);
129+
130+
JobNotification jobStateChanged = JobNotification.newBuilder()
131+
.setPubsubTopic(pubsubTopic)
132+
.setMessage(Message.newBuilder().setType(Type.JOB_STATE_CHANGED))
133+
.build();
134+
135+
JobNotification taskStateChanged = JobNotification.newBuilder()
136+
.setPubsubTopic(pubsubTopic)
137+
.setMessage(Message.newBuilder()
138+
.setType(Type.TASK_STATE_CHANGED)
139+
.setNewTaskState(State.FAILED))
140+
.build();
141+
142+
return Lists.newArrayList(jobStateChanged, taskStateChanged);
143+
}
144+
}
145+
// [END batch_notifications]
Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
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+
package com.example.batch;
16+
17+
// [START batch_create_using_secret_manager]
18+
19+
import com.google.cloud.batch.v1.BatchServiceClient;
20+
import com.google.cloud.batch.v1.CreateJobRequest;
21+
import com.google.cloud.batch.v1.Environment;
22+
import com.google.cloud.batch.v1.Job;
23+
import com.google.cloud.batch.v1.LogsPolicy;
24+
import com.google.cloud.batch.v1.LogsPolicy.Destination;
25+
import com.google.cloud.batch.v1.Runnable;
26+
import com.google.cloud.batch.v1.Runnable.Script;
27+
import com.google.cloud.batch.v1.TaskGroup;
28+
import com.google.cloud.batch.v1.TaskSpec;
29+
import com.google.protobuf.Duration;
30+
import java.io.IOException;
31+
import java.util.concurrent.ExecutionException;
32+
import java.util.concurrent.TimeUnit;
33+
import java.util.concurrent.TimeoutException;
34+
35+
public class CreateBatchUsingSecretManager {
36+
37+
public static void main(String[] args)
38+
throws IOException, ExecutionException, InterruptedException, TimeoutException {
39+
// TODO(developer): Replace these variables before running the sample.
40+
// Project ID or project number of the Google Cloud project you want to use.
41+
String projectId = "YOUR_PROJECT_ID";
42+
// Name of the region you want to use to run the job. Regions that are
43+
// available for Batch are listed on: https://cloud.google.com/batch/docs/get-started#locations
44+
String region = "europe-central2";
45+
// The name of the job that will be created.
46+
// It needs to be unique for each project and region pair.
47+
String jobName = "JOB_NAME";
48+
// The name of the secret variable.
49+
// This variable name is specified in this job's runnables
50+
// and is accessible to all of the runnables that are in the same environment.
51+
String secretVariableName = "VARIABLE_NAME";
52+
// The name of an existing Secret Manager secret.
53+
String secretName = "SECRET_NAME";
54+
// The version of the specified secret that contains the data you want to pass to the job.
55+
// This can be the version number or latest.
56+
String version = "VERSION";
57+
58+
createBatchUsingSecretManager(projectId, region,
59+
jobName, secretVariableName, secretName, version);
60+
}
61+
62+
// Create a basic script job to securely pass sensitive data.
63+
// The data is obtained from Secret Manager secrets
64+
// and set as custom environment variables in the job.
65+
public static Job createBatchUsingSecretManager(String projectId, String region,
66+
String jobName, String secretVariableName,
67+
String secretName, String version)
68+
throws IOException, ExecutionException, InterruptedException, TimeoutException {
69+
// Initialize client that will be used to send requests. This client only needs to be created
70+
// once, and can be reused for multiple requests.
71+
try (BatchServiceClient batchServiceClient = BatchServiceClient.create()) {
72+
// Define what will be done as part of the job.
73+
Runnable runnable =
74+
Runnable.newBuilder()
75+
.setScript(
76+
Script.newBuilder()
77+
.setText(
78+
String.format("echo This is the secret: ${%s}.", secretVariableName))
79+
// You can also run a script from a file. Just remember, that needs to be a
80+
// script that's already on the VM that will be running the job.
81+
// Using setText() and setPath() is mutually exclusive.
82+
// .setPath("/tmp/test.sh")
83+
.build())
84+
.build();
85+
86+
// Construct the resource path to the secret's version.
87+
String secretValue = String
88+
.format("projects/%s/secrets/%s/versions/%s", projectId, secretName, version);
89+
90+
// Set the secret as an environment variable.
91+
Environment.Builder environmentVariable = Environment.newBuilder()
92+
.putSecretVariables(secretVariableName, secretValue);
93+
94+
TaskSpec task = TaskSpec.newBuilder()
95+
// Jobs can be divided into tasks. In this case, we have only one task.
96+
.addRunnables(runnable)
97+
.setEnvironment(environmentVariable)
98+
.setMaxRetryCount(2)
99+
.setMaxRunDuration(Duration.newBuilder().setSeconds(3600).build())
100+
.build();
101+
102+
// Tasks are grouped inside a job using TaskGroups.
103+
// Currently, it's possible to have only one task group.
104+
TaskGroup taskGroup = TaskGroup.newBuilder()
105+
.setTaskSpec(task)
106+
.build();
107+
108+
Job job =
109+
Job.newBuilder()
110+
.addTaskGroups(taskGroup)
111+
.putLabels("env", "testing")
112+
.putLabels("type", "script")
113+
// We use Cloud Logging as it's an out of the box available option.
114+
.setLogsPolicy(
115+
LogsPolicy.newBuilder().setDestination(Destination.CLOUD_LOGGING))
116+
.build();
117+
118+
CreateJobRequest createJobRequest =
119+
CreateJobRequest.newBuilder()
120+
// The job's parent is the region in which the job will run.
121+
.setParent(String.format("projects/%s/locations/%s", projectId, region))
122+
.setJob(job)
123+
.setJobId(jobName)
124+
.build();
125+
126+
Job result =
127+
batchServiceClient
128+
.createJobCallable()
129+
.futureCall(createJobRequest)
130+
.get(5, TimeUnit.MINUTES);
131+
132+
System.out.printf("Successfully created the job: %s", result.getName());
133+
134+
return result;
135+
}
136+
}
137+
}
138+
// [END batch_create_using_secret_manager]

0 commit comments

Comments
 (0)