Adding S3 Support for Artifacts

Project.toml

@@ -4,10 +4,13 @@ authors = ["@deyandyankov and contributors"]
 version = "0.4.4"
 
 [deps]
+AWS = "fbe9abb3-538b-5e4e-ba9e-bc94f4f92ebc"
+AWSS3 = "1c724243-ef5b-51ab-93f4-b0a88ac62a95"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FilePathsBase = "48062228-2e41-5def-b9a4-89aafe57970f"
 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"

src/MLFlowClient.jl

@@ -19,6 +19,10 @@ using URIs
 using JSON
 using ShowCases
 using FilePathsBase: AbstractPath
+using AWSS3
+using Minio
+using AWS: AbstractAWSConfig, AWSConfig, AWSCredentials, global_aws_config
+
 
 include("types/core.jl")
 export

src/loggers.jl

@@ -74,7 +74,7 @@
 # Arguments
 - `mlf::MLFlow`: [`MLFlow`](@ref) onfiguration. Currently not used, but when this method is extended to support `S3`, information from `mlf` will be needed.
 - `run`: one of [`MLFlowRun`](@ref), [`MLFlowRunInfo`](@ref) or `String`.
-- `basefilename`: name of the file to be written.
+- `basefilename`: name of the file to be written; can contain a folder such as `model/mycode.jl` which the folder will be created in the artifact directory.
 - `data`: artifact content, an object that can be written directly to a file handle.
 
 # Throws
@@ -83,19 +83,54 @@
 # Returns
 path of the artifact that was created.
 """
-function logartifact(mlf::MLFlow, run_id::AbstractString, basefilename::AbstractString, data)
+function logartifact(mlf::MLFlow, run_id::AbstractString, basefilename::AbstractString, data; artifact_path="")
  mlflowrun = getrun(mlf, run_id)
- artifact_uri = mlflowrun.info.artifact_uri
- mkpath(artifact_uri)
- filepath = joinpath(artifact_uri, basefilename)
- try
- f = open(filepath, "w")
- write(f, data)
- close(f)
- catch e
- error("Unable to create artifact $(filepath): $e")
+ artifact_uri = joinpath(mlflowrun.info.artifact_uri,artifact_path,dirname(basefilename))
+ basefilename = basename(basefilename)
+
+ if !startswith(artifact_uri, "s3://")
+ mkpath(artifact_uri)
+ filepath = joinpath(artifact_uri, basefilename)
+ try
+ open(filepath, "w") do f
+ write(f, data)
+ end
+ catch e
+ error("Unable to create artifact $(filepath): $e")
+ end
+ else
+ region = get(ENV, "AWS_REGION", "") # Optional, defaults to empty if not set
+
+ if region == ""
+ region = get(ENV, "AWS_DEFAULT_REGION", "") 
+ end
+
+ if haskey(ENV, "MLFLOW_S3_ENDPOINT_URL")
+ s3creds = AWSCredentials()
+ s3config = MinioConfig(ENV["MLFLOW_S3_ENDPOINT_URL"], s3creds; region=region)
+ else
+ s3config = global_aws_config() # default AWS configuration
+ end
+
+ filepath = joinpath(artifact_uri, basefilename)
+ artifact_uri = rstrip(artifact_uri[6:end], '/')# get rid of s3:// so s3_put doesnt' complain
+
+ try
+ #TODO: Figure out the correct IO stream way of doing this
+ open(joinpath("/tmp/",basefilename), "w") do f
+ write(f, data)
+ end
+ open(joinpath("/tmp/",basefilename), "r") do f
+ file_data = read(f)
+ s3_put(s3config, artifact_uri, basefilename, file_data)
+ end
+ rm(joinpath("/tmp",basefilename))
+ catch e
+ error("Unable to upload artifact to S3 $(filepath): $e")
+ end
  end
- filepath
+
+ return filepath
 end
 logartifact(mlf::MLFlow, run::MLFlowRun, basefilename::AbstractString, data) =
  logartifact(mlf, run.info, basefilename, data)