Skip to content

Commit

Permalink
Resolved merge conflict
Browse files Browse the repository at this point in the history
  • Loading branch information
@TetyanaYahodska
TetyanaYahodska committed Dec 20, 2024
2 parents 99d4a4a + 772c39c commit 5a7af75
Show file tree
Hide file tree
Showing 4 changed files with 66 additions and 75 deletions.
44 changes: 10 additions & 34 deletions tpu/src/main/java/tpu/CreateQueuedResource.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,26 +17,25 @@
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;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

public class CreateQueuedResource {
public static void main(String[] args)
throws IOException, ExecutionException, InterruptedException {
throws IOException, ExecutionException, InterruptedException, TimeoutException {
// 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";
String zone = "us-central1-f";
// 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.
Expand All @@ -56,35 +55,19 @@ public static void main(String[] args)
// 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());
throws IOException, ExecutionException, InterruptedException, TimeoutException {
String resource = String.format("projects/%s/locations/%s/queuedResources/%s",
projectId, zone, queuedResourceId);
// 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())) {
try (TpuClient tpuClient = TpuClient.create()) {
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))
.setQueuedResource(resource)
.build();

QueuedResource queuedResource =
Expand All @@ -99,9 +82,6 @@ public static QueuedResource createQueuedResource(String projectId, String zone,
.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 =
Expand All @@ -111,11 +91,7 @@ public static QueuedResource createQueuedResource(String projectId, String zone,
.setQueuedResource(queuedResource)
.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 created: %s\n", response.getName());
return response;
return tpuClient.createQueuedResourceAsync(request).get(1, TimeUnit.MINUTES);
}
}
}
Expand Down
48 changes: 9 additions & 39 deletions tpu/src/main/java/tpu/DeleteQueuedResource.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,72 +17,42 @@
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) {
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.
String projectId = "YOUR_PROJECT_ID";
// The zone in which the TPU was created.
String zone = "europe-west4-a";
String zone = "us-central1-f";
// 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) {
public static void deleteQueuedResource(String projectId, String zone, String queuedResourceId)
throws ExecutionException, InterruptedException, IOException {
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();
try (TpuClient tpuClient = TpuClient.create()) {
// 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);
// For more information about deleting TPU
// see https://cloud.google.com/tpu/docs/managing-tpus-tpu-vm

DeleteQueuedResourceRequest request =
DeleteQueuedResourceRequest.newBuilder().setName(name).build();
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]
2 changes: 1 addition & 1 deletion tpu/src/main/java/tpu/GetQueuedResource.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ public static void main(String[] args) throws IOException {
// 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";
String zone = "us-central1-f";
// The name for your Queued Resource.
String queuedResourceId = "QUEUED_RESOURCE_ID";

Expand Down
47 changes: 46 additions & 1 deletion tpu/src/test/java/tpu/QueuedResourceIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import static com.google.common.truth.Truth.assertThat;
import static org.junit.Assert.assertEquals;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockStatic;
Expand All @@ -37,14 +38,15 @@
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.MockedStatic;

@RunWith(JUnit4.class)
@Timeout(value = 10)
@Timeout(value = 2, unit = TimeUnit.MINUTES)
public class QueuedResourceIT {
private static final String PROJECT_ID = "project-id";
private static final String ZONE = "europe-west4-a";
Expand All @@ -54,6 +56,30 @@ public class QueuedResourceIT {
private static final String QUEUED_RESOURCE_NAME = "queued-resource";
private static final String NETWORK_NAME = "default";

@Test
public void testCreateQueuedResource() throws Exception {
try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
QueuedResource mockQueuedResource = mock(QueuedResource.class);
TpuClient mockTpuClient = mock(TpuClient.class);
OperationFuture mockFuture = mock(OperationFuture.class);

mockedTpuClient.when(TpuClient::create).thenReturn(mockTpuClient);
when(mockTpuClient.createQueuedResourceAsync(any(CreateQueuedResourceRequest.class)))
.thenReturn(mockFuture);
when(mockFuture.get(anyLong(), any(TimeUnit.class))).thenReturn(mockQueuedResource);

QueuedResource returnedQueuedResource =
CreateQueuedResource.createQueuedResource(
PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME, NODE_NAME,
TPU_TYPE, TPU_SOFTWARE_VERSION);

verify(mockTpuClient, times(1))
.createQueuedResourceAsync(any(CreateQueuedResourceRequest.class));
verify(mockFuture, times(1)).get(anyLong(), any(TimeUnit.class));
assertEquals(returnedQueuedResource, mockQueuedResource);
}
}

@Test
public void testCreateQueuedResourceWithSpecifiedNetwork() throws Exception {
try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
Expand Down Expand Up @@ -144,6 +170,25 @@ public void testDeleteForceQueuedResource()
}
}

@Test
public void testDeleteQueuedResource()
throws IOException, ExecutionException, InterruptedException {
try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
TpuClient mockTpuClient = mock(TpuClient.class);
OperationFuture mockFuture = mock(OperationFuture.class);

mockedTpuClient.when(TpuClient::create).thenReturn(mockTpuClient);
when(mockTpuClient.deleteQueuedResourceAsync(any(DeleteQueuedResourceRequest.class)))
.thenReturn(mockFuture);
when(mockFuture.get()).thenReturn(null);

DeleteQueuedResource.deleteQueuedResource(PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME);

verify(mockTpuClient, times(1))
.deleteQueuedResourceAsync(any(DeleteQueuedResourceRequest.class));
}
}

@Test
public void testCreateQueuedResourceWithStartupScript() throws Exception {
try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
Expand Down

0 comments on commit 5a7af75

Please sign in to comment.