[CI] Parallelize notebooks & end2end tests

test.sh

@@ -12,8 +12,7 @@ if [ -z $1 ] || [ $1 = "quicktest" ]; then
 elif [ $1 = "full" ]; then
   pytest -m 'not (end2end or sanity_bnn or notebooks)' --junitxml=$CI_PROJECT_DIR/reports/main.xml --html=$CI_PROJECT_DIR/reports/main.html --reruns 1 --dist worksteal -n $PYTEST_PARALLEL
   # The following tests cannot be parallelized to the same degree as most tests:
-  pytest -m 'sanity_bnn' --junitxml=$CI_PROJECT_DIR/reports/sanity_bnn.xml --html=$CI_PROJECT_DIR/reports/sanity_bnn.html --reruns 1 --dist loadgroup -n $PYTEST_PARALLEL
-  pytest -m 'end2end or notebooks' --junitxml=$CI_PROJECT_DIR/reports/end2end.xml --html=$CI_PROJECT_DIR/reports/end2end.html --reruns 1 --dist loadfile -n $PYTEST_PARALLEL
+  pytest -m 'end2end or sanity_bnn or notebooks' --junitxml=$CI_PROJECT_DIR/reports/end2end.xml --html=$CI_PROJECT_DIR/reports/end2end.html --reruns 1 --dist loadgroup -n $PYTEST_PARALLEL
   pytest_html_merger -i $CI_PROJECT_DIR/reports/ -o $CI_PROJECT_DIR/reports/full_test_suite.html
 else
   echo "Unrecognized argument to test.sh: $1"

tests/end2end/test_end2end_bnn_pynq.py

@@ -458,7 +458,7 @@ def pytest_generate_tests(metafunc):
             argvalues_scenario = [x[1] for x in items]
             argvalues.append(
                 pytest.param(
-                    *argvalues_scenario, marks=pytest.mark.xdist_group(name="group_%d" % i)
+                    *argvalues_scenario, marks=pytest.mark.xdist_group(name="bnn_pynq_%d" % i)
                 )
             )
         metafunc.parametrize(argnames, argvalues, ids=idlist, scope="class")

tests/end2end/test_end2end_cybsec_mlp.py

@@ -79,6 +79,7 @@ def forward(self, x):
         return out_final
 
 
+@pytest.mark.xdist_group(name="end2end_cybsec")
 @pytest.mark.end2end
 def test_end2end_cybsec_mlp_export():
     assets_dir = os.environ["FINN_ROOT"] + "/src/finn/qnn-data/cybsec-mlp"
@@ -143,6 +144,7 @@ def test_end2end_cybsec_mlp_export():
     assert model.get_tensor_datatype(first_matmul_w_name) == DataType["INT2"]
 
 
+@pytest.mark.xdist_group(name="end2end_cybsec")
 @pytest.mark.slow
 @pytest.mark.vivado
 @pytest.mark.end2end

tests/end2end/test_end2end_mobilenet_v1.py

@@ -106,6 +106,7 @@
 first_layer_res_type = "dsp"
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.end2end
 def test_end2end_mobilenet_export():
     # export preprocessing
@@ -163,6 +164,7 @@ def test_end2end_mobilenet_export():
     assert os.path.isfile(build_dir + "/end2end_mobilenet_preproc.onnx")
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.end2end
 def test_end2end_mobilenet_tidy_and_merge_with_preproc():
     preproc_model = load_test_checkpoint_or_skip(build_dir + "/end2end_mobilenet_preproc.onnx")
@@ -185,6 +187,7 @@ def test_end2end_mobilenet_tidy_and_merge_with_preproc():
     model.save(build_dir + "/end2end_mobilenet_tidy.onnx")
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.end2end
 def test_end2end_mobilenet_streamline():
     model = load_test_checkpoint_or_skip(build_dir + "/end2end_mobilenet_tidy.onnx")
@@ -214,6 +217,7 @@ def test_end2end_mobilenet_streamline():
     assert len(model.get_nodes_by_op_type("Mul")) == 0  # no Mul ops remain
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.end2end
 def test_end2end_mobilenet_lowering():
     model = load_test_checkpoint_or_skip(build_dir + "/end2end_mobilenet_streamlined.onnx")
