Fix bad error message when PjRtComputationClient throws exception

[will-cromar]

The g_computation_client_initialized doesn't get reset in this case, leading to incorrect behavior in GetComputationClientIfInitialized. Make CreateClient anonymous so we can't call it twice by accident.

Example from @jonb377:

ptxla@t1v-n-f494979e-w-0:/workspaces/work/pytorch/xla$ python -c 'import torch_xla; torch_xla._XLAC._xla_get_runtime_devices()'
Traceback (most recent call last):
  File "<string>", line 1, in <module>
RuntimeError: torch_xla/csrc/runtime/runtime.cc:27 : $PJRT_DEVICE is not set.

Error in atexit._run_exitfuncs:
Traceback (most recent call last):
  File "/workspaces/work/pytorch/xla/torch_xla/__init__.py", line 127, in _prepare_to_exit
    _XLAC._prepare_to_exit()
RuntimeError: torch_xla/csrc/runtime/runtime.cc:17 : Check failed: !was_initialized 
*** Begin stack trace ***
        tsl::CurrentStackTrace[abi:cxx11]()

        torch_xla::runtime::GetComputationClient()


        PyCFunction_Call
        _PyObject_MakeTpCall
        _PyEval_EvalFrameDefault
        _PyFunction_Vectorcall
        PyVectorcall_Call

        Py_FinalizeEx
        Py_RunMain
        Py_BytesMain
        __libc_start_main
        _start
*** End stack trace ***
ComputationClient already initialized

ComputationClient was never initialized, so it doesn't make sense to print ComputationClient already initialized.

[JackCaoG]

what would the error message be after your change?

[jonb377]

Thanks Will!

[will-cromar]

what would the error message be after your change?

It would just be the first error, $PJRT_DEVICE is not set.

[will-cromar]

There seems to be a consistent error during teardown during the GPU tests. Creating a GPU dev environment to see what's going on

[will-cromar]

Setting up a GPU VM is taking too long, so I came up with a better idea. We don't have to carefully prevent CreateClient from being called twice if we just make it anonymous and thus impossible to call twice.

Let's see if this works on GPU...

[will-cromar]

Some of this GPU tests are starting to pass, so it looks like this tweak works.

[jonb377]

What's the difference between this and the original approach, since both rely on the static initializer? It seems like the logic has just been moved out of CreateClient and into the lambda.

[will-cromar]

Yeah, that's right. This lets me skip checking the case where this function gets called twice, the handling of which causes the bad error message. C++11 statics will ensure that it only completes once, and the function is anonymous now, which will prevent future code from erroneously calling CreateClient somewhere else in the future.

[jonb377]

Ah I see, the lambda is guaranteed to only ever be called once. Do we know how CreateClient was called twice in the first place?

[will-cromar]

I added that check to defend against future errors (likely by future me). There was never a case in the original code where CreateClient actually completed twice or was called concurrently.

The bad error messaging occurred when the constructor of PjRtComputationClient threw an exception. g_computation_client_initialized was never reset to false, GetComputationClientIfInitialized calls GetComputationClient during teardown, which re-runs CreateClient, which complains because g_computation_client_initialized never got reset.

In this PR, I only set g_computation_client_initialized after the actual runtime init completed without breaking, since there are guaranteed to be no other concurrent callers of CreateClient.