pyproject.py

from __future__ import annotations

import re
import subprocess
import sys
from collections import Counter
from enum import Enum
from pathlib import Path
from typing import Any, cast

import toml
import yaml
from packaging.version import Version
from pydantic import BaseModel, Field, RootModel


class PackageType(Enum):
    SKINNY = "skinny"
    RELEASE = "release"
    DEV = "dev"
    TRACING = "tracing"

    def description(self) -> str:
        WARNING = "# Auto-generated by dev/pyproject.py. Do not edit manually."

        if self is PackageType.TRACING:
            return f"""{WARNING}
# This file defines the package metadata of `mlflow-tracing`.
"""

        if self is PackageType.SKINNY:
            return f"""{WARNING}
# This file defines the package metadata of `mlflow-skinny`.
"""
        if self is PackageType.RELEASE:
            return f"""{WARNING}
# This file defines the package metadata of `mlflow`. `mlflow-skinny` and `mlflow-tracing`
# are included in the requirements to prevent a version mismatch between `mlflow` and those
# child packages. This file will replace `pyproject.toml` when releasing a new version.
"""
        if self is PackageType.DEV:
            return f"""{WARNING}
# This file defines the package metadata of `mlflow` **during development**. To install `mlflow`
# from the source code, `mlflow-skinny` and `mlflow-tracing` are NOT included in the requirements.
# This file will be replaced by `pyproject.release.toml` when releasing a new version.
"""
        raise ValueError(f"Unreachable: {self}")


SEPARATOR = """
# Package metadata: can't be updated manually, use dev/pyproject.py
# -----------------------------------------------------------------
# Dev tool settings: can be updated manually

"""

SKINNY_README = """
<!--  Autogenerated by dev/pyproject.py. Do not edit manually.  -->

📣 This is the `mlflow-skinny` package, a lightweight MLflow package without SQL storage, server, UI, or data science dependencies.
Additional dependencies can be installed to leverage the full feature set of MLflow. For example:

- To use the `mlflow.sklearn` component of MLflow Models, install `scikit-learn`, `numpy` and `pandas`.
- To use SQL-based metadata storage, install `sqlalchemy`, `alembic`, and `sqlparse`.
- To use serving-based features, install `flask` and `pandas`.

**Note:** When using `mlflow-skinny`, set the tracking URI to your remote MLflow server:

```bash
export MLFLOW_TRACKING_URI="http://your-mlflow-server:5000"
```

---

<br>
<br>

"""  # noqa: E501

# Tracing SDK should only include the minimum set of MLflow modules
# to minimize the size of the package.
TRACING_INCLUDE_FILES = [
    "mlflow",
    # Flavors that we support auto tracing
    "mlflow.agno*",
    "mlflow.anthropic*",
    "mlflow.autogen*",
    "mlflow.bedrock*",
    "mlflow.crewai*",
    "mlflow.dspy*",
    "mlflow.gemini*",
    "mlflow.groq*",
    "mlflow.langchain*",
    "mlflow.litellm*",
    "mlflow.llama_index*",
    "mlflow.mistral*",
    "mlflow.openai*",
    "mlflow.strands*",
    "mlflow.haystack*",
    # Other necessary modules
    "mlflow.azure*",
    "mlflow.entities*",
    "mlflow.environment_variables",
    "mlflow.exceptions",
    "mlflow.legacy_databricks_cli*",
    "mlflow.prompt*",
    "mlflow.protos*",
    "mlflow.pydantic_ai*",
    "mlflow.smolagents*",
    "mlflow.store*",
    "mlflow.telemetry*",
    "mlflow.tracing*",
    "mlflow.tracking*",
    "mlflow.types*",
    "mlflow.utils*",
    "mlflow.version",
]
TRACING_EXCLUDE_FILES = [
    # Large proto files that are not needed in the package
    "mlflow/protos/databricks_artifacts_pb2.py",
    "mlflow/protos/databricks_filesystem_service_pb2.py",
    "mlflow/protos/databricks_uc_registry_messages_pb2.py",
    "mlflow/protos/databricks_uc_registry_service_pb2.py",
    "mlflow/protos/model_registry_pb2.py",
    "mlflow/protos/unity_catalog_oss_messages_pb2.py",
    "mlflow/protos/unity_catalog_oss_service_pb2.py",
    # Test files
    "tests",
    "tests.*",
]


def find_duplicates(seq: list[str]) -> list[str]:
    counted = Counter(seq)
    return [item for item, count in counted.items() if count > 1]


def write_file_if_changed(file_path: Path, new_content: str) -> None:
    if file_path.exists():
        existing_content = file_path.read_text()
        if existing_content == new_content:
            print(f"No changes in {file_path}, skipping write.")
            return

    print(f"Writing changes to {file_path}.")
    file_path.write_text(new_content)


def format_content_with_taplo(content: str) -> str:
    return (
        subprocess.check_output(
            ["bin/taplo", "fmt", "-"],
            input=content,
            text=True,
        ).strip()
        + "\n"
    )


def write_toml_file_if_changed(
    file_path: Path, description: str, toml_data: dict[str, Any]
) -> None:
    """
    Write a TOML file with description only if content has changed.
    Formats content with taplo before comparison.
    """
    new_content = description + "\n" + toml.dumps(toml_data)
    formatted_content = format_content_with_taplo(new_content)
    write_file_if_changed(file_path, formatted_content)


class PackageRequirement(BaseModel):
    pip_release: str = Field(..., description="The pip package name")
    max_major_version: int = Field(..., description="Maximum major version allowed")
    minimum: str | None = Field(None, description="Minimum version required")
    unsupported: list[str] | None = Field(None, description="List of unsupported versions")
    markers: str | None = Field(
        None, description="Environment markers for conditional installation"
    )
    extras: list[str] | None = Field(None, description="Package extras to install")
    freeze: bool | None = Field(None, description="Whether to freeze this package version")


RequirementsYaml = RootModel[dict[str, PackageRequirement]]


def generate_requirements_from_yaml(requirements_yaml: RequirementsYaml) -> list[str]:
    """Generate pip requirement strings from validated YAML specification."""
    requirement_strs: list[str] = []
    for package_entry in requirements_yaml.root.values():
        pip_release = package_entry.pip_release
        version_specs: list[str] = []

        extras = f"[{','.join(package_entry.extras)}]" if package_entry.extras else ""

        max_major_version = package_entry.max_major_version
        version_specs.append(f"<{max_major_version + 1}")

        if package_entry.minimum:
            version_specs.append(f">={package_entry.minimum}")

        if package_entry.unsupported:
            version_specs.extend(f"!={version}" for version in package_entry.unsupported)

        markers = f"; {package_entry.markers}" if package_entry.markers else ""

        requirement_str = f"{pip_release}{extras}{','.join(version_specs)}{markers}"
        requirement_strs.append(requirement_str)

    requirement_strs.sort()
    return requirement_strs


def read_requirements_yaml(yaml_path: Path) -> list[str]:
    """Read and parse a YAML requirements file into pip requirement strings."""
    with yaml_path.open() as f:
        requirements_data = yaml.safe_load(f)

    return generate_requirements_from_yaml(RequirementsYaml(requirements_data))


def read_package_versions_yml() -> dict[str, Any]:
    with open("mlflow/ml-package-versions.yml") as f:
        return cast(dict[str, Any], yaml.safe_load(f))


