handle exception waiting for work (meta-pytorch#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

Reviewed By: d4l3k

Differential Revision: D84880993

Author: tushar00jain

.github/workflows/docs.yaml

@@ -25,6 +25,7 @@ jobs:
 
           sudo apt-get install -y protobuf-compiler
 
+          pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
           pip install .[dev] -v
 
           pip install -r docs/requirements.txt

torchft/device_mesh.py

@@ -10,6 +10,7 @@
     init_device_mesh,
     ProcessGroup as BaseProcessGroup,
 )
+from torch.distributed._mesh_layout import _MeshLayout
 from torch.distributed.tensor.device_mesh import _mesh_resources
 
 from torchft.manager import Manager
@@ -69,12 +70,20 @@ 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
+            mesh_tensor = (
+                mesh.detach().to(dtype=torch.int).contiguous()
+                if isinstance(mesh, torch.Tensor)
+                else torch.tensor(mesh, device="cpu", dtype=torch.int)
+            )
+            self._layout = _MeshLayout(mesh_tensor.size(), mesh_tensor.stride())
         else:
             assert parent is not None
             self._device_type = parent.device_type
+            self._layout = parent._layout
         self._flatten_mesh_list: tuple[DeviceMesh, ...] = tuple()
         self._thread_id: Optional[int] = None
         self._hash: Optional[int] = None

torchft/manager.py

@@ -1253,14 +1253,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()