Skip to content

Commit 5452ee2

Browse files
pebetopebeto
committed
Implementing `setexperimenttag
1 parent 22767d4 commit 5452ee2

File tree

5 files changed

+105
-109
lines changed

5 files changed

+105
-109
lines changed

‎src/MLFlowClient.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ export
5858
createexperiment,
5959
deleteexperiment,
6060
updateexperiment,
61+
setexperimenttag,
6162
restoreexperiment,
6263
searchexperiments,
6364
getexperimentbyname

‎src/api.jl‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,13 @@ function mlfget(mlf, endpoint; kwargs...)
3636

3737
try
3838
response = HTTP.get(apiuri, apiheaders)
39-
return JSON.parse(String(response.body))
39+
return response.body |> String |> JSON.parse
4040
catch e
41-
throw(e)
41+
error_response = e.response.body |> String |> JSON.parse
42+
error_message =
43+
"$(error_response["error_code"]) - $(error_response["message"])"
44+
@error error_message
45+
throw(ErrorException(error_message))
4246
end
4347
end
4448

@@ -54,8 +58,12 @@ function mlfpost(mlf, endpoint; kwargs...)
5458

5559
try
5660
response = HTTP.post(apiuri, apiheaders, body)
57-
return JSON.parse(String(response.body))
61+
return response.body |> String |> JSON.parse
5862
catch e
59-
throw(e)
63+
error_response = e.response.body |> String |> JSON.parse
64+
error_message =
65+
"$(error_response["error_code"]) - $(error_response["message"])"
66+
@error error_message
67+
throw(ErrorException(error_message))
6068
end
6169
end

‎src/services/experiment.jl‎

Lines changed: 47 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -21,20 +21,9 @@ The ID of the newly created experiment.
2121
function createexperiment(instance::MLFlow, name::String;
2222
artifact_location::Union{String, Missing}=missing,
2323
tags::MLFlowUpsertData{Tag}=Tag[])::String
24-
try
25-
result = mlfpost(instance, "experiments/create"; name=name,
26-
artifact_location=artifact_location,
27-
tags=parse(Tag, tags))
28-
return result["experiment_id"]
29-
catch e
30-
if isa(e, HTTP.ExceptionRequest.StatusError) && e.status == 400
31-
error_code = (e.response.body |> String |> JSON.parse)["error_code"]
32-
if error_code == MLFLOW_ERROR_CODES.RESOURCE_ALREADY_EXISTS
33-
error("Experiment with name \"$name\" already exists")
34-
end
35-
end
36-
throw(e)
37-
end
24+
result = mlfpost(instance, "experiments/create"; name=name,
25+
artifact_location=artifact_location, tags=parse(Tag, tags))
26+
return result["experiment_id"]
3827
end
3928

4029
"""
@@ -51,13 +40,8 @@ Get metadata for an experiment. This method works on deleted experiments.
5140
An instance of type [`Experiment`](@ref).
5241
"""
5342
function getexperiment(instance::MLFlow, experiment_id::String)::Experiment
54-
try
55-
arguments = (:experiment_id => experiment_id,)
56-
result = mlfget(instance, "experiments/get"; arguments...)
57-
return result["experiment"] |> Experiment
58-
catch e
59-
throw(e)
60-
end
43+
result = mlfget(instance, "experiments/get"; experiment_id=experiment_id)
44+
return result["experiment"] |> Experiment
6145
end
6246
getexperiment(instance::MLFlow, experiment_id::Integer)::Experiment =
6347
getexperiment(instance, string(experiment_id))
@@ -80,16 +64,9 @@ An instance of type [`Experiment`](@ref).
8064
"""
8165
function getexperimentbyname(instance::MLFlow,
8266
experiment_name::String)::Experiment
83-
try
84-
arguments = (:experiment_name => experiment_name,)
85-
result = mlfget(instance, "experiments/get-by-name"; arguments...)
86-
return result["experiment"] |> Experiment
87-
catch e
88-
if isa(e, HTTP.ExceptionRequest.StatusError) && e.status == 404
89-
return missing
90-
end
91-
throw(e)
92-
end
67+
result = mlfget(instance, "experiments/get-by-name";
68+
experiment_name=experiment_name)
69+
return result["experiment"] |> Experiment
9370
end
9471

9572
"""
@@ -106,21 +83,11 @@ experiment are also deleted.
10683
- `experiment_id`: ID of the associated experiment.
10784
10885
# Returns
109-
11086
`true` if successful. Otherwise, raises exception.
11187
"""
11288
function deleteexperiment(instance::MLFlow, experiment_id::String)
113-
endpoint = "experiments/delete"
114-
try
115-
mlfpost(instance, endpoint; experiment_id=experiment_id)
116-
return true
117-
catch e
118-
if isa(e, HTTP.ExceptionRequest.StatusError) && e.status == 404
119-
# experiment already deleted
120-
return true
121-
end
122-
throw(e)
123-
end
89+
mlfpost(instance, "experiments/delete"; experiment_id=experiment_id)
90+
return true
12491
end
12592
deleteexperiment(instance::MLFlow, experiment_id::Integer) =
12693
deleteexperiment(instance, string(experiment_id))
@@ -141,23 +108,11 @@ underlying artifacts associated with experiment are also restored.
141108
- `experiment_id`: ID of the associated experiment.
142109
143110
# Returns
144-
145111
`true` if successful. Otherwise, raises exception.
146112
"""
147113
function restoreexperiment(instance::MLFlow, experiment_id::String)
148-
endpoint = "experiments/restore"
149-
try
150-
mlfpost(instance, endpoint; experiment_id=experiment_id)
151-
return true
152-
catch e
153-
if isa(e, HTTP.ExceptionRequest.StatusError) && e.status == 404
154-
error_code = JSON.parse(String(e.response.body))["error_code"]
155-
if error_code == MLFLOW_ERROR_CODES.RESOURCE_DOES_NOT_EXIST
156-
error("Experiment with id \"$experiment_id\" does not exist")
157-
end
158-
end
159-
throw(e)
160-
end
114+
mlfpost(instance, "experiments/restore"; experiment_id=experiment_id)
115+
return true
161116
end
162117
restoreexperiment(instance::MLFlow, experiment_id::Integer) =
163118
restoreexperiment(instance, string(experiment_id))
@@ -184,13 +139,9 @@ The new name must be unique.
184139
"""
185140
function updateexperiment(instance::MLFlow, experiment_id::String,
186141
new_name::String)
187-
endpoint = "experiments/update"
188-
try
189-
mlfpost(instance, endpoint; experiment_id=experiment_id, new_name=new_name)
190-
return true
191-
catch e
192-
throw(e)
193-
end
142+
mlfpost(instance, "experiments/update"; experiment_id=experiment_id,
143+
new_name=new_name)
144+
return true
194145
end
195146
updateexperiment(instance::MLFlow, experiment_id::Integer, new_name::String) =
196147
updateexperiment(instance, string(experiment_id), new_name)
@@ -219,23 +170,45 @@ unspecified, return only active experiments.
219170
"""
220171
function searchexperiments(instance::MLFlow; max_results::Int64=20000,
221172
page_token::String="", filter::String="", order_by::Array{String}=String[],
222-
view_type::ViewType=ACTIVE_ONLY)::Tuple{Array{Experiment}, Union{String, Nothing}}
223-
endpoint = "experiments/search"
173+
view_type::ViewType=ACTIVE_ONLY
174+
)::Tuple{Array{Experiment}, Union{String, Nothing}}
224175
parameters = (; max_results, page_token, filter,
225176
:view_type => view_type |> Integer)
226177

227178
if order_by |> !isempty
228179
parameters = (; order_by, parameters...)
229180
end
230181

231-
try
232-
result = mlfget(instance, endpoint; parameters...)
182+
result = mlfget(instance, "experiments/search"; parameters...)
233183

234-
experiments = result["experiments"] |> (x -> [Experiment(y) for y in x])
235-
next_page_token = get(result, "next_page_token", nothing)
184+
experiments = result["experiments"] |> (x -> [Experiment(y) for y in x])
185+
next_page_token = get(result, "next_page_token", nothing)
236186

237-
return experiments, next_page_token
238-
catch e
239-
throw(e)
240-
end
187+
return experiments, next_page_token
241188
end
189+
190+
"""
191+
setexperimenttag(instance::MLFlow, experiment_id::String, key::String,
192+
value::String)
193+
setexperimenttag(instance::MLFlow, experiment_id::Integer, key::String,
194+
value::String)
195+
setexperimenttag(instance::MLFlow, experiment::Experiment, key::String,
196+
value::String)
197+
198+
Set a tag on an experiment. Experiment tags are metadata that can be updated.
199+
200+
# Arguments
201+
- `experiment_id`: ID of the experiment under which to log the tag.
202+
- `key`: Name of the tag.
203+
- `value`: String value of the tag being logged.
204+
"""
205+
setexperimenttag(instance::MLFlow, experiment_id::String, key::String,
206+
value::String) =
207+
mlfpost(instance, "experiments/set-experiment-tag";
208+
experiment_id=experiment_id, key=key, value=value)
209+
setexperimenttag(instance::MLFlow, experiment_id::Integer, key::String,
210+
value::String) =
211+
setexperimenttag(instance, string(experiment_id), key, value)
212+
setexperimenttag(instance::MLFlow, experiment::Experiment, key::String,
213+
value::String) =
214+
setexperimenttag(instance, experiment.experiment_id, key, value)

‎src/services/run.jl‎

Lines changed: 9 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,9 @@ function createrun(instance::MLFlow, experiment_id::String;
2121
run_name::Union{String, Missing}=missing,
2222
start_time::Union{Int64, Missing}=missing,
2323
tags::MLFlowUpsertData{Tag}=Tag[])
24-
try
25-
result = mlfpost(instance, "runs/create"; experiment_id=experiment_id,
26-
run_name=run_name, start_time=start_time, tags=parse(Tag, tags))
27-
return result["run"] |> Run
28-
catch e
29-
throw(e)
30-
end
24+
result = mlfpost(instance, "runs/create"; experiment_id=experiment_id,
25+
run_name=run_name, start_time=start_time, tags=parse(Tag, tags))
26+
return result["run"] |> Run
3127
end
3228
createrun(instance::MLFlow, experiment_id::Integer;
3329
run_name::Union{String, Missing}=missing,
@@ -53,17 +49,11 @@ Mark a run for deletion.
5349
- `run_id`: ID of the run to delete.
5450
5551
# Returns
56-
5752
`true` if successful. Otherwise, raises exception.
5853
"""
5954
function deleterun(instance::MLFlow, run_id::String)
60-
endpoint = "runs/delete"
61-
try
62-
mlfpost(instance, endpoint; run_id=run_id)
63-
return true
64-
catch e
65-
throw(e)
66-
end
55+
mlfpost(instance, "runs/delete"; run_id=run_id)
56+
return true
6757
end
6858
deleterun(instance::MLFlow, run::Run) = deleterun(instance, run.info.run_id)
6959

@@ -78,17 +68,11 @@ Restore a deleted run.
7868
- `run_id`: ID of the run to restore.
7969
8070
# Returns
81-
8271
`true` if successful. Otherwise, raises exception.
8372
"""
8473
function restorerun(instance::MLFlow, run_id::String)
85-
endpoint = "runs/restore"
86-
try
87-
mlfpost(instance, endpoint; run_id=run_id)
88-
return true
89-
catch e
90-
throw(e)
91-
end
74+
mlfpost(instance, "runs/restore"; run_id=run_id)
75+
return true
9276
end
9377
restorerun(instance::MLFlow, run::Run) = restorerun(instance, run.info.run_id)
9478

@@ -108,11 +92,6 @@ return the maximum of these values.
10892
An instance of type [`Run`](@ref).
10993
"""
11094
function getrun(instance::MLFlow, run_id::String)
111-
try
112-
arguments = (:run_id => run_id,)
113-
result = mlfget(instance, "runs/get"; arguments...)
114-
return result["run"] |> Run
115-
catch e
116-
throw(e)
117-
end
95+
result = mlfget(instance, "runs/get"; run_id=run_id)
96+
return result["run"] |> Run
11897
end

‎test/services/experiment.jl‎

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ end
9393

9494
@testset "delete already deleted" begin
9595
deleteexperiment(mlf, experiment_id)
96-
@test deleteexperiment(mlf, experiment_id)
96+
@test_throws ErrorException deleteexperiment(mlf, experiment_id)
9797
end
9898
end
9999

@@ -179,3 +179,38 @@ end
179179

180180
experiment_ids .|> (id -> deleteexperiment(mlf, id))
181181
end
182+
183+
@testset verbose = true "set experiment tag" begin
184+
@ensuremlf
185+
186+
@testset "set tag with string id" begin
187+
experiment_id = createexperiment(mlf, UUIDs.uuid4() |> string)
188+
setexperimenttag(mlf, experiment_id, "test_key", "test_value")
189+
experiment = getexperiment(mlf, experiment_id)
190+
@test experiment.tags |> !isempty
191+
@test (experiment.tags |> first).key == "test_key"
192+
@test (experiment.tags |> first).value == "test_value"
193+
deleteexperiment(mlf, experiment_id)
194+
end
195+
196+
@testset "set tag with integer id" begin
197+
experiment_id = createexperiment(mlf, UUIDs.uuid4() |> string)
198+
setexperimenttag(mlf, parse(Int, experiment_id), "test_key", "test_value")
199+
experiment = getexperiment(mlf, experiment_id)
200+
@test experiment.tags |> !isempty
201+
@test (experiment.tags |> first).key == "test_key"
202+
@test (experiment.tags |> first).value == "test_value"
203+
deleteexperiment(mlf, experiment_id)
204+
end
205+
206+
@testset "set tag with Experiment" begin
207+
experiment_id = createexperiment(mlf, UUIDs.uuid4() |> string)
208+
experiment = getexperiment(mlf, experiment_id)
209+
setexperimenttag(mlf, experiment, "test_key", "test_value")
210+
experiment = getexperiment(mlf, experiment_id)
211+
@test experiment.tags |> !isempty
212+
@test (experiment.tags |> first).key == "test_key"
213+
@test (experiment.tags |> first).value == "test_value"
214+
deleteexperiment(mlf, experiment_id)
215+
end
216+
end

0 commit comments

Comments
 (0)