Skip to content

Commit f4bd787

Browse files
[mlflow_utils] Mlflow upgrade 3.5 (#966) (#969)
* fix(mlflow): resolve deprecation warnings for mlflow 3.5 - Replace deprecated artifact_path with output_path parameter - Migrate tracking backend from filesystem to SQLite * update function and item yaml * deleted filename from function.yaml Co-authored-by: Omer Mimon <81911093+omermaim@users.noreply.github.com>
1 parent 5b5225e commit f4bd787

File tree

4 files changed

+55
-31
lines changed

4 files changed

+55
-31
lines changed
Lines changed: 18 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,30 @@
1+
kind: serving
12
verbose: false
3+
metadata:
4+
categories:
5+
- model-serving
6+
- utilities
7+
name: mlflow-utils
8+
tag: ''
29
spec:
3-
command: ''
4-
source: ''
5-
default_class: MLFlowModelServer
10+
image: mlrun/mlrun
611
function_kind: serving_v2
12+
disable_auto_mount: false
13+
max_replicas: 4
14+
min_replicas: 1
15+
function_handler: mlflow-utils-nuclio:handler
716
build:
817
functionSourceCode: aW1wb3J0IHppcGZpbGUKZnJvbSB0eXBpbmcgaW1wb3J0IEFueSwgRGljdAppbXBvcnQgbWxmbG93CmZyb20gbWxydW4uc2VydmluZy52Ml9zZXJ2aW5nIGltcG9ydCBWMk1vZGVsU2VydmVyCmltcG9ydCBwYW5kYXMgYXMgcGQKCgpjbGFzcyBNTEZsb3dNb2RlbFNlcnZlcihWMk1vZGVsU2VydmVyKToKICAgICIiIgogICAgTUxGbG93IHRyYWNrZXIgTW9kZWwgc2VydmluZyBjbGFzcywgaW5oZXJpdGluZyB0aGUgVjJNb2RlbFNlcnZlciBjbGFzcyBmb3IgYmVpbmcgaW5pdGlhbGl6ZWQgYXV0b21hdGljYWxseSBieSB0aGUgbW9kZWwKICAgIHNlcnZlciBhbmQgYmUgYWJsZSB0byBydW4gbG9jYWxseSBhcyBwYXJ0IG9mIGEgbnVjbGlvIHNlcnZlcmxlc3MgZnVuY3Rpb24sIG9yIGFzIHBhcnQgb2YgYSByZWFsLXRpbWUgcGlwZWxpbmUuCiAgICAiIiIKCiAgICBkZWYgbG9hZChzZWxmKToKICAgICAgICAiIiIKICAgICAgICBsb2FkcyBhIG1vZGVsIHRoYXQgd2FzIGxvZ2dlZCBieSB0aGUgTUxGbG93IHRyYWNrZXIgbW9kZWwKICAgICAgICAiIiIKICAgICAgICAjIFVuemlwIHRoZSBtb2RlbCBkaXIgYW5kIHRoZW4gdXNlIG1sZmxvdydzIGxvYWQgZnVuY3Rpb24KICAgICAgICBtb2RlbF9maWxlLCBfID0gc2VsZi5nZXRfbW9kZWwoIi56aXAiKQogICAgICAgIG1vZGVsX3BhdGhfdW56aXAgPSBtb2RlbF9maWxlLnJlcGxhY2UoIi56aXAiLCAiIikKCiAgICAgICAgd2l0aCB6aXBmaWxlLlppcEZpbGUobW9kZWxfZmlsZSwgInIiKSBhcyB6aXBfcmVmOgogICAgICAgICAgICB6aXBfcmVmLmV4dHJhY3RhbGwobW9kZWxfcGF0aF91bnppcCkKCiAgICAgICAgc2VsZi5tb2RlbCA9IG1sZmxvdy5weWZ1bmMubG9hZF9tb2RlbChtb2RlbF9wYXRoX3VuemlwKQoKICAgIGRlZiBwcmVkaWN0KHNlbGYsIHJlcXVlc3Q6IERpY3Rbc3RyLCBBbnldKSAtPiBsaXN0OgogICAgICAgICIiIgogICAgICAgIEluZmVyIHRoZSBpbnB1dHMgdGhyb3VnaCB0aGUgbW9kZWwuIFRoZSBpbmZlcnJlZCBkYXRhIHdpbGwKICAgICAgICBiZSByZWFkIGZyb20gdGhlICJpbnB1dHMiIGtleSBvZiB0aGUgcmVxdWVzdC4KCiAgICAgICAgOnBhcmFtIHJlcXVlc3Q6IFRoZSByZXF1ZXN0IHRvIHRoZSBtb2RlbCB1c2luZyB4Z2Jvb3N0J3MgcHJlZGljdC4KICAgICAgICAgICAgICAgIFRoZSBpbnB1dCB0byB0aGUgbW9kZWwgd2lsbCBiZSByZWFkIGZyb20gdGhlICJpbnB1dHMiIGtleS4KCiAgICAgICAgOnJldHVybjogVGhlIG1vZGVsJ3MgcHJlZGljdGlvbiBvbiB0aGUgZ2l2ZW4gaW5wdXQuCiAgICAgICAgIiIiCgogICAgICAgICMgR2V0IHRoZSBpbnB1dHMgYW5kIHNldCB0byBhY2NlcHRlZCB0eXBlOgogICAgICAgIGlucHV0cyA9IHBkLkRhdGFGcmFtZShyZXF1ZXN0WyJpbnB1dHMiXSkKCiAgICAgICAgIyBQcmVkaWN0IHVzaW5nIHRoZSBtb2RlbCdzIHByZWRpY3QgZnVuY3Rpb246CiAgICAgICAgcHJlZGljdGlvbnMgPSBzZWxmLm1vZGVsLnByZWRpY3QoaW5wdXRzKQoKICAgICAgICAjIFJldHVybiBhcyBsaXN0OgogICAgICAgIHJldHVybiBwcmVkaWN0aW9ucy50b2xpc3QoKQoKZnJvbSBtbHJ1bi5ydW50aW1lcyBpbXBvcnQgbnVjbGlvX2luaXRfaG9vawpkZWYgaW5pdF9jb250ZXh0KGNvbnRleHQpOgogICAgbnVjbGlvX2luaXRfaG9vayhjb250ZXh0LCBnbG9iYWxzKCksICdzZXJ2aW5nX3YyJykKCmRlZiBoYW5kbGVyKGNvbnRleHQsIGV2ZW50KToKICAgIHJldHVybiBjb250ZXh0Lm1scnVuX2hhbmRsZXIoY29udGV4dCwgZXZlbnQpCg==
918
requirements:
10-
- mlflow==2.12.2
11-
- lightgbm
12-
- xgboost
13-
code_origin: ''
19+
- mlflow~=3.5
1420
origin_filename: ''
15-
image: mlrun/mlrun
21+
code_origin: ''
22+
description: Mlflow model server, and additional utils.
23+
command: ''
1624
base_image_pull: false
25+
default_class: MLFlowModelServer
26+
source: ''
1727
default_handler: ''
18-
max_replicas: 4
19-
disable_auto_mount: false
20-
min_replicas: 1
21-
description: Mlflow model server, and additional utils.
22-
function_handler: mlflow-utils-nuclio:handler
2328
env:
2429
- name: MLRUN_HTTPDB__NUCLIO__EXPLICIT_ACK
2530
value: enabled
26-
metadata:
27-
categories:
28-
- model-serving
29-
- utilities
30-
name: mlflow-utils
31-
tag: ''
32-
kind: serving

functions/src/mlflow_utils/item.yaml

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ labels:
1212
author: Iguazio
1313
maintainers: []
1414
marketplaceType: ''
15-
mlrunVersion: 1.8.0
15+
mlrunVersion: 1.10.0
1616
name: mlflow_utils
1717
platformVersion: ''
1818
spec:
@@ -23,8 +23,6 @@ spec:
2323
image: mlrun/mlrun
2424
kind: serving
2525
requirements:
26-
- mlflow~=2.22
27-
- lightgbm
28-
- xgboost
26+
- mlflow~=3.5
2927
url: ''
30-
version: 1.1.0
28+
version: 1.2.0
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
1-
mlflow~=2.22
1+
mlflow~=3.5
22
lightgbm
33
xgboost

functions/src/mlflow_utils/test_mlflow_utils.py

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
# limitations under the License.
1414
#
1515
import tempfile
16-
16+
import shutil
1717
import lightgbm as lgb
1818
import mlflow
1919
import mlflow.environment_variables
@@ -132,7 +132,9 @@ def test_track_run_with_experiment_name(handler):
132132
# Set the mlflow experiment name
133133
mlflow.environment_variables.MLFLOW_EXPERIMENT_NAME.set(f"{handler}_test_track")
134134
with tempfile.TemporaryDirectory() as test_directory:
135-
mlflow.set_tracking_uri(test_directory) # Tell mlflow where to save logged data
135+
# Use SQLite backend instead of filesystem (filesystem will be deprecated in Feb 2026)
136+
db_uri = f"sqlite:///{os.path.join(test_directory, 'mlflow.db')}"
137+
mlflow.set_tracking_uri(db_uri) # Tell mlflow where to save logged data
136138

137139
# Create a project for this tester:
138140
project = mlrun.get_or_create_project(name="default", context=test_directory)
@@ -149,17 +151,43 @@ def test_track_run_with_experiment_name(handler):
149151
trainer_run = func.run(
150152
local=True,
151153
handler=handler,
152-
artifact_path=test_directory,
154+
output_path=test_directory,
153155
)
154156

157+
# Find the MLflow logged model and prepare it for serving
158+
# Note: In MLflow 2.24+, we must dynamically discover model paths since MLflow changed
159+
# its directory structure from predictable paths (e.g., experiment_name/0/model/) to
160+
# UUID-based paths (e.g., experiment_id/run_uuid/artifacts/model/).
161+
162+
# Create MLflow client to query the tracking server
163+
mlflow_client = mlflow.tracking.MlflowClient(tracking_uri=db_uri)
164+
165+
# Get the experiment by name to obtain its ID
166+
experiment = mlflow_client.get_experiment_by_name(f"{handler}_test_track")
167+
168+
# Search for runs in this experiment and get the run ID
169+
# (There should only be one run from our training above)
170+
run_id = mlflow_client.search_runs(experiment_ids=[experiment.experiment_id])[0].info.run_id
171+
172+
# Find all models logged in this run
173+
logged_models = mlflow.search_logged_models(filter_string=f"source_run_id = '{run_id}'")
174+
175+
# Extract the artifact location and remove the "file://" prefix
176+
model_artifacts_dir = logged_models["artifact_location"].tolist()[0].replace("file://", "")
177+
178+
# Package the model artifacts as a zip file for MLFlowModelServer
179+
# Note: MLFlowModelServer requires models to be packaged as zip archives
180+
# rather than loose directories for deployment
181+
model_path = os.path.join(test_directory, f"{handler}-model-serving")
182+
os.makedirs(model_path, exist_ok=True)
183+
shutil.make_archive(os.path.join(model_path, "model"), 'zip', model_artifacts_dir)
184+
155185
serving_func = project.set_function(
156186
func=os.path.abspath("function.yaml"),
157187
name=f"{handler}-server",
158188
)
159189
model_name = f"{handler}-model"
160190
# Add the model
161-
upper_handler = handler.replace("_", "-")
162-
model_path = test_directory + f"/{upper_handler}-test-{upper_handler}/0/model/"
163191
serving_func.add_model(
164192
model_name,
165193
class_name="MLFlowModelServer",

0 commit comments

Comments
 (0)