vllm-project / vllm-project/aibrix

[KVCache (Vineyard)] Do we need PodAffinity for model name in kvcache crd needs proper reconciliation?

Open
#876 2 comments 0 reactions 2 assignees Claimed by @Jeffwan View on GitHub
area/distributed area/kv-cache
Dominant language
Go
Stars
5.1k
Forks
694
Avg merge
1d 20h
Merged PRs (30d)
98

Description

### 🐛 Describe the bug

Originally, kvcache pod and vllm-engine pods were scheduled at the same node by the PodAffinity in kvcache crd.

![Image](https://github.com/user-attachments/assets/fc688c89-f4cc-493c-81d5-239d64a0b3fd)

However, after restarting llama-3-8b-instruct, kv cache pod and engine pod are scheduled in different pods. The new engine pod is scheduled in the new node but kv cache pod stays in the same pod.

![Image](https://github.com/user-attachments/assets/45912aa6-d2b8-43aa-b214-fb2dd41ab2a6)

Not sure which way is the best to make kv cache pod and engine pod colocated.

### Steps to Reproduce

```yaml
apiVersion: apps/v1
kind: Deployment
metadata:
name: llama-3-8b-instruct
labels:
model.aibrix.ai/name: llama-3-8b-instruct
model.aibrix.ai/port: "8000"
spec:
replicas: 1
strategy:
rollingUpdate:
maxSurge: 100%
maxUnavailable: 100%
type: RollingUpdate
selector:
matchLabels:
model.aibrix.ai/name: llama-3-8b-instruct
template:
metadata:
labels:
model.aibrix.ai/name: llama-3-8b-instruct
spec:
affinity:
nodeAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
nodeSelectorTerms:
- matchExpressions:
- key: machine.cluster.vke.volcengine.com/gpu-name
operator: In
values:
- NVIDIA-L20
containers:
- command:
- python3
- -m
- vllm.entrypoints.openai.api_server
- --port
- "8000"
- --uvicorn-log-level
- warning
- --model
- /models/llama-3.1-8b-instruct/
- --served-model-name
- llama-3-8b-instruct
- --trust-remote-code
- --enable-chunked-prefill
- "false"
- --max-model-len
- "100000"
- --dtype
- bfloat16
- --disable-log-requests
- --swap-space
- "0"
# - --enable-prefix-caching
env:
- name: VLLM_USE_VINEYARD_CACHE
value: "0"
- name: VINEYARD_CACHE_CPU_MEM_LIMIT_GB
value: "72"
- name: AIBRIX_LLM_KV_CACHE
value: "0"
- name: AIBRIX_LLM_KV_CACHE_KV_CACHE_NS
value: "aibrix"
- name: AIBRIX_LLM_KV_CACHE_CHUNK_SIZE
value: "16"
- name: AIBRIX_LLM_KV_CACHE_SOCKET
value: /var/run/vineyard.sock
- name: AIBRIX_LLM_KV_CACHE_RPC_ENDPOINT
value: "aibrix-kvcache-llama-3-8b-instruct-rpc:9600"
- name: VINEYARD_CACHE_ENABLE_ASYNC_UPDATE
value: "1"
- name: "VINEYARD_CACHE_METRICS_ENABLED"
value: "1"
- name: FLAGS_metrics
value: "1"
- name: GLOG_logtostderr
value: "1"
image: aibrix-container-registry-cn-beijing.cr.volces.com/aibrix/vllm-openai:v0.6.1-edb07092-20250118
imagePullPolicy: IfNotPresent
lifecycle:
preStop:
exec:
command:
- /bin/sh
- -c
- |
while true; do
RUNNING=$(curl -s http://localhost:8000/metrics | grep 'vllm:num_requests_running' | grep -v '#' | awk '{print $2}')
WAITING=$(curl -s http://localhost:8000/metrics | grep 'vllm:num_requests_waiting' | grep -v '#' | awk '{print $2}')
if [ "$RUNNING" = "0.0" ] && [ "$WAITING" = "0.0" ]; then
echo "Terminating: No active or waiting requests, safe to terminate" >> /proc/1/fd/1
exit 0
else
echo "Terminating: Running: $RUNNING, Waiting: $WAITING" >> /proc/1/fd/1
sleep 5
fi
done
livenessProbe:
failureThreshold: 3
httpGet:
path: /health
port: 8000
scheme: HTTP
initialDelaySeconds: 90
periodSeconds: 5
successThreshold: 1
timeoutSeconds: 1
name: vllm-openai
ports:
- containerPort: 8000
protocol: TCP
readinessProbe:
failureThreshold: 3
httpGet:
path: /health
port: 8000
scheme: HTTP
initialDelaySeconds: 90
periodSeconds: 5
successThreshold: 1
timeoutSeconds: 1
resources:
limits:
cpu: 1
nvidia.com/gpu: "1"
requests:
cpu: 1
nvidia.com/gpu: "1"
terminationMessagePath: /dev/termination-log
terminationMessagePolicy: File
volumeMounts:
- mountPath: /models
name: model-hostpath
- mountPath: /dev/shm
name: dshm
- mountPath: /var/run
name: kvcache-socket
- command:
- aibrix_runtime
- --port
- "8080"
env:
- name: INFERENCE_ENGINE
value: vllm
- name: INFERENCE_ENGINE_ENDPOINT
value: http://localhost:8000
image: aibrix-container-registry-cn-beijing.cr.volces.com/aibrix/runtime:v0.2.0
imagePullPolicy: IfNotPresent
livenessProbe:
failureThreshold: 3
httpGet:
path: /healthz
port: 8080
scheme: HTTP
initialDelaySeconds: 3
periodSeconds: 2
successThreshold: 1
timeoutSeconds: 1
name: aibrix-runtime
ports:
- containerPort: 8080
protocol: TCP
readinessProbe:
failureThreshold: 3
httpGet:
path: /ready
port: 8080
scheme: HTTP
initialDelaySeconds: 5
periodSeconds: 10
successThreshold: 1
timeoutSeconds: 1
resources: {}
terminationMessagePath: /dev/termination-log
terminationMessagePolicy: File
dnsPolicy: ClusterFirst
initContainers:
- command:
- aibrix_download
- --model-uri
- tos://aibrix-artifact-testing/models/llama-3.1-8b-instruct/
- --local-dir
- /models/
env:
- name: DOWNLOADER_MODEL_NAME
value: llama-3.1-8b-instruct
- name: DOWNLOADER_NUM_THREADS
value: "16"
- name: DOWNLOADER_ALLOW_FILE_SUFFIX
value: json, safetensors
- name: TOS_ACCESS_KEY
valueFrom:
secretKeyRef:
key: TOS_ACCESS_KEY
name: tos-credential
- name: TOS_SECRET_KEY
valueFrom:
secretKeyRef:
key: TOS_SECRET_KEY
name: tos-credential
- name: TOS_ENDPOINT
value: tos-cn-beijing.ivolces.com
- name: TOS_REGION
value: cn-beijing
image: aibrix-container-registry-cn-beijing.cr.volces.com/aibrix/runtime:v0.1.0
imagePullPolicy: IfNotPresent
name: init-model
resources: {}
volumeMounts:
- mountPath: /models
name: model-hostpath
terminationGracePeriodSeconds: 60
volumes:
- name: model-hostpath
hostPath:
path: /root/models
type: DirectoryOrCreate
- name: dshm
emptyDir:
medium: Memory
sizeLimit: "4Gi"
- name: kvcache-socket
hostPath:
path: /var/run/vineyard-kubernetes/default/aibrix-kvcache-llama-3-8b-instruct

---
apiVersion: v1
kind: Service
metadata:
labels:
model.aibrix.ai/name: llama-3-8b-instruct
prometheus-discovery: "true"
annotations:
prometheus.io/scrape: "true"
prometheus.io/port: "8080"
name: llama-3-8b-instruct # Note: The Service name must match the label value `model.aibrix.ai/name` in the Deployment
namespace: default
spec:
ports:
- name: serve
port: 8000
protocol: TCP
targetPort: 8000
- name: http
port: 8080
protocol: TCP
targetPort: 8080
selector:
model.aibrix.ai/name: llama-3-8b-instruct
type: ClusterIP

---

apiVersion: orchestration.aibrix.ai/v1alpha1
kind: KVCache
metadata:
name: aibrix-kvcache-llama-3-8b-instruct
namespace: default
annotations:
kvcache.orchestration.aibrix.ai/pod-affinity-workload: llama-3-8b-instruct
kvcache.orchestration.aibrix.ai/pod-anti-affinity: "true"
kvcache.orchestration.aibrix.ai/node-affinity-gpu-type: NVIDIA-L20
spec:
replicas: 1
service:
type: ClusterIP
port: 9600
cacheSpec:
image: aibrix-container-registry-cn-beijing.cr.volces.com/aibrix/vineyardd:20241120
imagePullPolicy: IfNotPresent
cpu: "4000m"
memory: 72Gi
```

### Expected behavior

kv cache pod and engine are in the same node even after restarting engine pod.

### Environment

main branch

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.