Add oneshot release snapshot for test/ondemand (#768)

* Add oneshot release snapshot for test/ondemand

Add some build scripts to test new release flow based on IREE.
Wont affect current builds, once this works well we can plumb it
in.

Build with manylinux docker

* Fixes a few issues found when debugging powderluv's setup.

* It is optional to link against Python3_LIBRARIES. Check that and don't do it if they don't exist for this config.
* Clean and auditwheel need to operate on sanitized package names. So "torch_mlir" vs "torch-mlir".
* Adds a pyproject.toml file that pins the build dependencies needed to detect both Torch and Python (the MLIR Python build was failing to detect because Numpy wasn't in the pip venv).
* Commented out auditwheel: These wheels are not PyPi compliant since they weak link to libtorch at runtime. However, they should be fine to deploy to users.
* Adds the --extra-index-url to the pip wheel command, allowing PyTorch to be found.
* Hack setup.py to remove the _mlir_libs dir before building. This keeps back-to-back versions from accumulating in the wheels for subsequent versions. IREE has a more principled way of doing this, but what I have here should work.

Co-authored-by: Stella Laurenzo <stellaraccident@gmail.com>

.github/workflows/buildManylinux.yml

@@ -0,0 +1,55 @@
+name: Manylinux Build
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  workflow_dispatch:
+    inputs:
+      release_id:
+        description: 'Release id to upload artifacts to'
+        default: ''
+      python_package_version:
+        description: 'Version to use for creating the Python package'
+        default: ''
+
+jobs:
+  build:
+    name: Manylinux Build
+    runs-on: ubuntu-latest
+    steps:
+    - name: Get torch-mlir
+      uses: actions/checkout@v2
+      with:
+        submodules: 'true'
+    - uses: ./.github/actions/setup-build
+      with:
+        cache-suffix: ''
+    - name: Build Python wheels and smoke test.
+      run: |
+        cd $GITHUB_WORKSPACE
+        python -m pip install wheel
+        ./build_tools/python_deploy/build_linux_packages.sh
+
+    # If we were given a release_id, then upload the package we just built
+    # to the github releases page.
+    - name: Upload Release Assets (if requested)
+      if: github.event.inputs.release_id != ''
+      id: upload-release-assets
+      uses: dwenegar/upload-release-assets@v1
+      env:
+        GITHUB_TOKEN: ${{ secrets.WORKFLOW_INVOCATION_TOKEN }}
+      with:
+        release_id: ${{ github.event.inputs.release_id }}
+        assets_path: ./wheelhouse/*.whl
+    # Publishing is necessary to make the release visible to `pip`
+    # on the github releases page.
+    - name: Publish Release (if requested)
+      if: github.event.inputs.release_id != ''
+      id: publish_release
+      uses: eregon/publish-release@v1
+      env:
+        GITHUB_TOKEN: ${{ secrets.WORKFLOW_INVOCATION_TOKEN }}
+      with:
+        release_id: ${{ github.event.inputs.release_id }}

.github/workflows/oneshotSnapshotPackage.yml

@@ -0,0 +1,64 @@
+name: Release oneshot snapshot package
+
+on:
+  workflow_dispatch:
+
+jobs:
+  release_snapshot_package:
+    name: "Tag snapshot release"
+    runs-on: ubuntu-latest
+    # Don't run this in everyone's forks.
+    if: github.repository == 'llvm/torch-mlir'
+    steps:
+      - name: Checking out repository
+        uses: actions/checkout@v2
+        with:
+          token: ${{ secrets.WORKFLOW_INVOCATION_TOKEN }}
+
+      - name: Compute version
+        run: |
+          git fetch --depth=1 origin +refs/tags/*:refs/tags/*
+          package_version="$(printf '%(%Y%m%d)T.${{ github.run_number }}')"
+          tag_name="snapshot-${package_version}"
+          echo "package_version=${package_version}" >> $GITHUB_ENV
+          echo "tag_name=${tag_name}" >> $GITHUB_ENV
+
+      - name: Updating snapshot tag
+        run: |
+          git tag "${tag_name}"
+
+      - name: Pushing changes
+        uses: ad-m/github-push-action@v0.6.0
+        with:
+          github_token: ${{ secrets.WORKFLOW_INVOCATION_TOKEN }}
+          branch: main
+          tags: true
+
+      - name: Create Release
+        id: create_release
+        uses: actions/create-release@v1
+        env:
+          GITHUB_TOKEN: ${{ secrets.WORKFLOW_INVOCATION_TOKEN }}
+        with:
+          tag_name: ${{ env.tag_name }}
+          release_name: torch-mlir snapshot ${{ env.tag_name }}
+          body: |
+            Automatic snapshot release of torch-mlir.
+          draft: true
+          prerelease: false
+
+      - name: "Invoke workflow :: Build and Test"
+        uses: benc-uk/workflow-dispatch@v1
+        with:
+          workflow: Build and Test
+          token: ${{ secrets.WORKFLOW_INVOCATION_TOKEN }}
+          ref: "${{ env.tag_name }}"
+          inputs: '{"release_id": "", "python_package_version": "${{ env.package_version }}"}'
+
+      - name: "Invoke workflow :: Manylinux Build"
+        uses: benc-uk/workflow-dispatch@v1
+        with:
+          workflow: Manylinux Build
+          token: ${{ secrets.WORKFLOW_INVOCATION_TOKEN }}
+          ref: "${{ env.tag_name }}"
+          inputs: '{"release_id": "${{ steps.create_release.outputs.id }}", "python_package_version": "${{ env.package_version }}"}'

build_tools/python_deploy/.gitignore

@@ -0,0 +1 @@
+manylinux2014-x64

build_tools/python_deploy/build_linux_packages.sh

@@ -0,0 +1,124 @@
+#!/bin/bash
+# Copyright 2022 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# build_linux_packages.sh
+# One stop build of IREE Python packages for Linux. The Linux build is
+# complicated because it has to be done via a docker container that has
+# an LTS glibc version, all Python packages and other deps.
+# This script handles all of those details.
+#
+# Usage:
+# Build everything (all packages, all python versions):
+#   ./build_tools/python_deploy/build_linux_packages.sh
+#
+# Build specific Python versions and packages to custom directory:
+#   python_versions="cp38-cp38 cp39-cp39" \
+#   packages="torch-mlir" \
+#   output_dir="/tmp/wheelhouse" \
+#   ./build_tools/python_deploy/build_linux_packages.sh
+#
+# Valid Python versions match a subdirectory under /opt/python in the docker
+# image. Typically:
+#   cp37-cp37m cp38-cp38 cp39-cp39 cp310-cp310
+#
+# Valid packages:
+#   torch-mlir
+#
+# Note that this script is meant to be run on CI and it will pollute both the
+# output directory and in-tree build/ directories (under runtime/ and
+# iree/compiler/) with docker created, root owned builds. Sorry - there is
+# no good way around it.
+#
+# It can be run on a workstation but recommend using a git worktree dedicated
+# to packaging to avoid stomping on development artifacts.
+set -eu -o errtrace
+
+this_dir="$(cd $(dirname $0) && pwd)"
+script_name="$(basename $0)"
+repo_root="$(cd $this_dir/../../ && pwd)"
+script_name="$(basename $0)"
+manylinux_docker_image="${manylinux_docker_image:-stellaraccident/manylinux2014_x86_64-bazel-5.1.0:latest}"
+python_versions="${python_versions:-cp37-cp37m cp38-cp38 cp39-cp39 cp310-cp310}"
+output_dir="${output_dir:-${this_dir}/wheelhouse}"
+packages="${packages:-torch-mlir}"
+
+function run_on_host() {
+  echo "Running on host"
+  echo "Launching docker image ${manylinux_docker_image}"
+  echo "Outputting to ${output_dir}"
+  rm -rf "${output_dir}"
+  mkdir -p "${output_dir}"
+  docker run --rm \
+    -v "${repo_root}:/main_checkout/torch-mlir" \
+    -v "${output_dir}:/wheelhouse" \
+    -e __MANYLINUX_BUILD_WHEELS_IN_DOCKER=1 \
+    -e "python_versions=${python_versions}" \
+    -e "packages=${packages}" \
+    ${manylinux_docker_image} \
+    -- bash /main_checkout/torch-mlir/build_tools/python_deploy/build_linux_packages.sh
+}
+
+function run_in_docker() {
+  echo "Running in docker"
+  echo "Using python versions: ${python_versions}"
+
+  local orig_path="$PATH"
+
+  # Build phase.
+  for package in $packages; do
+    echo "******************** BUILDING PACKAGE ${package} ********************"
+    for python_version in $python_versions; do
+      python_dir="/opt/python/$python_version"
+      if ! [ -x "$python_dir/bin/python" ]; then
+        echo "ERROR: Could not find python: $python_dir (skipping)"
+        continue
+      fi
+      export PATH=$python_dir/bin:$orig_path
+      echo ":::: Python version $(python --version)"
+      case "$package" in
+        torch-mlir)
+          clean_wheels torch_mlir $python_version
+          build_torch_mlir
+          #run_audit_wheel torch_mlir $python_version
+          ;;
+        *)
+          echo "Unrecognized package '$package'"
+          exit 1
+          ;;
+      esac
+    done
+  done
+}
+
+function build_torch_mlir() {
+  CMAKE_GENERATOR=Ninja \
+  python -m pip wheel -v -w /wheelhouse /main_checkout/torch-mlir/ \
+    --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+}
+
+function run_audit_wheel() {
+  local wheel_basename="$1"
+  local python_version="$2"
+  generic_wheel="/wheelhouse/${wheel_basename}-*-${python_version}-linux_x86_64.whl"
+  echo ":::: Auditwheel $generic_wheel"
+  auditwheel repair -w /wheelhouse $generic_wheel
+  rm $generic_wheel
+}
+
+function clean_wheels() {
+  local wheel_basename="$1"
+  local python_version="$2"
+  echo ":::: Clean wheels $wheel_basename $python_version"
+  rm -f /wheelhouse/${wheel_basename}-*-${python_version}-*.whl
+}
+
+# Trampoline to the docker container if running on the host.
+if [ -z "${__MANYLINUX_BUILD_WHEELS_IN_DOCKER-}" ]; then
+  run_on_host "$@"
+else
+  run_in_docker "$@"
+fi

build_tools/python_deploy/build_macos_packages.sh

@@ -0,0 +1,77 @@
+#!/bin/bash
+# Copyright 2022 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# build_macos_packages.sh
+# One stop build of IREE Python packages for MacOS. This presumes that
+# dependencies are installed from install_macos_deps.sh. This will build
+# for a list of Python versions synchronized with that script and corresponding
+# with directory names under:
+#   /Library/Frameworks/Python.framework/Versions
+#
+# MacOS convention is to refer to this as major.minor (i.e. "3.9", "3.10").
+# Valid packages:
+#   torch-mlir
+
+set -eu -o errtrace
+
+this_dir="$(cd $(dirname $0) && pwd)"
+repo_root="$(cd $this_dir/../../ && pwd)"
+python_versions="${python_versions:3.9 3.10}"
+output_dir="${output_dir:-${this_dir}/wheelhouse}"
+packages="${packages:-torch-mlir}"
+
+# Note that this typically is selected to match the version that the official
+# Python distributed is built at.
+#export MACOSX_DEPLOYMENT_TARGET=11.0
+#export CMAKE_OSX_ARCHITECTURES="arm64;x86_64"
+
+function run() {
+  echo "Using python versions: ${python_versions}"
+
+  local orig_path="$PATH"
+
+  # Build phase.
+  for package in $packages; do
+    echo "******************** BUILDING PACKAGE ${package} ********************"
+    for python_version in $python_versions; do
+      python_dir="/Library/Frameworks/Python.framework/Versions/$python_version"
+      if ! [ -x "$python_dir/bin/python3" ]; then
+        echo "ERROR: Could not find python3: $python_dir (skipping)"
+        continue
+      fi
+      export PATH=$python_dir/bin:$orig_path
+      echo ":::: Python version $(python3 --version)"
+      case "$package" in
+        torch-mlir)
+          clean_wheels torch-mlir $python_version
+          build_torch_mlir
+          ;;
+        *)
+          echo "Unrecognized package '$package'"
+          exit 1
+          ;;
+      esac
+    done
+  done
+}
+
+function build_torch_mlir() {
+  python -m pip install -r /main_checkout/torch-mlir/requirements.txt
+  CMAKE_GENERATOR=Ninja \
+  MACOSX_DEPLOYMENT_TARGET=11.0 \
+  CMAKE_OSX_ARCHITECTURES="arm64;x86_64" \
+  python -m pip wheel -v -w /wheelhouse /main_checkout/torch-mlir/
+}
+
+function clean_wheels() {
+  local wheel_basename="$1"
+  local python_version="$2"
+  echo ":::: Clean wheels $wheel_basename $python_version"
+  rm -f /wheelhouse/${wheel_basename}-*-${python_version}-*.whl
+}
+
+run

build_tools/python_deploy/install_macos_deps.sh

@@ -0,0 +1,61 @@
+#!/bin/zsh
+# Copyright 2022 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# Installs dependencies on MacOS necessary to build IREE.
+# Additional dependencies (i.e MoltenVK) may be needed to use all functionality.
+#
+# Usage:
+#   sudo install_macos_deps.sh
+
+set -e -o pipefail
+
+if [[ "$(whoami)" != "root" ]]; then
+  echo "ERROR: Must setup deps as root"
+  exit 1
+fi
+
+PYTHON_INSTALLER_URLS=(
+  "https://www.python.org/ftp/python/3.10.4/python-3.10.4-macos11.pkg"
+  "https://www.python.org/ftp/python/3.9.12/python-3.9.12-macos11.pkg"
+)
+
+PYTHON_SPECS=(
+  3.10@https://www.python.org/ftp/python/3.10.4/python-3.10.4-macos11.pkg
+  3.9@https://www.python.org/ftp/python/3.9.12/python-3.9.12-macos11.pkg
+)
+
+for python_spec in $PYTHON_SPECS; do
+  python_version="${python_spec%%@*}"
+  url="${python_spec##*@}"
+  echo "-- Installing Python $python_version from $url"
+  python_path="/Library/Frameworks/Python.framework/Versions/$python_version"
+  python_exe="$python_path/bin/python3"
+
+  # Install Python.
+  if ! [ -x "$python_exe" ]; then
+    package_basename="$(basename $url)"
+    download_path="/tmp/torch_mlir_python_install/$package_basename"
+    mkdir -p "$(dirname $download_path)"
+    echo "Downloading $url -> $download_path"
+    curl $url -o "$download_path"
+
+    echo "Installing $download_path"
+    installer -pkg "$download_path" -target /
+  else
+    echo ":: Python version already installed. Not reinstalling."
+  fi
+
+  echo ":: Python version $python_version installed:"
+  $python_exe --version
+  $python_exe -m pip --version
+
+  echo ":: Installing system pip packages"
+  $python_exe -m pip install --upgrade pip
+  $python_exe -m pip install --upgrade delocate
+done
+
+echo "*** All done ***"

pyproject.toml

@@ -0,0 +1,21 @@
+[build-system]
+requires = [
+    "setuptools>=42",
+    "wheel",
+    # There is no fundamental reason to pin this CMake version, beyond
+    # build stability.
+    "cmake==3.22.2",
+    "ninja==1.10.2",
+    "packaging",
+    # Version 2.7.0 excluded: https://github.com/pybind/pybind11/issues/3136
+    "pybind11>=2.6.0,!=2.7.0",
+    "PyYAML",
+
+    # The torch-mlir CMake build requires numpy and torch to be installed.
+    # Further, the setup.py will pin the version selected here into built
+    # artifacts.
+    # TODO: Come up with a better way to pin the version.
+    "numpy",
+    "torch==1.12.0.dev20220419+cpu",
+]
+build-backend = "setuptools.build_meta"

python/torch_mlir/dialects/torch/importer/jit_ir/csrc/CMakeLists.txt

@@ -24,10 +24,19 @@ add_library(TorchMLIRJITIRImporter MODULE
 target_link_libraries(TorchMLIRJITIRImporter
   TorchMLIRAggregateCAPI
   ${TORCH_LIBRARIES}
-  ${Python3_LIBRARIES}
   torch_python
 )
 
+# On static Python builds, there may not be Python libraries to link against
+# (they will late bind at runtime from the executable). We have to condition
+# this because in that case it is set to NOTFOUND and CMake will consider
+# this an error.
+if(Python3_LIBRARIES)
+  target_link_libraries(TorchMLIRJITIRImporter
+    ${Python3_LIBRARIES}
+  )
+endif()
+
 message(STATUS "TORCH_CXXFLAGS=${TORCH_CXXFLAGS}")
 set_target_properties(TorchMLIRJITIRImporter PROPERTIES
   LIBRARY_OUTPUT_DIRECTORY "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/torch_mlir/_mlir_libs"

requirements.txt

@@ -8,3 +8,9 @@ torchvision
 # Build requirements.
 pybind11
 wheel
+setuptools
+cmake
+ninja
+
+# Test Requirements
+pillow

setup.py

@@ -33,8 +33,10 @@ import os
 import shutil
 import subprocess
 import sys
+import sysconfig
 
 from distutils.command.build import build as _build
+from distutils.sysconfig import get_python_inc
 from setuptools import setup, Extension
 from setuptools.command.build_ext import build_ext
 from setuptools.command.build_py import build_py
@@ -51,7 +53,6 @@ class CustomBuild(_build):
         self.run_command("build_ext")
         self.run_command("build_scripts")
 
-
 class CMakeBuild(build_py):
 
     def run(self):
@@ -60,6 +61,10 @@ class CMakeBuild(build_py):
         if not cmake_build_dir:
             cmake_build_dir = os.path.abspath(
                 os.path.join(target_dir, "..", "cmake_build"))
+        python_package_dir = os.path.join(cmake_build_dir,
+                                          "tools", "torch-mlir", "python_packages",
+                                          "torch_mlir")
+
         if not os.getenv("TORCH_MLIR_CMAKE_BUILD_DIR_ALREADY_BUILT"):
             src_dir = os.path.abspath(os.path.dirname(__file__))
             llvm_dir = os.path.join(
@@ -82,15 +87,26 @@ class CMakeBuild(build_py):
             cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt")
             if os.path.exists(cmake_cache_file):
                 os.remove(cmake_cache_file)
+            # NOTE: With repeated builds for different Python versions, the
+            # prior version binaries will continue to accumulate. IREE uses
+            # a separate install step and cleans the install directory to
+            # keep this from happening. That is the most robust. Here we just
+            # delete the directory where we build native extensions to keep
+            # this from happening but still take advantage of most of the
+            # build cache.
+            mlir_libs_dir = os.path.join(python_package_dir, "torch_mlir", "_mlir_libs")
+            if os.path.exists(mlir_libs_dir):
+                print(f"Removing _mlir_mlibs dir to force rebuild: {mlir_libs_dir}")
+                shutil.rmtree(mlir_libs_dir)
+            else:
+                print(f"Not removing _mlir_libs dir (does not exist): {mlir_libs_dir}")
+
             subprocess.check_call(["cmake", llvm_dir] +
                                   cmake_args, cwd=cmake_build_dir)
             subprocess.check_call(["cmake",
                                    "--build",  ".",
                                    "--target", "TorchMLIRPythonModules"],
                                   cwd=cmake_build_dir)
-        python_package_dir = os.path.join(cmake_build_dir,
-                                          "tools", "torch-mlir", "python_packages",
-                                          "torch_mlir")
 
         if os.path.exists(target_dir):
             shutil.rmtree(target_dir, ignore_errors=False, onerror=None)
@@ -130,8 +146,11 @@ setup(
         CMakeExtension("torch_mlir._mlir_libs._jit_ir_importer"),
     ],
     install_requires=[
+        "numpy",
         # To avoid issues with drift for each nightly build, we pin to the
         # exact version we built against.
+        # TODO: This includes the +cpu specifier which is overly
+        # restrictive and a bit unfortunate.
         f"torch=={torch.__version__}",
     ],
     zip_safe=False,