Catching ChannelErrors on SFTP Failures (#1245)

mosaicml · Jul 1, 2022 · ef9d1d2 · ef9d1d2
1 parent ff0c9b5
commit ef9d1d2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/composer/utils/object_store/sftp_object_store.py b/composer/utils/object_store/sftp_object_store.py
@@ -169,7 +169,7 @@ def get_object_size(self, object_name: str) -> int:
 
     @contextlib.contextmanager
     def _handle_transient_errors(self):
-        from paramiko import SSHException
+        from paramiko import ChannelException, SSHException
         try:
             yield
         except Exception as e:
@@ -182,7 +182,7 @@ def _handle_transient_errors(self):
             if isinstance(e, SSHException):
                 if 'Server connection dropped:' in str(e):
                     raise ObjectStoreTransientError from e
-            if isinstance(e, (TimeoutError, ConnectionError, EOFError)):
+            if isinstance(e, (TimeoutError, ConnectionError, EOFError, ChannelException)):
                 raise ObjectStoreTransientError from e
             raise e