[core][RDT] Fix nixl garbage collection after the object is freed

[Qiaolin-Yu]

Why are these changes needed?

Previously, we didn't deregister the tensor from nixl_agent after the tensor has been freed, this will cause error in nixl agent.
In this pr, we invalidate the metadata after the tensor is freed.

Note that this pr is based on #57112, so it should be merged after that.

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run pre-commit jobs to lint the changes in this PR. (pre-commit setup)
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[gemini-code-assist]

Code Review

This pull request refactors the GPU object freeing mechanism by replacing the FreeActorObject RPC with a callback-based approach, which is a good simplification. It also introduces a garbage_collect method to the tensor transport interface to handle backend-specific cleanup, specifically for nixl to deregister memory. The changes are well-structured, but I've identified a potential AttributeError for non-NIXL transports and a minor improvement for exception handling. My review includes suggestions to address these points.

[dayshah]

i'm a little confused on why we would remove the agent but not deregister the memory if the send is sync

we'll re-register the memory on every send anyways

[Qiaolin-Yu]

The deregister_memory should be called by the same agent which calls register_memory. In our case, it should be called by the sender, I added it in the gc function.

[stephanie-wang]

This looks good assuming it passes manual testing for now :s

I just realized, though - I think there can still be a race condition since the GC runs asynchronously. Could try to fix it with a wait_tensor_free call when the user returns a tensor with NIXL transport?

[dayshah]

I just realized, though - I think there can still be a race condition since the GC runs asynchronously. Could try to fix it with a wait_tensor_free call when the user returns a tensor with NIXL transport?

Confused on what the race is, the NIXL recv is sync right, so as long as the borrower / consumer isn't done with the ref we won't GC/deregister on the sender.

[stephanie-wang]

I just realized, though - I think there can still be a race condition since the GC runs asynchronously. Could try to fix it with a wait_tensor_free call when the user returns a tensor with NIXL transport?

Confused on what the race is, the NIXL recv is sync right, so as long as the borrower / consumer isn't done with the ref we won't GC/deregister on the sender.

Ah as I understood it, the bug here is a bit different from the NCCL one. It happens because NIXL doesn't allow the same memory region to be registered twice. It was happening if we allocate a tensor on the sender, register it, then before we can GC it, torch allocates the same physical memory to another tensor.

But actually I think I was wrong, seems like this should work because torch shouldn't allocate the memory to another tensor until GC runs. So ignore me :D