ci: fix runtime shard JUnit metadata

Co-authored-by: mr-lee <1302416+mr-lee@users.noreply.github.com>

Author: datadog-datadog-prod-us1[bot]

.gitlab/scripts/get-riot-hashes.sh

@@ -2,14 +2,55 @@
 set -e -u -o pipefail
 
 SUITE_NAME="${1:-}"
+CI_ALLOCATION_SUITE="${2:-${CI_ALLOCATION_SUITE:-${SUITE_NAME}}}"
 strategy_args=()
 if [[ -n "${CI_ALLOCATION_STRATEGY:-}" ]]; then
     strategy_args=(--strategy "${CI_ALLOCATION_STRATEGY}")
 fi
 
-riot list --hash-only "${SUITE_NAME}" | sort | \
+mapfile -t available_hashes < <(riot list --hash-only "${SUITE_NAME}" | sort)
+if [[ -n "${CI_ALLOCATION_ASSIGNMENTS:-}" ]]; then
+    IFS=';' read -r -a planned_assignments <<< "${CI_ALLOCATION_ASSIGNMENTS}"
+    node_index="${CI_NODE_INDEX:-1}"
+    node_total="${CI_NODE_TOTAL:-1}"
+    if [[ ! "${node_index}" =~ ^[1-9][0-9]*$ || ! "${node_total}" =~ ^[1-9][0-9]*$ ]]; then
+        echo "CI node index and total must be positive integers" >&2
+        exit 1
+    fi
+    if [[ "${#planned_assignments[@]}" -ne "${node_total}" ]]; then
+        echo "Generated allocation count differs from CI_NODE_TOTAL" >&2
+        exit 1
+    fi
+    if [[ "${node_index}" -gt "${#planned_assignments[@]}" ]]; then
+        echo "CI_NODE_INDEX is outside the generated allocation" >&2
+        exit 1
+    fi
+    CI_ALLOCATION_UNITS="${planned_assignments[$((node_index - 1))]}"
+fi
+if [[ -n "${CI_ALLOCATION_UNITS:-}" ]]; then
+    declare -A available=()
+    for riot_hash in "${available_hashes[@]}"; do
+        available["${riot_hash}"]=1
+    done
+    IFS=',' read -r -a execution_units <<< "${CI_ALLOCATION_UNITS}"
+    for unit in "${execution_units[@]}"; do
+        if [[ ! "${unit}" =~ ^([0-9a-f]+)(@([1-9][0-9]*)/([1-9][0-9]*))?$ ]]; then
+            echo "Invalid Riot execution unit: ${unit}" >&2
+            exit 1
+        fi
+        riot_hash="${BASH_REMATCH[1]}"
+        if [[ -z "${available[${riot_hash}]:-}" ]]; then
+            echo "Generated Riot execution unit is not in ${SUITE_NAME}: ${unit}" >&2
+            exit 1
+        fi
+        printf '%s\n' "${unit}"
+    done
+    exit 0
+fi
+
+printf '%s\n' "${available_hashes[@]}" | \
     ./scripts/ci_allocation_cli.py select \
-        --suite "${SUITE_NAME}" \
+        --suite "${CI_ALLOCATION_SUITE}" \
         --node-index "${CI_NODE_INDEX:-1}" \
         --node-total "${CI_NODE_TOTAL:-1}" \
         "${strategy_args[@]}"

.gitlab/scripts/get-riot-pip-cache-key.sh

@@ -2,7 +2,9 @@
 set -e -u -o pipefail
 
 SUITE_NAME="${1:-}"
-hashes=( $(./.gitlab/scripts/get-riot-hashes.sh "${SUITE_NAME}") )
+# Cache identity covers the full semantic suite and must not depend on one
+# physical allocation job's node index or runtime test slice.
+hashes=( $(riot list --hash-only "${SUITE_NAME}" | sort -u) )
 # Get the sha256sum of all the requirements files combined
 for hash in "${hashes[@]}"; do
   req_file="./.riot/requirements/${hash}.txt"

.gitlab/testrunner.yml

@@ -25,6 +25,7 @@ variables:
     when: always
     paths:
       - core.*
+      - test-results/ci-test-shard-inventory.*.json
     reports:
       junit: test-results/junit*.xml
     expire_in: 1 week

.gitlab/tests.yml

@@ -36,13 +36,24 @@ include:
         echo "No riot hashes found for ${SUITE_NAME}"
         exit 1
       fi
-      for hash in "${hashes[@]}"
+      for unit in "${hashes[@]}"
       do
+        hash="${unit%%@*}"
+        if [[ "${unit}" =~ ^[0-9a-f]+@([1-9][0-9]*)/([1-9][0-9]*)$ ]]; then
+          test_shard_index="${BASH_REMATCH[1]}"
+          test_shard_total="${BASH_REMATCH[2]}"
+        else
+          test_shard_index=1
+          test_shard_total=1
+        fi
         echo "Running riot hash: ${hash}"
+        echo "Runtime test shard: ${test_shard_index}/${test_shard_total}"
         riot list "${hash}"
         export RIOT_HASH="${hash}"
+        export RIOT_TEST_SHARD_INDEX="${test_shard_index}"
+        export RIOT_TEST_SHARD_TOTAL="${test_shard_total}"
         export RIOT_CI_ALLOCATION_STRATEGY="${CI_ALLOCATION_STRATEGY:-legacy}"
-        export _CI_DD_TAGS="test.configuration.riot_hash:${RIOT_HASH},test.configuration.ci_allocation_strategy:${RIOT_CI_ALLOCATION_STRATEGY}"
+        export _CI_DD_TAGS="test.configuration.riot_hash:${RIOT_HASH},test.configuration.ci_allocation_strategy:${RIOT_CI_ALLOCATION_STRATEGY},test.configuration.runtime_shard_index:${RIOT_TEST_SHARD_INDEX},test.configuration.runtime_shard_total:${RIOT_TEST_SHARD_TOTAL}"
         ${RIOT_RUN_CMD} "${hash}" -- --ddtrace
       done
       ./scripts/check-diff ".riot/requirements/" \

ci/ci-allocation-policy.json

@@ -1,8 +1,9 @@
 {
   "allocation": {
     "active_strategy": "legacy",
-    "maximum_parallelism_per_suite": 25,
+    "maximum_parallelism_per_suite": 50,
     "maximum_runtime_model_bytes": 100000,
+    "maximum_slices_per_hash": 5,
     "target_jobs": 200,
     "target_shard_seconds": 300
   },

ci/ci-allocation-runtime-model.json

@@ -2062,6 +2062,57 @@
     "minimum_samples": 5,
     "sparse_safety_factor": 1.25
   },
