Repository: geldata/gel
Pull request: Run test workflow in shards
This PR adds 3 new options to `edb test`:

* `edb test --shard 2/5` separates all collected tests into 5 shards, and run the 2nd shard only. The split is supposed to be deterministic.
* `edb test --time-file` generates a CSV file of time spent on each test. If the file exists, it'll be taken as an input for `--shard` to generate a more even split, and the CSV file will be updated after the run with new average running times.
* `edb test --list` outputs a list of all selected tests and exits.

The pipeline will fetch from a Gist for `--time-file`, and update that Gist if the tests pass in the master branch.

- [x] Replace the Gist URL with a new one by @edgedb-ci 
- [x] Add @edgedb-ci 's Github Personal Access Token with `gist` scope as Secret `GIST_TOKEN`
- [x] Update the branch protection rules to check for the `verify` job (or update the `test.yml` to rename the final job back to `test`?)

Fixes #2201

Problem statement:
Improve test workflow performance
The current runtime of the test Github workflow has grown quite a bit, because we execute all tests in a single job (with `-j2`).  This can be improved by splitting the test suite into chunks that can then be executed in concurrent jobs.  Instance bootstrap is also quite slow and cannot be parallelized, so it'll have to be executed (and cached) in a separate job.

Hints from issue discussion:

Commit message: Run test workflow in shards

This PR adds 3 new options to `edb test`:

* `edb test --shard 2/5` separates all collected tests into 5 shards,
  and run the 2nd shard only. The split is supposed to be deterministic.
* `edb test --time-file` generates a CSV file of time spent on each
  test. If the file exists, it'll be taken as an input for `--shard` to
  generate a more even split, and the CSV file will be updated after the
  run with new average running times.
* `edb test --list` outputs a list of all selected tests and exits.

The pipeline will fetch from a Gist for `--time-file`, and update that
Gist if the tests pass in the master branch.

Fixes #2201

Patch to review:
diff --git a/.editorconfig b/.editorconfig
index 69692d36be8..562d6ce1ff4 100644
--- a/.editorconfig
+++ b/.editorconfig
@@ -4,10 +4,14 @@ root = true
 trim_trailing_whitespace = true
 insert_final_newline = true
 
-[*.{py,pyx,pxd,pxi,yml,h}]
+[*.{py,pyx,pxd,pxi,h}]
 indent_size = 4
 indent_style = space
 
