diff --git a/tpu/src/main/java/tpu/CreateQueuedResource.java b/tpu/src/main/java/tpu/CreateQueuedResource.java new file mode 100644 index 00000000000..3503d12c6fa --- /dev/null +++ b/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] \ No newline at end of file diff --git a/tpu/src/main/java/tpu/DeleteForceQueuedResource.java b/tpu/src/main/java/tpu/DeleteForceQueuedResource.java index 3de8567857d..91e75612954 100644 --- a/tpu/src/main/java/tpu/DeleteForceQueuedResource.java +++ b/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] \ No newline at end of file diff --git a/tpu/src/main/java/tpu/DeleteQueuedResource.java b/tpu/src/main/java/tpu/DeleteQueuedResource.java new file mode 100644 index 00000000000..0f592ff165f --- /dev/null +++ b/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] diff --git a/tpu/src/main/java/tpu/GetQueuedResource.java b/tpu/src/main/java/tpu/GetQueuedResource.java index 3a510e045fe..80c6762327c 100644 --- a/tpu/src/main/java/tpu/GetQueuedResource.java +++ b/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] \ No newline at end of file diff --git a/tpu/src/test/java/tpu/QueuedResourcesIT.java b/tpu/src/test/java/tpu/QueuedResourcesIT.java new file mode 100644 index 00000000000..487b29523f4 --- /dev/null +++ b/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)); + } +} \ No newline at end of file