Skip to content

Commit a2a797a

Browse files
committed
Pairing mlflow v3 models
1 parent c03c83b commit a2a797a

File tree

17 files changed

+285
-142
lines changed

17 files changed

+285
-142
lines changed

.github/workflows/CI.yml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,13 +26,14 @@ jobs:
2626
if: hashFiles('**/requirements.txt', '**/pyproject.toml') == ''
2727
run: |
2828
touch ./requirements.txt
29-
echo "mlflow==2.20.1" > ./requirements.txt
29+
echo "mlflow==3.1.2" > ./requirements.txt
3030
- uses: actions/setup-python@v4
3131
with:
3232
python-version: '3.12.3'
3333
cache: 'pip'
3434
- name: Setup mlflow locally
3535
run: |
36+
export MLFLOW_FLASK_SERVER_SECRET_KEY='mlflowclient.jl'
3637
pip install -r ./requirements.txt
3738
python3 /opt/hostedtoolcache/Python/3.12.3/x64/bin/mlflow server --app-name basic-auth --host 0.0.0.0 --port 5000 &
3839
sleep 5

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
[![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)
55
[![Coverage](https://codecov.io/gh/JuliaAI/MLFlowClient.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/JuliaAI/MLFlowClient.jl)
66

7-
Julia client for [MLFlow](https://www.mlflow.org/) `2.20.1` (but should work with other versions as well).
7+
Julia client for [MLFlow](https://www.mlflow.org/) `3.1.2` (but should work with other versions as well).
88

99
- [x] Supports tracking of metrics, parameters, tags, artifacts, and models.
1010
- [x] Compatible with latest MLFlow server capabilities.

src/MLFlowClient.jl

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,16 +26,17 @@ include("types/tag.jl")
2626
export Tag
2727

2828
include("types/enums.jl")
29-
export ViewType, RunStatus, ModelVersionStatus, Permission
29+
export ViewType, RunStatus, ModelVersionStatus, Permission, DeploymentJobRunState, State
3030

3131
include("types/dataset.jl")
3232
export Dataset, DatasetInput
3333

3434
include("types/artifact.jl")
3535
export FileInfo
3636

37-
include("types/model_version.jl")
38-
export ModelVersion
37+
include("types/model.jl")
38+
export ModelInput, ModelMetric, ModelOutput, ModelParam, ModelVersion,
39+
ModelVersionDeploymentJobState
3940

4041
include("types/registered_model.jl")
4142
export RegisteredModel, RegisteredModelAlias, RegisteredModelPermission

src/services/experiment.jl

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,7 @@ updateexperiment(instance::MLFlow, experiment::Experiment, new_name::String)::Bo
164164
"""
165165
function searchexperiments(instance::MLFlow; max_results::Int64=20000,
166166
page_token::String="", filter::String="", order_by::Array{String}=String[],
167-
view_type::ViewType=ACTIVE_ONLY)::Tuple{Array{Experiment},Union{String,Nothing}}
167+
view_type::ViewType.ViewTypeEnum=ViewType.ACTIVE_ONLY)::Tuple{Array{Experiment},Union{String,Nothing}}
168168
parameters = (; max_results, page_token, filter, :view_type => view_type |> Integer)
169169

170170
if order_by |> !isempty
@@ -226,16 +226,16 @@ setexperimenttag(instance::MLFlow, experiment::Experiment, key::String,
226226
An instance of type [`ExperimentPermission`](@ref).
227227
"""
228228
function createexperimentpermission(instance::MLFlow, experiment_id::String,
229-
username::String, permission::Permission)::ExperimentPermission
229+
username::String, permission::Permission.PermissionEnum)::ExperimentPermission
230230
result = mlfpost(instance, "experiments/permissions/create";
231231
experiment_id=experiment_id, username=username, permission=permission)
232232
return result["experiment_permission"] |> ExperimentPermission
233233
end
234234
createexperimentpermission(instance::MLFlow, experiment_id::Integer,
235-
username::String, permission::Permission)::ExperimentPermission =
235+
username::String, permission::Permission.PermissionEnum)::ExperimentPermission =
236236
createexperimentpermission(instance, experiment_id |> string, username, permission)
237237
createexperimentpermission(instance::MLFlow, experiment::Experiment,
238-
username::String, permission::Permission)::ExperimentPermission =
238+
username::String, permission::Permission.PermissionEnum)::ExperimentPermission =
239239
createexperimentpermission(instance, experiment.experiment_id, username, permission)
240240

241241
"""
@@ -282,16 +282,16 @@ getexperimentpermission(instance::MLFlow, experiment::Experiment,
282282
`true` if successful. Otherwise, raises exception.
283283
"""
284284
function updateexperimentpermission(instance::MLFlow, experiment_id::String,
285-
username::String, permission::Permission)::Bool
285+
username::String, permission::Permission.PermissionEnum)::Bool
286286
mlfpatch(instance, "experiments/permissions/update"; experiment_id=experiment_id,
287287
username=username, permission=permission)
288288
return true
289289
end
290290
updateexperimentpermission(instance::MLFlow, experiment_id::Integer,
291-
username::String, permission::Permission)::Bool =
291+
username::String, permission::Permission.PermissionEnum)::Bool =
292292
updateexperimentpermission(instance, experiment_id |> string, username, permission)
293293
updateexperimentpermission(instance::MLFlow, experiment::Experiment,
294-
username::String, permission::Permission)::Bool =
294+
username::String, permission::Permission.PermissionEnum)::Bool =
295295
updateexperimentpermission(instance, experiment.experiment_id, username, permission)
296296

