Skip to content

Commit

Permalink
Move POST /mlflow/registered-models/rename endpoint. (#88)
Browse files Browse the repository at this point in the history
Signed-off-by: Software Developer <7852635+dsuhinin@users.noreply.github.com>
  • Loading branch information
dsuhinin authored Nov 19, 2024
1 parent e005605 commit 7af6166
Show file tree
Hide file tree
Showing 8 changed files with 103 additions and 1 deletion.
2 changes: 1 addition & 1 deletion magefiles/generate/endpoints.go
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ var ServiceInfoMap = map[string]ServiceGenerationInfo{
ServiceName: "ModelRegistryService",
ImplementedEndpoints: []string{
// "createRegisteredModel",
// "renameRegisteredModel",
"renameRegisteredModel",
"updateRegisteredModel",
// "deleteRegisteredModel",
// "getRegisteredModel",
Expand Down
8 changes: 8 additions & 0 deletions mlflow_go/store/model_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from mlflow.entities.model_registry import ModelVersion, RegisteredModel
from mlflow.protos.model_registry_pb2 import (
GetLatestVersions,
RenameRegisteredModel,
UpdateRegisteredModel,
)

Expand Down Expand Up @@ -47,6 +48,13 @@ def update_registered_model(self, name, description):
)
return RegisteredModel.from_proto(response.registered_model)

def rename_registered_model(self, name, new_name):
request = RenameRegisteredModel(name=name, new_name=new_name)
response = self.service.call_endpoint(
get_lib().ModelRegistryServiceRenameRegisteredModel, request
)
return RegisteredModel.from_proto(response.registered_model)


def ModelRegistryStore(cls):
return type(cls.__name__, (_ModelRegistryStore, cls), {})
Expand Down
1 change: 1 addition & 0 deletions pkg/contract/service/model_registry.g.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

8 changes: 8 additions & 0 deletions pkg/lib/model_registry.g.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

21 changes: 21 additions & 0 deletions pkg/model_registry/service/model_versions.go
Original file line number Diff line number Diff line change
Expand Up @@ -32,3 +32,24 @@ func (m *ModelRegistryService) UpdateRegisteredModel(
RegisteredModel: registeredModel.ToProto(),
}, nil
}

func (m *ModelRegistryService) RenameRegisteredModel(
ctx context.Context, input *protos.RenameRegisteredModel,
) (*protos.RenameRegisteredModel_Response, *contract.Error) {
newName := input.GetNewName()
if newName == "" {
return nil, contract.NewError(
protos.ErrorCode_INVALID_PARAMETER_VALUE,
"Registered model name cannot be empty",
)
}

registeredModel, err := m.store.RenameRegisteredModel(ctx, input.GetName(), newName)
if err != nil {
return nil, err
}

return &protos.RenameRegisteredModel_Response{
RegisteredModel: registeredModel.ToProto(),
}, nil
}
52 changes: 52 additions & 0 deletions pkg/model_registry/store/sql/model_versions.go
Original file line number Diff line number Diff line change
Expand Up @@ -147,3 +147,55 @@ func (m *ModelRegistrySQLStore) UpdateRegisteredModel(

return registeredModel, nil
}

func (m *ModelRegistrySQLStore) RenameRegisteredModel(
ctx context.Context, name, newName string,
) (*entities.RegisteredModel, *contract.Error) {
registeredModel, err := m.GetRegisteredModelByName(ctx, name)
if err != nil {
return nil, err
}

if err := m.db.WithContext(ctx).Transaction(func(transaction *gorm.DB) error {
if err := transaction.Model(
&models.ModelVersion{},
).Where(
"name = ?", registeredModel.Name,
).Updates(&models.ModelVersion{
Name: newName,
LastUpdatedTime: time.Now().UnixMilli(),
}).Error; err != nil {
return err
}

if err := transaction.Model(
&models.RegisteredModel{},
).Where(
"name = ?", registeredModel.Name,
).Updates(&models.RegisteredModel{
Name: newName,
LastUpdatedTime: time.Now().UnixMilli(),
}).Error; err != nil {
return err
}

return nil
}); err != nil {
if errors.Is(err, gorm.ErrDuplicatedKey) {
return nil, contract.NewErrorWith(
protos.ErrorCode_RESOURCE_ALREADY_EXISTS,
fmt.Sprintf("Registered Model (name=%s) already exists", newName),
err,
)
}

return nil, contract.NewErrorWith(protos.ErrorCode_INTERNAL_ERROR, "failed to rename registered model", err)
}

registeredModel, err = m.GetRegisteredModelByName(ctx, newName)
if err != nil {
return nil, err
}

return registeredModel, nil
}
1 change: 1 addition & 0 deletions pkg/model_registry/store/store.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,4 +12,5 @@ type ModelRegistryStore interface {
contract.Destroyer
GetLatestVersions(ctx context.Context, name string, stages []string) ([]*protos.ModelVersion, *contract.Error)
UpdateRegisteredModel(ctx context.Context, name, description string) (*entities.RegisteredModel, *contract.Error)
RenameRegisteredModel(ctx context.Context, name, newName string) (*entities.RegisteredModel, *contract.Error)
}
11 changes: 11 additions & 0 deletions pkg/server/routes/model_registry.g.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 7af6166

Please sign in to comment.