[P/D] Support NIXL connector to disconnect during a clean shutdown

[chaunceyjiang]

Purpose

Part of #22295

Refer https://github.com/ai-dynamo/nixl/blob/main/examples/python/nixl_api_example.py#L205-L211

Test Plan

[gemini-code-assist]

Code Review

This pull request introduces an explicit shutdown method for the NixlConnector to ensure proper resource cleanup, moving away from relying on __del__. The new method handles the shutdown of background threads and releases NIXL resources like dlist handles, remote agents, and registered memory. This is a significant improvement for ensuring a clean shutdown. My main feedback is regarding the shutdown of the _nixl_handshake_listener_t thread, which could be made more graceful.

[chaunceyjiang]

@njhill PTAL

[njhill]

cc @wseaton @markmc

[wseaton]

Thanks @chaunceyjiang, this seems reasonable (saw it plugs into the shutdown API that you just added in another PR). Have you observed what happens with this change when a prefill node is killed or scaled after handshaking with many decode nodes? Do you see disconnect logs in the decode peers? Any ill effects?

[chaunceyjiang]

Hi @wseaton
The primary purpose of this PR is to trigger a call to nixl.backendEngines.disconnect().

Based on this PR, I tested two scenarios:

1. When killing the prefill node, the decode node enters an error callback, such as: UCX error handling callback was invoked with status -25 (Connection reset by remote peer). When the prefill node restarts, it establishes a new connection.
2. When killing the decode node, the prefill node shows no change.

I also noticed that all nixl test examples explicitly call disconnect(). Therefore, I believe it is reasonable to maintain consistent behavior.

[markmc]

I'm not clear which cleanup call is doing more than just releasing memory that would be freed by when objects are finalized anyway ... is it remove_remote_agent() does a disconnect?

I guess it's fine to follow the recommended API usage ... but it looks like we're missing doing release_xfer_handle() on any in-progress _recving_transfers ?

[markmc]

Hi @wseaton The primary purpose of this PR is to trigger a call to nixl.backendEngines.disconnect().

Would you mind explaining how exactly the changes trigger this call? Thanks in advance

[chaunceyjiang]

is it remove_remote_agent() does a disconnect?

@markmc Yes . If I’m not mistaken, it should be like this.
remove_remote_agent() -> invalidateRemoteMD() -> backendEngines[it.first]->disconnect() -> ucx.disconnect()

[NickLucche]

Thanks for adding this @chaunceyjiang !
I think this is a good practice change, in particular as we expand Nixl to support different backends other than UCX, which may require specific resource cleanup.

While we're at it though, as we're already releasing data structures memory nixl-side, I think we should also release on-going transfers prior to that, as @markmc pointed out.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @chaunceyjiang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[chaunceyjiang]

/cc @NickLucche PTAL.

[markmc]

Thank you for explaining that remove_remote_agent() calling ucx->disconnect() is the main purpose of this PR

I think we should ensure that shutdown() is idempotent, and it's probably good to test this also. I've pushed a commit to do this here: chaunceyjiang/vllm@nixl_shutdown...markmc:vllm:nixl_shutdown

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @chaunceyjiang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[NickLucche]

LGTM , thanks for adding tests as well @chaunceyjiang !