GPU implementation of hamming distance

.cirun.yml

@@ -0,0 +1,11 @@
+runners:
+  - name: aws-gpu-runner
+    cloud: aws
+    instance_type: g4dn.xlarge
+    machine_image: ami-067a4ba2816407ee9
+    region: eu-north-1
+    preemptible:
+      - true
+      - false
+    labels:
+      - cirun-aws-gpu

.github/workflows/test-gpu.yaml

@@ -0,0 +1,65 @@
+name: GPU-CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    types:
+      - labeled
+      - opened
+      - synchronize
+
+# Cancel the job if new commits are pushed
+# https://stackoverflow.com/questions/66335225/how-to-cancel-previous-runs-in-the-pr-when-you-push-new-commitsupdate-the-curre
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  check:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: flying-sheep/check@v1
+        with:
+          success: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'run-gpu-ci') }}
+  test:
+    name: GPU Tests
+    needs: check
+    runs-on: "cirun-aws-gpu--${{ github.run_id }}"
+    timeout-minutes: 30
+
+    defaults:
+      run:
+        shell: bash -el {0}
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Nvidia SMI sanity check
+        run: nvidia-smi
+
+      - name: Install Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install UV
+        uses: hynek/setup-cached-uv@v2
+        with:
+          cache-dependency-path: pyproject.toml
+
+      - name: Install scirpy
+        run: uv pip install --system -e ".[dev,test,rpack,dandelion,diversity,parasail,cupy]"
+      - name: Pip list
+        run: pip list
+
+      - name: Run test
+        run: pytest -m gpu
+
+      - name: Remove 'run-gpu-ci' Label
+        if: always()
+        uses: actions-ecosystem/action-remove-labels@v1
+        with:
+          labels: "run-gpu-ci"
+          github_token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/test.yaml

@@ -59,7 +59,7 @@ jobs:
           PLATFORM: ${{ matrix.os }}
           DISPLAY: :42
         run: |
-          coverage run -m pytest -v --color=yes
+          coverage run -m pytest -v --color=yes -m "not gpu"
       - name: Report coverage
         run: |
           coverage report

.pre-commit-config.yaml

@@ -2,8 +2,8 @@ fail_fast: false
 default_language_version:
   python: python3
 default_stages:
-  - commit
-  - push
+  - pre-commit
+  - pre-push
 minimum_pre_commit_version: 2.16.0
 repos:
   - repo: https://github.com/pre-commit/mirrors-prettier

pyproject.toml

@@ -76,6 +76,9 @@ test = [
     'coverage',
     'black',
 ]
+cupy = [
+    'cupy-cuda12x',
+]
 dandelion = [
     'sc-dandelion>=0.3.5',
 ]
@@ -107,7 +110,8 @@ xfail_strict = true
 # ]
 markers = [
     "conda: marks a subset of tests to be ran on the Bioconda CI.",
-    "extra: marks tests that require extra dependencies."
+    "extra: marks tests that require extra dependencies.",
+    "gpu: mark test to run on GPU",
 ]
 minversion = 6.0
 norecursedirs = [ '.*', 'build', 'dist', '*.egg', 'data', '__pycache__']

src/scirpy/ir_dist/__init__.py

@@ -60,6 +60,10 @@ def IrNeighbors(*args, **kwargs):
         BLOSUM62 matrix. Faster implementation of `alignment` with some loss.
         This option is incompatible with nucleotide sequences.
         See :class:`~scirpy.ir_dist.metrics.FastAlignmentDistanceCalculator`.
+      * `normalized_hamming` -- Normalized Hamming distance (in percent) for CDR3 sequences of equal length.
+        See :class:`~scirpy.ir_dist.metrics.HammingDistanceCalculator`.
+      * `tcrdist` -- Distance based on pairwise sequence alignments between TCR CDR3 sequences based on the tcrdist metric.
+        See :class:`~scirpy.ir_dist.metrics.TCRdistDistanceCalculator`.
       * any instance of :class:`~scirpy.ir_dist.metrics.DistanceCalculator`.
 """
 
@@ -105,6 +109,8 @@ def _get_distance_calculator(metric: MetricType, cutoff: int | None, *, n_jobs=-
         dist_calc = metrics.HammingDistanceCalculator(n_jobs=n_jobs, **kwargs)
     elif metric == "normalized_hamming":
         dist_calc = metrics.HammingDistanceCalculator(n_jobs=n_jobs, normalize=True, **kwargs)
+    elif metric == "gpu_hamming":
+        dist_calc = metrics.GPUHammingDistanceCalculator(**kwargs)
     elif metric == "tcrdist":
         dist_calc = metrics.TCRdistDistanceCalculator(n_jobs=n_jobs, **kwargs)
     else: