Adding artifact_location modification to getorcreateexperiment (#26)

* Adding artifact_location modification to getorcreateexperiment function (and updating LICENSE year)

* updating run creation/edition functions adding run_name. misc on test

* Updating CI mlflow image to the last version

Author: pebeto

.github/workflows/CI.yml

@@ -15,7 +15,7 @@ jobs:
 
     services:
       mlflow:
-        image: adacotechjp/mlflow:1.21.0
+        image: adacotechjp/mlflow:2.3.1
         ports:
           - 5000:5000
 

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2021-2021 @deyandyankov and contributors
+Copyright (c) 2021-2023 @deyandyankov and contributors
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

src/experiments.jl

@@ -83,22 +83,24 @@ function _getexperimentbyname(mlf::MLFlow, experiment_name::String)
 end
 
 """
-    getorcreateexperiment(mlf::MLFlow, experiment_name::String)
+    getorcreateexperiment(mlf::MLFlow, experiment_name::String; artifact_location=missing, tags=missing)
 
 Gets an experiment if one alrady exists, or creates a new one.
 
 # Arguments
 - `mlf`: [`MLFlow`](@ref) configuration.
 - `experiment_name`: Experiment name.
+- `artifact_location`: directory where artifacts of this experiment will be stored. If not specified, MLFlow uses its default configuration.
+- `tags`: a Dictionary of key-values which tag the experiment.
 
 # Returns
 An instance of type [`MLFlowExperiment`](@ref)
 
 """
-function getorcreateexperiment(mlf::MLFlow, experiment_name::String)
+function getorcreateexperiment(mlf::MLFlow, experiment_name::String; artifact_location=missing, tags=missing)
     exp = getexperiment(mlf, experiment_name)
     if ismissing(exp)
-        exp = createexperiment(mlf, name=experiment_name)
+        exp = createexperiment(mlf, name=experiment_name, artifact_location=artifact_location, tags=tags)
     end
     exp
 end

src/runs.jl

@@ -1,5 +1,5 @@
 """
-    createrun(mlf::MLFlow, experiment_id; start_time=missing, tags=missing)
+    createrun(mlf::MLFlow, experiment_id; run_name=missing, start_time=missing, tags=missing)
 
 Creates a run associated to an experiment.
 
@@ -8,27 +8,28 @@ Creates a run associated to an experiment.
 - `experiment_id`: experiment identifier.
 
 # Keywords
+- `run_name`: run name. If not specified, MLFlow sets it.
 - `start_time`: if provided, must be a UNIX timestamp in milliseconds. By default, set to current time.
 - `tags`: if provided, must be a key-value structure such as a dictionary.
 
 # Returns
 - an instance of type [`MLFlowRun`](@ref)
 """
-function createrun(mlf::MLFlow, experiment_id; start_time=missing, tags=missing)
+function createrun(mlf::MLFlow, experiment_id; run_name=missing, start_time=missing, tags=missing)
     endpoint = "runs/create"
     if ismissing(start_time)
         start_time = Int(trunc(datetime2unix(now(UTC)) * 1000))
     end
-    result = mlfpost(mlf, endpoint; experiment_id=experiment_id, start_time=start_time, tags=tags)
+    result = mlfpost(mlf, endpoint; experiment_id=experiment_id, run_name=run_name, start_time=start_time, tags=tags)
     MLFlowRun(result["run"]["info"], result["run"]["data"])
 end
 """
-    createrun(mlf::MLFlow, experiment::MLFlowExperiment; start_time=missing, tags=missing)
+    createrun(mlf::MLFlow, experiment::MLFlowExperiment; run_name=missing, start_time=missing, tags=missing)
 
-Dispatches to `createrun(mlf::MLFlow, experiment_id; start_time=start_time, tags=tags)`
+Dispatches to `createrun(mlf::MLFlow, experiment_id; run_name=run_name, start_time=start_time, tags=tags)`
 """
-createrun(mlf::MLFlow, experiment::MLFlowExperiment; start_time=missing, tags=missing) =
-    createrun(mlf, experiment.experiment_id; start_time=start_time, tags=tags)
+createrun(mlf::MLFlow, experiment::MLFlowExperiment; run_name=missing, start_time=missing, tags=missing) =
+    createrun(mlf, experiment.experiment_id; run_name=run_name, start_time=start_time, tags=tags)
 
 """
     getrun(mlf::MLFlow, run_id)
@@ -59,13 +60,15 @@ Updates the status of an experiment's run.
 - `status`: either `String` and one of ["RUNNING", "SCHEDULED", "FINISHED", "FAILED", "KILLED"], or an instance of `MLFlowRunStatus`
 
 # Keywords
+- `run_name`: if provided, must be a String. By default, not set.
 - `end_time`: if provided, must be a UNIX timestamp in milliseconds. By default, set to current time.
 """
-function updaterun(mlf::MLFlow, run_id::String, status::MLFlowRunStatus; end_time=missing)
+function updaterun(mlf::MLFlow, run_id::String, status::MLFlowRunStatus; run_name=missing, end_time=missing)
     endpoint = "runs/update"
     kwargs = Dict(
         :run_id => run_id,
         :status => status.status,
+        :run_name => run_name,
         :end_time => end_time
     )
     if ismissing(end_time) && status.status == "FINISHED"
@@ -75,16 +78,16 @@ function updaterun(mlf::MLFlow, run_id::String, status::MLFlowRunStatus; end_tim
     result = mlfpost(mlf, endpoint; kwargs...)
     MLFlowRun(result["run_info"])
 end
-updaterun(mlf::MLFlow, run_id::String, status::String; end_time=missing) =
-    updaterun(mlf, run_id, MLFlowRunStatus(status); end_time=end_time)
-updaterun(mlf::MLFlow, run_info::MLFlowRunInfo, status::String; end_time=missing) =
-    updaterun(mlf, run_info.run_id, MLFlowRunStatus(status); end_time=end_time)
-updaterun(mlf::MLFlow, run::MLFlowRun, status::String; end_time=missing) =
-    updaterun(mlf, run.info, MLFlowRunStatus(status), end_time=end_time)
-updaterun(mlf::MLFlow, run_info::MLFlowRunInfo, status::MLFlowRunStatus; end_time=missing) =
-    updaterun(mlf, run_info.run_id, status, end_time=end_time)
-updaterun(mlf::MLFlow, run::MLFlowRun, status::MLFlowRunStatus; end_time=missing) =
-    updaterun(mlf, run.info, status; end_time=end_time)
+updaterun(mlf::MLFlow, run_id::String, status::String; run_name=missing, end_time=missing) =
+    updaterun(mlf, run_id, MLFlowRunStatus(status); run_name=run_name, end_time=end_time)
+updaterun(mlf::MLFlow, run_info::MLFlowRunInfo, status::String; run_name=missing, end_time=missing) =
+    updaterun(mlf, run_info.run_id, MLFlowRunStatus(status); run_name=run_name, end_time=end_time)
+updaterun(mlf::MLFlow, run::MLFlowRun, status::String; run_name=missing, end_time=missing) =
+    updaterun(mlf, run.info, MLFlowRunStatus(status); run_name=run_name, end_time=end_time)
+updaterun(mlf::MLFlow, run_info::MLFlowRunInfo, status::MLFlowRunStatus; run_name=missing, end_time=missing) =
+    updaterun(mlf, run_info.run_id, status; run_name=run_name, end_time=end_time)
+updaterun(mlf::MLFlow, run::MLFlowRun, status::MLFlowRunStatus; run_name=missing, end_time=missing) =
+    updaterun(mlf, run.info, status; run_name=run_name, end_time=end_time)
 
 """
     deleterun(mlf::MLFlow, run)

src/types.jl

@@ -100,20 +100,22 @@ Represents run metadata.
 - `run_id::String`: run identifier.
 - `experiment_id::Integer`: experiment identifier.
 - `status::MLFlowRunStatus`: run status.
+- `run_name::String`: run name.
 - `start_time::Union{Int64,Missing}`: when was the run started, UNIX time in milliseconds.
 - `end_time::Union{Int64,Missing}`: when did the run end, UNIX time in milliseconds.
 - `artifact_uri::String`: where are artifacts from this run stored.
 - `lifecycle_stage::String`: one of `active` or `deleted`.
 
 # Constructors
 
-- `MLFlowRunInfo(run_id, experiment_id, status, start_time, end_time, artifact_uri, lifecycle_stage)`
+- `MLFlowRunInfo(run_id, experiment_id, status, run_name, start_time, end_time, artifact_uri, lifecycle_stage)`
 - `MLFlowRunInfo(info::Dict{String,Any})`
 """
 struct MLFlowRunInfo
     run_id::String
     experiment_id::Integer
     status::MLFlowRunStatus
+    run_name::String
     start_time::Union{Int64,Missing}
     end_time::Union{Int64,Missing}
     artifact_uri::String
@@ -123,6 +125,7 @@ function MLFlowRunInfo(info::Dict{String,Any})
     run_id = get(info, "run_id", missing)
     experiment_id = get(info, "experiment_id", missing)
     status = get(info, "status", missing)
+    run_name = get(info, "run_name", missing)
     start_time = get(info, "start_time", missing)
     end_time = get(info, "end_time", missing)
     artifact_uri = get(info, "artifact_uri", "")
@@ -138,7 +141,7 @@ function MLFlowRunInfo(info::Dict{String,Any})
     if !ismissing(end_time) && !(typeof(end_time) <: Int)
         end_time = parse(Int64, end_time)
     end
-    MLFlowRunInfo(run_id, experiment_id, status, start_time, end_time, artifact_uri, lifecycle_stage)
+    MLFlowRunInfo(run_id, experiment_id, status, run_name, start_time, end_time, artifact_uri, lifecycle_stage)
 end
 Base.show(io::IO, t::MLFlowRunInfo) = show(io, ShowCase(t, new_lines=true))
 get_run_id(runinfo::MLFlowRunInfo) = runinfo.run_id

test/issues/test_issue17.jl

@@ -1,2 +1 @@
 include("test_functional.jl")
-include.(filter(contains(r"\.jl$"), readdir("./issues"; join=true)))

test/issues/test_issue20.jl

@@ -61,25 +61,55 @@ end
 
 @testset "createrun" begin
     @ensuremlf
-    expname = "getorcreate-$(UUIDs.uuid4())"
+    expname = "createrun-$(UUIDs.uuid4())"
     e = getorcreateexperiment(mlf, expname)
-    r = createrun(mlf, e.experiment_id)
+    runname = "run-$(UUIDs.uuid4())"
+    r = createrun(mlf, e.experiment_id; run_name=runname)
+
     @test isa(r, MLFlowRun)
+    @test r.info.run_name == runname
+    deleteexperiment(mlf, e)
+end
+
+@testset "deleterun" begin
+    @ensuremlf
+    expname = "deleterun-$(UUIDs.uuid4())"
+    e = getorcreateexperiment(mlf, expname)
+    r = createrun(mlf, e.experiment_id)
+
     @test deleterun(mlf, r)
-    rr = createrun(mlf, e)
-    @test isa(rr, MLFlowRun)
-    @test deleterun(mlf, rr)
-    @test deleteexperiment(mlf, e)
+    deleteexperiment(mlf, e)
+end
+
+@testset "updaterun" begin
+    @ensuremlf
+    expname = "updaterun-$(UUIDs.uuid4())"
+    e = getorcreateexperiment(mlf, expname)
+    runname = "run-$(UUIDs.uuid4())"
+    r = createrun(mlf, e.experiment_id; run_name=runname)
+
+    new_runname = "new_updaterun-$(UUIDs.uuid4())"
+    new_status = "FINISHED"
+    r_updated = updaterun(mlf, r, new_status; run_name=new_runname)
+
+    @test isa(r_updated, MLFlowRun)
+    @test r_updated.info.run_name != r.info.run_name
+    @test r_updated.info.status.status != r.info.status
+    @test r_updated.info.run_name == new_runname
+    @test r_updated.info.status.status == new_status
+    deleteexperiment(mlf, e)
 end
 
 @testset "getorcreateexperiment" begin
     @ensuremlf
     expname = "getorcreate"
-    e = getorcreateexperiment(mlf, expname)
+    artifact_location = "test$(expname)"
+    e = getorcreateexperiment(mlf, expname; artifact_location=artifact_location)
     @test isa(e, MLFlowExperiment)
     ee = getorcreateexperiment(mlf, expname)
     @test isa(ee, MLFlowExperiment)
     @test e === ee
+    @test occursin(artifact_location, e.artifact_location)
     @test deleteexperiment(mlf, ee)
     @test deleteexperiment(mlf, ee)
 end
@@ -255,4 +285,5 @@ end
     deleterun(mlf, exprun2)
 
     deleteexperiment(mlf, exp)
+    deleteexperiment(mlf, experiment)
 end