Error Handling: refactor ComputationClient::TransferFromDevice to propagate status.

[ysiraichi]

This PR makes 2 main changes in order to standardize and improve error handling:

• ComputationClient::TransferFromDevice returns a StatusOr<T> instance
• Wrap xla::PjRtLoadedExecutable::Execute(Sharded) with GetValueOrThrow

These changes mainly affect the errors whenever an OOM occurs. The second one targets eager mode. As an example, the following is the result of running the file below (without eager mode):

device = torch_xla.device()
a = torch.rand(1024, 1024, 1024, 1024, 1024, device=device)
b = a.sum()
print(b)

Before this PR:

F0000 00:00:1751368998.014824    2835 pjrt_computation_client.cpp:525] Non-OK-status: status
Status: INTERNAL: Error preparing computation: Out of memory allocating 4503599761588224 bytes.
*** Begin stack trace ***
        tsl::CurrentStackTrace[abi:cxx11]()
        torch_xla::runtime::PjRtComputationClient::TransferFromDevice(absl::lts_20230802::Span<std::shared_ptr<torch_xla::runtime::ComputationClient::Data> const>)
        ...
        _start
*** End stack trace ***
Failed to await future from buffer to literal inTransferFromDevice
*** Check failure stack trace: ***
    @     0x7ddc923438f9  absl::lts_20230802::log_internal::LogMessage::PrepareToDie()
    ...
Aborted (core dumped)

After this PR:

Traceback (most recent call last):
  File "examples/mem.py", line 11, in <module>
    print(b)
  File "torch/_tensor.py", line 590, in __repr__
    return torch._tensor_str._str(self, tensor_contents=tensor_contents)
  File "torch/_tensor_str.py", line 726, in _str
    return _str_intern(self, tensor_contents=tensor_contents)
  File "torch/_tensor_str.py", line 462, in _str_intern
    self = self.to("cpu")
RuntimeError: Error preparing computation: Out of memory allocating 4503599627370496 bytes.

(with XLA_SHOW_CPP_ERROR_CONTEXT=1)

RuntimeError: Error preparing computation: Out of memory allocating 4503599627370496 bytes. (at torch_xla/csrc/runtime/pjrt_computation_client.cpp:524)

[zhanyong-wan]

Very nice!

Can you add a python test to ensure that OOM does result in a python exception with the expected error message as opposed to crashing?

[zhanyong-wan]

Great!

[zhanyong-wan]

* => +