Attach callback to a Tensor's underlying buffer

[will-cromar]

• New experimental high-level API torch_xla.experimental.callback.on_ready_callback to install callback ((Tensor) -> None) on Tensor, to be called when underlying buffer is ready
  • If Tensor is not assigned a buffer, throw an error
  • @JackCaoG and I chose this route instead of an explicit Future because Tensors are already async results in our case, and returning Futures would require expanding the API in several places.
• Low-level API takes void() to make Pybind plumbing easier
• Since this is a Python callable to be called asynchronously from C++, there is high potential for deadlocks, as callback does need to reacquire the GIL from the main thread
  • I tried to fix any cases I could find. If a deadlock happens to you, a thread apply all bt from gdb would be extremely helpful for me to debug.
• PrepareToExit now waits for pending runtime operations to complete. Otherwise, we start destroying the runtime client before operations are done. If the main thread is waiting for a callback to complete, a deadlock occurs.

[JackCaoG]

is it difficult to support the cases where call_back takes more than just the tensor value? call_back might be a general python function that tensor is just one of the inputs.

[will-cromar]

You should functools.partial in that case to add other arguments to the callable https://docs.python.org/3/library/functools.html#functools.partial

[JackCaoG]

mostly lgtm, minor nits