[Misc] Optimize ray worker initialization time

[ruisearch42]

This PR optimizes ray worker initialization time.

In the current code base, ray.get(worker.get_node_ip.remote()) is called for each worker right after we get its handle, and it takes ~3s. This call is expensive because when RayWorkerWrapper.remote() just returns, we get an actor handle, but the actor itself may not be fully initialized yet. At this time, any method call on the actor would need to wait for actor initialization to happen, which can take some time (~3s in this case).

And since we are calling ray.get(worker.get_node_ip.remote()) in a serialized manner for each newly created actor handle, this time adds up. For example, when we have TP=4, this would take ~12 seconds.

We optimize this by making ray.get(worker.get_node_ip.remote()) calls on all the actor handles after they are created. And since these run in parallel, the total time taken is ~3s. So for TP = 4, this reduces ~9 seconds.

I tested the following command:

python3 benchmarks/benchmark_latency.py --model meta-llama/Llama-3.1-8B-Instruct --tensor-parallel-size 4  --num-iters-warmup 5 --num-iters 20  --batch-size 8 --input-len 128 --output-len 256 --max-model-len 2048 --no-enable-prefix-caching --distributed-executor-backend ray

Without this PR, _init_workers_ray takes ~18 seconds. And with it, it takes ~9 seconds.

FIX #10283

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[comaniac]

LGTM

[youkaichao]

thanks for the fix!

[youkaichao]

@ruisearch42 this one looks concerning to me. we should change the tuple to sort, instead of using worker as the key. see the code from #11256

[ruisearch42]

I see. Can you elaborate a bit on the concern? The pattern of using an external dict for sorting is not uncommon.

[youkaichao]

using an arbitrary python object as a key introduces quite unpredictable behavior and can have silent bugs.

[youkaichao]

it's not about using an external dict, it's about using the worker object as a dict key, which implicitly calls its __hash__ function.

[ruisearch42]

I think the default behavior without a custom __hash__ function is to use the object's identity (memory address) as __hash__ and __eq__, so it's pretty safe unless there is some non-standard user overridden __hash__ and __eq__?

I think your implementation also makes sense.