feat: add NVIDIA_H100_80GB and TPU_V5_LITEPOD to AcceleratorType

feat: update field_behavior for `name` to be IMMUTABLE instead of OUTPUT_ONLY in Context, ModelMonitor, Schedule, DeploymentResourcePool
feat: Expose CreateDatasetVersionOperationMetadata and RestoreDatasetVersionOperationMetadata to DatasetService
feat: add Timestamp to Feature
feat: add disk_type and disk_size_gb to PersistentDiskSpec
feat: add schedule_name to PipelineJob
feat: add CountTokens to PredictionService
feat: add Parent to PublisherModel
feat: add open_evaluation_pipeline to PublisherModel
fix!: remove idle_replica_count from PersistentResource

PiperOrigin-RevId: 557186520

google/cloud/aiplatform/v1beta1/accelerator_type.proto

@@ -53,6 +53,9 @@ enum AcceleratorType {
   // Nvidia L4 GPU.
   NVIDIA_L4 = 11;
 
+  // Nvidia H100 80Gb GPU.
+  NVIDIA_H100_80GB = 13;
+
   // TPU v2.
   TPU_V2 = 6;
 
@@ -61,4 +64,7 @@ enum AcceleratorType {
 
   // TPU v4.
   TPU_V4_POD = 10;
+
+  // TPU v5.
+  TPU_V5_LITEPOD = 12;
 }

google/cloud/aiplatform/v1beta1/context.proto

@@ -36,8 +36,8 @@ message Context {
     pattern: "projects/{project}/locations/{location}/metadataStores/{metadata_store}/contexts/{context}"
   };
 
-  // Output only. The resource name of the Context.
-  string name = 1 [(google.api.field_behavior) = OUTPUT_ONLY];
+  // Immutable. The resource name of the Context.
+  string name = 1 [(google.api.field_behavior) = IMMUTABLE];
 
   // User provided display name of the Context.
   // May be up to 128 Unicode characters.

google/cloud/aiplatform/v1beta1/custom_job.proto

@@ -332,4 +332,9 @@ message Scheduling {
   // This feature can be used by distributed training jobs that are not
   // resilient to workers leaving and joining a job.
   bool restart_job_on_worker_restart = 3;
+
+  // Optional. Indicates if the job should retry for internal errors after the
+  // job starts running. If true, overrides
+  // `Scheduling.restart_job_on_worker_restart` to false.
+  bool disable_retries = 5 [(google.api.field_behavior) = OPTIONAL];
 }

google/cloud/aiplatform/v1beta1/dataset_service.proto

@@ -369,6 +369,20 @@ message ExportDataOperationMetadata {
   string gcs_output_directory = 2;
 }
 
+// Runtime operation information for
+// [DatasetService.CreateDatasetVersion][google.cloud.aiplatform.v1beta1.DatasetService.CreateDatasetVersion].
+message CreateDatasetVersionOperationMetadata {
+  // The common part of the operation metadata.
+  GenericOperationMetadata generic_metadata = 1;
+}
+
+// Runtime operation information for
+// [DatasetService.RestoreDatasetVersion][google.cloud.aiplatform.v1beta1.DatasetService.RestoreDatasetVersion].
+message RestoreDatasetVersionOperationMetadata {
+  // The common part of the operation metadata.
+  GenericOperationMetadata generic_metadata = 1;
+}
+
 // Request message for
 // [DatasetService.ListDataItems][google.cloud.aiplatform.v1beta1.DatasetService.ListDataItems].
 message ListDataItemsRequest {

google/cloud/aiplatform/v1beta1/deployment_resource_pool.proto

@@ -37,10 +37,10 @@ message DeploymentResourcePool {
     pattern: "projects/{project}/locations/{location}/deploymentResourcePools/{deployment_resource_pool}"
   };
 
-  // Output only. The resource name of the DeploymentResourcePool.
+  // Immutable. The resource name of the DeploymentResourcePool.
   // Format:
   // `projects/{project}/locations/{location}/deploymentResourcePools/{deployment_resource_pool}`
-  string name = 1 [(google.api.field_behavior) = OUTPUT_ONLY];
+  string name = 1 [(google.api.field_behavior) = IMMUTABLE];
 
   // Required. The underlying DedicatedResources that the DeploymentResourcePool
   // uses.

google/cloud/aiplatform/v1beta1/machine_resources.proto

@@ -181,6 +181,20 @@ message DiskSpec {
   int32 boot_disk_size_gb = 2;
 }
 
+// Represents the spec of [persistent
+// disk][https://cloud.google.com/compute/docs/disks/persistent-disks] options.
+message PersistentDiskSpec {
+  // Type of the disk (default is "pd-standard").
+  // Valid values: "pd-ssd" (Persistent Disk Solid State Drive)
+  // "pd-standard" (Persistent Disk Hard Disk Drive)
+  // "pd-balanced" (Balanced Persistent Disk)
+  // "pd-extreme" (Extreme Persistent Disk)
+  string disk_type = 1;
+
+  // Size in GB of the disk (default is 100GB).
+  int64 disk_size_gb = 2;
+}
+
 // Represents a mount configuration for Network File System (NFS) to mount.
 message NfsMount {
   // Required. IP address of the NFS server.

google/cloud/aiplatform/v1beta1/persistent_resource.proto

@@ -166,10 +166,10 @@ message ResourcePool {
         [(google.api.field_behavior) = OPTIONAL];
   }
 
-  // Optional. The unique ID in a PersistentResource to refer the this resource
+  // Immutable. The unique ID in a PersistentResource to refer the this resource
   // pool. User can specify it if need to use it, otherwise we will generate it
   // automatically.
-  string id = 1 [(google.api.field_behavior) = OPTIONAL];
+  string id = 1 [(google.api.field_behavior) = IMMUTABLE];
 
   // Required. Immutable. The specification of a single machine.
   MachineSpec machine_spec = 2 [
@@ -183,11 +183,6 @@ message ResourcePool {
   // Optional. Disk spec for the machine in this node pool.
   DiskSpec disk_spec = 4 [(google.api.field_behavior) = OPTIONAL];
 
-  // Output only. The number of machines currently not in use by training jobs
-  // for this resource pool. Deprecated. Use `used_replica_count` instead.
-  int64 idle_replica_count = 5
-      [deprecated = true, (google.api.field_behavior) = OUTPUT_ONLY];
-
   // Output only. The number of machines currently in use by training jobs for
   // this resource pool. Will replace idle_replica_count.
   int64 used_replica_count = 6 [(google.api.field_behavior) = OUTPUT_ONLY];
@@ -205,7 +200,7 @@ message ResourceRuntimeSpec {
   ServiceAccountSpec service_account_spec = 2
       [(google.api.field_behavior) = OPTIONAL];
 
-  // Ray cluster configuration.
+  // Optional. Ray cluster configuration.
   // Required when creating a dedicated RayCluster on the PersistentResource.
   RaySpec ray_spec = 1 [(google.api.field_behavior) = OPTIONAL];
 }

google/cloud/aiplatform/v1beta1/persistent_resource_service.proto

@@ -39,7 +39,7 @@ service PersistentResourceService {
   option (google.api.oauth_scopes) =
       "https://www.googleapis.com/auth/cloud-platform";
 
-  // Uploads a Model artifact into Vertex AI.
+  // Creates a PersistentResource.
   rpc CreatePersistentResource(CreatePersistentResourceRequest)
       returns (google.longrunning.Operation) {
     option (google.api.http) = {

google/cloud/aiplatform/v1beta1/pipeline_job.proto

@@ -207,6 +207,10 @@ message PipelineJob {
   // is from supported template registry.
   PipelineTemplateMetadata template_metadata = 20
       [(google.api.field_behavior) = OUTPUT_ONLY];
+
+  // Output only. The schedule resource name.
+  // Only returned if the Pipeline is created by Schedule API.
+  string schedule_name = 22 [(google.api.field_behavior) = OUTPUT_ONLY];
 }
 
 // Pipeline template metadata if

google/cloud/aiplatform/v1beta1/prediction_service.proto

@@ -108,6 +108,19 @@ service PredictionService {
     option (google.api.method_signature) =
         "endpoint,instances,parameters,deployed_model_id";
   }
+
+  // Perform a token counting.
+  rpc CountTokens(CountTokensRequest) returns (CountTokensResponse) {
+    option (google.api.http) = {
+      post: "/v1beta1/{endpoint=projects/*/locations/*/endpoints/*}:countTokens"
+      body: "*"
+      additional_bindings {
+        post: "/v1beta1/{endpoint=projects/*/locations/*/publishers/*/models/*}:countTokens"
+        body: "*"
+      }
+    };
+    option (google.api.method_signature) = "endpoint,instances";
+  }
 }
 
 // Request message for
@@ -316,3 +329,33 @@ message ExplainResponse {
   // [PredictResponse.predictions][google.cloud.aiplatform.v1beta1.PredictResponse.predictions].
   repeated google.protobuf.Value predictions = 3;
 }
+
+// Request message for
+// [PredictionService.CountTokens][google.cloud.aiplatform.v1beta1.PredictionService.CountTokens].
+message CountTokensRequest {
+  // Required. The name of the Endpoint requested to perform token counting.
+  // Format:
+  // `projects/{project}/locations/{location}/endpoints/{endpoint}`
+  string endpoint = 1 [
+    (google.api.field_behavior) = REQUIRED,
+    (google.api.resource_reference) = {
+      type: "aiplatform.googleapis.com/Endpoint"
+    }
+  ];
+
+  // Required. The instances that are the input to token counting call.
+  // Schema is identical to the prediction schema of the underlying model.
+  repeated google.protobuf.Value instances = 2
+      [(google.api.field_behavior) = REQUIRED];
+}
+
+// Response message for
+// [PredictionService.CountTokens][google.cloud.aiplatform.v1beta1.PredictionService.CountTokens].
+message CountTokensResponse {
+  // The total number of tokens counted across all instances from the request.
+  int32 total_tokens = 1;
+
+  // The total number of billable characters counted across all instances from
+  // the request.
+  int32 total_billable_characters = 2;
+}

google/cloud/aiplatform/v1beta1/publisher_model.proto

@@ -47,6 +47,16 @@ message PublisherModel {
     }
   }
 
+  // The information about the parent of a model.
+  message Parent {
+    // Required. The display name of the parent. E.g., LaMDA, T5, Vision API,
+    // Natural Language API.
+    string display_name = 1 [(google.api.field_behavior) = REQUIRED];
+
+    // Optional. The Google Cloud resource name or the URI reference.
+    ResourceReference reference = 2 [(google.api.field_behavior) = OPTIONAL];
+  }
+
   // A named piece of documentation.
   message Documentation {
     // Required. E.g., OVERVIEW, USE CASES, DOCUMENTATION, SDK & SAMPLES, JAVA,
@@ -154,6 +164,10 @@ message PublisherModel {
     // Optional. Request for access.
     RegionalResourceReferences request_access = 9
         [(google.api.field_behavior) = OPTIONAL];
+
+    // Optional. Open evaluation pipeline of the PublisherModel.
+    RegionalResourceReferences open_evaluation_pipeline = 11
+        [(google.api.field_behavior) = OPTIONAL];
   }
 
   // An enum representing the open source category of a PublisherModel.
@@ -214,6 +228,10 @@ message PublisherModel {
   OpenSourceCategory open_source_category = 7
       [(google.api.field_behavior) = REQUIRED];
 
+  // Optional. The parent that this model was customized from. E.g., Vision API,
+  // Natural Language API, LaMDA, T5, etc. Foundation models don't have parents.
+  Parent parent = 14 [(google.api.field_behavior) = OPTIONAL];
+
   // Optional. Supported call-to-action options.
   CallToAction supported_actions = 19 [(google.api.field_behavior) = OPTIONAL];
 

google/cloud/aiplatform/v1beta1/schedule.proto

@@ -88,8 +88,8 @@ message Schedule {
     CreatePipelineJobRequest create_pipeline_job_request = 14;
   }
 
-  // Output only. The resource name of the Schedule.
-  string name = 1 [(google.api.field_behavior) = OUTPUT_ONLY];
+  // Immutable. The resource name of the Schedule.
+  string name = 1 [(google.api.field_behavior) = IMMUTABLE];
 
   // Required. User provided name of the Schedule.
   // The name can be up to 128 characters long and can consist of any UTF-8

google/cloud/aiplatform/v1beta1/schedule_service.proto

@@ -222,6 +222,7 @@ message ListSchedulesRequest {
   // descending order.
   //
   // Supported fields:
+  //
   //   * `create_time`
   //   * `start_time`
   //   * `end_time`
@@ -296,7 +297,9 @@ message ResumeScheduleRequest {
 message UpdateScheduleRequest {
   // Required. The Schedule which replaces the resource on the server.
   // The following restrictions will be applied:
+  //
   //   * The scheduled request type cannot be changed.
+  //   * The non-empty fields cannot be unset.
   //   * The output_only fields will be ignored if specified.
   Schedule schedule = 1 [(google.api.field_behavior) = REQUIRED];