Skip to content

Commit 5a7af75

Browse files
Resolved merge conflict
2 parents 99d4a4a + 772c39c commit 5a7af75

File tree

4 files changed

+66
-75
lines changed

4 files changed

+66
-75
lines changed

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

Lines changed: 10 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -17,26 +17,25 @@
1717
package tpu;
1818

1919
//[START tpu_queued_resources_create]
20-
import com.google.api.gax.retrying.RetrySettings;
2120
import com.google.cloud.tpu.v2alpha1.CreateQueuedResourceRequest;
2221
import com.google.cloud.tpu.v2alpha1.Node;
2322
import com.google.cloud.tpu.v2alpha1.QueuedResource;
2423
import com.google.cloud.tpu.v2alpha1.TpuClient;
25-
import com.google.cloud.tpu.v2alpha1.TpuSettings;
2624
import java.io.IOException;
2725
import java.util.concurrent.ExecutionException;
28-
import org.threeten.bp.Duration;
26+
import java.util.concurrent.TimeUnit;
27+
import java.util.concurrent.TimeoutException;
2928

3029
public class CreateQueuedResource {
3130
public static void main(String[] args)
32-
throws IOException, ExecutionException, InterruptedException {
31+
throws IOException, ExecutionException, InterruptedException, TimeoutException {
3332
// TODO(developer): Replace these variables before running the sample.
3433
// Project ID or project number of the Google Cloud project you want to create a node.
3534
String projectId = "YOUR_PROJECT_ID";
3635
// The zone in which to create the TPU.
3736
// For more information about supported TPU types for specific zones,
3837
// see https://cloud.google.com/tpu/docs/regions-zones
39-
String zone = "europe-west4-a";
38+
String zone = "us-central1-f";
4039
// The name for your TPU.
4140
String nodeName = "YOUR_NODE_ID";
4241
// The accelerator type that specifies the version and size of the Cloud TPU you want to create.
@@ -56,35 +55,19 @@ public static void main(String[] args)
5655
// Creates a Queued Resource
5756
public static QueuedResource createQueuedResource(String projectId, String zone,
5857
String queuedResourceId, String nodeName, String tpuType, String tpuSoftwareVersion)
59-
throws IOException, ExecutionException, InterruptedException {
60-
// With these settings the client library handles the Operation's polling mechanism
61-
// and prevent CancellationException error
62-
TpuSettings.Builder clientSettings =
63-
TpuSettings.newBuilder();
64-
clientSettings
65-
.createQueuedResourceSettings()
66-
.setRetrySettings(
67-
RetrySettings.newBuilder()
68-
.setInitialRetryDelay(Duration.ofMillis(5000L))
69-
.setRetryDelayMultiplier(2.0)
70-
.setInitialRpcTimeout(Duration.ZERO)
71-
.setRpcTimeoutMultiplier(1.0)
72-
.setMaxRetryDelay(Duration.ofMillis(45000L))
73-
.setTotalTimeout(Duration.ofHours(24L))
74-
.build());
58+
throws IOException, ExecutionException, InterruptedException, TimeoutException {
59+
String resource = String.format("projects/%s/locations/%s/queuedResources/%s",
60+
projectId, zone, queuedResourceId);
7561
// Initialize client that will be used to send requests. This client only needs to be created
7662
// once, and can be reused for multiple requests.
77-
try (TpuClient tpuClient = TpuClient.create(clientSettings.build())) {
63+
try (TpuClient tpuClient = TpuClient.create()) {
7864
String parent = String.format("projects/%s/locations/%s", projectId, zone);
7965
Node node =
8066
Node.newBuilder()
8167
.setName(nodeName)
8268
.setAcceleratorType(tpuType)
8369
.setRuntimeVersion(tpuSoftwareVersion)
84-
.setQueuedResource(
85-
String.format(
86-
"projects/%s/locations/%s/queuedResources/%s",
87-
projectId, zone, queuedResourceId))
70+
.setQueuedResource(resource)
8871
.build();
8972

9073
QueuedResource queuedResource =
@@ -99,9 +82,6 @@ public static QueuedResource createQueuedResource(String projectId, String zone,
9982
.setNodeId(nodeName)
10083
.build())
10184
.build())
102-
// You can request a queued resource using a reservation by specifying it in code
103-
//.setReservationName(
104-
// "projects/YOUR_PROJECT_ID/locations/YOUR_ZONE/reservations/YOUR_RESERVATION_NAME")
10585
.build();
10686

10787
CreateQueuedResourceRequest request =
@@ -111,11 +91,7 @@ public static QueuedResource createQueuedResource(String projectId, String zone,
11191
.setQueuedResource(queuedResource)
11292
.build();
11393

114-
QueuedResource response = tpuClient.createQueuedResourceAsync(request).get();
115-
// You can wait until TPU Node is READY,
116-
// and check its status using getTpuVm() from "tpu_vm_get" sample.
117-
System.out.printf("Queued Resource created: %s\n", response.getName());
118-
return response;
94+
return tpuClient.createQueuedResourceAsync(request).get(1, TimeUnit.MINUTES);
11995
}
12096
}
12197
}

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

Lines changed: 9 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -17,72 +17,42 @@
1717
package tpu;
1818

1919
//[START tpu_queued_resources_delete]
20-
import com.google.api.gax.retrying.RetrySettings;
21-
import com.google.api.gax.rpc.UnknownException;
2220
import com.google.cloud.tpu.v2alpha1.DeleteQueuedResourceRequest;
23-
import com.google.cloud.tpu.v2alpha1.GetQueuedResourceRequest;
24-
import com.google.cloud.tpu.v2alpha1.QueuedResource;
2521
import com.google.cloud.tpu.v2alpha1.TpuClient;
26-
import com.google.cloud.tpu.v2alpha1.TpuSettings;
2722
import java.io.IOException;
2823
import java.util.concurrent.ExecutionException;
29-
import java.util.concurrent.TimeUnit;
30-
import org.threeten.bp.Duration;
3124

3225
public class DeleteQueuedResource {
33-
public static void main(String[] args) {
26+
public static void main(String[] args)
27+
throws IOException, ExecutionException, InterruptedException {
3428
// TODO(developer): Replace these variables before running the sample.
3529
// Project ID or project number of the Google Cloud project.
3630
String projectId = "YOUR_PROJECT_ID";
3731
// The zone in which the TPU was created.
38-
String zone = "europe-west4-a";
32+
String zone = "us-central1-f";
3933
// The name for your Queued Resource.
4034
String queuedResourceId = "QUEUED_RESOURCE_ID";
4135

4236
deleteQueuedResource(projectId, zone, queuedResourceId);
4337
}
4438

4539
// Deletes a Queued Resource asynchronously.
46-
public static void deleteQueuedResource(String projectId, String zone, String queuedResourceId) {
40+
public static void deleteQueuedResource(String projectId, String zone, String queuedResourceId)
41+
throws ExecutionException, InterruptedException, IOException {
4742
String name = String.format("projects/%s/locations/%s/queuedResources/%s",
4843
projectId, zone, queuedResourceId);
49-
// With these settings the client library handles the Operation's polling mechanism
50-
// and prevent CancellationException error
51-
TpuSettings.Builder clientSettings =
52-
TpuSettings.newBuilder();
53-
clientSettings
54-
.deleteQueuedResourceSettings()
55-
.setRetrySettings(
56-
RetrySettings.newBuilder()
57-
.setInitialRetryDelay(Duration.ofMillis(5000L))
58-
.setRetryDelayMultiplier(2.0)
59-
.setInitialRpcTimeout(Duration.ZERO)
60-
.setRpcTimeoutMultiplier(1.0)
61-
.setMaxRetryDelay(Duration.ofMillis(45000L))
62-
.setTotalTimeout(Duration.ofHours(24L))
63-
.build());
6444
// Initialize client that will be used to send requests. This client only needs to be created
6545
// once, and can be reused for multiple requests.
66-
try (TpuClient tpuClient = TpuClient.create(clientSettings.build())) {
67-
// Retrive node name
68-
GetQueuedResourceRequest getRequest =
69-
GetQueuedResourceRequest.newBuilder().setName(name).build();
70-
QueuedResource queuedResource = tpuClient.getQueuedResource(getRequest);
71-
String nodeName = queuedResource.getTpu().getNodeSpec(0).getNode().getName();
46+
try (TpuClient tpuClient = TpuClient.create()) {
7247
// Before deleting the queued resource it is required to delete the TPU VM.
73-
DeleteTpuVm.deleteTpuVm(projectId, zone, nodeName);
74-
// Wait until TpuVm is deleted
75-
TimeUnit.MINUTES.sleep(3);
48+
// For more information about deleting TPU
49+
// see https://cloud.google.com/tpu/docs/managing-tpus-tpu-vm
7650

7751
DeleteQueuedResourceRequest request =
78-
DeleteQueuedResourceRequest.newBuilder().setName(name).build();
52+
DeleteQueuedResourceRequest.newBuilder().setName(name).build();
7953

8054
tpuClient.deleteQueuedResourceAsync(request).get();
81-
82-
} catch (UnknownException | InterruptedException | ExecutionException | IOException e) {
83-
System.out.println(e.getMessage());
8455
}
85-
System.out.printf("Deleted Queued Resource: %s\n", name);
8656
}
8757
}
8858
//[END tpu_queued_resources_delete]

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

Lines changed: 1 addition & 1 deletion
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

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

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import static com.google.common.truth.Truth.assertThat;
2020
import static org.junit.Assert.assertEquals;
21+
import static org.mockito.ArgumentMatchers.anyLong;
2122
import static org.mockito.Mockito.any;
2223
import static org.mockito.Mockito.mock;
2324
import static org.mockito.Mockito.mockStatic;
@@ -37,14 +38,15 @@
3738
import java.util.Arrays;
3839
import java.util.List;
3940
import java.util.concurrent.ExecutionException;
41+
import java.util.concurrent.TimeUnit;
4042
import org.junit.jupiter.api.Test;
4143
import org.junit.jupiter.api.Timeout;
4244
import org.junit.runner.RunWith;
4345
import org.junit.runners.JUnit4;
4446
import org.mockito.MockedStatic;
4547

4648
@RunWith(JUnit4.class)
47-
@Timeout(value = 10)
49+
@Timeout(value = 2, unit = TimeUnit.MINUTES)
4850
public class QueuedResourceIT {
4951
private static final String PROJECT_ID = "project-id";
5052
private static final String ZONE = "europe-west4-a";
@@ -54,6 +56,30 @@ public class QueuedResourceIT {
5456
private static final String QUEUED_RESOURCE_NAME = "queued-resource";
5557
private static final String NETWORK_NAME = "default";
5658

59+
@Test
60+
public void testCreateQueuedResource() throws Exception {
61+
try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
62+
QueuedResource mockQueuedResource = mock(QueuedResource.class);
63+
TpuClient mockTpuClient = mock(TpuClient.class);
64+
OperationFuture mockFuture = mock(OperationFuture.class);
65+
66+
mockedTpuClient.when(TpuClient::create).thenReturn(mockTpuClient);
67+
when(mockTpuClient.createQueuedResourceAsync(any(CreateQueuedResourceRequest.class)))
68+
.thenReturn(mockFuture);
69+
when(mockFuture.get(anyLong(), any(TimeUnit.class))).thenReturn(mockQueuedResource);
70+
71+
QueuedResource returnedQueuedResource =
72+
CreateQueuedResource.createQueuedResource(
73+
PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME, NODE_NAME,
74+
TPU_TYPE, TPU_SOFTWARE_VERSION);
75+
76+
verify(mockTpuClient, times(1))
77+
.createQueuedResourceAsync(any(CreateQueuedResourceRequest.class));
78+
verify(mockFuture, times(1)).get(anyLong(), any(TimeUnit.class));
79+
assertEquals(returnedQueuedResource, mockQueuedResource);
80+
}
81+
}
82+
5783
@Test
5884
public void testCreateQueuedResourceWithSpecifiedNetwork() throws Exception {
5985
try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
@@ -144,6 +170,25 @@ public void testDeleteForceQueuedResource()
144170
}
145171
}
146172

173+
@Test
174+
public void testDeleteQueuedResource()
175+
throws IOException, ExecutionException, InterruptedException {
176+
try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
177+
TpuClient mockTpuClient = mock(TpuClient.class);
178+
OperationFuture mockFuture = mock(OperationFuture.class);
179+
180+
mockedTpuClient.when(TpuClient::create).thenReturn(mockTpuClient);
181+
when(mockTpuClient.deleteQueuedResourceAsync(any(DeleteQueuedResourceRequest.class)))
182+
.thenReturn(mockFuture);
183+
when(mockFuture.get()).thenReturn(null);
184+
185+
DeleteQueuedResource.deleteQueuedResource(PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME);
186+
187+
verify(mockTpuClient, times(1))
188+
.deleteQueuedResourceAsync(any(DeleteQueuedResourceRequest.class));
189+
}
190+
}
191+
147192
@Test
148193
public void testCreateQueuedResourceWithStartupScript() throws Exception {
149194
try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {

0 commit comments

Comments
 (0)