test_uv_model_logging.py
1 """ 2 Integration tests for uv package manager support in model logging and loading. 3 4 Tests the end-to-end workflow: 5 1. uv project detection during log_model() 6 2. Artifact generation (uv.lock, pyproject.toml, .python-version, requirements.txt) 7 3. Model loading with uv artifacts 8 9 These tests use REAL uv calls (not mocked) where possible, following MLflow best practices. 10 Tests requiring uv are skipped if uv is not installed or below minimum version. 11 """ 12 13 import platform 14 import shutil 15 import subprocess 16 import sys 17 from pathlib import Path 18 from unittest import mock 19 20 import pytest 21 22 import mlflow 23 import mlflow.pyfunc 24 from mlflow.utils.os import is_windows 25 from mlflow.utils.uv_utils import ( 26 _PYPROJECT_FILE, 27 _PYTHON_VERSION_FILE, 28 _UV_LOCK_FILE, 29 is_uv_available, 30 ) 31 32 # Constants for artifact file names 33 _REQUIREMENTS_FILE_NAME = "requirements.txt" 34 _PYTHON_ENV_FILE_NAME = "python_env.yaml" 35 36 # Skip marker for tests requiring uv 37 requires_uv = pytest.mark.skipif( 38 not is_uv_available(), 39 reason="uv is not installed or below minimum required version (0.6.10)", 40 ) 41 42 43 class SimplePythonModel(mlflow.pyfunc.PythonModel): 44 def predict(self, context, model_input, params=None): 45 return model_input 46 47 48 @pytest.fixture 49 def python_model(): 50 return SimplePythonModel() 51 52 53 PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" 54 55 56 @pytest.fixture 57 def tmp_uv_project(tmp_path): 58 """Create a real uv project with uv lock.""" 59 pyproject_content = """[project] 60 name = "test_uv_project" 61 version = "0.1.0" 62 requires-python = ">=3.10" 63 dependencies = [ 64 "numpy>=1.24.0", 65 ] 66 67 [build-system] 68 requires = ["hatchling"] 69 build-backend = "hatchling.build" 70 """ 71 (tmp_path / _PYPROJECT_FILE).write_text(pyproject_content) 72 73 # Create .python-version 74 (tmp_path / _PYTHON_VERSION_FILE).write_text(f"{PYTHON_VERSION}\n") 75 76 # Create minimal package structure for hatchling 77 pkg_dir = tmp_path / "test_uv_project" 78 pkg_dir.mkdir() 79 (pkg_dir / "__init__.py").write_text('"""Test uv project."""\n__version__ = "0.1.0"\n') 80 81 # Run uv lock to generate real uv.lock 82 result = subprocess.run( 83 ["uv", "lock"], 84 cwd=tmp_path, 85 capture_output=True, 86 text=True, 87 ) 88 if result.returncode != 0: 89 pytest.skip(f"uv lock failed: {result.stderr}") 90 91 return tmp_path 92 93 94 # --- Model Logging Tests with Real uv --- 95 96 97 @requires_uv 98 def test_pyfunc_log_model_copies_uv_artifacts(tmp_uv_project, python_model, monkeypatch): 99 monkeypatch.chdir(tmp_uv_project) 100 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 101 102 with mlflow.start_run() as run: 103 mlflow.pyfunc.log_model(name="model", python_model=python_model) 104 105 artifact_path = mlflow.artifacts.download_artifacts( 106 run_id=run.info.run_id, artifact_path="model" 107 ) 108 artifact_dir = Path(artifact_path) 109 110 # Verify uv artifacts are copied 111 assert (artifact_dir / _UV_LOCK_FILE).exists() 112 assert (artifact_dir / _PYPROJECT_FILE).exists() 113 assert (artifact_dir / _PYTHON_VERSION_FILE).exists() 114 115 # Verify content matches source 116 assert "version = 1" in (artifact_dir / _UV_LOCK_FILE).read_text() 117 assert "test_uv_project" in (artifact_dir / _PYPROJECT_FILE).read_text() 118 assert PYTHON_VERSION in (artifact_dir / _PYTHON_VERSION_FILE).read_text() 119 120 121 @requires_uv 122 def test_pyfunc_log_model_python_env_uses_current_python_version( 123 tmp_uv_project, python_model, monkeypatch 124 ): 125 monkeypatch.chdir(tmp_uv_project) 126 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 127 128 with mlflow.start_run() as run: 129 mlflow.pyfunc.log_model(name="model", python_model=python_model) 130 131 artifact_path = mlflow.artifacts.download_artifacts( 132 run_id=run.info.run_id, artifact_path="model" 133 ) 134 artifact_dir = Path(artifact_path) 135 136 python_env_file = artifact_dir / _PYTHON_ENV_FILE_NAME 137 assert python_env_file.exists() 138 python_env_content = python_env_file.read_text() 139 # python_env.yaml always uses the current interpreter version 140 assert platform.python_version() in python_env_content 141 142 143 @requires_uv 144 def test_pyfunc_log_model_respects_mlflow_log_uv_files_env_var( 145 tmp_uv_project, python_model, monkeypatch 146 ): 147 monkeypatch.chdir(tmp_uv_project) 148 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 149 monkeypatch.setenv("MLFLOW_LOG_UV_FILES", "false") 150 151 with mlflow.start_run() as run: 152 mlflow.pyfunc.log_model(name="model", python_model=python_model) 153 154 artifact_path = mlflow.artifacts.download_artifacts( 155 run_id=run.info.run_id, artifact_path="model" 156 ) 157 artifact_dir = Path(artifact_path) 158 159 # uv artifacts should NOT be copied when env var is false 160 assert not (artifact_dir / _UV_LOCK_FILE).exists() 161 assert not (artifact_dir / _PYPROJECT_FILE).exists() 162 163 # But requirements.txt should still exist (from uv export) 164 assert (artifact_dir / _REQUIREMENTS_FILE_NAME).exists() 165 requirements_content = (artifact_dir / _REQUIREMENTS_FILE_NAME).read_text() 166 assert "numpy" in requirements_content.lower() 167 168 169 @requires_uv 170 def test_pyfunc_log_model_with_explicit_uv_project_path_parameter( 171 tmp_path, tmp_uv_project, python_model, monkeypatch 172 ): 173 # Work from a different directory than the uv project 174 work_dir = tmp_path / "work" 175 work_dir.mkdir() 176 monkeypatch.chdir(work_dir) 177 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 178 179 with mlflow.start_run() as run: 180 mlflow.pyfunc.log_model( 181 name="model", 182 python_model=python_model, 183 uv_project_path=tmp_uv_project, 184 ) 185 186 artifact_path = mlflow.artifacts.download_artifacts( 187 run_id=run.info.run_id, artifact_path="model" 188 ) 189 artifact_dir = Path(artifact_path) 190 191 assert (artifact_dir / _UV_LOCK_FILE).exists() 192 assert (artifact_dir / _PYPROJECT_FILE).exists() 193 assert "test_uv_project" in (artifact_dir / _PYPROJECT_FILE).read_text() 194 195 196 @requires_uv 197 def test_pyfunc_log_model_generates_requirements_from_uv_export( 198 tmp_uv_project, python_model, monkeypatch 199 ): 200 monkeypatch.chdir(tmp_uv_project) 201 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 202 203 with mlflow.start_run() as run: 204 mlflow.pyfunc.log_model(name="model", python_model=python_model) 205 206 artifact_path = mlflow.artifacts.download_artifacts( 207 run_id=run.info.run_id, artifact_path="model" 208 ) 209 artifact_dir = Path(artifact_path) 210 211 requirements_file = artifact_dir / _REQUIREMENTS_FILE_NAME 212 assert requirements_file.exists() 213 requirements_content = requirements_file.read_text() 214 assert "numpy" in requirements_content.lower() 215 216 217 # --- Fallback Tests (mocking required to simulate uv unavailable) --- 218 219 220 def test_pyfunc_log_model_falls_back_when_uv_not_available(tmp_path, python_model, monkeypatch): 221 (tmp_path / _UV_LOCK_FILE).write_text('version = 1\nrequires-python = ">=3.10"\n') 222 (tmp_path / _PYPROJECT_FILE).write_text('[project]\nname = "test"\nversion = "0.1.0"\n') 223 monkeypatch.chdir(tmp_path) 224 225 with mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value=None): 226 with mlflow.start_run() as run: 227 mlflow.pyfunc.log_model(name="model", python_model=python_model) 228 229 artifact_path = mlflow.artifacts.download_artifacts( 230 run_id=run.info.run_id, artifact_path="model" 231 ) 232 artifact_dir = Path(artifact_path) 233 234 assert (artifact_dir / _REQUIREMENTS_FILE_NAME).exists() 235 236 237 def test_pyfunc_log_model_falls_back_when_uv_export_fails(tmp_path, python_model, monkeypatch): 238 (tmp_path / _UV_LOCK_FILE).write_text('version = 1\nrequires-python = ">=3.10"\n') 239 (tmp_path / _PYPROJECT_FILE).write_text('[project]\nname = "test"\nversion = "0.1.0"\n') 240 monkeypatch.chdir(tmp_path) 241 242 with ( 243 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 244 mock.patch( 245 "mlflow.utils.uv_utils.subprocess.run", 246 side_effect=subprocess.CalledProcessError(1, "uv"), 247 ), 248 ): 249 with mlflow.start_run() as run: 250 mlflow.pyfunc.log_model(name="model", python_model=python_model) 251 252 artifact_path = mlflow.artifacts.download_artifacts( 253 run_id=run.info.run_id, artifact_path="model" 254 ) 255 artifact_dir = Path(artifact_path) 256 257 assert (artifact_dir / _REQUIREMENTS_FILE_NAME).exists() 258 259 260 def test_pyfunc_log_model_non_uv_project_uses_standard_inference( 261 python_model, tmp_path, monkeypatch 262 ): 263 monkeypatch.chdir(tmp_path) 264 265 with mlflow.start_run() as run: 266 mlflow.pyfunc.log_model(name="model", python_model=python_model) 267 268 artifact_path = mlflow.artifacts.download_artifacts( 269 run_id=run.info.run_id, artifact_path="model" 270 ) 271 artifact_dir = Path(artifact_path) 272 273 assert (artifact_dir / _REQUIREMENTS_FILE_NAME).exists() 274 assert (artifact_dir / _PYTHON_ENV_FILE_NAME).exists() 275 assert not (artifact_dir / _UV_LOCK_FILE).exists() 276 assert not (artifact_dir / _PYPROJECT_FILE).exists() 277 278 279 # --- Model Loading Tests --- 280 281 282 @requires_uv 283 def test_load_pyfunc_model_with_uv_artifacts_and_predict(tmp_uv_project, python_model, monkeypatch): 284 monkeypatch.chdir(tmp_uv_project) 285 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 286 287 with mlflow.start_run() as run: 288 mlflow.pyfunc.log_model(name="model", python_model=python_model) 289 model_uri = f"runs:/{run.info.run_id}/model" 290 291 loaded_model = mlflow.pyfunc.load_model(model_uri) 292 293 assert loaded_model is not None 294 assert loaded_model.metadata is not None 295 296 import pandas as pd 297 298 test_input = pd.DataFrame({"a": [1, 2, 3]}) 299 predictions = loaded_model.predict(test_input) 300 assert predictions is not None 301 302 303 # --- Save Model Tests --- 304 305 306 @requires_uv 307 def test_pyfunc_save_model_with_uv_project(tmp_uv_project, python_model, tmp_path, monkeypatch): 308 monkeypatch.chdir(tmp_uv_project) 309 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 310 model_path = tmp_path / "saved_model" 311 312 mlflow.pyfunc.save_model(model_path, python_model=python_model) 313 314 assert (model_path / _REQUIREMENTS_FILE_NAME).exists() 315 assert (model_path / _UV_LOCK_FILE).exists() 316 assert (model_path / _PYPROJECT_FILE).exists() 317 assert (model_path / _PYTHON_VERSION_FILE).exists() 318 319 320 @requires_uv 321 def test_pyfunc_save_model_with_explicit_uv_project_path( 322 tmp_uv_project, python_model, tmp_path, monkeypatch 323 ): 324 work_dir = tmp_path / "work" 325 work_dir.mkdir() 326 model_path = tmp_path / "saved_model" 327 monkeypatch.chdir(work_dir) 328 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 329 330 mlflow.pyfunc.save_model( 331 model_path, 332 python_model=python_model, 333 uv_project_path=tmp_uv_project, 334 ) 335 336 assert (model_path / _UV_LOCK_FILE).exists() 337 assert (model_path / _PYPROJECT_FILE).exists() 338 339 340 # --- Environment Variable Variations --- 341 342 343 @requires_uv 344 @pytest.mark.parametrize("env_value", ["false", "0", "FALSE", "False"]) 345 def test_mlflow_log_uv_files_env_var_false_variants( 346 tmp_uv_project, python_model, monkeypatch, env_value 347 ): 348 monkeypatch.chdir(tmp_uv_project) 349 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 350 monkeypatch.setenv("MLFLOW_LOG_UV_FILES", env_value) 351 352 with mlflow.start_run() as run: 353 mlflow.pyfunc.log_model(name="model", python_model=python_model) 354 355 artifact_path = mlflow.artifacts.download_artifacts( 356 run_id=run.info.run_id, artifact_path="model" 357 ) 358 artifact_dir = Path(artifact_path) 359 360 assert not (artifact_dir / _UV_LOCK_FILE).exists() 361 assert not (artifact_dir / _PYPROJECT_FILE).exists() 362 assert (artifact_dir / _REQUIREMENTS_FILE_NAME).exists() 363 364 365 @requires_uv 366 @pytest.mark.parametrize("env_value", ["true", "1", "TRUE", "True"]) 367 def test_mlflow_log_uv_files_env_var_true_variants( 368 tmp_uv_project, python_model, monkeypatch, env_value 369 ): 370 monkeypatch.chdir(tmp_uv_project) 371 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 372 monkeypatch.setenv("MLFLOW_LOG_UV_FILES", env_value) 373 374 with mlflow.start_run() as run: 375 mlflow.pyfunc.log_model(name="model", python_model=python_model) 376 377 artifact_path = mlflow.artifacts.download_artifacts( 378 run_id=run.info.run_id, artifact_path="model" 379 ) 380 artifact_dir = Path(artifact_path) 381 382 assert (artifact_dir / _UV_LOCK_FILE).exists() 383 assert (artifact_dir / _PYPROJECT_FILE).exists() 384 385 386 # --- Dependency Groups Integration Tests --- 387 388 389 @pytest.fixture 390 def uv_project_with_groups(tmp_path): 391 """Create a uv project with dependency groups.""" 392 pyproject_content = """[project] 393 name = "test_uv_groups" 394 version = "0.1.0" 395 requires-python = ">=3.10" 396 dependencies = [ 397 "numpy>=1.24.0", 398 ] 399 400 [project.optional-dependencies] 401 gpu = ["scipy>=1.10.0"] 402 403 [dependency-groups] 404 serving = ["gunicorn>=21.0.0"] 405 dev = ["pytest>=7.0.0"] 406 407 [build-system] 408 requires = ["hatchling"] 409 build-backend = "hatchling.build" 410 """ 411 (tmp_path / _PYPROJECT_FILE).write_text(pyproject_content) 412 (tmp_path / _PYTHON_VERSION_FILE).write_text("3.11.5\n") 413 414 pkg_dir = tmp_path / "test_uv_groups" 415 pkg_dir.mkdir() 416 (pkg_dir / "__init__.py").write_text('"""Test uv groups project."""\n__version__ = "0.1.0"\n') 417 418 result = subprocess.run( 419 ["uv", "lock"], 420 cwd=tmp_path, 421 capture_output=True, 422 text=True, 423 ) 424 if result.returncode != 0: 425 pytest.skip(f"uv lock failed: {result.stderr}") 426 427 return tmp_path 428 429 430 @requires_uv 431 def test_export_uv_requirements_with_groups_real(uv_project_with_groups): 432 from mlflow.utils.uv_utils import export_uv_requirements 433 434 result = export_uv_requirements(uv_project_with_groups, groups=["serving"]) 435 436 assert result is not None 437 pkg_names = [r.split("==")[0].lower() for r in result] 438 assert "numpy" in pkg_names 439 assert "gunicorn" in pkg_names 440 assert "pytest" not in pkg_names 441 442 443 @requires_uv 444 def test_export_uv_requirements_with_extras_real(uv_project_with_groups): 445 from mlflow.utils.uv_utils import export_uv_requirements 446 447 result = export_uv_requirements(uv_project_with_groups, extras=["gpu"]) 448 449 assert result is not None 450 pkg_names = [r.split("==")[0].lower() for r in result] 451 assert "numpy" in pkg_names 452 assert "scipy" in pkg_names 453 454 455 # --- uv Sync Environment Setup Integration Tests --- 456 457 458 @requires_uv 459 def test_setup_uv_sync_environment_real(tmp_uv_project, tmp_path): 460 from mlflow.utils.uv_utils import has_uv_lock_artifact, setup_uv_sync_environment 461 462 model_artifacts = tmp_path / "model_artifacts" 463 model_artifacts.mkdir() 464 465 shutil.copy2(tmp_uv_project / _UV_LOCK_FILE, model_artifacts / _UV_LOCK_FILE) 466 shutil.copy2(tmp_uv_project / _PYTHON_VERSION_FILE, model_artifacts / _PYTHON_VERSION_FILE) 467 468 assert has_uv_lock_artifact(model_artifacts) 469 470 env_dir = tmp_path / "env" 471 result = setup_uv_sync_environment(env_dir, model_artifacts, "3.11.5") 472 473 assert result is True 474 assert (env_dir / _UV_LOCK_FILE).exists() 475 assert (env_dir / _PYPROJECT_FILE).exists() 476 assert (env_dir / _PYTHON_VERSION_FILE).exists() 477 478 # No pyproject.toml in model_artifacts, so create_uv_sync_pyproject 479 # generates one with pinned version 480 pyproject_content = (env_dir / _PYPROJECT_FILE).read_text() 481 assert 'name = "mlflow-model-env"' in pyproject_content 482 assert 'requires-python = "==3.11.5"' in pyproject_content 483 484 485 @requires_uv 486 def test_extract_index_urls_from_real_uv_lock(tmp_uv_project): 487 from mlflow.utils.uv_utils import extract_index_urls_from_uv_lock 488 489 result = extract_index_urls_from_uv_lock(tmp_uv_project / _UV_LOCK_FILE) 490 491 well_known_public = {"https://download.pytorch.org/whl/cpu"} 492 truly_private = [url for url in result if url not in well_known_public] 493 assert truly_private == [] 494 495 496 @pytest.mark.skipif(is_windows(), reason="This test fails on Windows") 497 @requires_uv 498 def test_run_uv_sync_real(tmp_uv_project, tmp_path): 499 from mlflow.utils.uv_utils import run_uv_sync 500 501 sync_dir = tmp_path / "sync_project" 502 503 # Create the virtual environment directly at sync_dir 504 subprocess.check_call(["uv", "venv", sync_dir, f"--python={PYTHON_VERSION}"]) 505 506 shutil.copytree(tmp_uv_project, sync_dir, dirs_exist_ok=True) 507 508 result = run_uv_sync(sync_dir, frozen=True, no_dev=True) 509 510 assert result is True 511 512 # Verify numpy is installed in the env at sync_dir 513 python_bin = sync_dir / "bin" / "python" 514 515 subprocess.check_call([python_bin, "-c", "import numpy"])