[ci] Generate frontend GPU CI shards from a template

This uses the templating to generate the frontend GPU shards so we can easily control the number of shards
apache · Mar 23, 2022 · 1536273 · 1536273
1 parent 120592b
commit 1536273
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 33 deletions.
diff --git a/Jenkinsfile b/Jenkinsfile
@@ -40,7 +40,9 @@
 //
 
 // ============================= IMPORTANT NOTE =============================
-// This file is generated by jenkins/generate.py. Do not edit this file directly!
+// This file is generated by 'jenkins/generate.py'. Do not edit this file directly!
+// Make edits to 'jenkins/Jenkinsfile.j2' and regenerate this with
+// 'python3 jenkins/generate.py'
 
 import org.jenkinsci.plugins.pipeline.modeldefinition.Utils
 // NOTE: these lines are scanned by docker/dev_common.sh. Please update the regex as needed. -->
@@ -744,15 +746,15 @@ stage('Test') {
               ci_setup(ci_gpu)
               sh (
                 script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_frontend.sh 1",
-                label: 'Run Python frontend tests (shard 1)',
+                label: 'Run Python frontend tests (shard 1 of 2)',
               )
             }
           } finally {
             junit 'build/pytest-results/*.xml'
           }
         }
       }
-     } else {
+    } else {
       Utils.markStageSkippedForConditional('frontend: GPU 1')
     }
   },
@@ -767,15 +769,15 @@ stage('Test') {
               ci_setup(ci_gpu)
               sh (
                 script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_frontend.sh 2",
-                label: 'Run Python frontend tests (shard 2)',
+                label: 'Run Python frontend tests (shard 2 of 2)',
               )
             }
           } finally {
             junit 'build/pytest-results/*.xml'
           }
         }
       }
-     } else {
+    } else {
       Utils.markStageSkippedForConditional('frontend: GPU 2')
     }
   },

diff --git a/jenkins/Jenkinsfile.j2 b/jenkins/Jenkinsfile.j2
@@ -732,7 +732,9 @@ stage('Test') {
       Utils.markStageSkippedForConditional('topi: GPU')
     }
   },
-  'frontend: GPU 1': {
+  {% set num_frontend_gpu_shards = 2 %}
+  {% for frontend_gpu_idx in range(1, num_frontend_gpu_shards + 1) %}
+  'frontend: GPU {{ frontend_gpu_idx }}': {
     if (!skip_ci && is_docs_only_build != 1) {
       node('GPU') {
         ws({{ m.per_exec_ws('tvm/frontend-python-gpu') }}) {
@@ -742,42 +744,20 @@ stage('Test') {
             timeout(time: max_time, unit: 'MINUTES') {
               ci_setup(ci_gpu)
               sh (
-                script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_frontend.sh 1",
-                label: 'Run Python frontend tests (shard 1)',
+                script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_frontend.sh {{ frontend_gpu_idx }}",
+                label: 'Run Python frontend tests (shard {{ frontend_gpu_idx }} of {{ num_frontend_gpu_shards }})',
               )
             }
           } finally {
             junit 'build/pytest-results/*.xml'
           }
         }
       }
-     } else {
-      Utils.markStageSkippedForConditional('frontend: GPU 1')
-    }
-  },
-  'frontend: GPU 2': {
-    if (!skip_ci && is_docs_only_build != 1) {
-      node('GPU') {
-        ws({{ m.per_exec_ws('tvm/frontend-python-gpu') }}) {
-          try {
-            init_git()
-            unpack_lib('gpu', tvm_multilib)
-            timeout(time: max_time, unit: 'MINUTES') {
-              ci_setup(ci_gpu)
-              sh (
-                script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_frontend.sh 2",
-                label: 'Run Python frontend tests (shard 2)',
-              )
-            }
-          } finally {
-            junit 'build/pytest-results/*.xml'
-          }
-        }
-      }
-     } else {
-      Utils.markStageSkippedForConditional('frontend: GPU 2')
+    } else {
+      Utils.markStageSkippedForConditional('frontend: GPU {{ frontend_gpu_idx }}')
     }
   },
+  {% endfor %}
   'frontend: CPU': {
     if (!skip_ci && is_docs_only_build != 1) {
       node('CPU') {