Support for pytest-xdist

.github/workflows/end2end_tests.yaml

@@ -48,7 +48,7 @@ jobs:
         pip install -e .[dev]
 
     - name: Run tests
-      run: pytest tests/test_end2end.py --download-weights --delete-weights-now -s -v
+      run: pytest -n 2 tests/test_end2end.py --download-weights -s -v
 
     - name: Authenticate to Google Cloud
       uses: google-github-actions/auth@v2
@@ -59,4 +59,4 @@ jobs:
 
     - name: Run private model tests
       if: github.event.pull_request.head.repo.full_name == github.repository
-      run: pytest tests/test_end2end.py::test_private_model_conversion --test-private --delete-weights-now -s -v
+      run: pytest -n 2 tests/test_end2end.py::test_private_model_conversion --test-private -s -v

.github/workflows/unittests.yaml

@@ -80,19 +80,14 @@ jobs:
     - name: Run tests with coverage [Ubuntu]
       if: matrix.os == 'ubuntu-latest' && matrix.version == '3.10'
       working-directory: ${{ inputs.tools_ref != '' && 'tools' || '' }}
-      env:
-        COVERAGE_PROCESS_START: .coveragerc
       run: |
-        coverage erase
-        coverage run -m pytest tests/test_unittests.py --download-weights --delete-weights-now -s -v --junit-xml pytest.xml
-        coverage combine
-        coverage xml -o coverage.xml -i # Recent changes to coverage 7.13 make it stricter where it was previously passing https://coverage.readthedocs.io/en/7.13.0/messages.html#error-no-source
+        pytest -n 2 tests/test_unittests.py --download-weights -s -v --junit-xml pytest.xml --cov=tools --cov-report=xml:coverage.xml
 
 
     - name: Run tests [Windows, macOS]
       if: matrix.os != 'ubuntu-latest' || matrix.version != '3.10'
       working-directory: ${{ inputs.tools_ref != '' && 'tools' || '' }}
-      run: pytest tests/test_unittests.py --download-weights --delete-weights-now -s -v --junit-xml pytest.xml
+      run: pytest -n 2 tests/test_unittests.py --download-weights -s -v --junit-xml pytest.xml
 
     - name: Generate coverage badge [Ubuntu]
       if: matrix.os == 'ubuntu-latest' && matrix.version == '3.10' && inputs.tools_ref == '' && inputs.ml_ref == ''

requirements-dev.txt

@@ -1,4 +1,5 @@
 pytest
+pytest-xdist>=3.6.1
 pytest-subtests
 pre-commit>=4.1.0
 pytest-cov>=4.1.0

tests/README.md

