[Feature] Use uv for package management and installation (#485)

* [Feature] Use `uv` for package management and installation

Author: nuzant

.pre-commit-config.yaml

@@ -68,3 +68,14 @@ repos:
         pass_filenames: false
         always_run: false
         require_serial: true
+
+  # Generate requirements.txt with uv
+  - repo: https://github.com/astral-sh/uv-pre-commit
+    # uv version.
+    rev: 0.9.5
+    hooks:
+      # Compile requirements
+      - id: pip-compile
+        name: Compile requirements.txt using uv
+        files: ^requirements\.txt|pyproject\.toml$
+        args: [pyproject.toml, -o, requirements.txt]

MANIFEST.in

@@ -0,0 +1 @@
+recursive-include evaluation/latex2sympy *

areal/tests/test_model_utils.py

@@ -23,8 +23,6 @@
 class InvalidConfig(BaseExperimentConfig):
     """Invalid config for testing - missing actor/model attribute."""
 
-    pass
-
 
 def _create_base_cluster_config(fileroot="/tmp/areal_test"):
     """Helper to create a cluster config with common test values."""

docs/tutorial/installation.md

@@ -65,19 +65,30 @@ conda create -n areal python=3.12
 conda activate areal
 ```
 
-3. Install pip dependencies:
+3. Install pip dependencies using uv:
 
 ```bash
 git clone https://github.com/inclusionAI/AReaL
 cd AReaL
-bash examples/env/setup-pip-deps.sh
+pip install uv
+uv pip install -e .[all]
 ```
 
-**Note**: Installing with `examples/env/setup-pip-deps.sh` will install
-`flash-attn==2.8.3` since it does not require compilation with torch version 2.8.0.
-However, `flash-attn==2.8.3` is not compatible with Megatron training backend. If you
-want to use Megatron training backend, please compile and install `flash-attn==2.8.1` in
-your custom environment, or use docker installation instead.
+**Note**: Directly install with `uv` and `pip` will install `flash-attn==2.8.3` since it
+does not require compilation with torch version 2.8.0. However, `flash-attn==2.8.3` is
+not compatible with Megatron training backend. If you want to use Megatron training
+backend, please compile and install `flash-attn==2.8.1` in your custom environment, or
+use docker installation instead.
+
+4. Validate your AReaL installation:
+
+We provide a script to validate AReaL installation. Simply run:
+
+```bash
+python3 examples/env/validate_installation.py
+```
+
+After installation validation passed, you are good to go!
 
 ## (Optional) Launch Ray Cluster for Distributed Training
 

examples/env/setup-eval-pip-deps.sh

@@ -10,7 +10,6 @@
 import importlib
 import sys
 from importlib.metadata import version as get_version
-from typing import Optional
 
 from packaging.version import Version
 
@@ -25,7 +24,7 @@ def test_import(
         self,
         module_name: str,
         required: bool = True,
-        test_func: Optional[callable] = None,
+        test_func: callable | None = None,
     ) -> bool:
         """Test importing a module and optionally run additional tests."""
         try:
@@ -78,9 +77,9 @@ def test_vllm_functionality(self, vllm_module):
         """Test vLLM basic functionality."""
         from vllm import LLM, SamplingParams  # noqa
 
-        assert Version(get_version("vllm")) == Version(
-            "0.10.2"
-        ), f"vLLM version should be 0.10.2, found {Version(get_version('vllm'))}"
+        assert Version(get_version("vllm")) == Version("0.10.2"), (
+            f"vLLM version should be 0.10.2, found {Version(get_version('vllm'))}"
+        )
 
         print("  - vLLM core classes imported successfully")
 
@@ -96,15 +95,15 @@ def test_sglang_functionality(self, sglang_module):
         )
         from sglang import Engine, launch_server  # noqa
 
-        assert Version(get_version("sglang")) >= Version(
-            "v0.5.2"
-        ), "SGLang version should be >= v0.5.2"
+        assert Version(get_version("sglang")) >= Version("v0.5.2"), (
+            "SGLang version should be >= v0.5.2"
+        )
         print("  - SGLang imported successfully")
 
     def test_transformers(self, transformers_module):
-        assert Version(get_version("transformers")) == Version(
-            "4.56.1"
-        ), "transformers version should be 4.56.1"
+        assert Version(get_version("transformers")) == Version("4.56.1"), (
+            "transformers version should be 4.56.1"
+        )
         print("  - transformers imported successfully")
 
     def validate_critical_dependencies(self):
@@ -121,7 +120,6 @@ def validate_critical_dependencies(self):
         self.test_import(
             "flash_attn", required=True, test_func=self.test_flash_attn_functionality
         )
-        self.test_import("cugae", required=True)
         # Inference engines
         self.test_import(
             "sglang", required=True, test_func=self.test_sglang_functionality
@@ -175,14 +173,15 @@ def validate_optional_dependencies(self):
         self.test_import("seaborn", required=False)
         self.test_import("numba", required=False)
         self.test_import("nltk", required=False)
+        self.test_import("cugae", required=False)
 
     def test_te_functionality(self, _):
         try:
             import torch
 
-            assert Version(get_version("transformer_engine")) >= Version(
-                "2.3.0"
-            ), "transformer_engine version must be larger than 2.3.0"
+            assert Version(get_version("transformer_engine")) >= Version("2.3.0"), (
+                "transformer_engine version must be larger than 2.3.0"
+            )
 
             if torch.cuda.is_available():
                 import transformer_engine.pytorch as te
@@ -307,14 +306,14 @@ def run_validation(self):
 
         # Determine overall result
         if self.critical_failures:
-            print(f"\n❌ INSTALLATION VALIDATION FAILED")
+            print("\n❌ INSTALLATION VALIDATION FAILED")
             print("Please check the critical failures above and ensure all required")
             print(
                 "dependencies are properly installed according to the installation guide."
             )
             return False
         else:
-            print(f"\n✅ INSTALLATION VALIDATION PASSED")
+            print("\n✅ INSTALLATION VALIDATION PASSED")
             if self.warnings:
                 print("Note: Some optional dependencies failed but this won't affect")
                 print("core functionality.")

examples/env/setup-pip-deps.sh

@@ -41,13 +41,20 @@ classifiers = [
 
 dependencies = [
     # Core ML/AI libraries
-    "torch>2.0.0",
+    "torch==2.8.0",
+    "torchaudio",
+    "torchvision",
+
     "huggingface_hub",
     "datasets",
     "accelerate",
     "transformers==4.56.1",
     "megatron-core==0.13.1",
     "mbridge==0.13.0",
+    "flashinfer-python==0.3.1",
+    "sglang[all]==0.5.2",
+    "flash-attn==2.8.3",
+    "vllm==0.10.2",
     "peft",
     "qwen_agent",
 
@@ -68,6 +75,7 @@ dependencies = [
     "orjson>=3.10.16",
     "pydantic",
     "PyYAML",
+    "omegaconf==2.4.0.dev2",
     "hydra-core==1.4.0.dev1",
     "packaging",
     "lark",
@@ -78,6 +86,7 @@ dependencies = [
     "tensordict",
     "pybase64",
     "msgspec",
+    "latex2sympy2",
     "openai==1.99.6",
     "dotenv",
     "json5",
@@ -89,7 +98,9 @@ dependencies = [
     "colorlog",
     "psutil",
     "pynvml",
-    "swanlab[dashboard]",
+    "nvidia-ml-py",
+    "swanboard==0.1.9b1",
+    "swanlab[dashboard]==0.6.12",
 
     # Performance and compression
     "ninja",
@@ -174,6 +185,19 @@ where = ["."]
 include = ["areal*"]
 exclude = ["tests*", "realhf*", "docs*", "examples*", "evaluation*", "benchmark*"]
 
+[tool.uv]
+override-dependencies=[
+    "openai==1.99.6",
+    "xgrammar==0.1.24",
+    "outlines-core==0.1.26",
+]
+
+[tool.uv.sources]
+latex2sympy2 = { path = "./evaluation/latex2sympy" }
+
+[tool.uv.extra-build-dependencies]
+flash-attn = ["torch==2.8.0"]
+
 [tool.pytest.ini_options]
 pythonpath = ["."]
 filterwarnings = [