Merge pull request #82 from apoorvkh/uv-v2

Switching to uv (attempt 2)

Author: pmcurtin

.github/workflows/main.yml

@@ -9,36 +9,21 @@ on:
 
 jobs:
 
-  check:
+  checks:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
-      - uses: prefix-dev/setup-[email protected]
+      - uses: astral-sh/setup-[email protected]
         with:
-          pixi-version: v0.27.1
-          frozen: true
-          cache: true
-          environments: dev
-          activate-environment: dev
-      - run: ruff check
+          version: "0.5.0"
+          python-version-file: ".python-version"
+          enable-cache: true
+      - run: uv sync
+      - run: uv run --frozen ruff check
         if: success() || failure()
-      - run: ruff format --check
+      - run: uv run --frozen ruff format --check
         if: success() || failure()
-
-  ##
-
-  typecheck:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - uses: prefix-dev/[email protected]
-        with:
-          pixi-version: v0.27.1
-          frozen: true
-          cache: true
-          environments: extra
-          activate-environment: extra
-      - run: pyright
+      - run: uv run --frozen pyright
         if: success() || failure()
 
   ##
@@ -64,26 +49,16 @@ jobs:
   test:
     runs-on: ubuntu-latest
     needs: get-pytorch-versions
-    strategy: 
+    strategy:
       fail-fast: false
       matrix:
         python: ["3.9", "3.10", "3.11", "3.12"] 
         pytorch: ${{fromJson(needs.get-pytorch-versions.outputs.versions)}}
-
     steps:
       - uses: actions/checkout@v4
-
-      - name: Inject (python / pytorch) test deps into pixi.toml / pyproject.toml
-        run: |
-          sed -i 's/python = "3.9"/python = "${{ matrix.python }}.*"/' pixi.toml
-          sed -i 's/torch>=2.0.0/torch~=${{ matrix.pytorch }}/' pyproject.toml
-
-      - uses: prefix-dev/[email protected]
+      - uses: astral-sh/[email protected]
         with:
-          pixi-version: v0.27.1
-          locked: false
-          frozen: false
-          cache: false
-          environments: default
-          activate-environment: default
-      - run: pytest tests/test_ci.py
+          version: "0.5.0"
+      - if: contains('2.0,2.1,2.2', matrix.pytorch)
+        run: echo "NUMPY_VERSION=--with \"numpy<2\"" >> $GITHUB_ENV
+      - run: uv run --python ${{ matrix.python }} --with torch~=${{ matrix.pytorch }} ${{ env.NUMPY_VERSION }} pytest --verbose tests/test_ci.py

.github/workflows/release.yml

@@ -9,15 +9,11 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
-      - uses: prefix-dev/setup-[email protected]
+      - uses: astral-sh/setup-uv@v3
         with:
-          pixi-version: v0.27.1
-          frozen: true
-          cache: true
-          environments: dev
-          activate-environment: dev
-      - run: python -m build
-      - run: twine upload dist/*
-        env:
-          TWINE_USERNAME: __token__
-          TWINE_PASSWORD: ${{ secrets.PYPI_TOKEN }}
+          version: "0.5.0"
+          python-version-file: ".python-version"
+          enable-cache: true
+      - run: uv sync
+      - run: uv build
+      - run: uv publish

.python-version

@@ -0,0 +1 @@
+3.9.20

CONTRIBUTING.md

@@ -1,10 +1,10 @@
 # Contributing
 
-We use the [`pixi`](https://pixi.sh) package manager. Simply [install `pixi`](https://pixi.sh/latest/#installation) and run `pixi shell` in this repository to activate the environment.
+We use the [`uv`](https://github.com/astral-sh/uv) package manager. Simply [install `uv`](https://github.com/astral-sh/uv#installation) and run `uv sync` in this repository to build the environment. Run `source .venv/bin/activate` to activate the environment.
 
 We use `ruff check` for linting, `ruff format` for formatting, `pyright` for static type checking, and `pytest` for testing.
 
-We build wheels with `python -m build` and upload to [PyPI](https://pypi.org/project/torchrunx) with [twine](https://twine.readthedocs.io). Our release pipeline is powered by Github Actions.
+We build wheels with `uv build` and upload to [PyPI](https://pypi.org/project/torchrunx) with `uv publish`. Our release pipeline is powered by Github Actions.
 
 ## Pull Requests
 

pixi.lock

@@ -22,16 +22,19 @@ classifiers = [
 ]
 requires-python = ">=3.9"
 dependencies = [
-    "cloudpickle>=3.0",
-    "fabric>=3.2",
-    "torch>=2.0",
-    # torch.distributed depends on numpy
-    "numpy>=1.20",
+  "cloudpickle>=3.0",
+  "fabric>=3.2",
+  "torch>=2.0",
+  # torch.distributed depends on numpy
+  "numpy>=1.20",
 ]
+[dependency-groups]
+dev = ["ruff", "pyright", "pytest", "build", "twine"]
+dev-extras = ["submitit", "transformers"]
 
-[project.urls]
-Repository = "https://github.com/apoorvkh/torchrunx.git"
-Documentation = "https://torchrunx.readthedocs.io"
+[tool.uv]
+managed = true
+python-preference = "only-managed"
 
 ## Development tools
 

pixi.toml

@@ -2,10 +2,14 @@
 
 import copy
 
-import submitit
+import submitit  # pyright: ignore [reportMissingImports]
 import torch
 from torch.utils.data import Dataset
-from transformers import BertForMaskedLM, Trainer, TrainingArguments
+from transformers import (  # pyright: ignore [reportMissingImports]
+    BertForMaskedLM,
+    Trainer,
+    TrainingArguments,
+)
 
 import torchrunx as trx
 

pyproject.toml

@@ -1,20 +1,21 @@
 import os
 
+import torch
+
 import torchrunx as trx
 
 
-def worker() -> None:
-    import torch
+class MLP(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.a = torch.nn.Linear(10, 10, bias=False)
+        self.b = torch.nn.Linear(10, 1, bias=False)
 
-    class MLP(torch.nn.Module):
-        def __init__(self) -> None:
-            super().__init__()
-            self.a = torch.nn.Linear(10, 10, bias=False)
-            self.b = torch.nn.Linear(10, 1, bias=False)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.b(self.a(x))
 
-        def forward(self, x: torch.Tensor) -> torch.Tensor:
-            return self.b(self.a(x))
 
+def worker() -> None:
     local_rank = int(os.environ["LOCAL_RANK"])
     print("init model")
     model = MLP().to(local_rank)
@@ -33,8 +34,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 def test_distributed_train() -> None:
     trx.launch(
         worker,
-        hostnames="slurm",
-        workers_per_host="slurm",
         backend="nccl",
     )