Add retry to the retrieve_dcv_session_url function

It could happen that the session is not started yet so we have to retry.

Signed-off-by: Enrico Usai <[email protected]>

Author: enrico-usai

Diff for: cli/pcluster/dcv/connect.py

@@ -22,11 +22,18 @@
     get_master_ip_and_username,
     get_stack,
     get_stack_name,
+    retry,
 )
 
 LOGGER = logging.getLogger(__name__)
 
 
+class DCVConnectionError(Exception):
+    """Error raised with DCV connection fails."""
+
+    pass
+
+
 def _check_command_output(cmd):
     return sub.check_output(cmd, shell=True, universal_newlines=True, stderr=sub.STDOUT).strip()
 
@@ -52,13 +59,35 @@ def dcv_connect(args):
         DCV_SHARED_DIR=shared_dir,
     )
 
-    # Connect by ssh to the master instance and prepare DCV session
     try:
-        LOGGER.debug("SSH command: {0}".format(cmd))
-        output = _check_command_output(cmd)
+        url = retry(_retrieve_dcv_session_url, func_args=[cmd, args.cluster_name, master_ip], attempts=4)
+        url_message = "Please use the following one-time URL in your browser within 30 seconds:\n{0}".format(url)
+    except DCVConnectionError as e:
+        error(
+            "Something went wrong during DCV connection.\n{0}"
+            "Please check the logs in the /var/log/parallelcluster/ folder "
+            "of the master instance and submit an issue {1}\n".format(e, PCLUSTER_ISSUES_LINK)
+        )
+
+    if args.show_url:
+        LOGGER.info(url_message)
+        return
+
+    try:
+        if not webbrowser.open_new(url):
+            raise webbrowser.Error("Unable to open the Web browser.")
+    except webbrowser.Error as e:
+        LOGGER.info("{0}\n{1}".format(e, url_message))
+
+
+def _retrieve_dcv_session_url(ssh_cmd, cluster_name, master_ip):
+    """Connect by ssh to the master instance, prepare DCV session and return the DCV session URL."""
+    try:
+        LOGGER.debug("SSH command: {0}".format(ssh_cmd))
+        output = _check_command_output(ssh_cmd)
         # At first ssh connection, the ssh command alerts it is adding the host to the known hosts list
         if re.search("Permanently added .* to the list of known hosts.", output):
-            output = _check_command_output(cmd)
+            output = _check_command_output(ssh_cmd)
 
         dcv_parameters = re.search(
             r"PclusterDcvServerPort=([\d]+) PclusterDcvSessionId=([\w]+) PclusterDcvSessionToken=([\w-]+)", output
@@ -71,30 +100,18 @@ def dcv_connect(args):
             error(
                 "Something went wrong during DCV connection. Please manually execute the command:\n{0}\n"
                 "If the problem persists, please check the logs in the /var/log/parallelcluster/ folder "
-                "of the master instance and submit an issue {1}.".format(cmd, PCLUSTER_ISSUES_LINK)
+                "of the master instance and submit an issue {1}".format(ssh_cmd, PCLUSTER_ISSUES_LINK)
             )
 
     except sub.CalledProcessError as e:
         if "{0}: No such file or directory".format(DCV_CONNECT_SCRIPT) in e.output:
             error(
                 "The cluster {0} has been created with an old version of ParallelCluster "
-                "without the DCV support.".format(args.cluster_name)
+                "without the DCV support.".format(cluster_name)
             )
         else:
-            error("Something went wrong during DCV connection.\n{0}".format(e.output))
+            raise DCVConnectionError(e.output)
 
-    # DCV URL
-    url = "https://{IP}:{PORT}?authToken={TOKEN}#{SESSION_ID}".format(
+    return "https://{IP}:{PORT}?authToken={TOKEN}#{SESSION_ID}".format(
         IP=master_ip, PORT=dcv_server_port, TOKEN=dcv_session_token, SESSION_ID=dcv_session_id
     )
-    url_message = "Please use the following one-time URL in your browser within 30 seconds:\n{0}".format(url)
-
-    if args.show_url:
-        LOGGER.info(url_message)
-        return
-
-    try:
-        if not webbrowser.open_new(url):
-            raise webbrowser.Error("Unable to open the Web browser.")
-    except webbrowser.Error as e:
-        LOGGER.info("{0}\n{1}".format(e, url_message))

Diff for: cli/pcluster/utils.py

@@ -463,3 +463,25 @@ def get_master_ip_and_username(cluster_name):
 
 def get_cli_log_file():
     return os.path.expanduser(os.path.join("~", ".parallelcluster", "pcluster-cli.log"))
+
+
+def retry(func, func_args, attempts=1, wait=0):
+    """
+    Call function and re-execute it if it raises an Exception.
+
+    :param func: the function to execute.
+    :param func_args: the positional arguments of the function.
+    :param attempts: the maximum number of attempts. Default: 1.
+    :param wait: delay between attempts. Default: 0.
+    :returns: the result of the function.
+    """
+    while attempts:
+        try:
+            return func(*func_args)
+        except Exception as e:
+            attempts -= 1
+            if not attempts:
+                raise e
+
+            LOGGER.debug("{0}, retrying in {1} seconds..".format(e, wait))
+            time.sleep(wait)