Skip to content

Commit

Permalink
feat: [tpu] Add ResetQueueResource to the TPU service (#9748)
Browse files Browse the repository at this point in the history
- [ ] Regenerate this pull request now.

feat: Add MultiNodeParams in TPU

PiperOrigin-RevId: 559244735

Source-Link: https://togithub.com/googleapis/googleapis/commit/75a3251882cca05096fe3a620eb67c0c52f45c8b

Source-Link: https://togithub.com/googleapis/googleapis-gen/commit/b2544705a21090f36b3e2f00a7c6c47000081573
Copy-Tag: eyJwIjoiamF2YS10cHUvLk93bEJvdC55YW1sIiwiaCI6ImIyNTQ0NzA1YTIxMDkwZjM2YjNlMmYwMGE3YzZjNDcwMDAwODE1NzMifQ==

BEGIN_NESTED_COMMIT
build: [tpu] do not create bazel-* symlinks
PiperOrigin-RevId: 555307834

Source-Link: https://togithub.com/googleapis/googleapis/commit/43f02766add4ac5e2e6557a9d2600f2101654dff

Source-Link: https://togithub.com/googleapis/googleapis-gen/commit/f340c3f9408f87a93ab676233181b2738de4a9fe
Copy-Tag: eyJwIjoiamF2YS10cHUvLk93bEJvdC55YW1sIiwiaCI6ImYzNDBjM2Y5NDA4Zjg3YTkzYWI2NzYyMzMxODFiMjczOGRlNGE5ZmUifQ==
END_NESTED_COMMIT
  • Loading branch information
gcf-owl-bot[bot] committed Aug 30, 2023
1 parent 48ac78c commit 9996b84
Show file tree
Hide file tree
Showing 23 changed files with 3,473 additions and 300 deletions.
4 changes: 2 additions & 2 deletions java-tpu/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ If you are using Maven with [BOM][libraries-bom], add this to your pom.xml file:
<dependency>
<groupId>com.google.cloud</groupId>
<artifactId>libraries-bom</artifactId>
<version>26.19.0</version>
<version>26.22.0</version>
<type>pom</type>
<scope>import</scope>
</dependency>
Expand Down Expand Up @@ -195,7 +195,7 @@ Java is a registered trademark of Oracle and/or its affiliates.
[kokoro-badge-link-5]: http://storage.googleapis.com/cloud-devrel-public/java/badges/google-cloud-java/java11.html
[stability-image]: https://img.shields.io/badge/stability-stable-green
[maven-version-image]: https://img.shields.io/maven-central/v/com.google.cloud/google-cloud-tpu.svg
[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-tpu/2.23.0
[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-tpu/2.24.0
[authentication]: https://github.com/googleapis/google-cloud-java#authentication
[auth-scopes]: https://developers.google.com/identity/protocols/oauth2/scopes
[predefined-iam-roles]: https://cloud.google.com/iam/docs/understanding-roles#predefined_roles
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1639,6 +1639,154 @@ public final OperationFuture<QueuedResource, OperationMetadata> deleteQueuedReso
return stub.deleteQueuedResourceCallable();
}

// AUTO-GENERATED DOCUMENTATION AND METHOD.
/**
* Resets a QueuedResource TPU instance
*
* <p>Sample code:
*
* <pre>{@code
* // This snippet has been automatically generated and should be regarded as a code template only.
* // It will require modifications to work:
* // - It may require correct/in-range values for request initialization.
* // - It may require specifying regional endpoints when creating the service client as shown in
* // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
* try (TpuClient tpuClient = TpuClient.create()) {
* QueuedResourceName name =
* QueuedResourceName.of("[PROJECT]", "[LOCATION]", "[QUEUED_RESOURCE]");
* QueuedResource response = tpuClient.resetQueuedResourceAsync(name).get();
* }
* }</pre>
*
* @param name Required. The name of the queued resource.
* @throws com.google.api.gax.rpc.ApiException if the remote call fails
*/
public final OperationFuture<QueuedResource, OperationMetadata> resetQueuedResourceAsync(
QueuedResourceName name) {
ResetQueuedResourceRequest request =
ResetQueuedResourceRequest.newBuilder()
.setName(name == null ? null : name.toString())
.build();
return resetQueuedResourceAsync(request);
}

// AUTO-GENERATED DOCUMENTATION AND METHOD.
/**
* Resets a QueuedResource TPU instance
*
* <p>Sample code:
*
* <pre>{@code
* // This snippet has been automatically generated and should be regarded as a code template only.
* // It will require modifications to work:
* // - It may require correct/in-range values for request initialization.
* // - It may require specifying regional endpoints when creating the service client as shown in
* // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
* try (TpuClient tpuClient = TpuClient.create()) {
* String name =
* QueuedResourceName.of("[PROJECT]", "[LOCATION]", "[QUEUED_RESOURCE]").toString();
* QueuedResource response = tpuClient.resetQueuedResourceAsync(name).get();
* }
* }</pre>
*
* @param name Required. The name of the queued resource.
* @throws com.google.api.gax.rpc.ApiException if the remote call fails
*/
public final OperationFuture<QueuedResource, OperationMetadata> resetQueuedResourceAsync(
String name) {
ResetQueuedResourceRequest request =
ResetQueuedResourceRequest.newBuilder().setName(name).build();
return resetQueuedResourceAsync(request);
}

// AUTO-GENERATED DOCUMENTATION AND METHOD.
/**
* Resets a QueuedResource TPU instance
*
* <p>Sample code:
*
* <pre>{@code
* // This snippet has been automatically generated and should be regarded as a code template only.
* // It will require modifications to work:
* // - It may require correct/in-range values for request initialization.
* // - It may require specifying regional endpoints when creating the service client as shown in
* // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
* try (TpuClient tpuClient = TpuClient.create()) {
* ResetQueuedResourceRequest request =
* ResetQueuedResourceRequest.newBuilder()
* .setName(
* QueuedResourceName.of("[PROJECT]", "[LOCATION]", "[QUEUED_RESOURCE]").toString())
* .build();
* QueuedResource response = tpuClient.resetQueuedResourceAsync(request).get();
* }
* }</pre>
*
* @param request The request object containing all of the parameters for the API call.
* @throws com.google.api.gax.rpc.ApiException if the remote call fails
*/
public final OperationFuture<QueuedResource, OperationMetadata> resetQueuedResourceAsync(
ResetQueuedResourceRequest request) {
return resetQueuedResourceOperationCallable().futureCall(request);
}

// AUTO-GENERATED DOCUMENTATION AND METHOD.
/**
* Resets a QueuedResource TPU instance
*
* <p>Sample code:
*
* <pre>{@code
* // This snippet has been automatically generated and should be regarded as a code template only.
* // It will require modifications to work:
* // - It may require correct/in-range values for request initialization.
* // - It may require specifying regional endpoints when creating the service client as shown in
* // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
* try (TpuClient tpuClient = TpuClient.create()) {
* ResetQueuedResourceRequest request =
* ResetQueuedResourceRequest.newBuilder()
* .setName(
* QueuedResourceName.of("[PROJECT]", "[LOCATION]", "[QUEUED_RESOURCE]").toString())
* .build();
* OperationFuture<QueuedResource, OperationMetadata> future =
* tpuClient.resetQueuedResourceOperationCallable().futureCall(request);
* // Do something.
* QueuedResource response = future.get();
* }
* }</pre>
*/
public final OperationCallable<ResetQueuedResourceRequest, QueuedResource, OperationMetadata>
resetQueuedResourceOperationCallable() {
return stub.resetQueuedResourceOperationCallable();
}

// AUTO-GENERATED DOCUMENTATION AND METHOD.
/**
* Resets a QueuedResource TPU instance
*
* <p>Sample code:
*
* <pre>{@code
* // This snippet has been automatically generated and should be regarded as a code template only.
* // It will require modifications to work:
* // - It may require correct/in-range values for request initialization.
* // - It may require specifying regional endpoints when creating the service client as shown in
* // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
* try (TpuClient tpuClient = TpuClient.create()) {
* ResetQueuedResourceRequest request =
* ResetQueuedResourceRequest.newBuilder()
* .setName(
* QueuedResourceName.of("[PROJECT]", "[LOCATION]", "[QUEUED_RESOURCE]").toString())
* .build();
* ApiFuture<Operation> future = tpuClient.resetQueuedResourceCallable().futureCall(request);
* // Do something.
* Operation response = future.get();
* }
* }</pre>
*/
public final UnaryCallable<ResetQueuedResourceRequest, Operation> resetQueuedResourceCallable() {
return stub.resetQueuedResourceCallable();
}

// AUTO-GENERATED DOCUMENTATION AND METHOD.
/**
* Generates the Cloud TPU service identity for the project.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,17 @@ public UnaryCallSettings<DeleteQueuedResourceRequest, Operation> deleteQueuedRes
return ((TpuStubSettings) getStubSettings()).deleteQueuedResourceOperationSettings();
}

/** Returns the object with the settings used for calls to resetQueuedResource. */
public UnaryCallSettings<ResetQueuedResourceRequest, Operation> resetQueuedResourceSettings() {
return ((TpuStubSettings) getStubSettings()).resetQueuedResourceSettings();
}

/** Returns the object with the settings used for calls to resetQueuedResource. */
public OperationCallSettings<ResetQueuedResourceRequest, QueuedResource, OperationMetadata>
resetQueuedResourceOperationSettings() {
return ((TpuStubSettings) getStubSettings()).resetQueuedResourceOperationSettings();
}

/** Returns the object with the settings used for calls to generateServiceIdentity. */
public UnaryCallSettings<GenerateServiceIdentityRequest, GenerateServiceIdentityResponse>
generateServiceIdentitySettings() {
Expand Down Expand Up @@ -450,6 +461,19 @@ public UnaryCallSettings.Builder<UpdateNodeRequest, Operation> updateNodeSetting
return getStubSettingsBuilder().deleteQueuedResourceOperationSettings();
}

/** Returns the builder for the settings used for calls to resetQueuedResource. */
public UnaryCallSettings.Builder<ResetQueuedResourceRequest, Operation>
resetQueuedResourceSettings() {
return getStubSettingsBuilder().resetQueuedResourceSettings();
}

/** Returns the builder for the settings used for calls to resetQueuedResource. */
public OperationCallSettings.Builder<
ResetQueuedResourceRequest, QueuedResource, OperationMetadata>
resetQueuedResourceOperationSettings() {
return getStubSettingsBuilder().resetQueuedResourceOperationSettings();
}

/** Returns the builder for the settings used for calls to generateServiceIdentity. */
public UnaryCallSettings.Builder<
GenerateServiceIdentityRequest, GenerateServiceIdentityResponse>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,9 @@
"ListRuntimeVersions": {
"methods": ["listRuntimeVersions", "listRuntimeVersions", "listRuntimeVersions", "listRuntimeVersionsPagedCallable", "listRuntimeVersionsCallable"]
},
"ResetQueuedResource": {
"methods": ["resetQueuedResourceAsync", "resetQueuedResourceAsync", "resetQueuedResourceAsync", "resetQueuedResourceOperationCallable", "resetQueuedResourceCallable"]
},
"SimulateMaintenanceEvent": {
"methods": ["simulateMaintenanceEventAsync", "simulateMaintenanceEventOperationCallable", "simulateMaintenanceEventCallable"]
},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@
import com.google.cloud.tpu.v2alpha1.Node;
import com.google.cloud.tpu.v2alpha1.OperationMetadata;
import com.google.cloud.tpu.v2alpha1.QueuedResource;
import com.google.cloud.tpu.v2alpha1.ResetQueuedResourceRequest;
import com.google.cloud.tpu.v2alpha1.RuntimeVersion;
import com.google.cloud.tpu.v2alpha1.SimulateMaintenanceEventRequest;
import com.google.cloud.tpu.v2alpha1.StartNodeRequest;
Expand Down Expand Up @@ -180,6 +181,16 @@ public class GrpcTpuStub extends TpuStub {
.setResponseMarshaller(ProtoUtils.marshaller(Operation.getDefaultInstance()))
.build();

private static final MethodDescriptor<ResetQueuedResourceRequest, Operation>
resetQueuedResourceMethodDescriptor =
MethodDescriptor.<ResetQueuedResourceRequest, Operation>newBuilder()
.setType(MethodDescriptor.MethodType.UNARY)
.setFullMethodName("google.cloud.tpu.v2alpha1.Tpu/ResetQueuedResource")
.setRequestMarshaller(
ProtoUtils.marshaller(ResetQueuedResourceRequest.getDefaultInstance()))
.setResponseMarshaller(ProtoUtils.marshaller(Operation.getDefaultInstance()))
.build();

private static final MethodDescriptor<
GenerateServiceIdentityRequest, GenerateServiceIdentityResponse>
generateServiceIdentityMethodDescriptor =
Expand Down Expand Up @@ -304,6 +315,9 @@ public class GrpcTpuStub extends TpuStub {
private final UnaryCallable<DeleteQueuedResourceRequest, Operation> deleteQueuedResourceCallable;
private final OperationCallable<DeleteQueuedResourceRequest, QueuedResource, OperationMetadata>
deleteQueuedResourceOperationCallable;
private final UnaryCallable<ResetQueuedResourceRequest, Operation> resetQueuedResourceCallable;
private final OperationCallable<ResetQueuedResourceRequest, QueuedResource, OperationMetadata>
resetQueuedResourceOperationCallable;
private final UnaryCallable<GenerateServiceIdentityRequest, GenerateServiceIdentityResponse>
generateServiceIdentityCallable;
private final UnaryCallable<ListAcceleratorTypesRequest, ListAcceleratorTypesResponse>
Expand Down Expand Up @@ -476,6 +490,16 @@ protected GrpcTpuStub(
return builder.build();
})
.build();
GrpcCallSettings<ResetQueuedResourceRequest, Operation> resetQueuedResourceTransportSettings =
GrpcCallSettings.<ResetQueuedResourceRequest, Operation>newBuilder()
.setMethodDescriptor(resetQueuedResourceMethodDescriptor)
.setParamsExtractor(
request -> {
RequestParamsBuilder builder = RequestParamsBuilder.create();
builder.add("name", String.valueOf(request.getName()));
return builder.build();
})
.build();
GrpcCallSettings<GenerateServiceIdentityRequest, GenerateServiceIdentityResponse>
generateServiceIdentityTransportSettings =
GrpcCallSettings
Expand Down Expand Up @@ -665,6 +689,17 @@ protected GrpcTpuStub(
settings.deleteQueuedResourceOperationSettings(),
clientContext,
operationsStub);
this.resetQueuedResourceCallable =
callableFactory.createUnaryCallable(
resetQueuedResourceTransportSettings,
settings.resetQueuedResourceSettings(),
clientContext);
this.resetQueuedResourceOperationCallable =
callableFactory.createOperationCallable(
resetQueuedResourceTransportSettings,
settings.resetQueuedResourceOperationSettings(),
clientContext,
operationsStub);
this.generateServiceIdentityCallable =
callableFactory.createUnaryCallable(
generateServiceIdentityTransportSettings,
Expand Down Expand Up @@ -841,6 +876,17 @@ public UnaryCallable<DeleteQueuedResourceRequest, Operation> deleteQueuedResourc
return deleteQueuedResourceOperationCallable;
}

@Override
public UnaryCallable<ResetQueuedResourceRequest, Operation> resetQueuedResourceCallable() {
return resetQueuedResourceCallable;
}

@Override
public OperationCallable<ResetQueuedResourceRequest, QueuedResource, OperationMetadata>
resetQueuedResourceOperationCallable() {
return resetQueuedResourceOperationCallable;
}

@Override
public UnaryCallable<GenerateServiceIdentityRequest, GenerateServiceIdentityResponse>
generateServiceIdentityCallable() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@
import com.google.cloud.tpu.v2alpha1.Node;
import com.google.cloud.tpu.v2alpha1.OperationMetadata;
import com.google.cloud.tpu.v2alpha1.QueuedResource;
import com.google.cloud.tpu.v2alpha1.ResetQueuedResourceRequest;
import com.google.cloud.tpu.v2alpha1.RuntimeVersion;
import com.google.cloud.tpu.v2alpha1.SimulateMaintenanceEventRequest;
import com.google.cloud.tpu.v2alpha1.StartNodeRequest;
Expand Down Expand Up @@ -167,6 +168,16 @@ public UnaryCallable<DeleteQueuedResourceRequest, Operation> deleteQueuedResourc
throw new UnsupportedOperationException("Not implemented: deleteQueuedResourceCallable()");
}

public OperationCallable<ResetQueuedResourceRequest, QueuedResource, OperationMetadata>
resetQueuedResourceOperationCallable() {
throw new UnsupportedOperationException(
"Not implemented: resetQueuedResourceOperationCallable()");
}

public UnaryCallable<ResetQueuedResourceRequest, Operation> resetQueuedResourceCallable() {
throw new UnsupportedOperationException("Not implemented: resetQueuedResourceCallable()");
}

public UnaryCallable<GenerateServiceIdentityRequest, GenerateServiceIdentityResponse>
generateServiceIdentityCallable() {
throw new UnsupportedOperationException("Not implemented: generateServiceIdentityCallable()");
Expand Down
Loading

0 comments on commit 9996b84

Please sign in to comment.