feat: Configure driver to support GPU tests

README.md

@@ -24,6 +24,7 @@ found in the [Run the tests](#run-the-tests) section.
       * [A subset of UATs](#run-a-subset-of-uats)
       * [Kubeflow UATs](#run-kubeflow-uats)
       * [MLflow UATs](#run-mlflow-uats)
+      * [Include NVIDIA GPU UATs](#include-nvidia-gpu-uats)
    * [Behind proxy](#run-behind-proxy)
       * [Prerequisites for KServe UATs](#prerequisites-for-kserve-uats)
       * [From inside a notebook](#running-using-notebook)
@@ -170,6 +171,19 @@ tox -e mlflow-remote
 tox -e mlflow-local
 ```
 
+#### Include NVIDIA GPU UATs
+
+By default, [GPU UATs](./tests/notebooks/gpu/) are not included in any of the `tox` environments since they require a cluster with a GPU. In order to include those, use the `--include-gpu-tests` flag, e.g.
+
+```bash
+# run all tests defined by tox environment `kubeflow` plus those under the 'gpu' directory
+tox -e kubeflow-remote -- --include-gpu-tests
+# run all tests containing 'kfp' in their name (both cpu and gpu ones)
+tox -e uats-remote -- --include-gpu-tests --filter "kfp"
+```
+
+As shown in the example above, tests under the `gpu` directory follow the same filters with the rest of the tests.
+
 ### Run behind proxy
 
 #### Prerequisites for KServe UATs

driver/conftest.py

@@ -7,8 +7,12 @@
 def pytest_addoption(parser: Parser):
     """Add pytest options.
 
+    * Add a `--proxy` option that enables setting `http_proxy`, `https_proxy` and
+      `no_proxy` environment variables.
     * Add a `--filter` option to (de)select test cases based on their name (see also
       https://docs.pytest.org/en/7.4.x/reference/reference.html#command-line-flags)
+    * Add an `--include-gpu-tests` flag to include the tests under the `gpu` directory
+      in the executed tests.
     """
     parser.addoption(
         "--proxy",
@@ -29,3 +33,9 @@ def pytest_addoption(parser: Parser):
         " any test that doesn't contain 'kserve' in its name. Essentially, the option simulates"
         " the behaviour of running `pytest -k '<filter>'` directly on the test suite.",
     )
+    parser.addoption(
+        "--include-gpu-tests",
+        action="store_true",
+        help="Defines whether to include the tests under the `gpu` directory in the executed tests."
+        "By default, it is set to False.",
+    )

driver/test_kubeflow_workloads.py

@@ -65,6 +65,12 @@ def pytest_filter(request):
     return f"-k '{filter}'" if filter else ""
 
 
+@pytest.fixture(scope="session")
+def include_gpu_tests(request):
+    """Retrieve the `--include-gpu-tests` flag from Pytest invocation."""
+    return True if request.config.getoption("--include-gpu-tests") else False
+
+
 @pytest.fixture(scope="session")
 def tests_checked_out_commit(request):
     """Retrieve active git commit."""
@@ -73,9 +79,14 @@ def tests_checked_out_commit(request):
 
 
 @pytest.fixture(scope="session")
-def pytest_cmd(pytest_filter):
+def pytest_cmd(pytest_filter, include_gpu_tests):
     """Format the Pytest command."""
-    return f"{PYTEST_CMD_BASE} {pytest_filter}" if pytest_filter else PYTEST_CMD_BASE
+    cmd = PYTEST_CMD_BASE
+    if pytest_filter:
+        cmd += f" {pytest_filter}"
+    if include_gpu_tests:
+        cmd += " --include-gpu-tests"
+    return cmd
 
 
 @pytest.fixture(scope="module")

tests/README.md

@@ -36,3 +36,10 @@ pytest -k "kfp or katib"
 # run any test that doesn't contain 'kserve' in its name
 pytest -k "not kserve"
 ```
+
+### NVIDIA GPU tests
+By default, [GPU UATs](./notebooks/gpu/) are not included when running `pytest` since they require a cluster with a GPU. In order to include those, use the `--include-gpu-tests` flag, e.g.
+
+```
+pytest --include-gpu-tests
+```

tests/conftest.py

@@ -0,0 +1,23 @@
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+import os
+
+from _pytest.config.argparsing import Parser
+
+
+def pytest_addoption(parser: Parser):
+    """Add pytest options.
+    * Add an `--include-gpu-tests` flag to include the tests under the `gpu` directory
+      in the executed tests.
+    """
+    parser.addoption(
+        "--include-gpu-tests",
+        action="store_true",
+        help="Defines whether to include tests under the `gpu` directory in the executed tests."
+        "By default, it is set to False.",
+    )
+
+
+def pytest_configure(config):
+    os.environ["include_gpu_tests"] = str(config.getoption("--include-gpu-tests"))

tests/notebooks/gpu/placeholder/placeholder.ipynb

@@ -0,0 +1,21 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\"Placeholder test here\")\n",
+    "assert True"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tests/test_notebooks.py

@@ -15,8 +15,12 @@
     save_notebook,
 )
 
-EXAMPLES_DIR = "notebooks"
-NOTEBOOKS = discover_notebooks(EXAMPLES_DIR)
+EXAMPLES_DIR = {"cpu": "notebooks/cpu", "gpu": "notebooks/gpu"}
+INCLUDE_GPU_TESTS = os.getenv("include_gpu_tests").lower() == "true"
+
+NOTEBOOKS = discover_notebooks(EXAMPLES_DIR["cpu"])
+if INCLUDE_GPU_TESTS:
+    NOTEBOOKS.update(discover_notebooks(EXAMPLES_DIR["gpu"]))
 
 log = logging.getLogger(__name__)
 
@@ -40,6 +44,13 @@ def test_notebook(test_notebook):
     )
     ep.skip_cells_with_tag = "pytest-skip"
 
+    if not INCLUDE_GPU_TESTS:
+        log.info(
+            "Note that only CPU tests will be run. In order to run tests that use an NVIDIA GPU,"
+            "use the `--include-gpu-tests` flag e.g. `tox -e kubeflow-local -- --include-gpu-tests`."
+            " To learn more, use `--help` or refer to the repository's README file."
+        )
+
     try:
         log.info(f"Running {os.path.basename(test_notebook)}...")
         output_notebook, _ = ep.preprocess(notebook, {"metadata": {"path": "./"}})

tests/utils.py

@@ -3,6 +3,7 @@
 
 import os
 import subprocess
+from typing import Dict
 
 import nbformat
 
@@ -17,7 +18,7 @@ def format_error_message(traceback: list):
     return "".join(traceback[-2:])
 
 
-def discover_notebooks(directory):
+def discover_notebooks(directory) -> Dict[str, str]:
     """Return a dictionary of notebooks in the provided directory.
 
     The dictionary contains a mapping between the notebook names (in alphabetical order) and the