Modify Tests

Author: AryanAhadinia

.github/workflows/ci.yml

@@ -58,4 +58,4 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - run: pip install -e ".[dev]"
-      - run: pytest --cov --cov-report=term-missing
+      - run: pytest --cov --cov-report=term-missing -n auto

Makefile

@@ -5,7 +5,7 @@ format:
 	ruff check --fix .
 
 test:
-	. ./.test.sh && pytest --cov --cov-report=term-missing
+	. ./.test.sh && pytest --cov --cov-report=term-missing -n auto
 
 clean-cache:
 	find . -type d -name "__pycache__" -exec rm -rf {} +

pyproject.toml

@@ -26,6 +26,7 @@ dev = [
     "pytest>=9.0",
     "pytest-asyncio>=1.0",
     "pytest-cov>=7.0",
+    "pytest-xdist>=3.0",
     "pre-commit>=4.0",
     "types-PyYAML>=6.0",
 ]
@@ -79,4 +80,7 @@ show_missing = true
 fail_under = 0  # TODO
 
 [dependency-groups]
-dev = ["pytest-asyncio>=1.3.0"]
+dev = [
+    "pytest-asyncio>=1.3.0",
+    "pytest-xdist>=3.8.0",
+]

src/swiss_ai_model_launch/launchers/firecrest_launcher.py

@@ -186,3 +186,9 @@ async def get_job_logs(self, job_id: int) -> tuple[str, str]:
                 err_log = ""
 
             return out_log, err_log
+
+    async def cancel_job(self, job_id: int) -> None:
+        await self.client.cancel_job(
+            system_name=self.system_name,
+            jobid=str(job_id),
+        )

src/swiss_ai_model_launch/launchers/launcher.py

@@ -23,3 +23,6 @@ async def get_job_status(self, job_id: int) -> JobStatus: ...
 
     @abstractmethod
     async def get_job_logs(self, job_id: int) -> tuple[str, str]: ...
+
+    @abstractmethod
+    async def cancel_job(self, job_id: int) -> None: ...

tests/integration/test_launch_apertus.py

@@ -139,5 +139,8 @@ async def test_launch_apertus_and_health(
     assert isinstance(job_id, int)
     assert served_model_name
 
-    await _wait_for_job_running(launcher, job_id, launch_timeout)
-    await _wait_for_model_healthy(served_model_name, cscs_api_key, health_timeout)
+    try:
+        await _wait_for_job_running(launcher, job_id, launch_timeout)
+        await _wait_for_model_healthy(served_model_name, cscs_api_key, health_timeout)
+    finally:
+        await launcher.cancel_job(job_id)