ENT-14394/ENT-14395: Fixed issues with resolving AMIs spawn & infinite hanging on unreachable hosts

cf_remote/aramid.py

@@ -49,6 +49,9 @@
 # just a named constant
 _DEFAULT_SSH_PORT = 22
 
+_DEFAULT_MAX_ATTEMPTS = 30
+_DEFAULT_WAIT_TIMEOUT = 30  # wall-clock seconds
+
 
 class AramidError(Exception):
     """Base exception class for the aramid module"""
@@ -258,8 +261,47 @@ def _hosts_to_host_specs(hosts):
     return host_specs
 
 
-def _wait_for_tasks(hosts, tasks, ignore_failed, echo, echo_action, out_flag=""):
+def _abort_tasks(tasks):
+    """Kill and reap the subprocess of every task that hasn't finished yet"""
+    for task in tasks:
+        if not task.done and task.proc.poll() is None:
+            task.proc.kill()
+            try:
+                task.proc.wait(timeout=2)
+            except subprocess.TimeoutExpired:
+                pass
+
+
+def _wait_for_tasks(
+    hosts,
+    tasks,
+    ignore_failed,
+    echo,
+    echo_action,
+    out_flag="",
+    max_attempts=_DEFAULT_MAX_ATTEMPTS,
+    timeout=_DEFAULT_WAIT_TIMEOUT,
+):
+    start = time.monotonic()
+    attempts = 0
     while not all(task.done for task in tasks):
+        attempts += 1
+        elapsed = time.monotonic() - start
+        pending = [t.host.host_name for t in tasks if not t.done]
+
+        if elapsed > timeout:
+            _abort_tasks(tasks)
+            raise ExecutionError(
+                "Timed out after %.1fs waiting for command(s) to finish on: %s"
+                % (elapsed, ", ".join(pending))
+            )
+        if attempts > max_attempts:
+            _abort_tasks(tasks)
+            raise ExecutionError(
+                "Exceeded maximum of %d attempts waiting for command(s) to finish on: %s"
+                % (max_attempts, ", ".join(pending))
+            )
+
         for task in (t for t in tasks if not t.done):
 
             if task.proc.args[0] == "scp":

cf_remote/spawn.py

@@ -351,12 +351,10 @@ def _get_image_criteria(platform_name):
     platform_parts = platform_name.split("-")
     platform = platform_parts[0]
     if platform == "ubuntu":
-        if len(platform_parts) == 2:
-            platform_version = platform_parts[1]
-        elif len(platform_parts) > 2:
+        if platform_parts[-1] in ("x64", "arm64"):
             platform_version = ".".join(platform_parts[1:-1])
         else:
-            platform_version = ""
+            platform_version = ".".join(platform_parts[1:3])
     else:
         platform_version = platform_name.count("-") > 0 and platform_parts[1] or "*"
     log.debug(

cf_remote/ssh.py

@@ -1,20 +1,37 @@
 import os
-import sys
 import pwd
 import shutil
 import signal
+import socket
 import subprocess
 from typing import Union
 from urllib.parse import urlparse
 
 from cf_remote import aramid
 from cf_remote import log
 from cf_remote import paths
-from cf_remote.utils import whoami, read_json
+from cf_remote.utils import whoami, read_json, CFRUserError
 from cf_remote.aramid import ExecutionResult
 from cf_remote.paths import SSH_CONFIG_FPATH, SSH_CONFIGS_JSON_FPATH, CLOUD_STATE_FPATH
 
 
+_PREFLIGHT_TIMEOUT = 1  # seconds
+
+
+class UnreachableHostError(aramid.AramidError):
+    pass
+
+
+def _check_reachable(host, port, timeout=_PREFLIGHT_TIMEOUT):
+    try:
+        with socket.create_connection((host, port), timeout=timeout):
+            pass
+    except OSError as e:
+        raise UnreachableHostError(
+            "Host '%s' is unreachable on port %s: %s" % (host, port, e)
+        ) from e
+
+
 class LocalConnection:
     is_local = True
     ssh_user = None
@@ -59,6 +76,11 @@ def __init__(self, host, user, connect_kwargs=None, port=aramid._DEFAULT_SSH_POR
         self.ssh_port = port
         self.ssh_user = user
         self._connect_kwargs = connect_kwargs
+        self._ssh_control_master = None
+
+        # Fail fast, before starting the Control Master or entering run()'s retry loop.
+        log.debug("Checking that '%s:%s' is reachable" % (host, port))
+        _check_reachable(host, port)
 
         # Create an SSH Control Master process (man:ssh_config(5)) so that
         # commands run on this host can reuse the same SSH connection.
@@ -199,9 +221,15 @@ def connect(host, users=None):
             c.ssh_port = port
             c.run("whoami", hide=True)
             return c
+        except UnreachableHostError as e:
+            # Host is down, trying other usernames won't help. Must raise
+            # rather than sys.exit(): install() calls connect() inside a
+            # multiprocessing.dummy.Pool worker thread, where a SystemExit
+            # is silently swallowed and pool.map() hangs forever instead.
+            raise CFRUserError(str(e)) from e
         except aramid.ExecutionError:
             continue
-    sys.exit("Could not ssh into '%s'" % host)
+    raise CFRUserError("Could not ssh into '%s'" % host)
 
 
 # Decorator to make a function automatically connect