[test] ci py 3.11 tests (#1099)

* [test] ci py 3.11 tests
Co-authored-by: Min Xu <min.xu.public@gmail.com>

* fixed setup.py

* fixed ci config

* fixed ci config's python 3.11 version

* fixed torch installs on cpu

* update pygit2 for 3.11

* we don't run benchmark on cpu, so no need to install the benchmark reqs

* update torch install

* try to install torchvision

* numpy version 311

* fix cpu test dependency installation

* pip git install cmd fix

* bypass some tests in 3.11. failure due to packages they use haven't been updated for 3.11 yet
Co-authored-by: Min Xu <min.xu.public@gmail.com>

.circleci/config.yml

@@ -38,6 +38,11 @@ cpu_py310: &cpu_py310
     - image: circleci/python:3.10
   resource_class: large
 
+cpu_py311: &cpu_py311
+  docker:
+    - image: cimg/python:3.11.0
+  resource_class: large
+
 # Here is the list of GPU images:
 # https://circleci.com/docs/2.0/configuration-reference/#available-linux-gpu-images
 # We need to use multiple gpus for several jobs. The resource_class
@@ -97,6 +102,22 @@ install_dep_pytorch_lts: &install_dep_pytorch_lts
         python -m torch.utils.collect_env
         wget -O /home/circleci/venv/check_version.py https://raw.githubusercontent.com/min-xu-ai/check_verion/main/check_version.py
 
+# most recent stable version on cpu
+# for testing 3.8/9/10/11, make sure torch has the right whl files.
+install_dep_pytorch_stable_cpu: &install_dep_pytorch_stable_cpu
+  - run:
+      name: Install Dependencies with torch 1.13.0 on CPU
+      command: |
+        # check if we have restored venv cache (/home/circleci/venv) correctly, if so, just skip
+        if [ -f /home/circleci/venv/check_version.py ]; then python /home/circleci/venv/check_version.py torch eq 1.13 && exit 0; fi
+        # start installing
+        pip install --progress-bar off torch torchvision
+        pip install --progress-bar off -r requirements-dev.txt
+        python -c 'import torch; print("Torch version:", torch.__version__)'
+        python -c 'import torch; assert torch.__version__.split(".")[:2] == ["1", "13"], f"wrong torch version {torch.__version__}"'
+        python -m torch.utils.collect_env
+        wget -O /home/circleci/venv/check_version.py https://raw.githubusercontent.com/min-xu-ai/check_verion/main/check_version.py
+
 # most recent stable version
 install_dep_pytorch_stable: &install_dep_pytorch_stable
   - run:
@@ -105,7 +126,8 @@ install_dep_pytorch_stable: &install_dep_pytorch_stable
         # check if we have restored venv cache (/home/circleci/venv) correctly, if so, just skip
         if [ -f /home/circleci/venv/check_version.py ]; then python /home/circleci/venv/check_version.py torch eq 1.13 && exit 0; fi
         # start installing
-        pip install --progress-bar off torch==1.13.0 torchvision==0.14.0 --extra-index-url https://download.pytorch.org/whl/cu113
+        pip install --progress-bar off torch==1.13.0 --extra-index-url https://download.pytorch.org/whl/cu117
+        pip install --progress-bar off torchvision==0.14.0 --extra-index-url https://download.pytorch.org/whl/cu117
         pip install --progress-bar off -r requirements-dev.txt
         pip install --progress-bar off -r requirements-benchmarks.txt
         python -c 'import torch; print("Torch version:", torch.__version__)'
@@ -282,13 +304,13 @@ jobs:
       # Cache the venv directory that contains dependencies
       - restore_cache:
           keys:
-            - cache-key-cpu-py38-torch-1-11-0-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
-      - <<: *install_dep_pytorch_stable
+            - cache-key-cpu-py38-torch-stable-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
+      - <<: *install_dep_pytorch_stable_cpu
 
       - save_cache:
           paths:
             - ~/venv
-          key: cache-key-cpu-py38-torch-1-11-0-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
+          key: cache-key-cpu-py38-torch-stable-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
 
       - <<: *run_unittests
       - <<: *run_doc_build
@@ -313,14 +335,14 @@ jobs:
       # Cache the venv directory that contains dependencies
       - restore_cache:
           keys:
-            - cache-key-cpu-py39-torch-1-11-0-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
+            - cache-key-cpu-py39-torch-stable-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
 
-      - <<: *install_dep_pytorch_stable
+      - <<: *install_dep_pytorch_stable_cpu
 
       - save_cache:
           paths:
             - ~/venv
-          key: cache-key-cpu-py39-torch-1-11-0-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
+          key: cache-key-cpu-py39-torch-stable-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
 
       - <<: *run_unittests
       - <<: *run_doc_build
@@ -345,14 +367,51 @@ jobs:
       # Cache the venv directory that contains dependencies
       - restore_cache:
           keys:
-            - cache-key-cpu-py310-torch-1-11-0-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
+            - cache-key-cpu-py310-torch-stable-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
 
-      - <<: *install_dep_pytorch_stable
+      - <<: *install_dep_pytorch_stable_cpu
+
+      - save_cache:
+          paths:
+            - ~/venv
+          key: cache-key-cpu-py310-torch-stable-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
+
+      - <<: *run_unittests
+      - <<: *run_doc_build
+
+      - store_test_results:
+          path: test-results
+
+  cpu_tests_py311:
+    <<: *cpu_py311
+
+    working_directory: ~/fairscale
+
+    steps:
+      - checkout
+      - <<: *check_test_list
+      - <<: *setup_venv
+
+      # Do this first to test repo dependencies. Basic import should work after this
+      # installation. See issue #1042 for an example.
+      - <<: *install_repo
+
+      # Cache the venv directory that contains dependencies
+      - restore_cache:
+          keys:
+            - cache-key-cpu-py311-torch-stable-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
+
+      - <<: *install_dep_pytorch_stable_cpu
+
+      # torchvision doesn't have cp311 binary whl yet and it is needed by some tests.
+      - run: pip install git+https://github.com/pytorch/vision@677fc939b21a8893f07db4c1f90482b648b6573f
+      # numpy for 3.11 needs a newer binary whl.
+      - run: pip install numpy==1.23.5
 
       - save_cache:
           paths:
             - ~/venv
-          key: cache-key-cpu-py310-torch-1-11-0-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
+          key: cache-key-cpu-py311-torch-stable-{{.Environment.CACHE_VERSION}}-{{checksum "setup.py"}}-{{checksum "requirements-dev.txt"}}
 
       - <<: *run_unittests
       - <<: *run_doc_build
@@ -591,6 +650,7 @@ workflows:
       - cpu_tests_py38
       - cpu_tests_py39
       - cpu_tests_py310
+      - cpu_tests_py311
       - gpu_tests_lts:
           test_list_file: tests/ci_test_list_1.txt
       - gpu_tests_stable:

requirements-dev.txt

@@ -36,5 +36,5 @@ scikit-learn == 1.1.3
 # For weigit. These are actually user requirements, not developer requirements.
 # However, due to the experimental nature of weigit, we don't expose to the
 # general users of fairscale yet. We check for them in weigit's init code.
-pygit2==1.9.2
+pygit2==1.11.1
 pgzip==0.3.1

setup.py

@@ -32,10 +32,12 @@ def find_version(version_file_path) -> str:
 
 extensions = []
 cmdclass = {}
+setup_requires = []
 
 if os.getenv("BUILD_CUDA_EXTENSIONS", "0") == "1":
     from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
+    setup_requires = ["ninja"]
     extensions.extend(
         [
             CUDAExtension(
@@ -58,16 +60,20 @@ if __name__ == "__main__":
         name="fairscale",
         description="FairScale: A PyTorch library for large-scale and high-performance training.",
         version=find_version("fairscale/version.py"),
-        setup_requires=["ninja"],  # ninja is required to build extensions
+        setup_requires=setup_requires,
         install_requires=fetch_requirements(),
         include_package_data=True,
-        packages=setuptools.find_packages(exclude=("tests", "tests.*")),
+        packages=setuptools.find_packages(include=["fairscale*"]),  # Only include code within fairscale.
         ext_modules=extensions,
         cmdclass=cmdclass,
         python_requires=">=3.8",
         author="Foundational AI Research @ Meta AI",
-        author_email="todo@fb.com",
-        long_description="FairScale is a PyTorch extension library for high performance and large scale training on one or multiple machines/nodes. This library extends basic PyTorch capabilities while adding new experimental ones.",
+        author_email="todo@meta.com",
+        long_description=(
+            "FairScale is a PyTorch extension library for high performance and "
+            "large scale training on one or multiple machines/nodes. This library "
+            "extends basic PyTorch capabilities while adding new experimental ones."
+        ),
         long_description_content_type="text/markdown",
         entry_points={"console_scripts": ["wgit = fairscale.experimental.wgit.__main__:main"]},
         classifiers=[

tests/experimental/nn/test_auto_shard.py

@@ -8,6 +8,7 @@ Testing Auto Shard functionality of non nn.Sequential models.
 """
 
 import math
+import sys
 
 import pytest
 import torch
@@ -78,6 +79,8 @@ dropout = 0.2  # the dropout value
 
 
 def test_single_run():
+    if sys.version_info.major == 3 and sys.version_info.minor > 10:
+        pytest.skip("torch.fx doesn't seem to work 3.11 yet")
     if torch_version() < (1, 8, 0):
         pytest.skip("requires torch version >= 1.8.0")
     from fairscale.experimental.nn.auto_shard import shard_model

tests/experimental/wgit/test_sha1_store.py

@@ -6,6 +6,7 @@
 import os
 from pathlib import Path
 import shutil
+import sys
 
 import pytest
 import torch
@@ -118,6 +119,9 @@ def test_sha1_add_tensor(sha1_store, compress):
 @pytest.mark.parametrize("compress", [True, False])
 def test_sha1_get(sha1_store, compress):
     """Testing the get() API: normal and exception cases."""
+    if sys.version_info.major == 3 and sys.version_info.minor > 10:
+        pytest.skip("pgzip package doesn't work with 3.11's gzip module")
+
     os.chdir(TESTING_STORE_DIR)
 
     # Add a file, a state dict and a tensor.