[primer] Parallelize primer workflow into package batches

Author: jacobtylerwalls

.github/workflows/primer_comment.yaml

@@ -103,8 +103,9 @@ jobs:
           . venv/bin/activate
           python tests/primer/__main__.py compare \
           --commit=${{ github.event.workflow_run.head_sha }} \
-          --base-file=output_${{ steps.python.outputs.python-version }}_main.txt \
-          --new-file=output_${{ steps.python.outputs.python-version }}_pr.txt
+          --base-file=output_${{ steps.python.outputs.python-version }}_main_BATCHIDX.txt \
+          --new-file=output_${{ steps.python.outputs.python-version }}_pr_BATCHIDX.txt \
+          --batches=4
       - name: Post comment
         id: post-comment
         uses: actions/[email protected]

.github/workflows/primer_run_main.yaml

@@ -24,12 +24,14 @@ permissions:
 
 jobs:
   run-primer:
-    name: Run / ${{ matrix.python-version }}
+    name: Run / ${{ matrix.python-version }} / batch index ${{ matrix.batchIdx }}
     runs-on: ubuntu-latest
-    timeout-minutes: 120
+    timeout-minutes: 45
     strategy:
       matrix:
         python-version: ["3.8", "3.11"]
+        batches: [4]
+        batchIdx: [0, 1, 2, 3]
     steps:
       - name: Check out code from GitHub
         uses: actions/[email protected]
@@ -91,7 +93,7 @@ jobs:
         run: |
           . venv/bin/activate
           pip install -e .
-          python tests/primer/__main__.py run --type=main 2>warnings.txt
+          python tests/primer/__main__.py run --type=main --batches=${{ matrix.batches }} --batchIdx=${{ matrix.batchIdx }} 2>warnings.txt
           WARNINGS=$(head -c 65000 < warnings.txt)
           if [[ $WARNINGS ]]
           then echo "::warning ::$WARNINGS"
@@ -102,4 +104,4 @@ jobs:
           name: primer_output
           path: >-
             tests/.pylint_primer_tests/output_${{ steps.python.outputs.python-version
-            }}_main.txt
+            }}_main_batch${{ matrix.batchIdx }}.txt

.github/workflows/primer_run_pr.yaml

@@ -33,12 +33,14 @@ permissions:
 
 jobs:
   run-primer:
-    name: Run / ${{ matrix.python-version }}
+    name: Run / ${{ matrix.python-version }} / batch index ${{ matrix.batchIdx }}
     runs-on: ubuntu-latest
-    timeout-minutes: 120
+    timeout-minutes: 45
     strategy:
       matrix:
         python-version: ["3.8", "3.11"]
+        batches: [4]
+        batchIdx: [0, 1, 2, 3]
     steps:
       - name: Check out code from GitHub
         uses: actions/[email protected]
@@ -156,7 +158,7 @@ jobs:
         run: |
           . venv/bin/activate
           pip install -e .
-          python tests/primer/__main__.py run --type=pr 2>warnings.txt
+          python tests/primer/__main__.py run --type=pr --batches=${{ matrix.batches }} --batchIdx=${{ matrix.batchIdx }} 2>warnings.txt
           WARNINGS=$(head -c 65000 < warnings.txt)
           if [[ $WARNINGS ]]
           then echo "::warning ::$WARNINGS"
@@ -167,12 +169,12 @@ jobs:
           name: primer_output_pr
           path:
             tests/.pylint_primer_tests/output_${{ steps.python.outputs.python-version
-            }}_pr.txt
+            }}_pr_batch${{ matrix.batchIdx }}.txt
       - name: Upload output of 'main'
         uses: actions/[email protected]
         with:
           name: primer_output_main
-          path: output_${{ steps.python.outputs.python-version }}_main.txt
+          path: output_${{ steps.python.outputs.python-version }}_main_batch${{ matrix.batchIdx }}.txt
 
       # Save PR number so we know which PR to comment on
       - name: Save PR number

pylint/testutils/_primer/primer.py

@@ -56,6 +56,18 @@ def __init__(self, primer_directory: Path, json_path: Path) -> None:
         run_parser.add_argument(
             "--type", choices=["main", "pr"], required=True, help="Type of primer run."
         )
+        run_parser.add_argument(
+            "--batches",
+            required=False,
+            type=int,
+            help="Number of batches",
+        )
+        run_parser.add_argument(
+            "--batchIdx",
+            required=False,
+            type=int,
+            help="Portion of primer packages to run.",
+        )
 
         # All arguments for the compare parser
         compare_parser = self._subparsers.add_parser("compare")
@@ -74,6 +86,12 @@ def __init__(self, primer_directory: Path, json_path: Path) -> None:
             required=True,
             help="Commit hash of the PR commit being checked.",
         )
+        compare_parser.add_argument(
+            "--batches",
+            required=False,
+            type=int,
+            help="Number of batches (filepaths with the placeholder BATCHIDX will be numbered)",
+        )
 
         # Storing arguments
         self.config = self._argument_parser.parse_args()

pylint/testutils/_primer/primer_compare_command.py

@@ -18,8 +18,20 @@
 
 class CompareCommand(PrimerCommand):
     def run(self) -> None:
-        main_data = self._load_json(self.config.base_file)
-        pr_data = self._load_json(self.config.new_file)
+        if self.config.batches is None:
+            main_data = self._load_json(self.config.base_file)
+            pr_data = self._load_json(self.config.new_file)
+        else:
+            main_data = {}
+            pr_data = {}
+            for idx in range(self.config.batches):
+                main_data.update(
+                    self._load_json(self.config.base_file.replace("BATCHIDX", "batch" + str(idx)))
+                )
+                pr_data.update(
+                    self._load_json(self.config.new_file.replace("BATCHIDX", "batch" + str(idx)))
+                )
+
         missing_messages_data, new_messages_data = self._cross_reference(
             main_data, pr_data
         )

pylint/testutils/_primer/primer_run_command.py

@@ -30,14 +30,19 @@ class RunCommand(PrimerCommand):
     def run(self) -> None:
         packages: PackageMessages = {}
         fatal_msgs: list[Message] = []
-        for package, data in self.packages.items():
+        package_data_iter = (
+            self.packages.items()
+            if self.config.batches is None
+            else list(self.packages.items())[self.config.batchIdx:: self.config.batches]
+        )
+        for package, data in package_data_iter:
             messages, p_fatal_msgs = self._lint_package(package, data)
             fatal_msgs += p_fatal_msgs
             local_commit = Repo(data.clone_directory).head.object.hexsha
             packages[package] = PackageData(commit=local_commit, messages=messages)
         path = (
             self.primer_directory
-            / f"output_{'.'.join(str(i) for i in sys.version_info[:3])}_{self.config.type}.txt"
+            / f"output_{'.'.join(str(i) for i in sys.version_info[:3])}_{self.config.type}_batch{self.config.batchIdx}.txt"
         )
         print(f"Writing result in {path}")
         with open(path, "w", encoding="utf-8") as f: