test_uv_utils.py
1 import subprocess 2 from unittest import mock 3 4 import pytest 5 from packaging.version import Version 6 7 from mlflow.environment_variables import MLFLOW_UV_AUTO_DETECT 8 from mlflow.utils.environment import infer_pip_requirements 9 from mlflow.utils.uv_utils import ( 10 _PYPROJECT_FILE, 11 _UV_LOCK_FILE, 12 copy_uv_project_files, 13 create_uv_sync_pyproject, 14 detect_uv_project, 15 export_uv_requirements, 16 extract_index_urls_from_uv_lock, 17 get_uv_version, 18 has_uv_lock_artifact, 19 is_uv_available, 20 run_uv_sync, 21 setup_uv_sync_environment, 22 ) 23 24 # --- get_uv_version tests --- 25 26 27 def test_get_uv_version_returns_none_when_uv_not_installed(): 28 with mock.patch("mlflow.utils.uv_utils.shutil.which", return_value=None): 29 assert get_uv_version() is None 30 31 32 def test_get_uv_version_returns_version_when_uv_installed(): 33 mock_result = mock.Mock() 34 mock_result.stdout = "uv 0.6.10 (abc123 2024-01-01)" 35 with ( 36 mock.patch("mlflow.utils.uv_utils.shutil.which", return_value="/usr/bin/uv"), 37 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result) as mock_run, 38 ): 39 version = get_uv_version() 40 assert version == Version("0.6.10") 41 mock_run.assert_called_once() 42 43 44 def test_get_uv_version_returns_none_on_subprocess_error(): 45 with ( 46 mock.patch("mlflow.utils.uv_utils.shutil.which", return_value="/usr/bin/uv"), 47 mock.patch( 48 "mlflow.utils.uv_utils.subprocess.run", 49 side_effect=subprocess.CalledProcessError(1, "uv"), 50 ), 51 ): 52 assert get_uv_version() is None 53 54 55 def test_get_uv_version_returns_none_on_parse_error(): 56 mock_result = mock.Mock() 57 mock_result.stdout = "invalid output" 58 with ( 59 mock.patch("mlflow.utils.uv_utils.shutil.which", return_value="/usr/bin/uv"), 60 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result), 61 ): 62 assert get_uv_version() is None 63 64 65 # --- is_uv_available tests --- 66 67 68 def test_is_uv_available_returns_false_when_uv_not_installed(): 69 with mock.patch("mlflow.utils.uv_utils.shutil.which", return_value=None): 70 assert is_uv_available() is False 71 72 73 def test_is_uv_available_returns_false_when_version_below_minimum(): 74 mock_result = mock.Mock() 75 mock_result.stdout = "uv 0.4.0 (abc123 2024-01-01)" 76 with ( 77 mock.patch("mlflow.utils.uv_utils.shutil.which", return_value="/usr/bin/uv"), 78 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result), 79 ): 80 assert is_uv_available() is False 81 82 83 @pytest.mark.parametrize("version_str", ["0.6.10", "1.0.0"]) 84 def test_is_uv_available_returns_true_when_version_meets_or_exceeds_minimum(version_str): 85 mock_result = mock.Mock() 86 mock_result.stdout = f"uv {version_str} (abc123 2024-01-01)" 87 with ( 88 mock.patch("mlflow.utils.uv_utils.shutil.which", return_value="/usr/bin/uv"), 89 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result), 90 ): 91 assert is_uv_available() is True 92 93 94 # --- detect_uv_project tests --- 95 96 97 @pytest.mark.parametrize( 98 ("create_uv_lock", "create_pyproject"), 99 [(False, True), (True, False)], 100 ids=["missing_uv_lock", "missing_pyproject"], 101 ) 102 def test_detect_uv_project_returns_none_when_file_missing( 103 tmp_path, create_uv_lock, create_pyproject 104 ): 105 if create_uv_lock: 106 (tmp_path / _UV_LOCK_FILE).touch() 107 if create_pyproject: 108 (tmp_path / _PYPROJECT_FILE).touch() 109 assert detect_uv_project(tmp_path) is None 110 111 112 def test_detect_uv_project_returns_paths_when_both_files_exist(tmp_path): 113 (tmp_path / _UV_LOCK_FILE).touch() 114 (tmp_path / _PYPROJECT_FILE).touch() 115 116 result = detect_uv_project(tmp_path) 117 assert result is not None 118 assert result.uv_lock == tmp_path / _UV_LOCK_FILE 119 assert result.pyproject == tmp_path / _PYPROJECT_FILE 120 121 122 def test_detect_uv_project_uses_cwd_when_directory_not_specified(tmp_path, monkeypatch): 123 monkeypatch.chdir(tmp_path) 124 (tmp_path / _UV_LOCK_FILE).touch() 125 (tmp_path / _PYPROJECT_FILE).touch() 126 127 result = detect_uv_project() 128 assert result is not None 129 assert result.uv_lock == tmp_path / _UV_LOCK_FILE 130 131 132 # --- export_uv_requirements tests --- 133 134 135 def test_export_uv_requirements_returns_none_when_uv_not_available(): 136 with mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value=None): 137 assert export_uv_requirements() is None 138 139 140 def test_export_uv_requirements_returns_requirements_list(tmp_path): 141 uv_output = """requests==2.28.0 142 numpy==1.24.0 143 pandas==2.0.0 144 """ 145 mock_result = mock.Mock() 146 mock_result.stdout = uv_output 147 148 with ( 149 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 150 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result) as mock_run, 151 ): 152 result = export_uv_requirements(tmp_path) 153 154 assert result == ["requests==2.28.0", "numpy==1.24.0", "pandas==2.0.0"] 155 mock_run.assert_called_once() 156 157 158 def test_export_uv_requirements_preserves_environment_markers(tmp_path): 159 uv_output = """requests==2.28.0 160 pywin32==306 ; sys_platform == 'win32' 161 numpy==1.24.0 162 """ 163 mock_result = mock.Mock() 164 mock_result.stdout = uv_output 165 166 with ( 167 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 168 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result), 169 ): 170 result = export_uv_requirements(tmp_path) 171 172 assert result is not None 173 assert len(result) == 3 174 assert "pywin32==306 ; sys_platform == 'win32'" in result 175 176 177 def test_export_uv_requirements_keeps_all_marker_variants(tmp_path): 178 uv_output = """numpy==2.2.6 ; python_version < '3.11' 179 numpy==2.4.1 ; python_version >= '3.11' 180 """ 181 mock_result = mock.Mock() 182 mock_result.stdout = uv_output 183 184 with ( 185 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 186 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result), 187 ): 188 result = export_uv_requirements(tmp_path) 189 190 assert result is not None 191 numpy_entries = [r for r in result if r.startswith("numpy")] 192 assert len(numpy_entries) == 2 193 194 195 def test_export_uv_requirements_returns_none_on_subprocess_error(tmp_path): 196 with ( 197 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 198 mock.patch( 199 "mlflow.utils.uv_utils.subprocess.run", 200 side_effect=subprocess.CalledProcessError(1, "uv"), 201 ), 202 ): 203 assert export_uv_requirements(tmp_path) is None 204 205 206 def test_export_uv_requirements_with_explicit_directory(tmp_path): 207 (tmp_path / _UV_LOCK_FILE).touch() 208 209 uv_output = """requests==2.28.0 210 numpy==1.24.0 211 """ 212 mock_result = mock.Mock() 213 mock_result.stdout = uv_output 214 215 with ( 216 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 217 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result) as mock_run, 218 ): 219 result = export_uv_requirements(directory=tmp_path) 220 221 assert result is not None 222 assert "requests==2.28.0" in result 223 assert "numpy==1.24.0" in result 224 mock_run.assert_called_once() 225 assert mock_run.call_args.kwargs["cwd"] == tmp_path 226 227 228 # --- export_uv_requirements groups/extras tests --- 229 230 231 def test_export_uv_requirements_passes_groups_to_command(tmp_path): 232 mock_result = mock.Mock() 233 mock_result.stdout = "requests==2.28.0\n" 234 235 with ( 236 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 237 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result) as mock_run, 238 ): 239 export_uv_requirements(tmp_path, groups=["serving", "ml"]) 240 241 call_args = mock_run.call_args[0][0] 242 assert "--group" in call_args 243 serving_idx = call_args.index("--group") 244 assert call_args[serving_idx + 1] == "serving" 245 second_group_idx = call_args.index("--group", serving_idx + 1) 246 assert call_args[second_group_idx + 1] == "ml" 247 248 249 def test_export_uv_requirements_passes_extras_to_command(tmp_path): 250 mock_result = mock.Mock() 251 mock_result.stdout = "uvicorn==0.29.0\n" 252 253 with ( 254 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 255 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result) as mock_run, 256 ): 257 export_uv_requirements(tmp_path, extras=["api", "gpu"]) 258 259 call_args = mock_run.call_args[0][0] 260 assert "--extra" in call_args 261 api_idx = call_args.index("--extra") 262 assert call_args[api_idx + 1] == "api" 263 second_extra_idx = call_args.index("--extra", api_idx + 1) 264 assert call_args[second_extra_idx + 1] == "gpu" 265 266 267 # --- copy_uv_project_files tests --- 268 269 270 def test_copy_uv_project_files_returns_false_when_not_uv_project(tmp_path): 271 dest_dir = tmp_path / "dest" 272 dest_dir.mkdir() 273 274 source_dir = tmp_path / "source" 275 source_dir.mkdir() 276 277 result = copy_uv_project_files(dest_dir, source_dir) 278 assert result is False 279 280 281 def test_copy_uv_project_files_copies_files_when_uv_project(tmp_path): 282 source_dir = tmp_path / "source" 283 source_dir.mkdir() 284 (source_dir / _UV_LOCK_FILE).write_text("lock content") 285 (source_dir / _PYPROJECT_FILE).write_text("pyproject content") 286 287 dest_dir = tmp_path / "dest" 288 dest_dir.mkdir() 289 290 result = copy_uv_project_files(dest_dir, source_dir) 291 292 assert result is True 293 assert (dest_dir / _UV_LOCK_FILE).read_text() == "lock content" 294 assert (dest_dir / _PYPROJECT_FILE).read_text() == "pyproject content" 295 296 297 @pytest.mark.parametrize( 298 ("has_python_version", "expected_exists"), 299 [(True, True), (False, False)], 300 ids=["with_python_version", "without_python_version"], 301 ) 302 def test_copy_uv_project_files_python_version_handling( 303 tmp_path, has_python_version, expected_exists 304 ): 305 source_dir = tmp_path / "source" 306 source_dir.mkdir() 307 (source_dir / _UV_LOCK_FILE).write_text("lock content") 308 (source_dir / _PYPROJECT_FILE).write_text("pyproject content") 309 if has_python_version: 310 (source_dir / ".python-version").write_text("3.11.5") 311 312 dest_dir = tmp_path / "dest" 313 dest_dir.mkdir() 314 315 result = copy_uv_project_files(dest_dir, source_dir) 316 317 assert result is True 318 assert (dest_dir / ".python-version").exists() == expected_exists 319 if expected_exists: 320 assert (dest_dir / ".python-version").read_text() == "3.11.5" 321 322 323 def test_copy_uv_project_files_respects_mlflow_log_uv_files_env_false(tmp_path, monkeypatch): 324 monkeypatch.setenv("MLFLOW_LOG_UV_FILES", "false") 325 326 source_dir = tmp_path / "source" 327 source_dir.mkdir() 328 (source_dir / _UV_LOCK_FILE).write_text("lock content") 329 (source_dir / _PYPROJECT_FILE).write_text("pyproject content") 330 331 dest_dir = tmp_path / "dest" 332 dest_dir.mkdir() 333 334 result = copy_uv_project_files(dest_dir, source_dir) 335 336 assert result is False 337 assert not (dest_dir / _UV_LOCK_FILE).exists() 338 assert not (dest_dir / _PYPROJECT_FILE).exists() 339 340 341 @pytest.mark.parametrize("env_value", ["0", "false", "FALSE", "False"]) 342 def test_copy_uv_project_files_env_var_false_variants(tmp_path, monkeypatch, env_value): 343 monkeypatch.setenv("MLFLOW_LOG_UV_FILES", env_value) 344 345 source_dir = tmp_path / "source" 346 source_dir.mkdir() 347 (source_dir / _UV_LOCK_FILE).write_text("lock content") 348 (source_dir / _PYPROJECT_FILE).write_text("pyproject content") 349 350 dest_dir = tmp_path / "dest" 351 dest_dir.mkdir() 352 353 result = copy_uv_project_files(dest_dir, source_dir) 354 assert result is False 355 356 357 @pytest.mark.parametrize("env_value", ["true", "1", "TRUE", "True"]) 358 def test_copy_uv_project_files_env_var_true_variants(tmp_path, monkeypatch, env_value): 359 monkeypatch.setenv("MLFLOW_LOG_UV_FILES", env_value) 360 361 source_dir = tmp_path / "source" 362 source_dir.mkdir() 363 (source_dir / _UV_LOCK_FILE).write_text("lock content") 364 (source_dir / _PYPROJECT_FILE).write_text("pyproject content") 365 366 dest_dir = tmp_path / "dest" 367 dest_dir.mkdir() 368 369 result = copy_uv_project_files(dest_dir, source_dir) 370 assert result is True 371 372 373 def test_copy_uv_project_files_with_monorepo_layout(tmp_path): 374 project_dir = tmp_path / "monorepo" / "subproject" 375 project_dir.mkdir(parents=True) 376 (project_dir / _UV_LOCK_FILE).write_text("lock content from monorepo") 377 (project_dir / _PYPROJECT_FILE).write_text("pyproject from monorepo") 378 (project_dir / ".python-version").write_text("3.12.0") 379 380 dest_dir = tmp_path / "dest" 381 dest_dir.mkdir() 382 383 result = copy_uv_project_files(dest_dir, source_dir=project_dir) 384 385 assert result is True 386 assert (dest_dir / _UV_LOCK_FILE).read_text() == "lock content from monorepo" 387 assert (dest_dir / ".python-version").read_text() == "3.12.0" 388 389 390 def test_copy_uv_project_files_with_nonexistent_source(tmp_path): 391 dest_dir = tmp_path / "dest" 392 dest_dir.mkdir() 393 nonexistent_dir = tmp_path / "nonexistent" 394 395 result = copy_uv_project_files(dest_dir, source_dir=nonexistent_dir) 396 assert result is False 397 398 399 def test_copy_uv_project_files_with_missing_pyproject(tmp_path): 400 project_dir = tmp_path / "incomplete_project" 401 project_dir.mkdir() 402 (project_dir / _UV_LOCK_FILE).write_text("lock content") 403 404 dest_dir = tmp_path / "dest" 405 dest_dir.mkdir() 406 407 result = copy_uv_project_files(dest_dir, source_dir=project_dir) 408 assert result is False 409 410 411 # --- Integration tests for infer_pip_requirements uv path --- 412 413 414 def test_infer_pip_requirements_uses_uv_when_project_detected(tmp_path, monkeypatch): 415 monkeypatch.chdir(tmp_path) 416 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 417 (tmp_path / _UV_LOCK_FILE).touch() 418 (tmp_path / _PYPROJECT_FILE).touch() 419 420 uv_output = "requests==2.28.0\nnumpy==1.24.0\n" 421 mock_result = mock.Mock() 422 mock_result.stdout = uv_output 423 424 with ( 425 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 426 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result), 427 ): 428 result = infer_pip_requirements("runs:/fake/model", "sklearn") 429 430 assert "requests==2.28.0" in result 431 assert "numpy==1.24.0" in result 432 433 434 def test_export_uv_requirements_returns_none_when_uv_binary_missing(tmp_path): 435 with mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value=None): 436 result = export_uv_requirements(tmp_path) 437 assert result is None 438 439 440 def test_infer_pip_requirements_passes_groups_and_extras_to_uv_export(tmp_path, monkeypatch): 441 monkeypatch.chdir(tmp_path) 442 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 443 (tmp_path / _UV_LOCK_FILE).touch() 444 (tmp_path / _PYPROJECT_FILE).touch() 445 446 uv_output = "fastapi==0.100.0\nuvicorn==0.29.0\n" 447 mock_result = mock.Mock() 448 mock_result.stdout = uv_output 449 450 with ( 451 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 452 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result) as mock_run, 453 ): 454 result = infer_pip_requirements( 455 str(tmp_path), 456 "sklearn", 457 uv_groups=["serving"], 458 uv_extras=["api"], 459 ) 460 461 assert "fastapi==0.100.0" in result 462 assert "uvicorn==0.29.0" in result 463 464 call_args = mock_run.call_args[0][0] 465 assert "--group" in call_args 466 group_idx = call_args.index("--group") 467 assert call_args[group_idx + 1] == "serving" 468 assert "--extra" in call_args 469 extra_idx = call_args.index("--extra") 470 assert call_args[extra_idx + 1] == "api" 471 472 473 def test_detect_uv_project_not_detected_when_files_missing(tmp_path, monkeypatch): 474 monkeypatch.chdir(tmp_path) 475 assert detect_uv_project() is None 476 477 478 # --- MLFLOW_UV_AUTO_DETECT Environment Variable Tests --- 479 480 481 def test_mlflow_uv_auto_detect_returns_true_by_default(monkeypatch): 482 monkeypatch.delenv("MLFLOW_UV_AUTO_DETECT", raising=False) 483 assert MLFLOW_UV_AUTO_DETECT.get() is True 484 485 486 @pytest.mark.parametrize("env_value", ["false", "0", "FALSE", "False"]) 487 def test_mlflow_uv_auto_detect_returns_false_when_disabled(monkeypatch, env_value): 488 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", env_value) 489 assert MLFLOW_UV_AUTO_DETECT.get() is False 490 491 492 @pytest.mark.parametrize("env_value", ["true", "1", "TRUE", "True"]) 493 def test_mlflow_uv_auto_detect_returns_true_when_enabled(monkeypatch, env_value): 494 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", env_value) 495 assert MLFLOW_UV_AUTO_DETECT.get() is True 496 497 498 def test_infer_pip_requirements_skips_uv_when_auto_detect_disabled(tmp_path, monkeypatch): 499 monkeypatch.chdir(tmp_path) 500 (tmp_path / "uv.lock").touch() 501 (tmp_path / "pyproject.toml").touch() 502 503 assert detect_uv_project() is not None 504 505 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "false") 506 507 with ( 508 mock.patch("mlflow.utils.environment.detect_uv_project") as mock_detect, 509 mock.patch("mlflow.utils.environment.export_uv_requirements") as mock_export, 510 mock.patch( 511 "mlflow.utils.environment._infer_requirements", 512 return_value=["scikit-learn==1.0"], 513 ), 514 ): 515 result = infer_pip_requirements(str(tmp_path), "sklearn") 516 517 mock_detect.assert_not_called() 518 mock_export.assert_not_called() 519 assert "scikit-learn==1.0" in result 520 521 522 def test_infer_pip_requirements_uses_explicit_uv_project_dir(tmp_path, monkeypatch): 523 work_dir = tmp_path / "work" 524 work_dir.mkdir() 525 monkeypatch.chdir(work_dir) 526 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 527 528 uv_project = tmp_path / "my_project" 529 uv_project.mkdir() 530 (uv_project / _UV_LOCK_FILE).touch() 531 (uv_project / _PYPROJECT_FILE).touch() 532 533 uv_output = "requests==2.28.0\n" 534 mock_result = mock.Mock() 535 mock_result.stdout = uv_output 536 537 with ( 538 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 539 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result), 540 ): 541 result = infer_pip_requirements(str(tmp_path), "sklearn", uv_project_dir=uv_project) 542 543 assert "requests==2.28.0" in result 544 545 546 def test_infer_pip_requirements_explicit_uv_project_dir_overrides_disabled_auto_detect( 547 tmp_path, monkeypatch 548 ): 549 monkeypatch.chdir(tmp_path) 550 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "false") 551 552 uv_project = tmp_path / "my_project" 553 uv_project.mkdir() 554 (uv_project / _UV_LOCK_FILE).touch() 555 (uv_project / _PYPROJECT_FILE).touch() 556 557 uv_output = "numpy==1.24.0\n" 558 mock_result = mock.Mock() 559 mock_result.stdout = uv_output 560 561 with ( 562 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 563 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result), 564 ): 565 result = infer_pip_requirements(str(tmp_path), "sklearn", uv_project_dir=uv_project) 566 567 assert "numpy==1.24.0" in result 568 569 570 def test_export_uv_requirements_strips_comment_lines(tmp_path): 571 uv_output = """requests==2.28.0 572 # via 573 # some-package 574 urllib3==1.26.0 575 # via requests 576 certifi==2023.7.22 577 """ 578 mock_result = mock.Mock() 579 mock_result.stdout = uv_output 580 581 with ( 582 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 583 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result), 584 ): 585 result = export_uv_requirements(tmp_path) 586 587 assert result == ["requests==2.28.0", "urllib3==1.26.0", "certifi==2023.7.22"] 588 589 590 def test_export_uv_requirements_returns_empty_list_on_empty_output(tmp_path): 591 mock_result = mock.Mock() 592 mock_result.stdout = "" 593 594 with ( 595 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 596 mock.patch("mlflow.utils.uv_utils.subprocess.run", return_value=mock_result), 597 ): 598 result = export_uv_requirements(tmp_path) 599 600 assert result == [] 601 602 603 # --- Private Index URL Extraction Tests --- 604 605 606 @pytest.mark.parametrize( 607 ("uv_lock_content", "expected_urls"), 608 [ 609 ( 610 """ 611 version = 1 612 requires-python = ">=3.11" 613 614 [[package]] 615 name = "my-private-pkg" 616 version = "1.0.0" 617 source = { registry = "https://internal.company.com/simple" } 618 619 [[package]] 620 name = "numpy" 621 version = "1.24.0" 622 source = { registry = "https://pypi.org/simple" } 623 """, 624 ["https://internal.company.com/simple"], 625 ), 626 ( 627 """ 628 version = 1 629 630 [[package]] 631 name = "pkg1" 632 source = { registry = "https://private1.com/simple" } 633 634 [[package]] 635 name = "pkg2" 636 source = { registry = "https://private2.com/simple" } 637 638 [[package]] 639 name = "pkg3" 640 source = { registry = "https://private1.com/simple" } 641 """, 642 ["https://private1.com/simple", "https://private2.com/simple"], 643 ), 644 ( 645 """ 646 version = 1 647 648 [[package]] 649 name = "numpy" 650 source = { registry = "https://pypi.org/simple" } 651 """, 652 [], 653 ), 654 ], 655 ids=["single_private", "multiple_private_deduped", "no_private"], 656 ) 657 def test_extract_index_urls_from_uv_lock(tmp_path, uv_lock_content, expected_urls): 658 uv_lock_path = tmp_path / "uv.lock" 659 uv_lock_path.write_text(uv_lock_content) 660 661 result = extract_index_urls_from_uv_lock(uv_lock_path) 662 assert result == expected_urls 663 664 665 def test_extract_index_urls_from_uv_lock_file_not_exists(tmp_path): 666 result = extract_index_urls_from_uv_lock(tmp_path / "nonexistent.lock") 667 assert result == [] 668 669 670 # --- uv Sync Environment Setup Tests --- 671 672 673 @pytest.mark.parametrize( 674 ("python_version", "project_name", "expected_name", "expected_python"), 675 [ 676 ("3.11.5", "mlflow-model-env", "mlflow-model-env", "==3.11.5"), 677 ("3.10.14", "my-custom-env", "my-custom-env", "==3.10.14"), 678 ], 679 ids=["default_name", "custom_name"], 680 ) 681 def test_create_uv_sync_pyproject( 682 tmp_path, python_version, project_name, expected_name, expected_python 683 ): 684 result_path = create_uv_sync_pyproject(tmp_path, python_version, project_name=project_name) 685 686 assert result_path.exists() 687 content = result_path.read_text() 688 assert f'name = "{expected_name}"' in content 689 assert f'requires-python = "{expected_python}"' in content 690 691 692 def test_setup_uv_sync_environment(tmp_path): 693 model_path = tmp_path / "model" 694 model_path.mkdir() 695 (model_path / "uv.lock").write_text('version = 1\nrequires-python = ">=3.11"') 696 (model_path / ".python-version").write_text("3.11.5") 697 698 env_dir = tmp_path / "env" 699 700 result = setup_uv_sync_environment(env_dir, model_path, "3.11.5") 701 702 assert result is True 703 assert (env_dir / "uv.lock").exists() 704 assert (env_dir / "pyproject.toml").exists() 705 assert (env_dir / ".python-version").exists() 706 707 708 def test_setup_uv_sync_environment_copies_existing_pyproject(tmp_path): 709 model_path = tmp_path / "model" 710 model_path.mkdir() 711 original_pyproject = '[project]\nname = "my-model"\nversion = "1.0.0"\n' 712 (model_path / "uv.lock").write_text('version = 1\nrequires-python = ">=3.11"') 713 (model_path / "pyproject.toml").write_text(original_pyproject) 714 715 env_dir = tmp_path / "env" 716 717 result = setup_uv_sync_environment(env_dir, model_path, "3.11.5") 718 719 assert result is True 720 # Copied from model, not generated (should have "my-model" not "mlflow-model-env") 721 pyproject_content = (env_dir / "pyproject.toml").read_text() 722 assert 'name = "my-model"' in pyproject_content 723 assert "mlflow-model-env" not in pyproject_content 724 725 726 def test_setup_uv_sync_environment_no_uv_lock(tmp_path): 727 model_path = tmp_path / "model" 728 model_path.mkdir() 729 730 env_dir = tmp_path / "env" 731 732 result = setup_uv_sync_environment(env_dir, model_path, "3.11") 733 734 assert result is False 735 assert not env_dir.exists() 736 737 738 def test_has_uv_lock_artifact(tmp_path): 739 model_path = tmp_path / "model" 740 model_path.mkdir() 741 742 assert has_uv_lock_artifact(model_path) is False 743 744 (model_path / "uv.lock").write_text("version = 1") 745 assert has_uv_lock_artifact(model_path) is True 746 747 748 def test_run_uv_sync_returns_false_when_uv_not_available(tmp_path): 749 with mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value=None): 750 result = run_uv_sync(tmp_path) 751 assert result is False 752 753 754 def test_run_uv_sync_builds_correct_command(tmp_path): 755 with ( 756 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 757 mock.patch("mlflow.utils.uv_utils.subprocess.run") as mock_run, 758 ): 759 run_uv_sync(tmp_path, frozen=True, no_dev=True) 760 761 mock_run.assert_called_once() 762 call_args = mock_run.call_args[0][0] 763 assert call_args[0] == "/usr/bin/uv" 764 assert call_args[1] == "sync" 765 assert "--frozen" in call_args 766 assert "--no-dev" in call_args 767 768 769 def test_run_uv_sync_returns_false_on_failure(tmp_path): 770 with ( 771 mock.patch("mlflow.utils.uv_utils._get_uv_binary", return_value="/usr/bin/uv"), 772 mock.patch( 773 "mlflow.utils.uv_utils.subprocess.run", 774 side_effect=subprocess.CalledProcessError(1, "uv sync"), 775 ), 776 ): 777 result = run_uv_sync(tmp_path) 778 assert result is False 779 780 781 def test_infer_pip_requirements_warns_when_groups_set_but_no_uv_project(tmp_path, monkeypatch): 782 monkeypatch.chdir(tmp_path) 783 monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "true") 784 785 with ( 786 mock.patch( 787 "mlflow.utils.environment._infer_requirements", 788 return_value=["scikit-learn==1.0"], 789 ), 790 mock.patch("mlflow.utils.environment._logger") as mock_logger, 791 ): 792 result = infer_pip_requirements( 793 str(tmp_path), 794 "sklearn", 795 uv_groups=["serving"], 796 uv_extras=["api"], 797 ) 798 799 assert "scikit-learn==1.0" in result 800 mock_logger.warning.assert_any_call( 801 "uv_groups and/or uv_extras were specified but no uv project was detected. " 802 "These parameters will be ignored. Falling back to package capture based inference." 803 )