Add CI to check with onnxruntime-traning and transformers

.github/workflows/main.yaml

@@ -106,6 +106,55 @@ jobs:
           name: IR profiling results
           path: tests/ir/serde_test_profiles
 
+  dort:
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest]
+        transformers: ["4.37.2", "4.41.2"]
+        torch: ["", "nightly"]
+        python_version: ["3.11"]
+        name:
+          - dort
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v4
+      - name: Setup Python ${{ matrix.python_version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python_version }}
+      - name: Pull Test Data
+        run: git lfs pull
+      - run: |
+          python -m pip install --upgrade pip setuptools wheel pytest onnx parameterized ml-dtypes pytest-coverage hypothesis expecttest==0.1.6 beartype!=0.16.0 pytest-cov pytest-randomly pytest-subtests pytest-xdist pyyaml
+        name: Install requirements.
+      - run: |
+          python -m pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu
+        if: matrix.torch == 'nightly'
+        name: Install torch nightly
+      - run: |
+          python -m pip install torch torchvision torchaudio
+        if: matrix.torch != 'nightly'
+        name: Install torch
+      - run: |
+          python -m pip install onnxruntime-training==1.17.1
+        name: Install onnxruntime-training
+      - run: |
+          python -m pip install transformers==${{ matrix.transformers }}
+        name: Install transformers
+      - run: |
+          python -m pip install .
+        name: Install onnxscript
+      - run: |
+          python -m pip freeze
+        name: pip freeze
+      - run: |
+          pytest onnxscript
+        name: Run onnxscript tests
+      - run: |
+          pytest tests
+        name: Run tests
+
   build_docs:
     strategy:
       fail-fast: false

onnxscript/tools/benchmark/export_model_test.py

@@ -62,6 +62,11 @@ def test_export_model_llama_cpu_eager(self):
 
     @unittest.skipIf(not has_transformers(), reason="transformers missing")
     @unittest.skipIf(not is_onnxruntime_training(), reason="onnxruntime-training is needed")
+    @unittest.skipIf(
+        torch_older_than("2.4"),
+        reason="TypeError: _functionalize_sync(): "
+        "argument 't' (position 1) must be Tensor, not NoneType",
+    )
     def test_export_model_phi_cpu_dynamo(self):
         args = [
             "--verbose",

tests/common/testutils.py

@@ -10,6 +10,7 @@
 import numpy as np
 import onnx
 import onnxruntime
+import torch
 
 from onnxscript import optimizer
 from onnxscript._legacy_ir import visitor
@@ -29,7 +30,7 @@ def skip_if_no_cuda(reason: str):
     def skip_dec(func):
         @functools.wraps(func)
         def wrapper(self, *args, **kwargs):
-            if not onnxruntime.get_device() == "GPU":
+            if not torch.cuda.is_available() or not onnxruntime.get_device() == "GPU":
                 raise unittest.SkipTest(f"GPU is not available. {reason}")
             return func(self, *args, **kwargs)
 

tests/ir/serde_roundtrip_test.py

@@ -1,5 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
+# pylint: disable=import-outside-toplevel
 from __future__ import annotations
 
 import pathlib
@@ -8,7 +9,6 @@
 import onnx
 import onnx.backend.test
 import parameterized
-import pyinstrument
 
 import onnxscript.testing
 from onnxscript import ir
@@ -26,15 +26,23 @@
 
 class SerdeTest(unittest.TestCase):
     def setUp(self) -> None:
-        self.profiler = pyinstrument.Profiler()
+        try:
+            import pyinstrument
+
+            self.profiler = pyinstrument.Profiler()
+        except ImportError:
+            self.profiler = None
 
     def tearDown(self) -> None:
-        self.profiler.reset()
+        if self.profiler:
+            self.profiler.reset()
 
     @parameterized.parameterized.expand(test_args)
     def test_serialization_deserialization_produces_same_model(
         self, _: str, model_path: pathlib.Path
     ) -> None:
+        if not self.profiler:
+            raise unittest.SkipTest("pyinstrument is not installed.")
         model = onnx.load(model_path)
         # Fix the missing graph name of some test models
         model.graph.name = "main_graph"