ci(tests): Run backend shards without Snuba where tests allow

.github/workflows/backend.yml

@@ -252,12 +252,22 @@ jobs:
       MATRIX_INSTANCE_TOTAL: ${{ needs.calculate-shards.outputs.shard-count }}
       TEST_GROUP_STRATEGY: roundrobin
       PYTHONHASHSEED: '0'
-      XDIST_PER_WORKER_SNUBA: '1'
+      # Full runs: shards [0, BACKEND_SNUBA_SHARDS) run the tests that need Snuba,
+      # Kafka or Symbolicator, the rest run without them. 0 runs every shard with
+      # every service and records what each test file uses. Pick-shard sets
+      # XDIST_PER_WORKER_SNUBA. See .github/workflows/scripts/backend-test-services.py.
+      BACKEND_SNUBA_SHARDS: '23'
       XDIST_WORKERS: '4'
 
     steps:
       - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
 
+      - name: Pick shard services
+        env:
+          SELECTIVE: ${{ needs.select-tests.outputs.has-selected-tests == 'true' }}
+          MATRIX_INSTANCE: ${{ matrix.instance }}
+        run: python3 .github/workflows/scripts/backend-test-services.py pick-shard
+
       - name: Start resource monitor
         run: python3 .github/workflows/scripts/monitor-resources.py --start
 
@@ -272,7 +282,7 @@ jobs:
         uses: ./.github/actions/setup-sentry
         id: setup
         with:
-          mode: backend-ci
+          mode: ${{ env.BACKEND_DEVSERVICES_MODE }}
           early-devservices: 'true'
 
       # Restart postgres with a Unix socket mount. Django's postgres driver uses TCP by
@@ -404,6 +414,16 @@ jobs:
             .artifacts/resource-usage-summary.json
           retention-days: 30
 
+      # backend-test-services.py builds backend-test-services.json from these.
+      - name: Upload service usage
+        if: ${{ !cancelled() && env.BACKEND_TEST_SERVICE_GUARD != '' && hashFiles('.artifacts/service-usage/*.json') != '' }}
+        continue-on-error: true
+        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
+        with:
+          name: backend-service-usage-${{ github.run_id }}-${{ matrix.instance }}
+          path: .artifacts/service-usage/*.json
+          retention-days: 30
+
       - name: Report failures
         if: ${{ !cancelled() && github.event_name == 'pull_request' }}
         continue-on-error: true

.github/workflows/scripts/backend-test-services.py

@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+"""Service-colored full backend runs: which shards start Snuba, and which tests need it.
+
+Full runs put the test files that use Snuba, ClickHouse, Kafka or Symbolicator
+on the first BACKEND_SNUBA_SHARDS shards (devservices mode `backend-ci`) and
+every other file on "lite" shards (`backend-ci-lite`: no Snuba, ClickHouse,
+Kafka or Symbolicator, and no per-worker Snuba bootstrap). pytest-selected.py
+plans the split from `backend-test-services.json`; pytest_service_guard.py
+records service use on Snuba shards and fails any test that tries one on a lite
+shard. Selective PR runs keep every service.
+
+    pick-shard    Write this matrix shard's devservices mode and guard mode to
+                  $GITHUB_ENV (backend.yml).
+    inventory     Every shard's `backend-service-usage-*` artifacts ->
+                  backend-test-services.json. Refresh it from a green full run:
+
+        gh run download RUN_ID --repo getsentry/sentry \\
+            --pattern 'backend-service-usage-*' --dir /tmp/usage
+        python3 .github/workflows/scripts/backend-test-services.py inventory \\
+            .github/workflows/scripts/backend-test-services.json /tmp/usage/*/*.json
+
+                  Files used services if any shard saw them do so. Files that no
+                  shard ran (new, renamed, or from a failed shard) drop out and
+                  go to Snuba shards until the next refresh.
+    suggest       Predicted per-shard test-seconds for each BACKEND_SNUBA_SHARDS.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+import os
+import sys
+from pathlib import Path
+
+
+def pick_shard(env: dict[str, str]) -> dict[str, str]:
+    selective = env["SELECTIVE"] == "true"
+    snuba_shards = int(env["BACKEND_SNUBA_SHARDS"])
+    instance = int(env["MATRIX_INSTANCE"])
+    if selective:
+        color, guard = "", ""
+    elif snuba_shards == 0:
+        # Not colored yet, but measure, to build backend-test-services.json.
+        color, guard = "", "record"
+    elif instance < snuba_shards:
+        color, guard = "snuba", "record"
+    else:
+        color, guard = "lite", "enforce"
+    return {
+        "BACKEND_TEST_SERVICE_COLOR": color,
+        "BACKEND_TEST_SERVICE_GUARD": guard,
+        "BACKEND_DEVSERVICES_MODE": "backend-ci-lite" if color == "lite" else "backend-ci",
+        "XDIST_PER_WORKER_SNUBA": "" if color == "lite" else "1",
+    }
+
+
+def inventory(usages: list[dict[str, object]]) -> dict[str, object]:
+    files: dict[str, set[str]] = {}
+    for usage in usages:
+        for path, services in usage["files"].items():  # type: ignore[attr-defined]
+            files.setdefault(path, set()).update(services)
+    return {"files": {path: sorted(files[path]) for path in sorted(files)}}
+
+
+def suggest(total: int) -> None:
+    sys.path.insert(0, str(Path(__file__).parent))
+    planner = importlib.import_module("pytest-selected")
+    files = planner.list_test_files(
+        ["tests", *(f"--ignore=tests/{d}" for d in ("acceptance", "apidocs", "js", "tools"))]
+    )
+    services = json.loads(planner.SERVICES_FILE.read_text())["files"]
+    snuba_files = {path for path in files if planner.needs_snuba(path, services)}
+    weights = planner.file_weights(files)
+    snuba = sum(weights[path] for path in snuba_files)
+    print(f"{len(snuba_files)}/{len(files)} files, {snuba:.0f}/{sum(weights.values()):.0f}s")
+    print("snuba_shards  max_snuba_shard_s  max_lite_shard_s")
+    for k in range(1, total):
+        try:
+            plan = planner.plan_colored_shards(weights, total, k, snuba_files)
+        except ValueError:
+            continue
+        loads = [
+            sum(weights[path] / shard.parts.get(path, (0, 1))[1] for path in shard.files)
+            for shard in plan
+        ]
+        print(f"{k:12}  {max(loads[:k]):17.0f}  {max(loads[k:]):16.0f}")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0])
+    sub = parser.add_subparsers(dest="command", required=True)
+    sub.add_parser("pick-shard", help="this shard's services -> $GITHUB_ENV")
+    p = sub.add_parser("inventory", help="service-usage artifacts -> the checked-in file")
+    p.add_argument("output", type=Path)
+    p.add_argument("inputs", type=Path, nargs="+")
+    p = sub.add_parser("suggest", help="predicted shard loads per BACKEND_SNUBA_SHARDS")
+    p.add_argument("--total", type=int, default=30)
+    args = parser.parse_args()
+
+    if args.command == "pick-shard":
+        picked = pick_shard(dict(os.environ))
+        with open(os.environ["GITHUB_ENV"], "a") as f:
+            f.writelines(f"{key}={value}\n" for key, value in picked.items())
+        print(picked)
+    elif args.command == "inventory":
+        result = inventory([json.loads(path.read_text()) for path in args.inputs])
+        args.output.write_text(json.dumps(result, indent=2) + "\n")
+    else:
+        suggest(args.total)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

.github/workflows/scripts/pytest-selected.py

@@ -17,6 +17,10 @@
   durations file, each shard gets whole selected files, balanced by the
   per-file test counts in `selected-tests.counts.json`, and a file too big for
   one shard gets several shards, which split it by test.
+- Service-colored full runs (BACKEND_TEST_SERVICE_COLOR set): the first
+  BACKEND_SNUBA_SHARDS shards run the files that need Snuba, Kafka or
+  Symbolicator, and the rest run every other file on "lite" shards without
+  those services. See backend-test-services.py.
 """
 
 from __future__ import annotations
@@ -33,6 +37,14 @@
 import pytest
 
 DURATIONS_FILE = Path(__file__).parent / "backend-test-durations.json"
+# Which services each test file used in the last measured full run.
+SERVICES_FILE = Path(__file__).parent / "backend-test-services.json"
+# Puts a file on Snuba shards whatever SERVICES_FILE says, e.g. while a new
+# Snuba dependency waits for the file to be measured again.
+SNUBA_PRAGMA = "# backend-ci: snuba"
+# Relay (which produces to Kafka) and Symbolicator run in containers, so
+# pytest_service_guard can't see them being used. These mark their tests.
+SNUBA_HINTS = ("RelayStoreHelper", "relay_server", "requires_symbolicator", ".mark.symbolicator")
 # Full runs split files slower than 1/PIECES_PER_SHARD of a shard by test.
 # Files run in collection order, so a big file that starts late on an xdist
 # worker holds up the whole shard; this bounds how big that can be.
@@ -122,6 +134,31 @@ def plan_full_shards(
     ]
 
 
+def needs_snuba(path: str, services: dict[str, list[str]]) -> bool:
+    """Whether a file must run on a shard with Snuba, Kafka and Symbolicator.
+
+    Only files measured as using none of them go to lite shards, so new and
+    renamed files start on Snuba shards until the next measurement.
+    """
+    if services.get(path) != []:
+        return True
+    source = Path(path).read_text()
+    return SNUBA_PRAGMA in source or any(hint in source for hint in SNUBA_HINTS)
+
+
+def plan_colored_shards(
+    weights: dict[str, float], total: int, snuba_shards: int, snuba_files: set[str]
+) -> list[Shard]:
+    """Snuba files on shards [0, snuba_shards), the rest on the others."""
+    snuba = {path: w for path, w in weights.items() if path in snuba_files}
+    lite = {path: w for path, w in weights.items() if path not in snuba_files}
+    snuba_plan = plan_full_shards(snuba, snuba_shards)
+    lite_plan = plan_full_shards(lite, total - snuba_shards)
+    if snuba_plan is None or lite_plan is None:
+        raise ValueError(f"BACKEND_SNUBA_SHARDS={snuba_shards} leaves a shard with no files")
+    return snuba_plan + lite_plan
+
+
 def list_test_files(args: list[str]) -> list[str]:
     """Every .py file pytest could collect from the paths in `args`, minus --ignore.
 
@@ -207,6 +244,58 @@ def run_full_shard(args: list[str], total: int, group: int) -> int:
     )
 
 