+  "suite_estimates": {
+    "contrib::integration_registry": {
+      "2e9f3b5": 127.23747
+    },
+    "integration_registry": {
+      "2e9f3b5": 129.210514
+    },
+    "tracer": {
+      "16f089d": 651.383996,
+      "190d82d": 651.383996,
+      "1c97cf2": 651.383996,
+      "3d924d3": 651.383996,
+      "f953f1c": 651.383996
+    },
+    "tracer-uwsgi": {
+      "16f089d": 75.349606,
+      "190d82d": 75.349606,
+      "1c97cf2": 75.349606,
+      "3d924d3": 75.349606,
+      "f953f1c": 75.349606
+    }
+  },
+  "test_sharding": {
+    "command_fingerprints": {
+      "21c32b52591cbd41dd16f7df087cb15963733fe59de2ce9b011bf033a77c76b7": {
+        "minimum_items": 820,
+        "observed_hash": "759749c",
+        "source_commit": "3bd95510c20255b89fa41e1ad9822d617c1ffe3d",
+        "source_pipeline_id": "131076783"
+      },
+      "3d6032832c5128ba1764834c1340c9e50a459fb876aab90bb20c6c13585cb148": {
+        "minimum_items": 6922,
+        "observed_hash": "107d2ec",
+        "source_commit": "3bd95510c20255b89fa41e1ad9822d617c1ffe3d",
+        "source_pipeline_id": "131076783"
+      },
+      "6c31901e916dd7d21856d6795591a3f3f0be7bbaf6d5c3d5e1d1d1661c08bef9": {
+        "minimum_items": 187,
+        "observed_hash": "dbf2d5f",
+        "source_commit": "aa623634dbf9e9ff55ba89f822ad7dda228ce087",
+        "source_pipeline_id": "131104148"
+      },
+      "986b1f7f162587fd9fda02ea4a2430b87751117fe67cfa6d3273aa8b301604c8": {
+        "minimum_items": 2086,
+        "observed_hash": "b783dae",
+        "source_commit": "3bd95510c20255b89fa41e1ad9822d617c1ffe3d",
+        "source_pipeline_id": "131076783"
+      }
+    },
+    "source": "datadog-test-visibility-items"
+  },
   "overheads": {
     "global_seconds": 0,
     "suite_seconds": {},

conftest.py

@@ -9,13 +9,18 @@
 import hashlib
 import json
 import os
+from pathlib import Path
 import re
 import sys
 from time import time
 
 import hypothesis
 import pytest
 
+from scripts.ci_allocation.planner import AllocationError
+from scripts.ci_allocation.runtime import build_runtime_inventory
+from scripts.ci_allocation.runtime import write_runtime_inventory
+
 
 # DEV: Enable "testdir" fixture https://docs.pytest.org/en/stable/reference.html#testdir
 pytest_plugins = ("pytester",)
@@ -63,11 +68,21 @@ def pytest_configure(config):
         fname, ext = os.path.splitext(config.option.xmlpath)
         strategy = os.getenv("RIOT_CI_ALLOCATION_STRATEGY")
         riot_hash = os.getenv("RIOT_HASH")
+        test_shard_index = os.getenv("RIOT_TEST_SHARD_INDEX")
+        test_shard_total = os.getenv("RIOT_TEST_SHARD_TOTAL")
+        test_shard_identity = None
+        if test_shard_index and test_shard_total and int(test_shard_total) > 1:
+            test_shard_identity = f"s{test_shard_index}of{test_shard_total}"
         execution_digest = None
         if riot_hash:
             execution = {}
             for env, value in os.environ.items():
-                if not env.startswith("RIOT_") or env in {"RIOT_HASH", "RIOT_CI_ALLOCATION_STRATEGY"}:
+                if not env.startswith("RIOT_") or env in {
+                    "RIOT_HASH",
+                    "RIOT_CI_ALLOCATION_STRATEGY",
+                    "RIOT_TEST_SHARD_INDEX",
+                    "RIOT_TEST_SHARD_TOTAL",
+                }:
                     continue
                 name = env[5:]
                 prefix, _, suffix = name.partition("_")
@@ -82,6 +97,7 @@ def pytest_configure(config):
             (
                 strategy,
                 riot_hash,
+                test_shard_identity,
                 execution_digest,
                 str(os.getpid()),
             ),
@@ -95,6 +111,52 @@ def pytest_configure(config):
         config.option.benchmark_save = str(time()).replace(".", "_") + gc + "_py%d_%d" % sys.version_info[:2]
 
 
+@pytest.hookimpl(trylast=True)
+def pytest_collection_modifyitems(config, items):
+    """Select one deterministic runtime slice after the Riot command collects tests."""
+    shard_index_value = os.getenv("RIOT_TEST_SHARD_INDEX")
+    shard_total_value = os.getenv("RIOT_TEST_SHARD_TOTAL")
+    if shard_index_value is None and shard_total_value is None:
+        return
+    if shard_index_value is None or shard_total_value is None:
+        raise pytest.UsageError("Riot runtime test sharding requires both shard index and total")
+    try:
+        shard_index = int(shard_index_value)
+        shard_total = int(shard_total_value)
+    except ValueError as exc:
+        raise pytest.UsageError("Riot runtime test shard index and total must be integers") from exc
+    if shard_total == 1 and shard_index == 1:
+        return
+
+    suite = os.getenv("CI_ALLOCATION_SUITE", "")
+    riot_hash = os.getenv("RIOT_HASH", "")
+    if not suite or not riot_hash:
+        raise pytest.UsageError("Riot runtime test sharding requires suite and hash identity")
+    try:
+        inventory = build_runtime_inventory(
+            suite=suite,
+            riot_hash=riot_hash,
+            shard_index=shard_index,
+            shard_total=shard_total,
+            collected_nodeids=[item.nodeid for item in items],
+        )
+    except AllocationError as exc:
+        raise pytest.UsageError(str(exc)) from exc
+
+    selected = set(inventory["selected_nodeids"])
+    deselected = [item for item in items if item.nodeid not in selected]
+    items[:] = [item for item in items if item.nodeid in selected]
+    if deselected:
+        config.hook.pytest_deselected(items=deselected)
+
+    # AIDEV-NOTE: xdist workers must make the same selection, but only one
+    # process writes the shared inventory artifact.
+    worker = os.getenv("PYTEST_XDIST_WORKER")
+    if worker in (None, "gw0"):
+        path = Path("test-results") / (f"ci-test-shard-inventory.{riot_hash}.{shard_index}-of-{shard_total}.json")
+        write_runtime_inventory(path, inventory)
+
+
 @pytest.hookimpl(tryfirst=True, hookwrapper=True)
 def pytest_runtest_makereport(item, call):
     # Attach the outcome of the test (failed, passed, skipped) to the test node so that fixtures

docs/contributing-testing.rst

@@ -131,20 +131,31 @@ How CI allocates Riot environments
 Semantic test ownership remains in ``tests/suitespec.py`` and the distributed
 ``suitespec.yml`` files. CI resolves every selected suite to its Riot environment
 hashes, then assigns those hashes to physical GitLab shards. A Riot environment is
-the smallest allocation unit, so duration-based balancing cannot change its command,
-Python version, services, environment, retry policy, or timeout.
+normally the smallest allocation unit, so duration-based balancing never changes its
+command, Python version, services, environment, retry policy, or timeout. For a
+measured long-running pytest command, CI may refine one hash into runtime execution
+units such as ``107d2ec@1/3``. Pytest still collects the suite-authored command, then
+a repository plugin deterministically selects one disjoint slice of the collected
+node IDs in each physical job. No item-level partition appears in ``suitespec.yml``.
 
 ``scripts/gen_gitlab_config.py`` writes ``.gitlab/ci-allocation-plan.json`` with the
 current round-robin plan and a duration-balanced plan. Generation fails unless both
 plans contain the exact same Riot hash set, with no duplicates or empty shards, and
-the execution metadata is represented by the same digest. The plan is retained as a
-CI artifact for review.
+the execution metadata is represented by the same digest. Runtime slices must also
+contain every index from one through their declared total. The plan is retained as a
+CI artifact for review. The generator embeds every balanced assignment into the
+suite's GitLab ``parallel`` job, and each physical job selects its assignment by
+``CI_NODE_INDEX``. The semantic suite and compact matrix representation therefore
+remain stable even when an assignment contains sub-hash execution units.
 
 The active strategy and promotion thresholds are in
 ``ci/ci-allocation-policy.json``. ``legacy`` is the rollback-safe default. The
 duration estimates in ``ci/ci-allocation-runtime-model.json`` are generated from
-Datadog Test Visibility session exports joined by
-``test.configuration.riot_hash``. The model uses a time-decayed p90, conservative
+Datadog Test Visibility session exports joined by semantic suite and
+``test.configuration.riot_hash``. Riot hashes identify environments rather than
+commands and can be shared by semantically different suites, so suite-scoped
+estimates override global hash estimates for those collisions. The model uses a
+time-decayed p90, conservative
 fallbacks for sparse hashes, and a recent holdout that is not used for fitting.
 CI job events are joined by pipeline and job identity. Riot setup and activation
 time outside the Test Visibility session is distributed over the atomic hashes in
@@ -155,6 +166,13 @@ model must remain below the size limit in the allocation policy. Failed and canc
 observations are retained as censored reliability evidence but are not treated as
 normal durations.
 
+Sub-hash expansion is fail-closed. A command is eligible only when its command
+fingerprint has real Test Visibility item-count evidence and the fitted runtime
+exceeds the target. Fallback estimates never create slices because they may describe
+a one-test or non-pytest command. Each slice writes a compact collection inventory.
+Across all slice artifacts, verification requires the same collection digest, an
+exact disjoint union, and no empty slice.
+
 The balanced strategy targets five minutes of modeled work per Riot shard. Promotion
 requires at least a 50 percent reduction in the median Riot critical path over paired
 live shadow runs. This objective covers the generated Riot child pipeline, not the
@@ -165,8 +183,10 @@ Balanced sizing is constrained to the legacy topology's total job count. When th
 duration target requests more jobs, the planner removes shards with the smallest
 modeled critical-path penalty, allowing capacity to move between semantic suites
 without increasing the job budget.
-If one atomic Riot environment exceeds the promotion target, the allocator must fail
-the gate; meeting that target requires a separately validated finer execution unit.
+Live critical-path measurement uses the actual interval from the first Riot job start
+to the last Riot job completion. It does not add semantic-stage maxima because the
+generated ``needs`` DAG lets those stages overlap. Modeled maxima are planning scores,
+not measured CI runtimes; promotion uses completed same-head shadow runs.
 
 Use the allocation helper to normalize an export, build a candidate model, and
 replay it against the untouched holdout:
@@ -221,6 +241,16 @@ filename as a fallback for parallel runs that omit test suite properties:
         --balanced balanced/test-results/junit*.xml \
         --output junit-parity.json
 
+For every sub-hash command, also download the runtime inventory artifacts and prove
+that all collected pytest items were executed exactly once:
+
+.. code-block:: bash
+
+    $ scripts/ci_allocation_cli.py verify-runtime-shards \
+        --plan ci-allocation-plan.json \
+        --manifests balanced/test-results/ci-test-shard-inventory.*.json \
+        --output runtime-shard-parity.json
+
 Promotion requires the checked thresholds for sample count, median improvement,
 p75, p90, runner time, clean-success rate, and retry rate. A scheduled retuning task
 may propose a new model, but it must not change the active strategy automatically.