[None][infra] take test durations into account to determine cbts splits num

jenkins/L0_MergeRequest.groovy

@@ -706,7 +706,8 @@ def getAutoTriggerTagList(pipeline, testFilter, globalVars) {
 //
 // Calls jenkins/scripts/cbts/main.py with PR changed_files + diffs and returns
 // a result map (or null = defer to existing filter chain). Result keys:
-// scope, affected_stages, reasons, test_db_dir_override, affected_stage_test_counts.
+// scope, affected_stages, reasons, test_db_dir_override,
+// affected_stage_test_counts, affected_stage_split_counts.
 // CBTS narrows test cases only — Build always runs. See cbts/README.md.
 // ============================================================================
 
@@ -780,8 +781,9 @@ def getCbtsResult(pipeline, testFilter, globalVars)
         // download fails.
         if (result.test_db_dir_override) {
             try {
-                sh "tar czf /tmp/cbts_test_db.tar.gz -C ${LLM_ROOT} ${result.test_db_dir_override}"
-                trtllm_utils.uploadArtifacts("/tmp/cbts_test_db.tar.gz", "${UPLOAD_PATH}/cbts/")
+                // Tar/upload from the workspace: rtUpload patterns are workspace-relative; /tmp uploads 0 files.
+                sh "tar czf ${LLM_ROOT}/cbts_test_db.tar.gz -C ${LLM_ROOT} ${result.test_db_dir_override}"
+                trtllm_utils.uploadArtifacts("${LLM_ROOT}/cbts_test_db.tar.gz", "${UPLOAD_PATH}/cbts/")
                 result.cbts_test_db_artifact_path = "${UPLOAD_PATH}/cbts/cbts_test_db.tar.gz"
                 pipeline.echo("CBTS Layer 3: uploaded cbts_test_db to ${result.cbts_test_db_artifact_path}")
             } catch (InterruptedException e) {
@@ -889,6 +891,7 @@ def _cbtsParseSelectionResult(String text)
         reasons: data.reasons ?: [],
         test_db_dir_override: data.test_db_dir_override,
         affected_stage_test_counts: data.affected_stage_test_counts ?: [:],
+        affected_stage_split_counts: data.affected_stage_split_counts ?: [:],
         // Explicit null check preserves `false`; default True is safe.
         sanity_required: data.sanity_required != null ? data.sanity_required : true,
         perfsanity_required: data.perfsanity_required != null ? data.perfsanity_required : true,

jenkins/L0_Test.groovy

@@ -1703,38 +1703,37 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG
     }
 }
 
-// CBTS Layer 2.5: a narrowed stage runs K = min(narrowed count, default splits) shards. Returns [skip, splits, splitId].
-def _cbtsMaybeCollapseSplits(stageName, splitId, splits) {
+// CBTS Layer 2.5: rename narrowed stages (reuse-safety) and resize their splits to k.
+def cbtsResizeSplits(configs) {
     def cbts = testFilter[(CBTS_RESULT)]
-    // No CBTS or tarball not uploaded (renderTestDB then runs full source) -> keep default splits.
     if (cbts == null || !cbts.cbts_test_db_artifact_path) {
-        return [skip: false, splits: splits, splitId: splitId]
+        return configs
     }
-    def counts = cbts.affected_stage_test_counts
-    def count = counts ? counts[stageName] : null
-    if (count == null) {
-        return [skip: false, splits: splits, splitId: splitId]
+    def kByStage = cbts.affected_stage_split_counts
+    if (!kByStage) {
+        return configs
     }
-    int k = Math.max(1, Math.min(count as int, splits as int))
-    if (k >= splits) {
-        return [skip: false, splits: splits, splitId: splitId]
-    }
-    if (splitId > k) {
-        echo "CBTS [${cbts.scope}]: ${stageName} narrowed to ${count} tests → ${k}/${splits} shards, skipping group ${splitId}"
-        return [skip: true, splits: splits, splitId: splitId]
+    def resized = [:]
+    configs.each { key, values ->
+        def k = kByStage[key]
+        if (k == null) {
+            resized[key] = values
+            return
+        }
+        int kk = Math.max(1, k as int)
+        if ((values[2] as int) > kk) {
+            echo "CBTS [${cbts.scope}]: ${key} narrowed -> ${kk} shard(s); dropping group ${values[2]}"
+            return
+        }
+        def v = values.collect()
+        v[3] = kk
+        resized[key + CBTS_STAGE_SUFFIX] = v
     }
-    return [skip: false, splits: k, splitId: splitId]
+    return resized
 }
 
 def runLLMTestlistOnSlurm(pipeline, platform, testList, config=VANILLA_CONFIG, perfMode=false, stageName="Undefined", splitId=1, splits=1, gpuCount=1, nodeCount=1, runWithSbatch=false, skipInstallWheel=false, cpver="cp312", String outerAttemptTag="", boolean useClusterDurations=false)
 {
-  def collapse = _cbtsMaybeCollapseSplits(stageName, splitId, splits)
-  if (collapse.skip) {
-    return
-  }
-  splits = collapse.splits
-  splitId = collapse.splitId
-
   echo "Run Slurm job with native sbatch: $runWithSbatch"
 
   def attempt = 0
@@ -1880,6 +1879,10 @@ def DEBUG_MODE = "debug"
 def DETAILED_LOG = "detailed_log"
 @Field
 def CBTS_RESULT = "cbts_result"
+// Suffix for CBTS-narrowed stages so their results aren't reused by non-CBTS runs.
+// A suffix (not prefix) keeps the GPU type as the first '-' token for positional parsers.
+@Field
+def CBTS_STAGE_SUFFIX = "-cbts"
 @Field
 def testFilter = [
     (REUSE_TEST): null,
@@ -2867,17 +2870,14 @@ def renderTestDB(pipeline, testContext, llmSrc, stageName, preDefinedMakoOpts=nu
     // renderTestDB falls back to the source test-db.
     def cbts = testFilter[(CBTS_RESULT)]
     if (cbts != null && cbts.test_db_dir_override && cbts.cbts_test_db_artifact_path) {
-        def overrideDir = "${llmSrc}/${cbts.test_db_dir_override}"
-        def dirExists = sh(returnStdout: true, script: "test -d ${overrideDir} && echo yes || echo no").trim()
-        if (dirExists != "yes") {
-            try {
-                def artifactUrl = "${URM_ARTIFACTORY_BASE}/${cbts.cbts_test_db_artifact_path}"
-                trtllm_utils.llmExecStepWithRetry(pipeline, script: "wget -nv '${artifactUrl}' -O /tmp/cbts_test_db.tar.gz && tar xzf /tmp/cbts_test_db.tar.gz -C ${llmSrc}")
-                echo "CBTS Layer 3: extracted cbts_test_db from artifact"
-            } catch (Exception e) {
-                echo "CBTS Layer 3: artifact download failed " +
-                     "(${e.class.simpleName}: ${e.message}); falling back to source test-db"
-            }
+        try {
+            // Always re-fetch: a reused workspace may hold a stale cbts_test_db/ shadowing this build's YAMLs.
+            def artifactUrl = "${URM_ARTIFACTORY_BASE}/${cbts.cbts_test_db_artifact_path}"
+            trtllm_utils.llmExecStepWithRetry(pipeline, script: "wget -nv '${artifactUrl}' -O /tmp/cbts_test_db.tar.gz && tar xzf /tmp/cbts_test_db.tar.gz -C ${llmSrc}")
+            echo "CBTS Layer 3: extracted cbts_test_db from artifact"
+        } catch (Exception e) {
+            echo "CBTS Layer 3: artifact download failed " +
+                 "(${e.class.simpleName}: ${e.message}); falling back to source test-db"
         }
     }
     def testDBPath = "${llmSrc}/tests/integration/test_lists/test-db"
@@ -3816,13 +3816,6 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO
 // and junit() for intermediate retryable failures).
 def runLLMTestlistOnPlatform(pipeline, platform, testList, config=VANILLA_CONFIG, perfMode=false, stageName="Undefined", splitId=1, splits=1, skipInstallWheel=false, cpver="cp312", postTag="", typeCheck=false, boolean isFinalAttempt=true, Map retryContext=null, boolean useClusterDurations=false)
 {
-    def collapse = _cbtsMaybeCollapseSplits(stageName, splitId, splits)
-    if (collapse.skip) {
-        return
-    }
-    splits = collapse.splits
-    splitId = collapse.splitId
-
     cacheErrorAndUploadResult(stageName, {
         runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config, perfMode, stageName, splitId, splits, skipInstallWheel, cpver, typeCheck, postTag, useClusterDurations)
     }, {
@@ -4432,6 +4425,7 @@ def launchTestJobs(pipeline, testFilter)
         "RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2": ["rtx-pro-6000d-x4", "l0_rtx_pro_6000", 2, 2, 4],
     ]
 
+    x86TestConfigs = cbtsResizeSplits(x86TestConfigs)
     parallelJobs = x86TestConfigs.collectEntries{key, values -> [key, [createKubernetesPodConfig(LLM_DOCKER_IMAGE, values[0], "amd64", values[4] ?: 1, key.contains("-Perf-")), { attemptTag, isFinalAttempt, retryContext = null ->
         def config = VANILLA_CONFIG
         if (key.contains("single-device")) {
@@ -4519,6 +4513,7 @@ def launchTestJobs(pipeline, testFilter)
         16,
         2
     )
+    x86SlurmTestConfigs = cbtsResizeSplits(x86SlurmTestConfigs)
     fullSet += x86SlurmTestConfigs.keySet()
 
     parallelSlurmJobs = x86SlurmTestConfigs.collectEntries{key, values -> [key, [createKubernetesPodConfig(LLM_DOCKER_IMAGE, "slurm", "amd64"), { attemptTag, isFinalAttempt, retryContext = null ->
@@ -4553,6 +4548,7 @@ def launchTestJobs(pipeline, testFilter)
         // DGX Spark is also named as GB10 Grace Blackwell Superchip.
         "GB10-PyTorch-1": ["gb10x", "l0_gb10", 1, 1],
     ]
+    SBSATestConfigs = cbtsResizeSplits(SBSATestConfigs)
     fullSet += SBSATestConfigs.keySet()
 
     SBSASlurmTestConfigs = [
@@ -4587,6 +4583,7 @@ def launchTestJobs(pipeline, testFilter)
         "GB300-4_GPUs-PyTorch-PerfSanity-Post-Merge-2": ["auto:gb300-x4", "l0_gb300_multi_gpus_perf_sanity", 2, 3, 4],
         "GB300-4_GPUs-PyTorch-PerfSanity-Post-Merge-3": ["auto:gb300-x4", "l0_gb300_multi_gpus_perf_sanity", 3, 3, 4],
     ]
+    SBSASlurmTestConfigs = cbtsResizeSplits(SBSASlurmTestConfigs)
     fullSet += SBSASlurmTestConfigs.keySet()
 
     multiNodesSBSAConfigs = [
@@ -4779,6 +4776,7 @@ def launchTestJobs(pipeline, testFilter)
         36,
         9
     )
+    multiNodesSBSAConfigs = cbtsResizeSplits(multiNodesSBSAConfigs)
     fullSet += multiNodesSBSAConfigs.keySet()
 
     if (env.targetArch == AARCH64_TRIPLE) {
@@ -5200,7 +5198,9 @@ def launchTestJobs(pipeline, testFilter)
     // they only run when explicitly listed in affected_stages.
     def cbts = testFilter[(CBTS_RESULT)]
     if (cbts != null) {
-        def affectedSet = (cbts.affected_stages ?: []) as Set
+        // Match the -cbts rename cbtsResizeSplits applies to narrowed stages.
+        def stageSuffix = cbts.cbts_test_db_artifact_path ? CBTS_STAGE_SUFFIX : ""
+        def affectedSet = (cbts.affected_stages ?: []).collect { it + stageSuffix } as Set
         def needsSanity = cbts.sanity_required
         def needsPerfSanity = cbts.perfsanity_required
         parallelJobsFiltered = parallelJobs.findAll { key, _ ->

jenkins/scripts/cbts/README.md

@@ -12,7 +12,7 @@ CBTS narrows test cases only; Build always runs.
 | Layer | Where | Action |
 |---|---|---|
 | **2. Stage** | `L0_Test.groovy::launchTestJobs` | Set `parallelJobsFiltered` to affected stages plus PackageSanityCheck (kept iff `sanity_required`) and PerfSanity (kept iff `perfsanity_required`). Pure `-Perf-` stages always excluded. Empty affectedSet + nothing force-kept → no-op. |
-| **2.5. Split-collapse** | `L0_Test.groovy::runLLMTestlistOn*` entries | Narrowed test count < 20 → collapse pytest-split to splits=1 (only group 1 runs); else default splits stand. |
+| **2.5. Split-resize** | `L0_Test.groovy::launchTestJobs` (`cbtsResizeSplits`) | Keep only shards `1..k` per narrowed stage, where `k` (duration-sized to ~2h/shard) is `affected_stage_split_counts`. |
 | **3. Within-stage tests** | `L0_Test.groovy::renderTestDB` | Point trt-test-db at the narrowed `cbts_test_db/`. Each affected block's `tests:` is restricted to entries in the filter prefix subtree; unaffected blocks are dropped. |
 
 CBTS only subtracts; anything it can't narrow → fallback to the existing
@@ -180,6 +180,7 @@ Decision JSON:
   ],
   "test_db_dir_override": "cbts_test_db",
   "affected_stage_test_counts": {"A10-PyTorch-1": 5, "A10-PyTorch-2": 5},
+  "affected_stage_split_counts": {"A10-PyTorch-1": 1, "A10-PyTorch-2": 1},
   "sanity_required": false,
   "perfsanity_required": false
 }
@@ -191,8 +192,8 @@ Decision JSON:
   `perfsanity_required`.
 - `test_db_dir_override: null` → no Layer 3 narrowing; trt-test-db reads
   the source test-db.
-- `affected_stage_test_counts` → per-stage post-keep-filter test count for
-  Layer 2.5 split-collapse.
+- `affected_stage_test_counts` → per-stage kept-entry count (telemetry).
+- `affected_stage_split_counts` → per-stage duration-sized split count (Layer 2.5).
 
 ## Cross-job seed for stage agents
 
@@ -211,15 +212,21 @@ absent and `renderTestDB` falls back to the source test-db. Layer 2 still
 applies. The tarball carries only the narrowed YAMLs, so no PR diff text
 travels between jobs.
 
-## Split-collapse heuristic (Layer 2.5)
-
-In `_cbtsMaybeCollapseSplits`, when the stage's narrowed count < 20:
-- `splitId == 1` → run as splits=1 (single agent runs the full list).
-- `splitId > 1` → early return; no agent allocated.
-
-At/above 20, default splits stand. The count is computed by
-`blocks.compute_stage_test_counts` using the same keep filter as
-`write_filtered_test_db`.
+## Split-resize heuristic (Layer 2.5)
+
+`blocks.compute_stage_split_counts` sizes each narrowed stage to
+`k = clamp(ceil(est_seconds / 2h), 1, stage.total_splits)` (cap parsed from
+`L0_Test.groovy`). `est_seconds` sums the `.test_durations` cache over the
+stage's kept entries (exact node-id, else subtree-sum, else average — over-counts
+toward the cap, never under-sizes). `launchTestJobs::cbtsResizeSplits` keeps only
+shards `1..k`; pytest-split then balances within them via `least_duration`.
+
+`cbtsResizeSplits` also renames each narrowed stage with a `-cbts` suffix so its
+narrowed result is never reused (whole-stage `REUSE_STAGE_LIST` or per-test
+`reusePassedTestResults`) by a non-CBTS full run on the same commit. A suffix (not
+a prefix) keeps the GPU type as the first `-` token, so positional stage-name
+parsers need no change; full sanity / PerfSanity stages keep their original names
+and reuse normally.
 
 ## Adding a new rule