Move POST /mlflow/model-versions/create endpoint.

magefiles/generate/endpoints.go

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

magefiles/generate/validations.go

@@ -73,4 +73,7 @@ var validations = map[string]string{
 	"SetModelVersionTag_Version":         "stringAsInteger",
 	"GetModelVersion_Version":            "stringAsInteger",
 	"GetModelVersionDownloadUri_Version": "stringAsInteger",
+	"CreateModelVersion_Name":            "notEmpty,required",
+	"ModelVersionTag_Key":                "required,max=250,validMetricParamOrTagName,pathIsUnique",
+	"ModelVersionTag_Value":              "omitempty,max=5000,truncate=5000",
 }

magefiles/tests.go

@@ -34,6 +34,7 @@ func runPythonTests(pytestArgs []string) error {
 		// "--log-cli-level=DEBUG",
 		"--confcutdir=.",
 		"-k", "not [file",
+		"-vv",
 	}
 	args = append(args, pytestArgs...)
 

mlflow_go_backend/store/model_registry.py

@@ -3,6 +3,7 @@
 
 from mlflow.entities.model_registry import ModelVersion, RegisteredModel
 from mlflow.protos.model_registry_pb2 import (
+    CreateModelVersion,
     CreateRegisteredModel,
     DeleteModelVersion,
     DeleteModelVersionTag,
@@ -176,6 +177,30 @@ def get_model_version_download_uri(self, name, version):
         )
         return response.artifact_uri
 
+    def create_model_version(
+        self,
+        name,
+        source,
+        run_id=None,
+        tags=None,
+        run_link=None,
+        description=None,
+        local_model_path=None,
+    ):
+        request = CreateModelVersion(
+            name=name,
+            source=source,
+            run_id=run_id,
+            tags=[tag.to_proto() for tag in tags] if tags else [],
+            run_link=run_link,
+            description=description,
+        )
+        response = self.service.call_endpoint(
+            get_lib().ModelRegistryServiceCreateModelVersion, request
+        )
+
+        return ModelVersion.from_proto(response.model_version)
+
 
 def ModelRegistryStore(cls):
     return type(cls.__name__, (_ModelRegistryStore, cls), {})

pkg/entities/model_version.go

