ci: add runtime-aware Riot sharding

.gitignore

@@ -203,6 +203,7 @@ cloned_venvs/
 .circleci/config.gen.yml
 # GitLab CI generated config
 .gitlab/**/*-gen.yml
+.gitlab/ci-allocation-plan.json
 .gitlab/benchmarks/bp-runner.microbenchmarks.fail-on-breach.yml
 .gitlab-ci-local/
 

.gitlab-ci.yml

@@ -55,6 +55,11 @@ workflow:
         DD_GIT_REPOSITORY_URL: "https://github.com/DataDog/dd-trace-py-release.git"
       auto_cancel:
         on_new_commit: none
+    # Temporary PR validation: run the balanced allocator beside legacy on the
+    # exact branch under review. Remove after the paired evidence is captured.
+    - if: '$CI_COMMIT_REF_NAME == "dd/ci-runtime-aware-riot-sharding-20260813-1915"'
+      variables:
+        CI_ALLOCATION_SHADOW: "true"
     - when: always
 
 include:
@@ -110,6 +115,7 @@ tests-gen:
   artifacts:
     paths:
       - .gitlab/tests-gen.yml
+      - .gitlab/ci-allocation-plan.json
       - .gitlab/benchmarks/microbenchmarks-gen.yml
       - .gitlab/benchmarks/bp-runner.microbenchmarks.fail-on-breach.yml
 

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

@@ -2,4 +2,55 @@
 set -e -u -o pipefail
 
 SUITE_NAME="${1:-}"
-riot list --hash-only "${SUITE_NAME}" | sort | ./.gitlab/ci-split-input.sh
+CI_ALLOCATION_SUITE="${2:-${CI_ALLOCATION_SUITE:-${SUITE_NAME}}}"
+strategy_args=()
+if [[ -n "${CI_ALLOCATION_STRATEGY:-}" ]]; then
+    strategy_args=(--strategy "${CI_ALLOCATION_STRATEGY}")
+fi
+
+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 "${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,11 +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 _CI_DD_TAGS="test.configuration.riot_hash:${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},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

@@ -0,0 +1,37 @@
+{
+  "allocation": {
+    "active_strategy": "legacy",
+    "maximum_parallelism_per_suite": 50,
+    "maximum_runtime_model_bytes": 100000,
+    "maximum_slices_per_hash": 5,
+    "target_jobs": 200,
+    "target_shard_seconds": 300
+  },
+  "model": {
+    "estimate_quantile": 0.9,
+    "half_life_days": 30,
+    "history_window_days": 90,
+    "holdout_days": 14,
+    "minimum_samples": 5,
+    "sparse_safety_factor": 1.25
+  },
+  "ratchets": {
+    "historical_replay": {
+      "maximum_runner_seconds_increase_ratio": 0.05,
+      "minimum_median_improvement_ratio": 0.05,
+      "minimum_runs": 30
+    },
+    "live_shadow": {
+      "maximum_queue_p90_increase_ratio": 0.05,
+      "maximum_runner_seconds_increase_ratio": 0.05,
+      "minimum_median_improvement_ratio": 0.5,
+      "minimum_runs": 15
+    },
+    "pr_shape_replay": {
+      "maximum_runner_seconds_increase_ratio": 0.05,
+      "minimum_median_improvement_ratio": 0.05,
+      "minimum_shapes": 30
+    }
+  },
+  "schema_version": 1
+}