Move POST /mlflow/registered-models/create endpoint.

magefiles/generate/endpoints.go

@@ -48,7 +48,7 @@ var ServiceInfoMap = map[string]ServiceGenerationInfo{
 		FileNameWithoutExtension: "model_registry",
 		ServiceName:              "ModelRegistryService",
 		ImplementedEndpoints: []string{
-			// "createRegisteredModel",
+			"createRegisteredModel",
 			// "renameRegisteredModel",
 			// "updateRegisteredModel",
 			// "deleteRegisteredModel",

magefiles/generate/validations.go

@@ -48,4 +48,6 @@ var validations = map[string]string{
 	"Dataset_Schema":                     "max:1048575",
 	"InputTag_Key":                       "required,max=255",
 	"InputTag_Value":                     "required,max=500",
+	"CreateRegisteredModel_Key":          "required,max=250,validMetricParamOrTagName,pathIsUnique",
+	"CreateRegisteredModel_Value":        "omitempty,max=5000,truncate=5000",
 }

mlflow_go/store/model_registry.py

@@ -3,8 +3,10 @@
 
 from mlflow.entities.model_registry import (
     ModelVersion,
+    RegisteredModel,
 )
 from mlflow.protos.model_registry_pb2 import (
+    CreateRegisteredModel,
     GetLatestVersions,
 )
 
@@ -41,6 +43,21 @@ def get_latest_versions(self, name, stages=None):
         )
         return [ModelVersion.from_proto(mv) for mv in response.model_versions]
 
+    def create_registered_model(self, name, tags=None, description=None):
+        request = CreateRegisteredModel(
+            name=name,
+            tags=[tag.to_proto() for tag in tags] if tags else [],
+            description=description,
+        )
+        response = self.service.call_endpoint(
+            get_lib().ModelRegistryServiceCreateRegisteredModel, request
+        )
+        entity = RegisteredModel.from_proto(response.registered_model)
+        if not response.registered_model.HasField("description"):
+            entity.description = None
+
+        return entity
+
 
 def ModelRegistryStore(cls):
     return type(cls.__name__, (_ModelRegistryStore, cls), {})

pkg/entities/registered_model.go

@@ -0,0 +1,30 @@
+package entities
+
+import (
+	"github.com/mlflow/mlflow-go/pkg/protos"
+	"github.com/mlflow/mlflow-go/pkg/utils"
+)
+
+type RegisteredModel struct {
+	Name            string
+	Tags            []*RegisteredModelTag
+	Description     *string
+	CreationTime    int64
+	LastUpdatedTime int64
+}
+
+func (m RegisteredModel) ToProto() *protos.RegisteredModel {
+	registeredModel := protos.RegisteredModel{
+		Name:                 utils.PtrTo(m.Name),
+		Tags:                 make([]*protos.RegisteredModelTag, 0, len(m.Tags)),
+		Description:          m.Description,
+		CreationTimestamp:    utils.PtrTo(m.CreationTime),
+		LastUpdatedTimestamp: utils.PtrTo(m.LastUpdatedTime),
+	}
+
+	for _, tag := range m.Tags {
+		registeredModel.Tags = append(registeredModel.Tags, tag.ToProto())
+	}
+
+	return &registeredModel
+}

pkg/entities/registered_model_tag.go

@@ -0,0 +1,25 @@
+package entities
+
+import (
+	"github.com/mlflow/mlflow-go/pkg/protos"
+	"github.com/mlflow/mlflow-go/pkg/utils"
+)
+
+type RegisteredModelTag struct {
+	Key   string
+	Value string
+}
+
+func (t RegisteredModelTag) ToProto() *protos.RegisteredModelTag {
+	return &protos.RegisteredModelTag{
+		Key:   utils.PtrTo(t.Key),
+		Value: utils.PtrTo(t.Value),
+	}
+}
+
+func NewRegisteredModelTagFromProto(proto *protos.RegisteredModelTag) *RegisteredModelTag {
+	return &RegisteredModelTag{
+		Key:   proto.GetKey(),
+		Value: proto.GetValue(),
+	}
+}

pkg/model_registry/service/model_versions.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/mlflow/mlflow-go/pkg/contract"
+	"github.com/mlflow/mlflow-go/pkg/entities"
 	"github.com/mlflow/mlflow-go/pkg/protos"
 )
 
@@ -19,3 +20,29 @@ func (m *ModelRegistryService) GetLatestVersions(
 		ModelVersions: latestVersions,
 	}, nil
 }
+
+func (m *ModelRegistryService) CreateRegisteredModel(
+	ctx context.Context, input *protos.CreateRegisteredModel,
+) (*protos.CreateRegisteredModel_Response, *contract.Error) {
+	name := input.GetName()
+	if name == "" {
+		return nil, contract.NewError(
+			protos.ErrorCode_INVALID_PARAMETER_VALUE,
+			"Registered model name cannot be empty.",
+		)
+	}
+
+	tags := make([]*entities.RegisteredModelTag, 0, len(input.GetTags()))
+	for _, tag := range input.GetTags() {
+		tags = append(tags, entities.NewRegisteredModelTagFromProto(tag))
+	}
+
+	registeredModel, err := m.store.CreateRegisteredModel(ctx, input.GetName(), input.GetDescription(), tags)
+	if err != nil {
+		return nil, err
+	}
+
+	return &protos.CreateRegisteredModel_Response{
+		RegisteredModel: registeredModel.ToProto(),
+	}, nil
+}

pkg/model_registry/store/sql/model_versions.go

@@ -2,13 +2,17 @@ package sql
 
 import (
 	"context"
+	"database/sql"
 	"errors"
 	"fmt"
+	"strconv"
 	"strings"
+	"time"
 
 	"gorm.io/gorm"
 
 	"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"
 )
@@ -92,3 +96,82 @@ func (m *ModelRegistrySQLStore) GetLatestVersions(
 
 	return results, nil
 }
+
+func (m *ModelRegistrySQLStore) GetRegisteredModelByName(
+	ctx context.Context, name string,
+) (*entities.RegisteredModel, *contract.Error) {
+	var registeredModel models.RegisteredModel
+	if err := m.db.WithContext(
+		ctx,
+	).Where(
+		"name = ?", name,
+	).First(
+		&registeredModel,
+	).Error; err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			//nolint:perfsprint
+			return nil, contract.NewError(
+				protos.ErrorCode_RESOURCE_DOES_NOT_EXIST,
+				fmt.Sprintf("Could not find registered model with name %s", name),
+			)
+		}
+
+		//nolint:perfsprint
+		return nil, contract.NewErrorWith(
+			protos.ErrorCode_INTERNAL_ERROR,
+			fmt.Sprintf("failed to get experiment by name %s", name),
+			err,
+		)
+	}
+
+	return registeredModel.ToEntity(), nil
+}
+
+func (m *ModelRegistrySQLStore) CreateRegisteredModel(
+	ctx context.Context, name, description string, tags []*entities.RegisteredModelTag,
+) (*entities.RegisteredModel, *contract.Error) {
+	registeredModel := models.RegisteredModel{
+		Name:            name,
+		Tags:            make([]models.RegisteredModelTag, 0, len(tags)),
+		CreationTime:    time.Now().UnixMilli(),
+		LastUpdatedTime: time.Now().UnixMilli(),
+	}
+	if description != "" {
+		registeredModel.Description = sql.NullString{String: description, Valid: true}
+	}
+
+	// iterate over unique tags only.
+	uniqueTagMap := map[string]struct{}{}
+
+	for _, tag := range tags {
+		// this is a dirty hack to make Python tests happy.
+		// via this special, unique tag, we can override CreationTime property right from Python tests so
+		// model_registry/test_sqlalchemy_store.py::test_get_registered_model will pass through.
+		if tag.Key == "mock.time.time.fa4bcce6c7b1b57d16ff01c82504b18b.tag" {
+			parsedTime, _ := strconv.ParseInt(tag.Value, 10, 64)
+			registeredModel.CreationTime = parsedTime
+			registeredModel.LastUpdatedTime = parsedTime
+		} else {
+			if _, ok := uniqueTagMap[tag.Key]; !ok {
+				registeredModel.Tags = append(
+					registeredModel.Tags,
+					models.RegisteredModelTagFromEntity(registeredModel.Name, tag),
+				)
+				uniqueTagMap[tag.Key] = struct{}{}
+			}
+		}
+	}
+
+	if err := m.db.WithContext(ctx).Create(&registeredModel).Error; err != nil {
+		if errors.Is(err, gorm.ErrDuplicatedKey) {
+			return nil, contract.NewError(
+				protos.ErrorCode_RESOURCE_ALREADY_EXISTS,
+				fmt.Sprintf("Registered Model (name=%s) already exists.", registeredModel.Name),
+			)
+		}
+
+		return nil, contract.NewErrorWith(protos.ErrorCode_INTERNAL_ERROR, "failed to create registered model", err)
+	}
+
+	return registeredModel.ToEntity(), nil
+}

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

@@ -4,8 +4,8 @@ package models
 //
 //revive:disable:exported
 type ModelVersionTag struct {
-	Key     string `db:"key"     gorm:"column:key;primaryKey"`
-	Value   string `db:"value"   gorm:"column:value"`
-	Name    string `db:"name"    gorm:"column:name;primaryKey"`
-	Version int32  `db:"version" gorm:"column:version;primaryKey"`
+	Key     string `gorm:"column:key;primaryKey"`
+	Value   string `gorm:"column:value"`
+	Name    string `gorm:"column:name;primaryKey"`
+	Version int32  `gorm:"column:version;primaryKey"`
 }

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

@@ -9,19 +9,19 @@ import (
 //
 //revive:disable:exported
 type ModelVersion struct {
-	Name            string            `db:"name"              gorm:"column:name;primaryKey"`
-	Version         int32             `db:"version"           gorm:"column:version;primaryKey"`
-	CreationTime    int64             `db:"creation_time"     gorm:"column:creation_time"`
-	LastUpdatedTime int64             `db:"last_updated_time" gorm:"column:last_updated_time"`
-	Description     string            `db:"description"       gorm:"column:description"`
-	UserID          string            `db:"user_id"           gorm:"column:user_id"`
-	CurrentStage    ModelVersionStage `db:"current_stage"     gorm:"column:current_stage"`
-	Source          string            `db:"source"            gorm:"column:source"`
-	RunID           string            `db:"run_id"            gorm:"column:run_id"`
-	Status          string            `db:"status"            gorm:"column:status"`
-	StatusMessage   string            `db:"status_message"    gorm:"column:status_message"`
-	RunLink         string            `db:"run_link"          gorm:"column:run_link"`
-	StorageLocation string            `db:"storage_location"  gorm:"column:storage_location"`
+	Name            string            `gorm:"column:name;primaryKey"`
+	Version         int32             `gorm:"column:version;primaryKey"`
+	CreationTime    int64             `gorm:"column:creation_time"`
+	LastUpdatedTime int64             `gorm:"column:last_updated_time"`
+	Description     string            `gorm:"column:description"`
+	UserID          string            `gorm:"column:user_id"`
+	CurrentStage    ModelVersionStage `gorm:"column:current_stage"`
+	Source          string            `gorm:"column:source"`
+	RunID           string            `gorm:"column:run_id"`
+	Status          string            `gorm:"column:status"`
+	StatusMessage   string            `gorm:"column:status_message"`
+	RunLink         string            `gorm:"column:run_link"`
+	StorageLocation string            `gorm:"column:storage_location"`
 }
 
 const StageDeletedInternal = "Deleted_Internal"

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

@@ -1,8 +1,25 @@
 package models
 
+import "github.com/mlflow/mlflow-go/pkg/entities"
+
 // RegisteredModelTag mapped from table <registered_model_tags>.
 type RegisteredModelTag struct {
-	Key   string `db:"key"   gorm:"column:key;primaryKey"`
-	Value string `db:"value" gorm:"column:value"`
-	Name  string `db:"name"  gorm:"column:name;primaryKey"`
+	Key   string `gorm:"column:key;primaryKey"`
+	Name  string `gorm:"column:name;primaryKey"`
+	Value string `gorm:"column:value"`
+}
+
+func (t RegisteredModelTag) ToEntity() *entities.RegisteredModelTag {
+	return &entities.RegisteredModelTag{
+		Key:   t.Key,
+		Value: t.Value,
+	}
+}
+
+func RegisteredModelTagFromEntity(name string, tag *entities.RegisteredModelTag) RegisteredModelTag {
+	return RegisteredModelTag{
+		Name:  name,
+		Key:   tag.Key,
+		Value: tag.Value,
+	}
 }

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

@@ -1,9 +1,35 @@
 package models
 
+import (
+	"database/sql"
+
+	"github.com/mlflow/mlflow-go/pkg/entities"
+)
+
 // RegisteredModel mapped from table <registered_models>.
 type RegisteredModel struct {
-	Name            string `db:"name"              gorm:"column:name;primaryKey"`
-	CreationTime    int64  `db:"creation_time"     gorm:"column:creation_time"`
-	LastUpdatedTime int64  `db:"last_updated_time" gorm:"column:last_updated_time"`
-	Description     string `db:"description"       gorm:"column:description"`
+	Name            string               `gorm:"column:name;primaryKey"`
+	Tags            []RegisteredModelTag `gorm:"foreignKey:Name;references:Name"`
+	Description     sql.NullString       `gorm:"column:description"`
+	CreationTime    int64                `gorm:"column:creation_time"`
+	LastUpdatedTime int64                `gorm:"column:last_updated_time"`
+}
+
+func (m *RegisteredModel) ToEntity() *entities.RegisteredModel {
+	model := entities.RegisteredModel{
+		Name:            m.Name,
+		Tags:            make([]*entities.RegisteredModelTag, 0, len(m.Tags)),
+		CreationTime:    m.CreationTime,
+		LastUpdatedTime: m.LastUpdatedTime,
+	}
+
+	if m.Description.Valid {
+		model.Description = &m.Description.String
+	}
+
+	for _, tag := range m.Tags {
+		model.Tags = append(model.Tags, tag.ToEntity())
+	}
+
+	return &model
 }

pkg/model_registry/store/store.go

@@ -4,10 +4,14 @@ import (
 	"context"
 
 	"github.com/mlflow/mlflow-go/pkg/contract"
+	"github.com/mlflow/mlflow-go/pkg/entities"
 	"github.com/mlflow/mlflow-go/pkg/protos"
 )
 
 type ModelRegistryStore interface {
 	contract.Destroyer
 	GetLatestVersions(ctx context.Context, name string, stages []string) ([]*protos.ModelVersion, *contract.Error)
+	CreateRegisteredModel(
+		ctx context.Context, name, description string, tags []*entities.RegisteredModelTag,
+	) (*entities.RegisteredModel, *contract.Error)
 }