feat(imperative): support python3.9 and fix some tests

GitOrigin-RevId: 87e56f9667

imperative/python/requires-test.txt

@@ -1,5 +1,6 @@
 pytest==5.3.0
 pytest-sphinx==0.3.1
 tensorboardX==2.4
+protobuf==3.20.0 ; python_version > '3.8'
 six==1.16.0
 redislite ; platform_system == "Linux" or platform_system == "Darwin"

imperative/python/setup.py

@@ -120,10 +120,11 @@ setup_kwargs.update(dict(
     'License :: OSI Approved :: Apache Software License',
     'Programming Language :: C++',
     'Programming Language :: Python :: 3',
-    'Programming Language :: Python :: 3.5',
     'Programming Language :: Python :: 3.6',
     'Programming Language :: Python :: 3.7',
     'Programming Language :: Python :: 3.8',
+    'Programming Language :: Python :: 3.9',
+    'Programming Language :: Python :: 3.10',
     'Topic :: Scientific/Engineering',
     'Topic :: Scientific/Engineering :: Mathematics',
     'Topic :: Scientific/Engineering :: Artificial Intelligence',

imperative/python/test/run.sh

@@ -25,19 +25,18 @@ if [[ "$TEST_PLAT" =~ "local" ]]; then
     megengine_dir=`python3 -c 'from pathlib import Path;import megengine;print(Path(megengine.__file__).resolve().parent)'`
     test_dirs="${megengine_dir} ."
 
-    # FIXME: at aarch64 env, run megengine_dir pytest have exit issue!!
     machine=$(uname -m)
     case ${machine} in
-        x86_64) test_dirs="${megengine_dir} ." ;;
-        aarch64) test_dirs="." ;;
+        x86_64) echo "test on ${machine}" ;;
+        aarch64) echo "test on ${machine}" ;;
         *) echo "nonsupport env!!!";exit -1 ;;
     esac
 
     echo "test local env at: ${test_dirs}"
-    PY_IGNORE_IMPORTMISMATCH=1 python3 -m pytest -s -v $test_dirs -m 'not isolated_distributed'
+    PY_IGNORE_IMPORTMISMATCH=1 python3 -m pytest --ignore=${megengine_dir}/_internal -s -v $test_dirs -m 'not isolated_distributed'
     if [[ "$TEST_PLAT" =~ "cuda" ]]; then
         echo "test GPU pytest now"
-        PY_IGNORE_IMPORTMISMATCH=1 python3 -m pytest -s -v $test_dirs -m 'isolated_distributed' --ignore=./integration/test_dtr.py
+        PY_IGNORE_IMPORTMISMATCH=1 python3 -m pytest --ignore=${megengine_dir}/_internal -s -v $test_dirs -m 'isolated_distributed' --ignore=./integration/test_dtr.py
     fi
 else
     cd $(dirname "${BASH_SOURCE[0]}")/..

scripts/cmake-build/BUILD_README.md

