多节点推理

分布式推理是指将LLM模型分割并跨多个节点或设备进行处理的技术。这种方法对于无法容纳在单台机器内存中的大型模型特别有用。此解决方案依赖KubeRay来编排Ray集群。

关键API设计

在分布式计算领域，高效编排多节点推理任务的需求已变得至关重要。Kubernetes已成为管理容器化应用程序的领先平台，提供强大的资源管理和可伸缩性。另一方面，Ray已成为构建和运行分布式应用程序的强大框架，特别适用于处理复杂的机器学习工作流。然而，现有的编排方法在灵活性和简单性方面往往不足。Kubernetes操作符虽然强大，但在处理分布式应用程序的细粒度编排时可能会变得过于复杂。Ray虽然在内部任务调度和资源管理方面表现出色，但缺乏Kubernetes提供的更广泛的资源编排能力。

为了解决这些挑战，我们提出了一种新的编排方法，该方法协同利用Kubernetes和Ray的优势。此方法利用Ray进行内部细粒度应用程序编排，允许用户利用Ray的API进行分布式计算。同时，Kubernetes将处理整体应用程序资源编排，专注于粗粒度资源分配和环境配置。这种职责划分简化了Kubernetes操作符的设计，并增强了编排过程的整体灵活性和效率。

我们引入了两个用于RayCluster管理的关键API，分别是RayClusterReplicaSet和RayClusterFleet。这类似于Kubernetes的核心概念ReplicaSet和Deployment。大多数情况下，您只需要使用RayClusterFleet。

• Ray框架焦点：在此模型中，Ray仅被强调用于其在应用内部编排中的作用。每个应用实例对应一个Ray集群，一个应用的多个服务实例等同于多个Ray集群。这确保Ray在内部处理应用的分布式特性，而不受外部编排系统的干扰。

• Kubernetes层：Kubernetes在外层操作，负责启动Ray集群并管理标准的Kubernetes功能，如自动扩缩和滚动更新。Kubernetes层不再编排应用程序内部的角色。这些功能在Kubernetes生态系统中已成熟，确保了健壮可靠的资源管理、扩展和更新过程。通过利用Kubernetes进行这些操作，我们可以实现Ray的分布式计算能力与Kubernetes成熟的运维管理的无缝集成。

• 服务封装和映射：在更高层次，服务以类似于Kubernetes Deployment和ReplicaSet的方式进行封装。关键区别在于映射：我们现在拥有代表应用程序实例的Ray集群，而不是Pod。传统上，单个Pod构成一个应用程序实例；然而，在这个分布式模型中，Ray集群充当此目的，将分布式执行的复杂性封装在自身内部。

注意

我们已经将我们的想法提交给KubeRay社区。希望我们很快就能合并到仓库中。

工作负载示例

注意

从v0.6.6版本开始，我们添加了必要的软件包，以开箱即用地运行vLLM官方容器镜像分发的分布式推理。如果您使用早期版本，您可以按照以下指南构建自己的兼容多节点推理的镜像。

这是一个RayClusterFleet示例，您可以将此yaml应用到您的集群中。

apiVersion: orchestration.aibrix.ai/v1alpha1
kind: RayClusterFleet
metadata:
  name: qwen-coder-7b-instruct
  labels:
    app.kubernetes.io/name: aibrix
    app.kubernetes.io/managed-by: kustomize
