Merge pull request #130 from VectorInstitute/feature/killarney-migration

* Adapted Slurm configuration changes to Killarney cluster while keeping support for Bon Echo
* Removed redundant code in generated script
* Added RDMA support
* Added more default value support in environment config
* Sorted list command
* Update config README to track cached model weights

Author: XkunW

.github/workflows/unit_tests.yml

@@ -71,6 +71,10 @@ jobs:
         run: |
           uv run pytest tests/test_imports.py
 
+      - name: Import Codecov GPG public key
+        run: |
+          gpg --keyserver keyserver.ubuntu.com --recv-keys 806BB28AED779869
+
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v5.5.0
         with:

.gitignore

@@ -152,3 +152,7 @@ collect_env.py
 
 # build files
 dist/
+
+# type stubs
+stubs/
+mypy.ini

Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04
+FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
 
 # Non-interactive apt-get commands
 ARG DEBIAN_FRONTEND=noninteractive
@@ -35,16 +35,29 @@ 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
+RUN apt-get update && apt-get install -y \
+    libibverbs1 libibverbs-dev ibverbs-utils \
+    librdmacm1 librdmacm-dev rdmacm-utils \
+    && 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
+
 # Set up project
 WORKDIR /vec-inf
 COPY . /vec-inf
 
 # Install project dependencies with build requirements
-RUN PIP_INDEX_URL="https://download.pytorch.org/whl/cu121" uv pip install --system -e .[dev]
+RUN PIP_INDEX_URL="https://download.pytorch.org/whl/cu128" uv pip install --system -e .[dev]
 
 # 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
 
 # Set the default command to start an interactive shell
 CMD ["bash"]

README.md

