[ci] Remove hardcoded test shards

Jenkinsfile

@@ -735,7 +735,7 @@ stage('Test') {
       Utils.markStageSkippedForConditional('topi: GPU')
     }
   },
-  'frontend: GPU 1': {
+  'frontend: GPU 1 of 2': {
     if (!skip_ci && is_docs_only_build != 1) {
       node('GPU') {
         ws("workspace/exec_${env.EXECUTOR_NUMBER}/tvm/frontend-python-gpu") {
@@ -745,7 +745,7 @@ 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",
+                script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_frontend.sh 1 2",
                 label: 'Run Python frontend tests (shard 1 of 2)',
               )
             }
@@ -758,7 +758,7 @@ stage('Test') {
       Utils.markStageSkippedForConditional('frontend: GPU 1')
     }
   },
-  'frontend: GPU 2': {
+  'frontend: GPU 2 of 2': {
     if (!skip_ci && is_docs_only_build != 1) {
       node('GPU') {
         ws("workspace/exec_${env.EXECUTOR_NUMBER}/tvm/frontend-python-gpu") {
@@ -768,7 +768,7 @@ stage('Test') {
             timeout(time: max_time, unit: 'MINUTES') {
               ci_setup(ci_gpu)
               sh (
-                script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_frontend.sh 2",
+                script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_frontend.sh 2 2",
                 label: 'Run Python frontend tests (shard 2 of 2)',
               )
             }

conftest.py

@@ -14,5 +14,34 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import hashlib
+import pytest
+import os
 
 pytest_plugins = ["tvm.testing.plugin"]
+
+
+def should_run(nodeid: str, num_shards: int, shard_index: int) -> bool:
+    """
+    Return true if this test should run on this shard
+    """
+    hash = hashlib.md5(nodeid.encode())
+    hash = int(hash.hexdigest(), 16)
+    print(nodeid, shard_index, hash % num_shards)
+    return hash % num_shards == shard_index
+
+
+def pytest_collection_modifyitems(config, items):
+    if "CI" not in os.environ:
+        return
+
+    # Only apportion tests if in CI and in a job that is set up for it
+    if "NUM_SHARDS" not in os.environ or "SHARD_INDEX" not in os.environ:
+        return
+
+    num_shards = int(os.environ["NUM_SHARDS"])
+    shard_index = int(os.environ["SHARD_INDEX"])
+
+    for item in items:
+        if not should_run(item.nodeid, num_shards=num_shards, shard_index=shard_index):
+            item.add_marker(pytest.mark.skip())

jenkins/Jenkinsfile.j2

@@ -734,7 +734,7 @@ stage('Test') {
   },
   {% set num_frontend_gpu_shards = 2 %}
   {% for frontend_gpu_idx in range(1, num_frontend_gpu_shards + 1) %}
-  'frontend: GPU {{ frontend_gpu_idx }}': {
+  'frontend: GPU {{ frontend_gpu_idx }} of {{ num_frontend_gpu_shards }}': {
     if (!skip_ci && is_docs_only_build != 1) {
       node('GPU') {
         ws({{ m.per_exec_ws('tvm/frontend-python-gpu') }}) {
@@ -744,7 +744,7 @@ stage('Test') {
             timeout(time: max_time, unit: 'MINUTES') {
               ci_setup(ci_gpu)
               sh (
-                script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_frontend.sh {{ frontend_gpu_idx }}",
+                script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_frontend.sh {{ frontend_gpu_idx }} {{ num_frontend_gpu_shards }}",
                 label: 'Run Python frontend tests (shard {{ frontend_gpu_idx }} of {{ num_frontend_gpu_shards }})',
               )
             }

jenkins/macros.j2

@@ -1,3 +1,20 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
 {% macro per_exec_ws(folder) -%}
   "workspace/exec_${env.EXECUTOR_NUMBER}/{{ folder }}"
 {%- endmacro -%}

tests/lint/check_file_type.py

@@ -87,6 +87,8 @@
     "ino",
     # linker scripts
     "ld",
+    # jinja templates
+    "j2",
 }
 
 # List of file names allowed

tests/scripts/task_lint.sh

@@ -29,6 +29,7 @@ echo "Convert scripts to Python..."
 tests/scripts/task_convert_scripts_to_python.sh
 
 echo "Check Jenkinsfile generation"
+python3 -m pip install --user -r jenkins/requirements.txt  # TODO: Remove once the docker images have been udpated
 python3 jenkins/generate.py --check
 
 echo "Checking file types..."

tests/scripts/task_python_frontend.sh

@@ -34,7 +34,7 @@ make cython3
 # The split is purely based on balancing the runtime of each shard so they should
 # be about the same. This may need rebalancing in the future if this is no longer
 # the case.
-function shard1 {
+function run_tests {
     echo "Running relay MXNet frontend test..."
     run_pytest cython python-frontend-mxnet tests/python/frontend/mxnet
 
@@ -43,9 +43,7 @@ function shard1 {
 
     echo "Running relay PyTorch frontend test..."
     run_pytest cython python-frontend-pytorch tests/python/frontend/pytorch
-}
 
-function shard2 {
     echo "Running relay Tensorflow frontend test..."
     # Note: Tensorflow tests often have memory issues, so invoke each one separately
     TENSORFLOW_TESTS=$(./tests/scripts/pytest_ids.py --folder tests/python/frontend/tensorflow)
@@ -74,12 +72,7 @@ if [ -z ${1+x} ]; then
     # TODO: This case can be removed once https://github.com/apache/tvm/pull/10413
     # is merged.
     # No sharding set, run everything
-    shard1
-    shard2
+    run_tests
 else
-    if [ "$1" == "1" ]; then
-        shard1
-    else
-        shard2
-    fi
+    SHARD_INDEX=$1 NUM_SHARDS=$2 run_tests
 fi