AWS Batch runner support runtime-configurable Docker images.

nextstrain/cli/errors.py

@@ -9,3 +9,8 @@ class NextstrainCliError(Exception):
 class UserError(NextstrainCliError):
     def __init__(self, message, *args, **kwargs):
         super().__init__("Error: " + message, *args, **kwargs)
+
+
+class AWSError(NextstrainCliError):
+    """Generic exception when interfacing with AWS."""
+    pass

nextstrain/cli/runner/__init__.py

@@ -107,6 +107,13 @@ def register_arguments(parser: ArgumentParser, runners: List, exec: List) -> Non
         metavar = "<prog>",
         default = exec_cmd)
 
+    # Docker image
+    development.add_argument(
+        "--image",
+        help    = "Docker name to either build or run nextstrain.",
+        metavar = "<image>",
+        default = docker.DEFAULT_IMAGE)
+
     # Static exec arguments; never modified directly by the user invocation,
     # but they won't be used if --exec is changed.
     parser.set_defaults(default_exec_cmd = exec_cmd)

nextstrain/cli/runner/aws_batch/__init__.py

@@ -14,6 +14,7 @@
 from ...types import RunnerTestResults, Tuple
 from ...util import colored, resolve_path, warn
 from ... import config
+from .. import docker
 from . import jobs, logs, s3
 
 
@@ -129,6 +130,7 @@ def run(opts, argv, working_volume = None, extra_env = {}) -> int:
         try:
             job = jobs.submit(
                 name       = run_id,
+                image      = opts.image,
                 queue      = opts.job_queue,
                 definition = opts.job_definition,
                 cpus       = opts.cpus,

nextstrain/cli/runner/aws_batch/jobs.py

@@ -5,7 +5,10 @@
 from time import time
 from typing import Callable, Generator, Iterable, Mapping, List, Optional
 from ... import hostenv, aws
+from ...util import warn, aws_job_definition_name
+from ...errors import AWSError
 from . import logs, s3
+import botocore.exceptions
 
 
 class JobState:
@@ -138,7 +141,42 @@ def stopped(self) -> bool:
         return self.status_reason == self.STOP_REASON
 
 
+def get_job_definition(definition_name: str, image: str) -> str:
+    """
+        Get the job definition according to the Docker image name.
+        Create the job definition if it doesn't exist yet.
+
+        Returns the name of the job definition function of the Docker image name.
+    """
+
+    batch = aws.client_with_default_region("batch")
+    base_definition_name = definition_name
+    definition_name = aws_job_definition_name(definition_name, image)
+    try:
+        if definition_exists(definition_name):
+            return definition_name
+
+        base_definition = batch.describe_job_definitions(
+            jobDefinitionName=base_definition_name, status='ACTIVE'
+            ).get("jobDefinitions")
+
+        if not base_definition:
+            raise AWSError("Job definition {} doesn't exist, cannot create {} from it".format(
+                base_definition_name, definition_name))
+
+        # Take the base job definition, change the image and create the image with it
+        base_definition['jobDefinitionName'] = definition_name
+        base_definition['containerProperties']['image'] = image
+        batch.register_job_definition(base_definition)
+
+    except botocore.exceptions.ClientError as error:
+        raise AWSError("Creation of job definition ({}) failed, {}".format(definition_name, error))
+
+    return definition_name
+
+
 def submit(name: str,
+           image: str,
            queue: str,
            definition: str,
            cpus: Optional[int],
@@ -153,10 +191,12 @@ def submit(name: str,
     """
     batch = aws.client_with_default_region("batch")
 
+    definition_name = get_job_definition(definition, image)
+
     submission = batch.submit_job(
         jobName = name,
         jobQueue = queue,
-        jobDefinition = definition,
+        jobDefinition = definition_name,
         containerOverrides = {
             "environment": [
                 {

nextstrain/cli/runner/docker.py

@@ -36,12 +36,6 @@ def register_arguments(parser) -> None:
     development = parser.add_argument_group(
         "development options for --docker")
 
-    development.add_argument(
-        "--image",
-        help    = "Container image in which to run the pathogen build",
-        metavar = "<name>",
-        default = DEFAULT_IMAGE)
-
     development.set_defaults(volumes = [])
 
     for name in COMPONENTS:

nextstrain/cli/runner/native.py

@@ -5,7 +5,8 @@
 import os
 import shutil
 from ..types import RunnerTestResults
-from ..util import exec_or_return
+from ..util import exec_or_return, warn
+from .docker import DEFAULT_IMAGE
 
 
 def register_arguments(parser) -> None:
@@ -16,9 +17,10 @@ def register_arguments(parser) -> None:
 
 
 def run(opts, argv, working_volume = None, extra_env = {}) -> int:
+    if opts.image != DEFAULT_IMAGE:
+        warn("Docker image specified ({}), but it won't be used".format(opts.image()))
     if working_volume:
         os.chdir(str(working_volume.src))
-
     return exec_or_return(argv, extra_env)
 
 

nextstrain/cli/util.py

@@ -213,3 +213,19 @@ def resolve_path(path: Path) -> Path:
         return path.resolve(strict = True) # type: ignore
     else:
         return path.resolve()
+
+
+def aws_job_definition_name(definition_name: str, docker_image: str) -> str:
+    """
+    Format the AWS Batch Job Definition name according to API restriction.
+
+    Returns a string.
+    """
+    # Import here to avoid cyclic imports
+    from .runner.docker import split_image_name
+
+    docker_image_repo, docker_image_tag = split_image_name(docker_image)
+    name = re.sub('[^0-9a-zA-Z-]+', '-', definition_name)
+    image = re.sub('[^0-9a-zA-Z-]+', '-', docker_image_repo)
+    tag = re.sub('[^0-9a-zA-Z-]+', '-', docker_image_tag)
+    return "{}_{}_{}".format(name, image, tag)