canonical · orfeas-k · Nov 28, 2024 · Nov 22, 2024 · Nov 26, 2024 · Nov 26, 2024
diff --git a/driver/conftest.py b/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.",
+    )
diff --git a/driver/test_kubeflow_workloads.py b/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")

diff --git a/tests/conftest.py b/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"))
diff --git a/tests/notebooks/e2e-wine/.gitignore → tests/notebooks/cpu/e2e-wine/.gitignore b/tests/notebooks/e2e-wine/.gitignore → tests/notebooks/cpu/e2e-wine/.gitignore
diff --git a/...e2e-wine/e2e-wine-kfp-mlflow-kserve.ipynb → ...e2e-wine/e2e-wine-kfp-mlflow-kserve.ipynb b/...e2e-wine/e2e-wine-kfp-mlflow-kserve.ipynb → ...e2e-wine/e2e-wine-kfp-mlflow-kserve.ipynb
diff --git a/tests/notebooks/e2e-wine/requirements.in → tests/notebooks/cpu/e2e-wine/requirements.in b/tests/notebooks/e2e-wine/requirements.in → tests/notebooks/cpu/e2e-wine/requirements.in
diff --git a/tests/notebooks/e2e-wine/requirements.txt → ...s/notebooks/cpu/e2e-wine/requirements.txt b/tests/notebooks/e2e-wine/requirements.txt → ...s/notebooks/cpu/e2e-wine/requirements.txt
diff --git a/...s/notebooks/katib/katib-integration.ipynb → ...tebooks/cpu/katib/katib-integration.ipynb b/...s/notebooks/katib/katib-integration.ipynb → ...tebooks/cpu/katib/katib-integration.ipynb
diff --git a/tests/notebooks/katib/requirements.in → tests/notebooks/cpu/katib/requirements.in b/tests/notebooks/katib/requirements.in → tests/notebooks/cpu/katib/requirements.in
diff --git a/tests/notebooks/katib/requirements.txt → tests/notebooks/cpu/katib/requirements.txt b/tests/notebooks/katib/requirements.txt → tests/notebooks/cpu/katib/requirements.txt
diff --git a/...notebooks/kfp_v1/kfp-v1-integration.ipynb → ...books/cpu/kfp_v1/kfp-v1-integration.ipynb b/...notebooks/kfp_v1/kfp-v1-integration.ipynb → ...books/cpu/kfp_v1/kfp-v1-integration.ipynb
diff --git a/tests/notebooks/kfp_v1/requirements.in → tests/notebooks/cpu/kfp_v1/requirements.in b/tests/notebooks/kfp_v1/requirements.in → tests/notebooks/cpu/kfp_v1/requirements.in
diff --git a/tests/notebooks/kfp_v1/requirements.txt → tests/notebooks/cpu/kfp_v1/requirements.txt b/tests/notebooks/kfp_v1/requirements.txt → tests/notebooks/cpu/kfp_v1/requirements.txt
diff --git a/...notebooks/kfp_v2/kfp-v2-integration.ipynb → ...books/cpu/kfp_v2/kfp-v2-integration.ipynb b/...notebooks/kfp_v2/kfp-v2-integration.ipynb → ...books/cpu/kfp_v2/kfp-v2-integration.ipynb
diff --git a/tests/notebooks/kfp_v2/requirements.in → tests/notebooks/cpu/kfp_v2/requirements.in b/tests/notebooks/kfp_v2/requirements.in → tests/notebooks/cpu/kfp_v2/requirements.in
diff --git a/tests/notebooks/kfp_v2/requirements.txt → tests/notebooks/cpu/kfp_v2/requirements.txt b/tests/notebooks/kfp_v2/requirements.txt → tests/notebooks/cpu/kfp_v2/requirements.txt
diff --git a/...notebooks/kserve/kserve-integration.ipynb → ...books/cpu/kserve/kserve-integration.ipynb b/...notebooks/kserve/kserve-integration.ipynb → ...books/cpu/kserve/kserve-integration.ipynb
diff --git a/tests/notebooks/kserve/requirements.in → tests/notebooks/cpu/kserve/requirements.in b/tests/notebooks/kserve/requirements.in → tests/notebooks/cpu/kserve/requirements.in
diff --git a/tests/notebooks/kserve/requirements.txt → tests/notebooks/cpu/kserve/requirements.txt b/tests/notebooks/kserve/requirements.txt → tests/notebooks/cpu/kserve/requirements.txt
diff --git a/tests/notebooks/mlflow-kserve/.gitignore → tests/notebooks/cpu/mlflow-kserve/.gitignore b/tests/notebooks/mlflow-kserve/.gitignore → tests/notebooks/cpu/mlflow-kserve/.gitignore
diff --git a/...tebooks/mlflow-kserve/mlflow-kserve.ipynb → ...oks/cpu/mlflow-kserve/mlflow-kserve.ipynb b/...tebooks/mlflow-kserve/mlflow-kserve.ipynb → ...oks/cpu/mlflow-kserve/mlflow-kserve.ipynb
diff --git a/...s/notebooks/mlflow-kserve/requirements.in → ...tebooks/cpu/mlflow-kserve/requirements.in b/...s/notebooks/mlflow-kserve/requirements.in → ...tebooks/cpu/mlflow-kserve/requirements.in
diff --git a/.../notebooks/mlflow-kserve/requirements.txt → ...ebooks/cpu/mlflow-kserve/requirements.txt b/.../notebooks/mlflow-kserve/requirements.txt → ...ebooks/cpu/mlflow-kserve/requirements.txt
diff --git a/...flow-minio/mlflow-minio-integration.ipynb → ...flow-minio/mlflow-minio-integration.ipynb b/...flow-minio/mlflow-minio-integration.ipynb → ...flow-minio/mlflow-minio-integration.ipynb
diff --git a/tests/notebooks/mlflow-minio/requirements.in → ...otebooks/cpu/mlflow-minio/requirements.in b/tests/notebooks/mlflow-minio/requirements.in → ...otebooks/cpu/mlflow-minio/requirements.in
diff --git a/...s/notebooks/mlflow-minio/requirements.txt → ...tebooks/cpu/mlflow-minio/requirements.txt b/...s/notebooks/mlflow-minio/requirements.txt → ...tebooks/cpu/mlflow-minio/requirements.txt
diff --git a/tests/notebooks/mlflow-minio/sample.csv → tests/notebooks/cpu/mlflow-minio/sample.csv b/tests/notebooks/mlflow-minio/sample.csv → tests/notebooks/cpu/mlflow-minio/sample.csv
diff --git a/tests/notebooks/mlflow-minio/sample.txt → tests/notebooks/cpu/mlflow-minio/sample.txt b/tests/notebooks/mlflow-minio/sample.txt → tests/notebooks/cpu/mlflow-minio/sample.txt
diff --git a/...notebooks/mlflow/mlflow-integration.ipynb → ...books/cpu/mlflow/mlflow-integration.ipynb b/...notebooks/mlflow/mlflow-integration.ipynb → ...books/cpu/mlflow/mlflow-integration.ipynb
diff --git a/tests/notebooks/mlflow/requirements.in → tests/notebooks/cpu/mlflow/requirements.in b/tests/notebooks/mlflow/requirements.in → tests/notebooks/cpu/mlflow/requirements.in
diff --git a/tests/notebooks/mlflow/requirements.txt → tests/notebooks/cpu/mlflow/requirements.txt b/tests/notebooks/mlflow/requirements.txt → tests/notebooks/cpu/mlflow/requirements.txt
diff --git a/tests/notebooks/training/requirements.in → tests/notebooks/cpu/training/requirements.in b/tests/notebooks/training/requirements.in → tests/notebooks/cpu/training/requirements.in
diff --git a/tests/notebooks/training/requirements.txt → ...s/notebooks/cpu/training/requirements.txt b/tests/notebooks/training/requirements.txt → ...s/notebooks/cpu/training/requirements.txt
diff --git a/...books/training/training-integration.ipynb → ...s/cpu/training/training-integration.ipynb b/...books/training/training-integration.ipynb → ...s/cpu/training/training-integration.ipynb
diff --git a/tests/notebooks/gpu/placeholder/placeholder.ipynb b/tests/notebooks/gpu/placeholder/placeholder.ipynb
diff --git a/tests/test_notebooks.py b/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__)