Implementing alias operations for RegisteredModel and ModelVersion

Author: pebeto

docs/src/reference/model_version.md

@@ -11,4 +11,5 @@ getdownloaduriformodelversionartifacts
 transitionmodelversionstage
 setmodelversiontag
 deletemodelversiontag
+getmodelversionbyalias
 ```

docs/src/reference/registered_model.md

@@ -8,4 +8,6 @@ deleteregisteredmodel
 searchregisteredmodels
 setregisteredmodeltag
 deleteregisteredmodeltag
+deleteregisteredmodelalias
+setregisteredmodelalias
 ```

src/MLFlowClient.jl

@@ -69,11 +69,13 @@ export refresh, getmetrichistory
 include("services/registered_model.jl")
 export getregisteredmodel, createregisteredmodel, deleteregisteredmodel,
     renameregisteredmodel, updateregisteredmodel, searchregisteredmodels,
-    setregisteredmodeltag, deleteregisteredmodeltag
+    setregisteredmodeltag, deleteregisteredmodeltag, deleteregisteredmodelalias,
+    setregisteredmodelalias
 
 include("services/model_version.jl")
 export getlatestmodelversions, getmodelversion, createmodelversion, deletemodelversion,
     updatemodelversion, searchmodelversions, getdownloaduriformodelversionartifacts,
-    transitionmodelversionstage, setmodelversiontag, deletemodelversiontag
+    transitionmodelversionstage, setmodelversiontag, deletemodelversiontag,
+    getmodelversionbyalias
 
 end

src/services/model_version.jl

@@ -207,3 +207,20 @@ function deletemodelversiontag(instance::MLFlow, name::String, version::String,
     mlfdelete(instance, "model-versions/delete-tag"; name=name, version=version, key=key)
     return true
 end
+
+"""
+    getmodelversionbyalias(instance::MLFlow, name::String, alias::String)
+
+# Arguments
+- `instance:` [`MLFlow`](@ref) configuration.
+- `name:` Name of the [`RegisteredModel`](@ref).
+- `alias:` Name of the alias.
+
+# Returns
+[`ModelVersion`](@ref) requested.
+"""
+function getmodelversionbyalias(instance::MLFlow, name::String,
+    alias::String)::ModelVersion
+    result = mlfget(instance, "registered-models/alias"; name=name, alias=alias)
+    return result["model_version"] |> ModelVersion
+end

src/services/registered_model.jl

@@ -156,3 +156,37 @@ function deleteregisteredmodeltag(instance::MLFlow, name::String, key::String)::
     mlfdelete(instance, "registered-models/delete-tag"; name=name, key=key)
     return true
 end
+
+"""
+    deleteregisteredmodelalias(instance::MLFlow, name::String, alias::String)
+
+# Arguments
+- `instance:` [`MLFlow`](@ref) configuration.
+- `name:` Name of the [`RegisteredModel`](@ref).
+- `alias:` Name of the alias.
+
+# Returns
+`true` if successful. Otherwise, raises exception.
+"""
+function deleteregisteredmodelalias(instance::MLFlow, name::String, alias::String)::Bool
+    mlfdelete(instance, "registered-models/alias"; name=name, alias=alias)
+    return true
+end
+
+"""
+    setregisteredmodelalias(instance::MLFlow, name::String, alias::String, version::String)
+
+# Arguments
+- `instance:` [`MLFlow`](@ref) configuration.
+- `name:` Name of the [`RegisteredModel`](@ref).
+- `alias:` Name of the alias.
+- `version:` [`ModelVersion`](@ref) number.
+
+# Returns
+`true` if successful. Otherwise, raises exception.
+"""
+function setregisteredmodelalias(instance::MLFlow, name::String, alias::String,
+    version::String)::Bool
+    mlfpost(instance, "registered-models/alias"; name=name, alias=alias, version=version)
+    return true
+end

test/services/model_version.jl

@@ -212,3 +212,22 @@ end
     deleteregisteredmodel(mlf, "missy")
     deleteexperiment(mlf, experiment)
 end
+
+@testset verbose = true "get model version by alias" begin
+    @ensuremlf
+
+    experiment = createexperiment(mlf, UUIDs.uuid4() |> string)
+    run = createrun(mlf, experiment)
+    registered_model = createregisteredmodel(mlf, "missy")
+
+    model_version = createmodelversion(mlf, "missy", run.info.artifact_uri)
+    setregisteredmodelalias(mlf, registered_model.name, "gala", model_version.version)
+
+    retrieved_model_version = getmodelversionbyalias(mlf, "missy", "gala")
+
+    @assert retrieved_model_version.aliases |> !isempty
+    @assert (retrieved_model_version.aliases |> first) == "gala"
+
+    deleteregisteredmodel(mlf, "missy")
+    deleteexperiment(mlf, experiment)
+end

test/services/registered_model.jl

@@ -54,13 +54,13 @@ end
     @ensuremlf
 
     registered_model = createregisteredmodel(mlf, "missy"; description="gala")
-    renamed_registered_model = renameregisteredmodel(mlf, registered_model.name, "mister")
+    renamed_registered_model = renameregisteredmodel(mlf, registered_model.name, "gala")
 
     @test renamed_registered_model isa RegisteredModel
-    @test renamed_registered_model.name == "mister"
+    @test renamed_registered_model.name == "gala"
     @test renamed_registered_model.description == registered_model.description
 
-    deleteregisteredmodel(mlf, "mister")
+    deleteregisteredmodel(mlf, "gala")
 end
 
 @testset verbose = true "update registered model" begin
@@ -137,3 +137,42 @@ end
 
     deleteregisteredmodel(mlf, "missy")
 end
+
+@testset verbose = true "delete registered model alias" begin
+    @ensuremlf
+
+    experiment = createexperiment(mlf, UUIDs.uuid4() |> string)
+    run = createrun(mlf, experiment)
+
+    registered_model = createregisteredmodel(mlf, "missy"; description="gala")
+    model_version = createmodelversion(mlf, "missy", run.info.artifact_uri)
+
+    setregisteredmodelalias(mlf, registered_model.name, "gala", model_version.version)
+    deleteregisteredmodelalias(mlf, registered_model.name, "gala")
+
+    retrieved_registered_model = getregisteredmodel(mlf, registered_model.name)
+    @test retrieved_registered_model.aliases |> isempty
+
+    deleteregisteredmodel(mlf, "missy")
+    deleteexperiment(mlf, experiment)
+end
+
+@testset verbose = true "delete registered model alias" begin
+    @ensuremlf
+
+    experiment = createexperiment(mlf, UUIDs.uuid4() |> string)
+    run = createrun(mlf, experiment)
+
+    registered_model = createregisteredmodel(mlf, "missy"; description="gala")
+    model_version = createmodelversion(mlf, "missy", run.info.artifact_uri)
+
+    setregisteredmodelalias(mlf, registered_model.name, "gala", model_version.version)
+    setregisteredmodelalias(mlf, registered_model.name, "missy", model_version.version)
+
+    retrieved_registered_model = getregisteredmodel(mlf, registered_model.name)
+    @test retrieved_registered_model.aliases |> !isempty
+    @test length(retrieved_registered_model.aliases) == 2
+
+    deleteregisteredmodel(mlf, "missy")
+    deleteexperiment(mlf, experiment)
+end