ci: switch to on-demand instances if spot instance is interrupted (#987)

cc @yongwww

Author: yzh119

Jenkinsfile

@@ -72,6 +72,13 @@ def cancel_previous_build() {
   }
 }
 
+def is_last_build() {
+  // whether it is last build
+  def job = Jenkins.instance.getItem(env.JOB_NAME)
+  def lastBuild = job.getLastBuild()
+  return lastBuild.getNumber() == env.BUILD_NUMBER
+}
+
 def init_git(submodule = false) {
   cleanWs()
   // add retry in case checkout timeouts
@@ -95,57 +102,116 @@ def init_git(submodule = false) {
 //   }
 // }
 
+def run_unittest_CPU_AOT_COMPILE(node_type) {
+  echo "Running CPU AOT Compile Unittest"
+  node(node_type) {
+    ws(per_exec_ws('flashinfer-aot')) {
+      init_git(true)
+      sh(script: "ls -alh", label: 'Show work directory')
+      sh(script: "./scripts/task_show_node_info.sh", label: 'Show node info')
+      sh(script: "${docker_run} --no-gpu ./scripts/task_test_aot_build_import.sh", label: 'Test AOT Build and Import')
+    }
+  }
+}
+
+def shard_run_unittest_GPU(node_type, shard_id) {
+  echo "Running unittest on ${node_type}, shard ${shard_id}"
+  node(node_type) {
+    ws(per_exec_ws('flashinfer-unittest')) {
+      init_git(true) // we need cutlass submodule
+      sh(script: "ls -alh", label: 'Show work directory')
+      sh(script: "./scripts/task_show_node_info.sh", label: 'Show node info')
+      sh(script: "${docker_run} ./scripts/task_jit_run_tests_part${shard_id}.sh", label: 'JIT Unittest Part ${shard_id}')
+    }
+  }
+}
+
 stage('Unittest') {
   cancel_previous_build()
   parallel(
     failFast: true,
     'AOT-Build-Import': {
-      node('CPU-LARGE-SPOT') {
-        ws(per_exec_ws('flashinfer-aot')) {
-          init_git(true)
-          sh(script: "ls -alh", label: 'Show work directory')
-          sh(script: "./scripts/task_show_node_info.sh", label: 'Show node info')
-          sh(script: "${docker_run} --no-gpu ./scripts/task_test_aot_build_import.sh", label: 'Test AOT Build and Import')
+      try {
+        run_unittest_CPU_AOT_COMPILE('CPU-LARGE-SPOT')
+      } catch (Throwable ex) {
+        if (is_last_build()) {
+          // retry if we are currently at last build
+          // mark the current stage as success
+          // and try again via on demand node
+          echo 'Exception during SPOT run ' + ex.toString() + ' retry on-demand'
+          currentBuild.result = 'SUCCESS'
+          run_unittest_CPU_AOT_COMPILE('CPU-LARGE')
+        } else {
+          echo 'Exception during SPOT run ' + ex.toString() + ' exit since it is not last build'
+          throw ex
         }
       }
     },
     'JIT-Unittest-1': {
-      node('GPU-G5-SPOT') {
-        ws(per_exec_ws('flashinfer-unittest')) {
-          init_git(true) // we need cutlass submodule
-          sh(script: "ls -alh", label: 'Show work directory')
-          sh(script: "./scripts/task_show_node_info.sh", label: 'Show node info')
-          sh(script: "${docker_run} ./scripts/task_jit_run_tests_part1.sh", label: 'JIT Unittest Part 1')
+      try {
+        shard_run_unittest_GPU('GPU-G5-SPOT', 1)
+      } catch (Throwable ex) {
+        if (is_last_build()) {
+          // retry if we are currently at last build
+          // mark the current stage as success
+          // and try again via on demand node
+          echo 'Exception during SPOT run ' + ex.toString() + ' retry on-demand'
+          currentBuild.result = 'SUCCESS'
+          shard_run_unittest_GPU('GPU-G5', 1)
+        } else {
+          echo 'Exception during SPOT run ' + ex.toString() + ' exit since it is not last build'
+          throw ex
         }
       }
     },
     'JIT-Unittest-2': {
-      node('GPU-G5-SPOT') {
-        ws(per_exec_ws('flashinfer-unittest')) {
-          init_git(true) // we need cutlass submodule
-          sh(script: "ls -alh", label: 'Show work directory')
-          sh(script: "./scripts/task_show_node_info.sh", label: 'Show node info')
-          sh(script: "${docker_run} ./scripts/task_jit_run_tests_part2.sh", label: 'JIT Unittest Part 2')
+      try {
+        shard_run_unittest_GPU('GPU-G5-SPOT', 2)
+      } catch (Throwable ex) {
+        if (is_last_build()) {
+          // retry if we are currently at last build
+          // mark the current stage as success
+          // and try again via on demand node
+          echo 'Exception during SPOT run ' + ex.toString() + ' retry on-demand'
+          currentBuild.result = 'SUCCESS'
+          shard_run_unittest_GPU('GPU-G5', 2)
+        } else {
+          echo 'Exception during SPOT run ' + ex.toString() + ' exit since it is not last build'
+          throw ex
         }
       }
     },
     'JIT-Unittest-3': {
-      node('GPU-G5-SPOT') {
-        ws(per_exec_ws('flashinfer-unittest')) {
-          init_git(true) // we need cutlass submodule
-          sh(script: "ls -alh", label: 'Show work directory')
-          sh(script: "./scripts/task_show_node_info.sh", label: 'Show node info')
-          sh(script: "${docker_run} ./scripts/task_jit_run_tests_part3.sh", label: 'JIT Unittest Part 3')
+      try {
+        shard_run_unittest_GPU('GPU-G5-SPOT', 3)
+      } catch (Throwable ex) {
+        if (is_last_build()) {
+          // retry if we are currently at last build
+          // mark the current stage as success
+          // and try again via on demand node
+          echo 'Exception during SPOT run ' + ex.toString() + ' retry on-demand'
+          currentBuild.result = 'SUCCESS'
+          shard_run_unittest_GPU('GPU-G5', 3)
+        } else {
+          echo 'Exception during SPOT run ' + ex.toString() + ' exit since it is not last build'
+          throw ex
         }
       }
     },
     'JIT-Unittest-4': {
-      node('GPU-G5-SPOT') {
-        ws(per_exec_ws('flashinfer-unittest')) {
-          init_git(true) // we need cutlass submodule
-          sh(script: "ls -alh", label: 'Show work directory')
-          sh(script: "./scripts/task_show_node_info.sh", label: 'Show node info')
-          sh(script: "${docker_run} ./scripts/task_jit_run_tests_part4.sh", label: 'JIT Unittest Part 4')
+      try {
+        shard_run_unittest_GPU('GPU-G5-SPOT', 4)
+      } catch (Throwable ex) {
+        if (is_last_build()) {
+          // retry if we are currently at last build
+          // mark the current stage as success
+          // and try again via on demand node
+          echo 'Exception during SPOT run ' + ex.toString() + ' retry on-demand'
+          currentBuild.result = 'SUCCESS'
+          shard_run_unittest_GPU('GPU-G5', 4)
+        } else {
+          echo 'Exception during SPOT run ' + ex.toString() + ' exit since it is not last build'
+          throw ex
         }
       }
     }