Prevent fork releases; add unit tests

.github/workflows/build.yml

@@ -4,9 +4,6 @@ on:
   push:
     branches:
       - main
-  pull_request:
-    branches:
-      - main
 jobs:
   build:
     runs-on: ubuntu-latest
@@ -29,10 +26,10 @@ jobs:
         run: pip wheel --no-deps -w dist .
   release:
     needs: build
+    if: github.repository == 'EleutherAI/bergson' && github.event_name == 'push' && github.ref == 'refs/heads/main' && !contains(github.event.head_commit.message, 'chore(release):')
     permissions:
       contents: write
       id-token: write
-    if: github.event_name == 'push' && github.ref == 'refs/heads/main' && !contains(github.event.head_commit.message, 'chore(release):')
     runs-on: ubuntu-latest
     concurrency: release
     steps:

.github/workflows/unit_tests.yml

@@ -0,0 +1,62 @@
+name: Unit Tests
+env:
+  UV_SYSTEM_PYTHON: 1
+on:
+  push:
+    branches:
+      - 'main'
+  pull_request:
+    branches:
+      - 'main'
+  workflow_dispatch:
+
+jobs:
+  linter:
+    name: Linters
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+
+    steps:
+      - uses: actions/checkout@v5
+      - name: "Set up Python"
+        uses: actions/setup-python@v6
+        with:
+          python-version: "3.10"
+      - name: Checkout
+        uses: astral-sh/setup-uv@v6
+        with:
+          enable-cache: true
+      - name: Install dependencies
+        run: uv pip install -e ".[dev]" --torch-backend=auto
+      - name: Pre-Commit
+        uses: pre-commit/action@v3.0.1
+      - name: Type Checking
+        uses: jakebailey/pyright-action@v2
+      - name: Cleanup
+        run: uv cache prune --ci
+
+  testcpu:
+    name: CPU Tests
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: true
+      matrix:
+        python-version: [ "3.10", "3.11", "3.12" ]
+
+    steps:
+      - uses: actions/checkout@v5
+
+      - name: Checkout
+        uses: astral-sh/setup-uv@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+          enable-cache: true
+
+      - name: Install dependencies
+        run: uv sync --extra dev
+
+      - name: Run tests
+        run: uv run pytest tests --showlocals -s -vv -n=auto
+
+      - name: Cleanup
+        run: uv cache prune --ci

.pre-commit-config.yaml

@@ -1,18 +1,16 @@
 # See https://pre-commit.com for more information
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
--   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v6.0.0
     hooks:
-    -   id: trailing-whitespace
-    -   id: end-of-file-fixer
-        exclude: \.txt$
--   repo: https://github.com/psf/black
-    rev: 25.1.0
+    - id: trailing-whitespace
+    - id: end-of-file-fixer
+      exclude: \.txt$
+    - id: no-commit-to-branch
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.14.5
     hooks:
-    -   id: black
--   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.11.9'
-    hooks:
-        - id: ruff
-          args: [--fix, --exit-non-zero-on-fix]
+    - id: ruff-check
+      args: [ --fix ]
+    - id: ruff-format

bergson/__main__.py

@@ -57,8 +57,7 @@ def execute(self):
 
         if self.index_cfg.projection_dim != 0:
             print(
-                "Warning: projection_dim is not 0. "
-                "Compressed gradients will be scored."
+                "Warning: projection_dim is not 0. Compressed gradients will be scored."
             )
 
         score_dataset(self.index_cfg, self.score_cfg)

bergson/gradients.py