def build(package_type: PackageType) -> None:
    requirements_dir = Path("requirements")
    tracing_requirements = read_requirements_yaml(requirements_dir / "tracing-requirements.yaml")
    skinny_requirements = read_requirements_yaml(requirements_dir / "skinny-requirements.yaml")
    _check_skinny_tracing_mismatch(
        skinny_reqs=skinny_requirements, tracing_reqs=tracing_requirements
    )
    core_requirements = read_requirements_yaml(requirements_dir / "core-requirements.yaml")
    gateways_requirements = read_requirements_yaml(requirements_dir / "gateway-requirements.yaml")
    genai_requirements = read_requirements_yaml(requirements_dir / "genai-requirements.yaml")
    version_match = re.search(
        r'^VERSION = "([a-z0-9\.]+)"$', Path("mlflow", "version.py").read_text(), re.MULTILINE
    )
    if version_match is None:
        raise ValueError(
            'Could not find VERSION in mlflow/version.py. Expected format: VERSION = "x.y.z"'
        )
    package_version = version_match.group(1)
    python_version = Path(".python-version").read_text().strip()
    versions_yaml = read_package_versions_yml()
    langchain_requirements = [
        "langchain>={},<={}".format(
            max(
                Version(versions_yaml["langchain"]["autologging"]["minimum"]),
                Version(versions_yaml["langchain"]["models"]["minimum"]),
            ),
            min(
                Version(versions_yaml["langchain"]["autologging"]["maximum"]),
                Version(versions_yaml["langchain"]["models"]["maximum"]),
            ),
        )
    ]

    match package_type:
        case PackageType.TRACING:
            dependencies = sorted(tracing_requirements)
        case PackageType.SKINNY:
            dependencies = sorted(skinny_requirements)
        case PackageType.RELEASE:
            dependencies = [
                f"mlflow-skinny=={package_version}",
                f"mlflow-tracing=={package_version}",
            ] + sorted(core_requirements)
        case PackageType.DEV:
            # skinny_requirements is an exact superset of tracing_requirements
            # (validated above), so we don't need to include both below.
            dependencies = sorted(core_requirements + skinny_requirements)
        case _:
            raise ValueError(f"Unreachable: {package_type}")

    if dep_duplicates := find_duplicates(dependencies):
        raise RuntimeError(f"Duplicated dependencies are found: {dep_duplicates}")

    match package_type:
        case PackageType.TRACING:
            package_name = "mlflow-tracing"
        case PackageType.SKINNY:
            package_name = "mlflow-skinny"
        case _:
            package_name = "mlflow"

    description = (
        "MLflow is an open source platform for the complete machine learning lifecycle"
        if package_type != PackageType.TRACING
        else (
            "MLflow Tracing SDK is an open-source, lightweight Python package that only "
            "includes the minimum set of dependencies and functionality to instrument "
            "your code/models/agents with MLflow Tracing."
        )
    )

    data = {
        "build-system": {
            "requires": ["setuptools<=82.0.1"],
            "build-backend": "setuptools.build_meta",
        },
        "project": {
            "name": package_name,
            "version": package_version,
            "maintainers": [
                {"name": "Databricks", "email": "mlflow-oss-maintainers@googlegroups.com"}
            ],
            "description": description,
            "readme": "README_SKINNY.md" if package_type == PackageType.SKINNY else "README.md",
            "license": {
                "file": "LICENSE.txt",
            },
            "keywords": ["mlflow", "ai", "databricks"],
            "classifiers": [
                "Development Status :: 5 - Production/Stable",
                "Intended Audience :: Developers",
                "Intended Audience :: End Users/Desktop",
                "Intended Audience :: Science/Research",
                "Intended Audience :: Information Technology",
                "Topic :: Scientific/Engineering :: Artificial Intelligence",
                "Topic :: Software Development :: Libraries :: Python Modules",
                "License :: OSI Approved :: Apache Software License",
                "Operating System :: OS Independent",
                f"Programming Language :: Python :: {python_version}",
            ],
            "requires-python": f">={python_version}",
            "dependencies": dependencies,
            "optional-dependencies": {
                "extras": [
                    # Required to log artifacts and models to HDFS artifact locations
                    "pyarrow",
                    # Required to sign outgoing request with SigV4 signature
                    "requests-auth-aws-sigv4",
                    # Required to log artifacts and models to AWS S3 artifact locations
                    "boto3",
                    "botocore",
                    # Required to log artifacts and models to GCS artifact locations
                    "google-cloud-storage>=1.30.0",
                    "azureml-core>=1.2.0",
                    # Required to log artifacts to SFTP artifact locations
                    "pysftp",
                    # Required by the mlflow.projects module, when running projects against
                    # a remote Kubernetes cluster
                    "kubernetes",
                    # Required for exporting metrics from the MLflow server to Prometheus
                    # as part of the MLflow server monitoring add-on
                    "prometheus-flask-exporter",
                ],
                "db": [
                    # Required to use MySQL, PostgreSQL, or SQL Server as the backend store
                    "PyMySQL",
                    "psycopg2-binary",
                    "pymssql",
                ],
                "databricks": [
                    # Required to write model artifacts to unity catalog locations
                    "azure-storage-file-datalake>12",
                    "google-cloud-storage>=1.30.0",
                    "boto3>1",
                    "botocore",
                    "databricks-agents>=1.2.0,<2.0",
                ],
                "gateway": gateways_requirements,
                "genai": genai_requirements,
                # click 8.3.0 causes MLflow MCP server to fail: https://github.com/mlflow/mlflow/issues/18747
                "mcp": ["fastmcp<4,>=2.7.0", "click!=8.3.0"],
                "azure": [
                    # Required to log artifacts and models to Azure Blob Storage
                    "azure-storage-blob>=12",
                    "azure-identity>=1.6.1",
                ],
                "sqlserver": ["mlflow-dbstore"],
                "aliyun-oss": ["aliyunstoreplugin"],
                "jfrog": ["mlflow-jfrog-plugin"],
                "kubernetes": ["kubernetes"],
                "langchain": langchain_requirements,
                "auth": ["Flask-WTF<2"],
            }
            # Tracing SDK does not support extras
            if package_type != PackageType.TRACING
            else None,
            "urls": {
                "homepage": "https://mlflow.org",
                "issues": "https://github.com/mlflow/mlflow/issues",
                "documentation": "https://mlflow.org/docs/latest",
                "repository": "https://github.com/mlflow/mlflow",
            },
            "scripts": {
                "mlflow": "mlflow.cli:cli",
            }
            if package_type != PackageType.TRACING
            else None,
            "entry-points": {
                "mlflow.app": {
                    "basic-auth": "mlflow.server.auth:create_app",
                },
                "mlflow.app.client": {
                    "basic-auth": "mlflow.server.auth.client:AuthServiceClient",
                },
                "mlflow.deployments": {
                    "databricks": "mlflow.deployments.databricks",
                    "http": "mlflow.deployments.mlflow",
                    "https": "mlflow.deployments.mlflow",
                    "openai": "mlflow.deployments.openai",
                },
            }
            if package_type != PackageType.TRACING
            else None,
        },
        "tool": {
            "setuptools": {
                "packages": {
                    "find": {
                        "where": ["."],
                        "include": ["mlflow", "mlflow.*"]
                        if package_type != PackageType.TRACING
                        else TRACING_INCLUDE_FILES,
                        "exclude": ["tests", "tests.*"]
                        if package_type != PackageType.TRACING
                        else TRACING_EXCLUDE_FILES,
                        "namespaces": False,
                    }
                },
                "package-data": _get_package_data(package_type),
            }
        },
    }

    if package_type == PackageType.TRACING:
        out_path = Path("libs/tracing/pyproject.toml")
        write_toml_file_if_changed(out_path, package_type.description(), data)
    elif package_type == PackageType.SKINNY:
        out_path = Path("libs/skinny/pyproject.toml")
        write_toml_file_if_changed(out_path, package_type.description(), data)

        skinny_readme_path = Path("libs/skinny/README_SKINNY.md")
        new_readme_content = SKINNY_README.lstrip() + Path("README.md").read_text()
        write_file_if_changed(skinny_readme_path, new_readme_content)

        for f in ["LICENSE.txt", "MANIFEST.in", "mlflow"]:
            symlink = Path("libs/skinny", f)
            if symlink.exists():
                symlink.unlink()
            target = Path("../..", f)
            symlink.symlink_to(target, target_is_directory=target.is_dir())
    elif package_type == PackageType.RELEASE:
        out_path = Path(f"pyproject.{package_type.value}.toml")
        write_toml_file_if_changed(out_path, package_type.description(), data)
    else:
        out_path = Path("pyproject.toml")
        original_manual_content = out_path.read_text().split(SEPARATOR)[1]
        generated_part = package_type.description() + "\n" + toml.dumps(data)
        formatted_generated_part = format_content_with_taplo(generated_part)
        formatted_full_content = formatted_generated_part + SEPARATOR + original_manual_content

        write_file_if_changed(out_path, formatted_full_content)
        subprocess.check_call(["uv", "lock"])


