Add torch_mlir snapshot packages.

This closely follows IREE's [schedule_snapshot_release.yml](https://github.com/google/iree/blob/f2f153d39472b3abb1b629517e3b2c0bb1812c77/.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" ```
Abhishek-TyRnT · Oct 6, 2021 · 4a8d05e · 4a8d05e
1 parent 712445e
commit 4a8d05e
Show file tree

Hide file tree

Showing 4 changed files with 129 additions and 30 deletions.
diff --git 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
diff --git 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 }}"}'
diff --git 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"
diff --git 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 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 @@ def run(self):
         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 @@ def run(self):
                         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 @@ def build_extension(self, ext):
         "built_ext": NoopBuildExtension,
         "build_py": CMakeBuild,
     },
+    ext_modules=[
+        CMakeExtension("torch_mlir._mlir_libs._jit_ir_importer"),
+    ],
     zip_safe=False,
 )