297297
"""

src/services/misc.jl

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,33 +10,31 @@ Get a list of all values for the specified [`Metric`](@ref) for a given [`Run`](
1010
- `metric_key`: Name of the [`Metric`](@ref) to fetch.
1111
- `page_token`: Token indicating the page of [`Metric`](@ref) history to fetch.
1212
- `max_results`: Maximum number of logged instances of a [`Metric`](@ref) for a
13-
[`Run`](@ref) to return per call.
13+
[`Run`](@ref) to return per call. Defaults to 50.
1414
1515
# Returns
1616
- A list of all historical values for the specified [`Metric`](@ref) in the specified
1717
[`Run`](@ref).
1818
- The next page token if there are more results.
1919
"""
2020
function getmetrichistory(instance::MLFlow, run_id::String, metric_key::String;
21-
page_token::String="", max_results::Union{Int64,Missing}=missing
21+
page_token::Union{String,Missing}=missing, max_results::Int64=50
2222
)::Tuple{Array{Metric},Union{String,Nothing}}
2323
result = mlfget(instance, "metrics/get-history"; run_id=run_id, metric_key=metric_key,
2424
page_token=page_token,
2525
max_results=(ismissing(max_results) ? max_results : (max_results |> Int32)))
2626

27-
metrics = result["metrics"] |> (x -> [Metric(y) for y in x])
27+
metrics = get(result, "metrics", []) |> (x -> [Metric(y) for y in x])
2828
next_page_token = get(result, "next_page_token", nothing)
2929

3030
return metrics, next_page_token
3131
end
3232
getmetrichistory(instance::MLFlow, run::Run, metric_key::String; page_token::String="",
33-
max_results::Union{Int64,Missing}=missing
34-
)::Tuple{Array{Metric},Union{String,Nothing}} =
33+
max_results::Int64=50)::Tuple{Array{Metric},Union{String,Nothing}} =
3534
getmetrichistory(instance, run.info.run_id, metric_key; page_token=page_token,
3635
max_results=max_results)
3736
getmetrichistory(instance::MLFlow, run::Run, metric::Metric; page_token::String="",
38-
max_results::Union{Int64,Missing}=missing
39-
)::Tuple{Array{Metric},Union{String,Nothing}} =
37+
max_results::Int64=50)::Tuple{Array{Metric},Union{String,Nothing}} =
4038
getmetrichistory(instance, run.info.run_id, metric.key; page_token=page_token,
4139
max_results=max_results)
4240