@@ -23,3 +23,11 @@ pytest --download-weights --log-cli-level=INFO --log-file=out.log --log-file-lev
 ```
 
 This will run the full test suite on all the supported models and store the DEBUG logs into the out.log file.
+
+To run the suite in parallel with `pytest-xdist`, use for example:
+
+```
+pytest -n auto --download-weights --log-cli-level=INFO --log-file=out.log --log-file-level=DEBUG .
+```
+
+The tests isolate generated `shared_with_container` artifacts per test process while reusing a shared `weights/` cache. Avoid combining `-n ...` with `--delete-weights-now`, since that option defeats the shared cache and is skipped under xdist to avoid cross-worker races.

tests/conftest.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import shutil
+from pathlib import Path
 
 import pytest
 
@@ -32,6 +33,7 @@ def pytest_addoption(parser):
             "v11",
             "v12",
             "v26",
+            "v26_nms",
         ],
         default=None,
         help="If set then test only that specific yolo version",
@@ -63,24 +65,36 @@ def test_config(pytestconfig):
         "test_case": pytestconfig.getoption("test_case"),
         "delete_weights_now": pytestconfig.getoption("delete_weights_now"),
         "test_private": pytestconfig.getoption("test_private"),
+        "weights_dir": Path(__file__).resolve().parents[1] / "weights",
     }
 
 
 @pytest.fixture(scope="function", autouse=True)
-def cleanup_output_after_tests(test_config):
+def isolate_test_workspace(monkeypatch, tmp_path):
+    shared_dir = tmp_path / "shared_with_container"
+    monkeypatch.setenv("TOOLS_SHARED_DIR", str(shared_dir))
+    return shared_dir
+
+
+@pytest.fixture(scope="function", autouse=True)
+def cleanup_output_after_tests(test_config, isolate_test_workspace):
     yield  # Tests run here
     if test_config["delete_output"]:
-        folder_to_delete = "shared_with_container"
-        if os.path.exists(folder_to_delete):
-            shutil.rmtree(folder_to_delete)
-            logger.info(f"Removed test artifacts from {folder_to_delete}")
+        if isolate_test_workspace.exists():
+            shutil.rmtree(isolate_test_workspace)
+            logger.info(f"Removed test artifacts from {isolate_test_workspace}")
 
 
 @pytest.fixture(scope="function", autouse=True)
 def cleanup_weights_after_tests(test_config):
     yield  # Tests run here
     if test_config["delete_weights_now"]:
-        folder_to_delete = "weights"
-        if os.path.exists(folder_to_delete):
+        if os.environ.get("PYTEST_XDIST_WORKER"):
+            logger.warning(
+                "Skipping `--delete-weights-now` cleanup under pytest-xdist to avoid cross-worker races."
+            )
+            return
+        folder_to_delete = test_config["weights_dir"]
+        if folder_to_delete.exists():
             shutil.rmtree(folder_to_delete)
             logger.info(f"Removed test artifacts from {folder_to_delete}")

tests/constants.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-SAVE_FOLDER = "./weights"
+from pathlib import Path
+
+SAVE_FOLDER = str(Path(__file__).resolve().parents[1] / "weights")
 
 TEST_MODELS = [
     {"name": "yolov6nr4", "version": "v6r4"},

tests/helper_functions.py

@@ -4,7 +4,9 @@
 import json
 import logging
 import os
+import sys
 import tarfile
+import tempfile
 from pathlib import Path
 
 import requests
@@ -14,18 +16,34 @@
 logger = logging.getLogger()
 
 
+def get_shared_dir() -> Path:
+    return Path(os.environ.get("TOOLS_SHARED_DIR", "shared_with_container"))
+
+
+def get_outputs_dir() -> Path:
+    return get_shared_dir() / "outputs"
+
+
+def get_tools_command(*args: str) -> list[str]:
+    return [sys.executable, "-m", "tools.main", *args]
+
+
 def download_model(model_name: str, folder: str):
     logger.info(f"Downloading model '{model_name}' into folder '{folder}'")
     url = MODEL_TYPE2URL.get(model_name)
     if not url:
         raise ValueError(f"No URL defined for model {model_name}")
     response = requests.get(url)
     os.makedirs(folder, exist_ok=True)
-    file_path = os.path.join(folder, f"{model_name}.pt")
-    with open(file_path, "wb") as f:
+    file_path = Path(folder) / f"{model_name}.pt"
+    with tempfile.NamedTemporaryFile(
+        dir=folder, prefix=f"{model_name}.", suffix=".tmp", delete=False
+    ) as f:
         f.write(response.content)
+        temp_path = Path(f.name)
+    os.replace(temp_path, file_path)
     logger.debug(f"Model downloaded and saved to {file_path}")
-    return file_path
+    return str(file_path)
 
 
 def download_private_model(model_name: str, filename: str, folder: str) -> str:
@@ -36,37 +54,43 @@ def download_private_model(model_name: str, filename: str, folder: str) -> str:
     dest_dir = Path(folder)
     dest_dir.mkdir(parents=True, exist_ok=True)
 
-    downloaded_path = LuxonisFileSystem.download(remote_path, dest=dest_dir)
+    temp_dir = dest_dir / f".tmp-{model_name}-{os.getpid()}"
+    temp_dir.mkdir(parents=True, exist_ok=True)
+    downloaded_path = LuxonisFileSystem.download(remote_path, dest=temp_dir)
     final_path = dest_dir / f"{model_name}.pt"
 
     # Cases where name != filename in the PRIVATE_TEST_MODELS entry
     if downloaded_path != final_path and downloaded_path.exists():
-        downloaded_path.rename(final_path)
+        os.replace(downloaded_path, final_path)
+
+    if temp_dir.exists():
+        temp_dir.rmdir()
 
     logger.debug(f"Private model downloaded and saved to {final_path}")
     return str(final_path)
 
 
 def nn_archive_checker(extra_keys_to_check: list = []):  # noqa: B006
     """Tests the content of the exported NNArchive."""
-    output_dir = "shared_with_container/outputs"
-    subdirs = [
-        d for d in os.listdir(output_dir) if os.path.isdir(os.path.join(output_dir, d))
-    ]
+    output_dir = get_outputs_dir()
+    assert output_dir.exists(), f"Output directory `{output_dir}` does not exist"
+    subdirs = [d for d in output_dir.iterdir() if d.is_dir()]
     assert subdirs, f"No folders found in `{output_dir}`"
 
     # Sort by modification time (most recent last)
-    subdirs.sort(key=lambda d: os.path.getmtime(os.path.join(output_dir, d)))
+    subdirs.sort(key=lambda d: d.stat().st_mtime)
     latest_subdir = subdirs[-1]
-    model_output_path = os.path.join(output_dir, latest_subdir)
+    model_output_path = latest_subdir
     logger.debug(f"Model output path: {model_output_path}")
 
     # Find .tar.xz archive
-    archive_files = [f for f in os.listdir(model_output_path) if f.endswith(".tar.xz")]
+    archive_files = [
+        f.name for f in model_output_path.iterdir() if f.name.endswith(".tar.xz")
+    ]
     assert len(archive_files) == 1, (
         f"Expected 1 .tar.xz file, found {len(archive_files)}: {archive_files}"
     )
-    archive_path = os.path.join(model_output_path, archive_files[0])
+    archive_path = model_output_path / archive_files[0]
 
     # Open and inspect the archive
     with tarfile.open(archive_path, "r:xz") as tar:
@@ -131,21 +155,22 @@ def nn_archive_checker(extra_keys_to_check: list = []):  # noqa: B006
 
 def load_latest_nn_archive_config() -> dict:
     """Load config.json from the most recently exported NNArchive."""
-    output_dir = "shared_with_container/outputs"
-    subdirs = [
-        d for d in os.listdir(output_dir) if os.path.isdir(os.path.join(output_dir, d))
-    ]
+    output_dir = get_outputs_dir()
+    assert output_dir.exists(), f"Output directory `{output_dir}` does not exist"
+    subdirs = [d for d in output_dir.iterdir() if d.is_dir()]
     assert subdirs, f"No folders found in `{output_dir}`"
 
-    subdirs.sort(key=lambda d: os.path.getmtime(os.path.join(output_dir, d)))
+    subdirs.sort(key=lambda d: d.stat().st_mtime)
     latest_subdir = subdirs[-1]
-    model_output_path = os.path.join(output_dir, latest_subdir)
+    model_output_path = latest_subdir
 
-    archive_files = [f for f in os.listdir(model_output_path) if f.endswith(".tar.xz")]
+    archive_files = [
+        f.name for f in model_output_path.iterdir() if f.name.endswith(".tar.xz")
+    ]
     assert len(archive_files) == 1, (
         f"Expected 1 .tar.xz file, found {len(archive_files)}: {archive_files}"
     )
-    archive_path = os.path.join(model_output_path, archive_files[0])
+    archive_path = model_output_path / archive_files[0]
 
     with tarfile.open(archive_path, "r:xz") as tar:
         file_names = [m.name for m in tar.getmembers() if m.isfile()]

tests/test_end2end.py

@@ -9,6 +9,7 @@
 from helper_functions import (
     download_model,
     download_private_model,
+    get_tools_command,
     load_latest_nn_archive_config,
     nn_archive_checker,
 )
@@ -60,7 +61,7 @@ def test_cli_conversion(model: dict, test_config: dict, subtests):
         else:
             pytest.skip("Weights not present and `download_weights` not set")
 
-    command = ["tools", model_path]
+    command = get_tools_command(model_path)
     if model.get("cli_version"):
         command += ["--version", model.get("cli_version")]
     if model.get("size"):  # edge case when stride=64 is needed
@@ -122,7 +123,7 @@ def test_n_variant_nnarchive_outputs(model_case: dict, test_config: dict):
         else:
             pytest.skip("Weights missing and `download_weights` not set")
 
-    command = ["tools", model_path]
+    command = get_tools_command(model_path)
     result = subprocess.run(
         command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
     )
@@ -172,7 +173,7 @@ def test_yolo26_semseg_nnarchive_head(test_config: dict):
         else:
             pytest.skip("Weights missing and `download_weights` not set")
 
-    command = ["tools", model_path]
+    command = get_tools_command(model_path)
     logger.debug(f"CLI command: {command}")
 
     result = subprocess.run(
@@ -233,7 +234,7 @@ def test_private_model_conversion(model: dict, test_config: dict, subtests):
             SAVE_FOLDER,
         )
 
-    command = ["tools", model_path]
+    command = get_tools_command(model_path)
     if model.get("size"):
         command += ["--imgsz", model.get("size")]
 

tests/test_unittests.py

@@ -6,7 +6,7 @@
 
 import pytest
 from constants import SAVE_FOLDER
-from helper_functions import download_model, nn_archive_checker
+from helper_functions import download_model, get_tools_command, nn_archive_checker
 
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
@@ -38,7 +38,7 @@ def test_explicit_version(model_info: tuple[str, str]):
     if not os.path.exists(model_path):
         download_model(model_name, SAVE_FOLDER)
 
-    command = ["tools", model_path, "--version", version]
+    command = get_tools_command(model_path, "--version", version)
     if version == "yolov5":  # edge case when stride=64 is needed
         command += ["--imgsz", "320"]
     logger.debug(f"CLI command: {command}")
@@ -64,7 +64,7 @@ def test_wrong_explicit_version(version: str, expected_exit_code: int):
     if not os.path.exists(model_path):
         download_model(model_name, SAVE_FOLDER)
 
-    command = ["tools", model_path, "--version", version]
+    command = get_tools_command(model_path, "--version", version)
     logger.debug(f"CLI command: {command}")
 
     result = subprocess.run(
@@ -83,7 +83,7 @@ def test_explicit_input_size(input_size: str):
     if not os.path.exists(model_path):
         download_model(model_name, SAVE_FOLDER)
 
-    command = ["tools", model_path, "--imgsz", input_size]
+    command = get_tools_command(model_path, "--imgsz", input_size)
     logger.debug(f"CLI command: {command}")
 
     result = subprocess.run(
@@ -112,7 +112,7 @@ def test_wrong_explicit_input_size(input_size: str):
     if not os.path.exists(model_path):
         download_model(model_name, SAVE_FOLDER)
 
-    command = ["tools", model_path, "--imgsz", input_size]
+    command = get_tools_command(model_path, "--imgsz", input_size)
     logger.debug(f"CLI command: {command}")
 
     result = subprocess.run(
@@ -131,7 +131,7 @@ def test_explicit_encoding(encoding: str):
     if not os.path.exists(model_path):
         download_model(model_name, SAVE_FOLDER)
 
-    command = ["tools", model_path, "--encoding", encoding]
+    command = get_tools_command(model_path, "--encoding", encoding)
     logger.debug(f"CLI command: {command}")
 
     result = subprocess.run(
@@ -157,7 +157,7 @@ def test_wrong_explicit_encoding(encoding: str):
     if not os.path.exists(model_path):
         download_model(model_name, SAVE_FOLDER)
 
-    command = ["tools", model_path, "--encoding", encoding]
+    command = get_tools_command(model_path, "--encoding", encoding)
     logger.debug(f"CLI command: {command}")
 
     result = subprocess.run(
@@ -177,7 +177,7 @@ def test_explicit_class_names():
 
     class_names = ["a"] * 80
     class_names_str = ", ".join(class_names)
-    command = ["tools", model_path, "--class-names", class_names_str]
+    command = get_tools_command(model_path, "--class-names", class_names_str)
     logger.debug(f"CLI command: {command}")
 
     result = subprocess.run(
@@ -204,7 +204,7 @@ def test_wrong_explicit_class_names():
         download_model(model_name, SAVE_FOLDER)
 
     class_names_str = "a"
-    command = ["tools", model_path, "--class-names", class_names_str]
+    command = get_tools_command(model_path, "--class-names", class_names_str)
     logger.debug(f"CLI command: {command}")
 
     result = subprocess.run(
@@ -222,7 +222,7 @@ def test_explicit_version_detection():
     if not os.path.exists(model_path):
         download_model(model_name, SAVE_FOLDER)
 
-    command = ["tools", model_path]
+    command = get_tools_command(model_path)
     logger.debug(f"CLI command: {command}")
 
     result = subprocess.run(