build: extras (cpu, gpu) and github action workflow with gpu runners

Author: radurogojanumai

.github/workflows/run-tests-cpu.yml

@@ -0,0 +1,38 @@
+name: "[CPU] mostlyai-qa Tests"
+
+on:
+  workflow_call:
+
+
+env:
+    PYTHON_KEYRING_BACKEND: keyring.backends.null.Keyring
+    FORCE_COLOR: "1"
+
+jobs:
+  run-tests-cpu:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+    steps:
+      - name: Setup | Checkout
+        uses: actions/checkout@v4
+
+      - name: Setup | Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
+
+      - name: Setup | uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: false
+          python-version: '3.10'
+
+      - name: Setup | dependencies
+        run: uv sync --frozen --extra cpu
+      - name: Run tests
+        run: uv run pytest tests/
+
+      - name: Build mkdocs
+        run: uv run mkdocs build --strict

.github/workflows/run-tests-gpu.yml

@@ -0,0 +1,39 @@
+name: "[GPU] mostlyai-qa Tests"
+
+on:
+  workflow_call:
+
+
+env:
+    PYTHON_KEYRING_BACKEND: keyring.backends.null.Keyring
+    FORCE_COLOR: "1"
+
+jobs:
+  run-tests-gpu:
+    runs-on: gha-gpu-public
+    permissions:
+      contents: read
+      packages: write
+    steps:
+      - name: Setup | Checkout
+        uses: actions/checkout@v4
+
+      - name: Setup | Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
+
+      - name: Setup | uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: false
+          python-version: '3.10'
+
+      - name: Setup | dependencies
+        run: uv sync --frozen --extra gpu
+
+      - name: Run tests
+        run: uv run pytest tests/
+
+      - name: Build mkdocs
+        run: uv run mkdocs build --strict

.github/workflows/run-tests.yml

@@ -13,9 +13,16 @@ jobs:
       (github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name != github.repository)
     uses: ./.github/workflows/pre-commit-check.yml
     secrets: inherit
-  run-tests:
+  run-tests-cpu:
     if: |
       github.event_name == 'push' ||
       (github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name != github.repository)
-    uses: ./.github/workflows/run-tests.yml
+    uses: ./.github/workflows/run-tests-cpu.yml
     secrets: inherit
+  run-tests-gpu:
+      if: |
+        github.ref == 'refs/heads/main' ||
+        startsWith(github.ref, 'refs/tags/') ||
+        contains(github.event.head_commit.message, '[gpu]')
+      uses: ./.github/workflows/run-tests-gpu.yml
+      secrets: inherit

.github/workflows/workflow.yml

@@ -20,6 +20,8 @@ The latest release of `mostlyai-qa` can be installed via pip:
 pip install -U mostlyai-qa
 ```
 
+On Linux, one can explicitly install `mostlyai-qa[cpu]` or `mostlyai-qa[gpu]`, for CPU-only or CUDA support respectively
+
 ## Quick Start
 
 ```python

README.md

@@ -41,6 +41,16 @@ dependencies = [
     "sentence-transformers>=3.1.0",
     "rich>=13.9.4,<14",
     "skops>=0.11.0",
+    "torch>=2.5.1",
+]
+
+[project.optional-dependencies]
+gpu = [
+    "torch>=2.5.1",
+]
+cpu = [
+    "torch==2.5.1+cpu; sys_platform == 'linux'",
+    "torch>=2.5.1; sys_platform != 'linux'",
 ]
 
 [project.urls]
@@ -68,6 +78,22 @@ docs = [
 
 [tool.uv]
 default-groups = ["dev", "docs"]
+conflicts = [
+  [
+    { extra = "cpu" },
+    { extra = "gpu" },
+  ],
+]
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cpu", extra = "cpu", marker = "sys_platform == 'linux'"},
+]
 
 [tool.hatch.build.targets.sdist]
 include = ["mostlyai/qa"]