add unit test + fix unit tests

Signed-off-by: Timothy Seah <[email protected]>

Author: TimothySeah

python/ray/train/constants.py

@@ -144,7 +144,6 @@ def _v2_migration_warnings_enabled() -> bool:
     RAY_TRAIN_ENABLE_STATE_TRACKING,
     TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE,
     TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S,
-    DEFAULT_TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S,
 }
 
 # Key for AIR Checkpoint metadata in TrainingResult metadata

python/ray/train/tests/test_backend.py

@@ -28,6 +28,7 @@
 from ray.train.constants import (
     ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV,
     ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV,
+    TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S,
     TRAIN_ENABLE_WORKER_SPREAD_ENV,
 )
 from ray.train.torch import TorchConfig
@@ -364,6 +365,24 @@ def check_process_group():
     assert not any(e.finish_training())
 
 
+@pytest.mark.parametrize(
+    "init_method, timeout_s", [("env", 5), ("tcp", 5), ("env", 0), ("tcp", 0)]
+)
+def test_torch_process_group_shutdown_timeout(
+    ray_start_2_cpus, monkeypatch, init_method, timeout_s
+):
+    monkeypatch.setenv(TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S, timeout_s)
+    torch_config = TorchConfig(backend="gloo", init_method=init_method)
+    e = BackendExecutor(torch_config, num_workers=2)
+    e.start()
+
+    _start_training(e, lambda: 1)
+    assert e.finish_training() == [1, 1]
+
+    # Verify that we do not raise an exception even if we time out
+    e._backend.on_shutdown(e.worker_group, e._backend_config)
+
+
 @pytest.mark.parametrize(
     "worker_results",
     [