1  import argparse
  2  import contextlib
  3  import shutil
  4  import subprocess
  5  import sys
  6  import zipfile
  7  from collections.abc import Generator
  8  from dataclasses import dataclass
  9  from pathlib import Path
 10  
 11  
 12  @dataclass(frozen=True)
 13  class Package:
 14      # name of the package on PyPI.
 15      pypi_name: str
 16      # type of the package, one of "dev", "skinny", "tracing", "release"
 17      type: str
 18      # path to the package relative to the root of the repository
 19      build_path: str
 20  
 21  
 22  DEV = Package("mlflow", "dev", ".")
 23  RELEASE = Package("mlflow", "release", ".")
 24  SKINNY = Package("mlflow-skinny", "skinny", "libs/skinny")
 25  TRACING = Package("mlflow-tracing", "tracing", "libs/tracing")
 26  
 27  PACKAGES = [
 28      DEV,
 29      SKINNY,
 30      RELEASE,
 31      TRACING,
 32  ]
 33  
 34  JS_BUILD_DIR = Path("mlflow/server/js/build")
 35  
 36  
 37  def parse_args() -> argparse.Namespace:
 38      parser = argparse.ArgumentParser(description="Build MLflow package.")
 39      parser.add_argument(
 40          "--package-type",
 41          help="Package type to build. Default is 'dev'.",
 42          choices=[p.type for p in PACKAGES],
 43          default="dev",
 44      )
 45      parser.add_argument(
 46          "--sha",
 47          help="If specified, include the SHA in the wheel name as a build tag.",
 48      )
 49      return parser.parse_args()
 50  
 51  
 52  @contextlib.contextmanager
 53  def restore_changes() -> Generator[None, None, None]:
 54      try:
 55          yield
 56      finally:
 57          subprocess.check_call([
 58              "git",
 59              "restore",
 60              "README.md",
 61              "pyproject.toml",
 62          ])
 63  
 64  
 65  def validate_ui_assets_pre_build(package: Package) -> None:
 66      if package != RELEASE:
 67          return
 68      if not JS_BUILD_DIR.exists() or not any(JS_BUILD_DIR.iterdir()):
 69          raise RuntimeError("Build the UI first before building the release package.")
 70  
 71  
 72  def validate_ui_assets_post_build(wheel_path: Path, package: Package) -> None:
 73      ui_asset_prefix = f"{JS_BUILD_DIR.as_posix()}/"
 74      with zipfile.ZipFile(wheel_path) as zf:
 75          has_ui_assets = any(name.startswith(ui_asset_prefix) for name in zf.namelist())
 76      if package == RELEASE:
 77          if not has_ui_assets:
 78              raise RuntimeError(
 79                  f"UI assets are missing from the release wheel: {wheel_path}. "
 80                  "Build the UI first before building the release package."
 81              )
 82      elif package in (SKINNY, TRACING):
 83          if has_ui_assets:
 84              raise RuntimeError(
 85                  f"UI assets should not be included in the {package.type} wheel: {wheel_path}."
 86              )
 87  
 88  
 89  def main() -> None:
 90      args = parse_args()
 91  
 92      # Initialize submodules (e.g., mlflow/assistant/skills)
 93      subprocess.check_call(["git", "submodule", "update", "--init", "--recursive"])
 94  
 95      # Clean up build artifacts generated by previous builds
 96      paths_to_clean_up = ["build"]
 97      for pkg in PACKAGES:
 98          paths_to_clean_up += [
 99              f"{pkg.build_path}/dist",
100              f"{pkg.build_path}/{pkg.pypi_name}.egg_info",
101          ]
102      for path in map(Path, paths_to_clean_up):
103          if not path.exists():
104              continue
105          if path.is_file():
106              path.unlink()
107          else:
108              shutil.rmtree(path)
109  
110      package = next(p for p in PACKAGES if p.type == args.package_type)
111  
112      validate_ui_assets_pre_build(package)
113  
114      with restore_changes():
115          pyproject = Path("pyproject.toml")
116          if package == RELEASE:
117              pyproject.write_text(Path("pyproject.release.toml").read_text())
118  
119          DIST_DIR = Path("dist").resolve()
120          DIST_DIR.mkdir(exist_ok=True)
121          subprocess.check_call([
122              sys.executable,
123              "-m",
124              "build",
125              package.build_path,
126              "--outdir",
127              DIST_DIR,
128          ])
129  
130      wheel = next(DIST_DIR.glob("mlflow*.whl"))
131      validate_ui_assets_post_build(wheel, package)
132  
133      if args.sha:
134          name, version, rest = wheel.name.split("-", 2)
135          build_tag = f"0.sha.{args.sha}"  # build tag must start with a digit
136          wheel.rename(wheel.with_name(f"{name}-{version}-{build_tag}-{rest}"))
137  
138  
139  if __name__ == "__main__":
140      main()