[PJRT] Use PJRT by default if XRT is not configured

[will-cromar]

If you don't set XRT configuration today, we print a warning that there is no XLA configuration. Instead, let's set the default device to PJRT:CPU or PJRT:TPU, depending on what is available.

Existing workloads that already set up XRT configuration will be unaffected by this change. However, we'll still print a warning for now to let the user know that they're using PJRT instead of XRT.

Example output:

# With libtpu installed:
$ python -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices('TPU'))"
WARNING:torch_xla:XRT configuration not detected. Defaulting to preview PJRT runtime. To silence this warning and continue using PJRT, explicitly set PJRT_DEVICE to a supported device. To use XRT, set any of the following environment variables: ['XRT_TPU_CONFIG', 'XRT_DEVICE_MAP', 'XRT_WORKERS']
WARNING:torch_xla:For more information about the status of PJRT, see https://github.com/pytorch/xla/blob/master/docs/pjrt.md
WARNING:torch_xla:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
['xla:0', 'xla:1', 'xla:2', 'xla:3']
$ PJRT_DEVICE=TPU python -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices('TPU'))"
['xla:0', 'xla:1', 'xla:2', 'xla:3']

# With libtpu uninstalled
$ python -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices('CPU'))"
WARNING:torch_xla:XRT configuration not detected. Defaulting to preview PJRT runtime. To silence this warning and continue using PJRT, explicitly set PJRT_DEVICE to a supported device. To use XRT, set any of the following environment variables: ['XRT_TPU_CONFIG', 'XRT_DEVICE_MAP', 'XRT_WORKERS']
WARNING:torch_xla:For more information about the status of PJRT, see https://github.com/pytorch/xla/blob/master/docs/pjrt.md
WARNING:torch_xla:Defaulting to PJRT_DEVICE=CPU
['xla:0']

[JackCaoG]

currently if user sets GPU_NUM_DEVICES it will default to xrt:gpu, user have to overwrite that with additional PJRT_DEVICE=GPU. I guess we intentionally don't change that behavior there?

[will-cromar]

Right. I left GPU out intentionally because we don't have the same easy way of knowing that the user intended XRT, and there are still some feature gaps on PJRT GPU.