ci: duration-weighted unit test batching from Actions-cached timings

.github/actions/build-test/action.yml

@@ -59,6 +59,13 @@ runs:
       name: junit-${{ inputs.alias }}-${{ inputs.container-id }}
       path: tmp/rspec/*.xml
 
+  - name: Upload task timings
+    if: ${{ !cancelled() }}
+    uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+    with:
+      name: timings-${{ inputs.alias }}-${{ inputs.container-id }}
+      path: tmp/ci-task-timings/*.json
+
   - name: Upload coverage data
     uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
     with:

.github/scripts/download_timing_artifacts.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+# Downloads the raw timing artifacts uploaded by Unit Tests batch jobs.
+#
+# Contract:
+#   Input:  GH_TOKEN (read access to GITHUB_REPOSITORY)
+#           $1 target directory
+#           $2 branch carrying successful Unit Tests runs (default: master)
+#   Output: $1 populated with one subdirectory per run, each holding the
+#           "timings-*.json" files written by TestBatching::TimingFiles
+#           and uploaded by .github/actions/build-test.
+#   Failure: nonzero exit when gh fails; a partial directory is left behind.
+
+set -euo pipefail
+
+directory="${1:?target directory required}"
+branch="${2:-master}"
+workflow="Unit Tests"
+
+run_ids="$(gh run list --repo "${GITHUB_REPOSITORY}" --workflow "${workflow}" --branch "${branch}" --status success --limit 5 --json databaseId --jq '.[].databaseId')"
+
+for run_id in ${run_ids}; do
+  names="$(gh api --paginate "repos/${GITHUB_REPOSITORY}/actions/runs/${run_id}/artifacts?per_page=100" --jq '.artifacts[] | select(.name | startswith("timings-")) | select(.expired == false) | .name')"
+
+  for name in ${names}; do
+    gh run download "${run_id}" --repo "${GITHUB_REPOSITORY}" --name "${name}" --dir "${directory}/${run_id}"
+  done
+done

.github/workflows/_unit_test.yml

@@ -61,6 +61,28 @@ jobs:
         id: bundle-cache
         uses: ./.github/actions/bundle-cache
 
+      - name: Restore task weights
+        uses: actions/cache/restore@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
+        with:
+          path: tmp/ci_task_timings.json
+          key: spec-task-weights-exact
+          restore-keys: |
+            spec-task-weights-
+
+      # TEMPORARY: source weights from the old repository until the daily
+      # weights job saves the first Actions cache. The file guard keeps the
+      # cache authoritative once it exists. Remove this step once the cache
+      # holds real timings.
+      - name: Fetch task weights
+        env:
+          WEIGHTS_URL: https://raw.githubusercontent.com/TonyCTHsu/dd-test-weights/main/ci_task_timings.json
+        run: |
+          if [ ! -f tmp/ci_task_timings.json ]; then
+            mkdir -p tmp
+            curl -fsSL "$WEIGHTS_URL" -o tmp/ci_task_timings.json
+            ruby -rjson -e 'JSON.parse(File.read("tmp/ci_task_timings.json"))'
+          fi
+
       - id: set-batches
         name: Distribute tasks into batches
         run: |
@@ -235,7 +257,7 @@ jobs:
           DD_DISABLE_ERROR_RESPONSES: "true"
           ENABLED_CHECKS: trace_content_length,trace_stall,meta_tracer_version_header,trace_count_header,trace_peer_service,trace_dd_service
       elasticsearch:
-        image: elasticsearch:8.19.11
+        image: ghcr.io/datadog/images-rb/services/elasticsearch:8.19.11
         env:
           discovery.type: single-node
           xpack.security.enabled: "false"

.github/workflows/update-task-weights.yml

@@ -0,0 +1,43 @@
+name: "Update Task Weights"
+
+on: # yamllint disable-line rule:truthy
+  schedule:
+    - cron: '0 6 * * *'
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}
+  cancel-in-progress: false
+
+# Default permissions for all jobs
+permissions: {}
+
+jobs:
+  update:
+    runs-on: ubuntu-24.04
+    permissions:
+      actions: read
+    env:
+      GH_TOKEN: ${{ github.token }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          persist-credentials: false
+
+      - name: Set up Ruby
+        uses: ruby/setup-ruby@a0102e0972be65f351c307e2d64b9314a57c8073 # v1.324.0
+        with:
+          ruby-version: '3.4'
+
+      - name: Download timing artifacts
+        run: .github/scripts/download_timing_artifacts.sh timings
+
+      - name: Aggregate task weights
+        run: rake -f tasks/github.rake "github:update_task_timings[$PWD/timings]"
+
+      - name: Save weights cache
+        uses: actions/cache/save@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
+        with:
+          path: tmp/ci_task_timings.json
+          key: spec-task-weights-${{ github.run_number }}-${{ github.run_attempt }}

spec/tasks/test_batching/batch_timing_spec.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+require_relative "../../../tasks/lib/test_batching/batch_timing"
+
+RSpec.describe TestBatching::BatchTiming do
+  def task(name, group = "")
+    TestBatching::Task.new(name: name, group: group, gemfile: "gemfiles/#{name}.gemfile")
+  end
+
+  def timing(task, build_seconds: nil, test_seconds: nil, junit_seconds: nil)
+    TestBatching::TaskTiming.new(
+      task: task, build_seconds: build_seconds, test_seconds: test_seconds, junit_seconds: junit_seconds
+    )
+  end
+
+  it "serializes to and from the timings artifact document" do
+    batch_timing = described_class.new(
+      ruby_version: "3.4",
+      timings: [timing(task("main"), build_seconds: 10.0, test_seconds: 30.0, junit_seconds: 25.0)]
+    )
+
+    expect(described_class.from_h(batch_timing.to_h)).to eq(batch_timing)
+  end
+
+  it "omits measurements the run has not made yet" do
+    batch_timing = described_class.new(
+      ruby_version: "3.4",
+      timings: [timing(task("main"), build_seconds: 10.0)]
+    )
+
+    expect(batch_timing.to_h.fetch("tasks").first.keys).to eq(["task", "group", "gemfile", "build_seconds"])
+  end
+
+  it "replaces the timing for a measured task and appends a new one" do
+    prior = described_class.new(
+      ruby_version: "3.4",
+      timings: [timing(task("main"), build_seconds: 10.0), timing(task("contrib"), build_seconds: 5.0)]
+    )
+
+    updated = prior.with_timing(timing(task("main"), build_seconds: 10.0, test_seconds: 30.0))
+
+    expect(updated.timings.map(&:test_seconds)).to eq([30.0, nil])
+    expect(updated.timings.map { |entry| entry.task.name }).to eq(["main", "contrib"])
+    expect(prior.timings.map(&:test_seconds)).to eq([nil, nil])
+  end
+end

spec/tasks/test_batching/manifest_store_spec.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+require "json"
+require "spec_helper"
+require "tmpdir"
+require_relative "../../../tasks/lib/test_batching/manifest_store"
+
+RSpec.describe TestBatching::ManifestStore do
+  it "loads weights from a manifest file" do
+    manifest = {"ruby_versions" => {"3.1" => {"tasks" => [], "gemfiles" => []}}}
+
+    Dir.mktmpdir do |directory|
+      path = File.join(directory, "manifest.json")
+      File.write(path, JSON.pretty_generate(manifest))
+
+      loaded = described_class.load(path: path)
+
+      expect(loaded.to_h).to eq(manifest)
+      expect(loaded).to be_a(TestBatching::Weights)
+    end
+  end
+
+  it "loads missing weights when the file is absent so batching stays static" do
+    Dir.mktmpdir do |directory|
+      loaded = described_class.load(path: File.join(directory, "missing.json"))
+
+      expect(loaded).to be_a(TestBatching::Weights::Missing)
+      expect(loaded).to be_empty
+    end
+  end
+
+  it "loads missing weights when the file is invalid" do
+    Dir.mktmpdir do |directory|
+      path = File.join(directory, "manifest.json")
+      File.write(path, "not json")
+
+      expect(described_class.load(path: path)).to be_a(TestBatching::Weights::Missing)
+    end
+  end
+
+  it "raises when aggregating a directory without timing files" do
+    Dir.mktmpdir do |directory|
+      expect { described_class.update_from(directory) }.to raise_error("no timing files found in #{directory}")
+    end
+  end
+
+  it "aggregates a directory of timing files into the manifest" do
+    tasks = [TestBatching::Task.new(name: "main", group: "", gemfile: "gemfiles/foo.gemfile")]
+    batch_timing = TestBatching::BatchTiming.new(
+      ruby_version: RUBY_VERSION[0..2],
+      timings: [TestBatching::TaskTiming.new(task: tasks.first, build_seconds: 10.0, test_seconds: 30.0)]
+    )
+
+    Dir.mktmpdir do |directory|
+      TestBatching::TimingFiles.write(batch_timing, directory: directory)
+
+      described_class.update_from(directory)
+
+      loaded = described_class.load
+
+      expect(loaded.for_ruby(RUBY_VERSION[0..2]).test_cost(tasks.first)).to eq(30.0)
+    ensure
+      File.delete(described_class::MANIFEST_PATH) if File.exist?(described_class::MANIFEST_PATH)
+    end
+  end
+end

spec/tasks/test_batching/samples_spec.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+require_relative "../../../tasks/lib/test_batching/samples"
+
+RSpec.describe TestBatching::Samples do
+  it "reports the median and 90th percentile of odd-sized values" do
+    samples = described_class.new([30, 10, 20])
+
+    expect(samples.to_h).to eq({"p50_seconds" => 20, "p90_seconds" => 30, "samples" => 3})
+  end
+
+  it "averages the middle pair for an even count" do
+    samples = described_class.new([10, 20, 30, 40])
+
+    expect(samples.to_h).to eq({"p50_seconds" => 25.0, "p90_seconds" => 40, "samples" => 4})
+  end
+
+  it "rounds to milliseconds when serializing" do
+    samples = described_class.new([10.12345, 20.0, 30.5])
+
+    expect(samples.to_h).to eq({"p50_seconds" => 20.0, "p90_seconds" => 30.5, "samples" => 3})
+  end
+end

spec/tasks/test_batching/static_spec.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+require_relative "../../../tasks/lib/test_batching/static"
+require_relative "../../../tasks/lib/test_batching/task_matrix"
+
+RSpec.describe TestBatching::Static do
+  def matrix
+    {
+      "main" => {"" => ["✅ 3.4"]},
+      "contrib" => {"" => ["✅ 3.4"]},
+      "mongodb" => {"" => ["✅ 3.4"]},
+      "skipped" => {"" => ["✅ 3.3"]},
+    }
+  end
+
+  it "even-splits the matching tasks and routes misc services out, with no estimates" do
+    plan = described_class.new.plan(TestBatching::TaskMatrix.new(matrix), "3.4")
+
+    expect(plan.misc_tasks.map(&:name)).to eq(["mongodb"])
+    # Two matching tasks, one per batch: master's batching shape exactly.
+    expect(plan.batches.map { |batch| batch.tasks.map(&:name) }).to eq([["main"], ["contrib"]])
+    expect(plan.batches.map(&:seconds)).to eq([nil, nil])
+  end
+
+  it "emits master's matrix JSON" do
+    data = described_class.new.plan(TestBatching::TaskMatrix.new(matrix), "3.4").to_matrix
+
+    expect(data.keys).to eq(["batches", "misc"])
+    expect(data["batches"]["include"].first.keys).to eq(["batch", "tasks"])
+    expect(data["misc"]["include"].first["tasks"].first.to_h).to include("task" => "mongodb")
+  end
+end

spec/tasks/test_batching/task_matrix_spec.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+require_relative "../../../tasks/lib/test_batching/task_matrix"
+
+RSpec.describe TestBatching::TaskMatrix do
+  before do
+    allow(AppraisalConversion).to receive(:to_bundle_gemfile) do |group|
+      raise "Gemfile not found" if group.to_s.empty?
+
+      "gemfiles/#{group}.gemfile"
+    end
+    allow(AppraisalConversion).to receive(:to_bundle_gemfile).with("missing") { raise "Gemfile not found" }
+    allow(AppraisalConversion).to receive(:parent_gemfile) { "gemfiles/Gemfile" }
+  end
+
+  def matrix
+    {
+      "main" => {"" => ["✅ 3.4", "✅ 3.5"]},
+      "core_with_rails" => {"rails8" => ["✅ 3.4"]},
+      "legacy" => {"missing" => ["✅ 3.4"]},
+      "mongodb" => {"" => ["✅ 3.4"]},
+      "skipped" => {"" => ["✅ 3.3"]},
+    }
+  end
+
+  it "collects the batchable tasks under a Ruby, mapped to their group Gemfiles" do
+    tasks = described_class.new(matrix).batchable_tasks("3.4")
+
+    expect(tasks.map(&:name)).to eq(["main", "core_with_rails", "legacy"])
+    expect(tasks.map(&:gemfile)).to eq(
+      ["gemfiles/Gemfile", "gemfiles/rails8.gemfile", "gemfiles/Gemfile"]
+    )
+  end
+
+  it "splits service-dependent misc tasks from the batchable ones" do
+    task_matrix = described_class.new(matrix)
+
+    expect(task_matrix.misc_tasks("3.4").map(&:name)).to eq(["mongodb"])
+    expect(task_matrix.misc_tasks("3.4").map(&:gemfile)).to eq(["gemfiles/Gemfile"])
+  end
+end

spec/tasks/test_batching/task_spec.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+require_relative "../../../tasks/lib/test_batching/task"
+
+RSpec.describe TestBatching::Task do
+  it "equals tasks with the same name and group, whatever their Gemfile" do
+    task = described_class.new(name: "main", group: "", gemfile: "gemfiles/foo.gemfile")
+    same_task = described_class.new(name: "main", group: "", gemfile: "gemfiles/other.gemfile")
+    other_group = described_class.new(name: "main", group: "rails8", gemfile: "gemfiles/foo.gemfile")
+
+    expect(task).to eq(same_task)
+    expect(task).not_to eq(other_group)
+  end
+
+  it "serializes to and from the matrix JSON shape" do
+    task = described_class.new(name: "main", group: "rails8", gemfile: "gemfiles/foo.gemfile")
+
+    expect(described_class.from_h(task.to_h)).to eq(task)
+  end
+
+  it "knows its Gemfile's name" do
+    task = described_class.new(name: "main", group: "", gemfile: "gemfiles/foo.gemfile")
+
+    expect(task.gemfile_name).to eq("foo.gemfile")
+  end
+end