Move to UV (#59)

Author: msaroufim

.github/workflows/ruff.yml

@@ -13,16 +13,17 @@ jobs:
     steps:
     - uses: actions/checkout@v4
 
+    - name: Install uv
+      uses: astral-sh/setup-uv@v3
+      
     - name: Set up Python
-      uses: actions/setup-python@v4
-      with:
-        python-version: '3.x'
+      run: uv python install
 
     - name: Install package with dev dependencies
-      run: pip install -e .[dev]
+      run: uv sync --dev
 
     - name: Run ruff check
-      run: ruff check .
+      run: uv run ruff check .
 
     - name: Run ruff format check
-      run: ruff format --check . 
+      run: uv run ruff format --check . 

.github/workflows/smoke-test.yml

@@ -15,19 +15,17 @@ jobs:
     steps:
     - uses: actions/checkout@v4
 
+    - name: Install uv
+      uses: astral-sh/setup-uv@v3
+    
     - name: Set up Python
-      uses: actions/setup-python@v4
-      with:
-        python-version: '3.x'
+      run: uv python install
 
     - name: Install package and dependencies
-      run: |
-        pip install -e .[dev]
+      run: uv sync --dev
 
     - name: Run smoke test
-      run: |
-        python -m BackendBench.scripts.main --suite smoke --backend aten 
+      run: uv run python -m BackendBench.scripts.main --suite smoke --backend aten 
 
     - name: Run pytest tests
-      run: |
-        pytest test/
+      run: uv run pytest test/

.gitignore

@@ -7,3 +7,4 @@ backendbench.egg-info/
 CLAUDE.md
 venv/
 ops/
+uv.lock

README.md

@@ -10,31 +10,43 @@ We provide both
 1. Comprehensive operator level correctness checks using the PyTorch OpInfo test suite
 2. Performance checks using the ops that show up in the most popular Hugging Face models with realistic tensor shapes
 
+# Installation:
+
+Install using uv (recommended):
+```bash
+uv add backendbench
+```
+
+Or install in development mode:
+```bash
+uv sync --dev
+```
+
 # Usage:
 
 Run a simple smoke test (relu) with the default ATen backend:
 ```bash
-python scripts/main.py --suite smoke --backend aten
+uv run python scripts/main.py --suite smoke --backend aten
 ```
 
 Run the smoke test with FlagGems:
 ```bash
-python scripts/main.py --suite smoke --backend flag_gems
+uv run python scripts/main.py --suite smoke --backend flag_gems
 ```
 
 Run opinfo tests (correctness only) with ATen
 ```bash
-python scripts/main.py --suite opinfo --backend aten
+uv run python scripts/main.py --suite opinfo --backend aten
 ```
 
 Run a filtered set of opinfo tests with FlagGems
 ```bash
-python scripts/main.py --suite opinfo --backend flag_gems --ops "add,sub"
+uv run python scripts/main.py --suite opinfo --backend flag_gems --ops "add,sub"
 ```
 
 Run all the opinfo tests with FlagGems (takes a few minutes)
 ```bash
-python scripts/main.py --suite opinfo --backend flag_gems
+uv run python scripts/main.py --suite opinfo --backend flag_gems
 ```
 
 ## LLM-Based Kernel Generation and Evaluation
@@ -44,7 +56,7 @@ Generate and evaluate PyTorch kernels using Claude API:
 Run LLM evaluation on smoke test (relu operation):
 ```bash
 export ANTHROPIC_API_KEY=your_api_key_here
-python scripts/main.py --suite smoke --backend llm
+uv run python scripts/main.py --suite smoke --backend llm
 ```
 
 ## KernelAgent-Based Triton Kernel Generation
@@ -59,19 +71,19 @@ git submodule update --init --recursive
 Run KernelAgent evaluation on smoke test (relu operation):
 ```bash
 export OPENAI_API_KEY=your_api_key_here
-python scripts/main.py --suite smoke --backend kernel_agent
+uv run python scripts/main.py --suite smoke --backend kernel_agent
 ```
 
 Run KernelAgent with custom configuration:
 ```bash
 export OPENAI_API_KEY=your_api_key_here
-python scripts/main.py --suite smoke --backend kernel_agent --kernel-agent-workers 6 --kernel-agent-max-rounds 15
+uv run python scripts/main.py --suite smoke --backend kernel_agent --kernel-agent-workers 6 --kernel-agent-max-rounds 15
 ```
 
 Run KernelAgent on opinfo tests with a specific operation:
 ```bash
 export OPENAI_API_KEY=your_api_key_here
-python scripts/main.py --suite opinfo --backend kernel_agent --ops "add"
+uv run python scripts/main.py --suite opinfo --backend kernel_agent --ops "add"
 ```
 
 ## Directory-Based Kernel Development
@@ -160,10 +172,10 @@ if __name__ == "__main__":
 
 Test individual implementations:
 ```bash
-python generated_kernels/relu/relu_implementation_1.py
+uv run python generated_kernels/relu/relu_implementation_1.py
 ```
 
 Test with BackendBench:
 ```bash
-python scripts/main.py --suite smoke --backend directory
+uv run python scripts/main.py --suite smoke --backend directory
 ```

pyproject.toml

@@ -1,6 +1,6 @@
 [build-system]
-requires = ["setuptools>=61.0", "wheel"]
-build-backend = "setuptools.build_meta"
+requires = ["hatchling"]
+build-backend = "hatchling.build"
 
 [project]
 name = "backendbench"
@@ -27,25 +27,27 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-dev = [
+flaggems = [
+    # flag_gems must be installed from source: https://github.com/FlagOpen/FlagGems
+]
+
+[project.scripts]
+backendbench = "BackendBench.scripts.main:cli"
+
+[tool.hatch.build.targets.wheel]
+packages = ["BackendBench"]
+
+[tool.uv]
+dev-dependencies = [
     "pytest",
-    "pytest-cov",
+    "pytest-cov", 
     "pytest-mock",
     "pytest-timeout",
     "ruff==0.12.1",
     "torch",
     "numpy",
-    "cupy-cuda12x",
+    # cupy-cuda12x is platform specific, install manually if needed
 ]
-flaggems = [
-    "flag_gems",
-]
-
-[project.scripts]
-backendbench = "BackendBench.scripts.main:cli"
-
-[tool.setuptools.packages.find]
-include = ["BackendBench*"]
 
 [tool.ruff]
 line-length = 100

test/test_utils.py

@@ -28,9 +28,9 @@ def func_with_pytorch_stream():
 
     @pytest.mark.skipif(not HAS_CUDA, reason="CUDA not available")
     def test_cupy_stream_creation(self):
-        import cupy
-
         """Test detection of CuPy CUDA stream creation."""
+        pytest.importorskip("cupy")
+        import cupy
 
         def func_with_cupy_stream():
             stream = cupy.cuda.Stream()