From 12f0be025f72f1ed477907275fc3a9d6c2590d0d Mon Sep 17 00:00:00 2001
From: Shantanu <12621235+hauntsaninja@users.noreply.github.com>
Date: Mon, 28 Dec 2020 04:17:39 -0600
Subject: [PATCH] mypy_primer: shard across two CI jobs (#4865)

Co-authored-by: hauntsaninja <>
---
 .github/workflows/mypy_primer.yml         | 17 ++++++++--
 .github/workflows/mypy_primer_comment.yml | 41 +++++++++++++----------
 2 files changed, 39 insertions(+), 19 deletions(-)

diff --git a/.github/workflows/mypy_primer.yml b/.github/workflows/mypy_primer.yml
index c9356c9cb..8fb950948 100644
--- a/.github/workflows/mypy_primer.yml
+++ b/.github/workflows/mypy_primer.yml
@@ -8,6 +8,10 @@ jobs:
   mypy_primer:
     name: Run
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        shard-index: [0, 1]
+      fail-fast: false
     steps:
       - uses: actions/checkout@v2
         with:
@@ -32,9 +36,18 @@ jobs:
           echo ''
           cd ..
           # fail action if exit code isn't zero or one
-          ( mypy_primer --new 0.790 --old 0.790 --custom-typeshed-repo typeshed_to_test --new-typeshed $GITHUB_SHA --old-typeshed upstream_master -o concise | tee diff.txt ) || [ $? -eq 1 ]
+          (
+            mypy_primer \
+            --new 0.790 --old 0.790 \
+            --custom-typeshed-repo typeshed_to_test \
+            --new-typeshed $GITHUB_SHA --old-typeshed upstream_master \
+            --num-shards 2 --shard-index ${{ matrix.shard-index }} \
+            --debug \
+            --output concise \
+            | tee diff.txt
+          ) || [ $? -eq 1 ]
       - name: Upload mypy_primer diff
         uses: actions/upload-artifact@v2
         with:
-          name: mypy_primer_diff
+          name: mypy_primer_diff_${{ matrix.shard-index }}
           path: diff.txt
diff --git a/.github/workflows/mypy_primer_comment.yml b/.github/workflows/mypy_primer_comment.yml
index 63299e6d3..15d259d91 100644
--- a/.github/workflows/mypy_primer_comment.yml
+++ b/.github/workflows/mypy_primer_comment.yml
@@ -48,38 +48,45 @@ jobs:
               return undefined
             }
 
-            const end_time = Number(new Date()) + 60 * 60 * 1000;
-            let primer_run = await check_mypy_primer();
+            const end_time = Number(new Date()) + 60 * 60 * 1000
+            let primer_run = await check_mypy_primer()
             while (!primer_run || primer_run.status != "completed") {
               if (Number(new Date()) > end_time) {
-                throw Error("Timed out waiting for mypy_primer");
+                throw Error("Timed out waiting for mypy_primer")
               }
               console.log("Waiting for mypy_primer to complete...")
-              await new Promise(r => setTimeout(r, 10000));
-              primer_run = await check_mypy_primer();
+              await new Promise(r => setTimeout(r, 10000))
+              primer_run = await check_mypy_primer()
             }
             console.log("Found mypy_primer run!")
             console.log(primer_run)
 
-            // Download artifact from the run
+            // Download artifact(s) from the run
             const artifacts = await github.actions.listWorkflowRunArtifacts({
                 owner: context.repo.owner,
                 repo: context.repo.repo,
                 run_id: primer_run.id,
             })
-            const artifact = artifacts.data.artifacts.find(a => a.name == "mypy_primer_diff")
-            console.log("Artifact from mypy_primer:")
-            console.log(artifact)
+            const filtered_artifacts = artifacts.data.artifacts.filter(
+              a => a.name.startsWith("mypy_primer_diff")
+            )
+            console.log("Artifacts from mypy_primer:")
+            console.log(filtered_artifacts)
 
-            const zip = await github.actions.downloadArtifact({
-                owner: context.repo.owner,
-                repo: context.repo.repo,
-                artifact_id: artifact.id,
-                archive_format: "zip",
-            })
-            const adm = new AdmZip(Buffer.from(zip.data))
+            async function get_artifact_data(artifact) {
+              const zip = await github.actions.downloadArtifact({
+                  owner: context.repo.owner,
+                  repo: context.repo.repo,
+                  artifact_id: artifact.id,
+                  archive_format: "zip",
+              })
+              const adm = new AdmZip(Buffer.from(zip.data))
+              return adm.readAsText(adm.getEntry("diff.txt"))
+            }
+
+            const all_data = await Promise.all(filtered_artifacts.map(get_artifact_data))
+            const data = all_data.join("\n")
 
-            const data = adm.readAsText(adm.getEntry("diff.txt"));
             console.log("Diff from mypy_primer:")
             console.log(data)
             try {