[Pallas] Introduce jax_import_guard

[alanwaketan]

Summary:
Importing JAX will lock the TPU devices and prevent any pytorxh/xla's TPU computations. To address it, we need to acquire the TPU first. Somehow xm.xla_device() is enough to acquire the TPU device.

Test Plan:
python test/test_pallas.py

[alanwaketan]

@will-cromar Do you have better ideas?

[JackCaoG]

torch_xla._XLAC._init_computation_client()

[alanwaketan]

When will this be called in general?

[JackCaoG]

we don't usually call this API directly I believe, runtime init usually happens when we try to get the device. However if all you need is to init the runtime this api is cleaner. I used it in https://github.com/pytorch/pytorch/blob/a04e7fca8eddde492b239da6ac23d6a056666a0e/benchmarks/dynamo/common.py#L91 as well

[miladm]

would this API work correctly on a multipod environment? @JackCaoG

[JackCaoG]

hmm it should. We just need to make sure that PyTorch/XLA init the runtime first and grab the libtpu, JAX init the runtime after pytorch doesn't cause any issue on single pod for us at least. We don't use JAX to execute any device program so it is OK.

In mutipod it sounds like init the runtime twice will cause issue? I never look into that too much

[miladm]

@jonb377 any thoughts on why this issue may happen on multi-pod?

[JackCaoG]

This is funny, @will-cromar and I were discusses how you handle the TPU ownership conflict between PyTorch/XLA and JAX this morning.

[will-cromar]

Will this break multiprocess? If you import this, then the current process will init the runtime, and you would not be able to call xmp.spawn after that point.

[alanwaketan]

Will this break multiprocess? If you import this, then the current process will init the runtime, and you would not be able to call xmp.spawn after that point.

What do you mean? xmp.spawn is supposed to be called before any torch-xla code, right? I hope they don't import pallas in the launcher...

[will-cromar]

What do you mean? xmp.spawn is supposed to be called before any torch-xla code, right? I hope they don't import pallas in the launcher...

This would be a problem if custom_kernels is imported at the global scope at all, e.g.

# Inits TPU
from torch_xla.experimental import custom_kernel

def main():
   # uses `custom_kernel`
   ...

if __name__ == "__main__":
  # Would fail because TPU is initialized
  xmp.spawn(main)

IMO it would be safer to import jax inside make_kernel_from_pallas where it is used.

[alanwaketan]

@will-cromar That works for me as well. Let me update it.

[JackCaoG]

hmm test failed with

Traceback (most recent call last):
  File "/tmp/pytorch/xla/test/test_pallas.py", line 10, in <module>
    from torch_xla.experimental.custom_kernel import jax_import_guard
  File "/opt/conda/lib/python3.8/site-packages/torch_xla-2.3.0+git22e6548-py3.8-linux-x86_64.egg/torch_xla/experimental/custom_kernel.py", line 18, in <module>
    import jax
ModuleNotFoundError: No module named 'jax'

[alanwaketan]

hmm test failed with

Traceback (most recent call last):
  File "/tmp/pytorch/xla/test/test_pallas.py", line 10, in <module>
    from torch_xla.experimental.custom_kernel import jax_import_guard
  File "/opt/conda/lib/python3.8/site-packages/torch_xla-2.3.0+git22e6548-py3.8-linux-x86_64.egg/torch_xla/experimental/custom_kernel.py", line 18, in <module>
    import jax
ModuleNotFoundError: No module named 'jax'

Yea, the CPU and GPU CI is not configured with JAX installed. Will fix it now.

[alanwaketan]

@JackCaoG @will-cromar I think this is ready for the new round of reviews.