diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index c50bbcfb7..b3b8c06b5 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -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" diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index fc9fea25c..7c4d1df37 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -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 diff --git a/.buildkite/test-nightly.yaml b/.buildkite/test-nightly.yml similarity index 96% rename from .buildkite/test-nightly.yaml rename to .buildkite/test-nightly.yml index 70c6b1fc4..1f1506870 100644 --- a/.buildkite/test-nightly.yaml +++ b/.buildkite/test-nightly.yml @@ -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 diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml new file mode 100644 index 000000000..e0cb91b08 --- /dev/null +++ b/.buildkite/test-ready.yml @@ -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 diff --git a/.buildkite/test-template-amd-omni.j2 b/.buildkite/test-template-amd-omni.j2 index 6442ff544..291ed0a9a 100644 --- a/.buildkite/test-template-amd-omni.j2 +++ b/.buildkite/test-template-amd-omni.j2 @@ -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: diff --git a/benchmarks/distributed/omni_connectors/README.md b/benchmarks/distributed/omni_connectors/README.md new file mode 100644 index 000000000..ab7346441 --- /dev/null +++ b/benchmarks/distributed/omni_connectors/README.md @@ -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 +``` + +--- + +## 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 \ + --remote-host \ + --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 \ + --remote-host \ + --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 \ + --remote-host \ + --tensor-size-mb 1024 \ + --num-transfers 20 \ + --benchmark + +# Consumer +python cross_node_mooncake_transfer_engine.py \ + --role consumer \ + --local-host \ + --remote-host \ + --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 | diff --git a/benchmarks/distributed/omni_connectors/cross_node_mooncake_transfer_engine.py b/benchmarks/distributed/omni_connectors/cross_node_mooncake_transfer_engine.py new file mode 100644 index 000000000..fd01739cc --- /dev/null +++ b/benchmarks/distributed/omni_connectors/cross_node_mooncake_transfer_engine.py @@ -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() diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index d0af83423..8fda12d13 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -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"] diff --git a/docs/design/feature/disaggregated_inference.md b/docs/design/feature/disaggregated_inference.md index 57cd39d16..83c35ac01 100644 --- a/docs/design/feature/disaggregated_inference.md +++ b/docs/design/feature/disaggregated_inference.md @@ -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 diff --git a/docs/design/feature/omni_connectors/mooncake_connector.md b/docs/design/feature/omni_connectors/mooncake_store_connector.md similarity index 96% rename from docs/design/feature/omni_connectors/mooncake_connector.md rename to docs/design/feature/omni_connectors/mooncake_store_connector.md index fee409ecf..5812a6abe 100644 --- a/docs/design/feature/omni_connectors/mooncake_connector.md +++ b/docs/design/feature/omni_connectors/mooncake_store_connector.md @@ -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://:8080/metadata" diff --git a/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md b/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md new file mode 100644 index 000000000..96f93ef9a --- /dev/null +++ b/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md @@ -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 ==> (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:-}" +echo "MC_IB_PCI_RELAXED_ORDERING=${MC_IB_PCI_RELAXED_ORDERING:-}" +``` + +### 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). diff --git a/docs/design/feature/ray_based_execution.md b/docs/design/feature/ray_based_execution.md index f69649d22..ae10d661f 100644 --- a/docs/design/feature/ray_based_execution.md +++ b/docs/design/feature/ray_based_execution.md @@ -45,7 +45,7 @@ ray start --address=: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. diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index 18348c792..5cfc725a7 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -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 diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index a3eece127..671cf1593 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -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 diff --git a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py new file mode 100755 index 000000000..84321f739 --- /dev/null +++ b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py @@ -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() diff --git a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py new file mode 100755 index 000000000..698804f88 --- /dev/null +++ b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py @@ -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() diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index 61927b1a8..d8ea61550 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -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"], }, } diff --git a/tests/e2e/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/stage_configs/qwen3_omni_ci.yaml index 5a16c340d..641fc77ae 100644 --- a/tests/e2e/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/stage_configs/qwen3_omni_ci.yaml @@ -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 diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index d2ed26d97..541a03716 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -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 diff --git a/vllm_omni/distributed/__init__.py b/vllm_omni/distributed/__init__.py index 3a7634219..f43a833f8 100644 --- a/vllm_omni/distributed/__init__.py +++ b/vllm_omni/distributed/__init__.py @@ -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 diff --git a/vllm_omni/distributed/omni_connectors/__init__.py b/vllm_omni/distributed/omni_connectors/__init__.py index fcec32eaa..383267928 100644 --- a/vllm_omni/distributed/omni_connectors/__init__.py +++ b/vllm_omni/distributed/omni_connectors/__init__.py @@ -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 diff --git a/vllm_omni/distributed/omni_connectors/connectors/base.py b/vllm_omni/distributed/omni_connectors/connectors/base.py index e3562b896..83edb2ab0 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/base.py +++ b/vllm_omni/distributed/omni_connectors/connectors/base.py @@ -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}" diff --git a/vllm_omni/distributed/omni_connectors/connectors/mooncake_connector.py b/vllm_omni/distributed/omni_connectors/connectors/mooncake_store_connector.py similarity index 76% rename from vllm_omni/distributed/omni_connectors/connectors/mooncake_connector.py rename to vllm_omni/distributed/omni_connectors/connectors/mooncake_store_connector.py index e28486481..c672e35f7 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/mooncake_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/mooncake_store_connector.py @@ -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) diff --git a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py new file mode 100644 index 000000000..b1dc8b898 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py @@ -0,0 +1,1202 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import queue +import socket +import threading +import time as _time_mod +import uuid +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import Any + +import msgspec +import torch +import zmq + +from ..utils.logging import get_connector_logger +from ..utils.serialization import OmniSerializer +from .base import OmniConnectorBase + +logger = get_connector_logger(__name__) + +try: + from mooncake.engine import TransferEngine +except ImportError: + TransferEngine = None + +# Stale buffer TTL: buffers older than this are automatically reclaimed +# to prevent memory leaks when receiver crashes or gives up. +_BUFFER_TTL_SECONDS = 300 # 5 minutes + +# ZMQ Message constants +TRANS_DONE = b"trans_done" +TRANS_ERROR = b"trans_error" +QUERY_INFO = b"query_info" +INFO_NOT_FOUND = b"info_not_found" + + +@dataclass +class QueryRequest: + """Request to query metadata for a specific key.""" + + request_id: str + + +@dataclass +class QueryResponse: + """Response containing metadata for a request.""" + + request_id: str + data_size: int + is_fast_path: bool + + +@dataclass +class MooncakeAgentMetadata: + """ + Metadata exchanged via ZMQ Handshake. + """ + + remote_hostname: str + remote_port: int # RDMA Port + request_id: str + dst_addrs: list[int] + lengths: list[int] + + +class BufferAllocator: + """ + Manages the allocation of memory segments within the registered pool. + Thread-safe implementation using a simple free list. + """ + + def __init__(self, total_size: int, alignment: int = 4096): + self.total_size = total_size + self.alignment = alignment + self.lock = threading.Lock() + # Free list: [(start, size), ...] sorted by start + self.free_blocks = [(0, total_size)] + + def alloc(self, size: int) -> int: + """ + Allocates a block of 'size' bytes. + Returns the starting offset. + """ + # Align size upwards + aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment + + with self.lock: + for i, (start, block_size) in enumerate(self.free_blocks): + if block_size >= aligned_size: + # Found a block + new_start = start + aligned_size + new_size = block_size - aligned_size + + if new_size > 0: + self.free_blocks[i] = (new_start, new_size) + else: + self.free_blocks.pop(i) + return start + + raise MemoryError(f"Out of memory in buffer pool. Requested {size} bytes (aligned {aligned_size}).") + + def free(self, offset: int, size: int): + """ + Frees a previously allocated block. + """ + aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment + + with self.lock: + # Check for double-free and corruption + for start, length in self.free_blocks: + # Case 1: Exact match = double free, safe to ignore + if offset == start and aligned_size == length: + logger.warning(f"Double free detected at offset {offset}, size {aligned_size}. Ignoring.") + return + # Case 2: Block is fully contained within an existing free block = also double free + # This happens when the block was freed and then merged with adjacent blocks + if offset >= start and offset + aligned_size <= start + length: + logger.warning( + f"Double free detected: block {offset}-{offset + aligned_size} " + f"is already within free block {start}-{start + length}. Ignoring." + ) + return + # Case 3: Partial overlap (but not fully contained) = memory corruption + if not (offset + aligned_size <= start or start + length <= offset): + raise RuntimeError( + f"Memory corruption detected: freeing {offset}-{offset + aligned_size} " + f"partially overlaps with free block {start}-{start + length}" + ) + + self.free_blocks.append((offset, aligned_size)) + self.free_blocks.sort() # Sort by offset + + # Merge adjacent blocks + i = 0 + while i < len(self.free_blocks) - 1: + curr_start, curr_size = self.free_blocks[i] + next_start, next_size = self.free_blocks[i + 1] + + if curr_start + curr_size == next_start: + self.free_blocks[i] = (curr_start, curr_size + next_size) + self.free_blocks.pop(i + 1) + else: + i += 1 + + +class ManagedBuffer: + """ + A temporary view into the global memory pool. + Must be kept alive while the data view is being used. + """ + + def __init__(self, allocator: BufferAllocator, offset: int, size: int, pool_tensor: torch.Tensor): + self.allocator = allocator + self.offset = offset + self.size = size + self.pool_tensor = pool_tensor + self._released = False + + def release(self): + """Explicitly release the buffer back to the pool.""" + if not self._released: + self.allocator.free(self.offset, self.size) + self._released = True + + def __del__(self): + self.release() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.release() + + @property + def tensor(self) -> torch.Tensor: + """ + Returns a 1D uint8 zero-copy view of the buffer. + """ + return self.pool_tensor[self.offset : self.offset + self.size] + + def as_tensor(self, dtype: torch.dtype, shape: tuple) -> torch.Tensor: + """ + Returns a typed, shaped zero-copy view. + Validates size, shape, and alignment. + """ + itemsize = torch.tensor([], dtype=dtype).element_size() + + # Calculate expected size + expected_bytes = itemsize + for dim in shape: + if dim < 0: + raise ValueError("Dynamic dimension (-1) is not supported in as_tensor") + expected_bytes *= dim + + if expected_bytes != self.size: + raise ValueError( + f"Shape {shape} with dtype {dtype} requires {expected_bytes} bytes, but buffer size is {self.size}" + ) + + # Check alignment (offset must be divisible by itemsize) + if self.offset % itemsize != 0: + raise RuntimeError(f"Buffer offset {self.offset} is not aligned for dtype {dtype} (itemsize {itemsize})") + + raw_view = self.tensor + # view() requires contiguous memory, slice of contiguous tensor is contiguous + typed_view = raw_view.view(dtype) + return typed_view.reshape(shape) + + def to_bytes(self) -> bytes: + """ + Returns a copy of the data as python bytes. + Performs D2H copy if pool is on GPU. + """ + t = self.tensor + if t.is_cuda: + t = t.cpu() + return t.numpy().tobytes() + + +class MooncakeTransferEngineConnector(OmniConnectorBase): + """ + OmniConnector implementation using Mooncake Transfer Engine with a managed memory pool. + Supports both CPU (Pinned) and GPU memory pools, and both RDMA and TCP protocols. + Topology limitations (current implementation): + Current design focuses on peer-to-peer communication between stages. + - **1 sender → 1 receiver per key**: After a successful RDMA write the + sender immediately cleans up the buffer (``cleanup()``), so only the + first receiver to pull a given key will succeed. Broadcast / multicast + (1 sender → N receivers sharing the same data) is not yet supported. + - **1 receiver → 1 sender**: ``update_sender_info()`` stores a single + ``(sender_host, sender_zmq_port)`` pair, so a receiver can only query + metadata from one sender at a time. + + Future work: + - Support 1 sender → N receivers (e.g. reference-counted buffers, or + explicit ``retain()`` / ``release()`` semantics so the buffer survives + multiple pulls). + - Support 1 receiver → N senders (e.g. a sender registry mapping + ``get_key`` prefixes to different sender endpoints). + """ + + # RDMA connector copies raw bytes/tensor directly to the memory pool + # without going through OmniSerializer, so the sender can use + # to_gpu_tensor() / to_bytes() fast-paths. + supports_raw_data: bool = True + + def __init__(self, config: dict[str, Any]): + if TransferEngine is None: + raise ImportError("Mooncake not available") + + self._closed = False + self._bind_error: Exception | None = None # fatal ZMQ bind error from listener thread + + # --- Early init of all teardown-related fields --- + # If __init__ fails later (engine init, pool alloc, etc.), __del__ → close() + # must not crash on missing attributes. Safe defaults here ensure close() + # can always run without AttributeError. + self._stop_event = threading.Event() + self._sender_executor: ThreadPoolExecutor | None = None + self._listener_thread: threading.Thread | None = None + self._listener_ready = threading.Event() + self._local_buffers: dict[str, Any] = {} + self._local_buffers_lock = threading.Lock() + self._req_local = threading.local() + self._worker_local = threading.local() + self._last_ttl_check: float = _time_mod.monotonic() + + self._metrics = { + "puts": 0, + "gets": 0, + "bytes_transferred": 0, + "errors": 0, + "timeouts": 0, + } + + self.config = config + host_config = config.get("host", "127.0.0.1") + # Support "auto" to auto-detect local IP address + if host_config.lower() == "auto": + self.host = self._get_local_ip() + logger.info(f"Auto-detected local IP for RDMA: {self.host}") + else: + self.host = host_config + self.zmq_port = config.get("zmq_port", 50051) + self.protocol = config.get("protocol", "rdma") + + # --- RDMA Device Configuration --- + # Specify device names to filter (comma-separated), or empty for all devices. + # Example: "mlx5_0,mlx5_1" to use only specific NICs. + # This is important for environments with mixed InfiniBand/RoCE NICs. + self.device_name = config.get("device_name", "") + if not self.device_name: + env_device = os.environ.get("RDMA_DEVICE_NAME", "") + if env_device: + self.device_name = env_device + logger.info(f"Using RDMA_DEVICE_NAME from env: {self.device_name}") + + # --- Memory Pool Configuration --- + self.pool_size = config.get("memory_pool_size", 1024**3) # Default 1GB + self.pool_device = config.get("memory_pool_device", "cpu") + + # --- Sender Configuration (for receiver to query without metadata) --- + # When receiver doesn't have metadata, it uses these to connect to sender + self.sender_host = config.get("sender_host", None) + self.sender_zmq_port = config.get("sender_zmq_port", None) + + # --- Role --- + # "sender": bind ZMQ listener, accept put() calls. + # "receiver": skip ZMQ bind, only accept get() calls. + # The orchestration layer (get_connectors_config_for_stage / + # kv_transfer_manager) is responsible for injecting the correct role. + role = str(config.get("role", "sender")).lower() + if role not in {"sender", "receiver"}: + raise ValueError( + f"Invalid role={role!r} for MooncakeTransferEngineConnector. Expected 'sender' or 'receiver'." + ) + self.can_put = role == "sender" + + self.engine_id = str(uuid.uuid4()) + + # --- Mooncake Engine Init --- + self.engine = TransferEngine() + # Note: For P2P handshake mode, local_hostname should be just the IP address. + # Mooncake will auto-assign an RPC port, retrievable via get_rpc_port(). + ret = self.engine.initialize(self.host, "P2PHANDSHAKE", self.protocol, self.device_name) + if ret != 0: + raise RuntimeError(f"Mooncake Engine initialization failed with code {ret}") + + self.rpc_port = self.engine.get_rpc_port() + logger.info(f"MooncakeTransferEngineConnector initialized at {self.host}:{self.rpc_port}") + + # --- Pool Allocation & Registration --- + logger.info(f"Allocating RDMA Memory Pool: {self.pool_size / 1024**2:.2f} MB on {self.pool_device}") + try: + if self.pool_device == "cpu": + self.pool = torch.empty(self.pool_size, dtype=torch.uint8).pin_memory() + self.base_ptr = self.pool.data_ptr() + else: + self.pool = torch.empty(self.pool_size, dtype=torch.uint8, device=self.pool_device) + self.base_ptr = self.pool.data_ptr() + + # Register the entire pool + ret = self.engine.register_memory(self.base_ptr, self.pool_size) + if ret != 0: + raise RuntimeError("Failed to register memory pool with Mooncake Engine") + + except Exception as e: + logger.error(f"Failed to allocate/register memory pool: {e}") + raise + + self.allocator = BufferAllocator(self.pool_size, alignment=4096) # 4KB alignment for safety + + # --- State Management & Background Threads --- + # Most fields already initialized at the top of __init__ for teardown + # safety. Only create the real ZMQ context and thread pool here. + self.zmq_ctx = zmq.Context() + self._sender_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="mooncake-sender") + + # Log complete connector configuration for debugging + logger.info( + f"MooncakeTransferEngineConnector config summary:\n" + f" Local: host={self.host}, zmq_port={self.zmq_port}, rpc_port={self.rpc_port}\n" + f" Remote: sender_host={self.sender_host}, sender_zmq_port={self.sender_zmq_port}\n" + f" Role: can_put={self.can_put}, configured_role={config.get('role', 'sender')}" + ) + + # Only sender needs ZMQ listener to handle pull requests + if self.can_put: + self._last_ttl_check = _time_mod.monotonic() # reset after slow init + self._listener_thread = threading.Thread(target=self._zmq_listener_loop, daemon=True) + self._listener_thread.start() + # Wait briefly for listener to bind (or fail) + self._listener_ready.wait(timeout=1.0) + # Propagate any bind error to the caller — sender must bind successfully. + if self._bind_error is not None: + raise RuntimeError( + f"MooncakeTransferEngineConnector failed to bind ZMQ on " + f"{self.host}:{self.zmq_port}: {self._bind_error}" + ) from self._bind_error + logger.info( + f"MooncakeTransferEngineConnector started as SENDER (ZMQ listener on {self.host}:{self.zmq_port})" + ) + else: + # Receiver mode — sender address is provided per-request via + # metadata from put() through the queue, not pre-configured. + if not self.sender_host or self.sender_host.lower() == "auto": + logger.info( + "MooncakeTransferEngineConnector receiver: sender_host='auto', " + "awaiting sender info via task from orchestrator." + ) + else: + logger.info( + f"MooncakeTransferEngineConnector started as RECEIVER " + f"(will query sender at {self.sender_host}:{self.sender_zmq_port})" + ) + + def get_connection_info(self) -> dict[str, Any]: + """Get connection info for this connector (useful for sender to share with receivers).""" + return { + "host": self.host, + "zmq_port": self.zmq_port, + "rpc_port": self.rpc_port, + "can_put": self.can_put, + } + + def update_sender_info(self, sender_host: str, sender_zmq_port: int) -> None: + """ + Inject the sender's ZMQ endpoint into the receiver connector. + Used for NO METADATA GET calls.(E.g: KV-cache transfer path) + Must be called before using get() without metadata! + Otherwise, get() will raise an error. + """ + self.sender_host = sender_host + self.sender_zmq_port = sender_zmq_port + logger.info(f"Sender info updated: host={sender_host!r}, zmq_port={sender_zmq_port}") + + def _get_local_ip(self) -> str: + """ + Auto-detect the local IP address that can be used for RDMA communication. + This tries to get the IP that would be used to communicate externally, + not the loopback address. + """ + try: + # Create a socket to determine the local IP used for external communication + # We don't actually connect, just use the socket to get routing info + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s: + # Use Google's DNS as a target (doesn't actually connect) + s.connect(("8.8.8.8", 80)) + local_ip = s.getsockname()[0] + return local_ip + except Exception as e: + logger.warning(f"Failed to auto-detect local IP: {e}, falling back to hostname lookup") + try: + # Fallback: use hostname resolution + hostname = socket.gethostname() + local_ip = socket.gethostbyname(hostname) + return local_ip + except Exception as e2: + logger.error(f"Failed to get local IP via hostname: {e2}, using 127.0.0.1") + return "127.0.0.1" + + def _get_req_socket(self, zmq_addr: str, timeout_ms: int) -> zmq.Socket: + """Get or create a thread-local cached ZMQ REQ socket. + + Each calling thread maintains its own ``{zmq_addr: socket}`` map + stored in ``self._req_local``, so concurrent ``get()`` / + ``_query_metadata_from_sender()`` calls from different threads + never share a REQ socket — which would violate ZMQ's strict + send→recv ordering constraint. + + If a call fails, use ``_invalidate_req_socket()`` to discard the + broken socket; the next call will transparently create a fresh one. + """ + cache: dict[str, zmq.Socket] = getattr(self._req_local, "cache", None) # type: ignore[assignment] + if cache is None: + cache = {} + self._req_local.cache = cache + + sock = cache.get(zmq_addr) + if sock is None: + sock = self.zmq_ctx.socket(zmq.REQ) + sock.connect(zmq_addr) + cache[zmq_addr] = sock + # Set timeouts per-call (safe to change between send/recv cycles). + # SNDTIMEO guards against send() blocking when the peer is unreachable + # or its receive buffer is full. + sock.setsockopt(zmq.SNDTIMEO, timeout_ms) + sock.setsockopt(zmq.RCVTIMEO, timeout_ms) + return sock + + def _invalidate_req_socket(self, zmq_addr: str) -> None: + """Close and remove a thread-local cached REQ socket (e.g., after recv timeout).""" + cache: dict[str, zmq.Socket] = getattr(self._req_local, "cache", None) # type: ignore[assignment] + if cache is None: + return + sock = cache.pop(zmq_addr, None) + if sock is not None: + try: + sock.close(linger=0) + except Exception: + pass # Best-effort; zmq_ctx.term() will reclaim if needed + + def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]: + """ + Producer Side. + Exposes data for RDMA transfer. + + Behavior by data type: + - ManagedBuffer (from this pool): Zero-copy, is_fast_path=True + - ManagedBuffer (different pool): Fallback to tensor copy path + - torch.Tensor / bytes: Copy to pool, is_fast_path=True + - Other (dict, etc.): Serialize to bytes, copy to pool, is_fast_path=False + (receiver will deserialize automatically) + """ + if self._closed: + raise RuntimeError("Cannot put data: MooncakeTransferEngineConnector is closed") + + if not self.can_put: + logger.warning( + f"Rejecting put for {put_key}: connector is in receiver mode " + f"(sender fallback occurred, ZMQ port likely occupied by another instance)" + ) + return False, 0, None + + put_key = self._make_key(put_key, from_stage, to_stage) + + try: + src_addr = 0 + size = 0 + holder = None + is_fast_path = True + should_release_holder = False + + # Reject empty data early — alloc(0) has undefined behaviour and + # would leave a zombie entry in _local_buffers that the receiver + # side can never consume (get() returns None for data_size==0). + if isinstance(data, bytes) and len(data) == 0: + logger.warning(f"Rejecting put for {put_key}: empty bytes payload") + return False, 0, None + if isinstance(data, torch.Tensor) and data.nbytes == 0: + logger.warning(f"Rejecting put for {put_key}: zero-size tensor") + return False, 0, None + + # Pre-process: serialize non-raw types before entering the copy path. + # This avoids the previous recursive put() pattern and keeps + # _local_buffers writes atomic (single write, no override needed). + if not isinstance(data, (ManagedBuffer, torch.Tensor, bytes)): + data = OmniSerializer.serialize(data) + is_fast_path = False # Receiver must deserialize + + if isinstance(data, ManagedBuffer): + # Zero-Copy Path + # Validate that the buffer belongs to this connector's pool. + if data.pool_tensor.data_ptr() != self.pool.data_ptr(): + # Fallback to copy path: extract tensor and let the + # torch.Tensor branch below handle the copy. + logger.warning("ManagedBuffer from different pool detected. Falling back to copy path.") + data = data.tensor.contiguous() + # Fall through to torch.Tensor branch + else: + src_addr = self.base_ptr + data.offset + size = data.size + holder = data # Keep the buffer alive + should_release_holder = False # Caller owns it + + # Use 'if' (not 'elif') so the ManagedBuffer fallback above can + # convert to Tensor and flow into this branch seamlessly. + if isinstance(data, (torch.Tensor, bytes)): + # Copy Path + # 1. Determine size + if isinstance(data, torch.Tensor): + size = data.nbytes + tensor_data = data + else: + size = len(data) + # Convert bytes to tensor for copy + tensor_data = torch.frombuffer(data, dtype=torch.uint8) + + # 2. Alloc from pool + try: + offset = self.allocator.alloc(size) + holder = ManagedBuffer(self.allocator, offset, size, self.pool) + should_release_holder = True # We created it, we release it + except MemoryError: + logger.error(f"Pool exhausted, cannot put data size {size}") + return False, 0, None + + # 3. Copy data to pool + # Handle device mismatch for copy + try: + dst_tensor = holder.tensor + if isinstance(data, torch.Tensor): + if not data.is_contiguous(): + data = data.contiguous() + + # View as flat uint8 + src_view = data.view(torch.uint8).flatten() + if src_view.device != dst_tensor.device: + dst_tensor.copy_(src_view, non_blocking=True) + # Ensure copy is complete before exposing buffer to RDMA + # Must sync on SOURCE device for D2H copies, or DST for H2D + if src_view.is_cuda: + with torch.cuda.device(src_view.device): + torch.cuda.current_stream().synchronize() + elif dst_tensor.is_cuda: + with torch.cuda.device(dst_tensor.device): + torch.cuda.current_stream().synchronize() + else: + dst_tensor.copy_(src_view) + if dst_tensor.is_cuda: + with torch.cuda.device(dst_tensor.device): + torch.cuda.current_stream().synchronize() + else: + # bytes -> tensor copy + # torch.frombuffer creates CPU tensor. If pool is GPU, copy_ handles H2D. + dst_tensor.copy_(tensor_data) + if dst_tensor.is_cuda: + with torch.cuda.device(dst_tensor.device): + torch.cuda.current_stream().synchronize() + except Exception as e: + # Copy failed, release the allocated buffer to prevent leak + holder.release() + logger.error(f"Failed to copy data to pool: {e}") + return False, 0, None + + src_addr = self.base_ptr + offset + + # Final guard: reject zero-size payloads regardless of how they + # got here (ManagedBuffer(size=0), empty serialization output, etc.). + # This is the unified catch-all that complements the early checks + # for bytes/Tensor above. + if size <= 0: + if should_release_holder and isinstance(holder, ManagedBuffer): + holder.release() + logger.warning(f"Rejecting put for {put_key}: final payload size is {size}") + return False, 0, None + + with self._local_buffers_lock: + # Release old buffer if put_key already exists (prevents pool leak) + old_item = self._local_buffers.pop(put_key, None) + if old_item: + _, _, old_holder, old_should_release, _, _ = old_item + if old_should_release and isinstance(old_holder, ManagedBuffer): + old_holder.release() + logger.warning(f"Released old buffer for duplicate put_key: {put_key}") + + # Store: (src_addrs, lengths, holder, should_release, is_fast_path, created_at) + self._local_buffers[put_key] = ( + [src_addr], + [size], + holder, + should_release_holder, + is_fast_path, + _time_mod.monotonic(), + ) + + # Metadata for Consumer + metadata = { + "source_host": self.host, + "source_port": self.zmq_port, + "data_size": size, + "is_fast_path": is_fast_path, # Hint: True = return ManagedBuffer, False = deserialize + } + + self._metrics["puts"] += 1 + self._metrics["bytes_transferred"] += size + + return True, size, metadata + + except Exception as e: + self._metrics["errors"] += 1 + logger.error(f"RDMA Put failed for {put_key}: {e}", exc_info=True) + return False, 0, None + + def _query_metadata_from_sender(self, get_key: str) -> dict[str, Any] | None: + """Query metadata from sender via ZMQ (fallback when ``metadata=None``). + + ``get()`` supports two metadata resolution paths:: + + get(metadata=?) + ├── metadata provided (adapter path) + │ → use metadata directly (source_host/port/data_size) + │ → RDMA pull + └── metadata=None (KV-transfer polling path) + → _query_metadata_from_sender(get_key) ← this method + │ + ├── sender_host resolved (via update_sender_info) + │ → ZMQ query → get data_size/is_fast_path + │ → construct metadata → RDMA pull + └── sender_host unresolved ("auto" / None) + → return None → caller retries or times out + + For the second path, the caller must call + :meth:`update_sender_info` before ``get()`` to resolve the sender's ZMQ endpoint. + Support the two paths in case that the orchestrator pushes the request info + to different stages at the same time knowing metadata or not. + """ + zmq_addr = f"tcp://{self.sender_host}:{self.sender_zmq_port}" + req_socket = self._get_req_socket(zmq_addr, timeout_ms=5000) + + try: + # Send query request + query = QueryRequest(request_id=get_key) + req_socket.send(QUERY_INFO + msgspec.msgpack.encode(query)) + resp = req_socket.recv() + + if resp == INFO_NOT_FOUND: + return None + + # Parse response + query_resp = msgspec.msgpack.decode(resp, type=QueryResponse) + return { + # source_host/source_port are used for verification + "source_host": self.sender_host, + "source_port": self.sender_zmq_port, + "data_size": query_resp.data_size, + "is_fast_path": query_resp.is_fast_path, + } + except Exception as e: + # Socket may be stuck in bad state after timeout; discard it + self._invalidate_req_socket(zmq_addr) + logger.debug(f"Failed to query metadata for {get_key}: {e}") + return None + + def get( + self, + from_stage: str, + to_stage: str, + get_key: str, + metadata: dict[str, Any] | None = None, + ) -> tuple[Any, int] | None: + """ + Consumer Side. + Allocates from local pool and pulls data via RDMA. + + If metadata is not provided, will attempt to query it from sender + using configured sender_host/sender_zmq_port. + + Returns: + ``(data, size)`` on success, ``None`` on failure. + + - **is_fast_path=True** (tensor *or* bytes payload): + Returns ``(ManagedBuffer, size)``. + **CALLER MUST call ``ManagedBuffer.release()`` after consuming.** + Note: even if the producer ``put()`` raw ``bytes``, the consumer + receives a ``ManagedBuffer`` — use ``buf.to_bytes()`` to obtain + a ``bytes`` copy, or ``buf.tensor`` for zero-copy access. + - **is_fast_path=False** (serialized Python object): + Returns ``(DeserializedObject, size)``. + Buffer is auto-released internally after deserialization. + """ + if self._closed: + raise RuntimeError("Cannot get data: MooncakeTransferEngineConnector is closed") + + get_key = self._make_key(get_key, from_stage, to_stage) + + _t0 = _time_mod.perf_counter() + + # If no metadata provided, try to query from sender + if not metadata: + # Must insert sender info before using get() without metadata. + if not self.sender_host or not self.sender_zmq_port or str(self.sender_host).lower() == "auto": + raise RuntimeError( + f"get(metadata=None) requires sender info to be resolved, " + f"but sender_host={self.sender_host!r}, sender_zmq_port={self.sender_zmq_port!r}. " + f"Call update_sender_info(host, port) before using get() without metadata." + ) + metadata = self._query_metadata_from_sender(get_key) + if not metadata: + return None + + _t1 = _time_mod.perf_counter() + _query_ms = (_t1 - _t0) * 1000 + + src_host = metadata.get("source_host") + src_port = metadata.get("source_port") + data_size = metadata.get("data_size", 0) + is_fast_path = metadata.get("is_fast_path", False) + + if not src_host or not src_port or str(src_host).lower() == "auto": + logger.error( + f"Invalid metadata for {get_key}: source_host={src_host!r}, " + f"source_port={src_port!r}. Cannot establish RDMA connection." + ) + return None + + if data_size == 0: + logger.warning(f"Skipping get for {get_key}: data_size is 0 (metadata={metadata})") + return None + + # 1. Allocate Destination Buffer from Pool + try: + offset = self.allocator.alloc(data_size) + recv_buffer = ManagedBuffer(self.allocator, offset, data_size, self.pool) + dst_ptr = self.base_ptr + offset + except MemoryError: + logger.error(f"Failed to allocate {data_size} bytes in receive pool") + return None + + _t2 = _time_mod.perf_counter() + _alloc_ms = (_t2 - _t1) * 1000 + + # 2. Prepare Handshake + agent_meta = MooncakeAgentMetadata( + remote_hostname=self.host, + remote_port=self.rpc_port, + request_id=get_key, + dst_addrs=[dst_ptr], + lengths=[data_size], + ) + + # 3. ZMQ Transaction (uses cached socket to avoid TCP reconnection) + # Timeout scales with data size: base 30s + 5s per 100MB. + # Large RDMA transfers (esp. loopback or cross-NIC) can take tens of + # seconds; if we time out too early the receiver closes and deregisters + # its memory, causing the sender's in-flight write to fail with ret=-1. + _base_timeout_ms = 30000 + _size_timeout_ms = max(0, (data_size // (100 * 1024 * 1024))) * 5000 + _total_timeout_ms = _base_timeout_ms + _size_timeout_ms + zmq_addr = f"tcp://{src_host}:{src_port}" + req_socket = self._get_req_socket(zmq_addr, timeout_ms=_total_timeout_ms) + + try: + req_socket.send(msgspec.msgpack.encode(agent_meta)) + resp = req_socket.recv() + + _t3 = _time_mod.perf_counter() + _rdma_ms = (_t3 - _t2) * 1000 + + if resp == TRANS_DONE: + # Success + # Ensure data is visible on GPU + # Note: RDMA write visibility on GPU usually requires some form of fence/sync. + # torch.cuda.synchronize() is a heavy hammer but safe. + # Ideally Mooncake engine provides a way to poll for completion that guarantees visibility. + # TODO(wzliu): Replace synchronize with cuda event in the future for better performance. + if self.pool.is_cuda: + with torch.cuda.device(self.pool.device): + torch.cuda.current_stream().synchronize() + + _t4 = _time_mod.perf_counter() + _sync_ms = (_t4 - _t3) * 1000 + + if is_fast_path: + # Return ManagedBuffer directly for ALL fast_path data + # (both bytes and tensor payloads). Caller is responsible + # for consuming the data and calling release() afterwards. + # This avoids the expensive to_bytes() copy (~54ms for 115MB). + # + # Usage patterns for callers: + # - Zero-copy read: mv = memoryview(buf.tensor.numpy()) + # - Copy to GPU: buf.tensor.to(device); buf.release() + # - Fallback: data = buf.to_bytes(); buf.release() + _t5 = _time_mod.perf_counter() + _total_ms = (_t5 - _t0) * 1000 + _mbps = (data_size / 1024 / 1024) / (_total_ms / 1000) if _total_ms > 0 else 0 + logger.info( + f"[RDMA GET] {get_key}: query={_query_ms:.1f}ms, alloc={_alloc_ms:.1f}ms, " + f"rdma={_rdma_ms:.1f}ms, sync={_sync_ms:.1f}ms, " + f"total={_total_ms:.1f}ms, {_mbps:.1f} MB/s (fast_path, zero-copy)" + ) + self._metrics["gets"] += 1 + self._metrics["bytes_transferred"] += data_size + return recv_buffer, data_size + else: + # If it was a serialized object or generic bytes, we assume standard Omni behavior: + # Deserialize and return object. This requires a copy (to_bytes). + # We MUST release the buffer after deserialization. + try: + _t_copy_start = _time_mod.perf_counter() + raw_bytes = recv_buffer.to_bytes() + _t_copy_end = _time_mod.perf_counter() + _copy_ms = (_t_copy_end - _t_copy_start) * 1000 + + _t_deser_start = _time_mod.perf_counter() + val = OmniSerializer.deserialize(raw_bytes) + _t_deser_end = _time_mod.perf_counter() + _deser_ms = (_t_deser_end - _t_deser_start) * 1000 + + _total_ms = (_t_deser_end - _t0) * 1000 + _mbps = (data_size / 1024 / 1024) / (_total_ms / 1000) if _total_ms > 0 else 0 + logger.info( + f"[RDMA GET] {get_key}: query={_query_ms:.1f}ms, alloc={_alloc_ms:.1f}ms, " + f"rdma={_rdma_ms:.1f}ms, sync={_sync_ms:.1f}ms, copy={_copy_ms:.1f}ms, " + f"deser={_deser_ms:.1f}ms, total={_total_ms:.1f}ms, {_mbps:.1f} MB/s" + ) + self._metrics["gets"] += 1 + self._metrics["bytes_transferred"] += data_size + return val, data_size + finally: + recv_buffer.release() + else: + self._metrics["errors"] += 1 + logger.error(f"RDMA Get failed: received {resp} instead of TRANS_DONE") + recv_buffer.release() + return None + except Exception as e: + # Socket may be stuck after timeout; discard it + self._invalidate_req_socket(zmq_addr) + self._metrics["timeouts"] += 1 + logger.error(f"RDMA Get error: {e}", exc_info=True) + recv_buffer.release() + return None + + def cleanup(self, request_id: str, from_stage: str | None = None, to_stage: str | None = None) -> None: + """Release the producer-side buffer associated with the request. + + Args: + request_id: The key used in ``put()``. + from_stage: Optional source stage. When both *from_stage* and + *to_stage* are provided the method applies ``_make_key()`` + internally, so callers can pass the same raw key they used + in ``put()`` / ``get()`` without knowing the internal format. + When omitted the *request_id* is used as-is (suitable for + internal callers that already hold the transformed key). + to_stage: Optional destination stage (see *from_stage*). + """ + if (from_stage is None) != (to_stage is None): + raise ValueError( + f"cleanup() requires both from_stage and to_stage, or neither. " + f"Got from_stage={from_stage!r}, to_stage={to_stage!r}" + ) + if from_stage is not None and to_stage is not None: + request_id = self._make_key(request_id, from_stage, to_stage) + with self._local_buffers_lock: + item = self._local_buffers.pop(request_id, None) + if item: + # item is (src_addrs, lengths, holder, should_release, is_fast_path, created_at) + _, _, holder, should_release, _, _ = item + if should_release and isinstance(holder, ManagedBuffer): + # We own this buffer (created internally), so we must release it. + holder.release() + # If holder was externally owned (should_release=False), we do nothing. + # If holder was something else (e.g. Tensor), GC handles it. + + def health(self) -> dict[str, Any]: + if self._closed: + return {"status": "unhealthy", "error": "Connector is closed"} + + return { + "status": "healthy", + "host": self.host, + "metadata_server": None, + "master": None, + "protocol": self.protocol, + "pool_device": self.pool_device, + "pool_size": self.pool_size, + **self._metrics, + } + + def close(self) -> None: + """ + Gracefully shutdown the connector and release all resources. + This method should be called when the connector is no longer needed. + Idempotent: safe to call multiple times. + """ + # Idempotent guard — use getattr for safety: if __init__ raised before + # setting _closed, __del__ → close() would hit AttributeError. + if getattr(self, "_closed", True): + return + self._closed = True + + logger.info("Closing MooncakeTransferEngineConnector...") + + # 1. Signal listener thread to stop + self._stop_event.set() + + # 2. Wait for listener thread to finish (only if sender mode) + if self._listener_thread is not None and self._listener_thread.is_alive(): + self._listener_thread.join(timeout=2.0) + if self._listener_thread.is_alive(): + logger.warning("Listener thread did not stop gracefully") + + # 3. Shutdown sender executor (may be None if __init__ failed early) + if self._sender_executor is not None: + self._sender_executor.shutdown(wait=True, cancel_futures=False) + + # 4. Release all pending buffers + with self._local_buffers_lock: + for req_id, item in list(self._local_buffers.items()): + _, _, holder, should_release, _, _ = item + if should_release and isinstance(holder, ManagedBuffer): + holder.release() + self._local_buffers.clear() + + # 5. Close thread-local cached REQ sockets (for the calling thread). + # Sockets cached in *other* threads are not directly accessible here + # because they live in per-thread ``threading.local()`` storage. + # They will be forcefully closed when ``zmq_ctx.term()`` is called + # in step 7 below, which is the intended cleanup path. + cache: dict[str, zmq.Socket] = getattr(self._req_local, "cache", None) # type: ignore[assignment] + if cache: + for addr, sock in cache.items(): + try: + sock.close(linger=0) + except Exception: + pass # Best-effort; zmq_ctx.term() below will reclaim + cache.clear() + + # 6. Unregister memory from engine (if supported) + try: + if hasattr(self, "engine") and hasattr(self.engine, "unregister_memory"): + # Mooncake API only takes address, not size + self.engine.unregister_memory(self.base_ptr) + except Exception as e: + logger.warning(f"Failed to unregister memory: {e}") + + # 7. Close ZMQ contexts (also reclaims sockets cached in other threads) + try: + if hasattr(self, "zmq_ctx"): + self.zmq_ctx.term() + except Exception as e: + logger.warning(f"Failed to terminate ZMQ context: {e}") + + # 8. Release pool tensor reference (let GC handle actual deallocation) + # Note: We set to None instead of del to avoid AttributeError on repeated access + self.pool = None + + logger.info("MooncakeTransferEngineConnector closed.") + + # ------------------------------------------------------- + # Listener Logic + # ------------------------------------------------------- + + def _cleanup_stale_buffers(self) -> None: + """Reclaim buffers older than ``_BUFFER_TTL_SECONDS``. + Prevents permanent memory leaks when a receiver crashes or times out + without ever pulling the data. + TODO(wzliu): In extreme rare case, long transfer time, there might exist + TTL cleanup vs in-flight transfer conflict, which will be handled in the next PR. + """ + now = _time_mod.monotonic() + with self._local_buffers_lock: + stale_keys = [k for k, v in self._local_buffers.items() if now - v[5] > _BUFFER_TTL_SECONDS] + for k in stale_keys: + item = self._local_buffers.pop(k) + _, _, holder, should_release, _, _ = item + if should_release and isinstance(holder, ManagedBuffer): + holder.release() + logger.warning(f"TTL expired ({_BUFFER_TTL_SECONDS}s): cleaned up stale buffer for {k}") + + def _zmq_listener_loop(self): + socket = self.zmq_ctx.socket(zmq.ROUTER) + try: + socket.bind(f"tcp://{self.host}:{self.zmq_port}") + except zmq.ZMQError as exc: + # Any bind failure (EADDRINUSE, EADDRNOTAVAIL, EACCES, etc.) + # is fatal for a sender — fail fast so __init__ propagates the error. + # There is no silent receiver fallback; roles are explicitly assigned. + logger.error(f"ZMQ bind failed on {self.host}:{self.zmq_port}: {exc} (errno={exc.errno})") + self.can_put = False + self._bind_error = exc + self._listener_ready.set() + return # let __init__ propagate the error + + # Successfully bound - signal ready + self._listener_ready.set() + + # Create inproc socket pair for worker thread notifications + # This allows workers to wake up the listener immediately when done + notify_addr = f"inproc://notify-{id(self)}" + notify_recv = self.zmq_ctx.socket(zmq.PULL) + notify_recv.bind(notify_addr) + + poller = zmq.Poller() + poller.register(socket, zmq.POLLIN) + poller.register(notify_recv, zmq.POLLIN) + + # Response queue for thread-safe socket operations + response_queue: queue.Queue = queue.Queue() + + try: + while not self._stop_event.is_set(): + try: + # Poll for incoming requests or notifications (with timeout) + events = dict(poller.poll(1000)) + + # Process notifications (drain all) + if notify_recv in events: + while True: + try: + notify_recv.recv(zmq.NOBLOCK) + except zmq.Again: + break + + # Process any pending responses (non-blocking) + while True: + try: + identity, response = response_queue.get_nowait() + socket.send_multipart([identity, b"", response]) + except queue.Empty: + break + + # Periodic TTL cleanup (~every 10s) + now = _time_mod.monotonic() + if now - self._last_ttl_check >= 10.0: + self._last_ttl_check = now + self._cleanup_stale_buffers() + + # Process incoming requests + if socket in events: + frames = socket.recv_multipart() + if len(frames) >= 2: + # Submit to thread pool + self._sender_executor.submit( + self._handle_pull_request, + response_queue, + notify_addr, + frames[0], + frames[-1], + ) + except zmq.ContextTerminated: + break + except Exception: + logger.debug("Listener loop error", exc_info=True) + finally: + try: + notify_recv.close(linger=0) + socket.close(linger=0) + except Exception: + pass # Best-effort cleanup during listener shutdown + + def _handle_pull_request(self, response_queue: queue.Queue, notify_addr: str, identity, payload): + """ + Handle pull request or query request in worker thread. + Results are put into response_queue and listener is notified via inproc. + """ + try: + # Check if this is a query request + if payload.startswith(QUERY_INFO): + self._handle_query_request(response_queue, notify_addr, identity, payload[len(QUERY_INFO) :]) + return + + # Normal RDMA transfer request + meta = msgspec.msgpack.decode(payload, type=MooncakeAgentMetadata) + + with self._local_buffers_lock: + item = self._local_buffers.get(meta.request_id) + + if not item: + response_queue.put((identity, TRANS_ERROR)) + self._notify_listener(notify_addr) + return + + src_addrs, src_lengths, _, _, _, _ = item + remote_session = f"{meta.remote_hostname}:{meta.remote_port}" + + # RDMA Write + ret = self.engine.batch_transfer_sync_write(remote_session, src_addrs, meta.dst_addrs, src_lengths) + + if ret == 0: + self.cleanup(meta.request_id) + response_queue.put((identity, TRANS_DONE)) + else: + # Keep buffer in _local_buffers so receiver can retry on transient failures. + # Buffer will be cleaned up when: (a) a retry succeeds, (b) close() is called, + # or (c) a future TTL mechanism reclaims stale entries. + logger.warning( + f"RDMA write failed for {meta.request_id} to {remote_session} " + f"(ret={ret}). Buffer retained for retry." + ) + response_queue.put((identity, TRANS_ERROR)) + + except Exception as e: + logger.error(f"Push failed: {e}") + response_queue.put((identity, TRANS_ERROR)) + + # Notify listener thread that response is ready + self._notify_listener(notify_addr) + + def _handle_query_request(self, response_queue: queue.Queue, notify_addr: str, identity, payload): + """Handle metadata query request.""" + try: + query = msgspec.msgpack.decode(payload, type=QueryRequest) + + with self._local_buffers_lock: + item = self._local_buffers.get(query.request_id) + + if not item: + response_queue.put((identity, INFO_NOT_FOUND)) + else: + src_addrs, src_lengths, _, _, is_fast_path, _ = item + + resp = QueryResponse( + request_id=query.request_id, + data_size=src_lengths[0] if src_lengths else 0, + is_fast_path=is_fast_path, + ) + response_queue.put((identity, msgspec.msgpack.encode(resp))) + + except Exception as e: + logger.error(f"Query request failed: {e}") + response_queue.put((identity, INFO_NOT_FOUND)) + + self._notify_listener(notify_addr) + + def _notify_listener(self, notify_addr: str): + """Send notification to wake up listener thread. + + Uses a per-worker-thread cached PUSH socket to avoid the overhead + of creating/connecting/closing a socket on every call. + """ + try: + local = self._worker_local + sock = getattr(local, "notify_socket", None) + cached_addr = getattr(local, "notify_addr", None) + if sock is None or cached_addr != notify_addr: + # First call on this thread, or address changed – create socket + if sock is not None: + sock.close(linger=0) + sock = self.zmq_ctx.socket(zmq.PUSH) + sock.connect(notify_addr) + local.notify_socket = sock + local.notify_addr = notify_addr + sock.send(b"", zmq.NOBLOCK) + except Exception: + # Socket may be broken; clear cache so next call recreates it + local.notify_socket = None + local.notify_addr = None diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 341bbcdba..812d8b9d5 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -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} diff --git a/vllm_omni/distributed/omni_connectors/factory.py b/vllm_omni/distributed/omni_connectors/factory.py index 34415f7fd..a680da953 100644 --- a/vllm_omni/distributed/omni_connectors/factory.py +++ b/vllm_omni/distributed/omni_connectors/factory.py @@ -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) diff --git a/vllm_omni/distributed/omni_connectors/utils/config.py b/vllm_omni/distributed/omni_connectors/utils/config.py index 5556296d9..f18e52c4d 100644 --- a/vllm_omni/distributed/omni_connectors/utils/config.py +++ b/vllm_omni/distributed/omni_connectors/utils/config.py @@ -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 diff --git a/vllm_omni/distributed/omni_connectors/utils/initialization.py b/vllm_omni/distributed/omni_connectors/utils/initialization.py index 91322e63e..aaa222b4c 100644 --- a/vllm_omni/distributed/omni_connectors/utils/initialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/initialization.py @@ -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 diff --git a/vllm_omni/entrypoints/cli/logo.py b/vllm_omni/entrypoints/cli/logo.py new file mode 100644 index 000000000..addfe3593 --- /dev/null +++ b/vllm_omni/entrypoints/cli/logo.py @@ -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) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index a59cc5673..e28637f6a 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml index c19302c76..862e51b83 100644 --- a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml @@ -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 diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml index 68f5817a7..f5d87aece 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml @@ -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" diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml index 30e894ae6..a52b20fcc 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml @@ -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