feat(tpu): add tpu queued resources create/get/delete/forceDelete samples

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

@@ -0,0 +1,121 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package tpu;
+
+//[START tpu_queued_resources_create]
+import com.google.api.gax.retrying.RetrySettings;
+import com.google.cloud.tpu.v2alpha1.CreateQueuedResourceRequest;
+import com.google.cloud.tpu.v2alpha1.Node;
+import com.google.cloud.tpu.v2alpha1.QueuedResource;
+import com.google.cloud.tpu.v2alpha1.TpuClient;
+import com.google.cloud.tpu.v2alpha1.TpuSettings;
+import java.io.IOException;
+import java.util.concurrent.ExecutionException;
+import org.threeten.bp.Duration;
+
+public class CreateQueuedResource {
+  public static void main(String[] args)
+      throws IOException, ExecutionException, InterruptedException {
+    // TODO(developer): Replace these variables before running the sample.
+    // Project ID or project number of the Google Cloud project you want to create a node.
+    String projectId = "YOUR_PROJECT_ID";
+    // The zone in which to create the TPU.
+    // For more information about supported TPU types for specific zones,
+    // see https://cloud.google.com/tpu/docs/regions-zones
+    String zone = "europe-west4-a";
+    // The name for your TPU.
+    String nodeName = "YOUR_NODE_ID";
+    // The accelerator type that specifies the version and size of the Cloud TPU you want to create.
+    // For more information about supported accelerator types for each TPU version,
+    // see https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#versions.
+    String tpuType = "v2-8";
+    // Software version that specifies the version of the TPU runtime to install.
+    // For more information see https://cloud.google.com/tpu/docs/runtimes
+    String tpuSoftwareVersion = "tpu-vm-tf-2.14.1";
+    // The name for your Queued Resource.
+    String queuedResourceId = "QUEUED_RESOURCE_ID";
+
+    createQueuedResource(
+        projectId, zone, queuedResourceId, nodeName, tpuType, tpuSoftwareVersion);
+  }
+
+  // Creates a Queued Resource
+  public static QueuedResource createQueuedResource(String projectId, String zone,
+      String queuedResourceId, String nodeName, String tpuType, String tpuSoftwareVersion)
+      throws IOException, ExecutionException, InterruptedException {
+    // With these settings the client library handles the Operation's polling mechanism
+    // and prevent CancellationException error
+    TpuSettings.Builder clientSettings =
+        TpuSettings.newBuilder();
+    clientSettings
+        .createQueuedResourceSettings()
+        .setRetrySettings(
+            RetrySettings.newBuilder()
+                .setInitialRetryDelay(Duration.ofMillis(5000L))
+                .setRetryDelayMultiplier(2.0)
+                .setInitialRpcTimeout(Duration.ZERO)
+                .setRpcTimeoutMultiplier(1.0)
+                .setMaxRetryDelay(Duration.ofMillis(45000L))
+                .setTotalTimeout(Duration.ofHours(24L))
+                .build());
+    // Initialize client that will be used to send requests. This client only needs to be created
+    // once, and can be reused for multiple requests.
+    try (TpuClient tpuClient = TpuClient.create(clientSettings.build())) {
+      String parent = String.format("projects/%s/locations/%s", projectId, zone);
+      Node node =
+          Node.newBuilder()
+              .setName(nodeName)
+              .setAcceleratorType(tpuType)
+              .setRuntimeVersion(tpuSoftwareVersion)
+              .setQueuedResource(
+                  String.format(
+                      "projects/%s/locations/%s/queuedResources/%s",
+                      projectId, zone, queuedResourceId))
+              .build();
+
+      QueuedResource queuedResource =
+          QueuedResource.newBuilder()
+              .setName(queuedResourceId)
+              .setTpu(
+                  QueuedResource.Tpu.newBuilder()
+                      .addNodeSpec(
+                          QueuedResource.Tpu.NodeSpec.newBuilder()
+                              .setParent(parent)
+                              .setNode(node)
+                              .setNodeId(nodeName)
+                              .build())
+                      .build())
+              // You can request a queued resource using a reservation by specifying it in code
+              //.setReservationName(
+              // "projects/YOUR_PROJECT_ID/locations/YOUR_ZONE/reservations/YOUR_RESERVATION_NAME")
+              .build();
+
+      CreateQueuedResourceRequest request =
+          CreateQueuedResourceRequest.newBuilder()
+              .setParent(parent)
+              .setQueuedResourceId(queuedResourceId)
+              .setQueuedResource(queuedResource)
+              .build();
+
+      // You can wait until TPU Node is READY,
+      // and check its status using getTpuVm() from "tpu_vm_get" sample.
+
+      return tpuClient.createQueuedResourceAsync(request).get();
+    }
+  }
+}
+//[END tpu_queued_resources_create]

