From 49f17ccc9b79544de84b8f368c329efe31201230 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Mon, 29 Apr 2024 15:57:33 -0400
Subject: [PATCH] ci(test): split each pytest job into 6 separated jobs (#3716)

Split each pytest job into 6 separated jobs in GHA.

Currently, each separated job takes:
- ~2.5 min for installation, which may need to be optimized in the
future;
- ~2 min for setup, which may be a defect and should be resolved in
#3715;
-  4~5 min for tests (expected).

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 .github/workflows/test_python.yml         | 55 +++++++++++++++++++----
 pyproject.toml                            |  1 +
 source/tests/tf/test_nvnmd_entrypoints.py | 13 +++---
 3 files changed, 54 insertions(+), 15 deletions(-)

diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml
index f499691683..91f1311c20 100644
--- a/.github/workflows/test_python.yml
+++ b/.github/workflows/test_python.yml
@@ -13,14 +13,10 @@ jobs:
     name: Test Python
     runs-on: ubuntu-22.04
     strategy:
+      fail-fast: false
       matrix:
-        include:
-          - python: 3.8
-            tf:
-            torch:
-          - python: "3.11"
-            tf:
-            torch:
+        group: [1, 2, 3, 4, 5, 6]
+        python: ["3.8", "3.11"]
 
     steps:
     - uses: actions/checkout@v4
@@ -45,15 +41,56 @@ jobs:
         HOROVOD_WITHOUT_PYTORCH: 1
         HOROVOD_WITHOUT_GLOO: 1
     - run: dp --version
-    - run: pytest --cov=deepmd source/tests --durations=0
+    - name: Get durations from cache
+      uses: actions/cache@v4
+      with:
+        path: test_durations
+        # the key must never match, even when restarting workflows, as that
+        # will cause durations to get out of sync between groups, the
+        # combined durations will be loaded if available
+        key: test-durations-split-${{ github.run_id }}-${{ github.run_number}}-${{ matrix.python }}-${{ matrix.group }}
+        restore-keys: |
+          test-durations-combined-${{ matrix.python }}-${{ github.sha }}
+          test-durations-combined-${{ matrix.python }}
+    - run: pytest --cov=deepmd source/tests --durations=0 --splits 6 --group ${{ matrix.group }} --store-durations --durations-path=.test_durations_${{ matrix.group }} --splitting-algorithm least_duration
       env:
         NUM_WORKERS: 0
+    - name: Upload partial durations
+      uses: actions/upload-artifact@v4
+      with:
+        name: split-${{ matrix.python }}-${{ matrix.group }}
+        path: .test_durations_${{ matrix.group }}
     - uses: codecov/codecov-action@v4
       env:
         CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+  update_durations:
+    name: Combine and update integration test durations
+    runs-on: ubuntu-22.04
+    strategy:
+      fail-fast: false
+      matrix:
+        python: ["3.8", "3.11"]
+    needs: testpython
+    steps:
+    - name: Get durations from cache
+      uses: actions/cache@v4
+      with:
+        path: .test_durations
+        # key won't match during the first run for the given commit, but
+        # restore-key will if there's a previous stored durations file,
+        # so cache will both be loaded and stored
+        key: test-durations-combined-${{ matrix.python }}-${{ github.sha }}
+        restore-keys: test-durations-combined-${{ matrix.python }}
+    - name: Download artifacts
+      uses: actions/download-artifact@v4
+      with:
+        pattern: split-${{ matrix.python }}-*
+        merge-multiple: true
+    - name: Combine test durations
+      run: jq '. + input' .test_durations_* > .test_durations
   pass:
     name: Pass testing Python
-    needs: [testpython]
+    needs: [testpython, update_durations]
     runs-on: ubuntu-latest
     if: always()
     steps:
diff --git a/pyproject.toml b/pyproject.toml
index edbc5e0b19..23d42e73d2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -73,6 +73,7 @@ test = [
     "pytest",
     "pytest-cov",
     "pytest-sugar",
+    "pytest-split",
     "dpgui",
     "mendeleev",
 ]
diff --git a/source/tests/tf/test_nvnmd_entrypoints.py b/source/tests/tf/test_nvnmd_entrypoints.py
index cc7a92c032..bf4b2288c0 100644
--- a/source/tests/tf/test_nvnmd_entrypoints.py
+++ b/source/tests/tf/test_nvnmd_entrypoints.py
@@ -391,8 +391,6 @@ def test_mapt_cnn_v0(self):
             2.24079704,
         ]
         np.testing.assert_almost_equal(pred, ref_dout, 8)
-        # close
-        nvnmd_cfg.enable = False
 
     @pytest.mark.run(order=1)
     def test_model_qnn_v0(self):
@@ -510,6 +508,8 @@ def test_model_qnn_v0(self):
         pred = valuedic["o_energy"]
         ref_dout = -62.60181403
         np.testing.assert_almost_equal(pred, ref_dout, 8)
+
+    def tearDown(self):
         # close
         nvnmd_cfg.enable = False
 
@@ -700,8 +700,6 @@ def test_mapt_cnn_v1(self):
             3.24079514,
         ]
         np.testing.assert_almost_equal(pred, ref_dout, 8)
-        # close
-        nvnmd_cfg.enable = False
 
     @pytest.mark.run(order=1)
     def test_model_qnn_v1(self):
@@ -846,11 +844,12 @@ def test_model_qnn_v1(self):
             pred = d2[key]
             ref_dout = d1[key]
             np.testing.assert_almost_equal(pred, ref_dout, 8)
-        # close
-        nvnmd_cfg.enable = False
 
     @pytest.mark.run(order=2)
     def test_wrap_qnn_v1(self):
+        # without calling test_mapt_cnn_v1, this test will fail when running individually
+        self.test_mapt_cnn_v1()
+
         tf.reset_default_graph()
         jdata = {}
         jdata["nvnmd_config"] = str(tests_path / "nvnmd" / "ref" / "config_v1_cnn.npy")
@@ -865,6 +864,8 @@ def test_wrap_qnn_v1(self):
         pred = [data[i] for i in idx]
         red_dout = [249, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 254, 95, 24, 176]
         np.testing.assert_equal(pred, red_dout)
+
+    def tearDown(self):
         # close
         nvnmd_cfg.enable = False