Skip to content

Commit 6b14833

Browse files
TetyanaYahodskaTetyanaYahodska
committed
Implemented tpu_queued_resources_time_bound sample, created test
1 parent 66ae551 commit 6b14833

File tree

3 files changed

+271
-2
lines changed

3 files changed

+271
-2
lines changed

‎tpu/src/main/java/tpu/CreateTimeBoundQueuedResource.java‎

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package tpu;
18+
19+
// [START tpu_queued_resources_time_bound]
20+
import com.google.api.gax.retrying.RetrySettings;
21+
import com.google.cloud.tpu.v2alpha1.CreateQueuedResourceRequest;
22+
import com.google.cloud.tpu.v2alpha1.Node;
23+
import com.google.cloud.tpu.v2alpha1.QueuedResource;
24+
import com.google.cloud.tpu.v2alpha1.TpuClient;
25+
import com.google.cloud.tpu.v2alpha1.TpuSettings;
26+
import com.google.protobuf.Duration;
27+
import com.google.protobuf.Timestamp;
28+
// Uncomment the following line to use Interval or Date
29+
//import com.google.type.Interval;
30+
//import java.util.Date;
31+
import java.io.IOException;
32+
import java.time.Instant;
33+
import java.util.concurrent.ExecutionException;
34+
35+
public class CreateTimeBoundQueuedResource {
36+
37+
public static void main(String[] args)
38+
throws IOException, ExecutionException, InterruptedException {
39+
// TODO(developer): Replace these variables before running the sample.
40+
// Project ID or project number of the Google Cloud project you want to create a node.
41+
String projectId = "YOUR_PROJECT_ID";
42+
// The zone in which to create the TPU.
43+
// For more information about supported TPU types for specific zones,
44+
// see https://cloud.google.com/tpu/docs/regions-zones
45+
String zone = "europe-west4-a";
46+
// The name for your TPU.
47+
String nodeName = "YOUR_NODE_ID";
48+
// The accelerator type that specifies the version and size of the Cloud TPU you want to create.
49+
// For more information about supported accelerator types for each TPU version,
50+
// see https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#versions.
51+
String tpuType = "v2-8";
52+
// Software version that specifies the version of the TPU runtime to install.
53+
// For more information see https://cloud.google.com/tpu/docs/runtimes
54+
String tpuSoftwareVersion = "tpu-vm-tf-2.14.1";
55+
// The name for your Queued Resource.
56+
String queuedResourceId = "QUEUED_RESOURCE_ID";
57+
58+
59+
createTimeBoundQueuedResource(projectId, nodeName,
60+
queuedResourceId, zone, tpuType, tpuSoftwareVersion);
61+
}
62+
63+
// Creates a Queued Resource with time bound configuration.
64+
public static QueuedResource createTimeBoundQueuedResource(
65+
String projectId, String nodeName, String queuedResourceName,
66+
String zone, String tpuType, String tpuSoftwareVersion)
67+
throws IOException, ExecutionException, InterruptedException {
68+
// With these settings the client library handles the Operation's polling mechanism
69+
// and prevent CancellationException error
70+
TpuSettings.Builder clientSettings =
71+
TpuSettings.newBuilder();
72+
clientSettings
73+
.createQueuedResourceSettings()
74+
.setRetrySettings(
75+
RetrySettings.newBuilder()
76+
.setInitialRetryDelay(org.threeten.bp.Duration.ofMillis(5000L))
77+
.setRetryDelayMultiplier(2.0)
78+
.setInitialRpcTimeout(org.threeten.bp.Duration.ZERO)
79+
.setRpcTimeoutMultiplier(1.0)
80+
.setMaxRetryDelay(org.threeten.bp.Duration.ofMillis(45000L))
81+
.setTotalTimeout(org.threeten.bp.Duration.ofHours(24L))
82+
.build());
83+
// Initialize client that will be used to send requests. This client only needs to be created
84+
// once, and can be reused for multiple requests.
85+
try (TpuClient tpuClient = TpuClient.create(clientSettings.build())) {
86+
// Define parent for requests
87+
String parent = String.format("projects/%s/locations/%s", projectId, zone);
88+
// Create a Duration object representing 6 hours.
89+
Duration validAfterDuration = Duration.newBuilder().setSeconds(6 * 3600).build();
90+
// Uncomment the appropriate lines to use other time bound configurations
91+
// Duration validUntilDuration = Duration.newBuilder().setSeconds(6 * 3600).build();
92+
// String validAfterTime = "2024-10-14T09:00:00Z";
93+
// String validUntilTime = "2024-12-14T09:00:00Z";
94+
95+
// Create a node
96+
Node node =
97+
Node.newBuilder()
98+
.setName(nodeName)
99+
.setAcceleratorType(tpuType)
100+
.setRuntimeVersion(tpuSoftwareVersion)
101+
.setQueuedResource(
102+
String.format(
103+
"projects/%s/locations/%s/queuedResources/%s",
104+
projectId, zone, queuedResourceName))
105+
.build();
106+
107+
// Create queued resource
108+
QueuedResource queuedResource =
109+
QueuedResource.newBuilder()
110+
.setName(queuedResourceName)
111+
.setTpu(
112+
QueuedResource.Tpu.newBuilder()
113+
.addNodeSpec(
114+
QueuedResource.Tpu.NodeSpec.newBuilder()
115+
.setParent(parent)
116+
.setNode(node)
117+
.setNodeId(nodeName)
118+
.build())
119+
.build())
120+
.setQueueingPolicy(
121+
QueuedResource.QueueingPolicy.newBuilder()
122+
// You can specify a duration after which a resource should be allocated.
123+
// corresponds --valid-after-duration flag
124+
.setValidAfterDuration(validAfterDuration)
125+
.build())
126+
// Uncomment the appropriate lines to use other time bound configurations
127+
//.setQueueingPolicy(
128+
//QueuedResource.QueueingPolicy.newBuilder()
129+
// You can specify a time after which a resource should be allocated.
130+
// corresponds --valid-until-duration flag
131+
//.setValidUntilDuration(validUntilDuration)
132+
//.build())
133+
//.setQueueingPolicy(
134+
//QueuedResource.QueueingPolicy.newBuilder()
135+
// You can specify a time before which the resource should be allocated.
136+
// corresponds --valid-after-time flag
137+
//.setValidAfterTime(convertStringToTimestamp(validAfterTime))
138+
//.build())
139+
//.setQueueingPolicy(
140+
//QueuedResource.QueueingPolicy.newBuilder()
141+
// You can specify a time after which the resource should be allocated.
142+
// corresponds --valid-until-time flag
143+
//.setValidUntilTime(convertStringToTimestamp(validUntilTime))
144+
//.build())
145+
//.setQueueingPolicy(
146+
//QueuedResource.QueueingPolicy.newBuilder()
147+
// You can specify a time interval before and after which
148+
// the resource should be allocated.
149+
//.setValidInterval(
150+
//Interval.newBuilder()
151+
//.setStartTime(convertStringToTimestamp(validAfterTime))
152+
//.setEndTime(convertStringToTimestamp(validUntilTime))
153+
//.build())
154+
//.build())
155+
.build();
156+
157+
CreateQueuedResourceRequest request =
158+
CreateQueuedResourceRequest.newBuilder()
159+
.setParent(parent)
160+
.setQueuedResource(queuedResource)
161+
.setQueuedResourceId(queuedResourceName)
162+
.build();
163+
164+
QueuedResource response = tpuClient.createQueuedResourceAsync(request).get();
165+
// You can wait until TPU Node is READY,
166+
// and check its status using getTpuVm() from "tpu_vm_get" sample.
167+
System.out.printf("Queued Resource created: %s\n", response.getName());
168+
return response;
169+
}
170+
}
171+
172+
// Method to convert a string timestamp to a Protobuf Timestamp object
173+
private static Timestamp convertStringToTimestamp(String timestampString) {
174+
Instant instant = Instant.parse(timestampString);
175+
return Timestamp.newBuilder()
176+
.setSeconds(instant.getEpochSecond())
177+
.setNanos(instant.getNano())
178+
.build();
179+
}
180+
}
181+
// [END tpu_queued_resources_time_bound]