tpu/src/main/java/tpu/DeleteForceQueuedResource.java

@@ -74,4 +74,4 @@ public static void deleteForceQueuedResource(
     System.out.printf("Deleted Queued Resource: %s\n", name);
   }
 }
-//[END tpu_queued_resources_delete_force]
+//[END tpu_queued_resources_delete_force]

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

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package tpu;
+
+//[START tpu_queued_resources_delete]
+import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.rpc.UnknownException;
+import com.google.cloud.tpu.v2alpha1.DeleteQueuedResourceRequest;
+import com.google.cloud.tpu.v2alpha1.GetQueuedResourceRequest;
+import com.google.cloud.tpu.v2alpha1.QueuedResource;
+import com.google.cloud.tpu.v2alpha1.TpuClient;
+import com.google.cloud.tpu.v2alpha1.TpuSettings;
+import java.io.IOException;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import org.threeten.bp.Duration;
+
+public class DeleteQueuedResource {
+  public static void main(String[] args) {
+    // TODO(developer): Replace these variables before running the sample.
+    // Project ID or project number of the Google Cloud project.
+    String projectId = "YOUR_PROJECT_ID";
+    // The zone in which the TPU was created.
+    String zone = "europe-west4-a";
+    // The name for your Queued Resource.
+    String queuedResourceId = "QUEUED_RESOURCE_ID";
+
+    deleteQueuedResource(projectId, zone, queuedResourceId);
+  }
+
+  // Deletes a Queued Resource asynchronously.
+  public static void deleteQueuedResource(String projectId, String zone, String queuedResourceId) {
+    String name = String.format("projects/%s/locations/%s/queuedResources/%s",
+        projectId, zone, queuedResourceId);
+    // With these settings the client library handles the Operation's polling mechanism
+    // and prevent CancellationException error
+    TpuSettings.Builder clientSettings =
+        TpuSettings.newBuilder();
+    clientSettings
+        .deleteQueuedResourceSettings()
+        .setRetrySettings(
+            RetrySettings.newBuilder()
+                .setInitialRetryDelay(Duration.ofMillis(5000L))
+                .setRetryDelayMultiplier(2.0)
+                .setInitialRpcTimeout(Duration.ZERO)
+                .setRpcTimeoutMultiplier(1.0)
+                .setMaxRetryDelay(Duration.ofMillis(45000L))
+                .setTotalTimeout(Duration.ofHours(24L))
+                .build());
+    // Initialize client that will be used to send requests. This client only needs to be created
+    // once, and can be reused for multiple requests.
+    try (TpuClient tpuClient = TpuClient.create(clientSettings.build())) {
+      // Retrive node name
+      GetQueuedResourceRequest getRequest =
+          GetQueuedResourceRequest.newBuilder().setName(name).build();
+      QueuedResource queuedResource = tpuClient.getQueuedResource(getRequest);
+      String nodeName = queuedResource.getTpu().getNodeSpec(0).getNode().getName();
+      // Before deleting the queued resource it is required to delete the TPU VM.
+      DeleteTpuVm.deleteTpuVm(projectId, zone, nodeName);
+      // Wait until TpuVm is deleted
+      TimeUnit.MINUTES.sleep(3);
+
+      DeleteQueuedResourceRequest request =
+          DeleteQueuedResourceRequest.newBuilder().setName(name).build();
+
+      tpuClient.deleteQueuedResourceAsync(request).get();
+
+    } catch (UnknownException | InterruptedException | ExecutionException | IOException e) {
+      System.out.println(e.getMessage());
+    }
+    System.out.printf("Deleted Queued Resource: %s\n", name);
+  }
+}
+//[END tpu_queued_resources_delete]

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

@@ -17,7 +17,6 @@
 package tpu;
 
 //[START tpu_queued_resources_get]
-
 import com.google.cloud.tpu.v2alpha1.GetQueuedResourceRequest;
 import com.google.cloud.tpu.v2alpha1.QueuedResource;
 import com.google.cloud.tpu.v2alpha1.TpuClient;
@@ -51,4 +50,4 @@ public static QueuedResource getQueuedResource(
     }
   }
 }
-//[END tpu_queued_resources_get]
+//[END tpu_queued_resources_get]

tpu/src/test/java/tpu/QueuedResourcesIT.java

