Merge pull request #44 from chyumin/ci-test-gpu

Ci test gpu

Author: alejoe91

.github/workflows/test_containers_gpu.yml

@@ -0,0 +1,48 @@
+name: CI-Test in AWS for GPU
+
+on:
+  workflow_dispatch:
+
+jobs:
+  deploy-runner:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: iterative/setup-cml@v1
+      - uses: actions/checkout@v3
+      - name: Deploy runner on EC2
+        env:
+          REPO_TOKEN: ${{ secrets.PERSONAL_ACCESS_TOKEN }}
+          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        run: |
+          cml runner \
+            --cloud=aws \
+            --cloud-region=us-east-2 \
+            --cloud-type=g4dn.xlarge \
+            --labels=cml-gpu
+  test-gpu-sorters:
+    needs: deploy-runner
+    runs-on: [self-hosted, cml-gpu]
+    timeout-minutes: 360 # 6h
+    container:
+      image: docker://iterativeai/cml:0-dvc2-base1-gpu
+      options: --privileged --gpus all
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Install dependencies
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y wget squashfs-tools
+          pip install -U pip
+          pip install -r requirements_test.txt
+
+      - uses: eWaterCycle/setup-singularity@v7
+        with:
+          singularity-version: 3.8.3
+
+      - name: Run test singularity containers with GPU
+        env:
+          REPO_TOKEN: ${{ secrets.PERSONAL_ACCESS_TOKEN }}
+        run: |
+          pytest -sv tests/test_containers_gpu.py

requirements_test.txt

@@ -1,4 +1,3 @@
 git+https://github.com/SpikeInterface/spikeinterface.git#egg=spikeinterface[full]
-docker==5.0.3
 spython==0.2.1
 pytest==7.1.2

tests/test_containers_gpu.py

@@ -0,0 +1,58 @@
+import os
+import shutil
+
+import pytest
+
+import spikeinterface.extractors as se
+import spikeinterface.sorters as ss
+
+os.environ['SINGULARITY_DISABLE_CACHE'] = 'true'
+
+
+@pytest.fixture(autouse=True)
+def work_dir(request, tmp_path):
+    """
+    This fixture, along with "run_kwargs" creates one folder per
+    test function using built-in tmp_path pytest fixture
+
+    The tmp_path will be the working directory for the test function
+
+    At the end of the each test function, a clean up will be done
+    """
+    os.chdir(tmp_path)
+    yield
+    os.chdir(request.config.invocation_dir)
+    shutil.rmtree(str(tmp_path))
+
+
+@pytest.fixture
+def run_kwargs(work_dir):
+    test_recording, _ = se.toy_example(
+        duration=30,
+        seed=0,
+        num_channels=64,
+        num_segments=1
+    )
+    test_recording = test_recording.save(name='toy')
+    return dict(recording=test_recording, verbose=True, singularity_image=True)
+
+
+def test_kilosort2(run_kwargs):
+    sorting = ss.run_kilosort2(output_folder="kilosort2", **run_kwargs)
+    print(sorting)
+
+
+def test_kilosort2_5(run_kwargs):
+    sorting = ss.run_kilosort2_5(output_folder="kilosort2_5", **run_kwargs)
+    print(sorting)
+
+
+def test_kilosort3(run_kwargs):
+    sorting = ss.run_kilosort3(output_folder="kilosort3", **run_kwargs)
+    print(sorting)
+
+
+@pytest.mark.skip(reason="Need to check errors")
+def test_pykilosort(run_kwargs):
+    sorting = ss.run_pykilosort(output_folder="pykilosort", **run_kwargs)
+    print(sorting)

tests/test_singularity_containers.py

@@ -76,27 +76,3 @@ def test_hdsort(run_kwargs):
 def test_kilosort1(run_kwargs):
     sorting = ss.run_kilosort(output_folder="kilosort", useGPU=False, **run_kwargs)
     print(sorting)
-
-
-@pytest.mark.skip(reason="GPU required")
-def test_kilosort2(run_kwargs):
-    sorting = ss.run_kilosort2(output_folder="kilosort2", **run_kwargs)
-    print(sorting)
-
-
-@pytest.mark.skip(reason="GPU required")
-def test_kilosort2_5(run_kwargs):
-    sorting = ss.run_kilosort2_5(output_folder="kilosort2_5", **run_kwargs)
-    print(sorting)
-
-
-@pytest.mark.skip(reason="GPU required")
-def test_kilosort3(run_kwargs):
-    sorting = ss.run_kilosort3(output_folder="kilosort3", **run_kwargs)
-    print(sorting)
-
-
-@pytest.mark.skip(reason="GPU required")
-def test_pykilosort(run_kwargs):
-    sorting = ss.run_pykilosort(output_folder="pykilosort", **run_kwargs)
-    print(sorting)