@@ -227,6 +231,7 @@ def test_end2end_mobilenet_lowering():
     model.save(build_dir + "/end2end_mobilenet_lowered.onnx")
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.end2end
 def test_end2end_mobilenet_convert_to_hw_layers():
     model = load_test_checkpoint_or_skip(build_dir + "/end2end_mobilenet_lowered.onnx")
@@ -243,6 +248,7 @@ def test_end2end_mobilenet_convert_to_hw_layers():
     model.save(build_dir + "/end2end_mobilenet_hw_layers.onnx")
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.end2end
 def test_end2end_mobilenet_specialize_layers():
     model = load_test_checkpoint_or_skip(build_dir + "/end2end_mobilenet_hw_layers.onnx")
@@ -252,6 +258,7 @@ def test_end2end_mobilenet_specialize_layers():
     model.save(build_dir + "/end2end_mobilenet_specialize_layers.onnx")
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.end2end
 def test_end2end_mobilenet_create_dataflow_partition():
     model = load_test_checkpoint_or_skip(build_dir + "/end2end_mobilenet_specialize_layers.onnx")
@@ -265,6 +272,7 @@ def test_end2end_mobilenet_create_dataflow_partition():
     dataflow_model.save(build_dir + "/end2end_mobilenet_dataflow_model.onnx")
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.end2end
 def test_end2end_mobilenet_folding():
     model = load_test_checkpoint_or_skip(build_dir + "/end2end_mobilenet_dataflow_model.onnx")
@@ -348,6 +356,7 @@ def test_end2end_mobilenet_folding():
     model.save(build_dir + "/end2end_mobilenet_folded.onnx")
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.end2end
 def test_end2end_mobilenet_minimize_bit_width():
     model = load_test_checkpoint_or_skip(build_dir + "/end2end_mobilenet_folded.onnx")
@@ -357,6 +366,7 @@ def test_end2end_mobilenet_minimize_bit_width():
     model.save(build_dir + "/end2end_mobilenet_minimize_bitwidth.onnx")
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.slow
 @pytest.mark.vivado
 @pytest.mark.end2end
@@ -393,6 +403,7 @@ def test_end2end_mobilenet_cppsim():
     # assert np.isclose(golden_prob, res_cppsim_prob[0, 0, 0, :5]).all()
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.slow
 @pytest.mark.vivado
 @pytest.mark.end2end
@@ -403,6 +414,7 @@ def test_end2end_mobilenet_ipgen():
     model.save(build_dir + "/end2end_mobilenet_hw_ipgen.onnx")
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.slow
 @pytest.mark.vivado
 @pytest.mark.end2end
@@ -439,6 +451,7 @@ def test_end2end_mobilenet_rtlsim():
     # assert np.isclose(golden_prob, res_rtlsim_prob[0, 0, 0, :5]).all()
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.slow
 @pytest.mark.vivado
 @pytest.mark.end2end
@@ -465,6 +478,7 @@ def test_end2end_mobilenet_set_fifo_depths():
     model.save(build_dir + "/end2end_mobilenet_set_fifo_depths.onnx")
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.slow
 @pytest.mark.vivado
 @pytest.mark.end2end
@@ -481,6 +495,7 @@ def test_end2end_mobilenet_stitched_ip():
     model.save(build_dir + "/end2end_mobilenet_stitched_ip.onnx")
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.slow
 @pytest.mark.vivado
 @pytest.mark.end2end
@@ -517,6 +532,7 @@ def test_end2end_mobilenet_stitched_ip_rtlsim():
     # assert np.isclose(golden_prob, res_rtlsim_ip_prob[0, 0, 0, :5]).all()
 
 
+@pytest.mark.xdist_group(name="end2end_mobilenet")
 @pytest.mark.slow
 @pytest.mark.vivado
 @pytest.mark.end2end

tests/end2end/test_ext_weights.py

@@ -66,6 +66,7 @@ def get_checkpoint_name(step):
         return build_dir + "/end2end_ext_weights_%s.onnx" % (step)
 
 
+@pytest.mark.xdist_group(name="end2end_ext_weights")
 @pytest.mark.end2end
 def test_end2end_ext_weights_download():
     if not os.path.isfile(onnx_zip_local):
@@ -75,6 +76,7 @@ def test_end2end_ext_weights_download():
     assert os.path.isfile(get_checkpoint_name("download"))
 
 
