Ray Train test_jax_trainer::test_minimal_multihost Flaky Test Fix (ray-project#56548)

A fix that addresses the failing flaky test
`test_jax_trainer.py::test_minimal_multihost`.

https://buildkite.com/ray-project/postmerge/builds/12941#01993f89-cc62-4e31-8de2-8b18f81ac177

Issue:
The `test_minimal_multihost` introduces a race condition by attempting
to initialize a virtualenv directory twice at the same directory path
during worker runtime environment setup. This test would not fail in a
true multi-host environment, but the tests simulate a multi-host
environment on a singular device. This might be a ray core issue
resulting in errors on runtime _env, but this PR will at least unblock
the test so it is no longer flaky.

Fix:
Move `worker_runtime_env` to the job level so that the `pip install jax`
only happens once

---------

Signed-off-by: JasonLi1909 <jasli1909@gmail.com>
Signed-off-by: Marco Stephan <marco@magic.dev>

Author: JasonLi1909

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

@@ -3,12 +3,27 @@
 import ray
 from ray.tests.conftest import _ray_start_cluster
 from ray.train import RunConfig, ScalingConfig
-from ray.train.v2._internal.constants import HEALTH_CHECK_INTERVAL_S_ENV_VAR
+from ray.train.v2._internal.constants import (
+    HEALTH_CHECK_INTERVAL_S_ENV_VAR,
+    is_v2_enabled,
+)
 from ray.train.v2.jax import JaxTrainer
 
+assert is_v2_enabled()
+
+
+@pytest.fixture
+def jax_runtime_env():
+    return {
+        "pip": ["jax"],
+        "env_vars": {
+            "JAX_PLATFORMS": "cpu",
+        },
+    }
+
 
 @pytest.fixture
-def ray_tpu_single_host(monkeypatch):
+def ray_tpu_single_host(monkeypatch, jax_runtime_env):
     """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")
@@ -19,14 +34,14 @@ def ray_tpu_single_host(monkeypatch):
             resources={"TPU": 8},
         )
 
-        ray.init(address=cluster.address)
+        ray.init(address=cluster.address, runtime_env=jax_runtime_env)
 
         yield cluster
         ray.shutdown()
 
 
 @pytest.fixture
-def ray_tpu_multi_host(monkeypatch):
+def ray_tpu_multi_host(monkeypatch, jax_runtime_env):
     """Start a simulated multi-host TPU Ray cluster."""
     with _ray_start_cluster() as cluster:
         monkeypatch.setenv("TPU_NAME", "test-slice-1")
@@ -44,7 +59,7 @@ def ray_tpu_multi_host(monkeypatch):
             resources={"TPU": 4},
         )
 
-        ray.init(address=cluster.address)
+        ray.init(address=cluster.address, runtime_env=jax_runtime_env)
 
         yield cluster
         ray.shutdown()
@@ -78,12 +93,6 @@ def test_minimal_singlehost(ray_tpu_single_host, tmp_path):
         ),
         run_config=RunConfig(
             storage_path=str(tmp_path),
-            worker_runtime_env={
-                "pip": ["jax"],
-                "env_vars": {
-                    "JAX_PLATFORMS": "cpu",
-                },
-            },
         ),
     )
     result = trainer.fit()
@@ -109,12 +118,6 @@ def test_minimal_multihost(ray_tpu_multi_host, tmp_path):
         ),
         run_config=RunConfig(
             storage_path=str(tmp_path),
-            worker_runtime_env={
-                "pip": ["jax"],
-                "env_vars": {
-                    "JAX_PLATFORMS": "cpu",
-                },
-            },
         ),
     )
     result = trainer.fit()