Skip to content

Commit 772c39c

Browse files
TetyanaYahodskaTetyanaYahodska
authored
feat(tpu): add tpu queued resources create/get/delete samples (#9613)
* Changed package, added information to CODEOWNERS * Added information to CODEOWNERS * Added timeout * Fixed parameters for test * Fixed DeleteTpuVm and naming * Added comment, created Util class * Fixed naming * Fixed whitespace * Split PR into smaller, deleted redundant code * Implemented tpu_queued_resources_create, tpu_queued_resources_get, tpu_queued_resources_delete_force and tpu_queued_resources_delete samples, created tests * Fixed test * Fixed tests * Fixed error massage * Fixed typo * Fixed zone * Fixed test * Fixed code * Deleted commented imports * Fixed code as requested in comments
1 parent 05b16a2 commit 772c39c

File tree

4 files changed

+204
-3
lines changed

4 files changed

+204
-3
lines changed

‎tpu/src/main/java/tpu/CreateQueuedResource.java

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
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_create]
20+
import com.google.cloud.tpu.v2alpha1.CreateQueuedResourceRequest;
21+
import com.google.cloud.tpu.v2alpha1.Node;
22+
import com.google.cloud.tpu.v2alpha1.QueuedResource;
23+
import com.google.cloud.tpu.v2alpha1.TpuClient;
24+
import java.io.IOException;
25+
import java.util.concurrent.ExecutionException;
26+
import java.util.concurrent.TimeUnit;
27+
import java.util.concurrent.TimeoutException;
28+
29+
public class CreateQueuedResource {
30+
public static void main(String[] args)
31+
throws IOException, ExecutionException, InterruptedException, TimeoutException {
32+
// TODO(developer): Replace these variables before running the sample.
33+
// Project ID or project number of the Google Cloud project you want to create a node.
34+
String projectId = "YOUR_PROJECT_ID";
35+
// The zone in which to create the TPU.
36+
// For more information about supported TPU types for specific zones,
37+
// see https://cloud.google.com/tpu/docs/regions-zones
38+
String zone = "us-central1-f";
39+
// The name for your TPU.
40+
String nodeName = "YOUR_NODE_ID";
41+
// The accelerator type that specifies the version and size of the Cloud TPU you want to create.
42+
// For more information about supported accelerator types for each TPU version,
43+
// see https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#versions.
44+
String tpuType = "v2-8";
45+
// Software version that specifies the version of the TPU runtime to install.
46+
// For more information see https://cloud.google.com/tpu/docs/runtimes
47+
String tpuSoftwareVersion = "tpu-vm-tf-2.14.1";
48+
// The name for your Queued Resource.
49+
String queuedResourceId = "QUEUED_RESOURCE_ID";
50+
51+
createQueuedResource(
52+
projectId, zone, queuedResourceId, nodeName, tpuType, tpuSoftwareVersion);
53+
}
54+
55+
// Creates a Queued Resource
56+
public static QueuedResource createQueuedResource(String projectId, String zone,
57+
String queuedResourceId, String nodeName, String tpuType, String tpuSoftwareVersion)
58+
throws IOException, ExecutionException, InterruptedException, TimeoutException {
59+
String resource = String.format("projects/%s/locations/%s/queuedResources/%s",
60+
projectId, zone, queuedResourceId);
61+
// Initialize client that will be used to send requests. This client only needs to be created
62+
// once, and can be reused for multiple requests.
63+
try (TpuClient tpuClient = TpuClient.create()) {
64+
String parent = String.format("projects/%s/locations/%s", projectId, zone);
65+
Node node =
66+
Node.newBuilder()
67+
.setName(nodeName)
68+
.setAcceleratorType(tpuType)
69+
.setRuntimeVersion(tpuSoftwareVersion)
70+
.setQueuedResource(resource)
71+
.build();
72+
73+
QueuedResource queuedResource =
74+
QueuedResource.newBuilder()
75+
.setName(queuedResourceId)
76+
.setTpu(
77+
QueuedResource.Tpu.newBuilder()
78+
.addNodeSpec(
79+
QueuedResource.Tpu.NodeSpec.newBuilder()
80+
.setParent(parent)
81+
.setNode(node)
82+
.setNodeId(nodeName)
83+
.build())
84+
.build())
85+
.build();
86+
87+
CreateQueuedResourceRequest request =
88+
CreateQueuedResourceRequest.newBuilder()
89+
.setParent(parent)
90+
.setQueuedResourceId(queuedResourceId)
91+
.setQueuedResource(queuedResource)
92+
.build();
93+
94+
return tpuClient.createQueuedResourceAsync(request).get(1, TimeUnit.MINUTES);
95+
}
96+
}
97+
}
98+
//[END tpu_queued_resources_create]

