Add settag

src/MLFlowClient.jl

@@ -79,7 +79,8 @@ export
     logparam,
     logmetric,
     logartifact,
-    listartifacts
+    listartifacts,
+    settag
 
 include("deprecated.jl")
 

src/loggers.jl

@@ -1,3 +1,36 @@
+"""
+    settag(mlf::MLFlow, run, key, value)
+    settag(mlf::MLFlow, run, kv)
+
+Associates a tag (a key and a value) to the particular run.
+
+Refer to [the official MLflow REST API
+docs](https://mlflow.org/docs/latest/rest-api.html#set-tag) for restrictions on
+`key` and `value`.
+
+# Arguments
+- `mlf`: [`MLFlow`](@ref) configuration.
+- `run`: one of [`MLFlowRun`](@ref), [`MLFlowRunInfo`](@ref), or `String`.
+- `key`: tag key (name). Automatically converted to string before sending to MLFlow because this is the only type that MLFlow supports.
+- `value`: parameter value. Automatically converted to string before sending to MLFlow because this is the only type that MLFlow supports.
+
+One could also specify `kv::Dict` instead of separate `key` and `value` arguments.
+"""
+function settag(mlf::MLFlow, run_id::String, key, value)
+    endpoint ="runs/set-tag"
+    mlfpost(mlf, endpoint; run_id=run_id, key=string(key), value=string(value))
+end
+settag(mlf::MLFlow, run_info::MLFlowRunInfo, key, value) =
+    settag(mlf, run_info.run_id, key, value)
+settag(mlf::MLFlow, run::MLFlowRun, key, value) =
+    settag(mlf, run.info, key, value)
+function settag(mlf::MLFlow, run::Union{String,MLFlowRun,MLFlowRunInfo}, kv)
+    for (k, v) in kv
+        logparam(mlf, run, k, v)
+    end
+end
+
+
 """
     logparam(mlf::MLFlow, run, key, value)
     logparam(mlf::MLFlow, run, kv)

test/test_loggers.jl

@@ -1,3 +1,41 @@
+@testset verbose = true "settag" begin
+    @ensuremlf
+    expname = "settag-$(UUIDs.uuid4())"
+    e = getorcreateexperiment(mlf, expname)
+    runname = "run-$(UUIDs.uuid4())"
+    r = createrun(mlf, e.experiment_id)
+
+    @testset "settag_by_run_id_and_key_value" begin
+        settag(mlf, r.info.run_id, "run_id_key_value", "test")
+        retrieved_run = searchruns(mlf, e; filter="tags.run_id_key_value = 'test'")
+        @test length(retrieved_run) == 1
+        @test retrieved_run[1].info.run_id == r.info.run_id
+    end
+
+    @testset "settag_by_run_info_and_key_value" begin
+        settag(mlf, r.info, "run_id_key_value", "test")
+        retrieved_run = searchruns(mlf, e; filter="tags.run_id_key_value = 'test'")
+        @test length(retrieved_run) == 1
+        @test retrieved_run[1].info.run_id == r.info.run_id
+    end
+
+    @testset "settag_by_run_and_key_value" begin
+        settag(mlf, r, "run_id_key_value", "test")
+        retrieved_run = searchruns(mlf, e; filter="tags.run_id_key_value = 'test'")
+        @test length(retrieved_run) == 1
+        @test retrieved_run[1].info.run_id == r.info.run_id
+    end
+
+    @testset "settag_by_union_and_dict_key_value" begin
+        settag(mlf, r, Dict("run_id_key_value" => "test"))
+        retrieved_run = searchruns(mlf, e; filter="tags.run_id_key_value = 'test'")
+        @test length(retrieved_run) == 1
+        @test retrieved_run[1].info.run_id == r.info.run_id
+    end
+
+    deleteexperiment(mlf, e)
+end
+
 @testset verbose = true "logparam" begin
     @ensuremlf
     expname = "logparam-$(UUIDs.uuid4())"