chore(ci): test with tox (#332)

This PR introduces tox for managing test environments. It allows us to
run the same tests locally and on CI on many different environment
configurations.

- fixes #329
- Fix a new error in `cat` from torch-nightly

Author: justinchuby

.github/workflows/lint.yaml

@@ -50,10 +50,6 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install --upgrade setuptools
           python -m pip install -q -r requirements-dev.txt
-          # Install ONNX
-          python -m pip uninstall -y onnx-function-experiment
-          python -m pip uninstall -y ort-function-experiment-nightly
-          python -m pip install -r requirements-onnx.txt
           # Install packages
           python -m pip install -e .
           lintrunner init

.github/workflows/main.yaml

@@ -10,80 +10,47 @@ on:
   # Allows you to run this workflow manually from the Actions tab
   workflow_dispatch:
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
+  cancel-in-progress: true
+
 jobs:
   test:
     strategy:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, windows-latest, macos-latest]
         name:
-          # - python311-onnx
-          # TODO: Enable python311 when onnx is updated
-          - python310-onnx
-          - python39-onnx
-          - python39-expe
-          - python38-expe
+          # names should match tox labels in tox.ini
+          - python310
+          - python39
+          - python38
+          - python310-experimental
+          - python310-torch-nightly
         include:
-          # - name: python311-onnx
-          #   python-version: "3.11"
-          #   onnx_standard: true
-          #   test_examples: true
-          - name: python310-onnx
+          - name: python310
             python-version: "3.10"
-            onnx_standard: true
-            test_examples: true
-          - name: python39-onnx
-            python-version: "3.9"
-            onnx_standard: true
-            test_examples: false
-          - name: python39-expe
+          - name: python39
             python-version: "3.9"
-            onnx_standard: false
-            test_examples: false
-          - name: python38-expe
+          - name: python38
             python-version: "3.8"
-            onnx_standard: false
-            test_examples: false
+          - name: python310-experimental
+            python-version: "3.10"
+          - name: python310-torch-nightly
+            python-version: "3.10"
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v3
       - name: Setup Python ${{ matrix.python-version }}
         uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python-version }}
-          cache: pip
-          cache-dependency-path: "**/requirements-dev.txt"
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip setuptools wheel build
-          python -m pip install -r requirements-dev.txt
-
-      - name: Install standard onnx
-        if: ${{ matrix.onnx_standard }}
-        run: |
-          python -m pip uninstall -y onnx-function-experiment
-          python -m pip uninstall -y ort-function-experiment-nightly
-          python -m pip install -r requirements-onnx.txt
-
-      - name: Versions
-        run: |
-          pip list | grep numpy
-          pip list | grep onnx
-          pip list | grep torch
-
-      - name: pytest
-        run: pytest -v onnxscript --cov=onnxscript --cov-report=xml -n=auto
-
-      - name: Install package
-        run: pip install .
-
-      - name: Test examples
-        if: ${{ matrix.test_examples }}
-        run: pytest -v docs/test -n=auto
-
-      - name: Build package
-        run: python -m build
-
+      - name: Install tox
+        run: python -m pip install tox tox-gh
+      - name: Setup test suite
+        run: tox -m ${{ matrix.name }} --colored yes -vv --notest
+      - name: Run test suite
+        run: tox -m ${{ matrix.name }} --colored yes --skip-pkg-install -- -v --cov=onnxscript --cov-report=xml -n=auto
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v3
 

azure-pipelines.yml

@@ -1,7 +1,5 @@
 # Python package
-# Create and test a Python package on multiple Python versions.
-# Add steps that analyze code, save the dist with the build record, publish to a PyPI-compatible index, and more:
-# https://docs.microsoft.com/azure/devops/pipelines/languages/python
+# Clone the repository for mandatory injected checks
 
 trigger:
 - main
@@ -12,43 +10,8 @@ strategy:
   matrix:
     Python310-onnx:
       python.version: '3.10'
-      onnx.standard: '1'
-      test.examples: '1'
 steps:
 - task: UsePythonVersion@0
   inputs:
     versionSpec: '$(python.version)'
   displayName: 'Use Python $(python.version)'
-
-- script: |
-    python -m pip install --upgrade pip setuptools wheel
-    python -m pip install -r requirements-dev.txt
-  displayName: 'Install dependencies'
-
-- script: |
-      if [ '$(onnx.standard)' == '1' ]
-        then
-          python -m pip uninstall -y onnx-function-experiment
-          python -m pip uninstall -y ort-function-experiment-nightly
-          python -m pip install -r requirements-onnx.txt
-      fi
-  displayName: 'Install onnx'
-
-- script: |
-    python setup.py sdist
-  displayName: 'make package'
-
-- script: |
-    pytest -v onnxscript --cov=onnxscript --cov-report term-missing
-  displayName: 'pytest'
-
-- script: |
-    python setup.py install
-  displayName: 'install package'
-
-- script: |
-      if [ '$(test.examples)' == '1' ]
-        then
-          pytest -v docs/test
-      fi
-  displayName: 'Test examples'

onnxscript/backend/onnx_export.py

