ci: balance CLI shards by duration

HavenDV · HavenDV · commit d83d02ab183d · 2026-07-13T23:28:04.000+04:00
diff --git a/scripts/plan-test-shards.py b/scripts/plan-test-shards.py
@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+import csv
+import statistics
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class TestEstimate:
+    name: str
+    duration_seconds: float
+
+
+def read_tests(path: Path) -> tuple[str, ...]:
+    return tuple(
+        line.strip()
+        for line in path.read_text(encoding="utf-8").splitlines()
+        if line.strip()
+    )
+
+
+def read_timings(path: Path) -> dict[str, float]:
+    if not path.exists():
+        return {}
+
+    with path.open(encoding="utf-8", newline="") as stream:
+        return {
+            row["test_name"]: float(row["duration_seconds"])
+            for row in csv.DictReader(stream, delimiter="\t")
+        }
+
+
+def plan_shards(
+    tests: tuple[str, ...],
+    timings: dict[str, float],
+    shard_count: int,
+) -> tuple[tuple[int, TestEstimate], ...]:
+    known_durations = [timings[name] for name in tests if name in timings]
+    fallback_duration = statistics.median(known_durations) if known_durations else 1.0
+    estimates = [
+        TestEstimate(name=name, duration_seconds=timings.get(name, fallback_duration))
+        for name in tests
+    ]
+    estimates.sort(key=lambda test: (-test.duration_seconds, test.name))
+
+    loads = [0.0] * shard_count
+    counts = [0] * shard_count
+    assignments = []
+    for estimate in estimates:
+        shard = min(range(shard_count), key=lambda index: (loads[index], counts[index], index))
+        assignments.append((shard, estimate))
+        loads[shard] += estimate.duration_seconds
+        counts[shard] += 1
+
+    return tuple(sorted(assignments, key=lambda assignment: (assignment[0], assignment[1].name)))
+
+
+def write_plan(path: Path, assignments: tuple[tuple[int, TestEstimate], ...]) -> None:
+    with path.open("w", encoding="utf-8", newline="") as stream:
+        writer = csv.writer(stream, delimiter="\t", lineterminator="\n")
+        writer.writerow(("shard_index", "estimated_seconds", "test_name"))
+        for shard, estimate in assignments:
+            writer.writerow((shard, f"{estimate.duration_seconds:.6f}", estimate.name))
+
+
+def main() -> None:
+    if len(sys.argv) != 5:
+        raise SystemExit(
+            "Usage: plan-test-shards.py <tests.txt> <timings.tsv> <shard-count> <output.tsv>"
+        )
+
+    tests_path = Path(sys.argv[1])
+    timings_path = Path(sys.argv[2])
+    shard_count = int(sys.argv[3])
+    output_path = Path(sys.argv[4])
+    if shard_count <= 0:
+        raise SystemExit("Shard count must be positive.")
+
+    tests = read_tests(tests_path)
+    if not tests:
+        raise SystemExit("At least one test is required.")
+    if len(set(tests)) != len(tests):
+        raise SystemExit("Test names must be unique.")
+
+    assignments = plan_shards(tests, read_timings(timings_path), shard_count)
+    write_plan(output_path, assignments)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/run-test-shard.sh b/scripts/run-test-shard.sh
@@ -13,6 +13,7 @@ shard_count=$3
 configuration=${4:-Release}
 artifacts_root=${TEST_SHARD_ARTIFACTS_DIR:-artifacts/test-shards}
 shard_artifacts="$artifacts_root/shard-$shard_index"
+timings_file=${TEST_SHARD_TIMINGS_FILE:-"$(dirname "$project")/test-timings.tsv"}
 
 if ! [[ $shard_index =~ ^[0-9]+$ && $shard_count =~ ^[1-9][0-9]*$ ]]; then
   echo "Shard index and count must be non-negative integers, with a positive count." >&2
@@ -57,26 +58,35 @@ if [[ -n $duplicate_names ]]; then
   exit 1
 fi
 
+mkdir -p "$shard_artifacts"
+rm -f \
+  "$shard_artifacts/discovered-tests.txt" \
+  "$shard_artifacts/membership.txt" \
+  "$shard_artifacts/metadata.tsv" \
+  "$shard_artifacts/plan.tsv" \
+  "$shard_artifacts/results.trx"
+printf '%s\n' "${tests[@]}" > "$shard_artifacts/discovered-tests.txt"
+python3 scripts/plan-test-shards.py \
+  "$shard_artifacts/discovered-tests.txt" \
+  "$timings_file" \
+  "$shard_count" \
+  "$shard_artifacts/plan.tsv"
+
 filters=()
 selected_tests=()
-for test_index in "${!tests[@]}"; do
-  if (( test_index % shard_count == shard_index )); then
-    selected_tests+=("${tests[$test_index]}")
-    filters+=("Name=${tests[$test_index]}")
+while IFS=$'\t' read -r planned_shard _ test_name; do
+  if [[ $planned_shard == "$shard_index" ]]; then
+    selected_tests+=("$test_name")
+    filters+=("Name=$test_name")
   fi
-done
+done < <(sed '1d' "$shard_artifacts/plan.tsv")
 
 if (( ${#filters[@]} == 0 )); then
   echo "Shard $shard_index has no tests." >&2
   exit 1
 fi
 
 filter=$(IFS='|'; printf '%s' "${filters[*]}")
-mkdir -p "$shard_artifacts"
-rm -f \
-  "$shard_artifacts/membership.txt" \
-  "$shard_artifacts/metadata.tsv" \
-  "$shard_artifacts/results.trx"
 printf '%s\n' "${selected_tests[@]}" > "$shard_artifacts/membership.txt"
 printf 'key\tvalue\nproject\t%s\nshard_index\t%s\nshard_count\t%s\ndiscovered_tests\t%s\nselected_tests\t%s\n' \
   "$project" \
diff --git a/src/tests/AutoSDK.IntegrationTests.Cli/test-timings.tsv b/src/tests/AutoSDK.IntegrationTests.Cli/test-timings.tsv
@@ -0,0 +1,153 @@
+test_name	duration_seconds
+CliProject_GeneratesPackAsToolApiCommands_AndBuilds	56.338681
+CliProject_RespectsMethodNamingConvention	24.345590
+Cli_ProtoInput_ShowsGrpcNotSupportedMessage	9.491544
+Generate_BlackForestLabsSdk	13.765671
+Generate_BlandAISdk	13.586970
+Generate_BraintrustSdk	43.345768
+Generate_BraveSearchSdk	12.709234
+Generate_BrowserbaseSdk	12.396770
+Generate_BufModuleInput_ScaffoldsGrpcClientProject	8.907099
+Generate_CartesiaSdk	21.886657
+Generate_ChromaSdk	14.632093
+Generate_CohereSdk	33.776420
+Generate_ComposioSdk	107.296004
+Generate_CozeSdk	17.891651
+Generate_CursorAgentsSdk	13.041061
+Generate_DIdSdk	107.415319
+Generate_DeepLSdk	17.225703
+Generate_DeepgramMultichannelSdk	11.277537
+Generate_DescriptSdk	13.273608
+Generate_DescriptorSetInput_ScaffoldsGrpcClientProject	11.271982
+Generate_DescriptorSetInput_WithInvalidBinary_ShowsExplicitMessage	5.475935
+Generate_DoclingServeSdk	25.352062
+Generate_DoesNotEmitDoubleDotGeneratedFileNames	8.588101
+Generate_E2BSdk	13.766403
+Generate_ElevenLabsStreamingSdk	82.569579
+Generate_ExaSdk	16.504456
+Generate_FalSdk	17.166083
+Generate_FirecrawlSdk	18.456578
+Generate_FishAudioSdk	12.805229
+Generate_FlowiseSdk	15.730806
+Generate_ForemSdk	11.461942
+Generate_FreeFormObjectSchemas_UseGeneratedWrapperTypes	9.579016
+Generate_GladiaSdk	22.758680
+Generate_GoogleGeminiDiscoverySdk	15.933686
+Generate_GroqSdk	22.502740
+Generate_GroundXSdk	11.229611
+Generate_GuardrailsSdk	13.222305
+Generate_HedraSdk	17.216163
+Generate_HeliconeSdk	64.631989
+Generate_HumanloopSdk	14.545482
+Generate_HumeAISdk	18.937675
+Generate_KlingAISdk	10.700919
+Generate_LabelStudioSdk	59.854471
+Generate_LakeraSdk	11.271660
+Generate_LangfuseSdk	38.671667
+Generate_LangsmithSdk	68.860621
+Generate_LettaSdk	74.055314
+Generate_LlamaParseSdk	62.814768
+Generate_Mem0Sdk	27.468112
+Generate_MeshcapadeSdk	15.221809
+Generate_MeshySdk	11.192649
+Generate_MilvusSdk	15.081205
+Generate_MistralSdk	25.278406
+Generate_MixedbreadSdk	25.295117
+Generate_MurfSdk	17.140185
+Generate_NanonetsSdk	9.110343
+Generate_NightfallAISdk	11.904269
+Generate_NixtlaSdk	11.650309
+Generate_OpenApiAndProtoInput_ProducesMixedModeLayout	10.740149
+Generate_OpenRouterSdk	71.945790
+Generate_OpikSdk	44.472350
+Generate_PhoenixSdk	27.813597
+Generate_PhotoroomSdk	12.434273
+Generate_PortkeySdk	78.529553
+Generate_PredictionGuardSdk	14.019341
+Generate_PromptLayerSdk	32.921365
+Generate_ProtoInput_ScaffoldsGrpcClientProject	13.174959
+Generate_ProtoInput_WithMissingImport_ShowsExplicitMessage	15.428375
+Generate_QdrantSdk	19.669681
+Generate_RekaSdk	12.946108
+Generate_RetellAISdk	53.579624
+Generate_RevAISdk	12.449747
+Generate_RoboflowSdk	15.155737
+Generate_SarvamAISdk	13.408680
+Generate_SerperSdk	10.090231
+Generate_ShotstackSdk	25.184466
+Generate_SpeechmaticsSdk	14.641383
+Generate_StabilityAISdk	13.209048
+Generate_SynthesiaSdk	14.582139
+Generate_TavilySdk	15.583399
+Generate_TavusSdk	11.897048
+Generate_TogetherSdk	13.909542
+Generate_TripoSdk	12.562399
+Generate_TurbopufferSdk	13.014473
+Generate_TwelveLabsSdk	23.819447
+Generate_UpstageSdk	12.515157
+Generate_VapiSdk	133.425813
+Generate_VoyageAISdk	8.805826
+Generate_WeaveSdk	31.820161
+Generate_WeaviateSdk	20.982182
+Generate_WithCaseInsensitiveEnumCollisions_Builds	10.676688
+Generate_WithCloudSigningHelpers_SignsAwsAzureAndTencentRequests	10.721230
+Generate_WithCompositeAllOfRequest_DoesNotFlattenWrapperProperties	9.762113
+Generate_WithDeprecatedGeneratedTypeReferences_SuppressesCS0618_AndBuilds	32.806732
+Generate_WithDerivedAllOfInheritedPropertyCollision_Builds	7.524596
+Generate_WithDiscriminatorOneOfVariantNamedLikeDiscriminatorWithoutMapping_RenamesVariantMember	11.447545
+Generate_WithDiscriminatorOneOf_EmitsTryPickSwitchHelpers_AndBuildsNullableFlow	10.691411
+Generate_WithDuplicateQueryParameterNames_Builds	7.906962
+Generate_WithDynamicMultipartHelpers_BuildsComputedAttachmentParts	15.230473
+Generate_WithEnumAndPrimitiveResponses_UsesSerializerDeserialization_AndBuilds	8.102986
+Generate_WithEnumPathParameter_UsesWireValue_AndBuilds	26.800282
+Generate_WithEnvironmentFactoryOptions_BuildsAndUpdatesSnippets	11.702345
+Generate_WithEvaluationWorkflowHelpers_RunsDatasetScoringFailuresAndBatchPublishCases	12.222121
+Generate_WithFernGroupNameMatchingTagPascalCase_UsesUnsuffixedClientName_AndBuilds	10.583719
+Generate_WithFernIdempotencyHeaders_EmitsFirstClassSdkSupport	16.419805
+Generate_WithFernRequestName_UsesRequestModelName	10.190216
+Generate_WithFernTypeLiteralAndDateTime_UsesTypedOutput	7.727887
+Generate_WithHugeUnionDeprecatedConverterTypes_SuppressesCS0618_AndBuilds	12.329631
+Generate_WithHugeUnionTypeInfoPropertyName_Builds	9.063271
+Generate_WithInjectedApiKeyHeaderSecurityScheme_UsesFriendlyConstructorFileName	10.046297
+Generate_WithInjectedCustomHttpAuthorizationScheme_UsesCustomSchemeName	13.146797
+Generate_WithInlineAllOf_DoesNotEmitReflectionValidation	10.432265
+Generate_WithInlineOneOfEnumReferences_UsesQualifiedOneOfJsonConverter_AndBuilds	10.006503
+Generate_WithInlineOneOfInlineEnumAndString_UsesQualifiedOneOfJsonConverter_AndBuilds	10.478485
+Generate_WithInlineOneOfInlineEnumsAndParentEnum_UsesQualifiedOneOfJsonConverter_AndBuilds	16.658950
+Generate_WithLargeJsonSerializerContext_SplitsIntoChunkResolvers	20.107380
+Generate_WithNamespaceDelimiterAndExcludedDottedModel_UsesExternalType	10.388765
+Generate_WithNamespaceDelimiter_GeneratesNestedModelNamespaces	10.519169
+Generate_WithNoEndpointsOrModels_WarnsAboutEmptyGeneratedSurface	4.845481
+Generate_WithObjectAnyOfAndAllOf_UsesNullableTypeInfo_AndBuilds	37.171824
+Generate_WithObservabilityLifecycleHelpers_BatchesFlushesAndShutsDown	7.752002
+Generate_WithOpenApi31NumericExclusiveBounds_Builds	10.524112
+Generate_WithOpenApi32QueryStringCookieAndAdditionalOperations_BuildsAndEmitsSupport	11.554149
+Generate_WithOpenApi32SequentialMediaTypesAndItemSchema_BuildsAndEmitsStreamingSupport	10.647561
+Generate_WithOpenApiExamples_EmitsGeneratedSnippetManifest	6.989902
+Generate_WithPredictionWorkflowHelpers_RunsWaitPollFailureTimeoutAndCancelCases	10.655617
+Generate_WithPromptTemplateHelpers_RendersAndCachesPromptVersions	10.570712
+Generate_WithQuotedEnumValues_EscapesGeneratedStringLiterals_AndBuilds	43.889339
+Generate_WithRawModelData_BuildsTypedExtensionDataHelpers	8.773767
+Generate_WithRemainingOfficialSecuritySchemes_BuildsAndEmitsSupport	11.250453
+Generate_WithRequestAndClientOptions_EmitsOptionsSupportAndBuilds	9.545742
+Generate_WithRequiredDeprecatedLeafFactoryInputs_DoesNotEmitIncompleteFactory	11.425902
+Generate_WithRequiredNonNullableObjectProperty_FromAllOfRequest_Builds	7.763784
+Generate_WithRequiredNullableAnyOfRequestProperty_Builds	9.742333
+Generate_WithRequiredTypelessQueryParameter_UsesNonNullableStringValue_AndBuilds	16.139352
+Generate_WithSecuritySchemeOverride_ReplacesAuthAndSuppressesDuplicateParameters	8.033548
+Generate_WithServerVariableDefaults_UsesResolvedBaseUrl	16.338994
+Generate_WithSpeakeasyPolling_EmitsWaitHelpersAndBuilds	10.436655
+Generate_WithSuccessResponseHeaders_EmitsResponseWrapperMethod_AndBuilds	8.046789
+Generate_WithSymbolicEnumValues_UsesReadableNamesAndBuilds	10.309868
+Generate_WithSystemNetHttpJsonFlag_OnNetstandard20_WithPackage_Builds	10.875241
+Generate_WithSystemNetHttpJsonFlag_UsesHttpJsonHelpers_AndBuilds	9.685873
+Generate_WithValidationFlag_EmitsDataAnnotations_AndBuilds	7.755394
+Generate_WithWarningRegressionModels_Builds	9.934825
+Generate_WithWebhookVerifier_RunsReplicateStyleVerificationCases	9.774297
+Generate_WriterSdk	14.731776
+Http_ProtoInput_ShowsGrpcNotSupportedMessage	7.317541
+SourceGeneratorsProps_Restore_WithCentralPackageManagement	1.275784
+SourceGeneratorsProps_Restore_WithCentralPackageManagement_AddsOptInMicrosoftExtensionsPackages	7.146797
+SourceGeneratorsProps_Restore_WithoutCentralPackageManagement_AddsOptInMicrosoftExtensionsPackages	2.164988
+Trim_WhenValidationBuildFails_ReturnsNonZeroExitCode	8.032452
+Trim_WithAnalysisModeAll_DoesNotFailGeneratedTemporaryProgram	14.812644
