From 330d3a688facce39a158d6c4f84e511f3b041307 Mon Sep 17 00:00:00 2001
From: Ganda <70550963+gandatchabana@users.noreply.github.com>
Date: Wed, 23 Aug 2023 16:14:09 -0400
Subject: [PATCH] Catch and handle Paramiko SSH errors (#48)

* Added a custom Exception class SSHConnectionError.
* Wraps Remote class with try except raise condition.
* Prints more explicit error message with possible workarounds instead of basic Exception error message which recommends to file a bug.

Co-authored-by: Ganda Tchabana <ganda.tchabana@mila.quebec>
Co-authored-by: Fabrice Normandin <normandf@mila.quebec>
---
 milatools/cli/commands.py |  4 ++++
 milatools/cli/remote.py   | 16 ++++++++++------
 milatools/cli/utils.py    | 27 +++++++++++++++++++++++++++
 3 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py
index d1336c68..fb59e22b 100644
--- a/milatools/cli/commands.py
+++ b/milatools/cli/commands.py
@@ -23,6 +23,7 @@
 from .utils import (
     CommandNotFoundError,
     MilatoolsUserError,
+    SSHConnectionError,
     T,
     get_fully_qualified_name,
     qualified,
@@ -45,6 +46,9 @@ def main():
     except MilatoolsUserError as exc:
         # These are user errors and should not be reported
         print("ERROR:", exc, file=sys.stderr)
+    except SSHConnectionError as err:
+        # These are errors coming from paramiko's failure to connect to the host
+        print("ERROR:", f"{err}", file=sys.stderr)
     except Exception:
         print(T.red(traceback.format_exc()), file=sys.stderr)
         options = {
diff --git a/milatools/cli/remote.py b/milatools/cli/remote.py
index cbf9a2d4..351fc0c6 100644
--- a/milatools/cli/remote.py
+++ b/milatools/cli/remote.py
@@ -5,10 +5,11 @@
 from pathlib import Path
 from queue import Empty, Queue
 
+import paramiko
 import questionary as qn
 from fabric import Connection
 
-from .utils import T, control_file_var, here, shjoin
+from .utils import SSHConnectionError, T, control_file_var, here, shjoin
 
 batch_template = """#!/bin/bash
 #SBATCH --output={output_file}
@@ -78,11 +79,14 @@ def get_first_node_name(node_names_out: str) -> str:
 class Remote:
     def __init__(self, hostname, connection=None, transforms=(), keepalive=60):
         self.hostname = hostname
-        if connection is None:
-            connection = Connection(hostname)
-            if keepalive:
-                connection.open()
-                connection.transport.set_keepalive(keepalive)
+        try:
+            if connection is None:
+                connection = Connection(hostname)
+                if keepalive:
+                    connection.open()
+                    connection.transport.set_keepalive(keepalive)
+        except paramiko.SSHException as err:
+            raise SSHConnectionError(node_hostname=self.hostname, error=err)
         self.connection = connection
         self.transforms = transforms
 
diff --git a/milatools/cli/utils.py b/milatools/cli/utils.py
index 4c67e9f2..b906af2c 100644
--- a/milatools/cli/utils.py
+++ b/milatools/cli/utils.py
@@ -10,6 +10,7 @@
 from pathlib import Path
 
 import blessed
+import paramiko
 import questionary as qn
 from invoke.exceptions import UnexpectedExit
 from sshconf import read_ssh_config
@@ -88,6 +89,32 @@ def __str__(self):
         return message
 
 
+class SSHConnectionError(paramiko.SSHException):
+    def __init__(self, node_hostname: str, error: paramiko.SSHException):
+        super().__init__()
+        self.node_hostname = node_hostname
+        self.error = error
+
+    def __str__(self):
+        return (
+            "An error happened while trying to establish a connection with {0}".format(
+                self.node_hostname
+            )
+            + "\n\t"
+            + "-The cluster might be under maintenance"
+            + "\n\t   "
+            + "Check #mila-cluster for updates on the state of the cluster"
+            + "\n\t"
+            + "-Check the status of your connection to the cluster by ssh'ing onto it."
+            + "\n\t"
+            + "-Retry connecting with mila"
+            + "\n\t"
+            + "-Try to exclude the node with -x {0} parameter".format(
+                self.node_hostname
+            )
+        )
+
+
 def yn(prompt: str, default: bool = True) -> bool:
     return qn.confirm(prompt, default=default).unsafe_ask()