wehub-resource-sync
520 lines
20 KiB
Python
520 lines
20 KiB
Python
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()
|