Created cleanup method for tpu

Author: TetyanaYahodska

compute/cloud-client/src/test/java/compute/Util.java

@@ -34,15 +34,21 @@
 import com.google.cloud.compute.v1.SnapshotsClient;
 import com.google.cloud.compute.v1.StoragePool;
 import com.google.cloud.compute.v1.StoragePoolsClient;
+import com.google.cloud.tpu.v2.Node;
+import com.google.cloud.tpu.v2.TpuClient;
+import com.google.protobuf.Timestamp;
 import compute.deleteprotection.SetDeleteProtection;
 import compute.disks.DeleteDisk;
 import compute.disks.DeleteSnapshot;
 import compute.reservation.DeleteReservation;
+import compute.tpu.DeleteTpuVm;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.security.SecureRandom;
 import java.time.Instant;
 import java.time.OffsetDateTime;
+import java.time.ZoneOffset;
+import java.time.format.DateTimeFormatter;
 import java.time.temporal.ChronoUnit;
 import java.util.Base64;
 import java.util.Random;
@@ -57,7 +63,7 @@ public abstract class Util {
   // resources
   // and delete the listed resources based on the timestamp.
 
-  private static final int DELETION_THRESHOLD_TIME_MINUTES = 45;
+  private static final int DELETION_THRESHOLD_TIME_MINUTES = 3;
   // comma separate list of zone names
   private static final String TEST_ZONES_NAME = "JAVA_DOCS_COMPUTE_TEST_ZONES";
   private static final String DEFAULT_ZONES = "us-central1-a,us-west1-a,asia-south1-a";
@@ -231,14 +237,31 @@ && isCreatedBeforeThresholdTime(snapshot.getCreationTimestamp())) {
     }
   }
 
+  // Delete tpu which starts with the given prefixToDelete and
+  // has creation timestamp >24 hours.
+  public static void cleanUpExistingTpu(String prefixToDelete, String projectId, String zone)
+      throws IOException, ExecutionException, InterruptedException {
+    try (TpuClient tpuClient = TpuClient.create()) {
+      String parent = String.format("projects/%s/locations/%s", projectId, zone);
+      for (Node node : tpuClient.listNodes(parent).iterateAll()) {
+        String creationTime = formatTimestamp(node.getCreateTime());
+        String name = node.getName().substring(node.getName().lastIndexOf("/") + 1);
+        if (containPrefixToDeleteAndZone(node, prefixToDelete, zone)
+            && isCreatedBeforeThresholdTime(creationTime)) {
+          DeleteTpuVm.deleteTpuVm(projectId, zone, name);
+        }
+      }
+    }
+  }
+
   // Delete storagePools which starts with the given prefixToDelete and
   // has creation timestamp >24 hours.
   public static void cleanUpExistingStoragePool(
       String prefixToDelete, String projectId, String zone)
       throws IOException, ExecutionException, InterruptedException, TimeoutException {
     try (StoragePoolsClient storagePoolsClient = StoragePoolsClient.create()) {
       for (StoragePool storagePool : storagePoolsClient.list(projectId, zone).iterateAll()) {
-        if (containPrefixToDeleteAndZone(projectId, prefixToDelete, zone)
+        if (containPrefixToDeleteAndZone(storagePool, prefixToDelete, zone)
             && isCreatedBeforeThresholdTime(storagePool.getCreationTimestamp())) {
           deleteStoragePool(projectId, zone, storagePool.getName());
         }
@@ -291,6 +314,10 @@ public static boolean containPrefixToDeleteAndZone(
         containPrefixAndZone = ((StoragePool) resource).getName().contains(prefixToDelete)
             && ((StoragePool) resource).getZone().contains(zone);
       }
+      if (resource instanceof Node) {
+        containPrefixAndZone = ((Node) resource).getName().contains(prefixToDelete)
+            && ((Node) resource).getName().split("/")[3].contains(zone);
+      }
     } catch (NullPointerException e) {
       System.out.println("Resource not found, skipping deletion:");
     }
@@ -312,4 +339,12 @@ public static boolean containPrefixToDelete(
     }
     return containPrefixToDelete;
   }
+
+  private static String formatTimestamp(Timestamp timestamp) {
+    DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSSXXX");
+    OffsetDateTime offsetDateTime = OffsetDateTime.ofInstant(
+        Instant.ofEpochSecond(timestamp.getSeconds(), timestamp.getNanos()),
+        ZoneOffset.UTC);
+    return formatter.format(offsetDateTime);
+  }
 }

compute/cloud-client/src/test/java/compute/tpu/TpuVmTest.java renamed to compute/cloud-client/src/test/java/compute/tpu/TpuVmIT.java

@@ -26,27 +26,31 @@
 import com.google.api.gax.rpc.NotFoundException;
 import com.google.cloud.tpu.v2.Node;
 import com.google.cloud.tpu.v2.TpuClient;
+import compute.Util;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.PrintStream;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeoutException;
 import org.junit.Assert;
 import org.junit.FixMethodOrder;
-import org.junit.Test;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.junit.runners.MethodSorters;
 
 @RunWith(JUnit4.class)
 @FixMethodOrder(MethodSorters.NAME_ASCENDING)
-public class TpuVmTest {
+public class TpuVmIT {
   private static final String PROJECT_ID = System.getenv("GOOGLE_CLOUD_PROJECT");
   private static final String ZONE = "europe-west4-a";
-  private static final String TPU_VM_NAME = "test-tpu-" + UUID.randomUUID();
+  static String javaVersion = System.getProperty("java.version").substring(0, 2);
+  private static final String TPU_VM_NAME = "test-tpu-" + javaVersion + "-"
+      + UUID.randomUUID().toString().substring(0, 8);
   private static final String ACCELERATOR_TYPE = "v2-8";
   private static final String VERSION = "tpu-vm-tf-2.14.1";
   private static final String TPU_VM_PATH_NAME =
@@ -58,15 +62,20 @@ public static void requireEnvVar(String envVarName) {
   }
 
   @BeforeAll
-  public static void setUp() throws IOException, ExecutionException, InterruptedException {
+  public static void setUp()
+      throws IOException, ExecutionException, InterruptedException, TimeoutException {
     requireEnvVar("GOOGLE_APPLICATION_CREDENTIALS");
     requireEnvVar("GOOGLE_CLOUD_PROJECT");
+
+    // Cleanup existing stale resources.
+    Util.cleanUpExistingTpu("test-tpu-" + javaVersion, PROJECT_ID, ZONE);
   }
 
   @AfterAll
   public static void cleanup() throws Exception {
     DeleteTpuVm.deleteTpuVm(PROJECT_ID, ZONE, TPU_VM_NAME);
     assertNull(GetTpuVm.getTpuVm(PROJECT_ID, ZONE, TPU_VM_NAME));
+
     // Test that reservations are deleted
     Assertions.assertThrows(
         NotFoundException.class,