spec:
  replicas: 1
  selector:
    matchLabels:
      model.aibrix.ai/name: qwen-coder-7b-instruct
  strategy:
    type: RollingUpdate
    rollingUpdate:
      maxSurge: 25%
      maxUnavailable: 25%
  template:
    metadata:
      labels:
        model.aibrix.ai/name: qwen-coder-7b-instruct
      annotations:
        ray.io/overwrite-container-cmd: "true"
    spec:
      rayVersion: "2.10.0"
      headGroupSpec:
        rayStartParams:
          dashboard-host: "0.0.0.0"
        template:
          metadata:
            labels:
              model.aibrix.ai/name: qwen-coder-7b-instruct
          spec:
            containers:
              - name: ray-head
                image: vllm/vllm-openai:v0.7.1
                command: ["/bin/bash", "-c"]
                args:
                  - >
                    ulimit -n 65536 &&
                    apt update && apt install -y wget net-tools && pip3 install ray[default] pyarrow pandas &&
                    echo "[INFO] Starting Ray head node..." &&
                    eval "$KUBERAY_GEN_RAY_START_CMD" &

                    echo "[INFO] Waiting for Ray dashboard to be ready..." &&
                    until curl --max-time 5 --fail http://127.0.0.1:8265 > /dev/null 2>&1; do
                      echo "[WAITING] $(date -u +'%Y-%m-%dT%H:%M:%SZ') - Ray dashboard not ready yet...";
                      sleep 2;
                    done &&
                    echo "[SUCCESS] Ray dashboard is available!" &&

                    vllm serve Qwen/Qwen2.5-Coder-7B-Instruct \
                      --served-model-name qwen-coder-7b-instruct \
                      --tensor-parallel-size 2 \
                      --distributed-executor-backend ray \
                      --host 0.0.0.0 \
                      --port 8000 \
                      --dtype half
                ports:
                  - containerPort: 6379
                    name: gcs-server
                  - containerPort: 8265
                    name: dashboard
                  - containerPort: 10001
                    name: client
                  - containerPort: 8000
                    name: service
                resources:
                  limits:
                    cpu: "4"
                    nvidia.com/gpu: 1
                  requests:
                    cpu: "4"
                    nvidia.com/gpu: 1
              - name: aibrix-runtime
                image: aibrix/runtime:v0.3.0
                command:
                  - aibrix_runtime
                  - --port
                  - "8080"
                env:
                  - name: INFERENCE_ENGINE
                    value: vllm
                  - name: INFERENCE_ENGINE_ENDPOINT
                    value: https://:8000
                  - name: PYTORCH_CUDA_ALLOC_CONF
                    value: "expandable_segments:True"
                ports:
                  - containerPort: 8080
                    protocol: TCP
                livenessProbe:
                  httpGet:
                    path: /healthz
                    port: 8080
                  initialDelaySeconds: 3
                  periodSeconds: 2
                readinessProbe:
                  httpGet:
                    path: /ready
                    port: 8080
                  initialDelaySeconds: 5
                  periodSeconds: 10
                resources:
                  limits:
                    cpu: "1"
                  requests:
                    cpu: "1"
      workerGroupSpecs:
        - groupName: small-group
          replicas: 1
          minReplicas: 1
          maxReplicas: 5
          rayStartParams: {}
          template:
            metadata:
              labels:
                model.aibrix.ai/name: qwen-coder-7b-instruct
            spec:
              containers:
                - name: ray-worker
                  image: vllm/vllm-openai:v0.7.1
                  env:
                    - name: MY_POD_IP
                      valueFrom:
                        fieldRef:
                          fieldPath: status.podIP
                  command: [ "/bin/bash", "-c" ]
                  args:
                    - >
                      ulimit -n 65536 &&
                      eval "$KUBERAY_GEN_RAY_START_CMD --node-ip-address=$MY_POD_IP" &&
                      tail -f /dev/null
                  lifecycle:
                    preStop:
                      exec:
                        command: [ "/bin/sh", "-c", "ray stop" ]
                  resources:
                    limits:
                      cpu: "4"
                      nvidia.com/gpu: 1
                    requests:
                      cpu: "4"
                      nvidia.com/gpu: 1

---

apiVersion: v1
kind: Service
metadata:
  name: qwen-coder-7b-instruct
  labels:
    model.aibrix.ai/name: qwen-coder-7b-instruct
    prometheus-discovery: "true"
  annotations:
    prometheus.io/scrape: "true"
    prometheus.io/port: "8080"
spec:
  selector:
    model.aibrix.ai/name: qwen-coder-7b-instruct
  ports:
    - name: serve
      port: 8000
      protocol: TCP
      targetPort: 8000
    - name: http
      port: 8080
      protocol: TCP
      targetPort: 8080

---

apiVersion: gateway.networking.k8s.io/v1
kind: HTTPRoute
metadata:
  name: qwen-coder-7b-instruct-router
  namespace: aibrix-system
spec:
  parentRefs:
    - group: gateway.networking.k8s.io
      kind: Gateway
      name: aibrix-eg
      namespace: aibrix-system
  rules:
    - backendRefs:
        - group: ""
          kind: Service
          name: qwen-coder-7b-instruct
          namespace: default
          port: 8000  # or 8000 if you're not using the runtime sidecar
          weight: 1
      matches:
        - headers:
            - name: model
              type: Exact
              value: qwen-coder-7b-instruct
          path:
            type: PathPrefix
            value: /v1/completions
        - headers:
            - name: model
              type: Exact
              value: qwen-coder-7b-instruct
          path:
            type: PathPrefix
            value: /v1/chat/completions
      timeouts:
        request: 120s

vLLM版本

如果您使用的是vLLM早期版本，您有两个选择。

• 使用我们构建的镜像aibrix/vllm-openai:v0.6.1.post2-distributed。

• 构建您自己的镜像并按照这里的步骤操作。

FROM vllm/vllm-openai:v0.6.1.post2
RUN apt update && apt install -y wget # important for future healthcheck
RUN pip3 install ray[default] # important for future healthcheck
ENTRYPOINT [""]

docker build -t aibrix/vllm-openai:v0.6.1.post2-distributed .