src/services/registered_model.jl

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,11 @@ name already exists.
1717
An instance of type [`RegisteredModel`](@ref).
1818
"""
1919
function createregisteredmodel(instance::MLFlow, name::String;
20-
tags::MLFlowUpsertData{Tag}=Tag[],
21-
description::Union{String,Missing}=missing)::RegisteredModel
20+
tags::MLFlowUpsertData{Tag}=Tag[], description::Union{String,Missing}=missing,
21+
deployment_job_id::Union{String,Missing}=missing)::RegisteredModel
2222
result = mlfpost(instance, "registered-models/create"; name=name,
23-
tags=parse(Tag, tags), description=description)
23+
tags=parse(Tag, tags), description=description,
24+
deployment_job_id=deployment_job_id)
2425
return result["registered_model"] |> RegisteredModel
2526
end
2627

@@ -205,7 +206,7 @@ end
205206
An instance of type [`RegisteredModelPermission`](@ref).
206207
"""
207208
function createregisteredmodelpermission(instance::MLFlow, name::String, username::String,
208-
permission::Permission)::RegisteredModelPermission
209+
permission::Permission.PermissionEnum)::RegisteredModelPermission
209210
result = mlfpost(instance, "registered-models/permissions/create"; name=name,
210211
username=username, permission=permission)
211212
return result["registered_model_permission"] |> RegisteredModelPermission
@@ -243,7 +244,7 @@ end
243244
`true` if successful. Otherwise, raises exception.
244245
"""
245246
function updateregisteredmodelpermission(instance::MLFlow, name::String, username::String,
246-
permission::Permission)::Bool
247+
permission::Permission.PermissionEnum)::Bool
247248
mlfpatch(instance, "registered-models/permissions/update"; name=name, username=username,
248249
permission=permission)
249250
return true

src/services/run.jl

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -170,8 +170,8 @@ Search for runs that satisfy expressions. Search expressions can use [`Metric`](
170170
- The next page token if there are more results.
171171
"""
172172
function searchruns(instance::MLFlow; experiment_ids::Array{String}=String[],
173-
filter::String="", run_view_type::ViewType=ACTIVE_ONLY, max_results::Int=1000,
174-
order_by::Array{String}=String[],
173+
filter::String="", run_view_type::ViewType.ViewTypeEnum=ViewType.ACTIVE_ONLY,
174+
max_results::Int=1000, order_by::Array{String}=String[],
175175
page_token::String="")::Tuple{Array{Run},Union{String,Nothing}}
176176
parameters = (; experiment_ids, filter, :run_view_type => run_view_type |> Integer,
177177
max_results, page_token)
@@ -207,13 +207,14 @@ Update [`Run`](@ref) metadata.
207207
- An instance of type [`RunInfo`](@ref) with the updated metadata.
208208
"""
209209
function updaterun(instance::MLFlow, run_id::String;
210-
status::Union{RunStatus,Missing}=missing, end_time::Union{Int64,Missing}=missing,
211-
run_name::Union{String,Missing})::RunInfo
210+
status::Union{RunStatus.RunStatusEnum,Missing}=missing,
211+
end_time::Union{Int64,Missing}=missing, run_name::Union{String,Missing})::RunInfo
212212
result = mlfpost(instance, "runs/update"; run_id=run_id, status=(status |> Integer),
213213
end_time=end_time, run_name=run_name)
214214
return result["run_info"] |> RunInfo
215215
end
216-
updaterun(instance::MLFlow, run::Run; status::Union{RunStatus,Missing}=missing,
216+
updaterun(instance::MLFlow, run::Run;
217+
status::Union{RunStatus.RunStatusEnum,Missing}=missing,
217218
end_time::Union{Int64,Missing}=missing, run_name::Union{String,Missing})::RunInfo =
218219
updaterun(instance, run.info.run_id; status=status, end_time=end_time,
219220
run_name=run_name)

src/types/enums.jl

Lines changed: 105 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -1,71 +1,114 @@
1-
"""
2-
ModelVersionStatus
3-
4-
# Members
5-
- `PENDING_REGISTRATION`: Request to register a new model version is pending as server
6-
performs background tasks.
7-
- `FAILED_REGISTRATION`: Request to register a new model version has failed.
8-
- `READY`: Model version is ready for use.
9-
"""
10-
@enum ModelVersionStatus begin
11-
PENDING_REGISTRATION = 1
12-
FAILED_REGISTRATION = 2
13-
READY = 3
1+
module ModelVersionStatus
2+
"""
3+
ModelVersionStatus
4+
5+
# Members
6+
- `PENDING_REGISTRATION`: Request to register a new model version is pending as server
7+
performs background tasks.
8+
- `FAILED_REGISTRATION`: Request to register a new model version has failed.
9+
- `READY`: Model version is ready for use.
10+
"""
11+
@enum ModelVersionStatusEnum begin
12+
PENDING_REGISTRATION = 1
13+
FAILED_REGISTRATION = 2
14+
READY = 3
15+
end
16+
parse(status::String) = Dict(value => key for (key, value) in ModelVersionStatusEnum |> Base.Enums.namemap)[status|>Symbol] |> ModelVersionStatusEnum
1417
end
15-
ModelVersionStatus(status::String) = Dict(value => key for (key, value) in ModelVersionStatus |> Base.Enums.namemap)[status|>Symbol] |> ModelVersionStatus
16-
17-
"""
18-
RunStatus
1918

