Implement setruntag

Author: pebeto

src/MLFlowClient.jl

@@ -57,22 +57,23 @@ export
     getexperiment,
     createexperiment,
     deleteexperiment,
-    updateexperiment,
     setexperimenttag,
+    updateexperiment,
     restoreexperiment,
     searchexperiments,
     getexperimentbyname
 
 include("services/run.jl")
 export 
+    getrun,
     createrun,
     deleterun,
-    restorerun,
-    getrun
+    setruntag,
+    restorerun
 include("services/loggers.jl")
 export
-    logmetric,
     logbatch,
-    loginputs
+    loginputs,
+    logmetric
 
 end

src/services/experiment.jl

@@ -85,13 +85,13 @@ experiment are also deleted.
 # Returns
 `true` if successful. Otherwise, raises exception.
 """
-function deleteexperiment(instance::MLFlow, experiment_id::String)
+function deleteexperiment(instance::MLFlow, experiment_id::String)::Bool
     mlfpost(instance, "experiments/delete"; experiment_id=experiment_id)
     return true
 end
-deleteexperiment(instance::MLFlow, experiment_id::Integer) =
+deleteexperiment(instance::MLFlow, experiment_id::Integer)::Bool =
     deleteexperiment(instance, string(experiment_id))
-deleteexperiment(instance::MLFlow, experiment::Experiment) =
+deleteexperiment(instance::MLFlow, experiment::Experiment)::Bool =
     deleteexperiment(instance, experiment.experiment_id)
 
 """
@@ -110,13 +110,13 @@ underlying artifacts associated with experiment are also restored.
 # Returns
 `true` if successful. Otherwise, raises exception.
 """
-function restoreexperiment(instance::MLFlow, experiment_id::String)
+function restoreexperiment(instance::MLFlow, experiment_id::String)::Bool
     mlfpost(instance, "experiments/restore"; experiment_id=experiment_id)
     return true
 end
-restoreexperiment(instance::MLFlow, experiment_id::Integer) =
+restoreexperiment(instance::MLFlow, experiment_id::Integer)::Bool =
     restoreexperiment(instance, string(experiment_id))
-restoreexperiment(instance::MLFlow, experiment::Experiment) =
+restoreexperiment(instance::MLFlow, experiment::Experiment)::Bool =
     restoreexperiment(instance, experiment.experiment_id)
 
 """
@@ -138,15 +138,17 @@ The new name must be unique.
 `true` if successful. Otherwise, raises exception.
 """
 function updateexperiment(instance::MLFlow, experiment_id::String,
-    new_name::String)
+    new_name::String)::Bool
     mlfpost(instance, "experiments/update"; experiment_id=experiment_id,
         new_name=new_name)
     return true
 end
-updateexperiment(instance::MLFlow, experiment_id::Integer, new_name::String) =
+updateexperiment(instance::MLFlow, experiment_id::Integer,
+    new_name::String)::Bool =
     updateexperiment(instance, string(experiment_id), new_name)
-updateexperiment(instance::MLFlow, experiment::Experiment, new_name::String) =
-    updateexperiment(instance, experiment.experiment_id, new_name::String)
+updateexperiment(instance::MLFlow, experiment::Experiment,
+    new_name::String)::Bool =
+    updateexperiment(instance, experiment.experiment_id, new_name)
 
 """
     searchexperiments(instance::MLFlow; max_results::Int64=20000,
@@ -201,14 +203,19 @@ Set a tag on an experiment. Experiment tags are metadata that can be updated.
 - `experiment_id`: ID of the experiment under which to log the tag.
 - `key`: Name of the tag.
 - `value`: String value of the tag being logged.
+
+# Returns
+`true` if successful. Otherwise, raises exception.
 """
-setexperimenttag(instance::MLFlow, experiment_id::String, key::String,
-    value::String) =
+function setexperimenttag(instance::MLFlow, experiment_id::String, key::String,
+    value::String)::Bool
     mlfpost(instance, "experiments/set-experiment-tag";
         experiment_id=experiment_id, key=key, value=value)
+    return true
+end
 setexperimenttag(instance::MLFlow, experiment_id::Integer, key::String,
-    value::String) =
+    value::String)::Bool =
     setexperimenttag(instance, string(experiment_id), key, value)
 setexperimenttag(instance::MLFlow, experiment::Experiment, key::String,
-    value::String) =
+    value::String)::Bool =
     setexperimenttag(instance, experiment.experiment_id, key, value)

src/services/loggers.jl

@@ -17,15 +17,22 @@ represent ML model accuracy. A metric can be logged multiple times.
 - `value`: Double value of the metric being logged.
 - `timestamp`: Unix timestamp in milliseconds at the time metric was logged.
 - `step`: Step at which to log the metric.
+
+# Returns
+`true` if successful. Otherwise, raises exception.
 """
-logmetric(instance::MLFlow, run_id::String, key::String, value::Float64;
-    timestamp::Int64=round(Int, now() |> datetime2unix),
-    step::Union{Int64, Missing}=missing) =
-    mlfpost(instance, "runs/log-metric"; run_id=run_id, key=key, value=value, timestamp=timestamp, step=step)
+function logmetric(instance::MLFlow, run_id::String, key::String,
+    value::Float64; timestamp::Int64=round(Int, now() |> datetime2unix),
+    step::Union{Int64, Missing}=missing)::Bool
+    mlfpost(instance, "runs/log-metric"; run_id=run_id, key=key, value=value,
+        timestamp=timestamp, step=step)
+    return true
+end
 logmetric(instance::MLFlow, run::Run, key::String, value::Float64;
     timestamp::Int64=round(Int, now() |> datetime2unix),
-    step::Union{Int64, Missing}=missing) =
-    logmetric(instance, run.info.run_id, key, value; timestamp=timestamp, step=step)
+    step::Union{Int64, Missing}=missing)::Bool =
+    logmetric(instance, run.info.run_id, key, value; timestamp=timestamp,
+        step=step)
 
 """
     logbatch(instance::MLFlow, run_id::String;
@@ -38,18 +45,33 @@ Log a batch of metrics, params, and tags for a run. In case of error, partial
 data may be written.
 
 For more information about this function, check [MLFlow official documentation](https://mlflow.org/docs/latest/rest-api.html#log-batch).
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `run_id`: ID of the run to log under.
+- `metrics`: Metrics to log.
+- `params`: Params to log.
+- `tags`: Tags to log.
+
+**Note**: A single request can contain up to 1000 metrics, and up to 1000
+metrics, params, and tags in total.
+
+# Returns
+`true` if successful. Otherwise, raises exception.
 """
