Merge remote-tracking branch 'origin/main' into dev/vllm-align
This commit is contained in:
tzhouam committed
commit
e22ae62396
33 files changed
+4275
-519
No files matched your search
+17
-299
@@ -9,314 +9,32 @@ steps:
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
- label: "Upload Nightly Pipeline"
|
||||
# L2 Test
|
||||
- label: "Upload Ready Pipeline"
|
||||
depends_on: image-build
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
key: upload-ready-pipeline
|
||||
if: build.branch != "main"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yaml
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
- label: "Simple Unit Test"
|
||||
# L3 Test
|
||||
- label: "Upload Merge Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
if: build.branch == "main"
|
||||
commands:
|
||||
- "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
timeout_in_minutes: 20
|
||||
# L4 Test
|
||||
- label: "Upload Nightly Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Audio Generation Model Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
timeout_in_minutes: 15
|
||||
depends_on: image-build
|
||||
commands:
|
||||
- pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
commands:
|
||||
- pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
|
||||
# - label: "Benchmark&Engine Test"
|
||||
# timeout_in_minutes: 15
|
||||
# depends_on: image-build
|
||||
# commands:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -s -v tests/benchmarks/test_serve_cli.py
|
||||
# - pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 2
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni Model Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
# - label: "Omni Model Test with H100"
|
||||
# timeout_in_minutes: 20
|
||||
# depends_on: image-build
|
||||
# commands:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
# - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
# - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 2
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
- label: "Qwen3-TTS E2E Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
# - label: "Diffusion Image Edit Test with H100 (1 GPU)"
|
||||
# timeout_in_minutes: 20
|
||||
# depends_on: image-build
|
||||
# commands:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
# - label: "Bagel Text2Img Model Test with H100"
|
||||
# timeout_in_minutes: 30
|
||||
# depends_on: image-build
|
||||
# commands:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
queue: "cpu_queue_premerge"
|
||||
+58
-69
@@ -1,16 +1,6 @@
|
||||
steps:
|
||||
- label: ":docker: Build image"
|
||||
key: image-build
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
|
||||
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
- label: "Simple Unit Test"
|
||||
depends_on: image-build
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
agents:
|
||||
@@ -27,7 +17,7 @@ steps:
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
|
||||
agents:
|
||||
@@ -44,7 +34,7 @@ steps:
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
|
||||
agents:
|
||||
@@ -61,7 +51,7 @@ steps:
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
@@ -79,7 +69,7 @@ steps:
|
||||
|
||||
- label: "Audio Generation Model Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
|
||||
agents:
|
||||
@@ -96,9 +86,9 @@ steps:
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
timeout_in_minutes: 15
|
||||
depends_on: image-build
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
|
||||
- pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
@@ -113,7 +103,7 @@ steps:
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
|
||||
agents:
|
||||
@@ -131,9 +121,9 @@ steps:
|
||||
|
||||
- label: "Diffusion Tensor Parallelism Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
|
||||
- pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
@@ -149,7 +139,7 @@ steps:
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
agents:
|
||||
@@ -165,10 +155,9 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
|
||||
- label: "Benchmark&Engine Test"
|
||||
- label: "Benchmark & Engine Test with H100"
|
||||
timeout_in_minutes: 15
|
||||
depends_on: image-build
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/benchmarks/test_serve_cli.py
|
||||
@@ -204,63 +193,63 @@ steps:
|
||||
|
||||
- label: "Omni Model Test"
|
||||
timeout_in_minutes: 15
|
||||
depends_on: image-build
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni Model Test with H100"
|
||||
timeout_in_minutes: 30
|
||||
depends_on: image-build
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion Image Edit Test with H100 (1 GPU)"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: image-build
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
|
||||
@@ -295,7 +284,7 @@ steps:
|
||||
|
||||
# - label: "Bagel Text2Img Model Test with H100"
|
||||
# timeout_in_minutes: 30
|
||||
# depends_on: image-build
|
||||
# depends_on: upload-merge-pipeline
|
||||
# commands:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
steps:
|
||||
- label: "Omni Model Test with H100"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: image-build
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
@@ -38,7 +38,7 @@ steps:
|
||||
|
||||
- label: "Omni Model Test"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: image-build
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
@@ -56,10 +56,9 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
|
||||
- label: "Omni Model Perf Test"
|
||||
timeout_in_minutes: 120
|
||||
depends_on: image-build
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
@@ -0,0 +1,303 @@
|
||||
steps:
|
||||
- label: "Simple Unit Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Audio Generation Model Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
timeout_in_minutes: 15
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
# - label: "Benchmark & Engine Test with H100"
|
||||
# timeout_in_minutes: 15
|
||||
# depends_on: upload-ready-pipeline
|
||||
# commands:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -s -v tests/benchmarks/test_serve_cli.py
|
||||
# - pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 2
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni Model Test"
|
||||
timeout_in_minutes: 15
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
# - label: "Omni Model Test with H100"
|
||||
# timeout_in_minutes: 20
|
||||
# depends_on: upload-ready-pipeline
|
||||
# commands:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
# # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
# - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 2
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
- label: "Qwen3-TTS E2E Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
# - label: "Diffusion Image Edit Test with H100 (1 GPU)"
|
||||
# timeout_in_minutes: 20
|
||||
# depends_on: upload-ready-pipeline
|
||||
# commands:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
# - label: "Bagel Text2Img Model Test with H100"
|
||||
# timeout_in_minutes: 30
|
||||
# depends_on: upload-ready-pipeline
|
||||
# commands:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
@@ -14,7 +14,7 @@
|
||||
soft_fail: false
|
||||
commands:
|
||||
- "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
|
||||
- "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target final --progress plain ."
|
||||
- "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
|
||||
- "docker push {{ docker_image_amd }}"
|
||||
key: "amd-build"
|
||||
env:
|
||||
|
||||
@@ -0,0 +1,397 @@
|
||||
# RDMA Test Configuration Guide
|
||||
|
||||
This document explains how to configure the RDMA environment and run tests for `MooncakeTransferEngineConnector`.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Docker Container Permissions](#docker-container-permissions)
|
||||
- [Single-Node Testing](#single-node-testing)
|
||||
- [Multi-Node Testing](#multi-node-testing)
|
||||
- [Running Tests](#running-tests)
|
||||
- [Cross-Node Testing](#cross-node-testing)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
|
||||
---
|
||||
|
||||
## Docker Container Permissions
|
||||
|
||||
RDMA tests require access to InfiniBand/RoCE devices and system topology. Add the following permissions when running `docker run`.
|
||||
|
||||
### Option 1: Minimal Permissions (Recommended)
|
||||
|
||||
```bash
|
||||
docker run -it \
|
||||
--cap-add=SYS_PTRACE \
|
||||
--cap-add=IPC_LOCK \
|
||||
--security-opt seccomp=unconfined \
|
||||
--network=host \
|
||||
--device=/dev/infiniband \
|
||||
-v /sys/class/infiniband:/sys/class/infiniband:ro \
|
||||
your-image:tag
|
||||
```
|
||||
|
||||
Parameter explanation:
|
||||
- `--cap-add=SYS_PTRACE`: Allow reading system topology information
|
||||
- `--cap-add=IPC_LOCK`: Allow memory locking (required for RDMA memory registration)
|
||||
- `--security-opt seccomp=unconfined`: Disable seccomp restrictions
|
||||
- `--network=host`: Use host network (required for RDMA)
|
||||
- `--device=/dev/infiniband`: Mount InfiniBand devices
|
||||
- `-v /sys/class/infiniband`: Mount IB device info (read-only)
|
||||
|
||||
### Option 2: Full Permissions (Quick but not recommended for production)
|
||||
|
||||
```bash
|
||||
docker run -it \
|
||||
--privileged \
|
||||
--network=host \
|
||||
your-image:tag
|
||||
```
|
||||
|
||||
`--privileged` grants full host permissions. Suitable for quick testing but not recommended for production.
|
||||
|
||||
---
|
||||
|
||||
## Single-Node Testing
|
||||
|
||||
When running single-node tests (producer and consumer on the same machine), ensure they use the **same RDMA device**.
|
||||
|
||||
### Problem Background
|
||||
|
||||
InfiniBand devices use LID (Local Identifier) for routing. Different devices have different LIDs and cannot communicate directly. If no device is specified, Mooncake may assign different devices to connectors, causing handshake failures.
|
||||
|
||||
Common error:
|
||||
```
|
||||
[Handshake] Failed to modify QP to RTR, check mtu, gid, peer lid, peer qp num: Invalid argument [22]
|
||||
```
|
||||
|
||||
### Solution
|
||||
|
||||
**Method 1: Set Environment Variable (Recommended)**
|
||||
|
||||
```bash
|
||||
# List available RDMA devices
|
||||
ibstat
|
||||
|
||||
# Select a device (e.g., mlx5_0)
|
||||
export RDMA_DEVICE_NAME='mlx5_0'
|
||||
|
||||
# Run tests
|
||||
pytest test_mooncake_transfer_engine_rdma.py -v -s
|
||||
```
|
||||
|
||||
**Method 2: Use RoCE Devices**
|
||||
|
||||
If the system has RoCE devices (using IPv4 routing), the test code will automatically detect and prefer them. RoCE device GIDs start with `00:00:00:00:00:00:00:00:00:00:ff:ff` (IPv4-mapped).
|
||||
|
||||
**Method 3: Ensure MTU Consistency**
|
||||
|
||||
Make sure both endpoints use the same MTU:
|
||||
|
||||
```bash
|
||||
# Check device MTU
|
||||
ibstatus mlx5_0
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Multi-Node Testing
|
||||
|
||||
For multi-node tests, producer and consumer run on different machines connected via InfiniBand switch.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
1. Both machines have Mooncake and RDMA drivers installed
|
||||
2. Both machines are in the same InfiniBand subnet
|
||||
3. Switch is properly configured
|
||||
|
||||
### Configuration
|
||||
|
||||
**Machine A (Producer):**
|
||||
|
||||
```bash
|
||||
# Set RDMA host IP (InfiniBand interface IP)
|
||||
export RDMA_TEST_HOST='10.0.0.1'
|
||||
|
||||
# Optional: Specify device
|
||||
export RDMA_DEVICE_NAME='mlx5_0'
|
||||
```
|
||||
|
||||
**Machine B (Consumer):**
|
||||
|
||||
```bash
|
||||
# Set RDMA host IP
|
||||
export RDMA_TEST_HOST='10.0.0.2'
|
||||
|
||||
# Optional: Specify device
|
||||
export RDMA_DEVICE_NAME='mlx5_0'
|
||||
```
|
||||
|
||||
### Verify Connectivity
|
||||
|
||||
```bash
|
||||
# Ping IB interface
|
||||
ping 10.0.0.2
|
||||
|
||||
# Test RDMA connectivity with ibping
|
||||
# On Machine B (server)
|
||||
ibping -S
|
||||
|
||||
# On Machine A (client)
|
||||
ibping -G <Machine_B_GID>
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Running Tests
|
||||
|
||||
### Run All RDMA Tests (Single-Node, fast suite)
|
||||
|
||||
Slow tests (large payloads, stress, concurrency integrity) are marked `@pytest.mark.slow`. Use `-m "not slow"` to skip them in quick CI or local fast iteration.
|
||||
|
||||
```bash
|
||||
cd tests/distributed/omni_connectors
|
||||
|
||||
# Fast suite only (excludes slow/stress tests)
|
||||
pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s -m "not slow"
|
||||
```
|
||||
|
||||
### Run Including Slow Tests
|
||||
|
||||
```bash
|
||||
# Run ALL tests including slow/stress tests
|
||||
pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s
|
||||
|
||||
# Run ONLY the slow/stress tests
|
||||
pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s -m slow
|
||||
```
|
||||
|
||||
### Run Buffer Management Tests
|
||||
|
||||
```bash
|
||||
# Fast only
|
||||
pytest test_mooncake_transfer_engine_buffer.py -v -s -m "not slow"
|
||||
|
||||
# Including allocator invariant tests (double-free, overlap, merge)
|
||||
pytest test_mooncake_transfer_engine_buffer.py -v -s
|
||||
```
|
||||
|
||||
### Run Specific Test Classes
|
||||
|
||||
```bash
|
||||
# Basic connector tests
|
||||
pytest test_mooncake_transfer_engine_rdma.py::TestBasicConnector -v -s
|
||||
|
||||
# End-to-end RDMA transfer tests
|
||||
pytest test_mooncake_transfer_engine_rdma.py::TestEndToEnd -v -s
|
||||
|
||||
# Lifecycle & resource management tests
|
||||
pytest test_mooncake_transfer_engine_rdma.py::TestLifecycle -v -s
|
||||
|
||||
# GPU memory pool tests (requires CUDA)
|
||||
pytest test_mooncake_transfer_engine_rdma.py::TestGPUPool -v -s
|
||||
|
||||
# Stress / correctness tests (slow)
|
||||
pytest test_mooncake_transfer_engine_rdma.py::TestStressCorrectness -v -s
|
||||
```
|
||||
|
||||
### RDMA Environment Diagnostics
|
||||
|
||||
For quick diagnostics (device status, Mooncake availability, env vars, etc.),
|
||||
see the [Troubleshooting section](../../../docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md#troubleshooting)
|
||||
in the connector documentation.
|
||||
|
||||
---
|
||||
|
||||
## Cross-Node Testing
|
||||
|
||||
The `cross_node_mooncake_transfer_engine.py` script enables testing RDMA transfers between two separate physical machines. This script is **not** auto-discovered by `pytest` (it does not start with `test_`) — it must be run manually on each node.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
1. Both machines have Mooncake installed
|
||||
2. Both machines are connected via InfiniBand/RoCE switch
|
||||
3. Firewall allows ZMQ ports (default: 15500, 15501)
|
||||
4. Same RDMA device name on both nodes (if multiple devices exist)
|
||||
|
||||
### Running Cross-Node Tests
|
||||
|
||||
**On Machine A (Producer) — start first:**
|
||||
|
||||
```bash
|
||||
cd benchmarks/distributed/omni_connectors/
|
||||
|
||||
# Optional: specify device if multiple exist
|
||||
export RDMA_DEVICE_NAME='mlx5_0'
|
||||
|
||||
python cross_node_mooncake_transfer_engine.py \
|
||||
--role producer \
|
||||
--local-host <PRODUCER_IP> \
|
||||
--remote-host <CONSUMER_IP> \
|
||||
--tensor-size-mb 100 \
|
||||
--num-transfers 3
|
||||
```
|
||||
|
||||
**On Machine B (Consumer) — start after producer:**
|
||||
|
||||
```bash
|
||||
cd benchmarks/distributed/omni_connectors/
|
||||
|
||||
export RDMA_DEVICE_NAME='mlx5_0'
|
||||
|
||||
python cross_node_mooncake_transfer_engine.py \
|
||||
--role consumer \
|
||||
--local-host <CONSUMER_IP> \
|
||||
--remote-host <PRODUCER_IP> \
|
||||
--tensor-size-mb 100 \
|
||||
--num-transfers 3
|
||||
```
|
||||
|
||||
### Transfer Modes
|
||||
|
||||
| Mode | Description | Example |
|
||||
|------|-------------|---------|
|
||||
| `copy` | Normal path — tensor copied to RDMA pool (default) | `--mode copy` |
|
||||
| `zerocopy` | Zero-copy path — data created directly in RDMA pool | `--mode zerocopy` |
|
||||
| `gpu` | GPU transfer — RDMA pool on GPU, uses GPUDirect | `--mode gpu --gpu-id 0` |
|
||||
|
||||
### Benchmark Mode
|
||||
|
||||
Skip MD5 verification and measure pure RDMA throughput:
|
||||
|
||||
```bash
|
||||
# Producer
|
||||
python cross_node_mooncake_transfer_engine.py \
|
||||
--role producer \
|
||||
--local-host <PRODUCER_IP> \
|
||||
--remote-host <CONSUMER_IP> \
|
||||
--tensor-size-mb 1024 \
|
||||
--num-transfers 20 \
|
||||
--benchmark
|
||||
|
||||
# Consumer
|
||||
python cross_node_mooncake_transfer_engine.py \
|
||||
--role consumer \
|
||||
--local-host <CONSUMER_IP> \
|
||||
--remote-host <PRODUCER_IP> \
|
||||
--tensor-size-mb 1024 \
|
||||
--num-transfers 20 \
|
||||
--benchmark
|
||||
```
|
||||
|
||||
### Cross-Node Test Options
|
||||
|
||||
| Option | Description | Default |
|
||||
|--------|-------------|---------|
|
||||
| `--role` | `producer` or `consumer` | Required |
|
||||
| `--local-host` | Local RDMA IP address | Required |
|
||||
| `--remote-host` | Remote RDMA IP address | Required |
|
||||
| `--local-port` | Local ZMQ port for RDMA data | 15500 |
|
||||
| `--remote-port` | Remote ZMQ port for RDMA data | 15500 |
|
||||
| `--ctrl-port` | Control channel port | 15501 |
|
||||
| `--tensor-size-mb` | Tensor size in MB | 100 |
|
||||
| `--num-transfers` | Number of transfers | 3 |
|
||||
| `--mode` | `copy`, `zerocopy`, or `gpu` | `copy` |
|
||||
| `--gpu-id` | GPU ID for GPU mode | 0 |
|
||||
| `--benchmark` | Skip MD5, pure performance test | off |
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### 1. "Failed to modify QP to RTR" Error
|
||||
|
||||
**Cause**: QP handshake failed, usually due to device configuration mismatch.
|
||||
|
||||
**Solution**:
|
||||
```bash
|
||||
# Force using the same device
|
||||
export RDMA_DEVICE_NAME='mlx5_0'
|
||||
```
|
||||
|
||||
### 2. "Mooncake TransferEngine is not available"
|
||||
|
||||
**Cause**: Mooncake not installed or import failed.
|
||||
|
||||
**Solution**:
|
||||
```bash
|
||||
# Check Mooncake installation
|
||||
python -c "from mooncake.engine import TransferEngine; print('OK')"
|
||||
|
||||
# Reinstall if needed
|
||||
pip install mooncake-transfer-engine
|
||||
# Or using uv
|
||||
uv pip install mooncake-transfer-engine
|
||||
|
||||
```
|
||||
|
||||
### 3. "Permission denied" accessing /dev/infiniband
|
||||
|
||||
**Cause**: Container lacks IB device access permissions.
|
||||
|
||||
**Solution**:
|
||||
```bash
|
||||
docker run --device=/dev/infiniband --cap-add=IPC_LOCK ...
|
||||
```
|
||||
|
||||
### 4. Test Timeout
|
||||
|
||||
**Cause**: RDMA connection establishment failed or network latency.
|
||||
|
||||
**Solution**:
|
||||
```bash
|
||||
# Check network status
|
||||
ibstat
|
||||
ibstatus
|
||||
```
|
||||
|
||||
### 5. GPU Test Failed "CUDA is not available"
|
||||
|
||||
**Cause**: CUDA environment not configured or GPU unavailable.
|
||||
|
||||
**Solution**:
|
||||
```bash
|
||||
# Check CUDA
|
||||
python -c "import torch; print(torch.cuda.is_available())"
|
||||
|
||||
# Docker needs NVIDIA runtime
|
||||
docker run --gpus all ...
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Environment Variables Reference
|
||||
|
||||
| Variable | Description | Example |
|
||||
|----------|-------------|---------|
|
||||
| `RDMA_DEVICE_NAME` | Specify RDMA device name | `mlx5_0` |
|
||||
| `RDMA_TEST_HOST` | Specify test host IP | `10.0.0.1` |
|
||||
| `MC_TE_METRIC` | Enable Mooncake metrics | `1` |
|
||||
| `MC_IB_PCI_RELAXED_ORDERING` | Enable PCIe relaxed ordering | `1` |
|
||||
|
||||
---
|
||||
|
||||
## Test Files Overview
|
||||
|
||||
| File | Description | Auto-discovered by pytest |
|
||||
|------|-------------|--------------------------|
|
||||
| `test_mooncake_transfer_engine_rdma.py` | Integration tests for MooncakeTransferEngineConnector (basic, E2E, lifecycle, GPU) | Yes |
|
||||
| `test_mooncake_transfer_engine_buffer.py` | Memory pool and buffer management unit tests | Yes |
|
||||
| `cross_node_mooncake_transfer_engine.py` | Cross-node (multi-machine) testing script — run manually | No (filename does not start with `test_`) |
|
||||
|
||||
### test_mooncake_transfer_engine_rdma.py — Test Classes
|
||||
|
||||
| Test Class | Memory Pool | Marker | Description |
|
||||
|------------|-------------|--------|-------------|
|
||||
| `TestBasicConnector` | CPU | — | Initialization, put tensor/bytes/object, cleanup, pool exhaustion |
|
||||
| `TestEndToEnd` | CPU | — | E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload (100MB), mixed types, concurrency |
|
||||
| `TestLifecycle` | CPU | — | Close, context manager, double-close safety |
|
||||
| `TestGPUPool` | GPU | — | GPU pool init, put CPU/GPU tensor, GPU E2E transfer |
|
||||
| `TestStressCorrectness` | CPU | `slow` | Concurrent put+get with MD5 integrity, bidirectional concurrency, edge cases (1-element tensor, empty bytes), 500MB payload, rapid alloc/free cycles |
|
||||
|
||||
### test_mooncake_transfer_engine_buffer.py — Test Classes
|
||||
|
||||
| Test Class | Marker | Description |
|
||||
|------------|--------|-------------|
|
||||
| `TestBufferAllocator` | — | Basic alloc/free, alignment, exhaustion/recovery, thread safety |
|
||||
| `TestAllocatorInvariants` | `slow` | Double-free safety, overlap corruption detection, adjacent-block merging, fragmentation/defrag |
|
||||
| `TestManagedBuffer` | — | Tensor views, context manager |
|
||||
@@ -0,0 +1,644 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Cross-Node RDMA Test Script (Automated Version)
|
||||
|
||||
This script enables testing RDMA transfers between two separate machines.
|
||||
Supports three transfer modes:
|
||||
- copy: Normal path - tensor copied to RDMA pool (default)
|
||||
- zerocopy: Zero-copy path - data created directly in RDMA pool
|
||||
- gpu: GPU transfer - RDMA pool on GPU, uses GPUDirect
|
||||
|
||||
Usage:
|
||||
# On Machine A (Producer) - start first:
|
||||
python cross_node_mooncake_transfer_engine.py --role producer --local-host hostname_A --remote-host hostname_B
|
||||
|
||||
# On Machine B (Consumer) - start after producer:
|
||||
python cross_node_mooncake_transfer_engine.py --role consumer --local-host hostname_B --remote-host hostname_A
|
||||
|
||||
# Zero-copy mode:
|
||||
python cross_node_mooncake_transfer_engine.py --role producer ... --mode zerocopy
|
||||
|
||||
# GPU mode (requires GPUDirect RDMA support):
|
||||
python cross_node_mooncake_transfer_engine.py --role producer ... --mode gpu --gpu-id 0
|
||||
|
||||
# Benchmark mode (skip random data generation and MD5 verification,
|
||||
# measures pure RDMA throughput):
|
||||
python cross_node_mooncake_transfer_engine.py --role producer ... --benchmark
|
||||
|
||||
Environment Variables:
|
||||
RDMA_DEVICE_NAME: Specify RDMA device (e.g., mlx5_0)
|
||||
MC_IB_PCI_RELAXED_ORDERING: Set to 1 to enable PCIe relaxed ordering
|
||||
for higher RDMA throughput
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import msgspec
|
||||
import torch
|
||||
import zmq
|
||||
|
||||
# Add parent path for imports
|
||||
sys.path.insert(
|
||||
0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
|
||||
)
|
||||
|
||||
from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
|
||||
ManagedBuffer,
|
||||
MooncakeTransferEngineConnector,
|
||||
TransferEngine,
|
||||
)
|
||||
|
||||
|
||||
def compute_md5(tensor: torch.Tensor) -> str:
|
||||
"""Compute MD5 checksum of a tensor."""
|
||||
if tensor.is_cuda:
|
||||
tensor = tensor.cpu()
|
||||
data = tensor.contiguous().view(torch.uint8).numpy().tobytes()
|
||||
return hashlib.md5(data).hexdigest()
|
||||
|
||||
|
||||
# Control channel message types
|
||||
class CtrlMsg(msgspec.Struct):
|
||||
"""Control channel message."""
|
||||
|
||||
msg_type: str # "READY", "TRANSFER", "ACK", "DONE", "ERROR"
|
||||
request_id: str = ""
|
||||
md5: str = ""
|
||||
data_size: int = 0
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransferConfig:
|
||||
"""Configuration for cross-node transfer test."""
|
||||
|
||||
local_host: str
|
||||
remote_host: str
|
||||
local_port: int
|
||||
remote_port: int
|
||||
ctrl_port: int
|
||||
num_transfers: int
|
||||
tensor_size_mb: int
|
||||
mode: str # "copy", "zerocopy", "gpu"
|
||||
gpu_id: int = 0
|
||||
benchmark: bool = False # Skip MD5 verification for pure performance test
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransferStats:
|
||||
"""Statistics for transfer operations."""
|
||||
|
||||
success_count: int = 0
|
||||
fail_count: int = 0
|
||||
total_bytes: int = 0
|
||||
elapsed_time: float = 0.0
|
||||
|
||||
@property
|
||||
def throughput_mbps(self) -> float:
|
||||
if self.elapsed_time > 0:
|
||||
return (self.total_bytes / (1024 * 1024)) / self.elapsed_time
|
||||
return 0.0
|
||||
|
||||
def print_summary(self, role: str):
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f" {role.upper()} SUMMARY")
|
||||
print(f" Successful: {self.success_count}/{self.success_count + self.fail_count}")
|
||||
print(f" Failed: {self.fail_count}/{self.success_count + self.fail_count}")
|
||||
print(f" Total: {self.total_bytes / (1024 * 1024):.2f} MB")
|
||||
print(f" Time: {self.elapsed_time:.2f} s")
|
||||
print(f" Throughput: {self.throughput_mbps:.2f} MB/s")
|
||||
print(f"{'=' * 60}")
|
||||
|
||||
|
||||
class CrossNodeTester(ABC):
|
||||
"""Abstract base class for cross-node RDMA testing."""
|
||||
|
||||
def __init__(self, config: TransferConfig):
|
||||
self.config = config
|
||||
self.connector: MooncakeTransferEngineConnector | None = None
|
||||
self.zmq_ctx: zmq.Context | None = None
|
||||
self.ctrl_socket: zmq.Socket | None = None
|
||||
self.stats = TransferStats()
|
||||
|
||||
def get_connector_config(self) -> dict:
|
||||
"""Get connector configuration based on mode."""
|
||||
pool_size = int(self.config.tensor_size_mb * 1.5) * 1024 * 1024
|
||||
pool_size = max(pool_size, 128 * 1024 * 1024)
|
||||
|
||||
conn_config = {
|
||||
"host": self.config.local_host,
|
||||
"zmq_port": self.config.local_port,
|
||||
"protocol": "rdma",
|
||||
"memory_pool_size": pool_size,
|
||||
}
|
||||
|
||||
# Set device based on mode
|
||||
if self.config.mode == "gpu":
|
||||
conn_config["memory_pool_device"] = f"cuda:{self.config.gpu_id}"
|
||||
else:
|
||||
conn_config["memory_pool_device"] = "cpu"
|
||||
|
||||
# RDMA device name from environment
|
||||
device_name = os.environ.get("RDMA_DEVICE_NAME")
|
||||
if device_name:
|
||||
conn_config["device_name"] = device_name
|
||||
print(f"[CONFIG] Using RDMA device: {device_name}")
|
||||
|
||||
return conn_config
|
||||
|
||||
def initialize(self):
|
||||
"""Initialize connector and ZMQ context."""
|
||||
print(f"[{self.role}] Initializing connector...")
|
||||
conn_config = self.get_connector_config()
|
||||
self.connector = MooncakeTransferEngineConnector(conn_config)
|
||||
self.zmq_ctx = zmq.Context()
|
||||
print(f"[{self.role}] Ready at {self.config.local_host}:{self.config.local_port}")
|
||||
|
||||
def cleanup(self):
|
||||
"""Cleanup resources."""
|
||||
if self.ctrl_socket:
|
||||
self.ctrl_socket.close()
|
||||
if self.zmq_ctx:
|
||||
self.zmq_ctx.term()
|
||||
if self.connector:
|
||||
self.connector.close()
|
||||
print(f"[{self.role}] Closed.")
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def role(self) -> str:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def run(self):
|
||||
pass
|
||||
|
||||
|
||||
class Producer(CrossNodeTester):
|
||||
"""Producer node - sends data to consumer."""
|
||||
|
||||
@property
|
||||
def role(self) -> str:
|
||||
return "PRODUCER"
|
||||
|
||||
def print_header(self):
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f" PRODUCER MODE ({self.config.mode.upper()})")
|
||||
print(f" Local: {self.config.local_host}:{self.config.local_port}")
|
||||
print(f" Remote: {self.config.remote_host}:{self.config.remote_port}")
|
||||
print(f" Control Port: {self.config.ctrl_port}")
|
||||
print(f" Transfer Mode: {self.config.mode}")
|
||||
if self.config.mode == "gpu":
|
||||
print(f" GPU ID: {self.config.gpu_id}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
def setup_control_channel(self):
|
||||
"""Setup ZMQ control channel as server (REP socket)."""
|
||||
self.ctrl_socket = self.zmq_ctx.socket(zmq.REP)
|
||||
self.ctrl_socket.bind(f"tcp://*:{self.config.ctrl_port}")
|
||||
print(f"[PRODUCER] Control channel listening on port {self.config.ctrl_port}")
|
||||
print("[PRODUCER] Waiting for consumer to connect...")
|
||||
|
||||
def wait_for_consumer(self) -> bool:
|
||||
"""Wait for consumer READY signal."""
|
||||
msg_data = self.ctrl_socket.recv()
|
||||
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
|
||||
if msg.msg_type != "READY":
|
||||
print(f"[PRODUCER] Unexpected message: {msg.msg_type}")
|
||||
return False
|
||||
print("[PRODUCER] Consumer connected!")
|
||||
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="ACK")))
|
||||
return True
|
||||
|
||||
def create_test_data(self, transfer_idx: int) -> tuple[Any, str, int]:
|
||||
"""
|
||||
Create test data based on transfer mode.
|
||||
|
||||
Returns:
|
||||
(data, md5, size) tuple
|
||||
"""
|
||||
num_elements = (self.config.tensor_size_mb * 1024 * 1024) // 4
|
||||
data_size = num_elements * 4
|
||||
|
||||
if self.config.mode == "zerocopy":
|
||||
# Zero-Copy Path: Allocate directly from connector's pool
|
||||
offset = self.connector.allocator.alloc(data_size)
|
||||
managed_buf = ManagedBuffer(self.connector.allocator, offset, data_size, self.connector.pool)
|
||||
|
||||
if self.config.benchmark:
|
||||
# In benchmark mode, skip random data generation (use uninitialized memory)
|
||||
return managed_buf, "", data_size
|
||||
else:
|
||||
# Fill buffer with random data using tensor view
|
||||
tensor_view = managed_buf.as_tensor(dtype=torch.float32, shape=(num_elements,))
|
||||
random_data = torch.randn(num_elements, dtype=torch.float32)
|
||||
if tensor_view.is_cuda:
|
||||
tensor_view.copy_(random_data.to(tensor_view.device))
|
||||
else:
|
||||
tensor_view.copy_(random_data)
|
||||
md5 = compute_md5(tensor_view)
|
||||
return managed_buf, md5, data_size
|
||||
|
||||
elif self.config.mode == "gpu":
|
||||
# GPU Path: Create tensor on GPU
|
||||
device = f"cuda:{self.config.gpu_id}"
|
||||
if self.config.benchmark:
|
||||
# In benchmark mode, use empty tensor (no random generation)
|
||||
gpu_tensor = torch.empty(num_elements, dtype=torch.float32, device=device)
|
||||
return gpu_tensor, "", data_size
|
||||
else:
|
||||
cpu_tensor = torch.randn(num_elements, dtype=torch.float32)
|
||||
md5 = compute_md5(cpu_tensor)
|
||||
gpu_tensor = cpu_tensor.to(device)
|
||||
return gpu_tensor, md5, data_size
|
||||
|
||||
else:
|
||||
# Copy Path (default): Create regular CPU tensor
|
||||
if self.config.benchmark:
|
||||
# In benchmark mode, use empty tensor (no random generation)
|
||||
tensor = torch.empty(num_elements, dtype=torch.float32)
|
||||
return tensor, "", data_size
|
||||
else:
|
||||
tensor = torch.randn(num_elements, dtype=torch.float32)
|
||||
md5 = compute_md5(tensor)
|
||||
return tensor, md5, data_size
|
||||
|
||||
def do_transfer(self, transfer_idx: int) -> bool:
|
||||
"""Perform a single transfer."""
|
||||
req_id = f"cross_node_transfer_{transfer_idx}"
|
||||
|
||||
if not self.config.benchmark:
|
||||
print(f"\n[PRODUCER] Transfer {transfer_idx + 1}/{self.config.num_transfers}")
|
||||
|
||||
# Create test data
|
||||
t0 = time.time()
|
||||
data, md5, data_size = self.create_test_data(transfer_idx)
|
||||
t_create = time.time() - t0
|
||||
|
||||
if not self.config.benchmark:
|
||||
print(f" Mode: {self.config.mode}")
|
||||
print(f" Size: {self.config.tensor_size_mb} MB")
|
||||
if md5:
|
||||
print(f" MD5: {md5[:16]}...")
|
||||
print(f" Create time: {t_create * 1000:.1f} ms")
|
||||
|
||||
# Put data
|
||||
t1 = time.time()
|
||||
success, size, metadata = self.connector.put("producer", "consumer", req_id, data)
|
||||
t_put = time.time() - t1
|
||||
|
||||
if not success:
|
||||
print(" [FAIL] Put failed")
|
||||
return False
|
||||
|
||||
if not self.config.benchmark:
|
||||
print(f" [OK] Put successful, {size} bytes ({t_put * 1000:.1f} ms)")
|
||||
|
||||
# Wait for consumer to request transfer info
|
||||
msg_data = self.ctrl_socket.recv()
|
||||
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
|
||||
|
||||
if msg.msg_type != "READY":
|
||||
print(f" [ERROR] Unexpected message: {msg.msg_type}")
|
||||
return False
|
||||
|
||||
# Send transfer metadata to consumer
|
||||
transfer_msg = CtrlMsg(
|
||||
msg_type="TRANSFER",
|
||||
request_id=req_id,
|
||||
md5=md5,
|
||||
data_size=data_size,
|
||||
)
|
||||
self.ctrl_socket.send(msgspec.msgpack.encode(transfer_msg))
|
||||
|
||||
# Wait for consumer ACK (this includes RDMA transfer time)
|
||||
t2 = time.time()
|
||||
msg_data = self.ctrl_socket.recv()
|
||||
t_rdma = time.time() - t2
|
||||
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
|
||||
|
||||
success = msg.msg_type == "ACK"
|
||||
if success:
|
||||
if not self.config.benchmark:
|
||||
print(f" [OK] RDMA transfer complete ({t_rdma * 1000:.1f} ms)")
|
||||
self.stats.success_count += 1
|
||||
self.stats.total_bytes += size
|
||||
else:
|
||||
print(f" [WARN] Consumer reported error: {msg.error}")
|
||||
self.stats.fail_count += 1
|
||||
|
||||
# Send ACK to allow consumer to continue
|
||||
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="ACK")))
|
||||
|
||||
# Cleanup buffer
|
||||
self.connector.cleanup(req_id)
|
||||
|
||||
return success
|
||||
|
||||
def run(self):
|
||||
"""Run the producer."""
|
||||
self.print_header()
|
||||
self.initialize()
|
||||
self.setup_control_channel()
|
||||
|
||||
try:
|
||||
if not self.wait_for_consumer():
|
||||
return
|
||||
|
||||
if self.config.benchmark:
|
||||
print(
|
||||
f"[BENCHMARK] Running {self.config.num_transfers} "
|
||||
f"transfers of {self.config.tensor_size_mb} MB each..."
|
||||
)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
for i in range(self.config.num_transfers):
|
||||
self.do_transfer(i)
|
||||
if self.config.benchmark and (i + 1) % 10 == 0:
|
||||
elapsed = time.time() - start_time
|
||||
current_throughput = (self.stats.total_bytes / (1024 * 1024)) / elapsed
|
||||
print(f" Progress: {i + 1}/{self.config.num_transfers}, Throughput: {current_throughput:.2f} MB/s")
|
||||
|
||||
self.stats.elapsed_time = time.time() - start_time
|
||||
self.stats.print_summary("PRODUCER")
|
||||
|
||||
# Wait for final consumer message and send DONE
|
||||
self.ctrl_socket.recv()
|
||||
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="DONE")))
|
||||
|
||||
finally:
|
||||
self.cleanup()
|
||||
|
||||
|
||||
class Consumer(CrossNodeTester):
|
||||
"""Consumer node - receives data from producer."""
|
||||
|
||||
@property
|
||||
def role(self) -> str:
|
||||
return "CONSUMER"
|
||||
|
||||
def print_header(self):
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f" CONSUMER MODE ({self.config.mode.upper()})")
|
||||
print(f" Local: {self.config.local_host}:{self.config.local_port}")
|
||||
print(f" Remote: {self.config.remote_host}:{self.config.remote_port}")
|
||||
print(f" Control Port: {self.config.ctrl_port}")
|
||||
print(f" Transfer Mode: {self.config.mode}")
|
||||
if self.config.mode == "gpu":
|
||||
print(f" GPU ID: {self.config.gpu_id}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
def setup_control_channel(self):
|
||||
"""Setup ZMQ control channel as client (REQ socket)."""
|
||||
self.ctrl_socket = self.zmq_ctx.socket(zmq.REQ)
|
||||
ctrl_addr = f"tcp://{self.config.remote_host}:{self.config.ctrl_port}"
|
||||
print(f"[CONSUMER] Connecting to producer control channel at {ctrl_addr}...")
|
||||
self.ctrl_socket.connect(ctrl_addr)
|
||||
|
||||
def connect_to_producer(self) -> bool:
|
||||
"""Connect to producer and send READY signal."""
|
||||
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
|
||||
msg_data = self.ctrl_socket.recv()
|
||||
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
|
||||
if msg.msg_type != "ACK":
|
||||
print(f"[CONSUMER] Unexpected response: {msg.msg_type}")
|
||||
return False
|
||||
print("[CONSUMER] Connected to producer! Starting transfers...")
|
||||
return True
|
||||
|
||||
def do_transfer(self, transfer_idx: int) -> bool:
|
||||
"""Perform a single transfer."""
|
||||
if not self.config.benchmark:
|
||||
print(f"\n[CONSUMER] Transfer {transfer_idx + 1}/{self.config.num_transfers}")
|
||||
|
||||
# Request next transfer info
|
||||
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
|
||||
msg_data = self.ctrl_socket.recv()
|
||||
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
|
||||
|
||||
if msg.msg_type == "DONE":
|
||||
print("[CONSUMER] Producer signaled completion")
|
||||
return False
|
||||
|
||||
if msg.msg_type != "TRANSFER":
|
||||
print(f"[CONSUMER] Unexpected message: {msg.msg_type}")
|
||||
return False
|
||||
|
||||
req_id = msg.request_id
|
||||
expected_md5 = msg.md5
|
||||
data_size = msg.data_size
|
||||
num_elements = data_size // 4
|
||||
|
||||
if not self.config.benchmark:
|
||||
print(f" Request ID: {req_id}")
|
||||
if expected_md5:
|
||||
print(f" Expected MD5: {expected_md5[:16]}...")
|
||||
print(f" Data Size: {data_size / (1024 * 1024):.2f} MB")
|
||||
|
||||
# Build metadata for get
|
||||
metadata = {
|
||||
"request_id": req_id,
|
||||
"source_host": self.config.remote_host,
|
||||
"source_port": self.config.remote_port,
|
||||
"data_size": data_size,
|
||||
"dtype": "float32",
|
||||
"shape": [num_elements],
|
||||
"is_fast_path": True,
|
||||
}
|
||||
|
||||
if not self.config.benchmark:
|
||||
print(f" [INFO] Requesting from {self.config.remote_host}:{self.config.remote_port}")
|
||||
|
||||
# Get data with timing
|
||||
t0 = time.time()
|
||||
result = self.connector.get("producer", "consumer", req_id, metadata)
|
||||
t_get = time.time() - t0
|
||||
|
||||
response_msg = CtrlMsg(msg_type="ERROR", error="Get failed")
|
||||
|
||||
if result is not None:
|
||||
recv_buffer, recv_size = result
|
||||
if not self.config.benchmark:
|
||||
print(f" [OK] Get successful, {recv_size} bytes ({t_get * 1000:.1f} ms)")
|
||||
|
||||
if isinstance(recv_buffer, ManagedBuffer):
|
||||
# In benchmark mode, skip MD5 verification
|
||||
if self.config.benchmark or not expected_md5:
|
||||
response_msg = CtrlMsg(msg_type="ACK")
|
||||
self.stats.success_count += 1
|
||||
self.stats.total_bytes += recv_size
|
||||
else:
|
||||
# Verify data
|
||||
t1 = time.time()
|
||||
reconstructed = recv_buffer.as_tensor(dtype=torch.float32, shape=(num_elements,))
|
||||
recv_md5 = compute_md5(reconstructed)
|
||||
t_md5 = time.time() - t1
|
||||
print(f" MD5: {recv_md5[:16]}... ({t_md5 * 1000:.1f} ms)")
|
||||
|
||||
if recv_md5 == expected_md5:
|
||||
print(" [PASS] MD5 checksum verified!")
|
||||
response_msg = CtrlMsg(msg_type="ACK")
|
||||
self.stats.success_count += 1
|
||||
self.stats.total_bytes += recv_size
|
||||
else:
|
||||
print(" [FAIL] MD5 mismatch!")
|
||||
response_msg = CtrlMsg(msg_type="ERROR", error="MD5 mismatch")
|
||||
self.stats.fail_count += 1
|
||||
|
||||
recv_buffer.release()
|
||||
else:
|
||||
response_msg = CtrlMsg(msg_type="ACK")
|
||||
self.stats.success_count += 1
|
||||
self.stats.total_bytes += recv_size
|
||||
else:
|
||||
print(" [FAIL] Get failed")
|
||||
self.stats.fail_count += 1
|
||||
|
||||
# Send response to producer
|
||||
self.ctrl_socket.send(msgspec.msgpack.encode(response_msg))
|
||||
# Wait for ACK
|
||||
self.ctrl_socket.recv()
|
||||
|
||||
return response_msg.msg_type == "ACK"
|
||||
|
||||
def run(self):
|
||||
"""Run the consumer."""
|
||||
self.print_header()
|
||||
self.initialize()
|
||||
self.setup_control_channel()
|
||||
|
||||
try:
|
||||
if not self.connect_to_producer():
|
||||
return
|
||||
|
||||
if self.config.benchmark:
|
||||
print(
|
||||
f"[BENCHMARK] Running {self.config.num_transfers} "
|
||||
f"transfers of {self.config.tensor_size_mb} MB each..."
|
||||
)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
for i in range(self.config.num_transfers):
|
||||
if not self.do_transfer(i):
|
||||
break
|
||||
if self.config.benchmark and (i + 1) % 10 == 0:
|
||||
elapsed = time.time() - start_time
|
||||
current_throughput = (self.stats.total_bytes / (1024 * 1024)) / elapsed
|
||||
print(f" Progress: {i + 1}/{self.config.num_transfers}, Throughput: {current_throughput:.2f} MB/s")
|
||||
|
||||
self.stats.elapsed_time = time.time() - start_time
|
||||
self.stats.print_summary("CONSUMER")
|
||||
|
||||
# Send final READY and wait for DONE
|
||||
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
|
||||
self.ctrl_socket.recv()
|
||||
|
||||
finally:
|
||||
self.cleanup()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Cross-Node RDMA Test (Automated)",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Transfer Modes:
|
||||
copy - Normal path: tensor copied to RDMA pool (default)
|
||||
zerocopy - Zero-copy path: data created directly in RDMA pool
|
||||
gpu - GPU transfer: RDMA pool on GPU, uses GPUDirect
|
||||
|
||||
Examples:
|
||||
# Copy mode (default):
|
||||
python cross_node_mooncake_transfer_engine.py --role producer \
|
||||
--local-host hostA --remote-host hostB
|
||||
|
||||
# Zero-copy mode:
|
||||
python cross_node_mooncake_transfer_engine.py --role producer \
|
||||
--local-host hostA --remote-host hostB --mode zerocopy
|
||||
|
||||
# GPU mode:
|
||||
python cross_node_mooncake_transfer_engine.py --role producer \
|
||||
--local-host hostA --remote-host hostB --mode gpu --gpu-id 0
|
||||
|
||||
# Benchmark mode (skip MD5, measure pure RDMA performance):
|
||||
python cross_node_mooncake_transfer_engine.py --role producer \
|
||||
--local-host hostA --remote-host hostB --benchmark
|
||||
|
||||
# With specific RDMA device:
|
||||
RDMA_DEVICE_NAME=mlx5_0 python cross_node_mooncake_transfer_engine.py --role producer ...
|
||||
""",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--role", required=True, choices=["producer", "consumer"], help="Role: producer (sends) or consumer (receives)"
|
||||
)
|
||||
parser.add_argument("--local-host", required=True, help="Local hostname or IP address")
|
||||
parser.add_argument("--remote-host", required=True, help="Remote hostname or IP address")
|
||||
parser.add_argument("--local-port", type=int, default=15500, help="Local ZMQ port for RDMA data (default: 15500)")
|
||||
parser.add_argument("--remote-port", type=int, default=15500, help="Remote ZMQ port for RDMA data (default: 15500)")
|
||||
parser.add_argument("--ctrl-port", type=int, default=15501, help="Control channel port (default: 15501)")
|
||||
parser.add_argument("--num-transfers", type=int, default=20, help="Number of transfers to perform (default: 3)")
|
||||
parser.add_argument("--tensor-size-mb", type=int, default=100, help="Tensor size in MB (default: 100)")
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=["copy", "zerocopy", "gpu"],
|
||||
default="copy",
|
||||
help="Transfer mode: copy, zerocopy, or gpu (default: copy)",
|
||||
)
|
||||
parser.add_argument("--gpu-id", type=int, default=0, help="GPU ID for GPU mode (default: 0)")
|
||||
parser.add_argument(
|
||||
"--benchmark", action="store_true", help="Benchmark mode: skip MD5 verification for pure performance test"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Check Mooncake
|
||||
if TransferEngine is None:
|
||||
print("[ERROR] Mooncake TransferEngine is not available.")
|
||||
print("Install with: pip install mooncake")
|
||||
sys.exit(1)
|
||||
|
||||
# Check CUDA for GPU mode
|
||||
if args.mode == "gpu":
|
||||
if not torch.cuda.is_available():
|
||||
print("[ERROR] CUDA is not available but GPU mode was requested.")
|
||||
sys.exit(1)
|
||||
if args.gpu_id >= torch.cuda.device_count():
|
||||
print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.cuda.device_count()} GPUs.")
|
||||
sys.exit(1)
|
||||
print(f"[INFO] Using GPU {args.gpu_id}: {torch.cuda.get_device_name(args.gpu_id)}")
|
||||
|
||||
config = TransferConfig(
|
||||
local_host=args.local_host,
|
||||
remote_host=args.remote_host,
|
||||
local_port=args.local_port,
|
||||
remote_port=args.remote_port,
|
||||
ctrl_port=args.ctrl_port,
|
||||
num_transfers=args.num_transfers,
|
||||
tensor_size_mb=args.tensor_size_mb,
|
||||
mode=args.mode,
|
||||
gpu_id=args.gpu_id,
|
||||
benchmark=args.benchmark,
|
||||
)
|
||||
|
||||
if args.role == "producer":
|
||||
tester = Producer(config)
|
||||
else:
|
||||
tester = Consumer(config)
|
||||
|
||||
tester.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+32
-7
@@ -1,8 +1,20 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.15.0
|
||||
FROM ${BASE_IMAGE} AS final
|
||||
FROM ${BASE_IMAGE} AS base
|
||||
|
||||
# Declare a variable to know if we want to use the nightly build or the stable build.
|
||||
# NOTE: REMINDER to vLLM-Omni rebase maintainer
|
||||
# Remember to set `USE_NIGHTLY_BUILD` to 0, when switching back to stable vLLM docker image
|
||||
# 1. If vLLM-omni maintainer is forced to use custom commits
|
||||
# during rebasing, they can change the two variables
|
||||
# 2. Whenever vLLM upstream has released a stable version,
|
||||
# we should swap over to use stable release ASAP.
|
||||
# We should avoid relying on custom commits.
|
||||
ARG USE_NIGHTLY_BUILD=1
|
||||
ARG VLLM_VERSION_OR_COMMIT_HASH=2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
|
||||
ARG ARG_PYTORCH_ROCM_ARCH
|
||||
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
|
||||
|
||||
ARG COMMON_WORKDIR=/app
|
||||
|
||||
WORKDIR ${COMMON_WORKDIR}
|
||||
|
||||
# Step 1: Setup - Install system dependencies
|
||||
@@ -11,18 +23,31 @@ RUN apt-get update && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
|
||||
# Step 2: Conditionally reinstall vllm from source for nightly builds
|
||||
RUN if [ "${USE_NIGHTLY_BUILD}" = "1" ]; then \
|
||||
python3 -m pip uninstall -y vllm && rm -rf vllm && \
|
||||
git clone https://github.com/vllm-project/vllm.git && \
|
||||
cd vllm && \
|
||||
git checkout ${VLLM_VERSION_OR_COMMIT_HASH} && \
|
||||
python3 -m pip install -r requirements/rocm.txt && \
|
||||
python3 setup.py clean --all && \
|
||||
PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} python3 setup.py develop && \
|
||||
cd ../ && \
|
||||
rm -rf vllm/.git; \
|
||||
fi
|
||||
|
||||
# Step 2: Copy vllm-omni code and install without uv
|
||||
# Step 3: Copy vllm-omni code and install without uv
|
||||
RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
|
||||
COPY . ${COMMON_WORKDIR}/vllm-omni
|
||||
RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]" --no-build-isolation
|
||||
|
||||
RUN ln -sf /usr/bin/python3 /usr/bin/python
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
FROM base AS test
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
ENTRYPOINT []
|
||||
|
||||
#Set entrypoint for vllm-openai official images
|
||||
FROM final AS vllm-openai
|
||||
# Set entrypoint for vllm-openai official images
|
||||
FROM base AS vllm-openai
|
||||
ENTRYPOINT ["vllm", "serve", "--omni"]
|
||||
@@ -6,7 +6,8 @@ This guide explains how to configure and use distributed connectors
|
||||
Backend-specific setup lives in separate docs:
|
||||
|
||||
- [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
|
||||
- [MooncakeConnector](omni_connectors/mooncake_connector.md)
|
||||
- [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
|
||||
- [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
|
||||
- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
|
||||
|
||||
## Overview
|
||||
@@ -19,7 +20,8 @@ Current connectors operate in D2H2D (device to host to device) mode.
|
||||
| Use Case | Recommended Connector | Notes |
|
||||
| :--- | :--- | :--- |
|
||||
| Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
|
||||
| Multi node (Mooncake) | MooncakeConnector | Requires Mooncake Master + metadata server. |
|
||||
| Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
|
||||
| Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
|
||||
| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
|
||||
|
||||
## Core API
|
||||
|
||||
+2
-2
@@ -1,4 +1,4 @@
|
||||
# MooncakeConnector
|
||||
# MooncakeStoreConnector
|
||||
|
||||
## When to Use
|
||||
|
||||
@@ -38,7 +38,7 @@ Define the connector in runtime:
|
||||
runtime:
|
||||
connectors:
|
||||
connector_of_mooncake:
|
||||
name: MooncakeConnector
|
||||
name: MooncakeStoreConnector
|
||||
extra:
|
||||
host: "127.0.0.1"
|
||||
metadata_server: "http://<MASTER_IP>:8080/metadata"
|
||||
@@ -0,0 +1,207 @@
|
||||
# MooncakeTransferEngineConnector
|
||||
|
||||
## When to Use
|
||||
|
||||
Best for high-performance multi-node data transfer between stages using Mooncake
|
||||
Transfer Engine. Supports both RDMA and TCP protocols with a managed memory pool,
|
||||
zero-copy deserialization, and optional GPUDirect RDMA. Applicable to any
|
||||
inter-stage data (KV caches, request payloads, etc.), not limited to KV cache transfer.
|
||||
|
||||
Compared to `MooncakeStoreConnector` (TCP key-value store), this connector
|
||||
provides **~60x faster** data transfer via RDMA direct memory access.
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install mooncake-transfer-engine
|
||||
```
|
||||
|
||||
Ensure RDMA drivers are installed on all nodes (e.g., Mellanox OFED for
|
||||
InfiniBand/RoCE NICs).
|
||||
|
||||
## Configuration
|
||||
|
||||
Define the connector in runtime:
|
||||
|
||||
```yaml
|
||||
runtime:
|
||||
connectors:
|
||||
rdma_connector:
|
||||
name: MooncakeTransferEngineConnector
|
||||
extra:
|
||||
host: "auto" # Auto-detect local RDMA IP
|
||||
zmq_port: 50051 # ZMQ base port (see "Port Offset Scheme" below)
|
||||
protocol: "rdma" # "rdma" or "tcp"
|
||||
device_name: "" # RDMA device (e.g., "mlx5_0"), empty for auto-detect
|
||||
memory_pool_size: 2147483648 # 2GB memory pool
|
||||
memory_pool_device: "cpu" # "cpu" for pinned memory, "cuda" for GPUDirect RDMA
|
||||
```
|
||||
|
||||
Wire stages to the connector:
|
||||
|
||||
```yaml
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
output_connectors:
|
||||
to_stage_1: rdma_connector
|
||||
|
||||
- stage_id: 1
|
||||
input_connectors:
|
||||
from_stage_0: rdma_connector
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
### Required
|
||||
|
||||
| Parameter | Description |
|
||||
|---|---|
|
||||
| `role` | **Internal, do not set manually.** Auto-injected by the orchestration layer (`"sender"` for `output_connectors`, `"receiver"` for `input_connectors`). Defaults to `"sender"` if omitted. |
|
||||
| `host` | Local IP address for RDMA. `"auto"` detects from network interfaces. |
|
||||
| `protocol` | Transport protocol: `"rdma"` (InfiniBand/RoCE) or `"tcp"`. |
|
||||
|
||||
### Memory Pool
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|---|---|---|
|
||||
| `memory_pool_size` | 1 GB | Total size of the RDMA-registered memory pool in bytes. |
|
||||
| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity). |
|
||||
|
||||
### Networking
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|---|---|---|
|
||||
| `zmq_port` | 50051 | ZMQ **base** port. The orchestration layer computes the actual port as `base + purpose_offset + stage_offset` (see table below). Users only set this base value. |
|
||||
| `sender_host` | `None` | **Internal.** Receiver-side only — dynamically resolved via `update_sender_info()`. Not needed in YAML. |
|
||||
| `sender_zmq_port` | `None` | **Internal.** Receiver-side only — defaults to the sender's adjusted port. Not needed in YAML. |
|
||||
| `device_name` | `""` | RDMA device name (e.g., `"mlx5_0"`). Empty for auto-detect. Can also be set via `RDMA_DEVICE_NAME` env var. |
|
||||
|
||||
#### ZMQ Port Offset Scheme
|
||||
|
||||
To avoid port conflicts when multiple edges, purposes, DP replicas, or TP ranks share the same node, the actual ZMQ port is computed as:
|
||||
|
||||
```
|
||||
side_channel_port = zmq_port + purpose_offset + stage_offset + dp_index * tp_size
|
||||
sender_listen = side_channel_port + tp_rank
|
||||
receiver_connect = remote_side_channel_port + tp_rank
|
||||
```
|
||||
|
||||
| Component | Value | Description |
|
||||
|---|---|---|
|
||||
| `zmq_port` | 50051 (default) | Base port from YAML config |
|
||||
| `purpose_offset` | `request_forwarding` = 0, `kv_transfer` = 100 | Separates control-plane vs KV-cache connections |
|
||||
| `stage_offset` | `int(from_stage)` (0, 1, 2...) | Separates edges from different source stages |
|
||||
| `dp_index * tp_size` | e.g., DP1 × TP2 = 2 | Each DP replica reserves a port range of size `tp_size` (following vLLM convention: `VLLM_MOONCAKE_BOOTSTRAP_PORT + dp_index * tp_size`) |
|
||||
| `tp_rank` | 0, 1, 2... | Each TP rank within a DP replica uses its own port |
|
||||
| orchestrator | +200 | Extra offset when caller is the orchestrator (avoids collision with stage workers on the same node) |
|
||||
|
||||
**Example** (base=50051, stage 0→1, DP=2, TP=2, kv_transfer):
|
||||
|
||||
| Caller | DP | TP rank | Port |
|
||||
|---|---|---|---|
|
||||
| Stage worker | DP0 | rank 0 | `50051 + 100 + 0 + 0×2 + 0 = 50151` |
|
||||
| Stage worker | DP0 | rank 1 | `50051 + 100 + 0 + 0×2 + 1 = 50152` |
|
||||
| Stage worker | DP1 | rank 0 | `50051 + 100 + 0 + 1×2 + 0 = 50153` |
|
||||
| Stage worker | DP1 | rank 1 | `50051 + 100 + 0 + 1×2 + 1 = 50154` |
|
||||
| Orchestrator | — | — | `50051 + 200 + 0 = 50251` |
|
||||
|
||||
## Memory Pool Modes
|
||||
|
||||
| Mode | Config | Data Flow | Best For |
|
||||
|---|---|---|---|
|
||||
| CPU Pinned | `memory_pool_device: "cpu"` | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
|
||||
| GPUDirect | `memory_pool_device: "cuda"` | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
|
||||
|
||||
> **Note**: GPUDirect RDMA requires the NIC and GPU to share a direct PCIe
|
||||
> switch (PIX topology). On systems where they are connected via PXB or NODE,
|
||||
> CPU pinned memory is faster due to GPU BAR1 bandwidth limitations.
|
||||
|
||||
## Environment Variables
|
||||
|
||||
| Variable | Description |
|
||||
|---|---|
|
||||
| `RDMA_DEVICE_NAME` | Override RDMA device name (e.g., `mlx5_0`). |
|
||||
| `MC_IB_PCI_RELAXED_ORDERING` | Set to `1` to enable PCIe relaxed ordering for GPUDirect. |
|
||||
|
||||
## Docker / Container Setup
|
||||
|
||||
RDMA requires host-level device access:
|
||||
|
||||
```bash
|
||||
docker run -it \
|
||||
--cap-add=SYS_PTRACE \
|
||||
--cap-add=IPC_LOCK \
|
||||
--security-opt seccomp=unconfined \
|
||||
--network=host \
|
||||
--device=/dev/infiniband \
|
||||
-v /sys/class/infiniband:/sys/class/infiniband:ro \
|
||||
your-image:tag
|
||||
```
|
||||
|
||||
## Performance
|
||||
|
||||
Benchmark results on H800 GPUs with mlx5_0 RDMA NIC (~186 MB KV cache):
|
||||
|
||||
| Metric | MooncakeStoreConnector | MooncakeTransferEngineConnector (CPU) |
|
||||
|---|---|---|
|
||||
| KV transfer wall time | ~810 ms | **~14 ms** |
|
||||
| RDMA throughput | N/A (TCP) | ~22 GB/s |
|
||||
| Speedup | 1x | **58x** |
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Quick Diagnostics
|
||||
|
||||
```bash
|
||||
# 1. Check RDMA devices and link status
|
||||
ibdev2netdev
|
||||
# Expected: "mlx5_X port 1 ==> <iface> (Up)"
|
||||
# RoCE devices map to Ethernet interfaces (e.g., enp75s0f0)
|
||||
# IB devices map to ib0, ib1, etc.
|
||||
|
||||
# 2. Check InfiniBand device details
|
||||
ibstat
|
||||
|
||||
# 3. Verify /dev/infiniband is accessible (required in containers)
|
||||
ls /dev/infiniband/
|
||||
|
||||
# 4. Check Mooncake installation
|
||||
python -c "from mooncake.engine import TransferEngine; print('OK')"
|
||||
|
||||
# 5. Check environment variables
|
||||
echo "RDMA_DEVICE_NAME=${RDMA_DEVICE_NAME:-<not set>}"
|
||||
echo "MC_IB_PCI_RELAXED_ORDERING=${MC_IB_PCI_RELAXED_ORDERING:-<not set>}"
|
||||
```
|
||||
|
||||
### Common Issues
|
||||
|
||||
| Symptom | Cause | Fix |
|
||||
|---|---|---|
|
||||
| `Failed to modify QP to RTR` | Cross-NIC QP handshake failure (multi-NIC DGX) | Set `device_name` to a single RoCE NIC (e.g., `mlx5_2`) or set `RDMA_DEVICE_NAME` env var |
|
||||
| `transport retry counter exceeded` | RDMA path between incompatible NICs | Same as above — restrict to one NIC |
|
||||
| `zmq.error.Again: Resource temporarily unavailable` | ZMQ recv timeout (transfer took too long) | Check NIC selection; increase data may need longer timeout |
|
||||
| `Mooncake Engine initialization failed` | Missing RDMA drivers or `/dev/infiniband` | Install Mellanox OFED; in Docker add `--device=/dev/infiniband` |
|
||||
| `MemoryError` in allocator | Memory pool too small for payload | Increase `memory_pool_size` |
|
||||
| GPU transfer slower than CPU | GPU BAR1 bandwidth limitation (PXB/NODE topology) | Use `memory_pool_device: "cpu"` instead of `"cuda"` |
|
||||
|
||||
### Multi-NIC Environments (DGX)
|
||||
|
||||
On DGX machines with 12+ RDMA NICs, only RoCE NICs (with a bound network
|
||||
interface) reliably support loopback. IB-only NICs may fail cross-NIC QP
|
||||
handshakes. To identify RoCE NICs:
|
||||
|
||||
```bash
|
||||
ibdev2netdev | grep -v "ib[0-9]"
|
||||
# RoCE devices show Ethernet interface names like enp75s0f0
|
||||
```
|
||||
|
||||
Then configure the connector:
|
||||
```yaml
|
||||
device_name: "mlx5_2" # or set RDMA_DEVICE_NAME=mlx5_2
|
||||
```
|
||||
|
||||
See the RDMA Test README in tests/distributed/omni_connectors/README.md
|
||||
for test-specific setup instructions.
|
||||
|
||||
For more details on the underlying engine, refer to the
|
||||
[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
|
||||
@@ -45,7 +45,7 @@ ray start --address=<HEAD_NODE_IP>:6399
|
||||
|
||||
When running on Ray, the system automatically adapts its communication strategy:
|
||||
|
||||
* **Cross-Node**: Recommended to use `MooncakeConnector` (requires separate configuration).
|
||||
* **Cross-Node**: Recommended to use `MooncakeTransferEngineConnector` (RDMA, fastest) or `MooncakeStoreConnector` (TCP fallback).
|
||||
* **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
|
||||
* **SHM threshold default differs**: when `worker_backend="ray"`, the SharedMemoryConnector default threshold is set to `sys.maxsize`, which forces payloads to go inline (no SHM). Override `shm_threshold_bytes` in the connector config if you want SHM for Ray runs.
|
||||
|
||||
|
||||
@@ -20,7 +20,20 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
```bash
|
||||
uv pip install vllm==0.16.0 --torch-backend=auto
|
||||
# vllm 0.16.0 is still under prerelease
|
||||
uv pip install --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
|
||||
|
||||
# vllm 0.16.0 may have some bugs for cuda 12.9, here is how we solve them:
|
||||
export FLASHINFER_CUDA_TAG="$(python3 -c 'import torch; print((torch.version.cuda or "12.4").replace(".", ""))')"
|
||||
|
||||
uv pip install --upgrade --force-reinstall \
|
||||
"flashinfer-python==0.6.3" \
|
||||
"flashinfer-cubin==0.6.3" \
|
||||
"flashinfer-jit-cache==0.6.3" \
|
||||
--extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}"
|
||||
|
||||
uv pip install --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4"
|
||||
uv pip install --upgrade --force-reinstall "numpy==2.2.6"
|
||||
```
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
|
||||
@@ -19,7 +19,18 @@ uv venv --python 3.12 --seed
|
||||
source .venv/bin/activate
|
||||
|
||||
# On CUDA
|
||||
uv pip install vllm==0.16.0 --torch-backend=auto
|
||||
# vllm 0.16.0 is still under prerelease
|
||||
uv pip install --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
|
||||
# vllm 0.16.0 may have some bugs for cuda 12.9, here is how we solve them:
|
||||
export FLASHINFER_CUDA_TAG="$(python3 -c 'import torch; print((torch.version.cuda or "12.4").replace(".", ""))')"
|
||||
uv pip install --upgrade --force-reinstall \
|
||||
"flashinfer-python==0.6.3" \
|
||||
"flashinfer-cubin==0.6.3" \
|
||||
"flashinfer-jit-cache==0.6.3" \
|
||||
--extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}"
|
||||
uv pip install --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4"
|
||||
uv pip install --upgrade --force-reinstall "numpy==2.2.6"
|
||||
|
||||
|
||||
# On ROCm
|
||||
uv pip install vllm==0.16.0 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
|
||||
|
||||
@@ -0,0 +1,230 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Unit tests for BufferAllocator and ManagedBuffer.
|
||||
These tests do NOT require Mooncake or RDMA environment.
|
||||
"""
|
||||
|
||||
import threading
|
||||
import unittest
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
|
||||
BufferAllocator,
|
||||
ManagedBuffer,
|
||||
)
|
||||
|
||||
# All tests in this file are pure-CPU unit tests for the memory allocator.
|
||||
pytestmark = [pytest.mark.cpu, pytest.mark.parallel]
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
class TestBufferAllocator(unittest.TestCase):
|
||||
"""Unit tests for BufferAllocator."""
|
||||
|
||||
def test_basic_alloc_free(self):
|
||||
"""Verify alloc, free, and reuse of freed space."""
|
||||
allocator = BufferAllocator(total_size=4096, alignment=64)
|
||||
|
||||
offset1 = allocator.alloc(512)
|
||||
self.assertEqual(offset1, 0)
|
||||
|
||||
offset2 = allocator.alloc(512)
|
||||
self.assertGreater(offset2, 0)
|
||||
|
||||
# Free first block, should be reusable
|
||||
allocator.free(offset1, 512)
|
||||
offset3 = allocator.alloc(512)
|
||||
self.assertEqual(offset3, 0)
|
||||
|
||||
def test_alignment(self):
|
||||
"""Verify allocation respects alignment."""
|
||||
allocator = BufferAllocator(total_size=4096, alignment=128)
|
||||
|
||||
_offset1 = allocator.alloc(100)
|
||||
offset2 = allocator.alloc(100)
|
||||
|
||||
self.assertEqual(offset2 % 128, 0)
|
||||
self.assertEqual(offset2, 128)
|
||||
|
||||
def test_exhaustion_and_recovery(self):
|
||||
"""Test that full allocation fails, then succeeds after free."""
|
||||
allocator = BufferAllocator(total_size=1024, alignment=64)
|
||||
|
||||
offset = allocator.alloc(1024)
|
||||
self.assertEqual(offset, 0)
|
||||
|
||||
with self.assertRaises(MemoryError):
|
||||
allocator.alloc(64)
|
||||
|
||||
allocator.free(offset, 1024)
|
||||
offset2 = allocator.alloc(1024)
|
||||
self.assertEqual(offset2, 0)
|
||||
|
||||
def test_thread_safety(self):
|
||||
"""Verify allocator is thread-safe under concurrent access."""
|
||||
allocator = BufferAllocator(total_size=1024 * 1024, alignment=64)
|
||||
errors = []
|
||||
|
||||
def worker(worker_id):
|
||||
try:
|
||||
for i in range(100):
|
||||
size = 1024 + (i % 10) * 64
|
||||
offset = allocator.alloc(size)
|
||||
allocator.free(offset, size)
|
||||
except Exception as e:
|
||||
errors.append(f"Worker {worker_id}: {e}")
|
||||
|
||||
threads = [threading.Thread(target=worker, args=(i,)) for i in range(10)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
self.assertEqual(len(errors), 0, f"Thread safety errors: {errors}")
|
||||
|
||||
|
||||
class TestAllocatorInvariants(unittest.TestCase):
|
||||
"""
|
||||
Defensive invariant tests for BufferAllocator: double-free, partial
|
||||
overlap corruption, and adjacent-block merging.
|
||||
|
||||
Marked @pytest.mark.slow so they are skipped in quick CI but retained
|
||||
as regression safety-net.
|
||||
"""
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_double_free_exact_is_safe(self):
|
||||
"""Double-free of exact same block should warn but NOT crash."""
|
||||
allocator = BufferAllocator(total_size=4096, alignment=64)
|
||||
offset = allocator.alloc(256)
|
||||
allocator.free(offset, 256)
|
||||
# Second free of the same block — should be silently ignored
|
||||
allocator.free(offset, 256)
|
||||
# Pool should still be consistent: allocate full size back
|
||||
offset2 = allocator.alloc(4096)
|
||||
self.assertEqual(offset2, 0)
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_double_free_after_merge_is_safe(self):
|
||||
"""
|
||||
Free A then B (adjacent → merged), then free A again.
|
||||
The allocator must detect A is already within the merged block.
|
||||
"""
|
||||
allocator = BufferAllocator(total_size=4096, alignment=64)
|
||||
a = allocator.alloc(64)
|
||||
b = allocator.alloc(64)
|
||||
allocator.free(a, 64)
|
||||
allocator.free(b, 64) # triggers merge with A
|
||||
# Now free A again — contained within the merged block
|
||||
allocator.free(a, 64) # should not raise
|
||||
# Pool should still be fully usable
|
||||
offset = allocator.alloc(4096)
|
||||
self.assertEqual(offset, 0)
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_partial_overlap_raises_corruption(self):
|
||||
"""Freeing a region that partially overlaps a free block must raise RuntimeError."""
|
||||
allocator = BufferAllocator(total_size=4096, alignment=64)
|
||||
a = allocator.alloc(128)
|
||||
b = allocator.alloc(128)
|
||||
allocator.free(a, 128) # [0, 128) is now free
|
||||
# Try to free a region that starts inside [0,128) but extends beyond
|
||||
with self.assertRaises(RuntimeError):
|
||||
allocator.free(64, 128) # [64, 192) overlaps with free [0, 128)
|
||||
|
||||
# b is still allocated; freeing b should be fine
|
||||
allocator.free(b, 128)
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_merge_adjacent_blocks(self):
|
||||
"""Free three adjacent blocks; they should merge into one contiguous region."""
|
||||
allocator = BufferAllocator(total_size=4096, alignment=64)
|
||||
a = allocator.alloc(64)
|
||||
b = allocator.alloc(64)
|
||||
c = allocator.alloc(64)
|
||||
# Free in non-sequential order to exercise sorting + merging
|
||||
allocator.free(b, 64)
|
||||
allocator.free(a, 64)
|
||||
allocator.free(c, 64)
|
||||
# After merge, free_blocks should contain one block starting at 0
|
||||
# covering at least 192 bytes (3 * 64).
|
||||
# Verify by allocating a contiguous block of 192 bytes.
|
||||
offset = allocator.alloc(192)
|
||||
self.assertEqual(offset, 0, "Adjacent blocks were not merged properly")
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_fragmentation_and_defrag(self):
|
||||
"""
|
||||
Allocate A B C D that exactly fill the pool, free B and D to
|
||||
create fragmentation, verify a large contiguous alloc fails,
|
||||
then free A and C — should result in full defrag.
|
||||
"""
|
||||
# Total pool = 4 * 64 = 256 bytes, so 4 allocs exhaust it completely
|
||||
allocator = BufferAllocator(total_size=256, alignment=64)
|
||||
a = allocator.alloc(64)
|
||||
b = allocator.alloc(64)
|
||||
c = allocator.alloc(64)
|
||||
d = allocator.alloc(64)
|
||||
|
||||
allocator.free(b, 64) # free blocks: [64, 128)
|
||||
allocator.free(d, 64) # free blocks: [64, 128) and [192, 256)
|
||||
|
||||
# Pool has two 64-byte holes; contiguous 128 is unavailable
|
||||
with self.assertRaises(MemoryError):
|
||||
allocator.alloc(128)
|
||||
|
||||
allocator.free(a, 64)
|
||||
allocator.free(c, 64)
|
||||
|
||||
# After freeing everything, full pool should be available
|
||||
offset = allocator.alloc(256)
|
||||
self.assertEqual(offset, 0)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
class TestManagedBuffer(unittest.TestCase):
|
||||
"""Unit tests for ManagedBuffer."""
|
||||
|
||||
def setUp(self):
|
||||
self.allocator = BufferAllocator(total_size=4096, alignment=64)
|
||||
self.pool = torch.zeros(4096, dtype=torch.uint8)
|
||||
|
||||
def test_tensor_view(self):
|
||||
"""Verify tensor property and as_tensor return correct views."""
|
||||
offset = self.allocator.alloc(64)
|
||||
buf = ManagedBuffer(self.allocator, offset, 64, self.pool)
|
||||
|
||||
# Write float32 data via pool
|
||||
src = torch.arange(16, dtype=torch.float32)
|
||||
self.pool[offset : offset + 64] = src.view(torch.uint8)
|
||||
|
||||
# Raw uint8 view
|
||||
self.assertEqual(buf.tensor.shape[0], 64)
|
||||
|
||||
# Typed view
|
||||
typed = buf.as_tensor(dtype=torch.float32, shape=(4, 4))
|
||||
self.assertEqual(typed.shape, (4, 4))
|
||||
self.assertTrue(torch.equal(typed.flatten(), src))
|
||||
|
||||
buf.release()
|
||||
|
||||
def test_context_manager_releases_buffer(self):
|
||||
"""Verify context manager releases buffer and space is reusable."""
|
||||
offset = self.allocator.alloc(128)
|
||||
|
||||
with ManagedBuffer(self.allocator, offset, 128, self.pool) as buf:
|
||||
self.assertFalse(buf._released)
|
||||
|
||||
self.assertTrue(buf._released)
|
||||
|
||||
# Space should be reusable
|
||||
new_offset = self.allocator.alloc(128)
|
||||
self.assertEqual(new_offset, offset)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,763 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Integration tests for MooncakeTransferEngineConnector.
|
||||
|
||||
These tests require Mooncake TransferEngine and an RDMA environment.
|
||||
Set RDMA_DEVICE_NAME to force a specific device for single-node testing.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
|
||||
ManagedBuffer,
|
||||
MooncakeTransferEngineConnector,
|
||||
TransferEngine,
|
||||
)
|
||||
|
||||
# All tests in this file require Mooncake TransferEngine and an RDMA environment.
|
||||
pytestmark = [pytest.mark.parallel, pytest.mark.gpu]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def get_rdma_host() -> str:
|
||||
"""Resolve RDMA-capable host IP (env > auto-detect IB > hostname > 127.0.0.1)."""
|
||||
import socket
|
||||
|
||||
env_host = os.environ.get("RDMA_TEST_HOST")
|
||||
if env_host:
|
||||
return env_host
|
||||
|
||||
for ip_cmd in ("ip", "/sbin/ip", "/usr/sbin/ip"):
|
||||
try:
|
||||
result = subprocess.run([ip_cmd, "addr", "show"], capture_output=True, text=True, timeout=5)
|
||||
if result.returncode != 0:
|
||||
continue
|
||||
lines = result.stdout.split("\n")
|
||||
rdma_patterns = ("ibp", "ib0", "ib1", "ib2", "mlx", "roce")
|
||||
for line in lines:
|
||||
if "inet " in line:
|
||||
for pat in rdma_patterns:
|
||||
if pat in line.lower():
|
||||
ip = line.strip().split()[1].split("/")[0]
|
||||
if not ip.startswith("127.") and not ip.startswith("169.254."):
|
||||
return ip
|
||||
for line in lines:
|
||||
if "inet " in line and "scope global" in line:
|
||||
ip = line.strip().split()[1].split("/")[0]
|
||||
if not ip.startswith("127.") and not ip.startswith("172.17."):
|
||||
return ip
|
||||
break
|
||||
except (FileNotFoundError, Exception):
|
||||
continue
|
||||
|
||||
try:
|
||||
host_ip = socket.gethostbyname(socket.gethostname())
|
||||
if host_ip and not host_ip.startswith("127."):
|
||||
return host_ip
|
||||
except Exception:
|
||||
pass # Hostname resolution failed; fall back to loopback below
|
||||
return "127.0.0.1"
|
||||
|
||||
|
||||
def _detect_rdma_device() -> str:
|
||||
"""Auto-detect a usable RDMA device for single-node loopback testing.
|
||||
|
||||
On DGX machines with 12+ RDMA NICs, only RoCE NICs (with a bound network
|
||||
interface) can reliably do loopback. IB-only NICs often fail cross-NIC
|
||||
QP handshakes, causing intermittent test failures.
|
||||
|
||||
Priority:
|
||||
1. RDMA_DEVICE_NAME / RDMA_DEVICE environment variable
|
||||
2. ibdev2netdev: pick first RoCE device whose netdev is UP
|
||||
3. Mooncake TransferEngine topology (if available)
|
||||
4. Empty string (Mooncake uses all devices — risky on multi-NIC)
|
||||
"""
|
||||
# 1. Explicit env override
|
||||
for env_key in ("RDMA_DEVICE_NAME", "RDMA_DEVICE"):
|
||||
env_val = os.environ.get(env_key, "").strip()
|
||||
if env_val:
|
||||
print(f"[_detect_rdma_device] Using {env_key}={env_val}")
|
||||
return env_val
|
||||
|
||||
# 2. ibdev2netdev: most reliable, no Mooncake dependency
|
||||
# Prefer RoCE (Ethernet-backed) over IB for loopback reliability.
|
||||
try:
|
||||
result = subprocess.run(["ibdev2netdev"], capture_output=True, text=True, timeout=5)
|
||||
if result.returncode == 0:
|
||||
up_lines = [ln.strip() for ln in result.stdout.strip().splitlines() if "(Up)" in ln]
|
||||
# Prefer Ethernet-backed (RoCE) — interface name is NOT ib*
|
||||
for line in up_lines:
|
||||
# Format: "mlx5_2 port 1 ==> enp75s0f0 (Up)"
|
||||
parts = line.split("==>")
|
||||
if len(parts) == 2:
|
||||
iface = parts[1].strip().split()[0] # e.g. "enp75s0f0"
|
||||
if not iface.startswith("ib"):
|
||||
dev = line.split()[0]
|
||||
print(f"[_detect_rdma_device] ibdev2netdev selected RoCE: {dev} ({line})")
|
||||
return dev
|
||||
# Fallback: any Up device (IB)
|
||||
if up_lines:
|
||||
dev = up_lines[0].split()[0]
|
||||
print(f"[_detect_rdma_device] ibdev2netdev selected IB: {dev} ({up_lines[0]})")
|
||||
return dev
|
||||
print("[_detect_rdma_device] ibdev2netdev: no device with (Up) link found")
|
||||
except FileNotFoundError:
|
||||
print("[_detect_rdma_device] ibdev2netdev not found, skipping")
|
||||
except Exception as exc:
|
||||
print(f"[_detect_rdma_device] ibdev2netdev failed: {exc}")
|
||||
|
||||
# 3. Mooncake topology query
|
||||
if TransferEngine is not None:
|
||||
try:
|
||||
import json
|
||||
|
||||
host = get_rdma_host()
|
||||
eng = TransferEngine()
|
||||
if eng.initialize(host, "P2PHANDSHAKE", "rdma", "") == 0:
|
||||
topo_str = getattr(eng, "get_local_topology", lambda: None)()
|
||||
if topo_str:
|
||||
topo = json.loads(topo_str)
|
||||
# Prefer RoCE (IPv4-mapped GID)
|
||||
for name, info in topo.items():
|
||||
if isinstance(info, dict):
|
||||
gid = info.get("gid", "")
|
||||
if gid.startswith("00:00:00:00:00:00:00:00:00:00:ff:ff"):
|
||||
print(f"[_detect_rdma_device] Mooncake topology selected RoCE: {name}")
|
||||
return name
|
||||
# Prefer bonded device
|
||||
for name in topo:
|
||||
if "bond" in name:
|
||||
print(f"[_detect_rdma_device] Mooncake topology selected bond: {name}")
|
||||
return name
|
||||
# First available
|
||||
for name, info in topo.items():
|
||||
if isinstance(info, dict):
|
||||
print(f"[_detect_rdma_device] Mooncake topology fallback: {name}")
|
||||
return name
|
||||
else:
|
||||
print("[_detect_rdma_device] get_local_topology() returned None/empty")
|
||||
except Exception as exc:
|
||||
print(f"[_detect_rdma_device] Mooncake topology query failed: {exc}")
|
||||
|
||||
print("[_detect_rdma_device] WARNING: no device detected, Mooncake will use ALL NICs")
|
||||
return ""
|
||||
|
||||
|
||||
RDMA_HOST = get_rdma_host()
|
||||
RDMA_DEVICE = _detect_rdma_device()
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
import socket
|
||||
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind((RDMA_HOST, 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def _connector_config(
|
||||
zmq_port: int,
|
||||
pool_size: int = 16 * 1024 * 1024,
|
||||
pool_device: str = "cpu",
|
||||
) -> dict:
|
||||
cfg = {
|
||||
"host": RDMA_HOST,
|
||||
"zmq_port": zmq_port,
|
||||
"protocol": "rdma",
|
||||
"memory_pool_size": pool_size,
|
||||
"memory_pool_device": pool_device,
|
||||
}
|
||||
if RDMA_DEVICE:
|
||||
cfg["device_name"] = RDMA_DEVICE
|
||||
return cfg
|
||||
|
||||
|
||||
def _md5(tensor: torch.Tensor) -> str:
|
||||
t = tensor.cpu() if tensor.is_cuda else tensor
|
||||
return hashlib.md5(t.contiguous().view(torch.uint8).numpy().tobytes()).hexdigest()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. Basic functionality (single connector, no RDMA transfer)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
|
||||
class TestBasicConnector(unittest.TestCase):
|
||||
"""Verify connector initialization, put, cleanup, and health check."""
|
||||
|
||||
def test_initialization(self):
|
||||
port = _free_port()
|
||||
with MooncakeTransferEngineConnector(_connector_config(port, pool_size=1024 * 1024)) as c:
|
||||
self.assertNotEqual(c.rpc_port, 0)
|
||||
self.assertEqual(c.pool_size, 1024 * 1024)
|
||||
self.assertTrue(c.pool.is_pinned())
|
||||
health = c.health()
|
||||
self.assertEqual(health["status"], "healthy")
|
||||
|
||||
def test_put_tensor_bytes_object(self):
|
||||
"""Put tensor / bytes / dict and verify metadata."""
|
||||
port = _free_port()
|
||||
with MooncakeTransferEngineConnector(_connector_config(port)) as c:
|
||||
ok, sz, meta = c.put("s0", "s1", "t", torch.randn(100))
|
||||
self.assertTrue(ok)
|
||||
self.assertTrue(meta["is_fast_path"])
|
||||
|
||||
ok, sz, meta = c.put("s0", "s1", "b", b"hello" * 100)
|
||||
self.assertTrue(ok)
|
||||
self.assertTrue(meta["is_fast_path"])
|
||||
|
||||
ok, sz, meta = c.put("s0", "s1", "d", {"k": [1, 2, 3]})
|
||||
self.assertTrue(ok)
|
||||
self.assertFalse(meta["is_fast_path"])
|
||||
|
||||
def test_cleanup_releases_buffer(self):
|
||||
port = _free_port()
|
||||
with MooncakeTransferEngineConnector(_connector_config(port)) as c:
|
||||
c.put("s0", "s1", "r1", torch.randn(100))
|
||||
key = MooncakeTransferEngineConnector._make_key("r1", "s0", "s1")
|
||||
self.assertIn(key, c._local_buffers)
|
||||
c.cleanup("r1", from_stage="s0", to_stage="s1")
|
||||
self.assertNotIn(key, c._local_buffers)
|
||||
|
||||
def test_pool_exhaustion_and_recovery(self):
|
||||
"""Fill pool, verify failure, free, verify recovery."""
|
||||
port = _free_port()
|
||||
with MooncakeTransferEngineConnector(_connector_config(port, pool_size=64 * 1024)) as c:
|
||||
ids = []
|
||||
for i in range(10):
|
||||
ok, _, _ = c.put("s", "s", f"f{i}", torch.randn(1000))
|
||||
if ok:
|
||||
ids.append(f"f{i}")
|
||||
else:
|
||||
break
|
||||
for rid in ids:
|
||||
c.cleanup(rid, from_stage="s", to_stage="s")
|
||||
ok, _, _ = c.put("s", "s", "recovery", torch.randn(1000))
|
||||
self.assertTrue(ok, "Pool recovery failed after cleanup")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. End-to-end RDMA transfer (producer + consumer, single node)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
|
||||
class TestEndToEnd(unittest.TestCase):
|
||||
"""E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload, mixed types."""
|
||||
|
||||
def _pair(self, pool_size=16 * 1024 * 1024):
|
||||
p = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size))
|
||||
c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size))
|
||||
return p, c
|
||||
|
||||
def test_tensor_e2e(self):
|
||||
p, c = self._pair()
|
||||
try:
|
||||
orig = torch.randn(1024, 1024, dtype=torch.float32)
|
||||
ok, sz, meta = p.put("s0", "s1", "t1", orig)
|
||||
self.assertTrue(ok)
|
||||
time.sleep(0.5)
|
||||
result = c.get("s0", "s1", "t1", meta)
|
||||
self.assertIsNotNone(result)
|
||||
buf, _ = result
|
||||
self.assertIsInstance(buf, ManagedBuffer)
|
||||
recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
|
||||
self.assertTrue(torch.equal(recon, orig))
|
||||
buf.release()
|
||||
finally:
|
||||
p.close()
|
||||
c.close()
|
||||
|
||||
def test_bytes_e2e(self):
|
||||
p, c = self._pair()
|
||||
try:
|
||||
orig = b"Hello RDMA! " * 1000
|
||||
ok, _, meta = p.put("s0", "s1", "b1", orig)
|
||||
self.assertTrue(ok)
|
||||
time.sleep(0.5)
|
||||
result = c.get("s0", "s1", "b1", meta)
|
||||
self.assertIsNotNone(result)
|
||||
buf, _ = result
|
||||
recv = buf.to_bytes() if hasattr(buf, "to_bytes") else buf
|
||||
if hasattr(buf, "release"):
|
||||
buf.release()
|
||||
self.assertEqual(recv, orig)
|
||||
finally:
|
||||
p.close()
|
||||
c.close()
|
||||
|
||||
def test_object_e2e(self):
|
||||
p, c = self._pair()
|
||||
try:
|
||||
orig = {"msg": "hello", "nums": [1, 2, 3], "nested": {"k": "v"}}
|
||||
ok, _, meta = p.put("s0", "s1", "o1", orig)
|
||||
self.assertTrue(ok)
|
||||
self.assertFalse(meta["is_fast_path"])
|
||||
time.sleep(1.0)
|
||||
result = c.get("s0", "s1", "o1", meta)
|
||||
self.assertIsNotNone(result)
|
||||
recv, _ = result
|
||||
self.assertEqual(recv, orig)
|
||||
finally:
|
||||
p.close()
|
||||
c.close()
|
||||
|
||||
def test_zero_copy_e2e(self):
|
||||
p, c = self._pair()
|
||||
try:
|
||||
shape = (1024, 1024)
|
||||
nbytes = 1024 * 1024 * 4
|
||||
offset = p.allocator.alloc(nbytes)
|
||||
sbuf = ManagedBuffer(p.allocator, offset, nbytes, p.pool)
|
||||
ref = torch.randn(*shape, dtype=torch.float32)
|
||||
sbuf.as_tensor(dtype=torch.float32, shape=shape).copy_(ref)
|
||||
|
||||
ok, sz, meta = p.put("s0", "s1", "zc", sbuf)
|
||||
self.assertTrue(ok)
|
||||
self.assertTrue(meta.get("is_fast_path"))
|
||||
time.sleep(0.5)
|
||||
|
||||
result = c.get("s0", "s1", "zc", meta)
|
||||
self.assertIsNotNone(result)
|
||||
rbuf, _ = result
|
||||
recon = rbuf.as_tensor(dtype=torch.float32, shape=shape)
|
||||
self.assertTrue(torch.equal(recon, ref))
|
||||
rbuf.release()
|
||||
finally:
|
||||
p.close()
|
||||
c.close()
|
||||
|
||||
def test_large_tensor_100mb(self):
|
||||
"""Transfer ~100 MB tensor and verify MD5 integrity."""
|
||||
p, c = self._pair(pool_size=128 * 1024 * 1024)
|
||||
try:
|
||||
orig = torch.randn(5000, 5000, dtype=torch.float32)
|
||||
md5_send = _md5(orig)
|
||||
ok, _, meta = p.put("s0", "s1", "lg", orig)
|
||||
self.assertTrue(ok)
|
||||
time.sleep(1.0)
|
||||
result = c.get("s0", "s1", "lg", meta)
|
||||
self.assertIsNotNone(result)
|
||||
buf, _ = result
|
||||
recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
|
||||
self.assertEqual(_md5(recon), md5_send)
|
||||
buf.release()
|
||||
finally:
|
||||
p.close()
|
||||
c.close()
|
||||
|
||||
def test_mixed_types_sequential(self):
|
||||
"""Transfer different dtypes and data types in sequence."""
|
||||
p, c = self._pair(pool_size=32 * 1024 * 1024)
|
||||
try:
|
||||
cases = [
|
||||
("f32", torch.randn(500, 500, dtype=torch.float32)),
|
||||
("f16", torch.randn(500, 500, dtype=torch.float16)),
|
||||
("i64", torch.randint(0, 100, (200, 200), dtype=torch.int64)),
|
||||
("byte", b"X" * (1024 * 1024)),
|
||||
("obj", {"key": "value", "list": [1, 2, 3]}),
|
||||
]
|
||||
time.sleep(0.5)
|
||||
for name, data in cases:
|
||||
ok, _, meta = p.put("s0", "s1", name, data)
|
||||
self.assertTrue(ok, f"put failed: {name}")
|
||||
time.sleep(0.3)
|
||||
result = c.get("s0", "s1", name, meta)
|
||||
self.assertIsNotNone(result, f"get failed: {name}")
|
||||
recv, _ = result
|
||||
if isinstance(data, torch.Tensor):
|
||||
self.assertIsInstance(recv, ManagedBuffer)
|
||||
recon = recv.as_tensor(dtype=data.dtype, shape=data.shape)
|
||||
self.assertTrue(torch.equal(recon, data), f"mismatch: {name}")
|
||||
recv.release()
|
||||
elif isinstance(data, bytes):
|
||||
got = recv.to_bytes() if hasattr(recv, "to_bytes") else recv
|
||||
if hasattr(recv, "release"):
|
||||
recv.release()
|
||||
self.assertEqual(got, data)
|
||||
else:
|
||||
self.assertEqual(recv, data)
|
||||
finally:
|
||||
p.close()
|
||||
c.close()
|
||||
|
||||
def test_concurrent_put(self):
|
||||
"""10 concurrent puts should all succeed."""
|
||||
port = _free_port()
|
||||
conn = MooncakeTransferEngineConnector(_connector_config(port, pool_size=64 * 1024 * 1024))
|
||||
errors: list[str] = []
|
||||
lock = threading.Lock()
|
||||
|
||||
def worker(rid, data):
|
||||
try:
|
||||
ok, _, _ = conn.put("s0", "s1", rid, data)
|
||||
if not ok:
|
||||
with lock:
|
||||
errors.append(f"{rid}: put failed")
|
||||
except Exception as e:
|
||||
with lock:
|
||||
errors.append(f"{rid}: {e}")
|
||||
|
||||
with conn:
|
||||
threads = []
|
||||
for i in range(10):
|
||||
t = threading.Thread(target=worker, args=(f"r{i}", torch.randn(256, 256)))
|
||||
threads.append(t)
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
self.assertEqual(len(errors), 0, f"errors: {errors}")
|
||||
|
||||
def test_auto_cleanup(self):
|
||||
"""Producer buffer should be released after consumer get."""
|
||||
p, c = self._pair()
|
||||
try:
|
||||
ok, _, meta = p.put("s0", "s1", "ac", torch.randn(100))
|
||||
self.assertTrue(ok)
|
||||
key = MooncakeTransferEngineConnector._make_key("ac", "s0", "s1")
|
||||
self.assertIn(key, p._local_buffers)
|
||||
time.sleep(0.5)
|
||||
result = c.get("s0", "s1", "ac", meta)
|
||||
self.assertIsNotNone(result)
|
||||
buf, _ = result
|
||||
if isinstance(buf, ManagedBuffer):
|
||||
buf.release()
|
||||
time.sleep(0.3)
|
||||
self.assertNotIn(key, p._local_buffers)
|
||||
finally:
|
||||
p.close()
|
||||
c.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. Lifecycle & resource management
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
|
||||
class TestLifecycle(unittest.TestCase):
|
||||
"""Close, context manager, double-close safety."""
|
||||
|
||||
def test_close_releases_resources(self):
|
||||
c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size=1024 * 1024))
|
||||
c.put("s0", "s1", "x", torch.randn(100))
|
||||
c.close()
|
||||
self.assertTrue(c._stop_event.is_set())
|
||||
self.assertEqual(len(c._local_buffers), 0)
|
||||
|
||||
def test_context_manager(self):
|
||||
with MooncakeTransferEngineConnector(_connector_config(_free_port())) as c:
|
||||
ok, _, _ = c.put("s0", "s1", "ctx", torch.randn(50))
|
||||
self.assertTrue(ok)
|
||||
self.assertTrue(c._stop_event.is_set())
|
||||
|
||||
def test_double_close_safe(self):
|
||||
c = MooncakeTransferEngineConnector(_connector_config(_free_port()))
|
||||
c.close()
|
||||
c.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. GPU memory pool (requires CUDA)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
|
||||
class TestGPUPool(unittest.TestCase):
|
||||
"""GPU memory pool: initialization, put (CPU/GPU tensor), E2E transfer."""
|
||||
|
||||
@staticmethod
|
||||
def _gpu_cfg(port, pool_size=32 * 1024 * 1024):
|
||||
return _connector_config(port, pool_size, pool_device="cuda:0")
|
||||
|
||||
def test_gpu_pool_init(self):
|
||||
with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c:
|
||||
self.assertEqual(c.pool_device, "cuda:0")
|
||||
self.assertTrue(c.pool.is_cuda)
|
||||
|
||||
def test_gpu_pool_put_cpu_and_gpu_tensor(self):
|
||||
with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c:
|
||||
ok, _, meta = c.put("s0", "s1", "h2d", torch.randn(256, 256))
|
||||
self.assertTrue(ok)
|
||||
self.assertTrue(meta["is_fast_path"])
|
||||
|
||||
ok, _, meta = c.put("s0", "s1", "d2d", torch.randn(256, 256, device="cuda:0"))
|
||||
self.assertTrue(ok)
|
||||
self.assertTrue(meta["is_fast_path"])
|
||||
|
||||
def test_gpu_e2e_transfer(self):
|
||||
p = MooncakeTransferEngineConnector(self._gpu_cfg(_free_port()))
|
||||
c = MooncakeTransferEngineConnector(self._gpu_cfg(_free_port()))
|
||||
try:
|
||||
orig = torch.randn(512, 512, dtype=torch.float32, device="cuda:0")
|
||||
ok, _, meta = p.put("s0", "s1", "ge", orig)
|
||||
self.assertTrue(ok)
|
||||
time.sleep(0.5)
|
||||
result = c.get("s0", "s1", "ge", meta)
|
||||
self.assertIsNotNone(result)
|
||||
buf, _ = result
|
||||
recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
|
||||
self.assertTrue(recon.is_cuda)
|
||||
self.assertTrue(torch.equal(recon.cpu(), orig.cpu()))
|
||||
buf.release()
|
||||
finally:
|
||||
p.close()
|
||||
c.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. Stress / Correctness tests (marked slow, skipped in quick CI)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
|
||||
class TestStressCorrectness(unittest.TestCase):
|
||||
"""
|
||||
Slow but high-value regression tests: concurrent put+get with data
|
||||
integrity, large payload, edge cases, and sustained stress.
|
||||
"""
|
||||
|
||||
def _pair(self, pool_size=64 * 1024 * 1024):
|
||||
p = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size))
|
||||
c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size))
|
||||
return p, c
|
||||
|
||||
# -- Concurrent put + get with data integrity --
|
||||
|
||||
def test_concurrent_put_get_integrity(self):
|
||||
"""
|
||||
10 threads each put a unique tensor; a single consumer thread gets all
|
||||
of them and verifies MD5 matches.
|
||||
"""
|
||||
p, c = self._pair(pool_size=128 * 1024 * 1024)
|
||||
num_workers = 10
|
||||
errors: list[str] = []
|
||||
lock = threading.Lock()
|
||||
md5_sent: dict[str, str] = {}
|
||||
meta_map: dict[str, dict] = {}
|
||||
|
||||
def producer_worker(rid: str, tensor: torch.Tensor):
|
||||
try:
|
||||
md5_val = _md5(tensor)
|
||||
with lock:
|
||||
md5_sent[rid] = md5_val
|
||||
ok, _, meta = p.put("s0", "s1", rid, tensor)
|
||||
if not ok:
|
||||
with lock:
|
||||
errors.append(f"{rid}: put failed")
|
||||
return
|
||||
with lock:
|
||||
meta_map[rid] = meta
|
||||
except Exception as e:
|
||||
with lock:
|
||||
errors.append(f"{rid}: put exception {e}")
|
||||
|
||||
try:
|
||||
threads = []
|
||||
for i in range(num_workers):
|
||||
t = threading.Thread(
|
||||
target=producer_worker,
|
||||
args=(f"cg{i}", torch.randn(512, 512, dtype=torch.float32)),
|
||||
)
|
||||
threads.append(t)
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
self.assertEqual(len(errors), 0, f"Producer errors: {errors}")
|
||||
self.assertEqual(len(meta_map), num_workers)
|
||||
|
||||
time.sleep(1.0)
|
||||
|
||||
# Consumer: get each tensor and verify integrity
|
||||
for rid, meta in meta_map.items():
|
||||
result = c.get("s0", "s1", rid, meta)
|
||||
self.assertIsNotNone(result, f"get returned None for {rid}")
|
||||
buf, _ = result
|
||||
self.assertIsInstance(buf, ManagedBuffer, f"{rid} not ManagedBuffer")
|
||||
recon = buf.as_tensor(dtype=torch.float32, shape=(512, 512))
|
||||
recv_md5 = _md5(recon)
|
||||
self.assertEqual(recv_md5, md5_sent[rid], f"MD5 mismatch for {rid}")
|
||||
buf.release()
|
||||
finally:
|
||||
p.close()
|
||||
c.close()
|
||||
|
||||
def test_concurrent_put_get_threaded_both_sides(self):
|
||||
"""
|
||||
True overlapping concurrency: producer threads put items one by
|
||||
one and signal the consumer via a queue; the consumer thread
|
||||
starts getting as soon as each item is available, so put and get
|
||||
execute in overlapping time windows.
|
||||
"""
|
||||
import queue
|
||||
|
||||
p, c = self._pair(pool_size=128 * 1024 * 1024)
|
||||
num_items = 8
|
||||
errors: list[str] = []
|
||||
lock = threading.Lock()
|
||||
md5_sent: dict[str, str] = {}
|
||||
item_queue: queue.Queue[tuple[str, dict] | None] = queue.Queue()
|
||||
|
||||
def producer_worker(rid, tensor):
|
||||
try:
|
||||
with lock:
|
||||
md5_sent[rid] = _md5(tensor)
|
||||
ok, _, meta = p.put("s0", "s1", rid, tensor)
|
||||
if not ok:
|
||||
with lock:
|
||||
errors.append(f"{rid}: put failed")
|
||||
return
|
||||
# Signal consumer immediately after put
|
||||
item_queue.put((rid, meta))
|
||||
except Exception as e:
|
||||
with lock:
|
||||
errors.append(f"{rid}: {e}")
|
||||
|
||||
consumed_count = [0] # mutable container for thread access
|
||||
|
||||
def consumer_thread():
|
||||
while consumed_count[0] < num_items:
|
||||
try:
|
||||
item = item_queue.get(timeout=30)
|
||||
except queue.Empty:
|
||||
with lock:
|
||||
errors.append(f"consumer: queue.Empty after 30s, consumed {consumed_count[0]}/{num_items}")
|
||||
break
|
||||
if item is None:
|
||||
break
|
||||
rid, meta = item
|
||||
try:
|
||||
# Small delay to let RDMA transfer initiate, but
|
||||
# overlap with other producers that are still putting
|
||||
time.sleep(0.1)
|
||||
result = c.get("s0", "s1", rid, meta)
|
||||
if result is None:
|
||||
with lock:
|
||||
errors.append(f"{rid}: get returned None")
|
||||
continue
|
||||
buf, _ = result
|
||||
if isinstance(buf, ManagedBuffer):
|
||||
recon = buf.as_tensor(dtype=torch.float32, shape=(256, 256))
|
||||
if _md5(recon) != md5_sent.get(rid, ""):
|
||||
with lock:
|
||||
errors.append(f"{rid}: MD5 mismatch")
|
||||
buf.release()
|
||||
except Exception as e:
|
||||
with lock:
|
||||
errors.append(f"{rid}: get exception {e}")
|
||||
finally:
|
||||
consumed_count[0] += 1
|
||||
|
||||
try:
|
||||
# Start consumer first so it is ready to get
|
||||
ct = threading.Thread(target=consumer_thread)
|
||||
ct.start()
|
||||
# Launch producers with staggered starts for overlap
|
||||
threads = []
|
||||
for i in range(num_items):
|
||||
t = threading.Thread(
|
||||
target=producer_worker,
|
||||
args=(f"bi{i}", torch.randn(256, 256, dtype=torch.float32)),
|
||||
)
|
||||
threads.append(t)
|
||||
t.start()
|
||||
time.sleep(0.05) # slight stagger so puts overlap with gets
|
||||
for t in threads:
|
||||
t.join()
|
||||
ct.join(timeout=60)
|
||||
self.assertEqual(len(errors), 0, f"errors: {errors}")
|
||||
self.assertEqual(
|
||||
consumed_count[0],
|
||||
num_items,
|
||||
f"Consumer only processed {consumed_count[0]}/{num_items} items",
|
||||
)
|
||||
finally:
|
||||
p.close()
|
||||
c.close()
|
||||
|
||||
# -- Edge cases --
|
||||
|
||||
def test_small_tensor_1_element(self):
|
||||
"""Transfer a single-element tensor."""
|
||||
p, c = self._pair()
|
||||
try:
|
||||
orig = torch.tensor([3.14], dtype=torch.float32)
|
||||
ok, _, meta = p.put("s0", "s1", "tiny", orig)
|
||||
self.assertTrue(ok)
|
||||
time.sleep(0.5)
|
||||
result = c.get("s0", "s1", "tiny", meta)
|
||||
self.assertIsNotNone(result)
|
||||
buf, _ = result
|
||||
recon = buf.as_tensor(dtype=torch.float32, shape=(1,))
|
||||
self.assertTrue(torch.equal(recon, orig))
|
||||
buf.release()
|
||||
finally:
|
||||
p.close()
|
||||
c.close()
|
||||
|
||||
def test_empty_bytes_rejected(self):
|
||||
"""Connector should gracefully reject empty bytes payload."""
|
||||
port = _free_port()
|
||||
with MooncakeTransferEngineConnector(_connector_config(port, pool_size=8 * 1024 * 1024)) as c:
|
||||
ok, sz, meta = c.put("s0", "s1", "empty_b", b"")
|
||||
self.assertFalse(ok, "Empty bytes should be rejected by connector")
|
||||
|
||||
# -- Large payload stress (500 MB) --
|
||||
|
||||
def test_large_tensor_500mb(self):
|
||||
"""Transfer ~500 MB tensor and verify MD5 integrity."""
|
||||
p, c = self._pair(pool_size=600 * 1024 * 1024)
|
||||
try:
|
||||
# ~500 MB: 11180 * 11180 * 4 bytes ≈ 500 MB
|
||||
side = 11180
|
||||
orig = torch.randn(side, side, dtype=torch.float32)
|
||||
md5_send = _md5(orig)
|
||||
ok, _, meta = p.put("s0", "s1", "500mb", orig)
|
||||
self.assertTrue(ok)
|
||||
time.sleep(2.0)
|
||||
result = c.get("s0", "s1", "500mb", meta)
|
||||
self.assertIsNotNone(result)
|
||||
buf, _ = result
|
||||
recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
|
||||
self.assertEqual(_md5(recon), md5_send)
|
||||
buf.release()
|
||||
finally:
|
||||
p.close()
|
||||
c.close()
|
||||
|
||||
# -- Allocator stress under connector --
|
||||
|
||||
def test_rapid_alloc_free_cycle(self):
|
||||
"""Put + cleanup in tight loop to stress allocator under real connector."""
|
||||
port = _free_port()
|
||||
with MooncakeTransferEngineConnector(_connector_config(port, pool_size=8 * 1024 * 1024)) as c:
|
||||
for i in range(50):
|
||||
rid = f"cycle_{i}"
|
||||
ok, _, _ = c.put("s0", "s1", rid, torch.randn(256, 256))
|
||||
self.assertTrue(ok, f"Put failed at iteration {i}")
|
||||
c.cleanup(rid, from_stage="s0", to_stage="s1")
|
||||
# After all cycles, pool should be fully recovered
|
||||
ok, _, _ = c.put("s0", "s1", "final", torch.randn(1024, 1024))
|
||||
self.assertTrue(ok, "Pool not fully recovered after rapid cycles")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -3,10 +3,6 @@ E2E Online tests for Qwen3-Omni model with video input and audio output.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
@@ -21,6 +17,10 @@ from tests.conftest import (
|
||||
from tests.utils import hardware_test
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
|
||||
|
||||
|
||||
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
|
||||
|
||||
|
||||
@@ -114,7 +114,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
"messages": messages,
|
||||
"stream": True,
|
||||
"key_words": {
|
||||
"audio": ["water", "chirping"],
|
||||
"audio": ["water", "chirping", "crackling"],
|
||||
"image": ["square", "quadrate"],
|
||||
},
|
||||
}
|
||||
|
||||
@@ -86,6 +86,7 @@ stage_args:
|
||||
max_num_batched_tokens: 100000
|
||||
hf_config_name: thinker_config
|
||||
async_scheduling: false
|
||||
load_format: dummy
|
||||
engine_input_source: [1]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
|
||||
final_output: true
|
||||
|
||||
@@ -1283,6 +1283,21 @@ class Bagel(nn.Module):
|
||||
|
||||
return generation_input
|
||||
|
||||
@staticmethod
|
||||
def _merge_naive_caches(caches: list) -> NaiveCache:
|
||||
"""Merge multiple NaiveCache objects by concatenating KV tensors per layer."""
|
||||
if not caches:
|
||||
# Handle empty list case gracefully if desired,
|
||||
# though original code also crashed on this.
|
||||
return NaiveCache(0)
|
||||
|
||||
num_layers = len(caches[0].key_cache)
|
||||
merged = NaiveCache(num_layers)
|
||||
for layer_idx in range(num_layers):
|
||||
merged.key_cache[layer_idx] = torch.cat([c.key_cache[layer_idx] for c in caches], dim=0)
|
||||
merged.value_cache[layer_idx] = torch.cat([c.value_cache[layer_idx] for c in caches], dim=0)
|
||||
return merged
|
||||
|
||||
def generate_image(
|
||||
self,
|
||||
packed_text_ids: torch.LongTensor,
|
||||
@@ -1324,6 +1339,64 @@ class Bagel(nn.Module):
|
||||
dts = timesteps[:-1] - timesteps[1:]
|
||||
timesteps = timesteps[:-1]
|
||||
|
||||
# ── Pre-compute batched CFG state (merged caches + indices) ──
|
||||
use_cfg_text = cfg_text_scale > 1.0
|
||||
use_cfg_img = cfg_img_scale > 1.0
|
||||
cfg_batched = None
|
||||
|
||||
if use_cfg_text:
|
||||
seq_len = int(packed_seqlens.sum())
|
||||
|
||||
# Branch 0: main (gen_context), always present
|
||||
branches_qi = [packed_indexes]
|
||||
branches_kvi = [packed_key_value_indexes]
|
||||
branches_kvl = [key_values_lens]
|
||||
branches_pid = [packed_position_ids]
|
||||
branches_cache = [past_key_values]
|
||||
|
||||
# Branch 1: cfg_text (unconditional text), always present when use_cfg_text
|
||||
branches_qi.append(cfg_text_packed_query_indexes)
|
||||
branches_kvi.append(cfg_text_packed_key_value_indexes)
|
||||
branches_kvl.append(cfg_text_key_values_lens)
|
||||
branches_pid.append(cfg_text_packed_position_ids)
|
||||
branches_cache.append(cfg_text_past_key_values)
|
||||
|
||||
# Branch 2: cfg_img (text-only, no image), optional
|
||||
if use_cfg_img:
|
||||
branches_qi.append(cfg_img_packed_query_indexes)
|
||||
branches_kvi.append(cfg_img_packed_key_value_indexes)
|
||||
branches_kvl.append(cfg_img_key_values_lens)
|
||||
branches_pid.append(cfg_img_packed_position_ids)
|
||||
branches_cache.append(cfg_img_past_key_values)
|
||||
|
||||
num_branches = len(branches_cache)
|
||||
|
||||
# Compute per-branch offsets in the merged KV+Q attention tensor
|
||||
merged_offsets = [0]
|
||||
for b_idx in range(num_branches):
|
||||
merged_offsets.append(merged_offsets[-1] + int(branches_kvl[b_idx].sum()) + seq_len)
|
||||
|
||||
cfg_batched = {
|
||||
"num_branches": num_branches,
|
||||
"seq_len": seq_len,
|
||||
"batched_query_lens": packed_seqlens.repeat(num_branches),
|
||||
"batched_position_ids": torch.cat(branches_pid),
|
||||
"batched_kv_lens": torch.cat(branches_kvl),
|
||||
"batched_query_indexes": torch.cat(
|
||||
[qi + merged_offsets[b_idx] for b_idx, qi in enumerate(branches_qi)]
|
||||
),
|
||||
"batched_kv_indexes": torch.cat(
|
||||
[kvi + merged_offsets[b_idx] for b_idx, kvi in enumerate(branches_kvi)]
|
||||
),
|
||||
"batched_text_indexes": torch.cat(
|
||||
[packed_text_indexes + b_idx * seq_len for b_idx in range(num_branches)]
|
||||
),
|
||||
"batched_vae_indexes": torch.cat(
|
||||
[packed_vae_token_indexes + b_idx * seq_len for b_idx in range(num_branches)]
|
||||
),
|
||||
"merged_cache": self._merge_naive_caches(branches_cache),
|
||||
}
|
||||
|
||||
for i, t in enumerate(timesteps):
|
||||
timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
|
||||
if t > cfg_interval[0] and t <= cfg_interval[1]:
|
||||
@@ -1347,24 +1420,13 @@ class Bagel(nn.Module):
|
||||
packed_key_value_indexes=packed_key_value_indexes,
|
||||
cfg_renorm_min=cfg_renorm_min,
|
||||
cfg_renorm_type=cfg_renorm_type,
|
||||
# cfg_text
|
||||
cfg_text_scale=cfg_text_scale_,
|
||||
cfg_text_packed_position_ids=cfg_text_packed_position_ids,
|
||||
cfg_text_packed_query_indexes=cfg_text_packed_query_indexes,
|
||||
cfg_text_key_values_lens=cfg_text_key_values_lens,
|
||||
cfg_text_past_key_values=cfg_text_past_key_values,
|
||||
cfg_text_packed_key_value_indexes=cfg_text_packed_key_value_indexes,
|
||||
# cfg_img
|
||||
cfg_img_scale=cfg_img_scale_,
|
||||
cfg_img_packed_position_ids=cfg_img_packed_position_ids,
|
||||
cfg_img_packed_query_indexes=cfg_img_packed_query_indexes,
|
||||
cfg_img_key_values_lens=cfg_img_key_values_lens,
|
||||
cfg_img_past_key_values=cfg_img_past_key_values,
|
||||
cfg_img_packed_key_value_indexes=cfg_img_packed_key_value_indexes,
|
||||
cfg_type=cfg_type,
|
||||
cfg_batched=cfg_batched,
|
||||
)
|
||||
|
||||
x_t = x_t - v_t.to(x_t.device) * dts[i] # velocity pointing from data to noise
|
||||
|
||||
unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
|
||||
return unpacked_latent
|
||||
|
||||
@@ -1384,22 +1446,11 @@ class Bagel(nn.Module):
|
||||
packed_key_value_indexes: torch.LongTensor,
|
||||
cfg_renorm_min: float = 0.0,
|
||||
cfg_renorm_type: str = "global",
|
||||
# cfg_text
|
||||
cfg_text_scale: float = 1.0,
|
||||
cfg_text_packed_position_ids: torch.LongTensor | None = None,
|
||||
cfg_text_packed_query_indexes: torch.LongTensor | None = None,
|
||||
cfg_text_key_values_lens: torch.Tensor | None = None,
|
||||
cfg_text_past_key_values: NaiveCache | None = None,
|
||||
cfg_text_packed_key_value_indexes: torch.LongTensor | None = None,
|
||||
# cfg_img
|
||||
cfg_img_scale: float = 1.0,
|
||||
cfg_img_packed_position_ids: torch.LongTensor | None = None,
|
||||
cfg_img_packed_query_indexes: torch.LongTensor | None = None,
|
||||
cfg_img_key_values_lens: torch.Tensor | None = None,
|
||||
cfg_img_past_key_values: NaiveCache | None = None,
|
||||
cfg_img_packed_key_value_indexes: torch.LongTensor | None = None,
|
||||
cfg_type: str = "parallel",
|
||||
cfg_batched: dict | None = None,
|
||||
):
|
||||
# Build query sequence (identical for all CFG branches)
|
||||
packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids)
|
||||
packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size))
|
||||
packed_sequence[packed_text_indexes] = packed_text_embedding
|
||||
@@ -1414,60 +1465,76 @@ class Bagel(nn.Module):
|
||||
|
||||
extra_inputs = {}
|
||||
if self.use_moe:
|
||||
extra_inputs = {
|
||||
"mode": "gen",
|
||||
"packed_vae_token_indexes": packed_vae_token_indexes,
|
||||
"packed_text_indexes": packed_text_indexes,
|
||||
}
|
||||
extra_inputs["mode"] = "gen"
|
||||
|
||||
output = self.language_model.forward(
|
||||
packed_query_sequence=packed_sequence,
|
||||
query_lens=packed_seqlens,
|
||||
packed_query_position_ids=packed_position_ids,
|
||||
packed_query_indexes=packed_indexes,
|
||||
past_key_values=past_key_values,
|
||||
key_values_lens=key_values_lens,
|
||||
packed_key_value_indexes=packed_key_value_indexes,
|
||||
update_past_key_values=False,
|
||||
is_causal=False,
|
||||
**extra_inputs,
|
||||
)
|
||||
v_t = self.llm2vae(output.packed_query_sequence)
|
||||
v_t = v_t[packed_vae_token_indexes]
|
||||
use_cfg = cfg_text_scale > 1.0
|
||||
cfg_text_v_t = None
|
||||
cfg_img_v_t = None
|
||||
|
||||
if cfg_text_scale > 1.0:
|
||||
cfg_text_output = self.language_model.forward(
|
||||
packed_query_sequence=packed_sequence,
|
||||
query_lens=packed_seqlens,
|
||||
packed_query_position_ids=cfg_text_packed_position_ids,
|
||||
packed_query_indexes=cfg_text_packed_query_indexes,
|
||||
past_key_values=cfg_text_past_key_values,
|
||||
key_values_lens=cfg_text_key_values_lens,
|
||||
packed_key_value_indexes=cfg_text_packed_key_value_indexes,
|
||||
if use_cfg and cfg_batched is not None:
|
||||
# ── Batched CFG: single LLM forward for all branches ──
|
||||
seq_len = cfg_batched["seq_len"]
|
||||
num_branches = cfg_batched["num_branches"]
|
||||
|
||||
batched_sequence = packed_sequence.repeat(num_branches, 1)
|
||||
|
||||
if self.use_moe:
|
||||
extra_inputs["packed_text_indexes"] = cfg_batched["batched_text_indexes"]
|
||||
extra_inputs["packed_vae_token_indexes"] = cfg_batched["batched_vae_indexes"]
|
||||
|
||||
output = self.language_model.forward(
|
||||
packed_query_sequence=batched_sequence,
|
||||
query_lens=cfg_batched["batched_query_lens"],
|
||||
packed_query_position_ids=cfg_batched["batched_position_ids"],
|
||||
packed_query_indexes=cfg_batched["batched_query_indexes"],
|
||||
past_key_values=cfg_batched["merged_cache"],
|
||||
key_values_lens=cfg_batched["batched_kv_lens"],
|
||||
packed_key_value_indexes=cfg_batched["batched_kv_indexes"],
|
||||
update_past_key_values=False,
|
||||
is_causal=False,
|
||||
**extra_inputs,
|
||||
)
|
||||
cfg_text_v_t = self.llm2vae(cfg_text_output.packed_query_sequence)
|
||||
cfg_text_v_t = cfg_text_v_t[packed_vae_token_indexes]
|
||||
|
||||
if cfg_img_scale > 1.0:
|
||||
cfg_img_output = self.language_model.forward(
|
||||
# Extract per-branch velocities from batched output
|
||||
all_hidden = output.packed_query_sequence
|
||||
assert all_hidden.shape[0] == seq_len * num_branches, (
|
||||
f"Expected packed sequence length {seq_len * num_branches}, but got {all_hidden.shape[0]}"
|
||||
)
|
||||
|
||||
v_t = self.llm2vae(all_hidden[:seq_len])[packed_vae_token_indexes]
|
||||
|
||||
branch_idx = 1
|
||||
cfg_text_v_t = self.llm2vae(all_hidden[branch_idx * seq_len : (branch_idx + 1) * seq_len])[
|
||||
packed_vae_token_indexes
|
||||
]
|
||||
branch_idx += 1
|
||||
if cfg_img_scale > 1.0:
|
||||
cfg_img_v_t = self.llm2vae(all_hidden[branch_idx * seq_len : (branch_idx + 1) * seq_len])[
|
||||
packed_vae_token_indexes
|
||||
]
|
||||
else:
|
||||
# ── Single forward (no CFG or outside cfg_interval) ──
|
||||
if self.use_moe:
|
||||
extra_inputs["packed_vae_token_indexes"] = packed_vae_token_indexes
|
||||
extra_inputs["packed_text_indexes"] = packed_text_indexes
|
||||
|
||||
output = self.language_model.forward(
|
||||
packed_query_sequence=packed_sequence,
|
||||
query_lens=packed_seqlens,
|
||||
packed_query_position_ids=cfg_img_packed_position_ids,
|
||||
packed_query_indexes=cfg_img_packed_query_indexes,
|
||||
past_key_values=cfg_img_past_key_values,
|
||||
key_values_lens=cfg_img_key_values_lens,
|
||||
packed_key_value_indexes=cfg_img_packed_key_value_indexes,
|
||||
packed_query_position_ids=packed_position_ids,
|
||||
packed_query_indexes=packed_indexes,
|
||||
past_key_values=past_key_values,
|
||||
key_values_lens=key_values_lens,
|
||||
packed_key_value_indexes=packed_key_value_indexes,
|
||||
update_past_key_values=False,
|
||||
is_causal=False,
|
||||
**extra_inputs,
|
||||
)
|
||||
cfg_img_v_t = self.llm2vae(cfg_img_output.packed_query_sequence)
|
||||
cfg_img_v_t = cfg_img_v_t[packed_vae_token_indexes]
|
||||
v_t = self.llm2vae(output.packed_query_sequence)
|
||||
v_t = v_t[packed_vae_token_indexes]
|
||||
|
||||
if cfg_text_scale > 1.0:
|
||||
# ── CFG combination ──
|
||||
if use_cfg:
|
||||
if cfg_renorm_type == "text_channel":
|
||||
v_t_text_ = cfg_text_v_t + cfg_text_scale * (v_t - cfg_text_v_t)
|
||||
norm_v_t = torch.norm(v_t, dim=-1, keepdim=True)
|
||||
@@ -1497,8 +1564,5 @@ class Bagel(nn.Module):
|
||||
raise NotImplementedError(f"{cfg_renorm_type} is not supported")
|
||||
scale = (norm_v_t / (norm_v_t_ + 1e-8)).clamp(min=cfg_renorm_min, max=1.0)
|
||||
v_t = v_t_ * scale
|
||||
else:
|
||||
# No CFG
|
||||
pass
|
||||
|
||||
return v_t
|
||||
@@ -4,6 +4,8 @@
|
||||
from .omni_connectors import (
|
||||
ConnectorSpec,
|
||||
MooncakeConnector,
|
||||
MooncakeStoreConnector,
|
||||
MooncakeTransferEngineConnector,
|
||||
OmniConnectorBase,
|
||||
OmniConnectorFactory,
|
||||
OmniTransferConfig,
|
||||
@@ -19,7 +21,9 @@ __all__ = [
|
||||
# Connectors
|
||||
"OmniConnectorBase",
|
||||
"OmniConnectorFactory",
|
||||
"MooncakeConnector",
|
||||
"MooncakeConnector", # compat alias
|
||||
"MooncakeStoreConnector",
|
||||
"MooncakeTransferEngineConnector",
|
||||
"SharedMemoryConnector",
|
||||
"YuanrongConnector",
|
||||
# Utilities
|
||||
|
||||
@@ -2,9 +2,14 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from .connectors.base import OmniConnectorBase
|
||||
from .connectors.mooncake_connector import MooncakeConnector
|
||||
from .connectors.mooncake_store_connector import MooncakeStoreConnector
|
||||
from .connectors.shm_connector import SharedMemoryConnector
|
||||
from .connectors.yuanrong_connector import YuanrongConnector
|
||||
|
||||
try:
|
||||
from .connectors.mooncake_transfer_engine_connector import MooncakeTransferEngineConnector
|
||||
except ImportError:
|
||||
MooncakeTransferEngineConnector = None # RDMA deps (msgspec/zmq/mooncake) not installed
|
||||
from .factory import OmniConnectorFactory
|
||||
from .utils.config import ConnectorSpec, OmniTransferConfig
|
||||
from .utils.initialization import (
|
||||
@@ -16,6 +21,10 @@ from .utils.initialization import (
|
||||
load_omni_transfer_config,
|
||||
)
|
||||
|
||||
# Backward-compatible alias: MooncakeConnector was renamed to MooncakeStoreConnector.
|
||||
# Keep this alias for at least one release cycle.
|
||||
MooncakeConnector = MooncakeStoreConnector
|
||||
|
||||
__all__ = [
|
||||
# Config
|
||||
"ConnectorSpec",
|
||||
@@ -25,7 +34,9 @@ __all__ = [
|
||||
# Factory
|
||||
"OmniConnectorFactory",
|
||||
# Specific implementations
|
||||
"MooncakeConnector",
|
||||
"MooncakeConnector", # compat alias → MooncakeStoreConnector
|
||||
"MooncakeStoreConnector",
|
||||
"MooncakeTransferEngineConnector",
|
||||
"SharedMemoryConnector",
|
||||
"YuanrongConnector",
|
||||
# Utilities
|
||||
|
||||
@@ -12,6 +12,11 @@ logger = get_connector_logger(__name__)
|
||||
class OmniConnectorBase(ABC):
|
||||
"""Base class for all OmniConnectors."""
|
||||
|
||||
# Whether the connector can handle raw bytes/torch.Tensor natively
|
||||
# without going through OmniSerializer. Connectors that copy raw
|
||||
# payloads directly (e.g. RDMA) should override this to True.
|
||||
supports_raw_data: bool = False
|
||||
|
||||
@abstractmethod
|
||||
def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]:
|
||||
"""Store Python object, internal serialization handled by connector.
|
||||
@@ -52,6 +57,29 @@ class OmniConnectorBase(ABC):
|
||||
"""Return health status and metrics."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def close(self) -> None:
|
||||
"""Release resources held by this connector.
|
||||
|
||||
Subclasses must implement this to clean up transport-specific
|
||||
resources (connections, memory pools, threads, etc.).
|
||||
Implementations should be idempotent (safe to call multiple times).
|
||||
"""
|
||||
pass
|
||||
|
||||
# --- Default resource-management protocol ---
|
||||
# Subclasses get context-manager and destructor support for free;
|
||||
# they only need to implement close().
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.close()
|
||||
|
||||
@staticmethod
|
||||
def serialize_obj(obj: Any) -> bytes:
|
||||
"""Serialize a Python object to bytes using centralized serializer."""
|
||||
@@ -65,3 +93,12 @@ class OmniConnectorBase(ABC):
|
||||
from ..utils.serialization import OmniSerializer
|
||||
|
||||
return OmniSerializer.deserialize(data)
|
||||
|
||||
@staticmethod
|
||||
def _make_key(key: str, from_stage: str, to_stage: str, separator: str = "@") -> str:
|
||||
"""Generate internal key with stage routing info.
|
||||
|
||||
Default format: ``{key}@{from_stage}_{to_stage}``.
|
||||
Connectors with different key conventions can override this method.
|
||||
"""
|
||||
return f"{key}{separator}{from_stage}_{to_stage}"
|
||||
+29
-18
@@ -19,7 +19,7 @@ except ImportError:
|
||||
ReplicateConfig = None
|
||||
|
||||
|
||||
class MooncakeConnector(OmniConnectorBase):
|
||||
class MooncakeStoreConnector(OmniConnectorBase):
|
||||
"""Mooncake-based distributed connector for OmniConnector."""
|
||||
|
||||
def __init__(self, config: dict[str, Any]):
|
||||
@@ -51,10 +51,6 @@ class MooncakeConnector(OmniConnectorBase):
|
||||
|
||||
self._init_store()
|
||||
|
||||
def _make_key(self, rid: str, from_stage: str, to_stage: str) -> str:
|
||||
"""Generate store key for request between stages."""
|
||||
return f"{rid}/{from_stage}_{to_stage}"
|
||||
|
||||
def _init_store(self):
|
||||
"""Initialize Mooncake store."""
|
||||
try:
|
||||
@@ -67,7 +63,7 @@ class MooncakeConnector(OmniConnectorBase):
|
||||
|
||||
self.pin = ReplicateConfig()
|
||||
self.pin.with_soft_pin = True
|
||||
logger.info("MooncakeConnector initialized successfully")
|
||||
logger.info("MooncakeStoreConnector initialized successfully")
|
||||
except Exception as e:
|
||||
logger.error("Failed to initialize Mooncake store: %s", e)
|
||||
raise
|
||||
@@ -88,7 +84,7 @@ class MooncakeConnector(OmniConnectorBase):
|
||||
self._metrics["bytes_transferred"] += len(serialized_data)
|
||||
|
||||
logger.debug(
|
||||
"MooncakeConnector: stored %s (%s -> %s) %d bytes",
|
||||
"MooncakeStoreConnector: stored %s (%s -> %s) %d bytes",
|
||||
key,
|
||||
from_stage,
|
||||
to_stage,
|
||||
@@ -98,7 +94,7 @@ class MooncakeConnector(OmniConnectorBase):
|
||||
|
||||
except Exception as e:
|
||||
self._metrics["errors"] += 1
|
||||
logger.error("MooncakeConnector put failed: %s", e)
|
||||
logger.error("MooncakeStoreConnector put failed: %s", e)
|
||||
return False, 0, None
|
||||
|
||||
def get(
|
||||
@@ -108,35 +104,47 @@ class MooncakeConnector(OmniConnectorBase):
|
||||
logger.error("Store not initialized")
|
||||
return None
|
||||
|
||||
import time as _time
|
||||
|
||||
_t0 = _time.perf_counter()
|
||||
|
||||
retries = 20
|
||||
sleep_s = 0.05
|
||||
key = self._make_key(get_key, from_stage, to_stage)
|
||||
|
||||
for attempt in range(retries):
|
||||
try:
|
||||
_t_fetch_start = _time.perf_counter()
|
||||
raw_data = self.store.get(key)
|
||||
_t_fetch_end = _time.perf_counter()
|
||||
|
||||
if raw_data:
|
||||
_fetch_ms = (_t_fetch_end - _t_fetch_start) * 1000
|
||||
|
||||
_t_deser_start = _time.perf_counter()
|
||||
data = self.deserialize_obj(raw_data)
|
||||
_t_deser_end = _time.perf_counter()
|
||||
_deser_ms = (_t_deser_end - _t_deser_start) * 1000
|
||||
|
||||
self._metrics["gets"] += 1
|
||||
payload_size = len(raw_data)
|
||||
logger.debug(
|
||||
"MooncakeConnector: retrieved %s (%s -> %s) %d bytes",
|
||||
key,
|
||||
from_stage,
|
||||
to_stage,
|
||||
payload_size,
|
||||
|
||||
_total_ms = (_t_deser_end - _t0) * 1000
|
||||
_mbps = (payload_size / 1024 / 1024) / (_total_ms / 1000) if _total_ms > 0 else 0
|
||||
logger.info(
|
||||
f"[TCP GET] {get_key}: fetch={_fetch_ms:.1f}ms, deser={_deser_ms:.1f}ms, "
|
||||
f"total={_total_ms:.1f}ms, {payload_size} bytes, {_mbps:.1f} MB/s"
|
||||
)
|
||||
return data, payload_size
|
||||
|
||||
except Exception as e:
|
||||
logger.debug("MooncakeConnector get attempt %s failed: %s", attempt, e)
|
||||
logger.debug("MooncakeStoreConnector get attempt %s failed: %s", attempt, e)
|
||||
|
||||
if attempt < retries - 1:
|
||||
time.sleep(sleep_s)
|
||||
|
||||
self._metrics["timeouts"] += 1
|
||||
logger.warning("MooncakeConnector: timeout waiting for %s", key)
|
||||
logger.warning("MooncakeStoreConnector: timeout waiting for %s", key)
|
||||
return None
|
||||
|
||||
def cleanup(self, request_id: str) -> None:
|
||||
@@ -145,7 +153,7 @@ class MooncakeConnector(OmniConnectorBase):
|
||||
|
||||
# Note: Mooncake doesn't have explicit delete, data will be garbage collected
|
||||
# We could implement a cleanup mechanism by storing deletion markers
|
||||
logger.debug("MooncakeConnector: cleanup requested for %s (no-op)", request_id)
|
||||
logger.debug("MooncakeStoreConnector: cleanup requested for %s (no-op)", request_id)
|
||||
|
||||
def health(self) -> dict[str, Any]:
|
||||
if not self.store:
|
||||
@@ -156,6 +164,9 @@ class MooncakeConnector(OmniConnectorBase):
|
||||
"host": self.host,
|
||||
"metadata_server": self.metadata,
|
||||
"master": self.master,
|
||||
"protocol": None,
|
||||
"pool_device": None,
|
||||
"pool_size": 0,
|
||||
**self._metrics,
|
||||
}
|
||||
|
||||
@@ -165,6 +176,6 @@ class MooncakeConnector(OmniConnectorBase):
|
||||
try:
|
||||
self.store.close()
|
||||
self.store = None
|
||||
logger.info("MooncakeConnector closed")
|
||||
logger.info("MooncakeStoreConnector closed")
|
||||
except Exception as e:
|
||||
logger.error("Error closing Mooncake store: %s", e)
|
||||
+1202
File diff suppressed because it is too large.
Load diff
@@ -162,5 +162,8 @@ class SharedMemoryConnector(OmniConnectorBase):
|
||||
# For now, we rely on the consumer to read and unlink.
|
||||
pass
|
||||
|
||||
def close(self) -> None:
|
||||
pass
|
||||
|
||||
def health(self) -> dict[str, Any]:
|
||||
return {"status": "healthy", "threshold": self.threshold, **self._metrics}
|
||||
@@ -1,6 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
@@ -11,7 +12,6 @@ try:
|
||||
from .utils.config import ConnectorSpec
|
||||
except ImportError:
|
||||
# Fallback for direct execution
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
||||
@@ -56,17 +56,16 @@ class OmniConnectorFactory:
|
||||
|
||||
|
||||
# Register built-in connectors with lazy imports
|
||||
def _create_mooncake_connector(config: dict[str, Any]) -> OmniConnectorBase:
|
||||
def _create_mooncake_store_connector(config: dict[str, Any]) -> OmniConnectorBase:
|
||||
try:
|
||||
from .connectors.mooncake_connector import MooncakeConnector
|
||||
from .connectors.mooncake_store_connector import MooncakeStoreConnector
|
||||
except ImportError:
|
||||
# Fallback import
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
||||
from omni_connectors.connectors.mooncake_connector import MooncakeConnector
|
||||
return MooncakeConnector(config)
|
||||
from omni_connectors.connectors.mooncake_store_connector import MooncakeStoreConnector
|
||||
return MooncakeStoreConnector(config)
|
||||
|
||||
|
||||
def _create_shm_connector(config: dict[str, Any]) -> OmniConnectorBase:
|
||||
@@ -74,7 +73,6 @@ def _create_shm_connector(config: dict[str, Any]) -> OmniConnectorBase:
|
||||
from .connectors.shm_connector import SharedMemoryConnector
|
||||
except ImportError:
|
||||
# Fallback import
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
||||
@@ -86,7 +84,6 @@ def _create_yuanrong_connector(config: dict[str, Any]) -> OmniConnectorBase:
|
||||
try:
|
||||
from .connectors.yuanrong_connector import YuanrongConnector
|
||||
except ImportError:
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
||||
@@ -94,7 +91,21 @@ def _create_yuanrong_connector(config: dict[str, Any]) -> OmniConnectorBase:
|
||||
return YuanrongConnector(config)
|
||||
|
||||
|
||||
def _create_mooncake_transfer_engine_connector(config: dict[str, Any]) -> OmniConnectorBase:
|
||||
try:
|
||||
from .connectors.mooncake_transfer_engine_connector import MooncakeTransferEngineConnector
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
||||
from omni_connectors.connectors.mooncake_transfer_engine_connector import MooncakeTransferEngineConnector
|
||||
return MooncakeTransferEngineConnector(config)
|
||||
|
||||
|
||||
# Register connectors
|
||||
OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_connector)
|
||||
OmniConnectorFactory.register_connector("MooncakeStoreConnector", _create_mooncake_store_connector)
|
||||
OmniConnectorFactory.register_connector("MooncakeTransferEngineConnector", _create_mooncake_transfer_engine_connector)
|
||||
OmniConnectorFactory.register_connector("SharedMemoryConnector", _create_shm_connector)
|
||||
OmniConnectorFactory.register_connector("YuanrongConnector", _create_yuanrong_connector)
|
||||
# Backward-compatible aliases – will be removed in the future
|
||||
OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_store_connector)
|
||||
@@ -13,7 +13,7 @@ logger = get_connector_logger(__name__)
|
||||
class ConnectorSpec:
|
||||
"""Specification for a connector instance."""
|
||||
|
||||
name: str # e.g., "MooncakeConnector", "SharedMemoryConnector", "YuanrongConnector"
|
||||
name: str # e.g., "MooncakeStoreConnector", "SharedMemoryConnector", "YuanrongConnector"
|
||||
extra: dict[str, Any] = field(default_factory=dict) # backend-specific config
|
||||
|
||||
|
||||
|
||||
@@ -85,14 +85,23 @@ def get_connectors_config_for_stage(transfer_config: OmniTransferConfig | None,
|
||||
stage_connectors_config = {}
|
||||
target_stage = str(stage_id)
|
||||
|
||||
# Iterate through all configured edges
|
||||
# Iterate through all configured edges and inject direction-specific role.
|
||||
# The shared edge-level ConnectorSpec is role-neutral; each stage gets
|
||||
# the correct role ("sender" or "receiver") based on its position in
|
||||
# the edge so that MooncakeTransferEngineConnector (and any future
|
||||
# role-aware connector) initializes correctly.
|
||||
for (from_stage, to_stage), spec in transfer_config.connectors.items():
|
||||
# We only care about incoming edges for the worker process
|
||||
# (Worker needs to create connectors to receive data)
|
||||
if to_stage == target_stage:
|
||||
stage_connectors_config[f"from_stage_{from_stage}"] = {"spec": {"name": spec.name, "extra": spec.extra}}
|
||||
# Incoming edge → this stage is the receiver
|
||||
extra = dict(spec.extra) if spec.extra else {}
|
||||
extra.setdefault("role", "receiver")
|
||||
stage_connectors_config[f"from_stage_{from_stage}"] = {"spec": {"name": spec.name, "extra": extra}}
|
||||
elif from_stage == target_stage and target_stage == "0":
|
||||
stage_connectors_config[f"to_stage_{to_stage}"] = {"spec": {"name": spec.name, "extra": spec.extra}}
|
||||
# Outgoing edge for stage 0 — included for async_chunk spec
|
||||
# extraction (omni_stage.py), NOT for connector instantiation.
|
||||
extra = dict(spec.extra) if spec.extra else {}
|
||||
extra.setdefault("role", "sender")
|
||||
stage_connectors_config[f"to_stage_{to_stage}"] = {"spec": {"name": spec.name, "extra": extra}}
|
||||
|
||||
return stage_connectors_config
|
||||
|
||||
@@ -142,42 +151,69 @@ def load_omni_transfer_config(
|
||||
for stage_config in stage_args:
|
||||
stage_id = str(stage_config["stage_id"])
|
||||
|
||||
# Input connectors
|
||||
# Input connectors (this stage is the receiver)
|
||||
# NOTE: role is NOT injected here — the shared edge-level ConnectorSpec
|
||||
# must remain role-neutral. Role is injected per-stage in
|
||||
# get_connectors_config_for_stage() / resolve_omni_kv_config_for_stage().
|
||||
for input_key, conn_ref in stage_config.get("input_connectors", {}).items():
|
||||
if isinstance(conn_ref, str):
|
||||
# Reference to global connector
|
||||
if conn_ref in global_connectors:
|
||||
conn_config = global_connectors[conn_ref]
|
||||
connector = ConnectorSpec(name=conn_config["name"], extra=conn_config.get("extra", {}))
|
||||
extra = dict(conn_config.get("extra", {}))
|
||||
else:
|
||||
raise ValueError(f"Undefined connector reference: {conn_ref}")
|
||||
connector = ConnectorSpec(name=conn_config["name"], extra=extra)
|
||||
else:
|
||||
# Inline connector definition
|
||||
connector = ConnectorSpec(name=conn_ref["name"], extra=conn_ref.get("extra", {}))
|
||||
extra = dict(conn_ref.get("extra", {}))
|
||||
connector = ConnectorSpec(name=conn_ref["name"], extra=extra)
|
||||
|
||||
# Parse from_stage from key (e.g., "from_stage_0" -> "0")
|
||||
from_stage = input_key.replace("from_stage_", "")
|
||||
edge_key = (from_stage, stage_id)
|
||||
connectors[edge_key] = connector
|
||||
# Both sides of an edge may define the same connector reference;
|
||||
# verify consistency if already registered.
|
||||
if edge_key in connectors:
|
||||
existing = connectors[edge_key]
|
||||
if existing.name != connector.name:
|
||||
raise ValueError(
|
||||
f"Connector type mismatch for edge {edge_key[0]}->{edge_key[1]}: "
|
||||
f"previously registered as '{existing.name}', "
|
||||
f"but input_connectors of stage {stage_id} specifies '{connector.name}'"
|
||||
)
|
||||
else:
|
||||
connectors[edge_key] = connector
|
||||
expected_edges.add(edge_key)
|
||||
|
||||
# Output connectors
|
||||
# Output connectors (this stage is the sender)
|
||||
for output_key, conn_ref in stage_config.get("output_connectors", {}).items():
|
||||
if isinstance(conn_ref, str):
|
||||
# Reference to global connector
|
||||
if conn_ref in global_connectors:
|
||||
conn_config = global_connectors[conn_ref]
|
||||
connector = ConnectorSpec(name=conn_config["name"], extra=conn_config.get("extra", {}))
|
||||
extra = dict(conn_config.get("extra", {}))
|
||||
else:
|
||||
raise ValueError(f"Undefined connector reference: {conn_ref}")
|
||||
connector = ConnectorSpec(name=conn_config["name"], extra=extra)
|
||||
else:
|
||||
# Inline connector definition
|
||||
connector = ConnectorSpec(name=conn_ref["name"], extra=conn_ref.get("extra", {}))
|
||||
extra = dict(conn_ref.get("extra", {}))
|
||||
connector = ConnectorSpec(name=conn_ref["name"], extra=extra)
|
||||
|
||||
# Parse to_stage from key (e.g., "to_stage_1" -> "1")
|
||||
to_stage = output_key.replace("to_stage_", "")
|
||||
edge_key = (stage_id, to_stage)
|
||||
connectors[edge_key] = connector
|
||||
if edge_key in connectors:
|
||||
existing = connectors[edge_key]
|
||||
if existing.name != connector.name:
|
||||
raise ValueError(
|
||||
f"Connector type mismatch for edge {edge_key[0]}->{edge_key[1]}: "
|
||||
f"previously registered as '{existing.name}', "
|
||||
f"but output_connectors of stage {stage_id} specifies '{connector.name}'"
|
||||
)
|
||||
else:
|
||||
connectors[edge_key] = connector
|
||||
expected_edges.add(edge_key)
|
||||
|
||||
# Auto-configure SharedMemoryConnector for missing edges based on runtime edges / engine_input_source
|
||||
@@ -294,13 +330,16 @@ def build_stage_connectors(
|
||||
from .config import ConnectorSpec
|
||||
|
||||
connectors: dict[tuple[str, str], Any] = {}
|
||||
# Convert dictionary-formatted config to ConnectorSpec objects
|
||||
# Convert dictionary-formatted config to ConnectorSpec objects.
|
||||
# Only instantiate INPUT connectors ("from_stage_X") — the stage worker
|
||||
# only receives via connectors. Output connectors are handled at the
|
||||
# orchestrator level (try_send_via_connector uses orchestrator connectors).
|
||||
stage_connector_specs = {}
|
||||
for input_key, config in connectors_config.items():
|
||||
if not input_key.startswith("from_stage_"):
|
||||
for key, config in connectors_config.items():
|
||||
if not key.startswith("from_stage_"):
|
||||
continue
|
||||
|
||||
from_stage = input_key.replace("from_stage_", "")
|
||||
from_stage = key.replace("from_stage_", "")
|
||||
spec_dict = config.get("spec", {})
|
||||
if not spec_dict:
|
||||
continue
|
||||
@@ -353,7 +392,8 @@ def resolve_omni_kv_config_for_stage(
|
||||
omni_from = None
|
||||
omni_to = None
|
||||
|
||||
# Prioritize outgoing (Sender) if exists, else check incoming (Receiver)
|
||||
# Prioritize outgoing (Sender) if exists, else check incoming (Receiver).
|
||||
# Inject direction-specific role so the connector initializes correctly.
|
||||
if outgoing:
|
||||
if len(outgoing) > 1:
|
||||
logger.debug(
|
||||
@@ -364,6 +404,7 @@ def resolve_omni_kv_config_for_stage(
|
||||
outgoing.sort(key=lambda x: int(x[0]) if str(x[0]).isdigit() else str(x[0]))
|
||||
to_s, spec = outgoing[0]
|
||||
omni_conn_cfg = {"type": spec.name, **(spec.extra or {})}
|
||||
omni_conn_cfg.setdefault("role", "sender")
|
||||
omni_from = stage_id_str
|
||||
omni_to = str(to_s)
|
||||
elif incoming:
|
||||
@@ -371,6 +412,7 @@ def resolve_omni_kv_config_for_stage(
|
||||
incoming.sort(key=lambda x: int(x[0]) if str(x[0]).isdigit() else str(x[0]))
|
||||
from_s, spec = incoming[0]
|
||||
omni_conn_cfg = {"type": spec.name, **(spec.extra or {})}
|
||||
omni_conn_cfg.setdefault("role", "receiver")
|
||||
omni_from = str(from_s)
|
||||
omni_to = stage_id_str
|
||||
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
import re
|
||||
|
||||
from vllm.logger import current_formatter_type, init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
ORANGE = "\033[38;5;208m"
|
||||
BLUE = "\033[34m"
|
||||
WHITE = "\033[97m"
|
||||
PURPLE = "\033[35m"
|
||||
RESET = "\033[0m"
|
||||
|
||||
# Official vLLM Logo Base
|
||||
VLLM_L1 = f" {WHITE}█ █ █▄ ▄█{RESET}"
|
||||
VLLM_L2 = f" {ORANGE}▄▄{RESET} {BLUE}▄█{RESET} {WHITE}█ █ █ ▀▄▀ █{RESET}"
|
||||
VLLM_L3 = f" {ORANGE}█{RESET}{BLUE}▄█▀{RESET} {WHITE}█ █ █ █{RESET}"
|
||||
VLLM_L4 = f" {BLUE}▀▀{RESET} {WHITE}▀▀▀▀▀ ▀▀▀▀▀ ▀ ▀{RESET}"
|
||||
|
||||
GAP_L1 = " "
|
||||
GAP_L2 = f"{WHITE} ▄▄▄ {RESET}"
|
||||
GAP_L3 = f"{WHITE} {RESET}"
|
||||
GAP_L4 = " "
|
||||
|
||||
O_L1 = f"{BLUE} ▄▀▀{ORANGE}▀▀▄ {RESET}"
|
||||
O_L2 = f"{BLUE} █ {ORANGE}█ {RESET}"
|
||||
O_L3 = f"{BLUE} █ {ORANGE}█ {RESET}"
|
||||
O_L4 = f"{PURPLE} ▀▀▀▀ {RESET}"
|
||||
MNI_L1 = f"{WHITE}█▄ ▄█ █▄ █ ▀█▀ {RESET}"
|
||||
MNI_L2 = f"{WHITE}█ ▀▄▀ █ █ ▀▄ █ █ {RESET}"
|
||||
MNI_L3 = f"{WHITE}█ █ █ ▀▄█ █ {RESET}"
|
||||
MNI_L4 = f"{WHITE}▀ ▀ ▀ ▀ ▀▀▀ {RESET}"
|
||||
|
||||
LOGO = f"""{VLLM_L1}{GAP_L1}{O_L1}{MNI_L1}
|
||||
{VLLM_L2}{GAP_L2}{O_L2}{MNI_L2}
|
||||
{VLLM_L3}{GAP_L3}{O_L3}{MNI_L3}
|
||||
{VLLM_L4}{GAP_L4}{O_L4}{MNI_L4}
|
||||
"""
|
||||
|
||||
|
||||
_ANSI_RE = re.compile(r"\033\[[0-9;]*m")
|
||||
|
||||
|
||||
def log_logo() -> None:
|
||||
logo = LOGO if current_formatter_type(logger) == "color" else _ANSI_RE.sub("", LOGO)
|
||||
logger.info(logo)
|
||||
@@ -6,6 +6,7 @@ diffusion models (e.g., Qwen-Image) through the same CLI interface.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
import uvloop
|
||||
from vllm.entrypoints.cli.types import CLISubcommand
|
||||
@@ -14,6 +15,7 @@ from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
|
||||
from vllm.logger import init_logger
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.entrypoints.cli.logo import log_logo
|
||||
from vllm_omni.entrypoints.openai.api_server import omni_run_server
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -45,6 +47,10 @@ class OmniServeCommand(CLISubcommand):
|
||||
|
||||
@staticmethod
|
||||
def cmd(args: argparse.Namespace) -> None:
|
||||
if not os.environ.get("VLLM_DISABLE_LOG_LOGO"):
|
||||
os.environ["VLLM_DISABLE_LOG_LOGO"] = "1"
|
||||
log_logo()
|
||||
|
||||
# If model is specified in CLI (as positional arg), it takes precedence
|
||||
if hasattr(args, "model_tag") and args.model_tag is not None:
|
||||
args.model = args.model_tag
|
||||
|
||||
@@ -78,7 +78,7 @@ runtime:
|
||||
connectors:
|
||||
# Mooncake connector for cross-node/intra-node communication
|
||||
mooncake_connector:
|
||||
name: MooncakeConnector
|
||||
name: MooncakeStoreConnector
|
||||
extra:
|
||||
host: "127.0.0.1"
|
||||
metadata_server: "http://10.90.67.86:8080/metadata"
|
||||
@@ -87,6 +87,24 @@ runtime:
|
||||
localbuf: 64000000 # 64MB
|
||||
proto: "tcp"
|
||||
|
||||
# PR1 (#1019) note:
|
||||
# - Keep this transfer-engine connector config as a ready-to-use template.
|
||||
# - Bagel does NOT consume this connector in PR1(#1019).
|
||||
# - output_connectors/input_connectors above still point to mooncake_connector.
|
||||
# - We will switch Bagel to this connector in the next PR.
|
||||
rdma_connector:
|
||||
name: MooncakeTransferEngineConnector
|
||||
extra:
|
||||
# NOTE:
|
||||
# - role/sender_host/sender_zmq_port are internal fields resolved by
|
||||
# orchestration logic and should not be set in user YAML.
|
||||
host: "auto" # Auto-detect local IP for RDMA
|
||||
zmq_port: 50051 # ZMQ base port (actual port uses runtime offsets)
|
||||
protocol: "rdma"
|
||||
device_name: "" # e.g. "mlx5_0"; empty for auto-detect
|
||||
memory_pool_size: 2147483648 # 2GB
|
||||
memory_pool_device: "cpu" # "cuda" for GPUDirect RDMA, "cpu" for pinned memory
|
||||
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
|
||||
@@ -108,7 +108,7 @@ runtime:
|
||||
connectors:
|
||||
# Mooncake connector for cross-node/intra-node communication
|
||||
mooncake_connector:
|
||||
name: MooncakeConnector
|
||||
name: MooncakeStoreConnector
|
||||
extra:
|
||||
host: "127.0.0.1"
|
||||
metadata_server: "http://10.90.67.86:8080/metadata"
|
||||
|
||||
@@ -36,7 +36,7 @@ stage_args:
|
||||
repetition_penalty: 1.05
|
||||
# Distributed connector configuration
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_mooncake
|
||||
to_stage_1: mooncake_connector
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm # Use llm stage type to launch OmniLLM
|
||||
@@ -70,9 +70,9 @@ stage_args:
|
||||
stop_token_ids: [2150]
|
||||
# Distributed connector configuration
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_mooncake
|
||||
from_stage_0: mooncake_connector
|
||||
output_connectors:
|
||||
to_stage_2: connector_of_mooncake
|
||||
to_stage_2: mooncake_connector
|
||||
|
||||
- stage_id: 2
|
||||
stage_type: llm # Use llm stage type to launch OmniLLM
|
||||
@@ -106,7 +106,7 @@ stage_args:
|
||||
repetition_penalty: 1.1
|
||||
# Distributed connector configuration
|
||||
input_connectors:
|
||||
from_stage_1: connector_of_mooncake
|
||||
from_stage_1: mooncake_connector
|
||||
|
||||
# Top-level runtime config: default windows and stage edges
|
||||
runtime:
|
||||
@@ -118,8 +118,8 @@ runtime:
|
||||
# Distributed connectors configuration
|
||||
connectors:
|
||||
# Mooncake connector for cross-node/intra-node communication
|
||||
connector_of_mooncake:
|
||||
name: MooncakeConnector
|
||||
mooncake_connector:
|
||||
name: MooncakeStoreConnector
|
||||
extra:
|
||||
host: "127.0.0.1"
|
||||
metadata_server: "http://10.90.67.86:8080/metadata"
|
||||
@@ -129,7 +129,7 @@ runtime:
|
||||
proto: "tcp"
|
||||
|
||||
# SharedMemory connector for intra-node communication
|
||||
connector_of_shared_memory:
|
||||
shared_memory_connector:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536 # 64KB threshold
|
||||
|
||||
Reference in new issue
Block a user