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

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

magefiles/generate/endpoints.go

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

mlflow_go/store/model_registry.py

@@ -9,6 +9,7 @@
     GetModelVersion,
     GetRegisteredModel,
     RenameRegisteredModel,
+    TransitionModelVersionStage,
     UpdateModelVersion,
     UpdateRegisteredModel,
 )
@@ -100,6 +101,17 @@ 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 transition_model_version_stage(self, name, version, stage, archive_existing_versions):
+        request = TransitionModelVersionStage(
+            name=name,
+            version=str(version),
+            stage=stage,
+            archive_existing_versions=archive_existing_versions,
+        )
+        self.service.call_endpoint(
+            get_lib().ModelRegistryServiceTransitionModelVersionStage, request
+        )
+
 
 def ModelRegistryStore(cls):
     return type(cls.__name__, (_ModelRegistryStore, cls), {})

pkg/model_registry/service/model_versions.go

@@ -3,8 +3,11 @@ package service
 import (
 	"context"
 	"strconv"
+	"fmt"
+	"strings"
 
 	"github.com/mlflow/mlflow-go/pkg/contract"
+	"github.com/mlflow/mlflow-go/pkg/model_registry/store/sql/models"
 	"github.com/mlflow/mlflow-go/pkg/protos"
 )
 
@@ -122,3 +125,36 @@ func (m *ModelRegistryService) UpdateModelVersion(
 		ModelVersion: modelVersion.ToProto(),
 	}, nil
 }
+
+func (m *ModelRegistryService) TransitionModelVersionStage(
+	ctx context.Context, input *protos.TransitionModelVersionStage,
+) (*protos.TransitionModelVersionStage_Response, *contract.Error) {
+	stage, ok := models.CanonicalMapping[strings.ToLower(input.GetStage())]
+	if !ok {
+		return nil, contract.NewError(
+			protos.ErrorCode_INVALID_PARAMETER_VALUE,
+			fmt.Sprintf(
+				"Invalid Model Version stage: unknown. Value must be one of %s, %s, %s, %s.",
+				models.ModelVersionStageNone,
+				models.ModelVersionStageStaging,
+				models.ModelVersionStageProduction,
+				models.ModelVersionStageArchived,
+			),
+		)
+	}
+
+	modelVersion, err := m.store.TransitionModelVersionStage(
+		ctx,
+		input.GetName(),
+		input.GetVersion(),
+		stage,
+		input.GetArchiveExistingVersions(),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return &protos.TransitionModelVersionStage_Response{
+		ModelVersion: modelVersion.ToProto(),
+	}, nil
+}

pkg/model_registry/store/sql/model_versions.go

@@ -57,7 +57,7 @@ func (m *ModelRegistrySQLStore) GetLatestVersions(
 		for idx, stage := range stages {
 			stages[idx] = strings.ToLower(stage)
 			if canonicalStage, ok := models.CanonicalMapping[stages[idx]]; ok {
-				stages[idx] = canonicalStage
+				stages[idx] = canonicalStage.String()
 
 				continue
 			}
@@ -402,3 +402,83 @@ func (m *ModelRegistrySQLStore) UpdateModelVersion(
 
 	return modelVersion, nil
 }
+
+//nolint:funlen,cyclop
+func (m *ModelRegistrySQLStore) TransitionModelVersionStage(
+	ctx context.Context, name, version string, stage models.ModelVersionStage, archiveExistingVersions bool,
+) (*entities.ModelVersion, *contract.Error) {
+	isActiveStage := false
+	if _, ok := models.DefaultStagesForGetLatestVersions[strings.ToLower(stage.String())]; ok {
+		isActiveStage = true
+	}
+
+	if archiveExistingVersions && !isActiveStage {
+		return nil, contract.NewError(
+			protos.ErrorCode_INVALID_PARAMETER_VALUE,
+			fmt.Sprintf(
+				`Model version transition cannot archive existing model versions because '%s' is not an Active stage. 
+Valid stages are %s`,
+				stage, models.AllModelVersionStages(),
+			),
+		)
+	}
+
+	modelVersion, err := m.GetModelVersion(ctx, name, version)
+	if err != nil {
+		return nil, err
+	}
+
+	registeredModel, err := m.GetRegisteredModel(ctx, name)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := m.db.Transaction(func(transaction *gorm.DB) error {
+		lastUpdatedTime := time.Now().UnixMilli()
+		if err := transaction.Model(
+			&models.RegisteredModel{},
+		).Where(
+			"name = ?", registeredModel.Name,
+		).Updates(&models.RegisteredModel{
+			LastUpdatedTime: lastUpdatedTime,
+		}).Error; err != nil {
+			return err
+		}
+
+		if err := transaction.Model(
+			&models.ModelVersion{},
+		).Where(
+			"name = ?", modelVersion.Name,
+		).Where(
+			"version = ?", modelVersion.Version,
+		).Updates(&models.ModelVersion{
+			CurrentStage:    stage,
+			LastUpdatedTime: lastUpdatedTime,
+		}).Error; err != nil {
+			return err
+		}
+
+		if archiveExistingVersions {
+			if err := transaction.Where(
+				"name = ?", name,
+			).Where(
+				"version != ?", version,
+			).Where(
+				"current_stage = ?", stage,
+			).Updates(&models.ModelVersion{
+				CurrentStage:    models.ModelVersionStageArchived,
+				LastUpdatedTime: lastUpdatedTime,
+			}).Error; err != nil {
+				return err
+			}
+		}
+
+		return nil
+	}); err != nil {
+		return nil, contract.NewErrorWith(
+			protos.ErrorCode_INTERNAL_ERROR, "error transitioning model version stage", err,
+		)
+	}
+
+	return modelVersion, nil
+}

pkg/model_registry/store/sql/models/model_version_stage.go

@@ -15,18 +15,23 @@ const (
 	ModelVersionStageArchived   = "Archived"
 )
 
-var CanonicalMapping = map[string]string{
+var CanonicalMapping = map[string]ModelVersionStage{
 	strings.ToLower(ModelVersionStageNone):       ModelVersionStageNone,
 	strings.ToLower(ModelVersionStageStaging):    ModelVersionStageStaging,
 	strings.ToLower(ModelVersionStageProduction): ModelVersionStageProduction,
 	strings.ToLower(ModelVersionStageArchived):   ModelVersionStageArchived,
 }
 
+var DefaultStagesForGetLatestVersions = map[string]ModelVersionStage{
+	strings.ToLower(ModelVersionStageStaging):    ModelVersionStageStaging,
+	strings.ToLower(ModelVersionStageProduction): ModelVersionStageProduction,
+}
+
 func AllModelVersionStages() string {
 	pairs := make([]string, 0, len(CanonicalMapping))
 
 	for _, v := range CanonicalMapping {
-		pairs = append(pairs, v)
+		pairs = append(pairs, v.String())
 	}
 
 	return strings.Join(pairs, ",")

pkg/model_registry/store/store.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/mlflow/mlflow-go/pkg/contract"
 	"github.com/mlflow/mlflow-go/pkg/entities"
+	"github.com/mlflow/mlflow-go/pkg/model_registry/store/sql/models"
 	"github.com/mlflow/mlflow-go/pkg/protos"
 )
 
@@ -18,4 +19,7 @@ type ModelRegistryStore interface {
 	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)
+	TransitionModelVersionStage(
+		ctx context.Context, name, version string, stage models.ModelVersionStage, archiveExistingVersions bool,
+	) (*entities.ModelVersion, *contract.Error)
 }