Repository: edgedb/edgedb

Issue:
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.


Discussion:

Resolving patch:
diff --git a/edb/testbase/server.py b/edb/testbase/server.py
--- 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


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
--- 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:
@@ -32,7 +32,8 @@ jobs:
       with:
         python-version: 3.8
 
-    - uses: actions/cache@v2
+    - name: Handle Pip dependency cache
+      uses: actions/cache@v2
       id: depcache
       with:
         path: deps
@@ -52,13 +53,32 @@ jobs:
         mkdir -p .tmp
         python setup.py --quiet gen_build_cache_key >.tmp/build_cache_key.txt
 
-    - uses: actions/cache@v2
+    - name: Handle build cache
+      uses: actions/cache@v2
       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: Download the running times log file
+      env:
+        GIST_TOKEN: ${{ secrets.CI_BOT_GIST_TOKEN }}
+      run: |
+        curl \
+          -H "Accept: application/vnd.github.v3+json" \
+          -u edgedb-ci:$GIST_TOKEN \
+          https://api.github.com/gists/8b722a65397f7c4c0df72f5394efa04c \
+        | jq '.files."time_stats.csv".raw_url' \
+        | xargs curl > .tmp/time_stats.csv
+
+    - name: Upload shared artifacts
+      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 +87,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 +95,308 @@ 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
+
+    - name: Handle Pip dependency cache
+      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
+
+    - name: Handle build cache
+      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
+
+  python-test:
+    needs: build
+    runs-on: ubuntu-18.04
+    strategy:
+      matrix:
+        shard: [
+             1/32,  2/32,  3/32,  4/32,
+             5/32,  6/32,  7/32,  8/32,
+             9/32, 10/32, 11/32, 12/32,
+            13/32, 14/32, 15/32, 16/32,
+            17/32, 18/32, 19/32, 20/32,
+            21/32, 22/32, 23/32, 24/32,
+            25/32, 26/32, 27/32, 28/32,
+            29/32, 30/32, 31/32, 32/32,
+        ]
+
+    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
+
+    - name: Handle Pip dependency cache
+      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 shared artifacts
+      uses: actions/download-artifact@v2
+      with:
+        name: build-cache-key-and-time-stats
+        path: .tmp
+
+    - name: Handle build cache
+      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: Install
+      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/new_${SHARD/\//_}.csv
+        edb test -j2 -v -s ${SHARD} --running-times-log=.tmp/new_${SHARD/\//_}.csv
+
+    - name: Update shared artifacts
+      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
+
+    - name: Handle Pip dependency cache
+      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 shared artifacts
+      uses: actions/download-artifact@v2
+      with:
+        name: build-cache-key-and-time-stats
+        path: .tmp
+
+    - name: Handle build cache
+      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: Install
+      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: Generate complete list of tests for verification
+      env:
+        SHARD: ${{ matrix.shard }}
       run: |
-        edb test -j2 -v
+        edb test --list > .tmp/all_tests.txt
+
+    - name: Update shared artifacts
+      uses: actions/upload-artifact@v2
+      with:
+        name: build-cache-key-and-time-stats
+        path: .tmp
+        retention-days: 1
+
+  test:
+    needs: [cargo-test, python-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 shared artifacts
+        uses: actions/download-artifact@v2
+        with:
+          name: build-cache-key-and-time-stats
+          path: .tmp
+
+      - name: Merge stats and verify tests completion
+        shell: python
+        env:
+          GIST_TOKEN: ${{ secrets.CI_BOT_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/8b722a65397f7c4c0df72f5394efa04c",
+                  headers={"Accept": "application/vnd.github.v3+json"},
+                  auth=("edgedb-ci", os.environ["GIST_TOKEN"]),
+                  json={"files": {"time_stats.csv": {"content": buf.getvalue()}}},
+              )
+              resp.raise_for_status()
diff --git a/edb/tools/test/__init__.py b/edb/tools/test/__init__.py
--- 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,14 @@
               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('--running-times-log', 'running_times_log_file',
+              type=click.File('a+'), metavar='FILEPATH',
+              help='Maintain a running time log 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,
+         running_times_log_file, list_tests):
     """Run EdgeDB test suite.
 
     Discovers and runs tests in the specified files or directories.
@@ -128,6 +137,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 +159,10 @@ def test(*, files, jobs, include, exclude, verbose, quiet, debug,
         failfast=failfast,
         shuffle=shuffle,
         repeat=repeat,
+        current_shard=current_shard,
+        total_shards=total_shards,
+        running_times_log_file=running_times_log_file,
+        list_tests=list_tests,
     )
 
     if cov:
@@ -214,7 +237,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,
+         running_times_log_file, list_tests):
     suite = unittest.TestSuite()
 
     total = 0
@@ -226,7 +250,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 +273,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 +299,9 @@ 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, running_times_log_file,
+        )
 
         if not result.wasSuccessful():
             return 1
diff --git a/edb/tools/test/runner.py b/edb/tools/test/runner.py
--- 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,19 @@ 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, running_times_log_file):
         session_start = time.monotonic()
         cases = tb.get_test_cases([test])
+        stats = {}
+        if running_times_log_file:
+            running_times_log_file.seek(0)
+            stats = {
+                k: (float(v), int(c))
+                for k, v, c in csv.reader(running_times_log_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 +816,17 @@ def run(self, test):
             self._echo()
             suite.run(result)
 
+            if running_times_log_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)
+                running_times_log_file.seek(0)
+                running_times_log_file.truncate()
+                writer = csv.writer(running_times_log_file)
+                for k, v in stats.items():
+                    writer.writerow((k, ) + v)
             tests_time_taken = time.monotonic() - start
 
         except KeyboardInterrupt: