Speed up tests

.github/workflows/pytest.yml

@@ -1,39 +1,30 @@
----
-name: Unit and integration tests
+name: Testing with pytest
 
 on:
   pull_request:
     branches: [main, dev]
 
 jobs:
-  test-itwinai:
-    name: Test itwinai with pytest
+  test-torch:
+    name: Testing with pytest
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
 
-      - name: Install compilers for Horovod
-        run: |
-          sudo apt-get update &&
-          sudo apt-get install -y gcc-11 g++-11 &&
-          sudo apt-get install -y cmake &&
-          sudo apt-get install openmpi-bin  openmpi-common openmpi-doc libopenmpi-dev &&
-          gcc --version &&
-          cmake --version &&
-          mpirun --version
-
+      - name: Set up Python 3.11.9
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.11.9
+
       - name: Make PyTorch virtualenv
         shell: bash -l {0}
         run: make torch-env-cpu
-
-      - name: Make Tensorflow virtualenv
-        shell: bash -l {0}
-        run: make tensorflow-env-cpu
 
-      # NOTE, to change the name in which tests are run, set custom TORCH_ENV and TF_ENV env variables.
-      # Default environment names are ".venv-pytorch" and ".venv-tf"
+      # Comment this back in to also build tensorflow env
+      # - name: Make Tensorflow virtualenv
+      #   shell: bash -l {0}
+      #   run: make tensorflow-env-cpu
+
       - name: Run pytest for workflows
         shell: bash -l {0}
-        run: .venv-pytorch/bin/pytest -v ./tests/ -m "not hpc"
-
-
+        run: .venv-pytorch/bin/pytest -v -n logical ./tests/ -m "not hpc and not tensorflow"

env-files/torch/generic_torch.sh

@@ -26,14 +26,23 @@ fi
 
 # Activate the venv and then install itwinai as editable
 source $ENV_NAME/bin/activate
-pip install -e ".[torch,tf,dev]" \
+pip install uv
+
+if [ -z "$NO_CUDA" ]; then
+  # Install with CUDA support
+  uv pip install -e ".[torch,dev]" \
     --no-cache-dir \
     --extra-index-url https://download.pytorch.org/whl/cu121
+else
+  # Install without CUDA support
+  uv pip install -e ".[torch,dev]" 
+fi
+
 
 # Install Prov4ML
 if [[ "$(uname)" == "Darwin" ]]; then
-  pip install --no-cache-dir  "prov4ml[apple]@git+https://github.com/matbun/ProvML@new-main"
+  uv pip install --no-cache-dir  "prov4ml[apple]@git+https://github.com/matbun/ProvML@new-main"
 else
   # Assuming Nvidia GPUs are available
-  pip install --no-cache-dir  "prov4ml[nvidia]@git+https://github.com/matbun/ProvML@new-main"
+  uv pip install --no-cache-dir  "prov4ml[nvidia]@git+https://github.com/matbun/ProvML@new-main"
 fi

pyproject.toml

@@ -61,6 +61,8 @@ dev = [
   "ipykernel>=6.29.5",
   "ipython>=8.30.0",
   "ruff>=0.8.3",
+  "psutil>=6.1.0",
+  "pytest-xdist>=3.6.1",
 ]
 docs = [
   "sphinx-rtd-theme>=2.0.0",

tests/loggers/conftest.py

@@ -8,7 +8,8 @@
 # - Anna Lappe <[email protected]> - CERN
 # -------------------------------------------------------------------------------------
 
-import shutil
+import tempfile
+from pathlib import Path
 from unittest.mock import MagicMock
 
 import pytest
@@ -27,37 +28,40 @@
 
 @pytest.fixture(scope="module")
 def console_logger():
-    yield ConsoleLogger(savedir="/tmp/console/test_mllogs", log_freq=1)
-    shutil.rmtree("/tmp/console/test_mllogs", ignore_errors=True)
+    with tempfile.TemporaryDirectory() as temp_dir:
+        save_dir = Path(temp_dir) / "console/test_mllogs"
+        yield ConsoleLogger(savedir=save_dir, log_freq=1)
 
 
 @pytest.fixture(scope="module")
 def mlflow_logger():
-    yield MLFlowLogger(
-        savedir="/tmp/mlflow/test_mllogs",
-        experiment_name="test_experiment",
-        tracking_uri="file:///tmp/mlruns",
-    )
-    shutil.rmtree("/tmp/mlflow/test_mllogs", ignore_errors=True)
-    shutil.rmtree("/tmp/mlruns", ignore_errors=True)
+    with tempfile.TemporaryDirectory() as temp_dir:
+        save_dir = Path(temp_dir) / "mlflow/test_mllogs"
+        yield MLFlowLogger(
+            savedir=save_dir,
+            experiment_name="test_experiment",
+        )
 
 
 @pytest.fixture(scope="module")
 def wandb_logger():
-    yield WandBLogger(savedir="/tmp/wandb/test_mllogs", project_name="test_project")
-    shutil.rmtree("/tmp/wandb/test_mllogs", ignore_errors=True)
+    with tempfile.TemporaryDirectory() as temp_dir:
+        save_dir = Path(temp_dir) / "wandb/test_mllogs"
+        yield WandBLogger(savedir=save_dir, project_name="test_project")
 
 
 @pytest.fixture(scope="module")
 def tensorboard_logger_tf():
-    yield TensorBoardLogger(savedir="/tmp/tf_tb/test_mllogs", framework="tensorflow")
-    shutil.rmtree("/tmp/tf_tb/test_mllogs", ignore_errors=True)
+    with tempfile.TemporaryDirectory() as temp_dir:
+        save_dir = Path(temp_dir) / "tf_tb/test_mllogs"
+        yield TensorBoardLogger(savedir=save_dir, framework="tensorflow")
 
 
 @pytest.fixture(scope="module")
 def tensorboard_logger_torch():
-    yield TensorBoardLogger(savedir="/tmp/torch_tb/test_mllogs", framework="pytorch")
-    shutil.rmtree("/tmp/torch_tb/test_mllogs", ignore_errors=True)
+    with tempfile.TemporaryDirectory() as temp_dir:
+        save_dir = Path(temp_dir) / "torch_tb/test_mllogs"
+        yield TensorBoardLogger(savedir=save_dir, framework="pytorch")
 
 
 @pytest.fixture(scope="module")

tests/use-cases/test_mnist.py

@@ -117,6 +117,7 @@ def test_mnist_train_torch_lightning(torch_env, install_requirements):
         subprocess.run(cmd.split(), check=True, cwd=temp_dir)
 
 
+@pytest.mark.tensorflow
 @pytest.mark.functional
 def test_mnist_train_tf(tf_env, install_requirements):
     """

use-cases/mnist/torch/dataloader.py

@@ -33,7 +33,7 @@ def __init__(
         self.save_path = save_path
 
     @monitor_exec
-    def execute(self) -> Tuple[Dataset, Dataset]:
+    def execute(self) -> Tuple[Dataset, Dataset, None]:
         train_dataset = datasets.MNIST(
             self.save_path,
             train=True,
@@ -95,8 +95,7 @@ def __getitem__(self, index: int) -> Tuple[Any, Any]:
 
     @staticmethod
     def generate_jpg_sample(root: str, max_items: int = 100):
-        """Generate a sample dataset of JPG images starting from
-            LeCun's test dataset.
+        """Generate a sample dataset of JPG images starting from LeCun's test dataset.
 
         Args:
             root (str): sample path on disk