+def run_colored_shard(args: list[str], total: int, group: int, color: str) -> int:
+    """Like run_full_shard, but no fallback: a lite shard can't run Snuba tests."""
+    try:
+        snuba_shards = int(os.environ.get("BACKEND_SNUBA_SHARDS", "0"))
+        if not 0 < snuba_shards < total or not 0 <= group < total:
+            raise ValueError(f"BACKEND_SNUBA_SHARDS={snuba_shards} must be in (0, {total})")
+        expected = "snuba" if group < snuba_shards else "lite"
+        if color != expected:
+            raise ValueError(f"shard {group} is a {expected} shard, got color {color}")
+        if color == "lite" and os.environ.get("BACKEND_TEST_SERVICE_GUARD") != "enforce":
+            raise ValueError("lite shards need BACKEND_TEST_SERVICE_GUARD=enforce")
+        files = list_test_files(args)
+        if not files or not per_file_plan_applies(args, files):
+            raise ValueError("needs test files under the working directory, and no -m/-k")
+        services = json.loads(SERVICES_FILE.read_text())["files"]
+        snuba_files = {path for path in files if needs_snuba(path, services)}
+        weights = file_weights(files)
+        plan = plan_colored_shards(weights, total, snuba_shards, snuba_files)
+    except (OSError, ValueError, KeyError) as exc:
+        print(f"pytest-selected: cannot plan service-colored shards: {exc}", file=sys.stderr)
+        return int(pytest.ExitCode.USAGE_ERROR)
+
+    snuba_seconds = sum(weights[path] for path in snuba_files)
+    print(
+        f"pytest-selected: {color} shard; {len(snuba_files)} of {len(files)} files and "
+        f"{snuba_seconds:.0f} of {sum(weights.values()):.0f} test-seconds need Snuba",
+        file=sys.stderr,
+    )
+    # Like TOTAL_TEST_GROUPS below: tests that run pytest-selected.py
+    # themselves must not inherit this shard's plan.
+    os.environ.pop("BACKEND_TEST_SERVICE_COLOR")
+    os.environ.pop("BACKEND_SNUBA_SHARDS", None)
+    if color == "lite":
+        # A guard failure is deterministic, so don't rerun it.
+        os.environ["PYTEST_ADDOPTS"] = (
+            f"{os.environ.get('PYTEST_ADDOPTS', '')} --rerun-except=ServiceGuardError"
+        )
+    artifacts = Path(".artifacts")
+    artifacts.mkdir(exist_ok=True)
+    return run_planned_shard(
+        args, plan, group, len(files), artifacts / f"full-tests.shard-{group}.txt"
+    )
+
+
+def add_plugin(name: str) -> None:
+    """Load a plugin from this directory in pytest and its xdist workers."""
+    os.environ["PYTHONPATH"] = os.pathsep.join(
+        filter(None, [str(Path(__file__).parent), os.environ.get("PYTHONPATH")])
+    )
+    os.environ["PYTEST_PLUGINS"] = ",".join(filter(None, [os.environ.get("PYTEST_PLUGINS"), name]))
+
+
 def run_planned_shard(
     args: list[str], plan: list[Shard], group: int, file_count: int, shard_file: Path
 ) -> int:
@@ -227,12 +316,7 @@ def run_planned_shard(
     os.environ["TOTAL_TEST_GROUPS"] = "1"
     os.environ["TEST_GROUP"] = "0"
     os.environ["PYTEST_SHARD_PARTS"] = str(parts_file)
-    os.environ["PYTHONPATH"] = os.pathsep.join(
-        filter(None, [str(Path(__file__).parent), os.environ.get("PYTHONPATH")])
-    )
-    os.environ["PYTEST_PLUGINS"] = ",".join(
-        filter(None, [os.environ.get("PYTEST_PLUGINS"), "pytest_shard_plugin"])
-    )
+    add_plugin("pytest_shard_plugin")
     start_assert_rewrite_prewarm(shard_file)
     return int(pytest.main(args))
 
@@ -241,6 +325,16 @@ def main(args: list[str]) -> int:
     selected_file = os.environ.get("SELECTED_TESTS_FILE")
     total = int(os.environ.get("TOTAL_TEST_GROUPS", 1))
     group = int(os.environ.get("TEST_GROUP", 0))
+    color = os.environ.get("BACKEND_TEST_SERVICE_COLOR", "")
+    if os.environ.get("BACKEND_TEST_SERVICE_GUARD"):
+        # Ahead of tests/conftest.py, so it can wrap Kafka clients before
+        # sentry imports them.
+        add_plugin("pytest_service_guard")
+    if color:
+        if selected_file:
+            print("pytest-selected: service colors are for full runs only", file=sys.stderr)
+            return int(pytest.ExitCode.USAGE_ERROR)
+        return run_colored_shard(args, total, group, color)
     if not selected_file and total > 1 and DURATIONS_FILE.is_file():
         return run_full_shard(args, total, group)
     if not selected_file or not os.path.isfile(selected_file) or total <= 1: