Skip to content

Commit d29a6b5

Browse files
TetyanaYahodskaTetyanaYahodska
committed
Added comment, created Util class
1 parent 2253b54 commit d29a6b5

File tree

3 files changed

+78
-50
lines changed

3 files changed

+78
-50
lines changed

‎tpu/src/main/java/tpu/DeleteTpuVm.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@ public static void main(String[] args)
4848
// Deletes a TPU VM with the specified name in the given project and zone.
4949
public static void deleteTpuVm(String projectId, String zone, String nodeName)
5050
throws IOException, ExecutionException, InterruptedException {
51+
// With these settings the client library handles the Operation's polling mechanism
52+
// and prevent CancellationException error
5153
TpuSettings.Builder clientSettings =
5254
TpuSettings.newBuilder();
5355
clientSettings

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

Lines changed: 1 addition & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,9 @@
2525
import com.google.api.gax.rpc.NotFoundException;
2626
import com.google.cloud.tpu.v2.Node;
2727
import com.google.cloud.tpu.v2.TpuClient;
28-
import com.google.protobuf.Timestamp;
2928
import java.io.ByteArrayOutputStream;
3029
import java.io.IOException;
3130
import java.io.PrintStream;
32-
import java.time.Instant;
33-
import java.time.OffsetDateTime;
34-
import java.time.ZoneOffset;
35-
import java.time.format.DateTimeFormatter;
36-
import java.time.temporal.ChronoUnit;
3731
import java.util.UUID;
3832
import java.util.concurrent.ExecutionException;
3933
import java.util.concurrent.TimeUnit;
@@ -75,13 +69,12 @@ public static void setUp()
7569
requireEnvVar("GOOGLE_CLOUD_PROJECT");
7670

7771
// Cleanup existing stale resources.
78-
cleanUpExistingTpu("test-tpu-" + javaVersion, PROJECT_ID, ZONE);
72+
Util.cleanUpExistingTpu("test-tpu-" + javaVersion, PROJECT_ID, ZONE);
7973
}
8074

8175
@AfterAll
8276
public static void cleanup() throws Exception {
8377
DeleteTpuVm.deleteTpuVm(PROJECT_ID, ZONE, TPU_VM_NAME);
84-
TimeUnit.MINUTES.sleep(5);
8578

8679
// Test that TPUs is deleted
8780
Assertions.assertThrows(
@@ -96,7 +89,6 @@ public void testCreateTpuVm() throws IOException, ExecutionException, Interrupte
9689
ByteArrayOutputStream stdOut = new ByteArrayOutputStream();
9790
System.setOut(new PrintStream(stdOut));
9891
CreateTpuVm.createTpuVm(PROJECT_ID, ZONE, TPU_VM_NAME, ACCELERATOR_TYPE, VERSION);
99-
TimeUnit.MINUTES.sleep(3);
10092

10193
assertThat(stdOut.toString()).contains("TPU VM created: " + TPU_VM_PATH_NAME);
10294
stdOut.close();
@@ -140,45 +132,4 @@ public void testStartTpuVm() throws IOException, ExecutionException, Interrupted
140132

141133
assertThat(node.getState()).isEqualTo(READY);
142134
}
143-
144-
public static void cleanUpExistingTpu(String prefixToDelete, String projectId, String zone)
145-
throws IOException, ExecutionException, InterruptedException {
146-
try (TpuClient tpuClient = TpuClient.create()) {
147-
String parent = String.format("projects/%s/locations/%s", projectId, zone);
148-
for (Node node : tpuClient.listNodes(parent).iterateAll()) {
149-
String creationTime = formatTimestamp(node.getCreateTime());
150-
String name = node.getName().substring(node.getName().lastIndexOf("/") + 1);
151-
if (containPrefixToDeleteAndZone(node, prefixToDelete, zone)
152-
&& isCreatedBeforeThresholdTime(creationTime)) {
153-
DeleteTpuVm.deleteTpuVm(projectId, zone, name);
154-
}
155-
}
156-
}
157-
}
158-
159-
public static boolean containPrefixToDeleteAndZone(
160-
Node node, String prefixToDelete, String zone) {
161-
boolean containPrefixAndZone = false;
162-
try {
163-
containPrefixAndZone = node.getName().contains(prefixToDelete)
164-
&& node.getName().split("/")[3].contains(zone);
165-
166-
} catch (NullPointerException e) {
167-
System.out.println("Resource not found, skipping deletion:");
168-
}
169-
return containPrefixAndZone;
170-
}
171-
172-
public static boolean isCreatedBeforeThresholdTime(String timestamp) {
173-
return OffsetDateTime.parse(timestamp).toInstant()
174-
.isBefore(Instant.now().minus(30, ChronoUnit.MINUTES));
175-
}
176-
177-
private static String formatTimestamp(Timestamp timestamp) {
178-
DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSSXXX");
179-
OffsetDateTime offsetDateTime = OffsetDateTime.ofInstant(
180-
Instant.ofEpochSecond(timestamp.getSeconds(), timestamp.getNanos()),
181-
ZoneOffset.UTC);
182-
return formatter.format(offsetDateTime);
183-
}
184135
}

‎tpu/src/test/java/tpu/Util.java

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
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 com.google.cloud.tpu.v2.Node;
20+
import com.google.cloud.tpu.v2.TpuClient;
21+
import com.google.protobuf.Timestamp;
22+
import java.io.IOException;
23+
import java.time.Instant;
24+
import java.time.OffsetDateTime;
25+
import java.time.ZoneOffset;
26+
import java.time.format.DateTimeFormatter;
27+
import java.time.temporal.ChronoUnit;
28+
import java.util.concurrent.ExecutionException;
29+
30+
public class Util {
31+
private static final int DELETION_THRESHOLD_TIME_MINUTES = 30;
32+
33+
// Delete TPU VMs which starts with the given prefixToDelete and
34+
// has creation timestamp >30 minutes.
35+
public static void cleanUpExistingTpu(String prefixToDelete, String projectId, String zone)
36+
throws IOException, ExecutionException, InterruptedException {
37+
try (TpuClient tpuClient = TpuClient.create()) {
38+
String parent = String.format("projects/%s/locations/%s", projectId, zone);
39+
for (Node node : tpuClient.listNodes(parent).iterateAll()) {
40+
String creationTime = formatTimestamp(node.getCreateTime());
41+
String name = node.getName().substring(node.getName().lastIndexOf("/") + 1);
42+
if (containPrefixToDeleteAndZone(node, prefixToDelete, zone)
43+
&& isCreatedBeforeThresholdTime(creationTime)) {
44+
DeleteTpuVm.deleteTpuVm(projectId, zone, name);
45+
}
46+
}
47+
}
48+
}
49+
50+
public static boolean containPrefixToDeleteAndZone(
51+
Node node, String prefixToDelete, String zone) {
52+
boolean containPrefixAndZone = false;
53+
try {
54+
containPrefixAndZone = node.getName().contains(prefixToDelete)
55+
&& node.getName().split("/")[3].contains(zone);
56+
57+
} catch (NullPointerException e) {
58+
System.out.println("Resource not found, skipping deletion:");
59+
}
60+
return containPrefixAndZone;
61+
}
62+
63+
public static boolean isCreatedBeforeThresholdTime(String timestamp) {
64+
return OffsetDateTime.parse(timestamp).toInstant()
65+
.isBefore(Instant.now().minus(DELETION_THRESHOLD_TIME_MINUTES, ChronoUnit.MINUTES));
66+
}
67+
68+
private static String formatTimestamp(Timestamp timestamp) {
69+
DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSSXXX");
70+
OffsetDateTime offsetDateTime = OffsetDateTime.ofInstant(
71+
Instant.ofEpochSecond(timestamp.getSeconds(), timestamp.getNanos()),
72+
ZoneOffset.UTC);
73+
return formatter.format(offsetDateTime);
74+
}
75+
}

0 commit comments

Comments
 (0)