@@ -427,9 +427,6 @@ def export_template(
     Returns:
         python code
     """
-    # delayed import to avoid raising an exception if not installed.
-    import autopep8  # pylint: disable=import-outside-toplevel
-
     # unique_function_domain_version
     unique_function_domain_version = set()
     if hasattr(model_onnx, "functions"):
@@ -521,6 +518,9 @@ def rename_variable(name):
     if "\nreturn" in final:
         raise SyntaxError(f"The produced code is wrong.\n{final}")
     if clean_code:
+        # delayed import to avoid raising an exception if not installed.
+        import autopep8  # pylint: disable=import-outside-toplevel
+
         cleaned_code = autopep8.fix_code(final, options=autopep_options)
         if "\nreturn" in cleaned_code:
             raise SyntaxError(f"The cleaned code is wrong.\n{final}\n------{cleaned_code}")

onnxscript/function_libs/torch_aten/ops/core.py

@@ -735,6 +735,9 @@ def aten_cat(tensors: Sequence[TTensor], dim: int = 0) -> TTensor:
     num_of_input = len(tensors)  # len() function not support yet
     a = op.SequenceEmpty()
     for i in range(num_of_input):
+        if op.Size(tensors[i]) == 0 and dim != 0:
+            # Skip empty tensors when concatenating along non-zero dimension
+            continue
         a = op.SequenceInsert(a, tensors[i])
     return op.ConcatFromSequence(a, axis=dim)
 

onnxscript/test/function_libs/torch_aten/ops_correctness_test.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 import onnx
+import packaging.version
 import torch
 from torch.testing._internal import common_device_type, common_methods_invocations
 from torch.testing._internal.opinfo import core as opinfo_core
@@ -148,9 +149,15 @@ def wrapped(fn):
 def duplicate_opinfo(opinfos: list[opinfo_core.OpInfo], name: str, new_names: tuple[str, ...]):
     """Duplicate an opinfo in the opinfo database and give it a new name."""
     duplicated = []
+    all_info_names = {opinfo.name for opinfo in opinfos}
     for opinfo in opinfos:
         if opinfo.name == name:
             for new_name in new_names:
+                if new_name in all_info_names:
+                    # NOTE: Avoid duplicating an opinfo that already exists in the database.
+                    # New opinfos are expected to be added in torch-nightly.
+                    warnings.warn(f"OpInfo {new_name} already exists in the database.")
+                    continue
                 new_opinfo = copy.deepcopy(opinfo)
                 new_opinfo.name = new_name
                 duplicated.append(new_opinfo)
@@ -187,7 +194,8 @@ def _full_input_wrangler(
     args: list[Any], kwargs: dict[str, Any]
 ) -> tuple[list[Any], dict[str, Any]]:
     # Remove the self argument
-    args.pop(0)
+    if packaging.version.parse(torch.__version__) <= packaging.version.parse("1.13.1"):
+        args.pop(0)
     return args, kwargs
 
 

pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
   "Programming Language :: Python :: 3.11",
   "License :: OSI Approved :: Apache Software License",
 ]
-dependencies = ["numpy", "protobuf<4", "onnx"]
+dependencies = ["numpy", "protobuf<4", "onnx", "typing_extensions"]
 
 [tool.setuptools.dynamic]
 version = {attr = "onnxscript.__version__"}

requirements-dev.txt

@@ -1,4 +1,3 @@
-autopep8
 click
 numpy==1.23.5
 matplotlib

requirements-onnx.txt

@@ -0,0 +1,96 @@
+[tox]
+requires =
+    tox>=4
+env_list =
+    py{37,38,39,310}
+    test-onnx-preview
+    examples
+    build
+labels =
+    python310 = py310, build
+    python39 = py39, build
+    python38 = py38, build
+    python37 = py37, build
+    python310-experimental = test-onnx-preview
+    python310-torch-nightly = test-torch-nightly
+
+[base]
+deps =
+    autopep8
+    click
+    jinja2
+    numpy==1.23.5
+    protobuf<4
+    typing_extensions
+
+    ; ATen lib
+    beartype
+    types-PyYAML
+
+    ; Testing
+    expecttest
+    packaging
+    parameterized
+    pytest-cov
+    pytest-subtests
+    pytest-xdist
+    pytest!=7.1.0
+    pyyaml
+commands =
+    pip list
+    pytest onnxscript {posargs}
+
+[testenv]
+description = run unit tests
+deps =
+    {[base]deps}
+    onnx==1.13
+    onnxruntime
+    torch==1.13
+commands =
+    {[base]commands}
+    pytest docs/test
+
+[testenv:test-onnx-preview]
+description = test with onnx experimental builds
+deps=
+    {[base]deps}
+    ort-function-experiment-nightly
+    torch==1.13
+pip_pre = true
+install_command = python -I -m pip install -f https://onnxruntimepackages.z14.web.core.windows.net/onnxruntime-function-experiment.html {opts} {packages}
+commands =
+    ; Uninstall onnx here, because the released onnx is required and installed when onnxscript is installed
+    python -I -m pip uninstall -y onnx
+    python -I -m pip install -f https://onnxruntimepackages.z14.web.core.windows.net/onnx-function-experiment.html --pre onnx-function-experiment
+    pip list
+    pytest onnxscript {posargs}
+    pytest docs/test
+
+[testenv:test-torch-nightly]
+description = test with pytorch nightly
+deps=
+    {[base]deps}
+    onnx==1.13
+    onnxruntime
+    torch
+pip_pre = true
+set_env =
+    PIP_EXTRA_INDEX_URL = {env:PIP_EXTRA_INDEX_URL:https://download.pytorch.org/whl/nightly/cpu}
+commands =
+    {[base]commands}
+
+[testenv:build]
+description = build package
+skip_install = true
+deps =
+    build
+    wheel
+commands =
+    python -m build
+
+[gh]
+python =
+    3.8: py38, build
+    3.9: py39, build
+    3.10: py310, build, test-onnx-preview, test-torch-nightly