Adding artifacts upload through S3/Minio

.github/workflows/CI.yml

@@ -22,21 +22,23 @@ jobs:
           - 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
-      - uses: actions/setup-python@v4
-        with:
-          python-version: '3.12.3'
-          cache: 'pip'
-      - 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 &
-          sleep 5
+      - name: Start services
+        run: docker-compose -f docker compose.test.yaml up -d
+#      - name: Setup custom python requirements
+#        if: hashFiles('**/requirements.txt', '**/pyproject.toml') == ''
+#        run: |
+#          touch ./requirements.txt
+#          echo "mlflow[auth]==3.2.0" > ./requirements.txt
+#      - uses: actions/setup-python@v4
+#        with:
+#          python-version: '3.12.3'
+#          cache: 'pip'
+#      - 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 &
+#          sleep 5
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}
@@ -54,8 +56,11 @@ 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'
+          MLFLOW_TRACKING_URI: "http://localhost:5050/api"
+          MLFLOW_S3_ENDPOINT_URL: "http://minio:9000"
+          AWS_ACCESS_KEY_ID: minioadmin
+          AWS_SECRET_ACCESS_KEY: minioadmin
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v3
         with:

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"

docker-compose.test.yaml

@@ -0,0 +1,43 @@
+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:v2.14.1
+    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
+    command: >
+      /bin/sh -c "
+      pip install boto3 &&
+      mlflow server
+      --host 0.0.0.0
+      --port 5050
+      --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,25 @@
 using Test
 using Dates
+using HTTP
 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", "$(TEST_MLFLOW_URI)/health", readtimeout=10)
+    return resp.status == 200 
 end
 
 # creates an instance of mlf
@@ -19,7 +28,27 @@ 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", "$(TEST_MLFLOW_S3_ENDPOINT_URL)/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,34 @@
 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")
+    error("""WARNING: MLFLOW_TRACKING_URI is not set. 
+To run the unit tests, you need to set the URI of your MLFlow server API.
+
+A test environment is provided in .devcontainers/compose.yaml, start it as
+`docker-compose .devcontainers/compose.yaml up`.
+
+Then set the environment variables
+MLFLOW_TRACKING_URI="http://localhost:5050/api"
+AWS_ACCESS_KEY_ID="minioadmin"
+AWS_SECRET_ACCESS_KEY="minioadmin"
+MLFLOW_S3_ENDPOINT_URL="http://localhost:9000"
+""")
 end
 
+
+# Set up access to testing environment defined in docker-compose.test.yaml
+const TEST_MLFLOW_URI = "http://localhost:5050/"
+const TEST_MLFLOW_TRACKING_URI = "http://127.0.0.1:5050/api"
+const TEST_MLFLOW_S3_ENDPOINT_URL = get(ENV, "MLFLOW_S3_ENDPOINT_URL", "http://127.0.0.1:9000")
+const TEST_AWS_ACCESS_KEY_ID = get(ENV, "AWS_ACCESS_KEY_ID", "minioadmin")
+const TEST_AWS_SECRET_ACCESS_KEY = get(ENV, "AWS_SECRET_ACCESS_KEY", "minioadmin")
+
+
+
+
 include("base.jl")
 
+const minio_cfg = MinioConfig(TEST_MLFLOW_S3_ENDPOINT_URL; username=TEST_AWS_ACCESS_KEY_ID, password=TEST_AWS_SECRET_ACCESS_KEY)
+include("setup.jl")
+
 include("types/mlflow.jl")
 
 include("services/run.jl")
@@ -13,4 +38,4 @@ include("services/artifact.jl")
 include("services/experiment.jl")
 include("services/registered_model.jl")
 include("services/model_version.jl")
-include("services/user.jl")
+#include("services/user.jl")