+[*.yml]
+indent_size = 2
+indent_style = space
+
 [ext/*.{c,cpp,h}]
 indent_size = 4
 indent_style = tab
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index bf13d745647..550c73af305 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -13,7 +13,7 @@ on:
     - cron: "0 */3 * * *"
 
 jobs:
-  test:
+  build:
     runs-on: ubuntu-18.04
 
     steps:
@@ -56,9 +56,24 @@ jobs:
       id: buildcache
       with:
         path: build
-        key: edbbuild-${{ hashFiles('.tmp/build_cache_key.txt') }}
+        key: edbbuild-v2-${{ hashFiles('.tmp/build_cache_key.txt') }}
         restore-keys: |
-          edbbuild-
+          edbbuild-v2-
+
+    - name: Prepare time stats file
+      run: |
+        curl \
+          -H "Accept: application/vnd.github.v3+json" \
+          https://api.github.com/gists/c61fd6a07e31cbab221ec5bde378b316 \
+        | jq '.files."time_stats.csv".raw_url' \
+        | xargs curl > .tmp/time_stats.csv
+
+    - name: Upload cache key and time stats file
+      uses: actions/upload-artifact@v2
+      with:
+        name: build-cache-key-and-time-stats
+        path: .tmp
+        retention-days: 1
 
     - name: Install system deps
       if: steps.buildcache.outputs.cache-hit != 'true'
@@ -67,6 +82,7 @@ jobs:
         sudo apt-get install -y uuid-dev libreadline-dev bison flex
 
     - name: Install rust toolchain
+      if: steps.buildcache.outputs.cache-hit != 'true'
       uses: actions-rs/toolchain@v1
       with:
         profile: minimal
@@ -74,15 +90,299 @@ jobs:
         default: true
 
     - name: Build
+      if: steps.buildcache.outputs.cache-hit != 'true'
       run: |
         # --no-use-pep517 because we have explicitly installed all deps
         # and don't want them to be reinstalled in an "isolated env".
         pip install --no-use-pep517 --no-deps -e .[test,docs]
 
+    - name: Bootstrap EdgeDB Server
+      if: steps.buildcache.outputs.cache-hit != 'true'
+      run: |
+        edb server --bootstrap-only
+
+  cargo-test:
+    needs: build
+    runs-on: ubuntu-18.04
+
+    steps:
+    - uses: actions/checkout@v2
+      with:
+        fetch-depth: 0
+        submodules: false
+
+    - uses: actions/checkout@v2
+      with:
+        fetch-depth: 50
+        submodules: true
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.8
+
+    - uses: actions/cache@v2
+      id: depcache
+      with:
+        path: deps
+        key: edbdeps-pip-${{ hashFiles('setup.py') }}
+
+    - name: Install Python deps
+      run: |
+        pip install -U --no-index --find-links=deps deps/*
+
+    - name: Download cache key
+      uses: actions/download-artifact@v2
+      with:
+        name: build-cache-key-and-time-stats
+        path: .tmp
+
+    - uses: actions/cache@v2
+      id: buildcache
+      with:
+        path: build
+        key: edbbuild-v2-${{ hashFiles('.tmp/build_cache_key.txt') }}
+
+    - name: Stop if we cannot retrieve the build cache
+      if: steps.buildcache.outputs.cache-hit != 'true'
+      run: |
+        echo ::error::Cannot retrieve build cache.
+        exit 1
+
+    - name: Install rust toolchain
+      uses: actions-rs/toolchain@v1
+      with:
+        profile: minimal
+        toolchain: stable
+        default: true
+
     - name: Cargo test
       uses: actions-rs/cargo@v1
       with:
         command: test
+
+  test:
+    needs: build
+    runs-on: ubuntu-18.04
+    strategy:
+      matrix:
+        shard: [
+             1/16,  2/16,  3/16,  4/16,
+             5/16,  6/16,  7/16,  8/16,
+             9/16, 10/16, 11/16, 12/16,
+            13/16, 14/16, 15/16, 16/16,
+        ]
+
+    steps:
+    - uses: actions/checkout@v2
+      with:
+        fetch-depth: 0
+        submodules: false
+
+    - uses: actions/checkout@v2
+      with:
+        fetch-depth: 50
+        submodules: true
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.8
+
+    - uses: actions/cache@v2
+      id: depcache
+      with:
+        path: deps
+        key: edbdeps-pip-${{ hashFiles('setup.py') }}
+
+    - name: Install Python deps
+      run: |
+        pip install -U --no-index --find-links=deps deps/*
+
+    - name: Download cache key and time stats file
+      uses: actions/download-artifact@v2
+      with:
+        name: build-cache-key-and-time-stats
+        path: .tmp
+
+    - uses: actions/cache@v2
+      id: buildcache
+      with:
+        path: build
+        key: edbbuild-v2-${{ hashFiles('.tmp/build_cache_key.txt') }}
+
+    - name: Stop if we cannot retrieve the build cache
+      if: steps.buildcache.outputs.cache-hit != 'true'
+      run: |
+        echo ::error::Cannot retrieve build cache.
+        exit 1
+
+    - name: Install rust toolchain
+      uses: actions-rs/toolchain@v1
+      with:
+        profile: minimal
+        toolchain: stable
+        default: true
+
+    - name: Build
+      run: |
+        # --no-use-pep517 because we have explicitly installed all deps
+        # and don't want them to be reinstalled in an "isolated env".
+        pip install --no-use-pep517 --no-deps -e .[test,docs]
+
     - name: Test
+      env:
+        SHARD: ${{ matrix.shard }}
+      run: |
+        cp .tmp/time_stats.csv .tmp/time_stats_copy.csv
+        edb test -j2 -v -s ${SHARD} --time-file=.tmp/time_stats_copy.csv
+        mv .tmp/time_stats_copy.csv .tmp/new_${SHARD/\//_}.csv
+
+    - name: Upload updated time stats file
+      uses: actions/upload-artifact@v2
+      with:
+        name: build-cache-key-and-time-stats
+        path: .tmp
+        retention-days: 1
+
+  collect:
+    needs: build
+    runs-on: ubuntu-18.04
+    steps:
+    - uses: actions/checkout@v2
+      with:
+        fetch-depth: 0
+        submodules: false
+
+    - uses: actions/checkout@v2
+      with:
+        fetch-depth: 50
+        submodules: true
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.8
+
+    - uses: actions/cache@v2
+      id: depcache
+      with:
+        path: deps
+        key: edbdeps-pip-${{ hashFiles('setup.py') }}
+
+    - name: Install Python deps
+      run: |
+        pip install -U --no-index --find-links=deps deps/*
+
+    - name: Download cache key and time stats file
+      uses: actions/download-artifact@v2
+      with:
+        name: build-cache-key-and-time-stats
+        path: .tmp
+
+    - uses: actions/cache@v2
+      id: buildcache
+      with:
+        path: build
+        key: edbbuild-v2-${{ hashFiles('.tmp/build_cache_key.txt') }}
+
+    - name: Stop if we cannot retrieve the build cache
+      if: steps.buildcache.outputs.cache-hit != 'true'
+      run: |
+        echo ::error::Cannot retrieve build cache.
+        exit 1
+
+    - name: Install rust toolchain
+      uses: actions-rs/toolchain@v1
+      with:
+        profile: minimal
+        toolchain: stable
+        default: true
+
+    - name: Build
       run: |
-        edb test -j2 -v
+        # --no-use-pep517 because we have explicitly installed all deps
+        # and don't want them to be reinstalled in an "isolated env".
+        pip install --no-use-pep517 --no-deps -e .[test,docs]
+
+    - name: Generate complete list of tests for verification
+      env:
+        SHARD: ${{ matrix.shard }}
+      run: |
+        edb test --list > .tmp/all_tests.txt
+
+    - name: Upload complete list of tests
+      uses: actions/upload-artifact@v2
+      with:
+        name: build-cache-key-and-time-stats
+        path: .tmp
+        retention-days: 1
+
+  verify:
+    needs: [test, collect]
+    runs-on: ubuntu-18.04
+    steps:
+      - name: Set up Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.8
+
+      - name: Install Python deps
+        run: |
+          pip install requests
+
+      - name: Download all time stats file
+        uses: actions/download-artifact@v2
+        with:
+          name: build-cache-key-and-time-stats
+          path: .tmp
+
+      - name: Merge stats and verify completion
+        shell: python
+        env:
+          GIST_TOKEN: ${{ secrets.GIST_TOKEN }}
+          GIT_REF: ${{ github.ref }}
+        run: |
+          import csv
+          import glob
+          import io
+          import os
+          import requests
+
+          orig = {}
+          new = {}
+          all_tests = set()
+          with open(".tmp/time_stats.csv") as f:
+              for name, t, c in csv.reader(f):
+                  assert name not in orig, "duplicate test name in original stats!"
+                  orig[name] = (t, int(c))
+
+          with open(".tmp/all_tests.txt") as f:
+              for line in f:
+                  assert line not in all_tests, "duplicate test name in this run!"
+                  all_tests.add(line.strip())
+
+          for new_file in glob.glob(".tmp/new*.csv"):
+              with open(new_file) as f:
+                  for name, t, c in csv.reader(f):
+                      if int(c) > orig.get(name, (0, 0))[1]:
+                          assert name not in new, f"duplicate test! {name}"
+                          new[name] = (t, c)
+                          all_tests.remove(name)
+
+          assert not all_tests, "Tests not run! \n" + "\n".join(all_tests)
+
+          if os.environ["GIT_REF"] == "refs/heads/master":
+              buf = io.StringIO()
+              writer = csv.writer(buf)
+              orig.update(new)
+              for k, v in sorted(orig.items()):
+                  writer.writerow((k,) + v)
+
+              resp = requests.patch(
+                  "https://api.github.com/gists/c61fd6a07e31cbab221ec5bde378b316",
+                  headers={"Accept": "application/vnd.github.v3+json"},
+                  auth=("fantix", os.environ["GIST_TOKEN"]),
+                  json={"files": {"time_stats.csv": {"content": buf.getvalue()}}},
+              )
+              resp.raise_for_status()
diff --git a/edb/testbase/server.py b/edb/testbase/server.py
index 849122bd3fb..6980a1d7c6e 100644
--- a/edb/testbase/server.py
+++ b/edb/testbase/server.py
@@ -28,6 +28,7 @@
 import contextlib
 import decimal
 import functools
+import heapq
 import inspect
 import json
 import math
@@ -65,6 +66,16 @@
     SetupScript = str
 
 
+def _add_test(result, test):
+    cls = type(test)
+    try:
+        methods = result[cls]
+    except KeyError:
+        methods = result[cls] = []
+
+    methods.append(test)
+
+
 def get_test_cases(tests):
     result = collections.OrderedDict()
 
@@ -72,13 +83,7 @@ def get_test_cases(tests):
         if isinstance(test, unittest.TestSuite):
             result.update(get_test_cases(test._tests))
         else:
-            cls = type(test)
-            try:
-                methods = result[cls]
-            except KeyError:
-                methods = result[cls] = []
-
-            methods.append(test)
+            _add_test(result, test)
 
     return result
 
@@ -1505,3 +1510,65 @@ def start_edgedb_server(
         compiler_pool_size=compiler_pool_size,
         debug=debug
     )
+
+
+def get_cases_by_shard(cases, current_shard, total_shards, verbosity, stats):
+    if total_shards <= 1:
+        return cases
+
+    current_shard -= 1
+    new_test_est = 0.1
+    tests_with_est = []
+    new_tests = []
+    total_tests = 0
+    selected_tests = 0
+    for tests in cases.values():
+        for test in tests:
+            total_tests += 1
+            est = stats.get(str(test))
+            if est:
+                heapq.heappush(tests_with_est, (-est[0], test))
+            else:
+                new_tests.append(test)
+    shard_est = (
+        sum(-est[0] for est in tests_with_est) + new_test_est * len(new_tests)
+    ) / total_shards
+    shards_est = [0] * total_shards
+    available = list(range(total_shards))
+    current = 0
+    cases = collections.OrderedDict()
+    while tests_with_est:
+        est, test = heapq.heappop(tests_with_est)
+        shard = available[current]
+        est_acc = shards_est[shard] = shards_est[shard] - est
+        if shard == current_shard:
+            _add_test(cases, test)
+            selected_tests += 1
+        if est_acc >= shard_est:
+            if shard == current_shard:
+                break
+            del available[current]
+        else:
+            current += 1
+        current %= len(available)
+    else:
+        offset = 0
+        while len(available) > 1:
+            shard = available.pop(0)
+            needs = int(round((shard_est - shards_est[shard]) / new_test_est))
+            if shard == current_shard:
+                for test in new_tests[offset:offset + needs]:
+                    _add_test(cases, test)
+                    selected_tests += 1
+                break
+            else:
+                offset += needs
+        else:
+            for test in new_tests[offset:]:
+                _add_test(cases, test)
+                selected_tests += 1
+
+    if verbosity >= 1:
+        print(f'Running {selected_tests}/{total_tests} for shard '
+              f'#{current_shard + 1} out of {total_shards} shards.')
+    return cases
diff --git a/edb/tools/test/__init__.py b/edb/tools/test/__init__.py
index 3d9250de211..fd40dc5fd45 100644
--- a/edb/tools/test/__init__.py
+++ b/edb/tools/test/__init__.py
@@ -32,7 +32,7 @@
 
 import edb
 from edb.common import devmode
-
+from edb.testbase.server import get_test_cases
 from edb.tools.edb import edbcommands
 
 from .decorators import not_implemented
@@ -66,6 +66,9 @@
 @click.option('-j', '--jobs', type=int,
               default=lambda: (os.cpu_count() or 0) // 2,
               help='number of parallel processes to use')
+@click.option('-s', '--shard', type=str,
+              default='1/1',
+              help='run tests in shards (current/total)')
 @click.option('-k', '--include', type=str, multiple=True, metavar='REGEXP',
               help='only run tests which match the given regular expression')
 @click.option('-e', '--exclude', type=str, multiple=True, metavar='REGEXP',
@@ -80,8 +83,13 @@
               help='package name to measure code coverage for, '
                    'can be specified multiple times '
                    '(e.g --cov edb.common --cov edb.server)')
-def test(*, files, jobs, include, exclude, verbose, quiet, debug,
-         output_format, warnings, failfast, shuffle, cov, repeat):
+@click.option('--time-file', type=click.File('a+'), metavar='FILEPATH',
+              help='Maintain a running time file at FILEPATH')
+@click.option('--list', 'list_tests', is_flag=True,
+              help='list all the tests and exit')
+def test(*, files, jobs, shard, include, exclude, verbose, quiet, debug,
+         output_format, warnings, failfast, shuffle, cov, repeat, time_file,
+         list_tests):
     """Run EdgeDB test suite.
 
     Discovers and runs tests in the specified files or directories.
@@ -128,6 +136,16 @@ def test(*, files, jobs, include, exclude, verbose, quiet, debug,
                 f'Error: test path {file!r} does not exist', fg='red')
             sys.exit(1)
 
+    try:
+        current_shard, total_shards = map(int, shard.split('/'))
+    except Exception:
+        click.secho(f'Error: --shard {shard} must match format e.g. 2/5')
+        sys.exit(1)
+
+    if current_shard < 1 or current_shard > total_shards:
+        click.secho(f'Error: --shard {shard} is out of bound')
+        sys.exit(1)
+
     run = functools.partial(
         _run,
         include=include,
@@ -140,6 +158,10 @@ def test(*, files, jobs, include, exclude, verbose, quiet, debug,
         failfast=failfast,
         shuffle=shuffle,
         repeat=repeat,
+        current_shard=current_shard,
+        total_shards=total_shards,
+        time_file=time_file,
+        list_tests=list_tests,
     )
 
     if cov:
@@ -214,7 +236,8 @@ def _coverage_wrapper(paths):
 
 
 def _run(*, include, exclude, verbosity, files, jobs, output_format,
-         warnings, failfast, shuffle, repeat):
+         warnings, failfast, shuffle, repeat, current_shard, total_shards,
+         time_file, list_tests):
     suite = unittest.TestSuite()
 
     total = 0
@@ -226,7 +249,8 @@ def _update_progress(n, unfiltered_n):
             total += n
             total_unfiltered += unfiltered_n
             click.echo(styles.status(
-                f'Collected {total}/{total_unfiltered} tests.\r'), nl=False)
+                f'Collected {total}/{total_unfiltered} tests.\r'),
+                nl=False, err=list_tests)
     else:
         _update_progress = None
 
@@ -248,6 +272,14 @@ def _update_progress(n, unfiltered_n):
 
         suite.addTest(tests)
 
+    if list_tests:
+        click.echo(err=True)
+        cases = get_test_cases([suite])
+        for tests in cases.values():
+            for test in tests:
+                click.echo(str(test))
+        return 0
+
     jobs = max(min(total, jobs), 1)
 
     if verbosity > 0:
@@ -266,7 +298,7 @@ def _update_progress(n, unfiltered_n):
             warnings=warnings, num_workers=jobs,
             failfast=failfast, shuffle=shuffle)
 
-        result = test_runner.run(suite)
+        result = test_runner.run(suite, current_shard, total_shards, time_file)
 
         if not result.wasSuccessful():
             return 1
diff --git a/edb/tools/test/runner.py b/edb/tools/test/runner.py
index 8bbd7e7ddd0..c22975262f7 100644
--- a/edb/tools/test/runner.py
+++ b/edb/tools/test/runner.py
@@ -21,6 +21,7 @@
 from typing import *
 
 import collections.abc
+import csv
 import enum
 import io
 import itertools
@@ -731,9 +732,18 @@ def __init__(self, *, stream=None, num_workers=1, verbosity=1,
         self.shuffle = shuffle
         self.output_format = output_format
 
-    def run(self, test):
+    def run(self, test, current_shard, total_shards, time_file):
         session_start = time.monotonic()
         cases = tb.get_test_cases([test])
+        stats = {}
+        if time_file:
+            time_file.seek(0)
+            stats = {
+                k: (float(v), int(c)) for k, v, c in csv.reader(time_file)
+            }
+        cases = tb.get_cases_by_shard(
+            cases, current_shard, total_shards, self.verbosity, stats,
+        )
         setup = tb.get_test_cases_setup(cases)
         bootstrap_time_taken = 0
         tests_time_taken = 0
@@ -805,6 +815,17 @@ def run(self, test):
             self._echo()
             suite.run(result)
 
+            if time_file:
+                for test, stat in result.test_stats:
+                    name = str(test)
+                    t = stat['running-time']
+                    at, c = stats.get(name, (0, 0))
+                    stats[name] = (at + (t - at) / (c + 1), c + 1)
+                time_file.seek(0)
+                time_file.truncate()
+                writer = csv.writer(time_file)
+                for k, v in stats.items():
+                    writer.writerow((k, ) + v)
             tests_time_taken = time.monotonic() - start
 
         except KeyboardInterrupt:


Code review comments:
diff_hunk: @@ -80,8 +83,13 @@
               help='package name to measure code coverage for, '
                    'can be specified multiple times '
                    '(e.g --cov edb.common --cov edb.server)')
-def test(*, files, jobs, include, exclude, verbose, quiet, debug,
-         output_format, warnings, failfast, shuffle, cov, repeat):
+@click.option('--time-file', type=click.File('a+'), metavar='FILEPATH',
path: edb/tools/test/__init__.py
text: @user1:
I'd call this `--log-running-times`

@user1:
Or, rather `--running-times-log`

@author:
Oh actually it also reads this file for calculating distribution - should we create a separate option for reading?

@user2:
Why a+?

@author:
It's both for reading and writing, and I'd like the file to be created if not exist. `r+` doesn't create the file, and `w+` truncates the file on open.