wehub-resource-sync
470 lines
15 KiB
Python
470 lines
15 KiB
Python
"""
|
|
Utilities for uv package manager integration.
|
|
|
|
This module provides functions for detecting uv projects and exporting dependencies
|
|
via ``uv export`` for automatic dependency inference during model logging.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
from typing import NamedTuple
|
|
|
|
from packaging.version import Version
|
|
|
|
from mlflow.environment_variables import MLFLOW_LOG_UV_FILES
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
# Minimum uv version required for ``uv export`` functionality
|
|
_MIN_UV_VERSION = Version("0.6.10")
|
|
|
|
# File names used for uv project detection and artifacts
|
|
_UV_LOCK_FILE = "uv.lock"
|
|
_PYPROJECT_FILE = "pyproject.toml"
|
|
_PYTHON_VERSION_FILE = ".python-version"
|
|
|
|
|
|
def get_uv_version() -> Version | None:
|
|
"""
|
|
Get the installed uv version.
|
|
|
|
Returns:
|
|
The uv version as a packaging.version.Version object, or None if uv is not installed
|
|
or version cannot be determined.
|
|
"""
|
|
uv_bin = shutil.which("uv")
|
|
if uv_bin is None:
|
|
return None
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
[uv_bin, "--version"],
|
|
capture_output=True,
|
|
text=True,
|
|
check=True,
|
|
)
|
|
# Output format: "uv 0.5.0 (abc123 2024-01-01)"
|
|
version_str = result.stdout.strip().split()[1]
|
|
return Version(version_str)
|
|
except (subprocess.CalledProcessError, IndexError, ValueError) as e:
|
|
_logger.debug(f"Failed to determine uv version: {e}")
|
|
return None
|
|
|
|
|
|
def _get_uv_binary() -> str | None:
|
|
"""Return the uv binary path if installed and meeting version requirements, else None."""
|
|
uv_bin = shutil.which("uv")
|
|
if uv_bin is None:
|
|
return None
|
|
|
|
version = get_uv_version()
|
|
if version is None:
|
|
return None
|
|
|
|
if version < _MIN_UV_VERSION:
|
|
_logger.debug(f"uv version {version} is below minimum required version {_MIN_UV_VERSION}")
|
|
return None
|
|
|
|
return uv_bin
|
|
|
|
|
|
def is_uv_available() -> bool:
|
|
"""
|
|
Check if uv is installed and meets the minimum version requirement.
|
|
|
|
Returns:
|
|
True if uv is installed and version >= 0.6.10, False otherwise.
|
|
"""
|
|
return _get_uv_binary() is not None
|
|
|
|
|
|
class UVProjectInfo(NamedTuple):
|
|
"""Paths for a detected uv project."""
|
|
|
|
uv_lock: Path
|
|
pyproject: Path
|
|
|
|
|
|
def detect_uv_project(directory: str | Path | None = None) -> UVProjectInfo | None:
|
|
"""
|
|
Detect if the given directory is a uv project.
|
|
|
|
A uv project is identified by the presence of BOTH ``uv.lock`` and ``pyproject.toml``
|
|
in the specified directory.
|
|
|
|
Args:
|
|
directory: The directory to check. Defaults to the current working directory.
|
|
|
|
Returns:
|
|
A UVProjectInfo with paths to uv.lock and pyproject.toml,
|
|
or None if the directory is not a uv project.
|
|
"""
|
|
directory = Path.cwd() if directory is None else Path(directory)
|
|
|
|
uv_lock_path = directory / _UV_LOCK_FILE
|
|
pyproject_path = directory / _PYPROJECT_FILE
|
|
|
|
if uv_lock_path.exists() and pyproject_path.exists():
|
|
_logger.info(
|
|
f"Detected uv project: found {_UV_LOCK_FILE} and {_PYPROJECT_FILE} in {directory}"
|
|
)
|
|
return UVProjectInfo(uv_lock=uv_lock_path, pyproject=pyproject_path)
|
|
|
|
return None
|
|
|
|
|
|
def export_uv_requirements(
|
|
directory: str | Path | None = None,
|
|
no_dev: bool = True,
|
|
no_hashes: bool = True,
|
|
frozen: bool = True,
|
|
groups: list[str] | None = None,
|
|
extras: list[str] | None = None,
|
|
) -> list[str] | None:
|
|
"""
|
|
Export dependencies from a uv project to pip-compatible requirements.
|
|
|
|
Runs ``uv export`` to generate a list of pinned dependencies from the uv lock file.
|
|
Environment markers (e.g. ``; python_version < '3.12'``) are preserved in the output
|
|
so that pip can evaluate them at install time.
|
|
|
|
Args:
|
|
directory: The uv project directory. Defaults to the current working directory.
|
|
no_dev: Exclude development dependencies. Defaults to True.
|
|
no_hashes: Omit hashes from output. Defaults to True.
|
|
frozen: Use frozen lockfile without updating. Defaults to True.
|
|
groups: Optional list of dependency groups to include (additive with project deps).
|
|
Maps to ``uv export --group <name>``.
|
|
extras: Optional list of optional dependency extras to include.
|
|
Maps to ``uv export --extra <name>``.
|
|
|
|
Returns:
|
|
A list of requirement strings (e.g., ``["requests==2.28.0", "numpy==1.24.0"]``),
|
|
or None if export fails.
|
|
"""
|
|
uv_bin = _get_uv_binary()
|
|
if uv_bin is None:
|
|
_logger.warning(
|
|
"uv is not available or version is below minimum required. "
|
|
"Falling back to pip-based inference."
|
|
)
|
|
return None
|
|
directory = Path.cwd() if directory is None else Path(directory)
|
|
|
|
cmd = [uv_bin, "export"]
|
|
|
|
if no_dev:
|
|
cmd.append("--no-dev")
|
|
if no_hashes:
|
|
cmd.append("--no-hashes")
|
|
if frozen:
|
|
cmd.append("--frozen")
|
|
|
|
if groups:
|
|
for group in groups:
|
|
cmd.extend(["--group", group])
|
|
|
|
if extras:
|
|
for extra in extras:
|
|
cmd.extend(["--extra", extra])
|
|
|
|
# Additional flags for cleaner output
|
|
cmd.extend([
|
|
"--no-header", # Omit the "autogenerated by uv" comment
|
|
"--no-emit-project", # Don't emit the project itself
|
|
"--no-annotate", # Omit "# via ..." comment annotations
|
|
])
|
|
|
|
try:
|
|
_logger.debug(f"Running uv export: {' '.join(str(c) for c in cmd)}")
|
|
result = subprocess.run(
|
|
cmd,
|
|
capture_output=True,
|
|
text=True,
|
|
check=True,
|
|
cwd=directory,
|
|
)
|
|
|
|
requirements = [
|
|
line.strip()
|
|
for line in result.stdout.strip().splitlines()
|
|
if line.strip() and not line.strip().startswith("#")
|
|
]
|
|
|
|
_logger.info(f"Exported {len(requirements)} dependencies via uv")
|
|
return requirements
|
|
|
|
except subprocess.CalledProcessError as e:
|
|
_logger.warning(
|
|
f"uv export failed with exit code {e.returncode}. "
|
|
f"stderr: {e.stderr}. Falling back to pip-based inference."
|
|
)
|
|
return None
|
|
except Exception as e:
|
|
_logger.warning(f"uv export failed: {e}. Falling back to pip-based inference.")
|
|
return None
|
|
|
|
|
|
def copy_uv_project_files(
|
|
dest_dir: str | Path,
|
|
source_dir: str | Path,
|
|
) -> bool:
|
|
"""
|
|
Copy uv project files to the model artifact directory.
|
|
|
|
Copies uv.lock, pyproject.toml, and .python-version (if present) to preserve
|
|
uv project configuration as model artifacts for reproducibility.
|
|
|
|
Can be disabled by setting MLFLOW_LOG_UV_FILES=false environment variable
|
|
for large projects where uv.lock size is a concern.
|
|
|
|
Args:
|
|
dest_dir: The destination directory (model artifact directory).
|
|
source_dir: The source directory containing uv project files.
|
|
|
|
Returns:
|
|
True if uv files were copied, False otherwise.
|
|
"""
|
|
if not MLFLOW_LOG_UV_FILES.get():
|
|
_logger.info("uv file logging disabled via MLFLOW_LOG_UV_FILES environment variable")
|
|
return False
|
|
|
|
dest_dir = Path(dest_dir)
|
|
source_dir = Path(source_dir)
|
|
|
|
uv_project = detect_uv_project(source_dir)
|
|
if uv_project is None:
|
|
return False
|
|
|
|
try:
|
|
shutil.copy2(uv_project.uv_lock, dest_dir / _UV_LOCK_FILE)
|
|
_logger.info(f"Copied {_UV_LOCK_FILE} to model artifacts")
|
|
|
|
shutil.copy2(uv_project.pyproject, dest_dir / _PYPROJECT_FILE)
|
|
_logger.info(f"Copied {_PYPROJECT_FILE} to model artifacts")
|
|
|
|
python_version_src = source_dir / _PYTHON_VERSION_FILE
|
|
if python_version_src.exists():
|
|
shutil.copy2(python_version_src, dest_dir / _PYTHON_VERSION_FILE)
|
|
_logger.info(f"Copied {_PYTHON_VERSION_FILE} to model artifacts")
|
|
|
|
return True
|
|
except Exception as e:
|
|
_logger.warning(f"Failed to copy uv project files: {e}")
|
|
return False
|
|
|
|
|
|
def extract_index_urls_from_uv_lock(uv_lock_path: str | Path) -> list[str]:
|
|
"""
|
|
Extract private index URLs from a uv.lock file.
|
|
|
|
Parses the uv.lock TOML file to find package sources that use non-PyPI registries.
|
|
These URLs are returned as `--extra-index-url` compatible strings.
|
|
|
|
Note: Credentials are NEVER stored in uv.lock. Users must provide credentials
|
|
at restore time via UV_INDEX_* environment variables or .netrc file.
|
|
|
|
Args:
|
|
uv_lock_path: Path to the uv.lock file.
|
|
|
|
Returns:
|
|
A list of unique index URLs (excluding PyPI). Empty list if none found
|
|
or if parsing fails.
|
|
"""
|
|
uv_lock_path = Path(uv_lock_path)
|
|
if not uv_lock_path.exists():
|
|
return []
|
|
|
|
try:
|
|
content = uv_lock_path.read_text()
|
|
|
|
# Extract registry URLs from uv.lock using regex
|
|
# Format in uv.lock: source = { registry = "https://..." }
|
|
registry_pattern = r'source\s*=\s*\{\s*registry\s*=\s*["\']([^"\']+)["\']'
|
|
urls = set(re.findall(registry_pattern, content))
|
|
|
|
# Filter out default PyPI URLs
|
|
pypi_urls = {
|
|
"https://pypi.org/simple",
|
|
"https://pypi.org/simple/",
|
|
"https://pypi.python.org/simple",
|
|
"https://pypi.python.org/simple/",
|
|
}
|
|
private_urls = [url for url in urls if url.lower() not in pypi_urls]
|
|
|
|
if private_urls:
|
|
_logger.info(f"Extracted {len(private_urls)} private index URL(s) from uv.lock")
|
|
_logger.warning(
|
|
"Private package indexes detected in uv lockfile. "
|
|
"Ensure credentials are available at model load time via "
|
|
"UV_INDEX_* environment variables or .netrc file."
|
|
)
|
|
|
|
return sorted(private_urls)
|
|
|
|
except Exception as e:
|
|
_logger.debug(f"Failed to extract index URLs from uv.lock: {e}")
|
|
return []
|
|
|
|
|
|
def create_uv_sync_pyproject(
|
|
dest_dir: str | Path,
|
|
python_version: str,
|
|
project_name: str = "mlflow-model-env",
|
|
) -> Path:
|
|
"""
|
|
Create a minimal pyproject.toml for uv sync.
|
|
|
|
This is required for `uv sync --frozen` to work when restoring a model
|
|
environment from a uv.lock artifact.
|
|
|
|
Args:
|
|
dest_dir: Directory where pyproject.toml will be created.
|
|
python_version: Python version requirement (e.g., "3.11" or "3.11.5").
|
|
project_name: Name for the temporary project. Defaults to "mlflow-model-env".
|
|
|
|
Returns:
|
|
Path to the created pyproject.toml file.
|
|
"""
|
|
dest_dir = Path(dest_dir)
|
|
|
|
# Pin the exact Python version including micro (e.g. "==3.11.5") so uv
|
|
# restores the environment with the same interpreter that was used at
|
|
# log time. "dependencies" is left empty because all deps come from
|
|
# uv.lock.
|
|
pyproject_content = f"""[project]
|
|
name = "{project_name}"
|
|
version = "0.0.0"
|
|
requires-python = "=={python_version}"
|
|
dependencies = []
|
|
|
|
[build-system]
|
|
requires = ["hatchling"]
|
|
build-backend = "hatchling.build"
|
|
"""
|
|
|
|
pyproject_path = dest_dir / _PYPROJECT_FILE
|
|
pyproject_path.write_text(pyproject_content)
|
|
_logger.debug(f"Created minimal pyproject.toml for uv sync at {pyproject_path}")
|
|
|
|
return pyproject_path
|
|
|
|
|
|
def setup_uv_sync_environment(
|
|
env_dir: str | Path,
|
|
model_path: str | Path,
|
|
python_version: str,
|
|
) -> bool:
|
|
"""
|
|
Set up a uv project structure for environment restoration via ``uv sync --frozen``.
|
|
|
|
Copies uv.lock from model artifacts and creates a minimal pyproject.toml
|
|
to enable exact environment restoration using uv.
|
|
|
|
Args:
|
|
env_dir: The environment directory where the uv project will be set up.
|
|
model_path: Path to the model artifacts containing uv.lock.
|
|
python_version: Python version for the environment.
|
|
|
|
Returns:
|
|
True if setup succeeded (uv.lock exists in model artifacts), False otherwise.
|
|
"""
|
|
env_dir = Path(env_dir)
|
|
model_path = Path(model_path)
|
|
|
|
uv_lock_artifact = model_path / _UV_LOCK_FILE
|
|
if not uv_lock_artifact.exists():
|
|
_logger.debug(f"No uv.lock found in model artifacts at {model_path}")
|
|
return False
|
|
|
|
env_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Copy uv.lock to environment directory
|
|
shutil.copy2(uv_lock_artifact, env_dir / _UV_LOCK_FILE)
|
|
_logger.debug(f"Copied uv.lock to {env_dir}")
|
|
|
|
# Copy pyproject.toml from model if available, otherwise create minimal one
|
|
pyproject_artifact = model_path / _PYPROJECT_FILE
|
|
if pyproject_artifact.exists():
|
|
shutil.copy2(pyproject_artifact, env_dir / _PYPROJECT_FILE)
|
|
_logger.debug(f"Copied pyproject.toml from model artifacts to {env_dir}")
|
|
else:
|
|
create_uv_sync_pyproject(env_dir, python_version)
|
|
|
|
# Copy .python-version if it exists in model artifacts
|
|
python_version_artifact = model_path / _PYTHON_VERSION_FILE
|
|
if python_version_artifact.exists():
|
|
shutil.copy2(python_version_artifact, env_dir / _PYTHON_VERSION_FILE)
|
|
_logger.debug(f"Copied .python-version to {env_dir}")
|
|
|
|
_logger.info(f"Set up uv sync environment at {env_dir}")
|
|
return True
|
|
|
|
|
|
def run_uv_sync(
|
|
project_dir: str | Path,
|
|
frozen: bool = True,
|
|
no_dev: bool = True,
|
|
capture_output: bool = False,
|
|
) -> bool:
|
|
"""
|
|
Run `uv sync` to install dependencies from a uv.lock file.
|
|
|
|
This provides exact cross-platform environment restoration, significantly
|
|
faster than pip-based installation.
|
|
|
|
Args:
|
|
project_dir: Directory containing pyproject.toml and uv.lock.
|
|
frozen: Use frozen lockfile without updating. Defaults to True.
|
|
no_dev: Exclude development dependencies. Defaults to True.
|
|
capture_output: Whether to capture stdout/stderr. Defaults to False.
|
|
|
|
Returns:
|
|
True if sync succeeded, False otherwise.
|
|
"""
|
|
uv_bin = _get_uv_binary()
|
|
if uv_bin is None:
|
|
_logger.warning("uv is not available for environment sync")
|
|
return False
|
|
project_dir = Path(project_dir)
|
|
|
|
cmd = [uv_bin, "sync"]
|
|
if frozen:
|
|
cmd.append("--frozen")
|
|
if no_dev:
|
|
cmd.append("--no-dev")
|
|
|
|
try:
|
|
_logger.info(f"Running uv sync in {project_dir}")
|
|
_logger.debug(f"uv sync command: {' '.join(str(c) for c in cmd)}")
|
|
|
|
subprocess.run(
|
|
cmd,
|
|
cwd=project_dir,
|
|
capture_output=capture_output,
|
|
check=True,
|
|
text=True,
|
|
# Set UV_PROJECT_ENVIRONMENT to install packages into python environment at
|
|
# `project_dir` path instead of the default `{cwd}/.venv` path.
|
|
env={**os.environ, "UV_PROJECT_ENVIRONMENT": str(project_dir)},
|
|
)
|
|
|
|
_logger.info("uv sync completed successfully")
|
|
return True
|
|
|
|
except subprocess.CalledProcessError as e:
|
|
_logger.warning(f"uv sync failed with exit code {e.returncode}: {e.stderr}")
|
|
return False
|
|
except Exception as e:
|
|
_logger.warning(f"uv sync failed: {e}")
|
|
return False
|
|
|
|
|
|
def has_uv_lock_artifact(model_path: str | Path) -> bool:
|
|
"""Check if a model has a uv.lock artifact."""
|
|
return (Path(model_path) / _UV_LOCK_FILE).exists()
|