mlflow--mlflow/examples/uv-dependency-management/log_model.py

"""
Example: Using uv for dependency management with MLflow models.

This script demonstrates three ways to use uv lockfile-based dependencies
when logging MLflow models:

1. Auto-detection: MLflow detects uv.lock + pyproject.toml in the current
   working directory and uses ``uv export`` to capture pinned dependencies.

2. Explicit path (uv_project_path): Point to a uv project directory when
   logging from a different working directory or in a monorepo layout.

3. Dependency groups and extras (uv_groups, uv_extras): Include additional
   dependency groups or optional extras defined in pyproject.toml.

Prerequisites:
    - uv >= 0.5.0 installed (``pip install uv`` or https://docs.astral.sh/uv/)
    - Run from this directory so auto-detection finds uv.lock and pyproject.toml

Usage:
    cd examples/uv-dependency-management
    uv run python log_model.py
"""

from pathlib import Path

from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split

import mlflow


def read_requirements(run_id, artifact_name="model"):
    """Read the model's requirements.txt from local run artifacts."""
    client = mlflow.tracking.MlflowClient()
    local_path = client.download_artifacts(run_id, f"{artifact_name}/requirements.txt")
    with open(local_path) as f:
        return [line.strip() for line in f if line.strip()]


def check_uv_artifacts(run_id, artifact_name="model"):
    """Check if uv project files were saved as model artifacts."""
    client = mlflow.tracking.MlflowClient()
    model_dir = client.download_artifacts(run_id, artifact_name)
    model_path = Path(model_dir)
    return {
        "uv.lock": (model_path / "uv.lock").exists(),
        "pyproject.toml": (model_path / "pyproject.toml").exists(),
    }


def train_model():
    """Train a simple RandomForest on the Iris dataset."""
    iris = load_iris()
    X_train, X_test, y_train, y_test = train_test_split(
        iris.data, iris.target, test_size=0.2, random_state=42
    )
    model = RandomForestClassifier(n_estimators=10, random_state=42)
    model.fit(X_train, y_train)
    accuracy = model.score(X_test, y_test)
    return model, X_test, accuracy


class SklearnWrapper(mlflow.pyfunc.PythonModel):
    """Wrap a scikit-learn model as a PythonModel for pyfunc logging."""

    def __init__(self, sklearn_model):
        self._model = sklearn_model

    def predict(self, context, model_input, params=None):
        return self._model.predict(model_input)


def example_auto_detection(model, input_example):
    """
    Example 1: Auto-detection.

    When run from a directory containing uv.lock and pyproject.toml,
    MLflow automatically uses uv export to capture pinned dependencies.
    No extra parameters needed.
    """
    print("=" * 60)
    print("Example 1: Auto-detection")
    print("=" * 60)

    with mlflow.start_run(run_name="uv-auto-detection") as run:
        model_info = mlflow.pyfunc.log_model(
            python_model=SklearnWrapper(model),
            name="model",
            input_example=input_example,
        )

        run_id = run.info.run_id
        reqs = read_requirements(run_id)
        print(f"Logged model: {model_info.model_uri}")
        print(f"Requirements ({len(reqs)} packages):")
        for req in reqs[:10]:
            print(f"  {req}")
        if len(reqs) > 10:
            print(f"  ... and {len(reqs) - 10} more")

        # Verify uv artifacts were saved
        uv_files = check_uv_artifacts(run_id)
        print(f"uv.lock saved as artifact: {uv_files['uv.lock']}")
        print(f"pyproject.toml saved as artifact: {uv_files['pyproject.toml']}")
        print()

    return model_info


def example_explicit_path(model, input_example):
    """
    Example 2: Explicit uv_project_path.

    Use uv_project_path to point to a uv project when logging from
    a different working directory. Useful in monorepos.
    """
    print("=" * 60)
    print("Example 2: Explicit uv_project_path")
    print("=" * 60)

    project_dir = Path(__file__).parent.resolve()

    with mlflow.start_run(run_name="uv-explicit-path") as run:
        model_info = mlflow.pyfunc.log_model(
            python_model=SklearnWrapper(model),
            name="model",
            input_example=input_example,
            uv_project_path=project_dir,
        )

        run_id = run.info.run_id
        reqs = read_requirements(run_id)
        print(f"Logged model: {model_info.model_uri}")
        print(f"uv_project_path: {project_dir}")
        print(f"Requirements ({len(reqs)} packages):")
        for req in reqs[:10]:
            print(f"  {req}")
        if len(reqs) > 10:
            print(f"  ... and {len(reqs) - 10} more")
        print()

    return model_info


def example_groups_and_extras(model, input_example):
    """
    Example 3: Dependency groups and extras.

    Include the 'ml' dependency group (xgboost) and 'serving' optional
    extra (flask) in the exported requirements.
    """
    print("=" * 60)
    print("Example 3: uv_groups and uv_extras")
    print("=" * 60)

    project_dir = Path(__file__).parent.resolve()

    with mlflow.start_run(run_name="uv-groups-and-extras") as run:
        model_info = mlflow.pyfunc.log_model(
            python_model=SklearnWrapper(model),
            name="model",
            input_example=input_example,
            uv_project_path=project_dir,
            uv_groups=["ml"],
            uv_extras=["serving"],
        )

        run_id = run.info.run_id
        reqs = read_requirements(run_id)
        print(f"Logged model: {model_info.model_uri}")
        print("uv_groups: ['ml']  (adds xgboost)")
        print("uv_extras: ['serving']  (adds flask)")
        print(f"Requirements ({len(reqs)} packages):")

        # Check that group and extra deps were included
        has_xgboost = any("xgboost" in r for r in reqs)
        has_flask = any("flask" in r.lower() for r in reqs)
        print(f"  xgboost included (from 'ml' group): {has_xgboost}")
        print(f"  flask included (from 'serving' extra): {has_flask}")
        print()

        for req in sorted(reqs):
            print(f"  {req}")
        print()

    return model_info


def main():
    print("MLflow uv Dependency Management Example")
    print()

    # Train a model
    model, X_test, accuracy = train_model()
    input_example = X_test[:2]
    print(f"Trained RandomForestClassifier (accuracy: {accuracy:.2f})")
    print()

    # Set up a local tracking URI for the example.
    # Use an absolute path so it works regardless of working directory.
    example_dir = Path(__file__).parent.resolve()
    db_path = example_dir / "mlflow.db"
    mlflow.set_tracking_uri(f"sqlite:///{db_path}")
    mlflow.set_experiment("uv-dependency-management")

    # Run all three examples
    example_auto_detection(model, input_example)
    example_explicit_path(model, input_example)
    example_groups_and_extras(model, input_example)

    print("All examples completed successfully.")


if __name__ == "__main__":
    main()