[Ray Core] Adds in Google Cloud TPUs as a native Resource (#38669)

[architkulkarni]

Cherry-pick of #38669 to the release branch.

The description of that PR is copy-pasted below.

---

Why are these changes needed?

The issue below has more details, but at a high level this change addresses the feature request of adding in TPUs as a native resource within Ray.

To a user, these changes intend to shift from:

@ray.remote(resources={"TPU": 1})
def my_function():
    ...

to

@ray.remote(num_tpus=4)
def my_function():
    ...

Since we're adding TPUs as a native resource, this gives users the added ability to access individual TPU chips in a TPU host, e.g. the following would be valid:

@ray.remote(num_tpus=1) # or 2
def my_function():
    import jax
    print(jax.device_count()) -> 1 (or 2)

This is enabled by setting environment variables as described in jax-ml/jax#14977.

Short overview of changes:

• Add in autodetect num TPU chips + TPU version (+ tests for each)
• Adds in num_tpus in ray_params and does input sanitization similar to how GPUs function (e.g. num_tpus instead of TPU in resources)
• Adds in several constants
• Since there are a lot of changes related to neuron cores (from Auto-detection of accelerator_type for aws_accelerators trn1_inf #37998) I tried to consolidate the neuron core specific logic with TPU logic where applicable, for instance changing _validate_neuron_core_accelerator -> _validate_accelerator in ray_option_utils.py
• And added in a lot of tests!

Related issue number

#38085

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[zhe-thoughts]

This is a must-have feature in 2.7. But let's be careful about possible issues on the release branch