# demo_vllm.py
from unittest.mock import patch
from accelerate.utils import gather_object
from accelerate import Accelerator
from vllm import LLM


def main():
    accelerator = Accelerator()
    if accelerator.is_main_process:
        # vLLM is not compatible with accelerate. So we need to patch it to make sure we can (1) place the vLLM model
        # on the desired device (world_size_patch) and (2) avoid a test that is not designed for our setting (profiling_patch).
        world_size_patch = patch("torch.distributed.get_world_size", return_value=1)
        profiling_patch = patch("vllm.worker.worker.Worker._assert_memory_footprint_increased_during_profiling", return_value=None)
        with world_size_patch, profiling_patch:
            LLM(model="Qwen/Qwen2.5-1.5B", device="cuda:7")

    # When using vLLM, the main process is responsible for loading the model weights. This can cause process desynchronization and seems
    # to lead to DeepSpeed hanging during initialization. To prevent this, we  synchronize all processes after vLLM has been fully initialized.
    accelerator.wait_for_everyone()

    prompts_text = ["Some text"]
    gather_object(prompts_text)  # it hangs here
    print("after gather")


if __name__ == "__main__":
    trainer = main()