Skip to main content

Efficient LLM inference on Slurm clusters using vLLM.

Project description

Vector Inference: Easy inference on Slurm clusters


PyPI downloads code checks docs codecov GitHub License

This repository provides an easy-to-use solution to run inference servers on Slurm-managed computing clusters using vLLM. All scripts in this repository runs natively on the Vector Institute cluster environment. To adapt to other environments, update the environment variables in vec_inf/client/slurm_vars.py, and the model config for cached model weights in vec_inf/config/models.yaml accordingly.

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:

pip install vec-inf

Otherwise, we recommend using the provided Dockerfile to set up your own environment with the package

Usage

Vector Inference provides 2 user interfaces, a CLI and an API

CLI

The launch command allows users to deploy a model as a slurm job. If the job successfully launches, a URL endpoint is exposed for the user to send requests for inference.

We will use the Llama 3.1 model as example, to launch an OpenAI compatible inference server for Meta-Llama-3.1-8B-Instruct, run:

vec-inf launch Meta-Llama-3.1-8B-Instruct

You should see an output like the following:

launch_image

Overrides

Models that are already supported by vec-inf would be launched using the cached configuration (set in slurm_vars.py) or default configuration. 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:

vec-inf launch Meta-Llama-3.1-8B-Instruct --qos <new_qos>

To overwrite default vLLM engine arguments, you can specify the engine arguments in a comma separated string:

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, make sure you select the correct vLLM version.

Custom models

You can also launch your own custom model as long as the model architecture is supported by vLLM, and make sure to follow the instructions below:

  • Your model weights directory naming convention should follow $MODEL_FAMILY-$MODEL_VARIANT ($MODEL_VARIANT is OPTIONAL).
  • Your model weights directory should contain HuggingFace format weights.
  • You should specify your model configuration by:
    • Creating a custom configuration file for your model and specify its path via setting the environment variable VEC_INF_CONFIG. Check the default parameters file for the format of the config file. All the parameters for the model should be specified in that config file.
    • Using launch command options to specify your model setup.
  • For other model launch parameters you can reference the default values for similar models using the list command .

Here is an example to deploy a custom Qwen2.5-7B-Instruct-1M model which is not supported in the default list of models using a user custom config. In this case, the model weights are assumed to be downloaded to a model-weights directory inside the user's home directory. The weights directory of the model follows the naming convention so it would be named Qwen2.5-7B-Instruct-1M. The following yaml file would need to be created, lets say it is named /h/<username>/my-model-config.yaml.

models:
  Qwen2.5-7B-Instruct-1M:
    model_family: Qwen2.5
    model_variant: 7B-Instruct-1M
    model_type: LLM
    gpus_per_node: 1
    num_nodes: 1
    vocab_size: 152064
    qos: m2
    time: 08:00:00
    partition: a40
    model_weights_parent_dir: /h/<username>/model-weights
    vllm_args:
      --max-model-len: 1010000
      --max-num-seqs: 256
      --compilation-confi: 3

You would then set the VEC_INF_CONFIG path using:

export VEC_INF_CONFIG=/h/<username>/my-model-config.yaml

Note that there are other parameters that can also be added to the config but not shown in this example, check the ModelConfig for details.

Other commands

  • status: Check the model status by providing its Slurm job ID, --json-mode supported.
  • 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, --json-mode supported.

For more details on the usage of these commands, refer to the User Guide

API

Example:

>>> from vec_inf.api import VecInfClient
>>> client = VecInfClient()
>>> response = client.launch_model("Meta-Llama-3.1-8B-Instruct")
>>> job_id = response.slurm_job_id
>>> status = client.get_status(job_id)
>>> if status.status == ModelStatus.READY:
...     print(f"Model is ready at {status.base_url}")
>>> client.shutdown_model(job_id)

For details on the usage of the API, refer to the API Reference

Check Job Configuration

With every model launch, a Slurm script will be generated dynamically based on the job and model configuration. Once the Slurm job is queued, the generated Slurm script will be moved to the log directory for reproducibility, located at $log_dir/$model_family/$model_name.$slurm_job_id/$model_name.$slurm_job_id.slurm. In the same directory you can also find a JSON file with the same name that captures the launch configuration, and will have an entry of server URL once the server is ready.

Send inference requests

Once the inference server is ready, you can start sending in inference requests. We provide example scripts for sending inference requests in examples folder. Make sure to update the model server URL and the model weights location in the scripts. For example, you can run python examples/inference/llm/chat_completions.py, and you should expect to see an output like the following:

{
    "id":"chatcmpl-387c2579231948ffaf66cdda5439d3dc",
    "choices": [
        {
            "finish_reason":"stop",
            "index":0,
            "logprobs":null,
            "message": {
                "content":"Arrr, I be Captain Chatbeard, the scurviest chatbot on the seven seas! Ye be wantin' to know me identity, eh? Well, matey, I be a swashbucklin' AI, here to provide ye with answers and swappin' tales, savvy?",
                "role":"assistant",
                "function_call":null,
                "tool_calls":[],
                "reasoning_content":null
            },
            "stop_reason":null
        }
    ],
    "created":1742496683,
    "model":"Meta-Llama-3.1-8B-Instruct",
    "object":"chat.completion",
    "system_fingerprint":null,
    "usage": {
        "completion_tokens":66,
        "prompt_tokens":32,
        "total_tokens":98,
        "prompt_tokens_details":null
    },
    "prompt_logprobs":null
}

NOTE: For multimodal models, currently only ChatCompletion is available, and only one image can be provided for each prompt.

SSH tunnel from your local device

If you want to run inference from your local device, you can open a SSH tunnel to your cluster environment like the following:

ssh -L 8081:172.17.8.29:8081 username@v.vectorinstitute.ai -N

Where the last number in the URL is the GPU number (gpu029 in this case). The example provided above is for the vector cluster, change the variables accordingly for your environment

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

vec_inf-0.6.0.tar.gz (315.8 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

vec_inf-0.6.0-py3-none-any.whl (39.4 kB view details)

Uploaded Python 3

File details

Details for the file vec_inf-0.6.0.tar.gz.

File metadata

  • Download URL: vec_inf-0.6.0.tar.gz
  • Upload date:
  • Size: 315.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.6.6

File hashes

Hashes for vec_inf-0.6.0.tar.gz
Algorithm Hash digest
SHA256 c02360de2fde5da9e93bcf702e97fd2ca0a49056831793285fb3788987cf7217
MD5 e2822aba100a0473cacaac3576f1e82c
BLAKE2b-256 30e4fb6a256e9be837af527d78a578c53de055aac4bd0362557f3d98695be3b1

See more details on using hashes here.

File details

Details for the file vec_inf-0.6.0-py3-none-any.whl.

File metadata

  • Download URL: vec_inf-0.6.0-py3-none-any.whl
  • Upload date:
  • Size: 39.4 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.6.6

File hashes

Hashes for vec_inf-0.6.0-py3-none-any.whl
Algorithm Hash digest
SHA256 cc0211957f1a4b12fba600ff456312ccb632cd7b3e423937c76ab89001a8ffc9
MD5 4473d7c3bc8795354affa17aabfede01
BLAKE2b-256 ef1d52cef4e49bfc53b4e091bc94d80183dfbfc84aac18cedaae0eb36400e208

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page