Adding artifacts upload through S3/Minio

.github/workflows/CI.yml

@@ -20,22 +20,59 @@ jobs:
           - ubuntu-latest
         arch:
           - x64
+
     steps:
       - uses: actions/checkout@v2
       - name: Setup custom python requirements
         if: hashFiles('**/requirements.txt', '**/pyproject.toml') == ''
         run: |
           touch ./requirements.txt
-          echo "mlflow[auth]==3.2.0" > ./requirements.txt
+          echo "mlflow[auth]==3.6.0" > ./requirements.txt
       - uses: actions/setup-python@v4
         with:
           python-version: '3.12.3'
           cache: 'pip'
+      - name: Setup rclone S3 server
+        run: |
+          # Download rclone
+          curl -O https://downloads.rclone.org/rclone-current-linux-amd64.zip
+          unzip rclone-current-linux-amd64.zip
+          cd rclone-*-linux-amd64
+          sudo cp rclone /usr/bin/
+          cd ..
+
+          # Create the data directory and the "mlflow" bucket (subdirectory)
+          mkdir -p ./s3-data/mlflow
+
+          # Start Rclone serving the local folder as S3
+          # --auth-key sets access_key,secret_key
+          rclone serve s3 ./s3-data \
+            --addr 127.0.0.1:8080 \
+            --auth-key minioadmin,minioadmin \
+            --no-modtime \
+            &
+
+          # Wait for port 8080 to be ready
+          sleep 3
+#      - name: Start MinIO
+#        run: |
+#          curl -O https://dl.min.io/server/minio/release/linux-amd64/minio
+#          chmod +x minio
+#          export MINIO_ROOT_USER=minioadmin
+#          export MINIO_ROOT_PASSWORD=minioadmin
+#          ./minio server ./data --console-address ":9001" &
+#          sleep 5
+#      - name: Create MLFlow bucket
+#        run: |
+#          curl -O https://dl.min.io/client/mc/release/linux-amd64/mc
+#          chmod +x mc
+#          ./mc alias set myminio http://localhost:9000 minioadmin minioadmin
+#          ./mc mb myminio/mlflow
       - name: Setup mlflow locally
         run: |
           export MLFLOW_FLASK_SERVER_SECRET_KEY='mlflowclient.jl'
           pip install -r ./requirements.txt
-          python3 /opt/hostedtoolcache/Python/3.12.3/x64/bin/mlflow server --app-name basic-auth --host 0.0.0.0 --port 5000 &
+          python3 /opt/hostedtoolcache/Python/3.12.3/x64/bin/mlflow server --app-name basic-auth --host 0.0.0.0 --port 5050 &
           sleep 5
       - uses: julia-actions/setup-julia@v1
         with:
@@ -54,8 +91,7 @@ jobs:
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         env:
-          JULIA_NUM_THREADS: '2'
-          MLFLOW_TRACKING_URI: "http://localhost:5000/api"
+          JULIA_NUM_THREADS: '1'
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v3
         with:

Project.toml

@@ -1,28 +1,36 @@
 name = "MLFlowClient"
 uuid = "64a0f543-368b-4a9a-827a-e71edb2a0b83"
-authors = ["@deyandyankov, @pebeto, and contributors"]
 version = "0.7.0"
