Adding artifacts upload through S3/Minio

Author: rkube

Project.toml

@@ -4,18 +4,24 @@ authors = ["@deyandyankov, @pebeto, and contributors"]
 version = "0.7.0"
 
 [deps]
+AWSS3 = "1c724243-ef5b-51ab-93f4-b0a88ac62a95"
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
+FileTypes = "b58e86d0-4a47-4fce-a54d-8006a143ed90"
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+Minio = "4281f0d9-7ae0-406e-9172-b7277c1efa20"
 ShowCases = "605ecd9f-84a6-4c9e-81e2-4798472b76a3"
 URIs = "5c2747f8-b7ea-4ff2-ba2e-563bfd36b1d4"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [compat]
+AWSS3 = "0.11.4"
 Base64 = "1.0"
+FileTypes = "0.1.1"
 HTTP = "1.0"
 JSON = "0.21"
+Minio = "0.2.2"
 ShowCases = "0.1"
 URIs = "1.0"
 julia = "1.0"

src/MLFlowClient.jl

@@ -11,13 +11,16 @@ and artifacts. If you are not familiar with `MLFlow` and its concepts, please re
 """
 module MLFlowClient
 
+using AWSS3
 using Dates
+using FileTypes
 using UUIDs
 using HTTP
 using Base64
 using URIs
 using JSON
 using ShowCases
+using Minio
 
 include("types/mlflow.jl")
 export MLFlow
@@ -65,7 +68,7 @@ export getrun, createrun, deleterun, setruntag, updaterun, restorerun, searchrun
     deleteruntag
 
 include("services/logger.jl")
-export logbatch, loginputs, logmetric, logmodel, logparam
+export logbatch, loginputs, logmetric, logmodel, logparam, logartifact
 
 include("services/artifact.jl")
 export listartifacts

src/services/logger.jl

@@ -139,3 +139,48 @@ function logmodel(instance::MLFlow, run_id::String, model_json::String)::Bool
 end
 logmodel(instance::MLFlow, run::Run, model_json::String)::Bool =
     logmodel(instance, run.info.run_id, model_json)
+
+"""
+    logartifact(s3_cfg::MinioConfig, run::Run, s3_cfg::MinioConfig, path::String="", artifact_name::String="")
+
+Log artifact for a run. Supports only S3 buckets.
+
+# Arguments
+- `s3_cfg`: Minio configuration
+- `Run`: ['Run'](@ref) instance
+- `path`: Path to the artifact
+- `artifact_name`: Name of the artifact in the bucket
+
+# Returns
+`true` if successful. Otherwise, raises exception.
+
+"""
+function logartifact(s3_cfg::MinioConfig, run::Run, path::String="", artifact_name::String="")
+    # Parse the URI
+    u = URI(run.info.artifact_uri)
+    u.scheme == "s3" || ArgumentError("The artifact URI for the run has to be a S3 bucket. Got: $(run.info.artifact_uri)")
+    isfile(path) || ArgumentError("Can not read file $(path).")
+    bucket_name = u.host
+    artifacts_base_path = u.path
+
+
+    # Determine MIME type to use
+    kind = matcher(path)
+    mime_type_str = if isnothing(kind)
+        @warn "FileTypes.jl could not determing the specific MIME type for $(path). Defaulting to application/octet-stream"
+        "application/octet-stream"
+    else 
+        string(kind.mime)
+    end
+
+    # Read the bytes of the file
+    content = read(path)
+
+    # Create the artifact path on the bucket
+    artifact_path = isempty(artifact_name) ? joinpath(artifacts_base_path, path) : joinpath(artifacts_base_path, artifact_name)
+
+    s3_put(s3_cfg, bucket_name, artifact_path, content, mime_type_str)
+    return true
+end
+
+

src/services/run.jl

@@ -20,6 +20,9 @@ An instance of type [`Run`](@ref).
 function createrun(instance::MLFlow, experiment_id::String;
     run_name::Union{String,Missing}=missing, start_time::Union{Int64,Missing}=missing,
     tags::MLFlowUpsertData{Tag}=Tag[])::Run
+    # Time returns system time in seconds, convert to ms.
+    start_time = ismissing(start_time) ? Int(floor(1e3 * time())) : start_time
+
     result = mlfpost(instance, "runs/create"; experiment_id=experiment_id,
         run_name=run_name, start_time=start_time, tags=parse(Tag, tags))
     return result["run"] |> Run

src/types/run.jl

@@ -140,5 +140,5 @@ struct Run
     outputs::RunOutputs
 end
 Run(data::Dict{String,Any}) = Run(RunInfo(data["info"]), RunData(data["data"]),
-    RunInputs(data["inputs"]), RunOutputs(data["outputs"]))
+    RunInputs(data["inputs"]), RunOutputs(get(data, "outputs", Dict{String, Any}())))
 Base.show(io::IO, t::Run) = show(io, ShowCase(t, new_lines=true))