Implementing tag operations for both RegisteredModel and ModelVersion

Author: pebeto

docs/src/reference/model_version.md

@@ -9,4 +9,6 @@ deletemodelversion
 searchmodelversions
 getdownloaduriformodelversionartifacts
 transitionmodelversionstage
+setmodelversiontag
+deletemodelversiontag
 ```

docs/src/reference/registered_model.md

@@ -6,4 +6,6 @@ renameregisteredmodel
 updateregisteredmodel
 deleteregisteredmodel
 searchregisteredmodels
+setregisteredmodeltag
+deleteregisteredmodeltag
 ```

src/MLFlowClient.jl

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

src/services/model_version.jl

@@ -169,3 +169,41 @@ function transitionmodelversionstage(instance::MLFlow, name::String, version::St
         version=version, stage=stage, archive_existing_versions=archive_existing_versions)
     return result["model_version"] |> ModelVersion
 end
+
+"""
+    setmodelversiontag(instance::MLFlow, name::String, key::String, value::String)
+
+# Arguments
+- `instance:` [`MLFlow`](@ref) configuration.
+- `name:` Unique name of the model.
+- `version:` Model version number.
+- `key:` Name of the [`Tag`](@ref).
+- `value:` String value of the tag being logged.
+
+# Returns
+`true` if successful. Otherwise, raises exception.
+"""
+function setmodelversiontag(instance::MLFlow, name::String, version::String, key::String,
+    value::String)::Bool
+    mlfpost(instance, "model-versions/set-tag"; name=name, version=version, key=key,
+        value=value)
+    return true
+end
+
+"""
+    deletemodelversiontag(instance::MLFlow, name::String, version::String, key::String)
+
+# Arguments
+- `instance:` [`MLFlow`](@ref) configuration.
+- `name:` Name of the [`RegisteredModel`](@ref) that the tag was logged under.
+- `version:` [`ModelVersion`](@ref) number that the tag was logged under.
+- `key:` Name of the [`Tag`](@ref).
+
+# Returns
+`true` if successful. Otherwise, raises exception.
+"""
+function deletemodelversiontag(instance::MLFlow, name::String, version::String,
+    key::String)::Bool
+    mlfdelete(instance, "model-versions/delete-tag"; name=name, version=version, key=key)
+    return true
+end

src/services/registered_model.jl