@@ -0,0 +1,115 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package tpu;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+import static org.junit.Assert.assertNotNull;
+
+import com.google.api.gax.rpc.NotFoundException;
+import com.google.cloud.tpu.v2alpha1.QueuedResource;
+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.TimeUnit;
+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.jupiter.api.Timeout;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+@Timeout(value = 6, unit = TimeUnit.MINUTES)
+public class QueuedResourcesIT {
+  private static final String PROJECT_ID = System.getenv("GOOGLE_CLOUD_PROJECT");
+  private static final String ZONE = "us-central1-f";
+  private static final String NODE_NAME = "test-tpu-queued-resource-" + UUID.randomUUID();
+  private static final String NODE_WITH_NETWORK_NAME =
+      "test-tpu-queued-resource-network-" + UUID.randomUUID();
+  private static final String TPU_TYPE = "v2-8";
+  private static final String TPU_SOFTWARE_VERSION = "tpu-vm-tf-2.17.0-pjrt";
+  private static final String QUEUED_RESOURCE_NAME = "queued-resource-" + UUID.randomUUID();
+  private static final String QUEUED_RESOURCE_PATH_NAME =
+      String.format("projects/%s/locations/%s/queuedResources/%s",
+          PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME);
+  private static final String QUEUED_RESOURCE_WITH_NETWORK =
+      "queued-resource-network-" + UUID.randomUUID();
+  private static final String NETWORK_NAME = "default";
+
+  public static void requireEnvVar(String envVarName) {
+    assertWithMessage(String.format("Missing environment variable '%s' ", envVarName))
+        .that(System.getenv(envVarName)).isNotEmpty();
+  }
+
+  @BeforeAll
+  public static void setUp() throws IOException, ExecutionException, InterruptedException {
+    requireEnvVar("GOOGLE_APPLICATION_CREDENTIALS");
+    requireEnvVar("GOOGLE_CLOUD_PROJECT");
+
+    QueuedResource queuedResource = CreateQueuedResource.createQueuedResource(PROJECT_ID, ZONE,
+        QUEUED_RESOURCE_NAME, NODE_NAME, TPU_TYPE, TPU_SOFTWARE_VERSION);
+
+    assertThat(queuedResource.getName()).isEqualTo(QUEUED_RESOURCE_PATH_NAME);
+    assertThat(queuedResource.getTpu().getNodeSpec(0).getNode().getName()).isEqualTo(NODE_NAME);
+  }
+
+  @AfterAll
+  public static void cleanup() throws IOException {
+    final PrintStream out = System.out;
+    ByteArrayOutputStream stdOut = new ByteArrayOutputStream();
+    System.setOut(new PrintStream(stdOut));
+    DeleteQueuedResource.deleteQueuedResource(PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME);
+    DeleteForceQueuedResource.deleteForceQueuedResource(
+        PROJECT_ID, ZONE, QUEUED_RESOURCE_WITH_NETWORK);
+
+    // Test that resources are deleted
+    assertThat(stdOut.toString()).contains("Deleted Queued Resource:");
+    Assertions.assertThrows(
+        NotFoundException.class,
+        () -> GetTpuVm.getTpuVm(PROJECT_ID, ZONE, NODE_NAME));
+
+    stdOut.close();
+    System.setOut(out);
+  }
+
+  @Test
+  public void testGetQueuedResource() throws IOException {
+    QueuedResource queuedResource = GetQueuedResource.getQueuedResource(
+        PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME);
+
+    assertNotNull(queuedResource);
+    assertThat(queuedResource.getName()).isEqualTo(QUEUED_RESOURCE_PATH_NAME);
+  }
+
+  @Test
+  public void testCreateQueuedResourceWithSpecifiedNetwork() throws Exception {
+    QueuedResource queuedResource = CreateQueuedResourceWithNetwork.createQueuedResourceWithNetwork(
+        PROJECT_ID, ZONE, QUEUED_RESOURCE_WITH_NETWORK, NODE_WITH_NETWORK_NAME,
+        TPU_TYPE, TPU_SOFTWARE_VERSION, NETWORK_NAME);
+
+    assertThat(queuedResource.getTpu().getNodeSpec(0).getNode().getName())
+        .isEqualTo(NODE_WITH_NETWORK_NAME);
+    assertThat(queuedResource.getTpu().getNodeSpec(0).getNode().getNetworkConfig().getNetwork()
+        .contains(NETWORK_NAME));
+    assertThat(queuedResource.getTpu().getNodeSpec(0).getNode().getNetworkConfig().getSubnetwork()
+        .contains(NETWORK_NAME));
+  }
+}