@@ -155,9 +155,9 @@ def to_adafactor(self) -> AdafactorNormalizer:
         and the factored second moments.
         """
         # We assume avg_sq is a square matrix of shape [O, I]
-        assert (
-            self.avg_sq.ndim == 2
-        ), f"Expected 2D tensor for avg_sq, got {self.avg_sq.ndim}D"
+        assert self.avg_sq.ndim == 2, (
+            f"Expected 2D tensor for avg_sq, got {self.avg_sq.ndim}D"
+        )
 
         # Compute row and column means
         return AdafactorNormalizer(
@@ -563,7 +563,6 @@ def _process_grad(self, module: nn.Module, _, grad_out):
         # If we are using AdamNormalizer, or including bias gradients
         # we need to materialize the full gradient and then project
         if isinstance(norm, AdamNormalizer) or include_bias:
-
             P = G.mT @ I  # [N, O, S] @ [N, S, I] → [N, O, I]
             if include_bias:
                 # Append the bias gradient to the input

bergson/huggingface.py

@@ -243,9 +243,9 @@ def on_step_end(
 
         # Record training order if enabled
         if self.order is not None:
-            assert (
-                self.batch_indices is not None
-            ), "Batch indices are not available for training order tracking"
+            assert self.batch_indices is not None, (
+                "Batch indices are not available for training order tracking"
+            )
 
             epoch = int(state.epoch or 0)
             global_step = state.global_step

bergson/query/faiss_index.py

@@ -14,7 +14,7 @@
 from bergson.config import FaissConfig
 
 if TYPE_CHECKING:
     import faiss  # noqa: F401
 
 
 class Index(Protocol):
@@ -242,9 +242,9 @@
         shard_sizes[-1] += remainder
 
         # Verify all gradients will be consumed
-        assert (
-            sum(shard_sizes) == total_grads
-        ), f"Shard sizes {shard_sizes} don't sum to total_grads {total_grads}"
+        assert sum(shard_sizes) == total_grads, (
+            f"Shard sizes {shard_sizes} don't sum to total_grads {total_grads}"
+        )
 
         dl = gradients_loader(gradients_path)
         buffer: list[NDArray] = []

pyproject.toml

@@ -21,10 +21,14 @@ dependencies = [
 version = "0.2.0"
 [project.optional-dependencies]
 dev = [
+    "matplotlib",
     "pre-commit",
     "pytest",
+    "pytest-xdist",
     "pyright",
+    "setuptools",
     "trl",
+    "wandb",
     # Generate documentation
     "furo",
     "myst-parser",
@@ -51,10 +55,12 @@ reportPrivateImportUsage = false
 include = ["bergson*"]
 
 [tool.ruff]
-lint.ignore = ["E741"]   # Ambiguous variable name
+lint.ignore = ["E741", # Ambiguous variable name
+               "E501", # line-too-long (formatter takes care of it)
+]
 # Enable pycodestyle (`E`), Pyflakes (`F`), and isort (`I`) codes
 # See https://beta.ruff.rs/docs/rules/ for more possible rules
-lint.select = ["E", "F", "I"]
+lint.extend-select = ["E", "F", "I"]
 # Same as Black.
 line-length = 88
 

tests/test_build.py

@@ -80,9 +80,9 @@ def test_split_attention_build(tmp_path: Path, model, dataset):
         attention_cfgs=attention_cfgs,
     )
 
-    assert any(
-        Path(cfg.partial_run_path).iterdir()
-    ), "Expected artifacts in the temp run_path"
+    assert any(Path(cfg.partial_run_path).iterdir()), (
+        "Expected artifacts in the temp run_path"
+    )
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
@@ -107,9 +107,9 @@ def test_conv1d_build(tmp_path: Path, dataset):
         cfg=cfg,
     )
 
-    assert any(
-        Path(cfg.partial_run_path).iterdir()
-    ), "Expected artifacts in the run path"
+    assert any(Path(cfg.partial_run_path).iterdir()), (
+        "Expected artifacts in the run path"
+    )
 
     index = load_gradients(cfg.partial_run_path)
 

tests/test_reduce.py

@@ -31,6 +31,6 @@ def test_reduce_e2e(tmp_path: Path):
         text=True,
     )
 
-    assert (
-        "error" not in result.stderr.lower()
-    ), f"Error found in stderr: {result.stderr}"
+    assert "error" not in result.stderr.lower(), (
+        f"Error found in stderr: {result.stderr}"
+    )

tests/test_score.py

@@ -64,9 +64,9 @@ def test_large_gradients_query(tmp_path: Path, dataset):
     )
 
     assert result.returncode == 0
-    assert (
-        "error" not in result.stderr.lower()
-    ), f"Error found in stderr: {result.stderr}"
+    assert "error" not in result.stderr.lower(), (
+        f"Error found in stderr: {result.stderr}"
+    )
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")