perf: split Databricks E2E into 3 parallel CPU partitions, tune timeouts (#2527)

Split the monolithic DatabricksCPUTests (~43 min) into 3 partitions
(DatabricksCPUTests1/2/3) that run as separate ADO matrix entries.
Each partition creates its own Databricks cluster and runs ~15 notebooks,
with all 3 starting simultaneously on different ADO agents.

Changes:
- DatabricksCPUTests.scala: Split into 3 test classes using partition indices
- DatabricksUtilities.scala: Add cpuNotebookPartition() with deterministic
  absolute-path-based sorting, add timeoutSeconds parameter threading
- DatabricksGPUTests.scala: Increase GPU workers 2->3 and concurrency 1->3
  (GPU notebooks are independent; parallel execution reduces E2E wall clock)
  Add 30-minute GPU-specific timeout override
- pipeline.yaml: Replace single databricks-cpu with cpu-1/cpu-2/cpu-3 matrix

Author: BrendanWalsh

core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/DatabricksCPUTests.scala

@@ -5,21 +5,31 @@ package com.microsoft.azure.synapse.ml.nbtest
 
 import com.microsoft.azure.synapse.ml.nbtest.DatabricksUtilities._
 
+import java.time.LocalDateTime
 import scala.language.existentials
 
-class DatabricksCPUTests extends DatabricksTestHelper {
+// Split CPU E2E tests into 3 partitions that run as separate ADO matrix entries.
+// Each partition creates its own cluster and contains ~15 notebooks,
+// running up to the configured worker concurrency.
+// All 3 partitions start simultaneously on different ADO agents.
 
-  val clusterId: String = createClusterInPool(ClusterName, AdbRuntime, NumWorkers, PoolId, memory = Some("7g"))
-
-  databricksTestHelper(clusterId, Libraries, CPUNotebooks, 5)
+class DatabricksCPUTests1 extends DatabricksTestHelper {
+  private val clusterName = s"mmlspark-build-cpu1-${LocalDateTime.now()}"
+  val clusterId: String = createClusterInPool(clusterName, AdbRuntime, NumWorkers, PoolId, memory = Some("7g"))
+  databricksTestHelper(clusterId, Libraries, cpuNotebookPartition(0), NumWorkers)
+  protected override def afterAll(): Unit = { afterAllHelper(clusterId, clusterName); super.afterAll() }
+}
 
-  protected override def afterAll(): Unit = {
-    afterAllHelper(clusterId, ClusterName)
-    super.afterAll()
-  }
+class DatabricksCPUTests2 extends DatabricksTestHelper {
+  private val clusterName = s"mmlspark-build-cpu2-${LocalDateTime.now()}"
+  val clusterId: String = createClusterInPool(clusterName, AdbRuntime, NumWorkers, PoolId, memory = Some("7g"))
+  databricksTestHelper(clusterId, Libraries, cpuNotebookPartition(1), NumWorkers)
+  protected override def afterAll(): Unit = { afterAllHelper(clusterId, clusterName); super.afterAll() }
+}
 
-  ignore("list running jobs for convenience") {
-    val obj = databricksGet("jobs/runs/list?active_only=true&limit=1000")
-    println(obj)
-  }
+class DatabricksCPUTests3 extends DatabricksTestHelper {
+  private val clusterName = s"mmlspark-build-cpu3-${LocalDateTime.now()}"
+  val clusterId: String = createClusterInPool(clusterName, AdbRuntime, NumWorkers, PoolId, memory = Some("7g"))
+  databricksTestHelper(clusterId, Libraries, cpuNotebookPartition(2), NumWorkers)
+  protected override def afterAll(): Unit = { afterAllHelper(clusterId, clusterName); super.afterAll() }
 }

core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/DatabricksGPUTests.scala

@@ -7,9 +7,12 @@ import com.microsoft.azure.synapse.ml.nbtest.DatabricksUtilities._
 
 class DatabricksGPUTests extends DatabricksTestHelper {
 
-  val clusterId: String = createClusterInPool(GPUClusterName, AdbGpuRuntime, 2, GpuPoolId)
+  // GPU fine-tuning notebooks can take up to 25 min; use 30 min timeout
+  private val gpuTimeoutMs = 30 * 60 * 1000
 
-  databricksTestHelper(clusterId, GPULibraries, GPUNotebooks, 1, List())
+  val clusterId: String = createClusterInPool(GPUClusterName, AdbGpuRuntime, 3, GpuPoolId)
+
+  databricksTestHelper(clusterId, GPULibraries, GPUNotebooks, 3, List(), gpuTimeoutMs)
 
   protected override def afterAll(): Unit = {
     afterAllHelper(clusterId, GPUClusterName)

core/src/test/scala/com/microsoft/azure/synapse/ml/nbtest/DatabricksUtilities.scala

@@ -97,7 +97,10 @@ object DatabricksUtilities {
   ).toJson.compactPrint
 
   // Execution Params
-  val TimeoutInMillis: Int = 50 * 60 * 1000
+  // Per-notebook timeout: 20 min for most notebooks.
+  // Observed max actual execution time is ~7 min for CPU notebooks.
+  // GPU fine-tuning notebooks can take up to 25 min, so GPU tests override this.
+  val TimeoutInMillis: Int = 20 * 60 * 1000
 
   val DocsDir = FileUtilities.join(BuildInfo.baseDirectory.getParent, "docs").getCanonicalFile()
   val NotebookFiles: Array[File] = FileUtilities.recursiveListFiles(DocsDir)
@@ -119,6 +122,15 @@ object DatabricksUtilities {
     .filterNot(_.getAbsolutePath.contains("Flooding Risk")) // Azure Maps Spatial API retired 9/30/2025
     .filterNot(_.getAbsolutePath.contains("Geospatial Services")) // Azure Maps Spatial API retired 9/30/2025
 
+  // Split CPU notebooks into 3 partitions for parallel ADO matrix jobs.
+  // Each partition creates its own cluster, so all 3 run simultaneously.
+  // Sort by absolute path for stable, deterministic partitioning across machines.
+  private val SortedCPUNotebooks = CPUNotebooks.sortBy(_.getAbsolutePath)
+  val NumCPUPartitions = 3
+  def cpuNotebookPartition(partIndex: Int): Seq[File] = {
+    SortedCPUNotebooks.zipWithIndex.filter(_._2 % NumCPUPartitions == partIndex).map(_._1)
+  }
+
   val GPUNotebooks: Seq[File] = ParallelizableNotebooks.filter { file =>
     file.getAbsolutePath.contains("Fine-tune") || file.getAbsolutePath.contains("Phi Model")
   }
@@ -273,13 +285,14 @@ object DatabricksUtilities {
     ()
   }
 
-  def submitRun(clusterId: String, notebookPath: String): Long = {
+  def submitRun(clusterId: String, notebookPath: String,
+                timeoutSeconds: Int = TimeoutInMillis / 1000): Long = {
     val body =
       s"""
          |{
          |  "run_name": "test1",
          |  "existing_cluster_id": "$clusterId",
-         |  "timeout_seconds": ${TimeoutInMillis / 1000},
+         |  "timeout_seconds": $timeoutSeconds,
          |  "notebook_task": {
          |    "notebook_path": "$notebookPath",
          |    "base_parameters": []
@@ -367,15 +380,16 @@ object DatabricksUtilities {
     }
   }
 
-  def runNotebook(clusterId: String, notebookFile: File): Unit = {
+  def runNotebook(clusterId: String, notebookFile: File,
+                  timeoutSeconds: Int = TimeoutInMillis / 1000): Unit = {
     val dirPaths = DocsDir.toURI.relativize(notebookFile.getParentFile.toURI).getPath
     val folderToCreate = Folder + "/" + dirPaths
     println(s"Creating folder $folderToCreate")
     workspaceMkDir(folderToCreate)
     val destination: String = folderToCreate + notebookFile.getName
     uploadNotebook(notebookFile, destination)
-    val runId: Long = submitRun(clusterId, destination)
-    val run: DatabricksNotebookRun = DatabricksNotebookRun(runId, notebookFile.getName)
+    val runId: Long = submitRun(clusterId, destination, timeoutSeconds)
+    val run: DatabricksNotebookRun = DatabricksNotebookRun(runId, notebookFile.getName, timeoutSeconds * 1000)
     println(s"Successfully submitted job run id ${run.runId} for notebook ${run.notebookName}")
     DatabricksState.JobIdsToCancel.append(run.runId)
     run.monitor(logLevel = 0)
@@ -439,17 +453,20 @@ abstract class DatabricksTestHelper extends TestBase {
                            libraries: String,
                            notebooks: Seq[File],
                            maxConcurrency: Int,
-                           retries: List[Int] = List(1000 * 15)): Unit = {
+                           retries: List[Int] = List(1000 * 15),
+                           timeoutMs: Int = TimeoutInMillis): Unit = {
 
     println("Checking if cluster is active")
-    tryWithRetries(Seq.fill(60 * 20)(1000).toArray) { () =>
+    // Pool-backed clusters start in ~1.5-3.5 min; allow up to 10 min
+    tryWithRetries(Seq.fill(60 * 10)(1000).toArray) { () =>
       assert(isClusterActive(clusterId))
     }
 
     Thread.sleep(1000) // Ensure cluster is not overwhelmed
     println("Installing libraries")
     installLibraries(clusterId, libraries)
-    tryWithRetries(Seq.fill(60 * 6)(1000).toArray) { () =>
+    // Library install typically takes ~1.5-3 min; allow up to 10 min
+    tryWithRetries(Seq.fill(60 * 10)(1000).toArray) { () =>
       assert(areLibrariesInstalled(clusterId))
     }
 
@@ -461,13 +478,13 @@ abstract class DatabricksTestHelper extends TestBase {
     val futures = notebooks.map { notebook =>
       Future {
         retry(retries, { () =>
-          runNotebook(clusterId, notebook)
+          runNotebook(clusterId, notebook, timeoutMs / 1000)
         })
       }
     }
     futures.zip(notebooks).foreach { case (f, nb) =>
       test(nb.getName) {
-        Await.result(f, Duration(TimeoutInMillis.toLong, TimeUnit.MILLISECONDS))
+        Await.result(f, Duration(timeoutMs.toLong, TimeUnit.MILLISECONDS))
       }
     }
 
@@ -482,8 +499,8 @@ abstract class DatabricksTestHelper extends TestBase {
   }
 }
 
-case class DatabricksNotebookRun(runId: Long, notebookName: String) {
+case class DatabricksNotebookRun(runId: Long, notebookName: String, timeoutMs: Int = TimeoutInMillis) {
   def monitor(logLevel: Int = 2): Unit = {
-    monitorJob(runId, TimeoutInMillis, logLevel)
+    monitorJob(runId, timeoutMs, logLevel)
   }
 }

pipeline.yaml

@@ -174,8 +174,12 @@ jobs:
     vmImage: $(UBUNTU_VERSION)
   strategy:
     matrix:
-      databricks-cpu:
-        TEST-CLASS: "com.microsoft.azure.synapse.ml.nbtest.DatabricksCPUTests"
+      databricks-cpu-1:
+        TEST-CLASS: "com.microsoft.azure.synapse.ml.nbtest.DatabricksCPUTests1"
+      databricks-cpu-2:
+        TEST-CLASS: "com.microsoft.azure.synapse.ml.nbtest.DatabricksCPUTests2"
+      databricks-cpu-3:
+        TEST-CLASS: "com.microsoft.azure.synapse.ml.nbtest.DatabricksCPUTests3"
       databricks-gpu:
         TEST-CLASS: "com.microsoft.azure.synapse.ml.nbtest.DatabricksGPUTests"
 #      databricks-rapids tests have been disabled because these tests are failing.