PJRT Computation Client Teardown Function #9679
Description
❓ Questions and Help
Is there a teardown function that can be hooked from PJRT Plugin implementers for system teardown purposes? For example, graceful device closure at session termination?
It seems like the PJRT Computation Client is instantiated with a leaky singleton pattern, so its destructor is not called, and we cannot leverage our PJRT Client's destructor.
Is there some client shutdown hook that can be used? It seems like PJRT_Client_Destroy would be a suitable candidate, except that I don't see it ever being called from pytorch/xla.
The reason for this is that we would like to have some automatic device cleanup / other system resource teardown implemented in our plugin that triggers at the end of a session. It would also be nice to have a user-accessible API that permits session teardown within PJRT, for example to reset devices between pytests within the same process.