+authors = ["@deyandyankov, @pebeto, and contributors"]
 
 [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"
+Plots = "1.41.2"
 ShowCases = "0.1"
 URIs = "1.0"
 julia = "1.0"
 
 [extras]
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [targets]
-test = ["Base64", "Test"]
+test = ["Base64", "Plots", "Test"]

docker-compose.test.yaml

@@ -0,0 +1,46 @@
+version: '3.8'
+
+services:
+  minio:
+    image: minio/minio
+    ports:
+      - "9000:9000"    # S3 API
+      - "9001:9001"    # Minio Console
+    environment:
+      - MINIO_ROOT_USER=minioadmin
+      - MINIO_ROOT_PASSWORD=minioadmin
+    command: server /data --console-address ":9001"
+
+  create-buckets:
+    image: minio/mc
+    depends_on:
+      - minio
+    entrypoint: >
+      /bin/sh -c "
+      sleep 10;
+      /usr/bin/mc alias set myminio http://minio:9000 minioadmin minioadmin;
+      /usr/bin/mc mb myminio/mlflow;
+      exit 0;
+      "
+
+  mlflow:
+    image: ghcr.io/mlflow/mlflow:v3.6.0
+    depends_on:
+      - create-buckets
+    ports:
+      - "5050:5050"
+    environment:
+      - AWS_ACCESS_KEY_ID=minioadmin
+      - AWS_SECRET_ACCESS_KEY=minioadmin
+      - MLFLOW_S3_ENDPOINT_URL=http://minio:9000
+      - MLFLOW_FLASK_SERVER_SECRET_KEY='mlflowclient.jl'
+    command: >
+      /bin/sh -c "
+      pip install boto3 &&
+      pip install mlflow[auth] && 
+      mlflow server
+      --host 0.0.0.0
+      --port 5050
+      --app-name basic-auth
+      --default-artifact-root s3://mlflow
+      "

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,49 @@ 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)
+    @info "logartifact: putting $(path) into $(artifact_path)"
+
+    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))

test/base.jl

@@ -1,16 +1,26 @@
 using Test
 using Dates
+using HTTP
+using Plots
 using UUIDs
 using Base64
 using MLFlowClient
+using Minio
+using URIs
 
-function mlflow_server_is_running(mlf::MLFlow)
-    try
-        response = MLFlowClient.mlfget(mlf, "experiments/list")
-        return isa(response, Dict)
-    catch e
-        return false
-    end
+# The route experiments/list was deprecated.
+# It is thee in f.ex. 1.30.1, but not in 2.0.1:
+# https://mlflow.org/docs/1.30.1/rest-api.html#list-experiments
+# https://mlflow.org/docs/2.0.1/rest-api.html#list-experiments
+# Query the health endpoint instead: https://github.com/mlflow/mlflow/pull/2725
+"""
+    mlflow_server_is_running(mlf::MLFlow)
+
+Check MLFlow health endpoint. Return true if healthy, false otherwise.
+"""
+function mlflow_server_is_running()
+    resp = HTTP.request("HEAD", "http://127.0.0.1:5050/health", readtimeout=10)
+    return resp.status == 200
 end
 
 # creates an instance of mlf
@@ -19,7 +29,31 @@ macro ensuremlf()
     e = quote
         encoded_credentials = Base64.base64encode("admin:password1234")
         mlf = MLFlow(headers=Dict("Authorization" => "Basic $(encoded_credentials)"))
-        mlflow_server_is_running(mlf) || return nothing
+        mlflow_server_is_running() || return nothing
+        mlf
+    end
+    eval(e)
+end
+
+"""
+    minio_is_running
+
+Check minio health endpoint. Return true if health, false otherwise
+"""
+function minio_is_running()
+    response = HTTP.request(
+        "HEAD",
+        "http://127.0.0.1:9000/minio/health/live",
+        readtimeout=10,
+    )
+    return response.status == 200
+end
+
+macro ensureminio()
+    e = quote
+        minio_is_running() || return nothing
     end
     eval(e)
 end
+
+

test/runtests.jl

@@ -1,9 +1,17 @@
-if ~haskey(ENV, "MLFLOW_TRACKING_URI")
-    error("WARNING: MLFLOW_TRACKING_URI is not set. To run this tests, you need to set the URI of your MLFlow server API")
-end
+@warn "To run this test suite, ensure to run `docker-compose -f docker-compose.test.yml up`"
+
+ENV["GKSwstype"]="nul" # to disable plotting windows during tests
+ENV["MLFLOW_TRACKING_URI"] = "http://127.0.0.1:5050/api"
 
 include("base.jl")
 
+const minio_cfg = MinioConfig(
+    "http://127.0.0.1:9000";
+    username="minioadmin",
+    password="minioadmin",
+)
+include("setup.jl")
+
 include("types/mlflow.jl")
 
 include("services/run.jl")