Merge branch 'main' into dependabot/github_actions/actions/upload-artifact-5

Author: XkunW

.github/workflows/code_checks.yml

@@ -30,7 +30,7 @@ jobs:
     steps:
       - uses: actions/checkout@v5.0.0
       - name: Install uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           # Install a specific version of uv.
           version: "0.5.21"
@@ -40,7 +40,7 @@ jobs:
         with:
           python-version-file: ".python-version"
       - name: Install the project
-        run: uv sync --dev
+        run: uv sync --dev --prerelease=allow
       - name: Install dependencies and check code
         run: |
           source .venv/bin/activate

.github/workflows/docs.yml

@@ -56,7 +56,7 @@ jobs:
           fetch-depth: 0  # Fetch all history for proper versioning
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           version: "0.5.21"
           enable-cache: true
@@ -67,10 +67,10 @@ jobs:
           python-version-file: ".python-version"
 
       - name: Install the project
-        run: uv sync --all-extras --group docs
+        run: uv sync --all-extras --group docs --prerelease=allow
 
       - name: Build docs
-        run: uv run mkdocs build
+        run: uv run --frozen mkdocs build
 
       - name: Create .nojekyll file
         run: touch site/.nojekyll
@@ -93,7 +93,7 @@ jobs:
           fetch-depth: 0  # Fetch all history for proper versioning
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           version: "0.5.21"
           enable-cache: true
@@ -104,7 +104,7 @@ jobs:
           python-version-file: ".python-version"
 
       - name: Install the project
-        run: uv sync --all-extras --group docs
+        run: uv sync --all-extras --group docs --frozen
 
       - name: Configure Git Credentials
         run: |

.github/workflows/publish.yml

@@ -16,7 +16,7 @@ jobs:
       - uses: actions/checkout@v5.0.0
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           version: "0.6.6"
           enable-cache: true

.github/workflows/unit_tests.yml

@@ -46,7 +46,7 @@ jobs:
       - uses: actions/checkout@v5.0.0
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           # Install a specific version of uv.
           version: "0.5.21"
@@ -58,18 +58,18 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Install the project
-        run: uv sync --dev
+        run: uv sync --dev --prerelease=allow
 
       - name: Install dependencies and check code
         run: |
-          uv run pytest -m "not integration_test" --cov vec_inf --cov-report=xml tests
+          uv run --frozen pytest -m "not integration_test" --cov vec_inf --cov-report=xml tests
 
       - name: Install the core package only
         run: uv sync --no-dev
 
       - name: Run package import tests
         run: |
-          uv run pytest tests/test_imports.py
+          uv run --frozen pytest tests/test_imports.py
 
       - name: Import Codecov GPG public key
         run: |
@@ -79,7 +79,7 @@ jobs:
         uses: codecov/codecov-action@v5.5.1
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
-          file: ./coverage.xml
+          files: ./coverage.xml
           name: codecov-umbrella
           fail_ci_if_error: true
           verbose: true

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
     - id: check-toml
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.13.2'
+    rev: 'v0.14.3'
     hooks:
     - id: ruff
       args: [--fix, --exit-non-zero-on-fix]

Dockerfile

@@ -35,29 +35,33 @@ RUN wget https://bootstrap.pypa.io/get-pip.py && \
     rm get-pip.py && \
     python3.10 -m pip install --upgrade pip setuptools wheel uv
 
-# Install Infiniband/RDMA support
+# Install RDMA support
 RUN apt-get update && apt-get install -y \
     libibverbs1 libibverbs-dev ibverbs-utils \
     librdmacm1 librdmacm-dev rdmacm-utils \
+    rdma-core ibverbs-providers infiniband-diags perftest \
     && rm -rf /var/lib/apt/lists/*
 
 # Set up RDMA environment (these will persist in the final container)
 ENV LD_LIBRARY_PATH="/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH"
-ENV UCX_NET_DEVICES=all
 ENV NCCL_IB_DISABLE=0
+ENV NCCL_SOCKET_IFNAME="^lo,docker0"
+ENV NCCL_NET_GDR_LEVEL=PHB
+ENV NCCL_IB_TIMEOUT=22
+ENV NCCL_IB_RETRY_CNT=7
+ENV NCCL_DEBUG=INFO
 
 # Set up project
 WORKDIR /vec-inf
 COPY . /vec-inf
 
 # Install project dependencies with build requirements
-RUN PIP_INDEX_URL="https://download.pytorch.org/whl/cu128" uv pip install --system -e .[dev]
+RUN uv pip install --system -e .[dev] --prerelease=allow
 
-# Final configuration
-RUN mkdir -p /vec-inf/nccl && \
-    mv /root/.config/vllm/nccl/cu12/libnccl.so.2.18.1 /vec-inf/nccl/libnccl.so.2.18.1
-ENV VLLM_NCCL_SO_PATH=/vec-inf/nccl/libnccl.so.2.18.1
-ENV NCCL_DEBUG=INFO
+# Install a single, system NCCL (from NVIDIA CUDA repo in base image)
+RUN apt-get update && apt-get install -y --allow-change-held-packages\
+    libnccl2 libnccl-dev \
+    && rm -rf /var/lib/apt/lists/*
 
 # Set the default command to start an interactive shell
 CMD ["bash"]

MODEL_TRACKING.md

@@ -40,6 +40,7 @@ This document tracks all model weights available in the `/model-weights` directo
 | `gemma-2b-it` | ❌ |
 | `gemma-7b` | ❌ |
 | `gemma-7b-it` | ❌ |
+| `gemma-2-2b-it` | ✅ |
 | `gemma-2-9b` | ✅ |
 | `gemma-2-9b-it` | ✅ |
 | `gemma-2-27b` | ✅ |

pyproject.toml

@@ -42,9 +42,10 @@ dev = [
     "xgrammar>=0.1.11",
     "torch>=2.7.0",
     "vllm>=0.10.0",
-    "vllm-nccl-cu12>=2.18,<2.19",
-    "ray>=2.40.0",
-    "cupy-cuda12x==12.1.0"
+    "ray[default]>=2.50.0",
+    "cupy-cuda12x==12.1.0",
+    "flashinfer-python>=0.4.0",
+    "sglang>=0.5.0",
 ]
 
 [project.scripts]

tests/vec_inf/cli/test_cli.py

@@ -39,6 +39,7 @@ def test_launch_command_success(runner):
             "mem_per_node": "32G",
             "model_weights_parent_dir": "/model-weights",
             "vocab_size": "128000",
+            "venv": "/path/to/venv",
             "vllm_args": {"max_model_len": 8192},
             "env": {"CACHE": "/cache"},
         }

tests/vec_inf/cli/test_helper.py

@@ -35,6 +35,7 @@ def test_format_table_output(self):
             "mem_per_node": "32G",
             "model_weights_parent_dir": "/model-weights",
             "log_dir": "/tmp/logs",
+            "venv": "/path/to/venv",
             "vllm_args": {"max_model_len": 8192, "enable_prefix_caching": True},
             "env": {"CACHE": "/cache"},
         }
@@ -63,6 +64,7 @@ def test_format_table_output_with_minimal_params(self):
             "mem_per_node": "16G",
             "model_weights_parent_dir": "/weights",
             "log_dir": "/logs",
+            "venv": "/path/to/venv",
             "vllm_args": {},
             "env": {},
         }