[train][jax] Skip py3.12 for jax test (ray-project#58979)

1. Jax dependency is introduced in
ray-project#58322
2. The current test environment is for CUDA 12.1, which limit jax
version below 0.4.14.
3. jax <= 0.4.14 does not support py 3.12.
4. skip jax test if it runs against py3.12+.

Signed-off-by: Lehui Liu <lehui@anyscale.com>

Author: liulehui

python/ray/train/v2/tests/test_jax_gpu.py

@@ -19,6 +19,10 @@ def reduce_health_check_interval(monkeypatch):
 
 
 @pytest.mark.skipif(sys.platform == "darwin", reason="JAX GPU not supported on macOS")
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12),
+    reason="Current jax version is not supported in python 3.12+",
+)
 def test_jax_distributed_gpu_training(ray_start_4_cpus_2_gpus, tmp_path):
     """Test multi-GPU JAX distributed training.
 

python/ray/train/v2/tests/test_jax_trainer.py

@@ -1,3 +1,5 @@
+import sys
+
 import pytest
 
 import ray
@@ -71,6 +73,10 @@ def train_func():
     train.report({"result": [str(d) for d in devices]})
 
 
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12),
+    reason="Current jax version is not supported in python 3.12+",
+)
 def test_minimal_singlehost(ray_tpu_single_host, tmp_path):
     trainer = JaxTrainer(
         train_loop_per_worker=train_func,
@@ -101,6 +107,10 @@ def test_minimal_singlehost(ray_tpu_single_host, tmp_path):
     assert len(labeled_nodes) == 1
 
 
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12),
+    reason="Current jax version is not supported in python 3.12+",
+)
 def test_minimal_multihost(ray_tpu_multi_host, tmp_path):
     trainer = JaxTrainer(
         train_loop_per_worker=train_func,