Merge pull request #60 from JuliaAI/pairing-mlflow-v3

Pairing mlflow v3

Author: pebeto

.github/workflows/CI.yml

@@ -26,21 +26,22 @@ jobs:
         if: hashFiles('**/requirements.txt', '**/pyproject.toml') == ''
         run: |
           touch ./requirements.txt
-          echo "mlflow==2.20.1" > ./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 }}
           arch: ${{ matrix.arch }}
-      - uses: actions/cache@v1
+      - uses: actions/cache@v4
         env:
           cache-name: cache-artifacts
         with:

README.md

@@ -4,7 +4,7 @@
 [![Build Status](https://github.com/JuliaAI/MLFlowClient.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/JuliaAI/MLFlowClient.jl/actions/workflows/CI.yml?query=branch%3Amain)
 [![Coverage](https://codecov.io/gh/JuliaAI/MLFlowClient.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/JuliaAI/MLFlowClient.jl)
 
-Julia client for [MLFlow](https://www.mlflow.org/) `2.20.1` (but should work with other versions as well).
+Julia client for [MLFlow](https://www.mlflow.org/) `3.2.0` (but should work with other versions as well).
 
 - [x] Supports tracking of metrics, parameters, tags, artifacts, and models.
 - [x] Compatible with latest MLFlow server capabilities.

docs/src/reference/logger.md

@@ -4,4 +4,5 @@ logmetric
 logbatch
 loginputs
 logparam
+logmodel
 ```

docs/src/reference/types.md

@@ -2,9 +2,6 @@
 ```@docs
 MLFlow
 Tag
-ViewType
-RunStatus
-ModelVersionStatus
 Dataset
 DatasetInput
 FileInfo
@@ -19,7 +16,6 @@ RunData
 RunInfo
 RunInputs
 User
-Permission
 ExperimentPermission
 RegisteredModelPermission
 ```

src/MLFlowClient.jl

@@ -26,16 +26,17 @@ include("types/tag.jl")
 export Tag
 
 include("types/enums.jl")
-export ViewType, RunStatus, ModelVersionStatus, Permission
+export ViewType, RunStatus, ModelVersionStatus, Permission, DeploymentJobRunState, State
 
 include("types/dataset.jl")
 export Dataset, DatasetInput
 
 include("types/artifact.jl")
 export FileInfo
 
-include("types/model_version.jl")
-export ModelVersion
+include("types/model.jl")
+export ModelInput, ModelMetric, ModelOutput, ModelParam, ModelVersion,
+    ModelVersionDeploymentJobState
 
 include("types/registered_model.jl")
 export RegisteredModel, RegisteredModelAlias, RegisteredModelPermission
@@ -64,7 +65,7 @@ export getrun, createrun, deleterun, setruntag, updaterun, restorerun, searchrun
     deleteruntag
 
 include("services/logger.jl")
-export logbatch, loginputs, logmetric, logparam
+export logbatch, loginputs, logmetric, logmodel, logparam
 
 include("services/artifact.jl")
 export listartifacts

src/services/experiment.jl

@@ -156,15 +156,15 @@ updateexperiment(instance::MLFlow, experiment::Experiment, new_name::String)::Bo
 - `order_by`: List of columns for ordering search results, which can include [`Experiment`](@ref)
     name and id with an optional “DESC” or “ASC” annotation, where “ASC” is the default.
 - `view_type`: Qualifier for type of experiments to be returned. If unspecified, return
-    only active experiments. For more values, see [`ViewType`](@ref).
+    only active experiments.
 
 # Returns
 - Vector of [`Experiment`](@ref) that were found in the [`MLFlow`](@ref) instance.
 - The next page token if there are more results.
 """
 function searchexperiments(instance::MLFlow; max_results::Int64=20000,
     page_token::String="", filter::String="", order_by::Array{String}=String[],
-    view_type::ViewType=ACTIVE_ONLY)::Tuple{Array{Experiment},Union{String,Nothing}}
+    view_type::ViewType.ViewTypeEnum=ViewType.ACTIVE_ONLY)::Tuple{Array{Experiment},Union{String,Nothing}}
     parameters = (; max_results, page_token, filter, :view_type => view_type |> Integer)
 
     if order_by |> !isempty
@@ -220,22 +220,22 @@ setexperimenttag(instance::MLFlow, experiment::Experiment, key::String,
 - `instance`: [`MLFlow`](@ref) configuration.
 - `experiment_id`: [`Experiment`](@ref) id.
 - `username`: [`User`](@ref) username.
-- `permission`: [`Permission`](@ref) to grant.
+- `permission`: Permission to grant.
 
 # Returns
 An instance of type [`ExperimentPermission`](@ref).
 """
 function createexperimentpermission(instance::MLFlow, experiment_id::String,
-    username::String, permission::Permission)::ExperimentPermission
+    username::String, permission::Permission.PermissionEnum)::ExperimentPermission
     result = mlfpost(instance, "experiments/permissions/create";
         experiment_id=experiment_id, username=username, permission=permission)
     return result["experiment_permission"] |> ExperimentPermission
 end
 createexperimentpermission(instance::MLFlow, experiment_id::Integer,
-    username::String, permission::Permission)::ExperimentPermission =
+    username::String, permission::Permission.PermissionEnum)::ExperimentPermission =
     createexperimentpermission(instance, experiment_id |> string, username, permission)
 createexperimentpermission(instance::MLFlow, experiment::Experiment,
-    username::String, permission::Permission)::ExperimentPermission =
+    username::String, permission::Permission.PermissionEnum)::ExperimentPermission =
     createexperimentpermission(instance, experiment.experiment_id, username, permission)
 
 """
@@ -276,22 +276,22 @@ getexperimentpermission(instance::MLFlow, experiment::Experiment,
 - `instance`: [`MLFlow`](@ref) configuration.
 - `experiment_id`: [`Experiment`](@ref) id.
 - `username`: [`User`](@ref) username.
-- `permission`: [`Permission`](@ref) to grant.
+- `permission`: Permission to grant.
 
 # Returns
 `true` if successful. Otherwise, raises exception.
 """
 function updateexperimentpermission(instance::MLFlow, experiment_id::String,
-    username::String, permission::Permission)::Bool
+    username::String, permission::Permission.PermissionEnum)::Bool
     mlfpatch(instance, "experiments/permissions/update"; experiment_id=experiment_id,
         username=username, permission=permission)
     return true
 end
 updateexperimentpermission(instance::MLFlow, experiment_id::Integer,
-    username::String, permission::Permission)::Bool =
+    username::String, permission::Permission.PermissionEnum)::Bool =
     updateexperimentpermission(instance, experiment_id |> string, username, permission)
 updateexperimentpermission(instance::MLFlow, experiment::Experiment,
-    username::String, permission::Permission)::Bool =
+    username::String, permission::Permission.PermissionEnum)::Bool =
     updateexperimentpermission(instance, experiment.experiment_id, username, permission)
 
 """

src/services/logger.jl

@@ -57,8 +57,9 @@ For more information about this function, check [MLFlow official documentation](
 - `params`: A collection of [`Param`](@ref) to log.
 - `tags`: A collection of [`Tag`](@ref) to log.
 
-**Note**: A single request can contain up to 1000 metrics, and up to 1000 metrics, params,
-and tags in total.
+!!! note
+    A single request can contain up to 1000 metrics, and up to 1000 metrics, params, and
+    tags in total.
 
 # Returns
 `true` if successful. Otherwise, raises exception.
@@ -123,3 +124,18 @@ logparam(instance::MLFlow, run_id::String, param::Param)::Bool =
     logparam(instance, run_id, param.key, param.value)
 logparam(instance::MLFlow, run::Run, param::Param)::Bool =
     logparam(instance, run.info.run_id, param.key, param.value)
+
+"""
+    logmodel(instance::MLFlow, run_id::String, model_json::String)
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `run_id`: ID of the [`Run`](@ref) to log under.
+- `model_json`: MLmodel file in json format.
+"""
+function logmodel(instance::MLFlow, run_id::String, model_json::String)::Bool
+    mlfpost(instance, "runs/log-model"; run_id=run_id, model_json=model_json)
+    return true
+end
+logmodel(instance::MLFlow, run::Run, model_json::String)::Bool =
+    logmodel(instance, run.info.run_id, model_json)

src/services/misc.jl

@@ -10,33 +10,31 @@ Get a list of all values for the specified [`Metric`](@ref) for a given [`Run`](
 - `metric_key`: Name of the [`Metric`](@ref) to fetch.
 - `page_token`: Token indicating the page of [`Metric`](@ref) history to fetch.
 - `max_results`: Maximum number of logged instances of a [`Metric`](@ref) for a
-    [`Run`](@ref) to return per call.
+    [`Run`](@ref) to return per call. Defaults to 50.
 
 # Returns
 - A list of all historical values for the specified [`Metric`](@ref) in the specified
     [`Run`](@ref).
 - The next page token if there are more results.
 """
 function getmetrichistory(instance::MLFlow, run_id::String, metric_key::String;
-    page_token::String="", max_results::Union{Int64,Missing}=missing
+    page_token::Union{String,Missing}=missing, max_results::Int64=50
 )::Tuple{Array{Metric},Union{String,Nothing}}
     result = mlfget(instance, "metrics/get-history"; run_id=run_id, metric_key=metric_key,
         page_token=page_token,
         max_results=(ismissing(max_results) ? max_results : (max_results |> Int32)))
 
-    metrics = result["metrics"] |> (x -> [Metric(y) for y in x])
+    metrics = get(result, "metrics", []) |> (x -> [Metric(y) for y in x])
     next_page_token = get(result, "next_page_token", nothing)
 
     return metrics, next_page_token
 end
 getmetrichistory(instance::MLFlow, run::Run, metric_key::String; page_token::String="",
-    max_results::Union{Int64,Missing}=missing
-)::Tuple{Array{Metric},Union{String,Nothing}} =
+    max_results::Int64=50)::Tuple{Array{Metric},Union{String,Nothing}} =
     getmetrichistory(instance, run.info.run_id, metric_key; page_token=page_token,
         max_results=max_results)
 getmetrichistory(instance::MLFlow, run::Run, metric::Metric; page_token::String="",
-    max_results::Union{Int64,Missing}=missing
-)::Tuple{Array{Metric},Union{String,Nothing}} =
+    max_results::Int64=50)::Tuple{Array{Metric},Union{String,Nothing}} =
     getmetrichistory(instance, run.info.run_id, metric.key; page_token=page_token,
         max_results=max_results)
 

src/services/registered_model.jl

@@ -17,10 +17,11 @@ name already exists.
 An instance of type [`RegisteredModel`](@ref).
 """
 function createregisteredmodel(instance::MLFlow, name::String;
-    tags::MLFlowUpsertData{Tag}=Tag[],
-    description::Union{String,Missing}=missing)::RegisteredModel
+    tags::MLFlowUpsertData{Tag}=Tag[], description::Union{String,Missing}=missing,
+    deployment_job_id::Union{String,Missing}=missing)::RegisteredModel
     result = mlfpost(instance, "registered-models/create"; name=name,
-        tags=parse(Tag, tags), description=description)
+        tags=parse(Tag, tags), description=description,
+        deployment_job_id=deployment_job_id)
     return result["registered_model"] |> RegisteredModel
 end
 
@@ -199,13 +200,13 @@ end
 - `instance:` [`MLFlow`](@ref) configuration.
 - `name:` [`RegisteredModel`](@ref) name.
 - `username:` [`User`](@ref) username.
-- `permission:` [`Permission`](@ref) to grant.
+- `permission:` Permission to grant.
 
 # Returns
 An instance of type [`RegisteredModelPermission`](@ref).
 """
 function createregisteredmodelpermission(instance::MLFlow, name::String, username::String,
-    permission::Permission)::RegisteredModelPermission
+    permission::Permission.PermissionEnum)::RegisteredModelPermission
     result = mlfpost(instance, "registered-models/permissions/create"; name=name,
         username=username, permission=permission)
     return result["registered_model_permission"] |> RegisteredModelPermission
@@ -237,13 +238,13 @@ end
 - `instance:` [`MLFlow`](@ref) configuration.
 - `name:` [`RegisteredModel`](@ref) name.
 - `username:` [`User`](@ref) username.
-- `permission:` New [`Permission`](@ref) to grant.
+- `permission:` New permission to grant.
 
 # Returns
 `true` if successful. Otherwise, raises exception.
 """
 function updateregisteredmodelpermission(instance::MLFlow, name::String, username::String,
-    permission::Permission)::Bool
+    permission::Permission.PermissionEnum)::Bool
     mlfpatch(instance, "registered-models/permissions/update"; name=name, username=username,
         permission=permission)
     return true

src/services/run.jl

@@ -170,8 +170,8 @@ Search for runs that satisfy expressions. Search expressions can use [`Metric`](
 - The next page token if there are more results.
 """
 function searchruns(instance::MLFlow; experiment_ids::Array{String}=String[],
-    filter::String="", run_view_type::ViewType=ACTIVE_ONLY, max_results::Int=1000,
-    order_by::Array{String}=String[],
+    filter::String="", run_view_type::ViewType.ViewTypeEnum=ViewType.ACTIVE_ONLY,
+    max_results::Int=1000, order_by::Array{String}=String[],
     page_token::String="")::Tuple{Array{Run},Union{String,Nothing}}
     parameters = (; experiment_ids, filter, :run_view_type => run_view_type |> Integer,
         max_results, page_token)
@@ -207,13 +207,14 @@ Update [`Run`](@ref) metadata.
 - An instance of type [`RunInfo`](@ref) with the updated metadata.
 """
 function updaterun(instance::MLFlow, run_id::String;
-    status::Union{RunStatus,Missing}=missing, end_time::Union{Int64,Missing}=missing,
-    run_name::Union{String,Missing})::RunInfo
+    status::Union{RunStatus.RunStatusEnum,Missing}=missing,
+    end_time::Union{Int64,Missing}=missing, run_name::Union{String,Missing})::RunInfo
     result = mlfpost(instance, "runs/update"; run_id=run_id, status=(status |> Integer),
         end_time=end_time, run_name=run_name)
     return result["run_info"] |> RunInfo
 end
-updaterun(instance::MLFlow, run::Run; status::Union{RunStatus,Missing}=missing,
+updaterun(instance::MLFlow, run::Run;
+    status::Union{RunStatus.RunStatusEnum,Missing}=missing,
     end_time::Union{Int64,Missing}=missing, run_name::Union{String,Missing})::RunInfo =
     updaterun(instance, run.info.run_id; status=status, end_time=end_time,
         run_name=run_name)