[train] Add try-except for pg.wait() (#60743)

## Description
placement_group() is async and returns a handle immediately, If the PG
gets removed between creation request and readiness, `pg.wait()` throws
“not found” during wait, which in turn creates a controller error, we
should treat this as WGStartupTimeoutError.

## Related issues
[#870](anyscale#870)

---------

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

Author: liulehui

python/ray/train/v2/_internal/execution/worker_group/placement_group_handle.py

@@ -1,3 +1,4 @@
+import logging
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Union
 
@@ -7,6 +8,8 @@
 if TYPE_CHECKING:
     from ray.util.tpu import SlicePlacementGroup
 
+logger = logging.getLogger(__name__)
+
 
 class PlacementGroupHandle(ABC):
     """Unified interface for placement groups in Ray Train.
@@ -63,7 +66,14 @@ def ready(self) -> ObjectRef:
         return self._pg.ready()
 
     def wait(self, timeout_seconds: Union[float, int] = 30) -> bool:
-        return self._pg.wait(timeout_seconds)
+        try:
+            return self._pg.wait(timeout_seconds)
+        except Exception:
+            logger.warning(
+                "Placement group wait failed; treating as not ready.",
+                exc_info=True,
+            )
+            return False
 
     def shutdown(self) -> None:
         remove_placement_group(self._pg)
@@ -83,7 +93,14 @@ def ready(self) -> ObjectRef:
         return self._spg.placement_group.ready()
 
     def wait(self, timeout_seconds: Union[float, int] = 30) -> bool:
-        return self._spg.placement_group.wait(timeout_seconds)
+        try:
+            return self._spg.placement_group.wait(timeout_seconds)
+        except Exception:
+            logger.warning(
+                "Slice placement group wait failed; treating as not ready.",
+                exc_info=True,
+            )
+            return False
 
     def shutdown(self) -> None:
         self._spg.shutdown()

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

@@ -56,6 +56,16 @@ def test_default_handle_wait():
     handle.shutdown()
 
 
+def test_default_handle_wait_not_found_returns_false():
+    """wait() should return False if the placement group no longer exists."""
+    mock_pg = MagicMock()
+    mock_pg.wait.side_effect = Exception(
+        "Placement group PlacementGroupID(abc) does not exist."
+    )
+    handle = DefaultPlacementGroupHandle(mock_pg)
+    assert handle.wait(timeout_seconds=0) is False
+
+
 def test_default_handle_shutdown():
     """shutdown() should remove the placement group."""
     pg = placement_group([{"CPU": 1}])