From 81d3170073d5d34a5b876fbe92aeac9c7eb928cd Mon Sep 17 00:00:00 2001
From: Ori Kronfeld <ori.kronfeld@weizmann.ac.il>
Date: Mon, 9 Dec 2024 15:19:48 +0200
Subject: [PATCH] adding cuda test

---
 .github/workflows/build.yml           |  8 ++-
 .github/workflows/test_linux_cuda.yml | 72 +++++++++++++++++++++++++++
 .readthedocs.yaml                     |  6 ++-
 tests/conftest.py                     | 28 +++++++++++
 4 files changed, 111 insertions(+), 3 deletions(-)
 create mode 100644 .github/workflows/test_linux_cuda.yml
 create mode 100644 tests/conftest.py

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 28e97ee..5267274 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -2,9 +2,13 @@ name: Check Build
 
 on:
   push:
-    branches: [main]
+    branches: [main, "[0-9]+.[0-9]+.x"]
   pull_request:
-    branches: [main]
+    branches: [main, "[0-9]+.[0-9]+.x"]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
 
 jobs:
   package:
diff --git a/.github/workflows/test_linux_cuda.yml b/.github/workflows/test_linux_cuda.yml
new file mode 100644
index 0000000..50c2c56
--- /dev/null
+++ b/.github/workflows/test_linux_cuda.yml
@@ -0,0 +1,72 @@
+name: PopV (cuda)
+
+on:
+  push:
+    branches: [main, "[0-9]+.[0-9]+.x"] #this is new
+  pull_request:
+    branches: [main, "[0-9]+.[0-9]+.x"]
+    types: [labeled, synchronize, opened]
+  schedule:
+    - cron: "0 10 * * *" # runs at 10:00 UTC (03:00 PST) every day
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test:
+    # if PR has label "cuda tests" or "all tests" or if scheduled or manually triggered
+    if: >-
+      (
+        contains(github.event.pull_request.labels.*.name, 'cuda tests') ||
+        contains(github.event.pull_request.labels.*.name, 'all tests') ||
+        contains(github.event_name, 'schedule') ||
+        contains(github.event_name, 'workflow_dispatch')
+      )
+
+    runs-on: [self-hosted, Linux, X64, CUDA]
+
+    defaults:
+      run:
+        shell: bash -e {0} # -e to fail on error
+
+    container:
+      image: ghcr.io/yoseflab/popv:py3.10-cu12-base
+      options: --user root --gpus all --pull always
+
+    name: integration
+
+    env:
+      OS: ${{ matrix.os }}
+      PYTHON: ${{ matrix.python }}
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+          cache: "pip"
+          cache-dependency-path: "**/pyproject.toml"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip wheel uv
+          python -m uv pip install --system "PopV[tests] @ ."
+          python -m pip install jax[cuda]
+          python -m pip install nvidia-nccl-cu12
+
+      - name: Run pytest
+        env:
+          MPLBACKEND: agg
+          PLATFORM: ${{ matrix.os }}
+          DISPLAY: :42
+          COLUMNS: 120
+        run: |
+          coverage run -m pytest -v --color=yes --accelerator cuda --devices auto
+          coverage report
+
+      - uses: codecov/codecov-action@v4
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
index 69897c3..58f1018 100644
--- a/.readthedocs.yaml
+++ b/.readthedocs.yaml
@@ -13,4 +13,8 @@ python:
     - method: pip
       path: .
       extra_requirements:
-        - doc
+        - docsbuild
+submodules:
+  include:
+    - "docs/notebooks"
+  recursive: true
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..b197c34
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,28 @@
+import shutil
+import pytest
+
+def pytest_addoption(parser):
+    """Docstring for pytest_addoption."""
+    parser.addoption(
+        "--accelerator",
+        action="store",
+        default="cpu",
+        help="Option to specify which accelerator to use for tests.",
+    )
+    parser.addoption(
+        "--devices",
+        action="store",
+        default="auto",
+        help="Option to specify which devices to use for tests.",
+    )
+
+@pytest.fixture(scope="session")
+def accelerator(request):
+    """Docstring for accelerator."""
+    return request.config.getoption("--accelerator")
+
+
+@pytest.fixture(scope="session")
+def devices(request):
+    """Docstring for devices."""
+    return request.config.getoption("--devices")