init

Author: kddubey

.envrc

@@ -0,0 +1,4 @@
+#!/usr/bin/env sh
+export VIRTUAL_ENV="$PWD/.venv"
+export PATH="$VIRTUAL_ENV/bin:$PATH"
+export DEV=1

.github/workflows/ci.yml

@@ -0,0 +1,36 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+      - uses: astral-sh/setup-uv@eac588ad8def6316056a12d4907a9d4d84ff7a3b # v7.3.0
+        with:
+          enable-cache: true
+      - run: uv tool run ruff check .
+      - run: uv tool run ruff format --check .
+
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+      - uses: astral-sh/setup-uv@eac588ad8def6316056a12d4907a9d4d84ff7a3b # v7.3.0
+        with:
+          enable-cache: true
+          python-version: "3.13"
+      - run: uv sync --extra dev
+      - uses: actions/cache@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0
+        with:
+          path: ~/.cache/huggingface
+          key: hf-${{ runner.os }}-gte-modernbert-base
+      - run: uv run pytest -vv

.gitignore

@@ -0,0 +1,11 @@
+.DS_Store
+*.pyc
+*.egg-info/
+__pycache__/
+.venv/
+
+# Data
+*.csv
+!tests/fixtures/train_tiny.csv
+eval/similarities/
+eval/similarities-uncompiled/

.pre-commit-config.yaml

@@ -0,0 +1,19 @@
+repos:
+  # Ruff lint + format. Config lives in pyproject.toml.
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.15.6
+    hooks:
+      - id: ruff-check
+        args: [--fix]
+      - id: ruff-format
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v5.0.0
+    hooks:
+      - id: check-case-conflict
+      - id: check-merge-conflict
+      - id: check-yaml
+      - id: check-toml
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+      - id: debug-statements

.python-version

@@ -0,0 +1 @@
+3.13

.vscode/settings.json

@@ -0,0 +1,19 @@
+{
+    "editor.rulers": [
+        120
+    ],
+    "autoDocstring.docstringFormat": "numpy",
+    "autoDocstring.startOnNewLine": true,
+    "editor.defaultFormatter": "charliermarsh.ruff",
+    "editor.formatOnSave": true,
+    "[markdown]": {
+        "editor.formatOnSave": false
+    },
+    "editor.codeActionsOnSave": {
+        "source.organizeImports.ruff": "explicit"
+    },
+    "rewrap.autoWrap.enabled": true,
+    "notebook.formatOnSave.enabled": true,
+    "notebook.defaultFormatter": "charliermarsh.ruff",
+    "python.defaultInterpreterPath": ".venv/bin/python"
+}

README.md

@@ -0,0 +1,94 @@
+# grouping-trainer
+
+Training code for Sentry's AI grouping model.
+
+Sampling and labeling is in [data-analysis](https://github.com/getsentry/data-analysis/tree/main/grouping/data).
+
+
+## Set up locally
+
+```bash
+./bin/set_up_local.sh
+```
+
+
+## Usage
+
+### Train
+
+```bash
+python train.py --gpu h100 --run_shortname my-run
+```
+
+Train and eval metrics are logged to https://wandb.ai/sentry-seer/grouping-trainer.
+
+For DDP:
+
+```bash
+python train.py --gpu h100-ddp-4 --run_shortname my-run
+```
+
+
+### Debug
+
+<details>
+<summary>Launch a bare instance to SSH into</summary>
+
+```bash
+python -m grouping_trainer.launch --gpu h100
+```
+
+</details>
+
+
+<details>
+<summary>SSH into an instance from local</summary>
+
+Add this function to your `~/.zshrc`:
+
+```bash
+gssh() {
+    gcloud compute ssh "$1" --zone="${2:-us-central1-a}" --project=ml-ai-420606 --tunnel-through-iap
+}
+```
+
+Then:
+
+```bash
+gssh grouping-trainer-h100
+# or override the zone:
+gssh grouping-trainer-a100 europe-west4-a
+```
+
+</details>
+
+
+<details>
+<summary>Check instance output</summary>
+
+SSH into the instance and run:
+
+```bash
+logs
+# shortcut for:
+# sudo tail -f /var/log/grouping_trainer_run.log
+```
+
+If that file doesn't exist, the startup script never reached the `eval $COMMAND` block. Check what it actually did:
+
+```bash
+sudo journalctl -u google-startup-scripts.service --no-pager
+```
+
+From local (use when you can't SSH in, e.g., the boot itself failed):
+
+```bash
+gcloud compute instances get-serial-port-output grouping-trainer-l4-eval --zone=us-central1-a --project=ml-ai-420606 | tail -100
+```
+
+</details>
+
+
+### Eval
+
+See [`./eval/`](./eval/).

benchmark/README.md

@@ -0,0 +1,43 @@
+# benchmark_compiled report
+
+## Run
+
+```
+run_gcs_dir=gs://grouping-data/runs/2026-04-10-12-39-45-large-no-prefix
+df_path=final_csvs/test_full2.csv
+stamp=2026-04-24-23-44-59
+sample_size=66753
+text_prefix=''
+model_kwargs={'dtype': torch.bfloat16, 'attn_implementation': 'sdpa'}
+```
+
+- Token bucket boundaries used for analysis: `(64, 128, 256, 512, 1024)`
+- Rows: 66,753
+
+## Headline
+
+- Median compiled: **14.7 ms**
+- Median base:     **36.8 ms**
+- Per-row speedup p10/p50/p90: **1.01x / 2.51x / 3.45x**
+- Compiled wins on **92.6%** of rows
+
+## Per-bucket
+
+| bucket   | n     | tok_p50 | compiled_ms_p50 | base_ms_p50 | compiled_ms_p90 | base_ms_p90 | speedup_p50 |
+|----------|-------|---------|-----------------|-------------|-----------------|-------------|-------------|
+| <=64     | 8301  | 35.0    | 10.58           | 35.17       | 11.72           | 35.99       | 3.32        |
+| 65-128   | 7628  | 93.0    | 10.68           | 35.4        | 11.65           | 36.21       | 3.31        |
+| 129-256  | 13930 | 194.0   | 11.62           | 36.05       | 12.85           | 37.02       | 3.1         |
+| 257-512  | 14804 | 369.0   | 15.54           | 36.97       | 17.3            | 38.1        | 2.38        |
+| 513-1024 | 13858 | 682.0   | 26.9            | 38.86       | 28.97           | 40.85       | 1.44        |
+| >1024    | 8232  | 1494.5  | 48.06           | 47.04       | 115.66          | 115.89      | 0.98        |
+
+## Worst 5 rows for compiled
+
+| num_tokens | compiled_ms | base_ms | speedup |
+|------------|-------------|---------|---------|
+| 1252       | 54.8        | 43.26   | 0.789   |
+| 1043       | 50.24       | 42.14   | 0.839   |
+| 2020       | 82.6        | 70.98   | 0.859   |
+| 1122       | 49.67       | 42.91   | 0.864   |
+| 2040       | 79.23       | 69.04   | 0.871   |