diff --git a/CHANGES.md b/CHANGES.md index 9fb57c0b..61fcaed3 100644 --- a/CHANGES.md +++ b/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 diff --git a/doc/changes.md b/doc/changes.md index ac2afeba..b1b47c30 100644 --- a/doc/changes.md +++ b/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 diff --git a/nextstrain/cli/runner/aws_batch/__init__.py b/nextstrain/cli/runner/aws_batch/__init__.py index ce0e24ed..91e578e3 100644 --- a/nextstrain/cli/runner/aws_batch/__init__.py +++ b/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): """ diff --git a/tests/runner-aws-batch.py b/tests/runner-aws-batch.py new file mode 100644 index 00000000..1833a788 --- /dev/null +++ b/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)