+@pytest.mark.xdist_group(name="end2end_ext_weights")
 @pytest.mark.slow
 @pytest.mark.vivado
 @pytest.mark.end2end

tests/notebooks/test_jupyter_notebooks.py

@@ -12,28 +12,67 @@
 notebook_bnn_dir = get_finn_root() + "/notebooks/end2end_example/bnn-pynq/"
 
 basics_notebooks = [
-    pytest.param(notebook_basic_dir + "0_how_to_work_with_onnx.ipynb"),
-    pytest.param(notebook_basic_dir + "1_brevitas_network_import_via_QONNX.ipynb"),
+    pytest.param(
+        notebook_basic_dir + "0_how_to_work_with_onnx.ipynb",
+        marks=pytest.mark.xdist_group(name="notebooks_general"),
+    ),
+    pytest.param(
+        notebook_basic_dir + "1_brevitas_network_import_via_QONNX.ipynb",
+        marks=pytest.mark.xdist_group(name="notebooks_general"),
+    ),
 ]
 
 advanced_notebooks = [
-    pytest.param(notebook_advanced_dir + "0_custom_analysis_pass.ipynb"),
-    pytest.param(notebook_advanced_dir + "1_custom_transformation_pass.ipynb"),
-    pytest.param(notebook_advanced_dir + "2_custom_op.ipynb"),
-    pytest.param(notebook_advanced_dir + "3_folding.ipynb"),
-    pytest.param(notebook_advanced_dir + "4_advanced_builder_settings.ipynb"),
+    pytest.param(
+        notebook_advanced_dir + "0_custom_analysis_pass.ipynb",
+        marks=pytest.mark.xdist_group(name="notebooks_general"),
+    ),
+    pytest.param(
+        notebook_advanced_dir + "1_custom_transformation_pass.ipynb",
+        marks=pytest.mark.xdist_group(name="notebooks_general"),
+    ),
+    pytest.param(
+        notebook_advanced_dir + "2_custom_op.ipynb",
+        marks=pytest.mark.xdist_group(name="notebooks_general"),
+    ),
+    pytest.param(
+        notebook_advanced_dir + "3_folding.ipynb",
+        marks=pytest.mark.xdist_group(name="notebooks_general"),
+    ),
+    pytest.param(
+        notebook_advanced_dir + "4_advanced_builder_settings.ipynb",
+        marks=pytest.mark.xdist_group(name="notebooks_general"),
+    ),
 ]
 
 cyber_notebooks = [
-    pytest.param(notebook_cyber_dir + "1-train-mlp-with-brevitas.ipynb"),
-    pytest.param(notebook_cyber_dir + "2-import-into-finn-and-verify.ipynb"),
-    pytest.param(notebook_cyber_dir + "3-build-accelerator-with-finn.ipynb"),
+    pytest.param(
+        notebook_cyber_dir + "1-train-mlp-with-brevitas.ipynb",
+        marks=pytest.mark.xdist_group(name="notebooks_cybsec"),
+    ),
+    pytest.param(
+        notebook_cyber_dir + "2-import-into-finn-and-verify.ipynb",
+        marks=pytest.mark.xdist_group(name="notebooks_cybsec"),
+    ),
+    pytest.param(
+        notebook_cyber_dir + "3-build-accelerator-with-finn.ipynb",
+        marks=pytest.mark.xdist_group(name="notebooks_cybsec"),
+    ),
 ]
 
 bnn_notebooks = [
-    pytest.param(notebook_bnn_dir + "cnv_end2end_example.ipynb"),
-    pytest.param(notebook_bnn_dir + "tfc_end2end_example.ipynb"),
-    pytest.param(notebook_bnn_dir + "tfc_end2end_verification.ipynb"),
+    pytest.param(
+        notebook_bnn_dir + "cnv_end2end_example.ipynb",
+        marks=pytest.mark.xdist_group(name="notebooks_cnv"),
+    ),
+    pytest.param(
+        notebook_bnn_dir + "tfc_end2end_example.ipynb",
+        marks=pytest.mark.xdist_group(name="notebooks_tfc"),
+    ),
+    pytest.param(
+        notebook_bnn_dir + "tfc_end2end_verification.ipynb",
+        marks=pytest.mark.xdist_group(name="notebooks_tfc"),
+    ),
 ]