Try uv for package management (245ms project installation)

Dockerfile

@@ -85,8 +85,12 @@ RUN python -m nltk.downloader punkt
 COPY open_instruct open_instruct
 
 # install the package in editable mode
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv
 COPY pyproject.toml .
-RUN pip install -e .
+COPY README.md .
+COPY uv.lock .
+RUN uv sync
+RUN uv sync --extra compile --no-build-isolation
 COPY .git .
 
 COPY eval eval

README.md

@@ -54,10 +54,10 @@ pip install -r weight-diff-requirements.txt
 For a second installation strategy, if you'd like to *run experiments within a Docker environment*, you can create one using:
 
 ```bash
-docker build --build-arg CUDA=12.1.0 --build-arg TARGET=cudnn8-devel --build-arg DIST=ubuntu20.04 --build-arg REQUIRE=requirements.txt . -t open_instruct
+docker build --build-arg CUDA=12.1.0 --build-arg TARGET=cudnn8-devel --build-arg DIST=ubuntu20.04 --build-arg REQUIRE=requirements.txt . -t open_instruct_dev_uv
 
 # if you are interally at AI2, you can create an image like this:
-beaker image create open_instruct -n open_instruct -w ai2/$(whoami)
+beaker image create open_instruct_dev_uv -n open_instruct_dev_uv -w ai2/$(whoami)
 ```
 
 If you are internally at AI2, you can use this pre-built beaker image `hamishivi/open-instruct-eval` (most recent version [here](https://beaker.org/im/01J2CKY81A6N1WG5QS08Y3WNM5/details)). For finetuning, you can use `hamishivi/open-instruct-public` (most recent version [here](https://beaker.org/im/01J2CQFX7076PDHZJR2GB0C3A9/details)). I will try to update these periodically.

pyproject.toml

@@ -1,17 +1,55 @@
-[build-system]
-requires = ["setuptools", "setuptools-scm"]
-build-backend = "setuptools.build_meta"
-
 [project]
-name = "open_instruct"
-authors = [
-    {name = "Example Author", email = "[email protected]"}
-]
-description = "package description"
-version = "0.0.1"
+name = "open-instruct"
+version = "0.1.0"
+description = "Add your description here"
 readme = "README.md"
-requires-python = ">=3.7"
-dependencies = []
+requires-python = ">=3.12"
+dependencies = [
+    # core dependencies
+    "torch>=2.4.0",
+    "datasets>=2.21.0",
+    "deepspeed>=0.14.5",
+    "accelerate>=0.33.0",
+    "peft>=0.12.0",
+    "bitsandbytes>=0.43.3",
+    "evaluate>=0.4.2",
+    "transformers>=4.44.1",
+    "openai>=1.42.0",
+    "tiktoken>=0.7.0",
+    "tensorboard>=2.17.1",
+    "wandb>=0.17.7",
+    "fire>=0.6.0",
+    "vllm>=0.5.4",
+    "hf-transfer>=0.1.8",
+    "hatchling>=1.25.0",
+    "einops>=0.8.0",
+    "editables>=0.5",
+    "rich>=13.7.1",
+    "matplotlib>=3.9.2",
+    "nltk>=3.9.1",
+    "langdetect>=1.0.9",
+    "immutabledict>=4.2.0",
+    "antlr4-python3-runtime>=4.13.2",
+    "mpmath>=1.3.0",
+    "sympy>=1.13.2",
+]
 
-[tool.setuptools]
-py-modules = ["open_instruct"]
+[project.optional-dependencies]
+compile = [
+    "flash-attn",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.uv]
+dev-dependencies = [
+    "black>=24.8.0",
+    "flake8>=7.1.1",
+    "isort>=5.13.2",
+    "autoflake>=2.3.1",
+    "pytest>=8.3.2",
+    "beaker-py>=1.31.2",
+]
+no-build-isolation-package = ["flash-attn"]

requirements.txt

@@ -47,5 +47,5 @@ flake8
 isort
 autoflake
 pytest
-hf_transfer
 beaker-py
+hf_transfer