@@ -7,18 +7,20 @@
 [![code checks](https://github.com/VectorInstitute/vector-inference/actions/workflows/code_checks.yml/badge.svg)](https://github.com/VectorInstitute/vector-inference/actions/workflows/code_checks.yml)
 [![docs](https://github.com/VectorInstitute/vector-inference/actions/workflows/docs.yml/badge.svg)](https://github.com/VectorInstitute/vector-inference/actions/workflows/docs.yml)
 [![codecov](https://codecov.io/github/VectorInstitute/vector-inference/branch/main/graph/badge.svg?token=NI88QSIGAC)](https://app.codecov.io/github/VectorInstitute/vector-inference/tree/main)
-[![vLLM](https://img.shields.io/badge/vllm-0.9.2)](https://docs.vllm.ai/en/v0.9.2/index.html)
+[![vLLM](https://img.shields.io/badge/vLLM-0.10.1.1-blue)](https://docs.vllm.ai/en/v0.10.1.1/)
 ![GitHub License](https://img.shields.io/github/license/VectorInstitute/vector-inference)
 
-This repository provides an easy-to-use solution to run inference servers on [Slurm](https://slurm.schedmd.com/overview.html)-managed computing clusters using [vLLM](https://docs.vllm.ai/en/latest/). **All scripts in this repository runs natively on the Vector Institute cluster environment**. To adapt to other environments, follow the instructions in [Installation](#installation).
+This repository provides an easy-to-use solution to run inference servers on [Slurm](https://slurm.schedmd.com/overview.html)-managed computing clusters using [vLLM](https://docs.vllm.ai/en/latest/). **This package runs natively on the Vector Institute cluster environments**. To adapt to other environments, follow the instructions in [Installation](#installation).
+
+**NOTE**: Supported models on Killarney are tracked [here](vec_inf/config/README.md)
 
 ## Installation
 If you are using the Vector cluster environment, and you don't need any customization to the inference server environment, run the following to install package:
 
 ```bash
 pip install vec-inf
 ```
-Otherwise, we recommend using the provided [`Dockerfile`](Dockerfile) to set up your own environment with the package. The latest image has `vLLM` version `0.9.2`.
+Otherwise, we recommend using the provided [`Dockerfile`](Dockerfile) to set up your own environment with the package. The latest image has `vLLM` version `0.10.1.1`.
 
 If you'd like to use `vec-inf` on your own Slurm cluster, you would need to update the configuration files, there are 3 ways to do it:
 * Clone the repository and update the `environment.yaml` and the `models.yaml` file in [`vec_inf/config`](vec_inf/config/), then install from source by running `pip install .`.
@@ -42,23 +44,26 @@ You should see an output like the following:
 
 <img width="600" alt="launch_image" src="https://github.com/user-attachments/assets/a72a99fd-4bf2-408e-8850-359761d96c4f">
 
+**NOTE**: On Vector Killarney Cluster environment, the following fields are required:
+  * `--account`, `-A`: The Slurm account, this argument can be set to default by setting environment variable `VEC_INF_ACCOUNT`.
+  * `--work-dir`, `-D`: A working directory other than your home directory, this argument can be set to default by seeting environment variable `VEC_INF_WORK_DIR`.
 
 #### Overrides
 
 Models that are already supported by `vec-inf` would be launched using the cached configuration (set in [slurm_vars.py](vec_inf/client/slurm_vars.py)) or [default configuration](vec_inf/config/models.yaml). You can override these values by providing additional parameters. Use `vec-inf launch --help` to see the full list of parameters that can be
-overriden. For example, if `qos` is to be overriden:
+overriden. For example, if `resource-type` is to be overriden:
 
 ```bash
-vec-inf launch Meta-Llama-3.1-8B-Instruct --qos <new_qos>
+vec-inf launch Meta-Llama-3.1-8B-Instruct --resource-type <new_resource_type>
 ```
 
-To overwrite default vLLM engine arguments, you can specify the engine arguments in a comma separated string:
+To overwrite default `vllm serve` arguments, you can specify the arguments in a comma separated string:
 
 ```bash
 vec-inf launch Meta-Llama-3.1-8B-Instruct --vllm-args '--max-model-len=65536,--compilation-config=3'
 ```
 
-For the full list of vLLM engine arguments, you can find them [here](https://docs.vllm.ai/en/stable/serving/engine_args.html), make sure you select the correct vLLM version.
+For the full list of `vllm serve` arguments, you can find them [here](https://docs.vllm.ai/en/stable/cli/serve.html), make sure you select the correct vLLM version.
 
 #### Custom models
 
@@ -85,14 +90,12 @@ models:
     gpus_per_node: 1
     num_nodes: 1
     vocab_size: 152064
-    qos: m2
     time: 08:00:00
-    partition: a40
+    resource_type: l40s # You can also leave this field empty if your environment has a default type of resource to use
     model_weights_parent_dir: /h/<username>/model-weights
     vllm_args:
       --max-model-len: 1010000
       --max-num-seqs: 256
-      --compilation-config: 3
 ```
 
 You would then set the `VEC_INF_MODEL_CONFIG` path using:
@@ -103,9 +106,8 @@ export VEC_INF_MODEL_CONFIG=/h/<username>/my-model-config.yaml
 
 **NOTE**
 * There are other parameters that can also be added to the config but not shown in this example, check the [`ModelConfig`](vec_inf/client/config.py) for details.
-* Check [vLLM Engine Arguments](https://docs.vllm.ai/en/stable/serving/engine_args.html) for the full list of available vLLM engine arguments, the default parallel size for any parallelization is default to 1, so none of the sizes were set specifically in this example
+* Check [`vllm serve`](https://docs.vllm.ai/en/stable/cli/serve.html) for the full list of available vLLM engine arguments, the default parallel size for any parallelization is default to 1, so none of the sizes were set specifically in this example
 * For GPU partitions with non-Ampere architectures, e.g. `rtx6000`, `t4v2`, BF16 isn't supported. For models that have BF16 as the default type, when using a non-Ampere GPU, use FP16 instead, i.e. `--dtype: float16`
-* Setting `--compilation-config` to `3` currently breaks multi-node model launches, so we don't set them for models that require multiple nodes of GPUs.
 
 #### Other commands
 
@@ -114,7 +116,7 @@ export VEC_INF_MODEL_CONFIG=/h/<username>/my-model-config.yaml
 * `metrics`: Streams performance metrics to the console.
 * `shutdown`: Shutdown a model by providing its Slurm job ID.
 * `list`: List all available model names, or view the default/cached configuration of a specific model.
-* `cleanup`: Remove old log directories. You can filter by `--model-family`, `--model-name`, `--job-id`, and/or `--before-job-id`. Use `--dry-run` to preview what would be deleted.
+* `cleanup`: Remove old log directories, use `--help` to see the supported filters. Use `--dry-run` to preview what would be deleted.
 
 For more details on the usage of these commands, refer to the [User Guide](https://vectorinstitute.github.io/vector-inference/user_guide/)
 
@@ -125,6 +127,7 @@ Example:
 ```python
 >>> from vec_inf.api import VecInfClient
 >>> client = VecInfClient()
+>>> # Assume VEC_INF_ACCOUNT and VEC_INF_WORK_DIR is set
 >>> response = client.launch_model("Meta-Llama-3.1-8B-Instruct")
 >>> job_id = response.slurm_job_id
 >>> status = client.get_status(job_id)

docs/index.md

@@ -1,6 +1,8 @@
 # Vector Inference: Easy inference on Slurm clusters
 
-This repository provides an easy-to-use solution to run inference servers on [Slurm](https://slurm.schedmd.com/overview.html)-managed computing clusters using [vLLM](https://docs.vllm.ai/en/latest/). **All scripts in this repository runs natively on the Vector Institute cluster environment**. To adapt to other environments, follow the instructions in [Installation](#installation).
+This repository provides an easy-to-use solution to run inference servers on [Slurm](https://slurm.schedmd.com/overview.html)-managed computing clusters using [vLLM](https://docs.vllm.ai/en/stable/). **This package runs natively on the Vector Institute cluster environment**. To adapt to other environments, follow the instructions in [Installation](#installation).
+
+**NOTE**: Supported models on Killarney are tracked [here](vec_inf/config/README.md)
 
 ## Installation
 
@@ -10,7 +12,7 @@ If you are using the Vector cluster environment, and you don't need any customiz
 pip install vec-inf
 ```
 
-Otherwise, we recommend using the provided [`Dockerfile`](https://github.com/VectorInstitute/vector-inference/blob/main/Dockerfile) to set up your own environment with the package. The latest image has `vLLM` version `0.9.2`.
+Otherwise, we recommend using the provided [`Dockerfile`](https://github.com/VectorInstitute/vector-inference/blob/main/Dockerfile) to set up your own environment with the package. The latest image has `vLLM` version `0.10.1.1`.
 
 If you'd like to use `vec-inf` on your own Slurm cluster, you would need to update the configuration files, there are 3 ways to do it:
 * Clone the repository and update the `environment.yaml` and the `models.yaml` file in [`vec_inf/config`](vec_inf/config/), then install from source by running `pip install .`.