Add RasterVisionBatchClient abstraction

This class allows us to avoid passing the job_def, job_queue, branch_name,
and attempts to every method that starts a Raster Vision job.

Author: lewfish

rasterfoundry/api.py

@@ -8,9 +8,9 @@
 
 from .models import Project, MapToken
 from .exceptions import RefreshTokenException
-from .utils import start_raster_vision_job, upload_raster_vision_config
 from .settings import (
     RV_CPU_JOB_DEF, RV_CPU_QUEUE, DEVELOP_BRANCH, RV_CONFIG_URI_ROOT)
+from .utils import upload_raster_vision_config
 
 SPEC_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)),
                          'spec.yml')
@@ -156,15 +156,17 @@ def get_project_configs(self, project_ids, annotation_uris):
 
         return project_configs
 
-    def start_prep_train_data_job(self, project_ids, annotation_uris,
                                   output_zip_uri,
                                   config_uri_root=RV_CONFIG_URI_ROOT,
                                   job_queue=RV_CPU_QUEUE,
                                   job_definition=RV_CPU_JOB_DEF,
                                   branch_name=DEVELOP_BRANCH, attempts=1):
+    def start_prep_train_data_job(self, rv_batch_client, project_ids,
         """Start a Batch job to prepare object detection training data.
 
         Args:
+            rv_batch_client: a RasterVisionBatchClient object used to start
+                Batch jobs
             project_ids (list of str): ids of projects to make train data for
             annotation_uris (list of str): annotation URIs for projects
             output_zip_uri (str): URI of output zip file
@@ -188,7 +190,5 @@ def start_prep_train_data_job(self, project_ids, annotation_uris,
                        config_uri, output_zip_uri)
 
         job_name = 'prep_train_data_{}'.format(uuid.uuid1())
-        job_id = start_raster_vision_job(
-            job_name, command, job_queue, job_definition, branch_name,
-            attempts=attempts)
+        job_id = rv_batch_client.start_raster_vision_job(job_name, command)
         return job_id

rasterfoundry/models/project.py

@@ -5,8 +5,6 @@
 from .. import NOTEBOOK_SUPPORT
 from ..decorators import check_notebook
 from ..exceptions import GatewayTimeoutException
-from ..utils import start_raster_vision_job
-from ..settings import RV_CPU_JOB_DEF, RV_CPU_QUEUE, DEVELOP_BRANCH
 from .map_token import MapToken
 
 
@@ -176,13 +174,13 @@ def get_image_source_uris(self):
 
         return source_uris
 
-    def start_predict_job(self, inference_graph_uri, label_map_uri,
-                          predictions_uri, job_queue=RV_CPU_QUEUE,
-                          job_definition=RV_CPU_JOB_DEF,
-                          branch_name=DEVELOP_BRANCH, attempts=1):
+    def start_predict_job(self, rv_batch_client, inference_graph_uri,
+                          label_map_uri, predictions_uri):
         """Start a Batch job to perform object detection on this project.
 
         Args:
+            rv_batch_client: instance of RasterVisionBatchClient used to start
+                Batch jobs
             inference_graph_uri (str): file with exported object detection
                 model file
             label_map_uri (str): file with mapping from class id to display name
@@ -203,10 +201,7 @@ def start_predict_job(self, inference_graph_uri, label_map_uri,
         command = 'python -m rv.run predict {} {} {} {}'.format(
             inference_graph_uri, label_map_uri, source_uris_str,
             predictions_uri)
-        job_id = start_raster_vision_job(
-            job_name, command, job_queue=job_queue,
-            job_definition=job_definition, branch_name=branch_name,
-            attempts=attempts)
+        job_id = rv_batch_client.start_raster_vision_job(job_name, command)
 
         return job_id
 

rasterfoundry/utils.py

@@ -8,36 +8,51 @@
 
 import boto3
 
+from .settings import RV_CPU_JOB_DEF, RV_CPU_QUEUE, DEVELOP_BRANCH
 
-def start_raster_vision_job(job_name, command, job_queue, job_definition,
-                            branch_name, attempts=1):
-    """Start a raster-vision Batch job.
 
-    Args:
-        job_name (str): name of the Batch job
-        command (str): command to run inside the Docker container
-        job_queue (str): name of the Batch job queue to run the job in
-        job_definition (str): name of the Batch job definition
-        branch_name (str): branch of the raster-vision repo to use
-        attempts (int): number of attempts for the Batch job
+class RasterVisionBatchClient():
+    def __init__(self, job_queue=RV_CPU_QUEUE, job_definition=RV_CPU_JOB_DEF,
+                 branch_name=DEVELOP_BRANCH, attempts=1):
+        """Create a Raster Vision Batch Client
 
-    Returns:
-        job_id (str): job_id of job started on Batch
-    """
-    batch_client = boto3.client('batch')
-    # `run_script.sh $branch_name $command` downloads a branch of the
-    # raster-vision repo and then runs the command.
-    job_command = ['run_script.sh', branch_name, command]
-    job_id = batch_client.submit_job(
-        jobName=job_name, jobQueue=job_queue, jobDefinition=job_definition,
-        containerOverrides={
-            'command': job_command
-        },
-        retryStrategy={
-            'attempts': attempts
-        })['jobId']
+        Args:
+            job_queue (str): name of the Batch job queue to run the job in
+            job_definition (str): name of the Batch job definition
+            branch_name (str): branch of the raster-vision repo to use
+            attempts (int): number of attempts for each job
+        """
+
+        self.job_queue = job_queue
+        self.job_definition = job_definition
+        self.branch_name = branch_name
+        self.attempts = attempts
+        self.batch_client = boto3.client('batch')
+
+    def start_raster_vision_job(self, job_name, command):
+        """Start a raster-vision Batch job.
+
+        Args:
+            job_name (str): name of the Batch job
+            command (str): command to run inside the Docker container
+
+        Returns:
+            job_id (str): job_id of job started on Batch
+        """
+        # `run_script.sh $branch_name $command` downloads a branch of the
+        # raster-vision repo and then runs the command.
+        job_command = ['run_script.sh', self.branch_name, command]
+        job_id = self.batch_client.submit_job(
+            jobName=job_name, jobQueue=self.job_queue,
+            jobDefinition=self.job_definition,
+            containerOverrides={
+                'command': job_command
+            },
+            retryStrategy={
+                'attempts': self.attempts
+            })['jobId']
 
-    return job_id
+        return job_id
 
 
 def upload_raster_vision_config(config_dict, config_uri_root):