handle exception waiting for work (#287)

Summary:

work.wait() can throw so wrap that in a try/catch to handle it gracefully by reporting error to the manager, leading the should_commit to fail

Differential Revision: D84880993

Author: tushar00jain

torchft/device_mesh.py

@@ -69,6 +69,7 @@ def __init__(
         self.replicate_dim_name: str = mesh_dim_names[replicate_dim]
         self.parent = parent
         self.flatten_meshes: Dict[str, DeviceMesh] = {}
+        self._flatten_mapping: Dict[str, "DeviceMesh"] = {}
         self._device_type: str
         if mesh is not None:
             self._device_type = mesh.device_type

torchft/manager.py

@@ -1244,14 +1244,19 @@ def _assert_same_stream(self) -> None:
     def wait(self, timeout: Optional[timedelta] = None) -> bool:
         self._assert_same_stream()
 
-        with get_stream_context(self._stream):
-            self._work.wait()
-            self._set_future_callback()
+        try:
+            with get_stream_context(self._stream):
+                self._work.wait()
+                self._set_future_callback()
 
-        with get_stream_context(self._stream):
-            self._managed_fut_tail.wait()
+            with get_stream_context(self._stream):
+                self._managed_fut_tail.wait()
 
-        return True
+            return True
+        except Exception as e:
+            self._manager._logger.exception(f"got exception waiting for work {e}")
+            self._manager.report_error(e)
+            return False
 
     def block_current_stream(self, timeout: Optional[timedelta] = None) -> None:
         self._assert_same_stream()