Implementing registered_model service

Author: pebeto

docs/make.jl

@@ -10,6 +10,7 @@ makedocs(;
         "Experiment operations" => "reference/experiment.md",
         "Logging operations" => "reference/loggers.md",
         "Miscellaneous operations" => "reference/misc.md",
-        "Run operations" => "reference/run.md"]])
+        "Run operations" => "reference/run.md",
+        "Registered model operations" => "reference/registered_model.md"]])
 
 deploydocs(; repo="github.com/JuliaAI/MLFlowClient.jl", devbranch="main")

docs/src/reference/registered_model.md

@@ -0,0 +1,8 @@
+# Registered model operations
+```@docs
+createregisteredmodel
+getregisteredmodel
+renameregisteredmodel
+updateregisteredmodel
+deleteregisteredmodel
+```

src/MLFlowClient.jl

@@ -66,4 +66,8 @@ export listartifacts
 include("services/misc.jl")
 export refresh, getmetrichistory
 
+include("services/registered_model.jl")
+export getregisteredmodel, createregisteredmodel, deleteregisteredmodel,
+    renameregisteredmodel, updateregisteredmodel
+
 end

src/api.jl

@@ -66,3 +66,48 @@ function mlfpost(mlf, endpoint; kwargs...)
         throw(ErrorException(error_message))
     end
 end
+
+"""
+    mlfpatch(mlf, endpoint; kwargs...)
+
+Performs a HTTP PATCH to the specified endpoint. kwargs are converted to JSON and become
+the PATCH body.
+"""
+function mlfpatch(mlf, endpoint; kwargs...)
+    apiuri = uri(mlf, endpoint;)
+    apiheaders = headers(mlf, Dict("Content-Type" => "application/json"))
+    body = JSON.json(kwargs)
+
+    try
+        response = HTTP.patch(apiuri, apiheaders, body)
+        return response.body |> String |> JSON.parse
+    catch e
+        error_response = e.response.body |> String |> JSON.parse
+        error_message = "$(error_response["error_code"]) -  $(error_response["message"])"
+        @error error_message
+        throw(ErrorException(error_message))
+    end
+end
+
+"""
+    mlfdelete(mlf, endpoint; kwargs...)
+
+Performs a HTTP DELETE to the specified endpoint. kwargs are converted to JSON and become
+the DELETE body.
+"""
+function mlfdelete(mlf, endpoint; kwargs...)
+    apiuri = uri(mlf, endpoint;
+        parameters=Dict(k => v for (k, v) in kwargs if v !== missing))
+    apiheaders = headers(mlf, Dict("Content-Type" => "application/json"))
+    body = JSON.json(kwargs)
+
+    try
+        response = HTTP.delete(apiuri, apiheaders, body)
+        return response.body |> String |> JSON.parse
+    catch e
+        error_response = e.response.body |> String |> JSON.parse
+        error_message = "$(error_response["error_code"]) -  $(error_response["message"])"
+        @error error_message
+        throw(ErrorException(error_message))
+    end
+end

src/services/registered_model.jl

@@ -0,0 +1,91 @@
+"""
+    createregisteredmodel(instance::MLFlow, name::String;
+        tags::MLFlowUpsertData{Tag}=Tag[], description::Union{String, Missing}=missing)
+
+Create a [`RegisteredModel`](@ref) with a name. Returns the newly created
+[`RegisteredModel`](@ref). Validates that another [`RegisteredModel`](@ref) with the same
+name does not already exist and fails if another [`RegisteredModel`](@ref) with the same
+name already exists.
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `name`: Register models under this name.
+- `tags`: A collection of [`Tag`](@ref).
+- `description`: Optional description for [`RegisteredModel`](@ref).
+
+# Returns
+An instance of type [`RegisteredModel`](@ref).
+"""
+function createregisteredmodel(instance::MLFlow, name::String;
+    tags::MLFlowUpsertData{Tag}=Tag[],
+    description::Union{String, Missing}=missing)::RegisteredModel
+    result = mlfpost(instance, "registered-models/create"; name=name,
+        tags=parse(Tag, tags), description=description)
+    return result["registered_model"] |> RegisteredModel
+end
+
+"""
+    getregisteredmodel(instance::MLFlow, name::String)
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `name`: [`RegisteredModel`](@ref) model unique name identifier.
+
+# Returns
+An instance of type [`RegisteredModel`](@ref).
+"""
+function getregisteredmodel(instance::MLFlow, name::String)::RegisteredModel
+    result = mlfget(instance, "registered-models/get"; name=name)
+    return result["registered_model"] |> RegisteredModel
+end
+
+"""
+    renameregisteredmodel(instance::MLFlow, name::String, new_name::String)
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `name`: [`RegisteredModel`](@ref) unique name identifier.
+- `new_name`: If provided, updates the name for this [`RegisteredModel`](@ref).
+
+# Returns
+An instance of type [`RegisteredModel`](@ref).
+"""
+function renameregisteredmodel(instance::MLFlow, name::String,
+    new_name::String)::RegisteredModel
+    result = mlfpost(instance, "registered-models/rename"; name=name, new_name=new_name)
+    return result["registered_model"] |> RegisteredModel
+end
+
+"""
+    updateregisteredmodel(instance::MLFlow, name::String;
+        description::Union{String, Missing}=missing)
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `name`: [`RegisteredModel`](@ref) unique name identifier.
+- `description`: If provided, updates the description for this [`RegisteredModel`](@ref).
+
+# Returns
+An instance of type [`RegisteredModel`](@ref).
+"""
+function updateregisteredmodel(instance::MLFlow, name::String;
+    description::Union{String, Missing}=missing)::RegisteredModel
+    result = mlfpatch(instance, "registered-models/update"; name=name,
+        description=description)
+    return result["registered_model"] |> RegisteredModel
+end
+
+"""
+    deleteregisteredmodel(instance::MLFlow, name::String)
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `name`: [`RegisteredModel`](@ref) unique name identifier.
+
+# Returns
+`true` if successful. Otherwise, raises exception.
+"""
+function deleteregisteredmodel(instance::MLFlow, name::String)::Bool
+    mlfdelete(instance, "registered-models/delete"; name=name)
+    return true
+end

src/types/enums.jl

@@ -12,6 +12,7 @@
     FAILED_REGISTRATION=2
     READY=3
 end
+ModelVersionStatus(status::String) = Dict(value => key for (key, value) in ModelVersionStatus |> Base.Enums.namemap)[status |> Symbol] |> ModelVersionStatus
 
 """
     RunStatus

src/types/model_version.jl

@@ -38,4 +38,10 @@ struct ModelVersion
     run_link::String
     aliases::Array{String}
 end
+ModelVersion(data::Dict{String, Any}) = ModelVersion(data["name"], data["version"],
+    data["creation_timestamp"], data["last_updated_timestamp"], data["user_id"],
+    data["current_stage"], data["description"], data["source"], data["run_id"],
+    ModelVersionStatus(data["status"]), data["status_message"],
+    [Tag(tag) for tag in get(data, "tags", [])], data["run_link"],
+    get(data, "aliases", []))
 Base.show(io::IO, t::ModelVersion) = show(io, ShowCase(t, new_lines=true))

src/types/registered_model.jl

@@ -11,6 +11,8 @@ struct RegisteredModelAlias
     alias::String
     version::String
 end
+RegisteredModelAlias(data::Dict{String, Any}) = RegisteredModelAlias(data["alias"],
+    data["version"])
 Base.show(io::IO, t::RegisteredModelAlias) = show(io, ShowCase(t, new_lines=true))
 
 """
