From b4b27fafbd89c70dbf098d664252cd22bafe01a9 Mon Sep 17 00:00:00 2001
From: Junru Shao <junrushao@apache.org>
Date: Fri, 30 Jun 2023 20:17:30 -0700
Subject: [PATCH] [Python] Enhance Wheel Packaging (#15167)

* [Python] Enhance Wheel Packaging

This PR allows us to package some extra directories with TVM wheels for
developing other language bindings.

* Update gen_requirements.py

* Disable Android Camera test for unmaintained MXNet
---
 .github/workflows/main.yml | 21 --------------------
 python/gen_requirements.py |  6 ++----
 python/setup.py            | 39 +++++++++++++++++++++++++++++---------
 3 files changed, 32 insertions(+), 34 deletions(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 977bc189100b..a4a30fe19aba 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -179,24 +179,3 @@ jobs:
         with:
           name: android_deploy-debug.apk
           path: ./apps/android_deploy/app/build/outputs/apk/debug/app-debug.apk
-      - name: Build android_camera
-        working-directory: apps/android_camera
-        run: |
-          export TVM_HOME=~/work/tvm/tvm
-          export PYTHONPATH=$TVM_HOME/python:${PYTHONPATH}
-          set -eux
-          mkdir -p app/src/main/assets/models/
-          export TVM_NDK_CC=${ANDROID_NDK_LATEST_HOME}/toolchains/llvm/prebuilt/linux-x86_64/bin/aarch64-linux-android30-clang++
-          python3 ${TVM_HOME}/python/gen_requirements.py
-          pip3 install -r ${TVM_HOME}/python/requirements/core.txt
-          cd models
-          pip3 install -r requirements.txt
-          python3 prepare_model.py
-          cd ..
-          export PATH="${ANDROID_NDK_LATEST_HOME}:$PATH"
-          gradle clean build
-      - name: Upload android_camera APK
-        uses: actions/upload-artifact@v2
-        with:
-          name: android_camera-debug.apk
-          path: ./apps/android_camera/app/build/outputs/apk/debug/app-debug.apk
diff --git a/python/gen_requirements.py b/python/gen_requirements.py
index c9eb5d9cc660..0c8200f60b10 100644
--- a/python/gen_requirements.py
+++ b/python/gen_requirements.py
@@ -47,11 +47,10 @@
 import collections
 import os
 import re
-import textwrap
 import sys
+import textwrap
 import typing
 
-
 RequirementsByPieceType = typing.List[typing.Tuple[str, typing.Tuple[str, typing.List[str]]]]
 
 
@@ -253,8 +252,7 @@
     ("h5py", "==2.10.0"),
     ("image", None),
     ("matplotlib", None),
-    # Workaround, see https://github.com/apache/tvm/issues/13647
-    ("numpy", "<=1.23"),
+    ("numpy", None),
     ("onnx", None),
     ("onnxoptimizer", None),
     ("onnxruntime", None),
diff --git a/python/setup.py b/python/setup.py
index f63d52bd3806..b4c4c40943b0 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -17,11 +17,10 @@
 # pylint: disable=invalid-name, exec-used
 """Setup TVM package."""
 import os
+import pathlib
 import shutil
 import sys
 import sysconfig
-import pathlib
-import platform
 
 from setuptools import find_packages
 from setuptools.dist import Distribution
@@ -77,6 +76,26 @@ def get_lib_path():
                 libs.append(candidate_path)
                 break
 
+        for dir in [
+            "3rdparty",
+            "jvm",
+            "web",
+            "rust",
+            "golang",
+            "include",
+            "src",
+            "cmake",
+            "CMakeLists.txt",
+        ]:
+            for name in lib_path:
+                candidate_path = os.path.abspath(os.path.join(os.path.dirname(name), "..", dir))
+                if os.path.exists(candidate_path):
+                    libs.append(candidate_path)
+                    if dir == "3rdparty":
+                        # remove large files
+                        _remove_path(os.path.join(candidate_path, "cutlass", "docs"))
+                        _remove_path(os.path.join(candidate_path, "cutlass", "media"))
+                    break
     else:
         libs = None
 
@@ -94,6 +113,14 @@ def git_describe_version(original_version):
     return gd_version
 
 
+def _remove_path(path):
+    if os.path.exists(path):
+        if os.path.isfile(path):
+            os.remove(path)
+        elif os.path.isdir(path):
+            shutil.rmtree(path)
+
+
 LIB_LIST, __version__ = get_lib_path()
 __version__ = git_describe_version(__version__)
 
@@ -245,10 +272,4 @@ def long_description_contents():
     os.remove("MANIFEST.in")
     for path in LIB_LIST:
         _, libname = os.path.split(path)
-        path_to_be_removed = f"tvm/{libname}"
-
-        if os.path.isfile(path_to_be_removed):
-            os.remove(path_to_be_removed)
-
-        if os.path.isdir(path_to_be_removed):
-            shutil.rmtree(path_to_be_removed)
+        _remove_path(f"tvm/{libname}")