Implementing searchregisteredmodels

Author: pebeto

docs/src/reference/registered_model.md

@@ -5,4 +5,5 @@ getregisteredmodel
 renameregisteredmodel
 updateregisteredmodel
 deleteregisteredmodel
+searchregisteredmodels
 ```

src/MLFlowClient.jl

@@ -68,7 +68,7 @@ export refresh, getmetrichistory
 
 include("services/registered_model.jl")
 export getregisteredmodel, createregisteredmodel, deleteregisteredmodel,
-    renameregisteredmodel, updateregisteredmodel
+    renameregisteredmodel, updateregisteredmodel, searchregisteredmodels
 
 include("services/model_version.jl")
 export getlatestmodelversions, getmodelversion, createmodelversion, deletemodelversion,

src/services/model_version.jl

@@ -11,7 +11,7 @@ Latest [`ModelVersion`](@ref) for each requests stage.
 """
 function getlatestmodelversions(instance::MLFlow, name::String;
     stages::Array{String}=String[])::Array{ModelVersion}
-result = mlfpost(instance, "registered-models/get-latest-versions"; name=name,
+    result = mlfpost(instance, "registered-models/get-latest-versions"; name=name,
         stages=stages)
     return result["model_versions"] .|> ModelVersion
 end
@@ -104,8 +104,8 @@ end
 - `filter`: String filter condition. See [MLFlow documentation](https://mlflow.org/docs/latest/rest-api.html#search-modelversions).
 - `max_results`: Maximum number of models desired.
 - `order_by`: List of columns to be ordered by including model name, version, stage with an
-optional “DESC” or “ASC” annotation, where “ASC” is the default. Tiebreaks are done by
-latest stage transition timestamp, followed by name ASC, followed by version DESC.
+    optional “DESC” or “ASC” annotation, where “ASC” is the default. Tiebreaks are done by
+    latest stage transition timestamp, followed by name ASC, followed by version DESC.
 - `page_token`: Pagination token to go to next page based on previous search query.
 
 # Returns
@@ -156,9 +156,9 @@ end
 - `version:` [`ModelVersion`](@ref) number.
 - `stage:` Transition [`ModelVersion`](@ref) to new stage.
 - `archive_existing_versions:` When transitioning a model version to a particular stage,
-this flag dictates whether all existing model versions in that stage should be atomically
-moved to the “archived” stage. This ensures that at-most-one model version exists in the
-target stage.
+    this flag dictates whether all existing model versions in that stage should be atomically
+    moved to the “archived” stage. This ensures that at-most-one model version exists in the
+    target stage.
 
 # Returns
 Updated [`ModelVersion`](@ref).

src/services/registered_model.jl

@@ -89,3 +89,37 @@ function deleteregisteredmodel(instance::MLFlow, name::String)::Bool
     mlfdelete(instance, "registered-models/delete"; name=name)
     return true
 end
+
+"""
+    searchregisteredmodels(instance::MLFlow, filter::String, max_results::Int64,
+        order_by::String, page_token::String)
+
+# Arguments
+- `instance:` [`MLFlow`](@ref) configuration.
+- `filter`: String filter condition. See [MLFlow documentation](https://mlflow.org/docs/latest/rest-api.html#search-registeredmodels).
+- `max_results`: Maximum number of models desired.
+- `order_by`: List of columns for ordering search results, which can include model name
+    and last updated timestamp with an optional “DESC” or “ASC” annotation, where “ASC” is
+    the default. Tiebreaks are done by model name ASC.
+- `page_token`: Pagination token to go to the next page based on a previous search query.
+
+# Returns
+- Vector of [`RegisteredModel`](@ref) that were found in the [`MLFlow`](@ref) instance.
+- The next page token if there are more results.
+"""
+function searchregisteredmodels(instance::MLFlow; filter::String="",
+    max_results::Int64=100, order_by::Array{String}=String[],
+    page_token::String="")::Tuple{Array{RegisteredModel}, Union{String, Nothing}}
+    parameters = (; max_results, page_token, filter)
+
+    if order_by |> !isempty
+        parameters = (; order_by, parameters...)
+    end
+
+    result = mlfget(instance, "registered-models/search"; parameters...)
+
+    registered_models = get(result, "registered_models", []) |> (x -> [RegisteredModel(y) for y in x])
+    next_page_token = get(result, "next_page_token", nothing)
+
+    return registered_models, next_page_token
+end

test/services/registered_model.jl

@@ -85,3 +85,28 @@ end
 
     @test_throws ErrorException getregisteredmodel(mlf, "missy")
 end
+
+@testset verbose = true "search registered models" begin
+    @ensuremlf
+
+    createregisteredmodel(mlf, "missy"; description="gala")
+    createregisteredmodel(mlf, "gala"; description="missy")
+
+    @testset "default search" begin
+        registered_models, next_page_token = searchregisteredmodels(mlf)
+
+        @test length(registered_models) == 2 # four because of the default experiment
+        @test next_page_token |> isnothing
+    end
+
+    @testset "with pagination" begin
+        registered_models, next_page_token = searchregisteredmodels(mlf; max_results=1)
+
+        @test length(registered_models) == 1
+        @test next_page_token |> !isnothing
+        @test next_page_token isa String
+    end
+
+    deleteregisteredmodel(mlf, "missy")
+    deleteregisteredmodel(mlf, "gala")
+end