Implementing updaterun

Author: pebeto

src/MLFlowClient.jl

@@ -45,7 +45,8 @@ export getexperiment, createexperiment, deleteexperiment, setexperimenttag,
     updateexperiment, restoreexperiment, searchexperiments, getexperimentbyname
 
 include("services/run.jl")
-export getrun, createrun, deleterun, setruntag, restorerun, searchruns, deleteruntag
+export getrun, createrun, deleterun, setruntag, updaterun, restorerun, searchruns,
+    deleteruntag
 
 include("services/loggers.jl")
 export logbatch, loginputs, logmetric, logparam

src/services/run.jl

@@ -184,3 +184,33 @@ function searchruns(instance::MLFlow; experiment_ids::Array{String}=String[],
 
     return runs, next_page_token
 end
+
+"""
+    updaterun(instance::MLFlow, run_id::String; status::Union{RunStatus, Missing}=missing,
+        end_time::Union{Int64, Missing}=missing, run_name::Union{String, Missing}=missing)
+    updaterun(instance::MLFlow, run::Run; status::Union{RunStatus, Missing}=missing,
+        end_time::Union{Int64, Missing}=missing, run_name::Union{String, Missing}=missing)
+
+Update run metadata.
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `run_id`: ID of the run to update.
+- `status`: Updated status of the run.
+- `end_time`: Unix timestamp in milliseconds of when the run ended.
+- `run_name`: Updated name of the run.
+
+# Returns
+- An instance of type [`RunInfo`](@ref) with the updated metadata.
+"""
+function updaterun(instance::MLFlow, run_id::String;
+    status::Union{RunStatus, Missing}=missing, end_time::Union{Int64, Missing}=missing,
+    run_name::Union{String, Missing})::RunInfo
+    result = mlfpost(instance, "runs/update"; run_id=run_id, status=(status |> Integer),
+        end_time=end_time, run_name=run_name)
+    return result["run_info"] |> RunInfo
+end
+updaterun(instance::MLFlow, run::Run; status::Union{RunStatus, Missing}=missing,
+    end_time::Union{Int64, Missing}=missing, run_name::Union{String, Missing})::RunInfo =
+    updaterun(instance, run.info.run_id; status=status, end_time=end_time,
+        run_name=run_name)

test/services/run.jl

@@ -185,3 +185,35 @@ end
 
     experiment_ids .|> (id -> deleteexperiment(mlf, id))
 end
+
+@testset verbose = true "update run" begin
+    @ensuremlf
+    experiment_id = createexperiment(mlf, UUIDs.uuid4() |> string)
+    run = createrun(mlf, experiment_id)
+
+    @testset "update with string id" begin
+        status = MLFlowClient.FINISHED
+        end_time = 123
+        run_name = "missy"
+
+        run_info = updaterun(mlf, run.info.run_id; status=status, end_time=end_time, run_name=run_name)
+
+        @test run_info.status == status
+        @test run_info.end_time == end_time
+        @test run_info.run_name == run_name
+    end
+
+    @testset "update with Run" begin
+        status = MLFlowClient.FAILED
+        end_time = 456
+        run_name = "gala"
+
+        run_info = updaterun(mlf, run.info.run_id; status=status, end_time=end_time, run_name=run_name)
+
+        @test run_info.status == status
+        @test run_info.end_time == end_time
+        @test run_info.run_name == run_name
+    end
+
+    deleteexperiment(mlf, experiment_id)
+end