‎tpu/src/main/java/tpu/DeleteQueuedResource.java

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
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_delete]
20+
import com.google.cloud.tpu.v2alpha1.DeleteQueuedResourceRequest;
21+
import com.google.cloud.tpu.v2alpha1.TpuClient;
22+
import java.io.IOException;
23+
import java.util.concurrent.ExecutionException;
24+
25+
public class DeleteQueuedResource {
26+
public static void main(String[] args)
27+
throws IOException, ExecutionException, InterruptedException {
28+
// TODO(developer): Replace these variables before running the sample.
29+
// Project ID or project number of the Google Cloud project.
30+
String projectId = "YOUR_PROJECT_ID";
31+
// The zone in which the TPU was created.
32+
String zone = "us-central1-f";
33+
// The name for your Queued Resource.
34+
String queuedResourceId = "QUEUED_RESOURCE_ID";
35+
36+
deleteQueuedResource(projectId, zone, queuedResourceId);
37+
}
38+
39+
// Deletes a Queued Resource asynchronously.
40+
public static void deleteQueuedResource(String projectId, String zone, String queuedResourceId)
41+
throws ExecutionException, InterruptedException, IOException {
42+
String name = String.format("projects/%s/locations/%s/queuedResources/%s",
43+
projectId, zone, queuedResourceId);
44+
// Initialize client that will be used to send requests. This client only needs to be created
45+
// once, and can be reused for multiple requests.
46+
try (TpuClient tpuClient = TpuClient.create()) {
47+
// Before deleting the queued resource it is required to delete the TPU VM.
48+
// For more information about deleting TPU
49+
// see https://cloud.google.com/tpu/docs/managing-tpus-tpu-vm
50+
51+
DeleteQueuedResourceRequest request =
52+
DeleteQueuedResourceRequest.newBuilder().setName(name).build();
53+
54+
tpuClient.deleteQueuedResourceAsync(request).get();
55+
}
56+
}
57+
}
58+
//[END tpu_queued_resources_delete]

‎tpu/src/main/java/tpu/GetQueuedResource.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ public static void main(String[] args) throws IOException {
2828
// Project ID or project number of the Google Cloud project.
2929
String projectId = "YOUR_PROJECT_ID";
3030
// The zone in which the TPU was created.
31-
String zone = "europe-west4-a";
31+
String zone = "us-central1-f";
3232
// The name for your Queued Resource.
3333
String queuedResourceId = "QUEUED_RESOURCE_ID";
3434

@@ -50,4 +50,4 @@ public static QueuedResource getQueuedResource(
5050
}
5151
}
5252
}
53-
//[END tpu_queued_resources_get]
53+
//[END tpu_queued_resources_get]

‎tpu/src/test/java/tpu/QueuedResourceIT.java

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package tpu;
1818

1919
import static org.junit.Assert.assertEquals;
20+
import static org.mockito.ArgumentMatchers.anyLong;
2021
import static org.mockito.Mockito.any;
2122
import static org.mockito.Mockito.mock;
2223
import static org.mockito.Mockito.mockStatic;
@@ -33,14 +34,15 @@
3334
import com.google.cloud.tpu.v2alpha1.TpuSettings;
3435
import java.io.IOException;
3536
import java.util.concurrent.ExecutionException;
37+
import java.util.concurrent.TimeUnit;
3638
import org.junit.jupiter.api.Test;
3739
import org.junit.jupiter.api.Timeout;
3840
import org.junit.runner.RunWith;
3941
import org.junit.runners.JUnit4;
4042
import org.mockito.MockedStatic;
4143

4244
@RunWith(JUnit4.class)
43-
@Timeout(value = 10)
45+
@Timeout(value = 2, unit = TimeUnit.MINUTES)
4446
public class QueuedResourceIT {
4547
private static final String PROJECT_ID = "project-id";
4648
private static final String ZONE = "europe-west4-a";
@@ -50,6 +52,30 @@ public class QueuedResourceIT {
5052
private static final String QUEUED_RESOURCE_NAME = "queued-resource";
5153
private static final String NETWORK_NAME = "default";
5254

55+
@Test
56+
public void testCreateQueuedResource() throws Exception {
57+
try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
58+
QueuedResource mockQueuedResource = mock(QueuedResource.class);
59+
TpuClient mockTpuClient = mock(TpuClient.class);
60+
OperationFuture mockFuture = mock(OperationFuture.class);
61+
62+
mockedTpuClient.when(TpuClient::create).thenReturn(mockTpuClient);
63+
when(mockTpuClient.createQueuedResourceAsync(any(CreateQueuedResourceRequest.class)))
64+
.thenReturn(mockFuture);
65+
when(mockFuture.get(anyLong(), any(TimeUnit.class))).thenReturn(mockQueuedResource);
66+
67+
QueuedResource returnedQueuedResource =
68+
CreateQueuedResource.createQueuedResource(
69+
PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME, NODE_NAME,
70+
TPU_TYPE, TPU_SOFTWARE_VERSION);
71+
72+
verify(mockTpuClient, times(1))
73+
.createQueuedResourceAsync(any(CreateQueuedResourceRequest.class));
74+
verify(mockFuture, times(1)).get(anyLong(), any(TimeUnit.class));
75+
assertEquals(returnedQueuedResource, mockQueuedResource);
76+
}
77+
}
78+
5379
@Test
5480
public void testCreateQueuedResourceWithSpecifiedNetwork() throws Exception {
5581
try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
@@ -113,6 +139,25 @@ public void testDeleteForceQueuedResource()
113139
}
114140
}
115141

142+
@Test
143+
public void testDeleteQueuedResource()
144+
throws IOException, ExecutionException, InterruptedException {
145+
try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
146+
TpuClient mockTpuClient = mock(TpuClient.class);
147+
OperationFuture mockFuture = mock(OperationFuture.class);
148+
149+
mockedTpuClient.when(TpuClient::create).thenReturn(mockTpuClient);
150+
when(mockTpuClient.deleteQueuedResourceAsync(any(DeleteQueuedResourceRequest.class)))
151+
.thenReturn(mockFuture);
152+
when(mockFuture.get()).thenReturn(null);
153+
154+
DeleteQueuedResource.deleteQueuedResource(PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME);
155+
156+
verify(mockTpuClient, times(1))
157+
.deleteQueuedResourceAsync(any(DeleteQueuedResourceRequest.class));
158+
}
159+
}
160+
116161
@Test
117162
public void testCreateQueuedResourceWithStartupScript() throws Exception {
118163
try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {

0 commit comments

Comments
 (0)