@@ -21,8 +23,8 @@ Base.show(io::IO, t::RegisteredModelAlias) = show(io, ShowCase(t, new_lines=true
 - `creation_timestamp::Int64`: Timestamp recorded when this RegisteredModel was created.
 - `last_updated_timestamp::Int64`: Timestamp recorded when metadata for this
     RegisteredModel was last updated.
-- `user_id::String`: User that created this RegisteredModel.
-- `description::String`: Description of this RegisteredModel.
+- `user_id::Union{String, Nothing}`: User that created this RegisteredModel.
+- `description::Union{String, Nothing}`: Description of this RegisteredModel.
 - `latest_versions::Array{ModelVersion}`: Collection of latest model versions for each
     stage. Only contains models with current READY status.
 - `tags::Array{Tag}`: Additional metadata key-value pairs.
@@ -33,10 +35,16 @@ struct RegisteredModel
     name::String
     creation_timestamp::Int64
     last_updated_timestamp::Int64
-    user_id::String
-    description::String
+    user_id::Union{String, Nothing}
+    description::Union{String, Nothing}
     latest_versions::Array{ModelVersion}
     tags::Array{Tag}
     aliases::Array{RegisteredModelAlias}
 end
+RegisteredModel(data::Dict{String, Any}) = RegisteredModel(data["name"],
+    data["creation_timestamp"], data["last_updated_timestamp"],
+    get(data, "user_id", nothing), get(data, "description", nothing),
+    [ModelVersion(version) for version in get(data, "latest_versions", [])],
+    [Tag(tag) for tag in get(data, "tags", [])],
+    [RegisteredModelAlias(alias) for alias in get(data, "aliases", [])])
 Base.show(io::IO, t::RegisteredModel) = show(io, ShowCase(t, new_lines=true))

test/runtests.jl

@@ -9,3 +9,4 @@ include("services/misc.jl")
 include("services/loggers.jl")
 include("services/artifact.jl")
 include("services/experiment.jl")
+include("services/registered_model.jl")

test/services/registered_model.jl

@@ -0,0 +1,87 @@
+@testset verbose = true "create registered model" begin
+    @ensuremlf
+
+    @testset "base" begin
+        registered_model = createregisteredmodel(mlf, "missy"; description="gala")
+
+        @test registered_model isa RegisteredModel
+        @test registered_model.name == "missy"
+        @test registered_model.description == "gala"
+    end
+
+    @testset "name exists" begin
+        registered_model = getregisteredmodel(mlf, "missy")
+        @test_throws ErrorException createregisteredmodel(mlf, registered_model.name)
+        deleteregisteredmodel(mlf, "missy")
+    end
+
+    @testset "with tags as array of tags" begin
+        createregisteredmodel(mlf, "missy"; tags=[Tag("test_key", "test_value")])
+        deleteregisteredmodel(mlf, "missy")
+    end
+
+    @testset "with tags as array of pairs" begin
+        createregisteredmodel(mlf, "missy"; tags=["test_key" => "test_value"])
+        deleteregisteredmodel(mlf, "missy")
+    end
+
+    @testset "with tags as array of dicts" begin
+        createregisteredmodel(mlf, "missy";
+            tags=[Dict("key" => "test_key", "value" => "test_value")])
+        deleteregisteredmodel(mlf, "missy")
+    end
+
+    @testset "with tags as dict" begin
+        createregisteredmodel(mlf, "missy"; tags=Dict("test_key" => "test_value"))
+        deleteregisteredmodel(mlf, "missy")
+    end
+end
+
+@testset verbose = true "get registered model" begin
+    @ensuremlf
+
+    registered_model = createregisteredmodel(mlf, "missy"; description="gala")
+    retrieved_registered_model = getregisteredmodel(mlf, registered_model.name)
+
+    @test retrieved_registered_model isa RegisteredModel
+    @test retrieved_registered_model.name == registered_model.name
+    @test retrieved_registered_model.description == registered_model.description
+
+    deleteregisteredmodel(mlf, "missy")
+end
+
+@testset verbose = true "rename registered model" begin
+    @ensuremlf
+    
+    registered_model = createregisteredmodel(mlf, "missy"; description="gala")
+    renamed_registered_model = renameregisteredmodel(mlf, registered_model.name, "mister")
+
+    @test renamed_registered_model isa RegisteredModel
+    @test renamed_registered_model.name == "mister"
+    @test renamed_registered_model.description == registered_model.description
+
+    deleteregisteredmodel(mlf, "mister")
+end
+
+@testset verbose = true "update registered model" begin
+    @ensuremlf
+
+    registered_model = createregisteredmodel(mlf, "missy"; description="gala")
+    updated_registered_model = updateregisteredmodel(mlf, registered_model.name;
+        description="ana")
+
+    @test updated_registered_model isa RegisteredModel
+    @test updated_registered_model.name == registered_model.name
+    @test updated_registered_model.description == "ana"
+
+    deleteregisteredmodel(mlf, "missy")
+end
+
+@testset verbose = true "delete registered model" begin
+    @ensuremlf
+
+    registered_model = createregisteredmodel(mlf, "missy"; description="gala")
+    deleteregisteredmodel(mlf, "missy")
+
+    @test_throws ErrorException getregisteredmodel(mlf, "missy")
+end