@@ -44,7 +44,7 @@ But some dependencies need to be installed manually:
 * [CUDA](https://developer.nvidia.com/cuda-toolkit-archive)(>=10.1), [cuDNN](https://developer.nvidia.com/cudnn)(>=7.6) are required when building MegBrain with CUDA support.
 * [TensorRT](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)(>=5.1.5) is required when building with TensorRT support.
 * LLVM/Clang(>=6.0) is required when building with Halide JIT support.
-* Python(>=3.5) and numpy are required to build Python modules.
+* Python(>=3.6) and numpy are required to build Python modules.
 ## Package install
 ### Windows host build
 * commands:
@@ -65,15 +65,17 @@ But some dependencies need to be installed manually:
 	* our ci use LLVM 12.0.1, if u install other version, please modify LLVM_PATH
     * install 12.0.1 to /c/Program\ Files/LLVM_12_0_1
 4: install python3 (Windows GUI)
-	* download python 64-bit install exe (we support python3.5-python3.8 now)
-	     https://www.python.org/ftp/python/3.5.4/python-3.5.4-amd64.exe
+	* download python 64-bit install exe (we support python3.6-python3.9 now)
 	     https://www.python.org/ftp/python/3.6.8/python-3.6.8-amd64.exe
 	     https://www.python.org/ftp/python/3.7.7/python-3.7.7-amd64.exe
 	     https://www.python.org/ftp/python/3.8.3/python-3.8.3-amd64.exe
-	* install 3.5.4 to /c/Users/${USER}/mge_whl_python_env/3.5.4
+	     https://www.python.org/ftp/python/3.9.4/python-3.9.4-amd64.exe
+	     https://www.python.org/ftp/python/3.10.1/python-3.10.1-amd64.exe
 	* install 3.6.8 to /c/Users/${USER}/mge_whl_python_env/3.6.8
 	* install 3.7.7 to /c/Users/${USER}/mge_whl_python_env/3.7.7
 	* install 3.8.3 to /c/Users/${USER}/mge_whl_python_env/3.8.3
+	* install 3.9.4 to /c/Users/${USER}/mge_whl_python_env/3.9.4
+	* install 3.10.1 to /c/Users/${USER}/mge_whl_python_env/3.10.1
 	* cp python.exe to python3.exe
 		loop cd /c/Users/${USER}/mge_whl_python_env/*
 		copy python.exe to python3.exe

scripts/whl/BUILD_PYTHON_WHL_README.md

@@ -12,16 +12,6 @@
 1: please refer to https://docs.docker.com/engine/security/rootless/ to enable rootless docker env
 2: cd ./scripts/whl/manylinux2014
 3: ./build_image.sh
-4: as aarch64-linux python3.5 pip do not provide megengine depends prebuild binary package, which definition
-    in imperative/python/requires.txt, before install megengine wheel package(just python3.5), you need
-    do follow step to init python3.5 arrow build env:
-        4a: wget https://apache.bintray.com/arrow/debian/apache-arrow-archive-keyring-latest-buster.deb
-        4b: sudo apt install ./apache-arrow-archive-keyring-latest-buster.deb
-        4c: sudo apt update
-        4d: sudo apt install libarrow-dev
-        4c: sudo apt install libarrow-python-dev
-        4e: sudo ln -s /usr/include/locale.h /usr/include/xlocale.h
-
 ```
 
 ## MacOS

scripts/whl/macos/macos_build_whl.sh

@@ -38,7 +38,7 @@ SRC_DIR=$($READLINK -f "`dirname $0`/../../../")
 source ${SRC_DIR}/scripts/whl/utils/utils.sh
 
 ALL_PYTHON=${ALL_PYTHON}
-FULL_PYTHON_VER="3.5.9 3.6.10 3.7.7 3.8.3"
+FULL_PYTHON_VER="3.6.10 3.7.7 3.8.3 3.9.4 3.10.1"
 if [[ -z ${ALL_PYTHON} ]]
 then
     ALL_PYTHON=${FULL_PYTHON_VER}
@@ -70,10 +70,7 @@ function config_python_env() {
     fi
     echo ${ver}
 
-    if [ "$1" = "3.5.9" ]; then
-        PYTHON_INCLUDE_DIR=${PYTHON_DIR}include/python3.5m
-        PYTHON_LIBRARY=${PYTHON_DIR}/lib/libpython3.5m.dylib
-    elif [ "$1" = "3.6.10" ]; then
+    if [ "$1" = "3.6.10" ]; then
         PYTHON_INCLUDE_DIR=${PYTHON_DIR}include/python3.6m
         PYTHON_LIBRARY=${PYTHON_DIR}/lib/libpython3.6m.dylib
     elif [ "$1" = "3.7.7" ]; then
@@ -82,6 +79,12 @@ function config_python_env() {
     elif [ "$1" = "3.8.3" ]; then
         PYTHON_INCLUDE_DIR=${PYTHON_DIR}include/python3.8
         PYTHON_LIBRARY=${PYTHON_DIR}/lib/libpython3.8.dylib
+    elif [ "$1" = "3.9.4" ]; then
+        PYTHON_INCLUDE_DIR=${PYTHON_DIR}include/python3.9
+        PYTHON_LIBRARY=${PYTHON_DIR}/lib/libpython3.9.dylib
+    elif [ "$1" = "3.10.1" ]; then
+        PYTHON_INCLUDE_DIR=${PYTHON_DIR}include/python3.10
+        PYTHON_LIBRARY=${PYTHON_DIR}/lib/libpython3.10.dylib
     else
         echo "ERR: DO NOT SUPPORT PYTHON VERSION"
         echo "now support list: ${FULL_PYTHON_VER}"

scripts/whl/macos/macos_whl_env_prepare.sh

@@ -45,7 +45,7 @@ fi
 SRC_DIR=$($READLINK -f "`dirname $0`/../../../")
 
 echo ${SRC_DIR}
-ALL_PYTHON="3.5.9 3.6.10 3.7.7 3.8.3"
+ALL_PYTHON="3.6.10 3.7.7 3.8.3 3.9.4 3.10.1"
 USER=$(whoami)
 
 function install_python_package() {
@@ -56,7 +56,7 @@ function install_python_package() {
             echo "FOUND install /Users/${USER}/.pyenv/versions/${pak} strip it..."
         else
             os_ver=$(sw_vers -productVersion | awk '{print int($0)}')
-            if [ $a -lt 11 ];then
+            if [ $os_ver -lt 11 ];then
                 env PYTHON_CONFIGURE_OPTS="--enable-shared" pyenv install ${pak}
             else
                 PYTHON_CONFIGURE_OPTS="--enable-shared" \

scripts/whl/manylinux2014/do_build_common.sh

@@ -74,7 +74,7 @@ function patch_elf_depend_lib_megenginelite() {
 SRC_DIR=$(readlink -f "`dirname $0`/../../../")
 source ${SRC_DIR}/scripts/whl/utils/utils.sh
 
-SUPPORT_ALL_VERSION="35m 36m 37m 38"
+SUPPORT_ALL_VERSION="36m 37m 38 39 310"
 ALL_PYTHON=${ALL_PYTHON}
 if [[ -z ${ALL_PYTHON} ]]
 then
@@ -109,15 +109,21 @@ do
         rm -rf ${BUILD_DIR}
     fi
 
-    python_ver=${ver:0:2}
+    python_ver=`echo $ver | tr -d m`
     MAJOR=${python_ver:0:1}
-    MINOR=${ver:1}
-    PYTHON_DIR=/opt/python/cp${python_ver}-cp${ver}/
+    MINOR=${python_ver:1}
+    PYTHON_DIR=/opt/python/cp${python_ver}-cp${ver}
+
+    SUFFIX=
+    if [[ $MINOR -lt 8 ]];then
+        SUFFIX="m"
+    fi
+
     export EXTRA_CMAKE_ARGS="${ORG_EXTRA_CMAKE_FLAG} -DCMAKE_BUILD_TYPE=RelWithDebInfo"
     export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DMGE_WITH_CUSTOM_OP=ON"
     export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DPYTHON_EXECUTABLE=${PYTHON_DIR}/bin/python3"
-    export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DPYTHON_LIBRARY=${PYTHON_DIR}lib/"
-    export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DPYTHON_INCLUDE_DIR=${PYTHON_DIR}include/python${MAJOR}.${MINOR}"
+    export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DPYTHON_LIBRARY=${PYTHON_DIR}/lib/"
+    export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DPYTHON_INCLUDE_DIR=${PYTHON_DIR}/include/python${MAJOR}.${MINOR}${SUFFIX}"
     export EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DMGE_WITH_ATLAS=ON"
 
     if [ -d "${BUILD_DIR}" ]; then

scripts/whl/manylinux2014/init_image.sh

@@ -1,7 +1,6 @@
 #!/bin/bash -e
 
 GET_PIP_URL='https://bootstrap.pypa.io/get-pip.py'
-GET_PIP_URL_35='https://bootstrap.pypa.io/pip/3.5/get-pip.py'
 GET_PIP_URL_36='https://bootstrap.pypa.io/pip/3.6/get-pip.py'
 SWIG_URL='https://codeload.github.com/swig/swig/tar.gz/refs/tags/rel-3.0.12'
 LLVM_URL='https://github.com/llvm-mirror/llvm/archive/release_60.tar.gz' 
@@ -15,35 +14,26 @@ yum install -y pcre-devel devtoolset-9-libatomic-devel.${ARCH}
 yum install -y devtoolset-8 devtoolset-8-libatomic-devel.${ARCH}
 # install a default python3 for cmake PYTHON3_EXECUTABLE_WITHOUT_VERSION
 yum install -y python3 python3-devel
-python3 -m pip install cython
-python3 -m pip install numpy
+python3 -m pip install cython -i https://mirrors.aliyun.com/pypi/simple
+python3 -m pip install numpy -i https://mirrors.aliyun.com/pypi/simple
 
-ALL_PYTHON="35m 36m 37m 38"
-numpy_version="1.18.1"
-if [ ${ARCH} = "aarch64" ];then
-    # numpy do not have 1.18.1 on aarch64 linux, so we use another fix version
-    numpy_version="1.19.5"
-fi
+# FIXME: failed when install pip with python3.10 because python3.10
+# is not installed on aarch64, so we remove 310 from ALL_PYTHON version now
+ALL_PYTHON="36m 37m 38 39"
+numpy_version="1.19.5"
 for ver in ${ALL_PYTHON}
 do
-    python_ver=${ver:0:2}
+    python_ver=`echo $ver | tr -d m`
     PIP_URL=${GET_PIP_URL}
-    if [ ${ver} = "35m" ];then
-        PIP_URL=${GET_PIP_URL_35}
-    else if [ ${ver} = "36m" ];then
+    if [ ${ver} = "36m" ];then
     	PIP_URL=${GET_PIP_URL_36}
-      fi
     fi
     echo "use pip url: ${PIP_URL}"
     curl ${PIP_URL} | /opt/python/cp${python_ver}-cp${ver}/bin/python - \
 	--no-cache-dir --only-binary :all:
-    if [ ${ARCH} = "aarch64" ] && [ ${ver} = "35m" ];then
-        # aarch64 linux python3.5 pip do not provide binary package
-        /opt/python/cp${python_ver}-cp${ver}/bin/pip install --no-cache-dir numpy setuptools==46.1.3
-    else
-        /opt/python/cp${python_ver}-cp${ver}/bin/pip install \
-        --no-cache-dir --only-binary :all: numpy==${numpy_version} setuptools==46.1.3
-    fi
+    /opt/python/cp${python_ver}-cp${ver}/bin/pip install \
+    --no-cache-dir --only-binary :all: numpy==${numpy_version} setuptools==46.1.3 \
+    -i https://mirrors.aliyun.com/pypi/simple
 done
 
 pushd /home >/dev/null

scripts/whl/utils/utils.sh

@@ -76,7 +76,7 @@ function check_build_ninja_python_api() {
         echo "org args: ${ver}"
         if [[ $OS =~ "NT" ]]; then
             INCLUDE_KEYWORD="${ver}\\\\include"
-            PYTHON_API_INCLUDES="3.5.4\\\\include 3.6.8\\\\include 3.7.7\\\\include 3.8.3\\\\include"
+            PYTHON_API_INCLUDES="3.6.8\\\\include 3.7.7\\\\include 3.8.3\\\\include 3.9.4\\\\include 3.10.1\\\\include"
         elif [[ $OS =~ "Linux" ]]; then
             INCLUDE_KEYWORD="include/python3.${ver:1:1}"
             info=`command -v termux-info || true`

scripts/whl/windows/windows_build_whl.sh

@@ -26,7 +26,7 @@ SRC_DIR=$(READLINK -f "`dirname $0`/../../../")
 source ${SRC_DIR}/scripts/whl/utils/utils.sh
 
 ALL_PYTHON=${ALL_PYTHON}
-FULL_PYTHON_VER="3.5.4 3.6.8 3.7.7 3.8.3"
+FULL_PYTHON_VER="3.6.8 3.7.7 3.8.3 3.9.4 3.10.1"
 if [[ -z ${ALL_PYTHON} ]]
 then
     ALL_PYTHON=${FULL_PYTHON_VER}