Make Use of CPU and GPU Queues

docs/changelog.rst

@@ -14,6 +14,7 @@ Raster Vision 0.9.0
 - Add ability to shift raster images by given numbers of meters.  `#573 <https://github.com/azavea/raster-vision/pull/573>`_
 - Add ability to generate GeoJSON segmentation predictions.  `#575 <https://github.com/azavea/raster-vision/pull/575>`_
 - Add ability to run the DeepLab eval script.  `#653 <https://github.com/azavea/raster-vision/pull/653>`_
+- Submit CPU-only stages to a CPU queue on Aws.  `#668 <https://github.com/azavea/raster-vision/pull/668>`_
 
 Raster Vision 0.8
 -----------------

docs/setup.rst

@@ -209,13 +209,32 @@ Set the appropriate configuration in your :ref:`raster vision config`:
 
    [AWS_BATCH]
    job_queue=rasterVisionQueue
+   cpu_job_queue=rasterVisionQueue
    job_definition=raster-vision-gpu
+   cpu_job_definition=raster-vision-gpu
    attempts=1
 
 
-* ``job_queue`` - Job Queue to submit Batch jobs to.
-* ``job_definition`` - The Job Definition that define the Batch jobs to run.
+* ``job_queue`` - Job Queue to submit GPU Batch jobs to.
+* ``cpu_job_queue`` - Job Queue to submit CPU-only jobs to.
+* ``job_definition`` - The Job Definition that defines the Batch jobs to run.
+* ``cpu_job_definition`` - The Job Definition that defines the CPU Batch jobs to run.
 * ``attempts`` - Optional number of attempts to retry failed jobs.
 
+Note that in the sample file given above, the same respective resources are used for the GPU-capable and the CPU-capable jobs.
+That is because the  `Raster Vision for AWS Batch setup repository <https://github.com/azavea/raster-vision-aws>`_ only creates one GPU-capable queue and one GPU-capable job definition.
+That repository is currently under revision.
+
+If you have created a CPU-only job queue and a CPU-only job definition, the configuration will look something like this:
+
+.. code:: ini
+
+   [AWS_BATCH]
+   job_queue=rasterVisionQueue
+   cpu_job_queue=rasterVisionQueueCPU
+   job_definition=raster-vision-gpu
+   cpu_job_definition=raster-vision-cpu
+   attempts=1
+
 .. seealso::
    For more information about how Raster Vision uses AWS Batch, see the section: :ref:`aws batch`.

rastervision/runner/aws_batch_experiment_runner.py

@@ -3,6 +3,9 @@
 
 from rastervision.runner import OutOfProcessExperimentRunner
 from rastervision.rv_config import RVConfig
+import rastervision as rv
+
+CPU_STAGES = {rv.ANALYZE, rv.CHIP, rv.EVAL, rv.BUNDLE}
 
 
 class AwsBatchExperimentRunner(OutOfProcessExperimentRunner):
@@ -24,6 +27,14 @@ def __init__(self):
                 job_queue = 'raster-vision-cpu'
         self.job_queue = job_queue
 
+        cpu_job_queue = batch_config('cpu_job_queue', default='')
+        if not cpu_job_queue:
+            if self.gpu:
+                cpu_job_queue = 'raster-vision-cpu'
+            else:
+                cpu_job_queue = job_queue
+        self.cpu_job_queue = cpu_job_queue
+
         job_definition = batch_config('job_definition', default='')
         if not job_definition:
             if self.gpu:
@@ -32,6 +43,14 @@ def __init__(self):
                 job_definition = 'raster-vision-cpu'
         self.job_definition = job_definition
 
+        cpu_job_definition = batch_config('cpu_job_definition', default='')
+        if not cpu_job_definition:
+            if self.gpu:
+                cpu_job_definition = 'raster-vision-cpu'
+            else:
+                cpu_job_definition = job_definition
+        self.cpu_job_definition = cpu_job_definition
+
         self.submit = self.batch_submit
         self.execution_environment = 'Batch'
 
@@ -65,10 +84,17 @@ def batch_submit(self,
         job_name = '{}_{}_{}'.format(command_type, exp, uuid_part)
         depends_on = [{'jobId': job_id} for job_id in parent_job_ids]
 
+        if command_type in CPU_STAGES:
+            job_queue = self.cpu_job_queue
+            job_definition = self.cpu_job_definition
+        else:
+            job_queue = self.job_queue
+            job_definition = self.job_definition
+
         kwargs = {
             'jobName': job_name,
-            'jobQueue': self.job_queue,
-            'jobDefinition': self.job_definition,
+            'jobQueue': job_queue,
+            'jobDefinition': job_definition,
             'containerOverrides': {
                 'command': full_command
             },