def _get_package_data(package_type: PackageType) -> dict[str, list[str]] | None:
    if package_type == PackageType.TRACING:
        return None

    package_data = {
        "mlflow": [
            "store/db_migrations/alembic.ini",
            "temporary_db_migrations_for_pre_1_users/alembic.ini",
            "pyspark/ml/log_model_allowlist.txt",
            "server/auth/basic_auth.ini",
            "server/auth/db/migrations/alembic.ini",
            "server/uvicorn_log_config.yaml",
            "models/notebook_resources/**/*",
            "ai_commands/**/*.md",
            "assistant/skills/**/*",
            "agent/setup/templates/**/*.md",
        ]
    }

    if package_type != PackageType.SKINNY:
        package_data["mlflow"] += [
            "models/container/**/*",
            "server/js/build/**/*",
            "utils/model_catalog/*.json",
        ]

    return package_data


def _check_skinny_tracing_mismatch(*, skinny_reqs: list[str], tracing_reqs: list[str]) -> None:
    """
    Check if the tracing requirements are a subset of the skinny requirements.
    NB: We don't make mlflow-tracing as a hard dependency of mlflow-skinny because
    it will complicate the package management (need another .release.toml file
    that is dependent by pyproject.release.toml)
    """
    if diff := set(tracing_reqs) - set(skinny_reqs):
        raise RuntimeError(
            "Tracing requirements must be a subset of skinny requirements. "
            "Please check the requirements/skinny-requirements.yaml and "
            "requirements/tracing-requirements.yaml files.\n"
            f"Diff: {diff}"
        )


def main() -> None:
    if not Path("bin/taplo").exists():
        print(
            "taplo is required to generate pyproject.toml. "
            "Please run 'python bin/install.py' to install it."
        )
        sys.exit(1)

    for package_type in PackageType:
        build(package_type)


if __name__ == "__main__":
    main()