Implementing searchruns

Author: pebeto

src/MLFlowClient.jl

@@ -70,6 +70,7 @@ export
     deleterun,
     setruntag,
     restorerun,
+    searchruns,
     deleteruntag
 
 include("services/loggers.jl")

src/services/experiment.jl

@@ -168,7 +168,8 @@ is the default.
 unspecified, return only active experiments.
 
 # Returns
-- vector of [`MLFlowExperiment`](@ref) experiments that were found in the MLFlow instance
+- Vector of [`Experiment`](@ref) that were found in the MLFlow instance.
+- The next page token if there are more results.
 """
 function searchexperiments(instance::MLFlow; max_results::Int64=20000,
     page_token::String="", filter::String="", order_by::Array{String}=String[],
@@ -183,7 +184,7 @@ function searchexperiments(instance::MLFlow; max_results::Int64=20000,
 
     result = mlfget(instance, "experiments/search"; parameters...)
 
-    experiments = result["experiments"] |> (x -> [Experiment(y) for y in x])
+    experiments = get(result, "experiments", []) |> (x -> [Experiment(y) for y in x])
     next_page_token = get(result, "next_page_token", nothing)
 
     return experiments, next_page_token

src/services/run.jl

@@ -141,3 +141,48 @@ function deleteruntag(instance::MLFlow, run_id::String, key::String)::Bool
 end
 deleteruntag(instance::MLFlow, run::Run, key::String)::Bool =
     deleteruntag(instance, run.info.run_id, key)
+
+"""
+    searchruns(instance::MLFlow; experiment_ids::Array{String}=String[],
+        filter::String="", run_view_type::ViewType=ACTIVE_ONLY,
+        max_results::Int=1000, order_by::Array{String}=String[],
+        page_token::String="")
+
+Search for runs that satisfy expressions. Search expressions can use Metric and
+Param keys.
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `experiment_ids`: List of experiment IDs to search over.
+- `filter`: A filter expression over params, metrics, and tags, that allows
+returning a subset of runs. See [MLFlow documentation](https://mlflow.org/docs/latest/rest-api.html#search-runs).
+- `run_view_type`: Whether to display only active, only deleted, or all runs.
+Defaults to only active runs.
+- `max_results`: Maximum number of runs desired.
+- `order_by`: List of columns to be ordered by, including attributes, params,
+metrics, and tags with an optional “DESC” or “ASC” annotation, where “ASC” is
+the default.
+- `page_token`: Token indicating the page of runs to fetch.
+
+# Returns
+- Vector of [`Run`](@ref) that were found in the specified experiments.
+- The next page token if there are more results.
+"""
+function searchruns(instance::MLFlow; experiment_ids::Array{String}=String[],
+    filter::String="", run_view_type::ViewType=ACTIVE_ONLY,
+    max_results::Int=1000, order_by::Array{String}=String[],
+    page_token::String="")::Tuple{Array{Run}, Union{String, Nothing}}
+    parameters = (; experiment_ids, filter,
+        :run_view_type => run_view_type |> Integer, max_results, page_token)
+
+    if order_by |> !isempty
+        parameters = (; order_by, parameters...)
+    end
+
+    result = mlfpost(instance, "runs/search"; parameters...)
+
+    runs = get(result, "runs", []) |> (x -> [Run(y) for y in x])
+    next_page_token = get(result, "next_page_token", nothing)
+
+    return runs, next_page_token
+end

src/types/mlflow.jl

@@ -32,13 +32,10 @@ struct MLFlow
     headers::Dict
 end
 MLFlow(apiroot; apiversion=2.0, headers=Dict()) = MLFlow(apiroot, apiversion, headers)
-function MLFlow()
-    apiroot = "http://localhost:5000/api"
-    if haskey(ENV, "MLFLOW_TRACKING_URI")
-        apiroot = ENV["MLFLOW_TRACKING_URI"]
-    end
-    return MLFlow(apiroot)
-end
+MLFlow(; apiroot="http://localhost:5000/api", apiversion=2.0, headers=Dict()) =
+    MLFlow((haskey(ENV, "MLFLOW_TRACKING_URI") ?
+            ENV["MLFLOW_TRACKING_URI"] : apiroot), apiversion, headers)
+
 Base.show(io::IO, t::MLFlow) = show(io, ShowCase(t, [:apiroot,:apiversion], new_lines=true))
 
 abstract type LoggingData end

test/services/loggers.jl

@@ -102,7 +102,8 @@ end
         run = refresh(mlf, run)
         last_metric = run.data.metrics |> last
         last_param = run.data.params |> last
-        last_tag = run.data.tags |> first
+        last_tag = run.data.tags[
+            findall(x -> !occursin("mlflow.runName", x.key), run.data.tags)[1]]
 
         @test last_metric isa Metric
         @test last_metric.key == "ana"
@@ -127,7 +128,8 @@ end
         run = refresh(mlf, run)
         last_metric = run.data.metrics |> last
         last_param = run.data.params |> last
-        last_tag = run.data.tags |> first
+        last_tag = run.data.tags[
+            findall(x -> !occursin("mlflow.runName", x.key), run.data.tags)[1]]
 
         @test last_metric isa Metric
         @test last_metric.key == "ana"
@@ -152,7 +154,8 @@ end
         run = refresh(mlf, run)
         last_metric = run.data.metrics |> last
         last_param = run.data.params |> last
-        last_tag = run.data.tags |> first
+        last_tag = run.data.tags[
+            findall(x -> !occursin("mlflow.runName", x.key), run.data.tags)[1]]
 
         @test last_metric isa Metric
         @test last_metric.key == "ana"
@@ -179,7 +182,8 @@ end
         run = refresh(mlf, run)
         last_metric = run.data.metrics |> last
         last_param = run.data.params |> last
-        last_tag = run.data.tags |> first
+        last_tag = run.data.tags[
+            findall(x -> !occursin("mlflow.runName", x.key), run.data.tags)[1]]
 
         @test last_metric isa Metric
         @test last_metric.key == "ana"

test/services/run.jl

@@ -102,8 +102,10 @@ end
         run = refresh(mlf, run)
 
         @test run.data.tags |> !isempty
-        @test (run.data.tags |> first).key == "tag"
-        @test (run.data.tags |> first).value == "value"
+        last_tag = run.data.tags[
+            findall(x -> !occursin("mlflow.runName", x.key), run.data.tags)[1]]
+        @test last_tag.key == "tag"
+        @test last_tag.value == "value"
 
         deleterun(mlf, run)
     end
@@ -115,8 +117,10 @@ end
         run = refresh(mlf, run)
 
         @test run.data.tags |> !isempty
-        @test (run.data.tags |> first).key == "tag"
-        @test (run.data.tags |> first).value == "value"
+        last_tag = run.data.tags[
+            findall(x -> !occursin("mlflow.runName", x.key), run.data.tags)[1]]
+        @test last_tag.key == "tag"
+        @test last_tag.value == "value"
 
         deleterun(mlf, run)
     end
@@ -151,3 +155,33 @@ end
     end
     deleteexperiment(mlf, experiment_id)
 end
+
+@testset verbose = true "search runs" begin
+    @ensuremlf
+
+    experiment_ids = [
+        createexperiment(mlf, UUIDs.uuid4() |> string),
+        createexperiment(mlf, UUIDs.uuid4() |> string),
+    ]
+    for experiment_id in experiment_ids
+        createrun(mlf, experiment_id)
+    end
+
+    @testset "default search" begin
+        runs, next_page_token = searchruns(mlf; experiment_ids=experiment_ids)
+
+        @test length(runs) == 2
+        @test next_page_token |> isnothing
+    end
+
+    @testset "with pagination" begin
+        runs, next_page_token = searchruns(mlf; experiment_ids=experiment_ids,
+            max_results=1)
+
+        @test length(runs) == 1
+        @test next_page_token |> !isnothing
+        @test next_page_token isa String
+    end
+
+    experiment_ids .|> (id -> deleteexperiment(mlf, id))
+end