python: separate build- and test-related pip dependencies (#1874)

We want to ensure that pip packages required for building torch-mlir
should be included in the dependencies of torch-mlir, but we don't want
the pip packages required for _testing_ of torch-mlir to be included
among the dependencies.  To be able to specify and install one set of
dependencies and not the other, this patch separates the pip packages
into two files: build-requirements.txt and test-requirements.txt.

This patch also updates references to the requirements.txt file so that
CI builds that run end-to-end tests install test-related pip
dependencies while everything else (including WHL builds) sticks to just
the build-related pip dependencies.

Despite this change, this patch should not affect a torch-mlir
developer's workflow.  More precisely, since this patch makes the
top-level requirements.txt file refer to both build-requirements.txt and
test-requirements.txt files, a torch-mlir developer should be able to
continue referring to the requirements.txt file without any impact.

.github/actions/setup-build/action.yml

@@ -26,7 +26,8 @@ runs:
 
   - name: Install PyTorch nightly depends
     run: |
-      python -m pip install -r requirements.txt
+      python -m pip install -r pytorch-requirements.txt
+      python -m pip install -r build-requirements.txt
     shell: bash
 
   - name: Install prerequisites (Linux)

.github/workflows/RollPyTorch.yml

@@ -44,7 +44,8 @@ jobs:
         # Read the version from the downloaded whl file without extracting it
         PT_RELEASE=$(unzip -p torch-*.whl 'torch-*/METADATA' | grep "^Version:" | awk '{ print $2 }' | sed 's/\([^+]*\).*/\1/')
         echo "Found torch release ${PT_RELEASE}"
-        printf -- "-f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html\n--pre\ntorch==%s\ntorchvision==%s\n" "${PT_RELEASE}" "${VISION_RELEASE}" > pytorch-requirements.txt
+        printf -- "-f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html\n--pre\ntorch==%s\n" "${PT_RELEASE}" > pytorch-requirements.txt
+        printf -- "-f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html\n--pre\ntorchvision==%s\n" "${VISION_RELEASE}" > torchvision-requirements.txt
 
         # Read the commit hash from the downloaded whl file without extracting it
         PT_HASH=$(unzip -p torch-"${PT_RELEASE}"*.whl torch/version.py | grep git_version | awk '{ print $3 }' | tr -d "'")
@@ -96,7 +97,7 @@ jobs:
         git fetch --recurse-submodules=no
         git checkout main
         git pull origin main
-        git add pytorch-hash.txt pytorch-requirements.txt lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
+        git add pytorch-hash.txt pytorch-requirements.txt torchvision-requirements.txt lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
         git diff --cached --exit-code || (git commit -m "update PyTorch version to ${{ env.PT_RELEASE }}" && git push --set-upstream origin main)
 
     - name: Update PyTorch Build Cache (if running on main branch)

build_tools/python_deploy/build_macos_packages.sh

@@ -77,7 +77,8 @@ function build_torch_mlir() {
   python"${python_version}" -m venv "$output_dir"/build_venv
   source "$output_dir"/build_venv/bin/activate
   python"${python_version}" -m pip install -U pip
-  python"${python_version}" -m pip install -r "$repo_root"/requirements.txt --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+  python"${python_version}" -m pip install -r "$repo_root"/pytorch-requirements.txt --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+  python"${python_version}" -m pip install -r "$repo_root"/build-requirements.txt
   CMAKE_GENERATOR=Ninja \
   TORCH_MLIR_PYTHON_PACKAGE_VERSION=${TORCH_MLIR_PYTHON_PACKAGE_VERSION} \
   MACOSX_DEPLOYMENT_TARGET=$MACOSX_DEPLOYMENT_TARGET \
@@ -107,7 +108,8 @@ function run_audit_wheel() {
     python"${python_version}" -m venv "$output_dir"/test_venv
     source "$output_dir"/test_venv/bin/activate
     python"${python_version}" -m pip install -U pip
-    python"${python_version}" -m pip install -r "$repo_root"/requirements.txt --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+    python"${python_version}" -m pip install -r "$repo_root"/pytorch-requirements.txt --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+    python"${python_version}" -m pip install -r "$repo_root"/build-requirements.txt
     python"${python_version}" -m pip install "$generic_wheel" --extra-index-url https://download.pytorch.org/whl/nightly/cpu
     DYLD_LIBRARY_PATH="$output_dir"/test_venv/lib/python"${python_version}"/site-packages/torch/lib delocate-wheel -v "$generic_wheel"
     deactivate

build_tools/python_deploy/build_windows.ps1

@@ -13,7 +13,8 @@
 Write-Host "Installing Build Dependencies"
 python -m venv .\mlir_venv\
 .\mlir_venv\Scripts\Activate.PS1
-pip install -r .\requirements.txt
+pip install -r .\pytorch-requirements.txt
+pip install -r .\build-requirements.txt
 pip install delvewheel
 Write-Host "Build Deps installation completed successfully"
 

pytorch-requirements.txt

@@ -1,4 +1,3 @@
 -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
 --pre
 torch==2.0.0.dev20230212
-torchvision==0.15.0.dev20230213

requirements.txt

@@ -1,7 +1,3 @@
--r build-requirements.txt
 -r pytorch-requirements.txt
-
-# Test Requirements
-pillow
-dill
-multiprocess
+-r build-requirements.txt
+-r test-requirements.txt

test-requirements.txt

@@ -0,0 +1,5 @@
+-r torchvision-requirements.txt
+
+pillow
+dill
+multiprocess

torchvision-requirements.txt

@@ -0,0 +1,3 @@
+-f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+--pre
+torchvision==0.15.0.dev20230213