From 20eed33a519a8ba1d4214969260dc0024161b332 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20P=C3=A4tzel?= <david.paetzel@posteo.de>
Date: Sat, 6 Jul 2024 16:52:57 +0200
Subject: [PATCH] Add `settag`

---
 src/MLFlowClient.jl  |  3 ++-
 src/loggers.jl       | 33 +++++++++++++++++++++++++++++++++
 test/test_loggers.jl | 38 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/src/MLFlowClient.jl b/src/MLFlowClient.jl
index b9a179a..18c7ddb 100644
--- a/src/MLFlowClient.jl
+++ b/src/MLFlowClient.jl
@@ -79,7 +79,8 @@ export
     logparam,
     logmetric,
     logartifact,
-    listartifacts
+    listartifacts,
+    settag
 
 include("deprecated.jl")
 
diff --git a/src/loggers.jl b/src/loggers.jl
index 183befd..b8f0639 100644
--- a/src/loggers.jl
+++ b/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)
diff --git a/test/test_loggers.jl b/test/test_loggers.jl
index d9871b6..bf3ea00 100644
--- a/test/test_loggers.jl
+++ b/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())"