aws-batch: Handle connection errors gracefully

CHANGES.md

@@ -13,6 +13,12 @@ development source code and as such may not be routinely kept up to date.
 
 # __NEXT__
 
+## Bug fixes
+
+* `nextstrain build --aws-batch` now handles network connection errors
+  gracefully with a clear error message instructing how to re-attach to the job,
+  rather than failing with an unhandled traceback.
+
 ## Development
 
 * We no longer test the macOS x86 standalone archive on aarch64 machines (Apple

doc/changes.md

@@ -16,6 +16,13 @@ development source code and as such may not be routinely kept up to date.
 (v-next)=
 ## __NEXT__
 
+(v-next-bug-fixes)=
+### Bug fixes
+
+* `nextstrain build --aws-batch` now handles network connection errors
+  gracefully with a clear error message instructing how to re-attach to the job,
+  rather than failing with an unhandled traceback.
+
 (v-next-development)=
 ### Development
 

nextstrain/cli/runner/aws_batch/__init__.py

@@ -89,6 +89,7 @@
 from time import sleep, time
 from typing import Iterable, Optional, cast
 from uuid import uuid4
+from ...errors import UserError
 from ...types import Env, RunnerModule, SetupStatus, SetupTestResults, UpdateStatus
 from ...util import colored, prose_list, runner_name, warn
 from ... import config
@@ -191,7 +192,16 @@ def run(opts, argv, working_volume = None, extra_env: Env = {}, cpus: int = None
 
     if opts.attach:
         print_stage("Attaching to Nextstrain AWS Batch Job ID:", opts.attach)
-        job = jobs.lookup(opts.attach)
+        try:
+            job = jobs.lookup(opts.attach)
+        except botocore.exceptions.ConnectionError as error:
+            raise UserError(f"""
+                Lost connection with AWS Batch.
+
+                Re-attach with:
+
+                    {reattach_cmd(opts.attach, local_workdir)}
+                """) from error
 
         # Read remote workdir from the job description
         remote_workdir = job.workdir
@@ -394,7 +404,16 @@ def interrupt_signaled(sig, frame):
     log_watcher = None
 
     while True:
-        job.update()
+        try:
+            job.update()
+        except botocore.exceptions.ConnectionError as error:
+            raise UserError(f"""
+                Lost connection with AWS Batch.
+
+                Re-attach with:
+
+                    {reattach_cmd(job.id, local_workdir)}
+                """) from error
 
         # Inform the user of intermediate status changes.  Final status changes
         # are messaged separately below.
@@ -419,7 +438,7 @@ def interrupt_signaled(sig, frame):
                     try:
                         for entry in job.log_entries():
                             print_job_log(entry)
-                    except botocore.exceptions.ClientError as error:
+                    except (botocore.exceptions.ClientError, botocore.exceptions.ConnectionError) as error:
                         warn(f"Unable to fetch job logs: {error}")
 
             print_stage(
@@ -511,24 +530,29 @@ def detach(job: jobs.JobState, local_workdir: Optional[Path]) -> int:
     print("")
     print_stage("Detaching from job, as requested")
 
-    reattach_cmd = " ".join([
+    print(dedent("""
+        Run the following command to re-attach to this job later to see output
+        and download results:
+
+        %s""") % (reattach_cmd(job.id, local_workdir),))
+
+    return 0
+
+
+def reattach_cmd(job_id: str, local_workdir: Optional[Path]) -> str:
+    """
+    Format a command to re-attach to an AWS Batch job.
+    """
+    return " ".join([
         "nextstrain",
         "build",
         "--aws-batch",
-        "--attach", shlex.quote(job.id),
+        "--attach", shlex.quote(job_id),
 
         # Preserve the local workdir, which has been resolved to an absolute path
         shlex.quote(str(local_workdir) if local_workdir else ".")
     ])
 
-    print(dedent("""
-        Run the following command to re-attach to this job later to see output
-        and download results:
-
-        %s""") % (reattach_cmd,))
-
-    return 0
-
 
 def print_stage(stage, *args):
     """

tests/runner-aws-batch.py

@@ -0,0 +1,31 @@
+import botocore.exceptions
+import pytest
+from nextstrain.cli.errors import UserError
+from nextstrain.cli.runner.aws_batch import run
+from unittest.mock import MagicMock, patch
+
+
+def pytest_connection_error_on_attach():
+    opts = MagicMock(attach="12345678-9abc-def0-1234-56789abcdef0", volumes=[])
+    err = botocore.exceptions.EndpointConnectionError(endpoint_url="https://batch.us-east-1.amazonaws.com")
+
+    with patch("nextstrain.cli.runner.aws_batch.jobs.lookup", side_effect=err):
+        with pytest.raises(UserError) as exc_info:
+            run(opts, [])
+
+        assert "Lost connection with AWS Batch" in str(exc_info.value)
+        assert "--attach 12345678-9abc-def0-1234-56789abcdef0" in str(exc_info.value)
+
+
+def pytest_connection_error_on_update():
+    opts = MagicMock(attach="12345678-9abc-def0-1234-56789abcdef0", volumes=[])
+    mock_job = MagicMock(id="12345678-9abc-def0-1234-56789abcdef0", workdir="s3://bucket/dir")
+    err = botocore.exceptions.EndpointConnectionError(endpoint_url="https://batch.us-east-1.amazonaws.com")
+    mock_job.update.side_effect = err
+
+    with patch("nextstrain.cli.runner.aws_batch.jobs.lookup", return_value=mock_job):
+        with pytest.raises(UserError) as exc_info:
+            run(opts, [])
+
+        assert "Lost connection with AWS Batch" in str(exc_info.value)
+        assert "--attach 12345678-9abc-def0-1234-56789abcdef0" in str(exc_info.value)