chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,519 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user