From 62d266f7cfbd41dc885b75705d4e08af84e368fe Mon Sep 17 00:00:00 2001
From: Harutaka Kawamura <hkawamura0130@gmail.com>
Date: Mon, 10 May 2021 17:17:08 +0900
Subject: [PATCH] Cache catboost wheel for dev version (#4285)

* Cache catboost wheel for dev version since it takes long to build it

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>

* Build wheel and cache

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>

* Fix incorrect directory path

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>

* split step

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>

* use -

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>

* Use matrix values

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>

* Clarify long

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
---
 .github/workflows/cross-version-tests.yml | 12 ++++++++----
 mlflow/ml-package-versions.yml            | 19 ++++++++++++++++++-
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/cross-version-tests.yml b/.github/workflows/cross-version-tests.yml
index 5f0a824511ac8..b3c59d332a008 100644
--- a/.github/workflows/cross-version-tests.yml
+++ b/.github/workflows/cross-version-tests.yml
@@ -69,19 +69,23 @@ jobs:
           hash=$(echo -n "$INSTALL_COMMAND" | sha256sum)
           echo "::set-output name=key::$ImageOS-$ImageVersion-wheels-$date-$hash"
       - uses: actions/cache@v2
-        if: ${{ matrix.package == 'pyspark' && matrix.version == 'dev' }}
+        # We only cache wheels that take long (> 20 min) to build
+        if: ${{ contains('pyspark, catboost', matrix.package) && matrix.version == 'dev' }}
         with:
           path: /home/runner/.cache/wheels
           key: ${{ steps.get-cache-key.outputs.key }}
           restore-keys: |
             ${{ steps.get-cache-key.outputs.key }}
-      - name: Install dependencies
-        env:
-          CACHE_DIR: /home/runner/.cache/wheels
+      - name: Install mlflow & test dependencies
         run: |
           set -x
           pip install -e .
           pip install -r dev/small-requirements.txt
+      - name: Install ${{ matrix.package }} ${{ matrix.version }}
+        env:
+          CACHE_DIR: /home/runner/.cache/wheels
+        run: |
+          set -x
           ${{ matrix.install }}
       - name: Check package versions
         run: |
diff --git a/mlflow/ml-package-versions.yml b/mlflow/ml-package-versions.yml
index 3229d8628e775..930b1902cf216 100644
--- a/mlflow/ml-package-versions.yml
+++ b/mlflow/ml-package-versions.yml
@@ -176,7 +176,24 @@ catboost:
   package_info:
     pip_release: "catboost"
     install_dev: |
-      pip install git+https://github.com/catboost/catboost.git#subdirectory=catboost/python-package
+      # The cross-version-tests workflow runs this command with the environment variable `CACHE_DIR`
+      if [ -d "$CACHE_DIR" ] && [ ! -z $(find $CACHE_DIR -type f -name "catboost-*.whl") ]; then
+        pip install $(find $CACHE_DIR -type f -name "catboost-*.whl")
+      else
+        # Build wheel from source
+        temp_dir=$(mktemp -d)
+        git clone --depth 1 https://github.com/catboost/catboost $temp_dir
+        cd $temp_dir/catboost/python-package
+        python setup.py bdist_wheel
+
+        # Copy wheel in cache directory
+        wheel_path=$(find dist -type f -name "*.whl")
+        mkdir -p $CACHE_DIR
+        cp $wheel_path $CACHE_DIR
+
+        # Install wheel
+        pip install $wheel_path
+      fi
 
   models:
     minimum: "0.23.1"