@@ -28,16 +28,19 @@ type ModelVersion struct {
 func (mv ModelVersion) ToProto() *protos.ModelVersion {
 	modelVersion := protos.ModelVersion{
 		Name:                 utils.PtrTo(mv.Name),
+		Tags:                 make([]*protos.ModelVersionTag, 0, len(mv.Tags)),
+		Source:               utils.PtrTo(mv.Source),
+		Status:               utils.PtrTo(protos.ModelVersionStatus(protos.ModelVersionStatus_value[mv.Status])),
+		RunLink:              utils.PtrTo(mv.RunLink),
 		Version:              utils.PtrTo(strconv.Itoa(int(mv.Version))),
+		Description:          utils.PtrTo(mv.Description),
 		CurrentStage:         utils.PtrTo(mv.CurrentStage),
 		CreationTimestamp:    utils.PtrTo(mv.CreationTime),
 		LastUpdatedTimestamp: utils.PtrTo(mv.LastUpdatedTime),
-		Description:          utils.PtrTo(mv.Description),
-		UserId:               utils.PtrTo(mv.UserID),
-		Source:               utils.PtrTo(mv.Source),
-		Status:               utils.PtrTo(protos.ModelVersionStatus(protos.ModelVersionStatus_value[mv.Status])),
-		Tags:                 make([]*protos.ModelVersionTag, 0, len(mv.Tags)),
-		RunLink:              utils.PtrTo(mv.RunLink),
+	}
+
+	if mv.UserID != "" {
+		modelVersion.UserId = utils.PtrTo(mv.UserID)
 	}
 
 	if mv.RunID != "" {

pkg/entities/model_version_tag.go

@@ -16,3 +16,10 @@ func (mvt ModelVersionTag) ToProto() *protos.ModelVersionTag {
 		Value: utils.PtrTo(mvt.Value),
 	}
 }
+
+func NewModelVersionTag(proto *protos.ModelVersionTag) *ModelVersionTag {
+	return &ModelVersionTag{
+		Key:   proto.GetKey(),
+		Value: proto.GetValue(),
+	}
+}

pkg/model_registry/service/model_versions.go

@@ -3,10 +3,10 @@ package service
 import (
 	"context"
 	"fmt"
-	"strconv"
 	"strings"
 
 	"github.com/mlflow/mlflow-go-backend/pkg/contract"
+	"github.com/mlflow/mlflow-go-backend/pkg/entities"
 	"github.com/mlflow/mlflow-go-backend/pkg/model_registry/store/sql/models"
 	"github.com/mlflow/mlflow-go-backend/pkg/protos"
 	"github.com/mlflow/mlflow-go-backend/pkg/utils"
@@ -38,16 +38,7 @@ func (m *ModelRegistryService) DeleteModelVersion(
 func (m *ModelRegistryService) GetModelVersion(
 	ctx context.Context, input *protos.GetModelVersion,
 ) (*protos.GetModelVersion_Response, *contract.Error) {
-	// by some strange reason GetModelVersion.Version has a string type so we can't apply our validation,
-	// that's why such a custom validation exists to satisfy Python tests.
-	version := input.GetVersion()
-	if _, err := strconv.Atoi(version); err != nil {
-		return nil, contract.NewErrorWith(
-			protos.ErrorCode_INVALID_PARAMETER_VALUE, "Model version must be an integer", err,
-		)
-	}
-
-	modelVersion, err := m.store.GetModelVersion(ctx, input.GetName(), version, true)
+	modelVersion, err := m.store.GetModelVersion(ctx, input.GetName(), input.GetVersion(), true)
 	if err != nil {
 		return nil, err
 	}
@@ -157,3 +148,29 @@ func (m *ModelRegistryService) GetModelVersionDownloadUri(
 		ArtifactUri: utils.PtrTo(artifactURI),
 	}, nil
 }
+
+func (m *ModelRegistryService) CreateModelVersion(
+	ctx context.Context, input *protos.CreateModelVersion,
+) (*protos.CreateModelVersion_Response, *contract.Error) {
+	tags := make([]*entities.ModelVersionTag, 0, len(input.GetTags()))
+	for _, tag := range input.GetTags() {
+		tags = append(tags, entities.NewModelVersionTag(tag))
+	}
+
+	modelVersion, err := m.store.CreateModelVersion(
+		ctx,
+		input.GetName(),
+		input.GetSource(),
+		input.GetRunId(),
+		tags,
+		input.GetRunLink(),
+		input.GetDescription(),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return &protos.CreateModelVersion_Response{
+		ModelVersion: modelVersion.ToProto(),
+	}, nil
+}

pkg/model_registry/store/sql/helpers.go

@@ -0,0 +1,81 @@
+package sql
+
+import (
+	"errors"
+	"net/url"
+	"strconv"
+	"strings"
+
+	"github.com/mlflow/mlflow-go-backend/pkg/entities"
+)
+
+// GetNextVersion returns the next version number for a given registered model.
+func GetNextVersion(sqlRegisteredModel *entities.RegisteredModel) int32 {
+	if len(sqlRegisteredModel.Versions) > 0 {
+		maxVersion := sqlRegisteredModel.Versions[0].Version
+		for _, mv := range sqlRegisteredModel.Versions {
+			if mv.Version > maxVersion {
+				maxVersion = mv.Version
+			}
+		}
+
+		return maxVersion + 1
+	}
+
+	return 1
+}
+
+type ParsedModelURI struct {
+	Name    string
+	Version string
+	Stage   string
+	Alias   string
+}
+
+//nolint:cyclop,err113,mnd,wrapcheck
+func ParseModelURI(uri string) (*ParsedModelURI, error) {
+	parsed, err := url.Parse(uri)
+	if err != nil {
+		return nil, err
+	}
+
+	if parsed.Scheme != "models" {
+		return nil, errors.New("invalid model URI scheme")
+	}
+
+	path := strings.TrimLeft(parsed.Path, "/")
+	if path == "" {
+		return nil, errors.New("invalid model URI path")
+	}
+
+	parts := strings.Split(path, "/")
+	if len(parts) > 2 {
+		return nil, errors.New("invalid model URI format")
+	}
+
+	if len(parts) == 2 {
+		name, suffix := parts[0], parts[1]
+		if suffix == "" {
+			return nil, errors.New("invalid model URI suffix")
+		}
+
+		if _, err := strconv.Atoi(suffix); err == nil {
+			// The suffix is a specific version
+			return &ParsedModelURI{Name: name, Version: suffix}, nil
+		} else if strings.EqualFold(suffix, "latest") {
+			// The suffix is "latest"
+			return &ParsedModelURI{Name: name}, nil
+		}
+
+		// The suffix is a specific stage
+		return &ParsedModelURI{Name: name, Stage: suffix}, nil
+	}
+
+	// The URI is an alias URI
+	aliasParts := strings.SplitN(parts[0], "@", 2)
+	if len(aliasParts) != 2 || aliasParts[1] == "" {
+		return nil, errors.New("invalid model alias format")
+	}
+
+	return &ParsedModelURI{Name: aliasParts[0], Alias: aliasParts[1]}, nil
+}

pkg/model_registry/store/sql/model_versions.go

@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
+	"net/url"
 	"strconv"
 	"strings"
 	"time"
@@ -18,6 +19,10 @@ import (
 	"github.com/mlflow/mlflow-go-backend/pkg/protos"
 )
 
+const (
+	CreateModelVersionRetries = 3
+)
+
 func (m *ModelRegistrySQLStore) GetLatestVersions(
 	ctx context.Context, name string, stages []string,
 ) ([]*protos.ModelVersion, *contract.Error) {
@@ -78,6 +83,104 @@ func (m *ModelRegistrySQLStore) GetLatestVersions(
 	return results, nil
 }
 
+//nolint:funlen,cyclop,staticcheck
+func (m *ModelRegistrySQLStore) CreateModelVersion(
+	ctx context.Context, name, source, runID string, tags []*entities.ModelVersionTag, runLink, description string,
+) (*entities.ModelVersion, *contract.Error) {
+	storageLocation := source
+
+	parsedSource, parsedSourceErr := url.Parse(source)
+	if parsedSourceErr != nil {
+		return nil, contract.NewErrorWith(
+			protos.ErrorCode_INTERNAL_ERROR,
+			fmt.Sprintf("failed to parse source %q", source),
+			parsedSourceErr,
+		)
+	}
+
+	if parsedSource.Scheme == "models" {
+		parsedModelURI, err := ParseModelURI(source)
+		if err != nil {
+			return nil, contract.NewErrorWith(
+				protos.ErrorCode_INTERNAL_ERROR,
+				fmt.Sprintf("Unable to fetch model from model URI source artifact location '%s'.", source),
+				parsedSourceErr,
+			)
+		}
+
+		downloadURI, contractErr := m.GetModelVersionDownloadURI(ctx, parsedModelURI.Name, parsedModelURI.Version)
+		if contractErr != nil {
+			return nil, contractErr
+		}
+
+		storageLocation = downloadURI
+	}
+
+	registeredModel, err := m.GetRegisteredModel(ctx, name)
+	if err != nil {
+		return nil, err
+	}
+
+	creationTime := time.Now().UnixMilli()
+	for range CreateModelVersionRetries {
+		registeredModel.LastUpdatedTime = creationTime
+		nextVersion := GetNextVersion(registeredModel)
+		modelVersion := &models.ModelVersion{
+			Name:            name,
+			Version:         nextVersion,
+			CreationTime:    creationTime,
+			LastUpdatedTime: creationTime,
+			Description:     sql.NullString{Valid: description != "", String: description},
+			Status:          models.ModelVersionStatusReady,
+			Source:          source,
+			RunID:           runID,
+			RunLink:         runLink,
+			CurrentStage:    models.ModelVersionStageNone,
+			StorageLocation: storageLocation,
+		}
+
+		modelTags := make([]models.ModelVersionTag, 0, len(tags))
+		for _, tag := range tags {
+			modelTags = append(modelTags, models.ModelVersionTag{
+				Key:     tag.Key,
+				Value:   tag.Value,
+				Name:    registeredModel.Name,
+				Version: nextVersion,
+			})
+		}
+
+		if err := m.db.Transaction(func(tx *gorm.DB) error {
+			if err := tx.Create(modelVersion).Error; err != nil {
+				return err
+			}
+			if len(modelTags) > 0 {
+				if err := tx.Create(modelTags).Error; err != nil {
+					return err
+				}
+				modelVersion.Tags = modelTags
+			}
+
+			return nil
+		}); err != nil {
+			return nil, contract.NewErrorWith(
+				protos.ErrorCode_INTERNAL_ERROR,
+				fmt.Sprintf("Model Version creation error (name=%s).", name),
+				err,
+			)
+		}
+
+		return modelVersion.ToEntity(), nil
+	}
+
+	return nil, contract.NewError(
+		protos.ErrorCode_INTERNAL_ERROR,
+		fmt.Sprintf(
+			"Model Version creation error (name=%s). Giving up after %d attempts.",
+			name, CreateModelVersionRetries,
+		),
+	)
+}
+
 func (m *ModelRegistrySQLStore) GetModelVersion(
 	ctx context.Context, name, version string, eager bool,
 ) (*entities.ModelVersion, *contract.Error) {

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

@@ -10,6 +10,8 @@ import (
 
 const StageDeletedInternal = "Deleted_Internal"
 
+const ModelVersionStatusReady = "READY"
+
 // ModelVersion mapped from table <model_versions>.
 //
 //revive:disable:exported