feat: Model Registry Checkpoint API

Google APIs · copybara-github · commit 2c21b3abc001 · 2025-01-27T15:32:49.000-08:00
PiperOrigin-RevId: 720330032
diff --git a/google/cloud/aiplatform/v1beta1/model.proto b/google/cloud/aiplatform/v1beta1/model.proto
@@ -180,6 +180,9 @@ message Model {
   // The description of this version.
   string version_description = 30;
 
+  // The default checkpoint id of a model version.
+  string default_checkpoint_id = 53;
+
   // The schemata that describe formats of the Model's predictions and
   // explanations as given and returned via
   // [PredictionService.Predict][google.cloud.aiplatform.v1beta1.PredictionService.Predict]
diff --git a/google/cloud/aiplatform/v1beta1/model_service.proto b/google/cloud/aiplatform/v1beta1/model_service.proto
@@ -84,6 +84,15 @@ service ModelService {
     option (google.api.method_signature) = "name";
   }
 
+  // Lists checkpoints of the specified model version.
+  rpc ListModelVersionCheckpoints(ListModelVersionCheckpointsRequest)
+      returns (ListModelVersionCheckpointsResponse) {
+    option (google.api.http) = {
+      get: "/v1beta1/{name=projects/*/locations/*/models/*}:listCheckpoints"
+    };
+    option (google.api.method_signature) = "name";
+  }
+
   // Updates a Model.
   rpc UpdateModel(UpdateModelRequest) returns (Model) {
     option (google.api.http) = {
@@ -460,6 +469,61 @@ message ListModelVersionsResponse {
   string next_page_token = 2;
 }
 
+// Request message for
+// [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints].
+message ListModelVersionCheckpointsRequest {
+  // Required. The name of the model version to list checkpoints for.
+  // `projects/{project}/locations/{location}/models/{model}@{version}`
+  // Example: `projects/{project}/locations/{location}/models/{model}@2`
+  // or
+  // `projects/{project}/locations/{location}/models/{model}@golden`
+  // If no version ID or alias is specified, the latest version will be
+  // used.
+  string name = 1 [
+    (google.api.field_behavior) = REQUIRED,
+    (google.api.resource_reference) = {
+      type: "aiplatform.googleapis.com/Model"
+    }
+  ];
+
+  // Optional. The standard list page size.
+  int32 page_size = 2 [(google.api.field_behavior) = OPTIONAL];
+
+  // Optional. The standard list page token.
+  // Typically obtained via
+  // [next_page_token][google.cloud.aiplatform.v1beta1.ListModelVersionCheckpointsResponse.next_page_token]
+  // of the previous
+  // [ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints]
+  // call.
+  string page_token = 3 [(google.api.field_behavior) = OPTIONAL];
+}
+
+// A proto representation of a Spanner-stored ModelVersionCheckpoint.
+// The meaning of the fields is equivalent to their in-Spanner counterparts.
+message ModelVersionCheckpoint {
+  // The ID of the checkpoint.
+  string checkpoint_id = 1;
+
+  // The epoch of the checkpoint.
+  int64 epoch = 2;
+
+  // The step of the checkpoint.
+  int64 step = 3;
+}
+
+// Response message for
+// [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints]
+message ListModelVersionCheckpointsResponse {
+  // List of Model Version checkpoints.
+  repeated ModelVersionCheckpoint checkpoints = 1;
+
+  // A token to retrieve the next page of results.
+  // Pass to
+  // [ListModelVersionCheckpointsRequest.page_token][google.cloud.aiplatform.v1beta1.ListModelVersionCheckpointsRequest.page_token]
+  // to obtain that page.
+  string next_page_token = 2;
+}
+
 // Request message for
 // [ModelService.UpdateModel][google.cloud.aiplatform.v1beta1.ModelService.UpdateModel].
 message UpdateModelRequest {
