Allow logging to remote artifact location

Project.toml

@@ -6,6 +6,7 @@ version = "0.5.1"
 [deps]
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FilePathsBase = "48062228-2e41-5def-b9a4-89aafe57970f"
+FileTypes = "b58e86d0-4a47-4fce-a54d-8006a143ed90"
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 ShowCases = "605ecd9f-84a6-4c9e-81e2-4798472b76a3"
@@ -25,4 +26,3 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Test"]
-

src/MLFlowClient.jl

@@ -19,6 +19,7 @@ using URIs
 using JSON
 using ShowCases
 using FilePathsBase: AbstractPath
+using FileTypes
 
 include("types/mlflow.jl")
 export

src/api.jl

@@ -31,4 +31,26 @@ function mlfpost(mlf, endpoint; kwargs...)
     catch e
         throw(e)
     end
-end
+end
+
+"""
+    mlfput_artifact(mlf, artifact_uri, filename, data)
+
+Performs a HTTP PUT to upload the specified artifact.
+Assumes that the artifact store is hosted in the mlflow server.
+"""
+function mlfput_artifact(mlf, artifact_uri, filename, data)
+    artifact_path = chopprefix(artifact_uri, "mlflow-artifacts:/")
+    if artifact_path == artifact_uri
+        error("Artifact URI must start with `mlflow-artifacts:/`")
+    end
+    content_type = guess_mime(filename).mime
+    apiuri = URI("$(mlf.apiroot)/$(mlf.apiversion)/mlflow-artifacts/artifacts/$(artifact_path)/$(filename)")
+    apiheaders = headers(mlf, Dict("Content-Type" => content_type))
+    try
+        response = HTTP.put(apiuri, apiheaders, data)
+        return JSON.parse(String(response.body))
+    catch e
+        throw(e)
+    end
+end

src/loggers.jl

@@ -68,8 +68,7 @@ end
 Stores an artifact (file) in the run's artifact location.
 
 !!! note
-    Assumes that artifact_uri is mapped to a local directory.
-    At the moment, this only works if both MLFlow and the client are running on the same host or they map a directory that leads to the same location over NFS, for example.
+    Only supports local and proxied artifact stores.
 
 # Arguments
 - `mlf::MLFlow`: [`MLFlow`](@ref) onfiguration. Currently not used, but when this method is extended to support `S3`, information from `mlf` will be needed.
@@ -86,8 +85,12 @@ path of the artifact that was created.
 function logartifact(mlf::MLFlow, run_id::AbstractString, basefilename::AbstractString, data)
     mlflowrun = getrun(mlf, run_id)
     artifact_uri = mlflowrun.info.artifact_uri
-    mkpath(artifact_uri)
     filepath = joinpath(artifact_uri, basefilename)
+    if startswith(artifact_uri, "mlflow-artifacts:/")
+        mlfput_artifact(mlf, artifact_uri, basefilename, data)
+        return filepath
+    end
+    mkpath(artifact_uri)
     try
         f = open(filepath, "w")
         write(f, data)

src/utils.jl

@@ -51,6 +51,17 @@ function generatefilterfromentity_type(filter_params::AbstractDict{K,V}, entity_
     join(filters, " and ")
 end
 
+function guess_mime(data)::FileTypes.FileType.Type
+    for i in FileTypes.Types
+        for (key, value) in i
+            if value(data)
+                return key
+            end
+        end
+    end
+    FileTypes.FileType.Type("", MIME("application/octet-stream"))
+end
+
 """
     generatefilterfromparams(filter_params::AbstractDict{K,V}) where {K,V}
 

test/test_experiments.jl

@@ -47,38 +47,40 @@ end
 @testset "deleteexperiment" begin
     @ensuremlf
     exp = createexperiment(mlf)
+    experiments_before = searchexperiments(mlf)
     deleteexperiment(mlf, exp)
 
-    experiments = searchexperiments(mlf)
-    @test length(experiments) == 1 # 1 for the default experiment
+    experiments_after = searchexperiments(mlf)
+    @test length(experiments_after) == length(experiments_before) - 1 # 1 for the default experiment
 end
 
 @testset "restoreexperiment" begin
     @ensuremlf
     exp = createexperiment(mlf)
+    experiments_before = searchexperiments(mlf)
     deleteexperiment(mlf, exp)
 
-    experiments = searchexperiments(mlf)
-    @test length(experiments) == 1 # 1 for the default experiment
+    experiments_after = searchexperiments(mlf)
+    @test length(experiments_after) == length(experiments_before) - 1 # 1 for the default experiment
 
     restoreexperiment(mlf, exp)
-    experiments = searchexperiments(mlf)
-    @test length(experiments) == 2 # the restored experiment and the default one
+    experiments_after_2 = searchexperiments(mlf)
+    @test length(experiments_after_2) == length(experiments_after) + 1 # the restored experiment and the default one
 
     deleteexperiment(mlf, exp)
 end
 
 @testset verbose = true "searchexperiments" begin
     @ensuremlf
-    n_experiments = 3
-    for i in 2:n_experiments
+    n_experiments_before = length(searchexperiments(mlf))
+    for i in 1:2
         createexperiment(mlf)
     end
     createexperiment(mlf; name="test")
     experiments = searchexperiments(mlf)
 
     @testset "searchexperiments_get_all" begin
-        @test length(experiments) == (n_experiments + 1) # Adding one for the default experiment
+        @test length(experiments) == (n_experiments_before + 3) # Adding one for the default experiment
     end
 
     @testset "searchexperiments_by_filter" begin

test/test_functional.jl

@@ -60,8 +60,18 @@ end
     # when running mlflow in a container, the below tests will be skipped
     # this is what happens in github actions - mlflow runs in a container, the artifact_uri is not immediately available, and tests are skipped
     artifact_uri = exprun.info.artifact_uri
+    if startswith(artifact_uri, "mlflow-artifacts:/")
+        tmp = tempdir()
+        mkpath(tmp)
+        tmpfiletoupload = "$(tmp)/sometempfilename.txt"
+        f = open(tmpfiletoupload, "w")
+        write(f, "samplecontents")
+        close(f)
+        logartifact(mlf, exprun, tmpfiletoupload)
+    end
+
     if isdir(artifact_uri)
-        @test_throws SystemError logartifact(mlf, exprun, "/etc/shadow")
+        @test_throws ErrorException logartifact(mlf, exprun, "/etc/shadow")
 
         tmpfiletoupload = "sometempfilename.txt"
         f = open(tmpfiletoupload, "w")