feat: Add ResetQueueResource to the TPU service

Google APIs · copybara-github · commit 75a3251882cc · 2023-08-22T15:42:40.000-07:00
feat: Add MultiNodeParams in TPU

PiperOrigin-RevId: 559244735
diff --git a/google/cloud/tpu/v2alpha1/cloud_tpu.proto b/google/cloud/tpu/v2alpha1/cloud_tpu.proto
@@ -164,6 +164,20 @@ service Tpu {
     };
   }
 
+  // Resets a QueuedResource TPU instance
+  rpc ResetQueuedResource(ResetQueuedResourceRequest)
+      returns (google.longrunning.Operation) {
+    option (google.api.http) = {
+      post: "/v2alpha1/{name=projects/*/locations/*/queuedResources/*}:reset"
+      body: "*"
+    };
+    option (google.api.method_signature) = "name";
+    option (google.longrunning.operation_info) = {
+      response_type: "QueuedResource"
+      metadata_type: "OperationMetadata"
+    };
+  }
+
   // Generates the Cloud TPU service identity for the project.
   rpc GenerateServiceIdentity(GenerateServiceIdentityRequest)
       returns (GenerateServiceIdentityResponse) {
@@ -544,6 +558,23 @@ message QueuedResource {
     // single node or multiple nodes.
     // NodeSpec provides the specification for node(s) to be created.
     message NodeSpec {
+      // Parameters to specify for multi-node QueuedResource requests. This
+      // field must be populated in case of multi-node requests instead of
+      // node_id. It's an error to specify both node_id and multi_node_params.
+      message MultiNodeParams {
+        // Required. Number of nodes with this spec. The system will attempt
+        // to provison "node_count" nodes as part of the request.
+        // This needs to be > 1.
+        int32 node_count = 1 [(google.api.field_behavior) = REQUIRED];
+
+        // Prefix of node_ids in case of multi-node request
+        // Should follow the `^[A-Za-z0-9_.~+%-]+$` regex format.
+        // If node_count = 3 and node_id_prefix = "np", node ids of nodes
+        // created will be "np-0", "np-1", "np-2". If this field is not
+        // provided we use queued_resource_id as the node_id_prefix.
+        string node_id_prefix = 2;
+      }
+
       // Required. The parent resource name.
       string parent = 1 [
         (google.api.field_behavior) = REQUIRED,
@@ -558,6 +589,10 @@ message QueuedResource {
       // instead. It's an error to specify both node_id and multi_node_params.
       string node_id = 2;
 
+      // Optional. Fields to specify in case of multi-node request.
+      MultiNodeParams multi_node_params = 6
+          [(google.api.field_behavior) = OPTIONAL];
+
       // Required. The node.
       Node node = 3 [(google.api.field_behavior) = REQUIRED];
     }
@@ -946,6 +981,18 @@ message DeleteQueuedResourceRequest {
   bool force = 3;
 }
 
+// Request for
+// [ResetQueuedResource][google.cloud.tpu.v2alpha1.Tpu.ResetQueuedResource].
+message ResetQueuedResourceRequest {
+  // Required. The name of the queued resource.
+  string name = 1 [
+    (google.api.field_behavior) = REQUIRED,
+    (google.api.resource_reference) = {
+      type: "tpu.googleapis.com/QueuedResource"
+    }
+  ];
+}
+
 // The per-product per-project service identity for Cloud TPU service.
 message ServiceIdentity {
   // The email address of the service identity.