20-
Status of a run.
21-
22-
# Members
23-
- `RUNNING`: Run has been initiated.
24-
- `SCHEDULED`: Run is scheduled to run at a later time.
25-
- `FINISHED`: Run has completed.
26-
- `FAILED`: Run execution failed.
27-
- `KILLED`: Run killed by user.
28-
"""
29-
@enum RunStatus begin
30-
RUNNING = 1
31-
SCHEDULED = 2
32-
FINISHED = 3
33-
FAILED = 4
34-
KILLED = 5
19+
module RunStatus
20+
"""
21+
RunStatus
22+
23+
Status of a run.
24+
25+
# Members
26+
- `RUNNING`: Run has been initiated.
27+
- `SCHEDULED`: Run is scheduled to run at a later time.
28+
- `FINISHED`: Run has completed.
29+
- `FAILED`: Run execution failed.
30+
- `KILLED`: Run killed by user.
31+
"""
32+
@enum RunStatusEnum begin
33+
RUNNING = 1
34+
SCHEDULED = 2
35+
FINISHED = 3
36+
FAILED = 4
37+
KILLED = 5
38+
end
39+
parse(status::String) = Dict(value => key for (key, value) in RunStatusEnum |> Base.Enums.namemap)[status|>Symbol] |> RunStatusEnum
3540
end
36-
RunStatus(status::String) = Dict(value => key for (key, value) in RunStatus |> Base.Enums.namemap)[status|>Symbol] |> RunStatus
37-
38-
"""
39-
ViewType
4041

