Revising test_jax_trainer flaky test (ray-project#56854)

Revisiting ray-project#56548 as test
continues to be flaky on CI

**Solution**: The previous attempt to deflake this test still used a
`pip install jax` via the `ray.init` runtime_env args. Hence, the pip
install related error persisted. This PR instead adds `jax` and `jaxlib`
as a dependency of CI train tests, avoiding the need to `pip install
jax` via the runtime_env.

---------

Signed-off-by: JasonLi1909 <jasli1909@gmail.com>
Signed-off-by: Jason Li <57246540+JasonLi1909@users.noreply.github.com>

Author: JasonLi1909

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

@@ -13,17 +13,7 @@
 
 
 @pytest.fixture
-def jax_runtime_env():
-    return {
-        "pip": ["jax"],
-        "env_vars": {
-            "JAX_PLATFORMS": "cpu",
-        },
-    }
-
-
-@pytest.fixture
-def ray_tpu_single_host(monkeypatch, jax_runtime_env):
+def ray_tpu_single_host(monkeypatch):
     """Start a mock single-host TPU Ray cluster with 2x4 v6e (8 chips per host)."""
     with _ray_start_cluster() as cluster:
         monkeypatch.setenv("TPU_ACCELERATOR_TYPE", "v6e-8")
@@ -34,14 +24,14 @@ def ray_tpu_single_host(monkeypatch, jax_runtime_env):
             resources={"TPU": 8},
         )
 
-        ray.init(address=cluster.address, runtime_env=jax_runtime_env)
+        ray.init(address=cluster.address)
 
         yield cluster
         ray.shutdown()
 
 
 @pytest.fixture
-def ray_tpu_multi_host(monkeypatch, jax_runtime_env):
+def ray_tpu_multi_host(monkeypatch):
     """Start a simulated multi-host TPU Ray cluster."""
     with _ray_start_cluster() as cluster:
         monkeypatch.setenv("TPU_NAME", "test-slice-1")
@@ -59,7 +49,7 @@ def ray_tpu_multi_host(monkeypatch, jax_runtime_env):
             resources={"TPU": 4},
         )
 
-        ray.init(address=cluster.address, runtime_env=jax_runtime_env)
+        ray.init(address=cluster.address)
 
         yield cluster
         ray.shutdown()

python/requirements/ml/train-test-requirements.txt

@@ -1,4 +1,6 @@
 evaluate==0.4.3
 mosaicml; python_version < "3.12"
 sentencepiece==0.1.96
+jax==0.4.25
+jaxlib==0.4.25
 s3torchconnector==1.4.3

python/requirements_compiled.txt

@@ -859,6 +859,10 @@ isoduration==20.11.0
     # via jsonschema
 itsdangerous==2.1.2
     # via flask
+jax==0.4.25
+    # via -r python/requirements/ml/train-test-requirements.txt
+jaxlib==0.4.25
+    # via -r python/requirements/ml/train-test-requirements.txt
 jedi==0.19.1
     # via ipython
 jinja2==3.1.6
@@ -1081,7 +1085,10 @@ mistune==0.8.4
 ml-collections==0.1.1
     # via open-spiel
 ml-dtypes==0.3.2
-    # via tensorflow
+    # via
+    #   jax
+    #   jaxlib
+    #   tensorflow
 mlagents-envs==0.28.0
     # via -r python/requirements/ml/rllib-test-requirements.txt
 mlflow==2.22.0
@@ -1253,6 +1260,8 @@ numpy==1.26.4
     #   hpbandster
     #   hyperopt
     #   imageio
+    #   jax
+    #   jaxlib
     #   labmaze
     #   lightgbm
     #   matplotlib
@@ -1379,6 +1388,7 @@ opentelemetry-util-http==0.55b1
     #   opentelemetry-instrumentation-fastapi
 opt-einsum==3.3.0
     # via
+    #   jax
     #   pyro-ppl
     #   tensorflow
 optuna==4.1.0
@@ -2024,6 +2034,8 @@ scipy==1.11.4
     #   gpy
     #   hpbandster
     #   hyperopt
+    #   jax
+    #   jaxlib
     #   lightgbm
     #   linear-operator
     #   mlflow