Merge branch 'main' into dsuhinin/move-get-model-version-endpoint

Signed-off-by: Software Developer <[email protected]>

magefiles/generate/endpoints.go

@@ -56,7 +56,7 @@ var ServiceInfoMap = map[string]ServiceGenerationInfo{
 			// "searchRegisteredModels",
 			"getLatestVersions",
 			// "createModelVersion",
-			// "updateModelVersion",
+			"updateModelVersion",
 			// "transitionModelVersionStage",
 			"deleteModelVersion",
 			"getModelVersion",

mlflow_go/store/model_registry.py

@@ -9,6 +9,7 @@
     GetModelVersion,
     GetRegisteredModel,
     RenameRegisteredModel,
+    UpdateModelVersion,
     UpdateRegisteredModel,
 )
 
@@ -92,6 +93,10 @@ def get_model_version(self, name, version):
         )
         return ModelVersion.from_proto(response.model_version)
 
+      def update_model_version(self, name, version, description=None):
+        request = UpdateModelVersion(name=name, version=str(version), description=description)
+        self.service.call_endpoint(get_lib().ModelRegistryServiceUpdateModelVersion, request)
+
 
 def ModelRegistryStore(cls):
     return type(cls.__name__, (_ModelRegistryStore, cls), {})

pkg/model_registry/service/model_versions.go

@@ -101,11 +101,20 @@ func (m *ModelRegistryService) GetModelVersion(
 	}
 
 	modelVersion, err := m.store.GetModelVersion(ctx, input.GetName(), version, true)
+  return &protos.GetModelVersion_Response{
+		ModelVersion: modelVersion.ToProto(),
+	}, nil
+}
+
+  func (m *ModelRegistryService) UpdateModelVersion(
+	ctx context.Context, input *protos.UpdateModelVersion,
+) (*protos.UpdateModelVersion_Response, *contract.Error) {
+	modelVersion, err := m.store.UpdateModelVersion(ctx, input.GetName(), input.GetVersion(), input.GetDescription())
 	if err != nil {
 		return nil, err
 	}
 
-	return &protos.GetModelVersion_Response{
+	return &protos.UpdateModelVersion_Response{
 		ModelVersion: modelVersion.ToProto(),
 	}, nil
 }

pkg/model_registry/store/sql/model_versions.go

@@ -378,3 +378,27 @@ func (m *ModelRegistrySQLStore) DeleteModelVersion(ctx context.Context, name, ve
 
 	return nil
 }
+
+func (m *ModelRegistrySQLStore) UpdateModelVersion(
+	ctx context.Context, name, version, description string,
+) (*entities.ModelVersion, *contract.Error) {
+	modelVersion, err := m.GetModelVersion(ctx, name, version)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := m.db.WithContext(ctx).Model(
+		&models.ModelVersion{},
+	).Where(
+		"name = ?", modelVersion.Name,
+	).Where(
+		"version = ?", modelVersion.Version,
+	).Updates(&models.ModelVersion{
+		Description:     sql.NullString{String: description, Valid: description != ""},
+		LastUpdatedTime: time.Now().UnixMilli(),
+	}).Error; err != nil {
+		return nil, contract.NewErrorWith(protos.ErrorCode_INTERNAL_ERROR, "error updating model version", err)
+	}
+
+	return modelVersion, nil
+}

pkg/model_registry/store/store.go

@@ -17,4 +17,5 @@ type ModelRegistryStore interface {
 	DeleteRegisteredModel(ctx context.Context, name string) *contract.Error
 	GetModelVersion(ctx context.Context, name, version string, eager bool) (*entities.ModelVersion, *contract.Error)
 	DeleteModelVersion(ctx context.Context, name, version string) *contract.Error
+	UpdateModelVersion(ctx context.Context, name, version, description string) (*entities.ModelVersion, *contract.Error)
 }