41-
View type for ListExperiments query.
42-
43-
# Members
44-
- `ACTIVE_ONLY`: Default. Return only active experiments.
45-
- `DELETED_ONLY`: Return only deleted experiments.
46-
- `ALL`: Get all experiments.
47-
"""
48-
@enum ViewType begin
49-
ACTIVE_ONLY = 1
50-
DELETED_ONLY = 2
51-
ALL = 3
42+
module ViewType
43+
"""
44+
ViewType
45+
46+
View type for ListExperiments query.
47+
48+
# Members
49+
- `ACTIVE_ONLY`: Default. Return only active experiments.
50+
- `DELETED_ONLY`: Return only deleted experiments.
51+
- `ALL`: Get all experiments.
52+
"""
53+
@enum ViewTypeEnum begin
54+
ACTIVE_ONLY = 1
55+
DELETED_ONLY = 2
56+
ALL = 3
57+
end
5258
end
5359

54-
"""
55-
Permission
60+
module Permission
61+
"""
62+
Permission
63+
64+
Permission of a user to an experiment or a registered model.
65+
66+
# Members
67+
- `READ`: Can read.
68+
- `EDIT`: Can read and update.
69+
- `MANAGE`: Can read, update, delete and manage.
70+
- `NO_PERMISSIONS`: No permissions.
71+
"""
72+
@enum PermissionEnum begin
73+
READ = 1
74+
EDIT = 2
75+
MANAGE = 3
76+
NO_PERMISSIONS = 4
77+
end
78+
parse(permission::String) = Dict(value => key for (key, value) in PermissionEnum |> Base.Enums.namemap)[permission|>Symbol] |> PermissionEnum
79+
end
5680

57-
Permission of a user to an experiment or a registered model.
81+
module DeploymentJobRunState
82+
@enum DeploymentJobRunStateEnum begin
83+
DEPLOYMENT_JOB_RUN_STATE_UNSPECIFIED = 1
84+
NO_VALID_DEPLOYMENT_JOB_FOUND = 2
85+
RUNNING = 3
86+
SUCCEEDED = 4
87+
FAILED = 5
88+
PENDING = 6
89+
APPROVAL = 7
90+
end
91+
parse(state::String) = Dict(value => key for (key, value) in DeploymentJobRunStateEnum |> Base.Enums.namemap)[state|>Symbol] |> DeploymentJobRunStateEnum
92+
end
5893

59-
# Members
60-
- `READ`: Can read.
61-
- `EDIT`: Can read and update.
62-
- `MANAGE`: Can read, update, delete and manage.
63-
- `NO_PERMISSIONS`: No permissions.
64-
"""
65-
@enum Permission begin
66-
READ = 1
67-
EDIT = 2
68-
MANAGE = 3
69-
NO_PERMISSIONS = 4
94+
module State
95+
"""
96+
State
97+
98+
# Members
99+
- `DEPLOYMENT_JOB_CONNECTION_STATE_UNSPECIFIED`
100+
- `NOT_SET_UP`: Default state.
101+
- `CONNECTED`: Connected job: job exists, owner has ACLs, and required job parameters are
102+
present.
103+
- `NOT_FOUND`: Job was deleted or owner had job ACLs removed.
104+
- `REQUIRED_PARAMETERS_CHANGED`: Required job parameters were changed.
105+
"""
106+
@enum StateEnum begin
107+
DEPLOYMENT_JOB_CONNECTION_STATE_UNSPECIFIED = 1
108+
NOT_SET_UP = 2
109+
CONNECTED = 3
110+
NOT_FOUND = 4
111+
REQUIRED_PARAMETERS_CHANGED = 5
112+
end
113+
parse(state::String) = Dict(value => key for (key, value) in StateEnum |> Base.Enums.namemap)[state|>Symbol] |> StateEnum
70114
end
71-
Permission(permission::String) = Dict(value => key for (key, value) in Permission |> Base.Enums.namemap)[permission|>Symbol] |> Permission

src/types/experiment.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,8 @@ Base.show(io::IO, t::Experiment) = show(io, ShowCase(t, new_lines=true))
3636
struct ExperimentPermission
3737
experiment_id::String
3838
user_id::String
39-
permission::Permission
39+
permission::Permission.PermissionEnum
4040
end
4141
ExperimentPermission(data::Dict{String,Any}) = ExperimentPermission(data["experiment_id"],
42-
data["user_id"] |> string, Permission(data["permission"]))
42+
data["user_id"] |> string, Permission.parse(data["permission"]))
4343
Base.show(io::IO, t::ExperimentPermission) = show(io, ShowCase(t, new_lines=true))

0 commit comments

Comments
 (0)