@@ -123,3 +123,36 @@ function searchregisteredmodels(instance::MLFlow; filter::String="",
 
     return registered_models, next_page_token
 end
+
+"""
+    setregisteredmodeltag(instance::MLFlow, name::String, key::String, value::String)
+
+# Arguments
+- `instance:` [`MLFlow`](@ref) configuration.
+- `name:` Unique name of the model.
+- `key:` Name of the [`Tag`](@ref).
+- `value:` String value of the [`Tag`](@ref) being logged.
+
+# Returns
+`true` if successful. Otherwise, raises exception.
+"""
+function setregisteredmodeltag(instance::MLFlow, name::String, key::String, value::String)::Bool
+    mlfpost(instance, "registered-models/set-tag"; name=name, key=key, value=value)
+    return true
+end
+
+"""
+    deleteregisteredmodeltag(instance::MLFlow, name::String, key::String)
+
+# Arguments
+- `instance:` [`MLFlow`](@ref) configuration.
+- `name:` Name of the [`RegisteredModel`](@ref) that the tag was logged under.
+- `key:` Name of the [`Tag`](@ref).
+
+# Returns
+`true` if successful. Otherwise, raises exception.
+"""
+function deleteregisteredmodeltag(instance::MLFlow, name::String, key::String)::Bool
+    mlfdelete(instance, "registered-models/delete-tag"; name=name, key=key)
+    return true
+end

test/services/model_version.jl

@@ -142,7 +142,8 @@ end
     createregisteredmodel(mlf, "missy")
 
     model_version = createmodelversion(mlf, "missy", run.info.artifact_uri)
-    download_uri = getdownloaduriformodelversionartifacts(mlf, "missy", model_version.version)
+    download_uri = getdownloaduriformodelversionartifacts(mlf, model_version.name,
+        model_version.version)
 
     @test download_uri isa String
     @test download_uri == run.info.artifact_uri
@@ -159,7 +160,8 @@ end
     createregisteredmodel(mlf, "missy")
 
     model_version = createmodelversion(mlf, "missy", run.info.artifact_uri)
-    updated_model_version = transitionmodelversionstage(mlf, "missy", model_version.version, "Production", true)
+    updated_model_version = transitionmodelversionstage(mlf, model_version.name,
+        model_version.version, "Production", true)
 
     @test updated_model_version isa ModelVersion
     @test updated_model_version.name == model_version.name
@@ -171,3 +173,42 @@ end
     deleteregisteredmodel(mlf, "missy")
     deleteexperiment(mlf, experiment)
 end
+
+@testset verbose = true "set model version tag" begin
+    @ensuremlf
+
+    experiment = createexperiment(mlf, UUIDs.uuid4() |> string)
+    run = createrun(mlf, experiment)
+    createregisteredmodel(mlf, "missy")
+
+    model_version = createmodelversion(mlf, "missy", run.info.artifact_uri)
+    setmodelversiontag(mlf, model_version.name, model_version.version, "test_key",
+        "test_value")
+    retrieved_model_version = getmodelversion(mlf, "missy", model_version.version)
+
+    @test retrieved_model_version.tags |> length == 1
+    @test (retrieved_model_version.tags |> first).key == "test_key"
+    @test (retrieved_model_version.tags |> first).value == "test_value"
+
+    deleteregisteredmodel(mlf, "missy")
+    deleteexperiment(mlf, experiment)
+end
+
+@testset verbose = true "delete model version tag" begin
+    @ensuremlf
+
+    experiment = createexperiment(mlf, UUIDs.uuid4() |> string)
+    run = createrun(mlf, experiment)
+    createregisteredmodel(mlf, "missy")
+
+    model_version = createmodelversion(mlf, "missy", run.info.artifact_uri)
+    setmodelversiontag(mlf, model_version.name, model_version.version, "test_key",
+        "test_value")
+    deletemodelversiontag(mlf, model_version.name, model_version.version, "test_key")
+    retrieved_model_version = getmodelversion(mlf, "missy", model_version.version)
+
+    @test isempty(retrieved_model_version.tags)
+
+    deleteregisteredmodel(mlf, "missy")
+    deleteexperiment(mlf, experiment)
+end

test/services/registered_model.jl

@@ -110,3 +110,30 @@ end
     deleteregisteredmodel(mlf, "missy")
     deleteregisteredmodel(mlf, "gala")
 end
+
+@testset verbose = true "set registered model tag" begin
+    @ensuremlf
+
+    registered_model = createregisteredmodel(mlf, "missy"; description="gala")
+    setregisteredmodeltag(mlf, registered_model.name, "test_key", "test_value")
+
+    retrieved_registered_model = getregisteredmodel(mlf, registered_model.name)
+    @test retrieved_registered_model.tags |> !isempty
+    @test (retrieved_registered_model.tags |> first).key == "test_key"
+    @test (retrieved_registered_model.tags |> first).value == "test_value"
+
+    deleteregisteredmodel(mlf, "missy")
+end
+
+@testset verbose = true "delete registered model tag" begin
+    @ensuremlf
+
+    registered_model = createregisteredmodel(mlf, "missy"; description="gala")
+    setregisteredmodeltag(mlf, registered_model.name, "test_key", "test_value")
+    deleteregisteredmodeltag(mlf, registered_model.name, "test_key")
+
+    retrieved_registered_model = getregisteredmodel(mlf, registered_model.name)
+    @test retrieved_registered_model.tags |> isempty
+
+    deleteregisteredmodel(mlf, "missy")
+end