-logbatch(instance::MLFlow, run_id::String;
+function logbatch(instance::MLFlow, run_id::String;
     metrics::MLFlowUpsertData{Metric}=Metric[],
     params::MLFlowUpsertData{Param}=Param[],
-    tags::MLFlowUpsertData{Tag}=Tag[]) =
+    tags::MLFlowUpsertData{Tag}=Tag[])::Bool
     mlfpost(instance, "runs/log-batch"; run_id=run_id,
         metrics=parse(Metric, metrics), params=parse(Param, params),
         tags=parse(Tag, tags))
+    return true
+end
 logbatch(instance::MLFlow, run::Run;
     metrics::MLFlowUpsertData{Metric}=Metric[],
     params::MLFlowUpsertData{Param}=Param[],
-    tags::MLFlowUpsertData{Tag}=Tag[]) =
+    tags::MLFlowUpsertData{Tag}=Tag[])::Bool =
     logbatch(instance, run.info.run_id; metrics=metrics, params=params,
         tags=tags)
 
@@ -61,8 +83,14 @@ logbatch(instance::MLFlow, run::Run;
 - `instance`: [`MLFlow`](@ref) configuration.
 - `run_id`: ID of the run to log under This field is required.
 - `datasets`: Dataset inputs.
+
+# Returns
+`true` if successful. Otherwise, raises exception.
 """
-loginputs(instance::MLFlow, run_id::String, datasets::Array{DatasetInput}) =
+function loginputs(instance::MLFlow, run_id::String,
+    datasets::Array{DatasetInput})::Bool
     mlfpost(instance, "runs/log-inputs"; run_id=run_id, datasets=datasets)
-loginputs(instance::MLFlow, run::Run, datasets::Array{DatasetInput}) =
+    return true
+end
+loginputs(instance::MLFlow, run::Run, datasets::Array{DatasetInput})::Bool =
     loginputs(instance, run.info.run_id, datasets)

src/services/run.jl

@@ -20,21 +20,21 @@ An instance of type [`Run`](@ref).
 function createrun(instance::MLFlow, experiment_id::String;
     run_name::Union{String, Missing}=missing,
     start_time::Union{Int64, Missing}=missing,
-    tags::MLFlowUpsertData{Tag}=Tag[])
+    tags::MLFlowUpsertData{Tag}=Tag[])::Run
     result = mlfpost(instance, "runs/create"; experiment_id=experiment_id,
         run_name=run_name, start_time=start_time, tags=parse(Tag, tags))
     return result["run"] |> Run
 end
 createrun(instance::MLFlow, experiment_id::Integer;
     run_name::Union{String, Missing}=missing,
     start_time::Union{Integer, Missing}=missing,
-    tags::MLFlowUpsertData{Tag}=Tag[]) =
+    tags::MLFlowUpsertData{Tag}=Tag[])::Run =
     createrun(instance, string(experiment_id); run_name=run_name,
         start_time=start_time, tags=tags)
 createrun(instance::MLFlow, experiment::Experiment;
     run_name::Union{String, Missing}=missing,
     start_time::Union{Integer, Missing}=missing,
-    tags::MLFlowUpsertData{Tag}=Tag[]) =
+    tags::MLFlowUpsertData{Tag}=Tag[])::Run =
     createrun(instance, string(experiment.experiment_id); run_name=run_name,
         start_time=start_time, tags=tags)
 
@@ -51,11 +51,12 @@ Mark a run for deletion.
 # Returns
 `true` if successful. Otherwise, raises exception.
 """
-function deleterun(instance::MLFlow, run_id::String)
+function deleterun(instance::MLFlow, run_id::String)::Bool
     mlfpost(instance, "runs/delete"; run_id=run_id)
     return true
 end
-deleterun(instance::MLFlow, run::Run) = deleterun(instance, run.info.run_id)
+deleterun(instance::MLFlow, run::Run)::Bool =
+    deleterun(instance, run.info.run_id)
 
 """
     restorerun(instance::MLFlow, run_id::String)
@@ -70,11 +71,12 @@ Restore a deleted run.
 # Returns
 `true` if successful. Otherwise, raises exception.
 """
-function restorerun(instance::MLFlow, run_id::String)
+function restorerun(instance::MLFlow, run_id::String)::Bool
     mlfpost(instance, "runs/restore"; run_id=run_id)
     return true
 end
-restorerun(instance::MLFlow, run::Run) = restorerun(instance, run.info.run_id)
+restorerun(instance::MLFlow, run::Run)::Bool =
+    restorerun(instance, run.info.run_id)
 
 """
     getrun(instance::MLFlow, run_id::String)
@@ -91,7 +93,31 @@ return the maximum of these values.
 # Returns
 An instance of type [`Run`](@ref).
 """
-function getrun(instance::MLFlow, run_id::String)
+function getrun(instance::MLFlow, run_id::String)::Run
     result = mlfget(instance, "runs/get"; run_id=run_id)
     return result["run"] |> Run
 end
+
+"""
+    setruntag(instance::MLFlow, run_id::String, key::String, value::String)
+    setruntag(instance::MLFlow, run::Run, key::String, value::String)
+
+Set a tag on a run. Tags are run metadata that can be updated during a run and
+after a run completes.
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `run_id`: ID of the run under which to log the tag.
+- `key`: Name of the tag.
+- `value`: String value of the tag being logged.
+
+# Returns
+`true` if successful. Otherwise, raises exception.
+"""
+function setruntag(instanceL::MLFlow, run_id::String, key::String,
+    value::String):Bool
+    mlfpost(instanceL, "runs/set-tag"; run_id=run_id, key=key, value=value)
+    return true
+end
+setruntag(instance::MLFlow, run::Run, key::String, value::String)::Bool =
+    setruntag(instance, run.info.run_id, key, value)

test/services/run.jl

@@ -90,3 +90,36 @@ end
 
     deleteexperiment(mlf, experiment_id)
 end
+
+@testset verbose = true "set run tag" begin
+    @ensuremlf
+    experiment_id = createexperiment(mlf, UUIDs.uuid4() |> string)
+
+    @testset "set tag with run string id" begin
+        run = createrun(mlf, experiment_id)
+        setruntag(mlf, run.info.run_id, "tag", "value")
+
+        run = refresh(mlf, run)
+
+        @test run.data.tags |> !isempty
+        @test (run.data.tags |> first).key == "tag"
+        @test (run.data.tags |> first).value == "value"
+
+        deleterun(mlf, run)
+    end
+
+    @testset "set tag with run" begin
+        run = createrun(mlf, experiment_id)
+        setruntag(mlf, run, "tag", "value")
+
+        run = refresh(mlf, run)
+
+        @test run.data.tags |> !isempty
+        @test (run.data.tags |> first).key == "tag"
+        @test (run.data.tags |> first).value == "value"
+
+        deleterun(mlf, run)
+    end
+
+    deleteexperiment(mlf, experiment_id)
+end