wehub-resource-sync
141 lines
3.9 KiB
Python
141 lines
3.9 KiB
Python
import argparse
|
|
import contextlib
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import zipfile
|
|
from collections.abc import Generator
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Package:
|
|
# name of the package on PyPI.
|
|
pypi_name: str
|
|
# type of the package, one of "dev", "skinny", "tracing", "release"
|
|
type: str
|
|
# path to the package relative to the root of the repository
|
|
build_path: str
|
|
|
|
|
|
DEV = Package("mlflow", "dev", ".")
|
|
RELEASE = Package("mlflow", "release", ".")
|
|
SKINNY = Package("mlflow-skinny", "skinny", "libs/skinny")
|
|
TRACING = Package("mlflow-tracing", "tracing", "libs/tracing")
|
|
|
|
PACKAGES = [
|
|
DEV,
|
|
SKINNY,
|
|
RELEASE,
|
|
TRACING,
|
|
]
|
|
|
|
JS_BUILD_DIR = Path("mlflow/server/js/build")
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Build MLflow package.")
|
|
parser.add_argument(
|
|
"--package-type",
|
|
help="Package type to build. Default is 'dev'.",
|
|
choices=[p.type for p in PACKAGES],
|
|
default="dev",
|
|
)
|
|
parser.add_argument(
|
|
"--sha",
|
|
help="If specified, include the SHA in the wheel name as a build tag.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def restore_changes() -> Generator[None, None, None]:
|
|
try:
|
|
yield
|
|
finally:
|
|
subprocess.check_call([
|
|
"git",
|
|
"restore",
|
|
"README.md",
|
|
"pyproject.toml",
|
|
])
|
|
|
|
|
|
def validate_ui_assets_pre_build(package: Package) -> None:
|
|
if package != RELEASE:
|
|
return
|
|
if not JS_BUILD_DIR.exists() or not any(JS_BUILD_DIR.iterdir()):
|
|
raise RuntimeError("Build the UI first before building the release package.")
|
|
|
|
|
|
def validate_ui_assets_post_build(wheel_path: Path, package: Package) -> None:
|
|
ui_asset_prefix = f"{JS_BUILD_DIR.as_posix()}/"
|
|
with zipfile.ZipFile(wheel_path) as zf:
|
|
has_ui_assets = any(name.startswith(ui_asset_prefix) for name in zf.namelist())
|
|
if package == RELEASE:
|
|
if not has_ui_assets:
|
|
raise RuntimeError(
|
|
f"UI assets are missing from the release wheel: {wheel_path}. "
|
|
"Build the UI first before building the release package."
|
|
)
|
|
elif package in (SKINNY, TRACING):
|
|
if has_ui_assets:
|
|
raise RuntimeError(
|
|
f"UI assets should not be included in the {package.type} wheel: {wheel_path}."
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
|
|
# Initialize submodules (e.g., mlflow/assistant/skills)
|
|
subprocess.check_call(["git", "submodule", "update", "--init", "--recursive"])
|
|
|
|
# Clean up build artifacts generated by previous builds
|
|
paths_to_clean_up = ["build"]
|
|
for pkg in PACKAGES:
|
|
paths_to_clean_up += [
|
|
f"{pkg.build_path}/dist",
|
|
f"{pkg.build_path}/{pkg.pypi_name}.egg_info",
|
|
]
|
|
for path in map(Path, paths_to_clean_up):
|
|
if not path.exists():
|
|
continue
|
|
if path.is_file():
|
|
path.unlink()
|
|
else:
|
|
shutil.rmtree(path)
|
|
|
|
package = next(p for p in PACKAGES if p.type == args.package_type)
|
|
|
|
validate_ui_assets_pre_build(package)
|
|
|
|
with restore_changes():
|
|
pyproject = Path("pyproject.toml")
|
|
if package == RELEASE:
|
|
pyproject.write_text(Path("pyproject.release.toml").read_text())
|
|
|
|
DIST_DIR = Path("dist").resolve()
|
|
DIST_DIR.mkdir(exist_ok=True)
|
|
subprocess.check_call([
|
|
sys.executable,
|
|
"-m",
|
|
"build",
|
|
package.build_path,
|
|
"--outdir",
|
|
DIST_DIR,
|
|
])
|
|
|
|
wheel = next(DIST_DIR.glob("mlflow*.whl"))
|
|
validate_ui_assets_post_build(wheel, package)
|
|
|
|
if args.sha:
|
|
name, version, rest = wheel.name.split("-", 2)
|
|
build_tag = f"0.sha.{args.sha}" # build tag must start with a digit
|
|
wheel.rename(wheel.with_name(f"{name}-{version}-{build_tag}-{rest}"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|