Fixed samples and tests

Author: TetyanaYahodska

tpu/src/main/java/tpu/CreateQueuedResourceWithNetwork.java

@@ -1,7 +1,22 @@
+/*
+ * 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_network]
-
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.cloud.tpu.v2alpha1.CreateQueuedResourceRequest;
 import com.google.cloud.tpu.v2alpha1.NetworkConfig;
@@ -14,22 +29,17 @@
 import org.threeten.bp.Duration;
 
 public class CreateQueuedResourceWithNetwork {
-
-  private final TpuClient tpuClient;
-  public CreateQueuedResourceWithNetwork(TpuClient tpuClient) {
-    this.tpuClient = tpuClient;
-  }
   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 = "tyaho-softserve-project";//"your-project";
+    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";
+    String nodeName = "YOUR_TPY_NAME";
     // 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.
@@ -38,22 +48,24 @@ public static void main(String[] args)
     // 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 queuedResourceName = "your-queued-resource-id";
-    TpuClient tpuClient = TpuClient.create();
-    CreateQueuedResourceWithNetwork creator = new CreateQueuedResourceWithNetwork(tpuClient);
-
-    creator.createQueuedResourceWithNetwork(projectId, zone, queuedResourceName, nodeName, tpuType, tpuSoftwareVersion);
+    String queuedResourceId = "QUEUED_RESOURCE_ID";
+    // The name of the network you want the node to connect to.
+    // The network should be assigned to your project.
+    String networkName = "YOUR_COMPUTE_TPU_NETWORK";
 
+    createQueuedResourceWithNetwork(projectId, zone, queuedResourceId, nodeName,
+        tpuType, tpuSoftwareVersion, networkName);
   }
 
   // Creates a Queued Resource with network configuration.
-  public static void createQueuedResourceWithNetwork(
+  public static QueuedResource createQueuedResourceWithNetwork(
       String projectId,
       String zone,
-      String queuedResourceName,
+      String queuedResourceId,
       String nodeName,
       String tpuType,
-      String tpuSoftwareVersion)
+      String tpuSoftwareVersion,
+      String networkName)
       throws IOException, ExecutionException, InterruptedException {
     // With these settings the client library handles the Operation's polling mechanism
     // and prevent CancellationException error
@@ -75,10 +87,8 @@ public static void createQueuedResourceWithNetwork(
     try (TpuClient tpuClient = TpuClient.create(clientSettings.build())) {
       String parent = String.format("projects/%s/locations/%s", projectId, zone);
       String region = zone.substring(0, zone.length() - 2);
-      // The name of the network you want the node to connect to. The network should be assigned to your project.
-      String networkName = "compute-tpu-network";
-      // Specify the network and subnetwork that you want to connect your TPU to.
 
+      // Specify the network and subnetwork that you want to connect your TPU to.
       NetworkConfig networkConfig =
           NetworkConfig.newBuilder()
               .setEnableExternalIps(true)
@@ -98,13 +108,13 @@ public static void createQueuedResourceWithNetwork(
               .setQueuedResource(
                   String.format(
                       "projects/%s/locations/%s/queuedResources/%s",
-                      projectId, zone, queuedResourceName))
+                      projectId, zone, queuedResourceId))
               .build();
 
       // Create queued resource
       QueuedResource queuedResource =
           QueuedResource.newBuilder()
-              .setName(queuedResourceName)
+              .setName(queuedResourceId)
               .setTpu(
                   QueuedResource.Tpu.newBuilder()
                       .addNodeSpec(
@@ -120,15 +130,14 @@ public static void createQueuedResourceWithNetwork(
           CreateQueuedResourceRequest.newBuilder()
               .setParent(parent)
               .setQueuedResource(queuedResource)
-              .setQueuedResourceId(queuedResourceName)
+              .setQueuedResourceId(queuedResourceId)
               .build();
 
       QueuedResource response = tpuClient.createQueuedResourceAsync(request).get();
       // You can wait until TPU Node is READY,
       // and check its status using getTpuVm() from "tpu_vm_get" sample.
-      System.out.printf(
-          "Queued resource %s with specified network created.\n", queuedResourceName);
-      System.out.println(response);
+      System.out.println("Queued Resource created: " + queuedResourceId);
+      return response;
     }
   }
 }

tpu/src/main/java/tpu/CreateQueuedResourceWithStartupScript.java

@@ -1,3 +1,19 @@
+/*
+ * 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 com.google.api.gax.retrying.RetrySettings;
@@ -6,25 +22,24 @@
 import com.google.cloud.tpu.v2alpha1.QueuedResource;
 import com.google.cloud.tpu.v2alpha1.TpuClient;
 import com.google.cloud.tpu.v2alpha1.TpuSettings;
-import org.threeten.bp.Duration;
-
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
+import org.threeten.bp.Duration;
 
 public class CreateQueuedResourceWithStartupScript {
   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 = "tyaho-softserve-project";//"your-project";
+    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";
+    String nodeName = "YOUR_TPY_NAME";
     // 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.
@@ -33,14 +48,14 @@ public static void main(String[] args)
     // 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 = "your-queued-resource-id";
+    String queuedResourceId = "QUEUED_RESOURCE_ID";
 
     createQueuedResource(
         projectId, zone, queuedResourceId, nodeName, tpuType, tpuSoftwareVersion);
   }
 
   // Creates a Queued Resource
-  public static void createQueuedResource(
+  public static QueuedResource createQueuedResource(
       String projectId,
       String zone,
       String queuedResourceId,
@@ -110,6 +125,7 @@ public static void createQueuedResource(
       // You can wait until TPU Node is READY,
       // and check its status using getTpuVm() from "tpu_vm_get" sample.
       System.out.printf("Queued Resource created: %s\n", queuedResourceId);
+      return response;
     }
   }
 }

tpu/src/main/java/tpu/CreateSpotQueuedResource.java

@@ -1,3 +1,19 @@
+/*
+ * 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 com.google.api.gax.retrying.RetrySettings;
@@ -7,19 +23,16 @@
 import com.google.cloud.tpu.v2alpha1.SchedulingConfig;
 import com.google.cloud.tpu.v2alpha1.TpuClient;
 import com.google.cloud.tpu.v2alpha1.TpuSettings;
-import org.threeten.bp.Duration;
-
 import java.io.IOException;
-import java.util.HashMap;
-import java.util.Map;
 import java.util.concurrent.ExecutionException;
+import org.threeten.bp.Duration;
 
 public class CreateSpotQueuedResource {
   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 = "tyaho-softserve-project";//"your-project";
+    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
@@ -41,7 +54,7 @@ public static void main(String[] args)
   }
 
   // Creates a Queued Resource
-  public static void createQueuedResource(
+  public static QueuedResource createQueuedResource(
       String projectId,
       String zone,
       String queuedResourceId,
@@ -109,6 +122,7 @@ public static void createQueuedResource(
       // You can wait until TPU Node is READY,
       // and check its status using getTpuVm() from "tpu_vm_get" sample.
       System.out.printf("Queued Resource created: %s\n", queuedResourceId);
+      return response;
     }
   }
 }