Add torch_mlir snapshot packages.

This closely follows IREE's [schedule_snapshot_release.yml](f2f153d394/.github/workflows/schedule_snapshot_release.yml (L1)) workflow. The snapshot releases can be installed with: ``` python -m pip install torch_mlir -f "https://github.com/llvm/torch-mlir/releases" ```
2021-10-05 23:26:26 +00:00 · 2021-10-05 23:26:26 +00:00 · 4a8d05e4a5
parent 712445eaa8
commit 4a8d05e4a5
4 changed files with 129 additions and 30 deletions
--- a/.github/workflows/buildAndTest.yml
+++ b/.github/workflows/buildAndTest.yml
@ -1,6 +1,13 @@
 name: Build and Test

-on: [push, pull_request]
+on: 
+  push:
+  pull_request:
+  workflow_dispatch:
+    inputs:
+      release_id:
+        description: 'Release id to upload artifacts to'
+        default: ''

 jobs:
  build:
@ -15,10 +22,10 @@ jobs:
      uses: actions/checkout@v2
      with:
        submodules: 'true'
-    - name: Install python depends
+    - name: Install MLIR Python depends
      run: |
        python -m pip install -r $GITHUB_WORKSPACE/external/llvm-project/mlir/python/requirements.txt
-    - name: Install pytorch_nightly depends
+    - name: Install PyTorch nightly depends
      run: |
        python -m pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
    - name: Install Ninja
@ -54,3 +61,23 @@ jobs:
        cd $GITHUB_WORKSPACE
        export PYTHONPATH="$GITHUB_WORKSPACE/build/tools/torch-mlir/python_packages/torch_mlir"
        python -m e2e_testing.torchscript.main --config=refbackend -v
+
+    # TODO: Only build packages in full Release mode.
+    # On the other hand, having assertions on isn't too bad of an idea at this
+    # early stage.
+    - name: Build Python wheels and smoke test.
+      run: |
+        cd $GITHUB_WORKSPACE
+        python -m pip install wheel
+        TORCH_MLIR_CMAKE_BUILD_DIR="$GITHUB_WORKSPACE/build" \
+          TORCH_MLIR_CMAKE_BUILD_DIR_ALREADY_BUILT=1 \
+          ./build_tools/build_python_wheels.sh
+    - 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
--- a/.github/workflows/releaseSnapshotPackage.yml
+++ b/.github/workflows/releaseSnapshotPackage.yml
@ -0,0 +1,59 @@
+name: Release snapshot package
+
+on:
+  schedule:
+    - cron: '0 10,22 * * *'
+
+  workflow_dispatch:
+
+jobs:
+  release_snapshot_package:
+    name: "Tag snapshot release"
+    runs-on: ubuntu-20.04
+    # 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: true
+
+      - 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": "${{ steps.create_release.outputs.id }}"}'
--- a/build_tools/build_python_wheels.sh
+++ b/build_tools/build_python_wheels.sh
@ -15,9 +15,7 @@ cd $wheelhouse

 echo "---- BUILDING torch-mlir ----"
 CMAKE_GENERATOR=Ninja CMAKE_C_COMPILER_LAUNCHER=ccache CMAKE_CXX_COMPILER_LAUNCHER=ccache \
-$PYTHON "${repo_root}/setup.py" bdist_wheel \
-  --dist-dir "$wheelhouse" -v
-
+$PYTHON "${repo_root}/setup.py" bdist_wheel --dist-dir "$wheelhouse" -v

 # Smoke test: create a venv, install the package, and run an example.

@ -27,10 +25,10 @@ VENV_PYTHON="$package_test_venv/bin/python"

 echo "---- INSTALLING torch ----"
 $VENV_PYTHON -m pip install --pre torch torchvision pybind11 -f "https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html"
-echo "---- INSTALLING other deps for smoke test script ----"
+echo "---- INSTALLING other deps for smoke test ----"
 $VENV_PYTHON -m pip install requests pillow
 echo "---- INSTALLING torch-mlir ----"
-$VENV_PYTHON -m pip install -f "$wheelhouse" --force-reinstall torch-mlir
+$VENV_PYTHON -m pip install -f "$wheelhouse" --force-reinstall torch_mlir

 echo "---- RUNNING SMOKE TEST ----"
 $VENV_PYTHON "$repo_root/examples/torchscript_resnet18_e2e.py"
--- a/setup.py
+++ b/setup.py
@ -16,6 +16,9 @@
 #
 # On CIs, it is often advantageous to re-use/control the CMake build directory.
 # This can be set with the TORCH_MLIR_CMAKE_BUILD_DIR env var.
+# Additionally, the TORCH_MLIR_CMAKE_BUILD_DIR_ALREADY_BUILT env var will
+# prevent this script from attempting to build the directory, and will simply
+# use the (presumed already built) directory as-is.
 #
 # Implementation notes:
 # The contents of the wheel is just the contents of the `python_packages`
@ -28,7 +31,7 @@ import subprocess
 import sys

 from distutils.command.build import build as _build
-from setuptools import setup
+from setuptools import setup, Extension
 from setuptools.command.build_ext import build_ext
 from setuptools.command.build_py import build_py

@ -49,27 +52,29 @@ class CMakeBuild(build_py):
        cmake_build_dir = os.getenv("TORCH_MLIR_CMAKE_BUILD_DIR")
        if not cmake_build_dir:
            cmake_build_dir = os.path.join(target_dir, "..", "cmake_build")
-        src_dir = os.path.abspath(os.path.dirname(__file__))
-        llvm_dir = os.path.join(src_dir, "external", "llvm-project", "llvm")
-        cmake_args = [
-            f"-DCMAKE_BUILD_TYPE=Release",
-            f"-DPython3_EXECUTABLE={sys.executable}",
-            f"-DLLVM_TARGETS_TO_BUILD=host",
-            f"-DMLIR_ENABLE_BINDINGS_PYTHON=ON",
-            f"-DLLVM_ENABLE_PROJECTS=mlir",
-            f"-DLLVM_EXTERNAL_PROJECTS=torch-mlir",
-            f"-DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR={src_dir}",
-        ]
-        os.makedirs(cmake_build_dir, exist_ok=True)
-        cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt")
-        if os.path.exists(cmake_cache_file):
-            os.remove(cmake_cache_file)
-        subprocess.check_call(["cmake", llvm_dir] +
-                              cmake_args, cwd=cmake_build_dir)
-        subprocess.check_call(["cmake",
-                               "--build",  ".",
-                               "--target", "TorchMLIRPythonModules"],
-                              cwd=cmake_build_dir)
+        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(
+                src_dir, "external", "llvm-project", "llvm")
+            cmake_args = [
+                f"-DCMAKE_BUILD_TYPE=Release",
+                f"-DPython3_EXECUTABLE={sys.executable}",
+                f"-DLLVM_TARGETS_TO_BUILD=host",
+                f"-DMLIR_ENABLE_BINDINGS_PYTHON=ON",
+                f"-DLLVM_ENABLE_PROJECTS=mlir",
+                f"-DLLVM_EXTERNAL_PROJECTS=torch-mlir",
+                f"-DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR={src_dir}",
+            ]
+            os.makedirs(cmake_build_dir, exist_ok=True)
+            cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt")
+            if os.path.exists(cmake_cache_file):
+                os.remove(cmake_cache_file)
+            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")
@ -79,6 +84,13 @@ class CMakeBuild(build_py):
                        dirs_exist_ok=True)


+class CMakeExtension(Extension):
+
+  def __init__(self, name, sourcedir=""):
+    Extension.__init__(self, name, sources=[])
+    self.sourcedir = os.path.abspath(sourcedir)
+
+
 class NoopBuildExtension(build_ext):

    def build_extension(self, ext):
@ -98,5 +110,8 @@ setup(
        "built_ext": NoopBuildExtension,
        "build_py": CMakeBuild,
    },
+    ext_modules=[
+        CMakeExtension("torch_mlir._mlir_libs._jit_ir_importer"),
+    ],
    zip_safe=False,
 )