‎tpu/src/test/java/tpu/CreateTimeBoundQueuedResourceIT.java‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package tpu;
18+
19+
import static com.google.common.truth.Truth.assertThat;
20+
import static com.google.common.truth.Truth.assertWithMessage;
21+
22+
import com.google.api.gax.rpc.NotFoundException;
23+
import com.google.cloud.tpu.v2alpha1.QueuedResource;
24+
import com.google.protobuf.Duration;
25+
import java.io.IOException;
26+
import java.util.UUID;
27+
import java.util.concurrent.ExecutionException;
28+
import java.util.concurrent.TimeUnit;
29+
import org.junit.jupiter.api.AfterAll;
30+
import org.junit.jupiter.api.Assertions;
31+
import org.junit.jupiter.api.BeforeAll;
32+
import org.junit.jupiter.api.Test;
33+
import org.junit.jupiter.api.Timeout;
34+
import org.junit.runner.RunWith;
35+
import org.junit.runners.JUnit4;
36+
37+
@RunWith(JUnit4.class)
38+
@Timeout(value = 6, unit = TimeUnit.MINUTES)
39+
public class CreateTimeBoundQueuedResourceIT {
40+
private static final String PROJECT_ID = System.getenv("GOOGLE_CLOUD_PROJECT");
41+
private static final String ZONE = "us-central1-f";
42+
static String javaVersion = System.getProperty("java.version").substring(0, 2);
43+
private static final String NODE_NAME = "test-tpu-time-bound-queued-resource-" + javaVersion + "-"
44+
+ UUID.randomUUID().toString().substring(0, 8);
45+
private static final String TPU_TYPE = "v2-8";
46+
private static final String TPU_SOFTWARE_VERSION = "tpu-vm-tf-2.17.0-pjrt";
47+
private static final String QUEUED_RESOURCE_NAME = "queued-resource-time-bound-" + javaVersion
48+
+ "-" + UUID.randomUUID().toString().substring(0, 8);
49+
private static final String QUEUED_RESOURCE_PATH_NAME =
50+
String.format("projects/%s/locations/%s/queuedResources/%s",
51+
PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME);
52+
53+
public static void requireEnvVar(String envVarName) {
54+
assertWithMessage(String.format("Missing environment variable '%s' ", envVarName))
55+
.that(System.getenv(envVarName)).isNotEmpty();
56+
}
57+
58+
@BeforeAll
59+
public static void setUp() throws IOException {
60+
requireEnvVar("GOOGLE_APPLICATION_CREDENTIALS");
61+
requireEnvVar("GOOGLE_CLOUD_PROJECT");
62+
63+
// Cleanup existing stale resources.
64+
Util.cleanUpExistingQueuedResources("queued-resource-time-bound-", PROJECT_ID, ZONE);
65+
}
66+
67+
@AfterAll
68+
public static void cleanup() {
69+
DeleteForceQueuedResource.deleteForceQueuedResource(PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME);
70+
71+
// Test that resources are deleted
72+
Assertions.assertThrows(
73+
NotFoundException.class,
74+
() -> GetTpuVm.getTpuVm(PROJECT_ID, ZONE, NODE_NAME));
75+
}
76+
77+
@Test
78+
public void testCreateTimeBoundQueuedResource()
79+
throws IOException, ExecutionException, InterruptedException {
80+
QueuedResource queuedResource = CreateTimeBoundQueuedResource.createTimeBoundQueuedResource(
81+
PROJECT_ID, NODE_NAME, QUEUED_RESOURCE_NAME, ZONE, TPU_TYPE, TPU_SOFTWARE_VERSION);
82+
Duration validAfterDuration = Duration.newBuilder().setSeconds(6 * 3600).build();
83+
84+
assertThat(queuedResource.getName()).isEqualTo(QUEUED_RESOURCE_PATH_NAME);
85+
assertThat(queuedResource.getQueueingPolicy().getValidAfterDuration()
86+
.equals(validAfterDuration));
87+
}
88+
}

‎tpu/src/test/java/tpu/TpuVmIT.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,11 +42,11 @@
4242
@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
4343
public class TpuVmIT {
4444
private static final String PROJECT_ID = System.getenv("GOOGLE_CLOUD_PROJECT");
45-
private static final String ZONE = "us-central1-a";
45+
private static final String ZONE = "us-east5-a";
4646
static String javaVersion = System.getProperty("java.version").substring(0, 2);
4747
private static final String NODE_NAME = "test-tpu-" + javaVersion + "-"
4848
+ UUID.randomUUID().toString().substring(0, 8);
49-
private static final String TPU_TYPE = "v3-8";
49+
private static final String TPU_TYPE = "v5p-8";
5050
private static final String TPU_SOFTWARE_VERSION = "tpu-vm-tf-2.14.1";
5151
private static final String NODE_PATH_NAME =
5252
String.format("projects/%s/locations/%s/nodes/%s", PROJECT_ID, ZONE, NODE_NAME);

0 commit comments

Comments
 (0)