Added remaining data source and recordio split

Author: Ignacio Quintero

src/sagemaker/amazon/common.py

@@ -153,7 +153,7 @@ def write_spmatrix_to_sparse_tensor(file, array, labels=None):
 def read_records(file):
     """Eagerly read a collection of amazon Record protobuf objects from file."""
     records = []
-    for record_data in _read_recordio(file):
+    for record_data in read_recordio(file):
         record = Record()
         record.ParseFromString(record_data)
         records.append(record)
@@ -183,7 +183,7 @@ def _write_recordio(f, data):
     f.write(padding[pad])
 
 
-def _read_recordio(f):
+def read_recordio(f):
     while(True):
         try:
             read_kmagic, = struct.unpack('I', f.read(4))

src/sagemaker/local/data.py

@@ -14,36 +14,50 @@
 
 import os
 import sys
+import tempfile
 from six.moves.urllib.parse import urlparse
 
+from sagemaker.amazon.common import read_recordio
+from sagemaker.local.utils import download_folder
+from sagemaker.utils import get_config_value
+
 
 class DataSourceFactory(object):
 
     @staticmethod
-    def get_instance(data_source):
+    def get_instance(data_source, sagemaker_session):
         parsed_uri = urlparse(data_source)
         if parsed_uri.scheme == 'file':
             return LocalFileDataSource(parsed_uri.path)
-        else:
-            # TODO Figure S3 and S3Manifest.
-            return None
+        elif parsed_uri.scheme == 's3':
+            return S3DataSource(parsed_uri.netloc, parsed_uri.path, sagemaker_session)
+
 
 class DataSource(object):
 
     def get_file_list(self):
         pass
 
+    def get_root_dir(self):
+        pass
+
 
 class LocalFileDataSource(DataSource):
+    """
+    Represents a data source within the local filesystem.
+    """
 
     def __init__(self, root_path):
-        self.root_path = root_path
-
-    def get_file_list(self):
+        self.root_path = os.path.abspath(root_path)
         if not os.path.exists(self.root_path):
             raise RuntimeError('Invalid data source: %s Does not exist.' % self.root_path)
 
-        files = []
+    def get_file_list(self):
+        """Retrieve the list of absolute paths to all the files in this data source.
+
+        Returns:
+             List[string] List of absolute paths.
+        """
         if os.path.isdir(self.root_path):
             files = [os.path.join(self.root_path, f) for f in os.listdir(self.root_path)
                      if os.path.isfile(os.path.join(self.root_path, f))]
@@ -52,12 +66,47 @@ def get_file_list(self):
 
         return files
 
+    def get_root_dir(self):
+        """Retrieve the absolute path to the root directory of this data source.
+
+        Returns:
+            string: absolute path to the root directory of this data source.
+        """
+        if os.path.isdir(self.root_path):
+            return self.root_path
+        else:
+            return os.path.dirname(self.root_path)
+
+
 class S3DataSource(DataSource):
-    pass
+    """Defines a data source given by a bucket and s3 prefix. The contents will be downloaded
+    and then processed as local data.
+    """
+
+    def __init__(self, bucket, prefix, sagemaker_session):
+        """Create an S3DataSource instance
+
+        Args:
+            bucket (str): s3 bucket name
+            prefix (str): s3 prefix path to the data
+            sagemaker_session (sagemaker.Session): a sagemaker_session with the desired settings to talk to s3
+
+        """
+
+        # Create a temporary dir to store the S3 contents
+        root_dir = get_config_value('local.container_root', sagemaker_session.config)
+        if root_dir:
+            root_dir = os.path.abspath(root_dir)
 
+        working_dir = tempfile.mkdtemp(dir=root_dir)
+        download_folder(bucket, prefix, working_dir, sagemaker_session)
+        self.files = LocalFileDataSource(working_dir)
 
-class S3ManifestDataSource(DataSource):
-    pass
+    def get_file_list(self):
+        return self.files.get_file_list()
+
+    def get_root_dir(self):
+        return self.files.get_root_dir()
 
 
 class SplitterFactory(object):
@@ -79,23 +128,37 @@ class Splitter(object):
     def split(self, file):
         pass
 
+
 class NoneSplitter(Splitter):
+    """Does not split records, essentially reads the whole file.
+    """
 
     def split(self, file):
         with open(file, 'r') as f:
             yield f.read()
 
+
 class LineSplitter(Splitter):
+    """Split records by new line.
+
+    """
 
     def split(self, file):
         with open(file, 'r') as f:
             for line in f:
                 yield line
 
+
 class RecordIOSplitter(Splitter):
+    """Split using Amazon Recordio.
+
+    Not useful for string content.
 
+    """
     def split(self, file):
-        pass
+        with open(file, 'rb') as f:
+            for record in read_recordio(f):
+                yield record
 
 
 class BatchStrategyFactory(object):
@@ -109,13 +172,19 @@ def get_instance(strategy, splitter):
         else:
             return None
 
+
 class BatchStrategy(object):
 
     def pad(self, file, size):
         pass
 
+
 class MultiRecordStrategy(BatchStrategy):
+    """Feed multiple records at a time for batch inference.
 
+    Will group up as many records as possible within the payload specified.
+
+    """
     def __init__(self, splitter):
         self.splitter = splitter
 
@@ -133,7 +202,10 @@ def pad(self, file, size=6):
 
 
 class SingleRecordStrategy(BatchStrategy):
+    """Feed a single record at a time for batch inference.
 
+    If a single record does not fit within the payload specified it will throw a Runtime error.
+    """
     def __init__(self, splitter):
         self.splitter = splitter
 
@@ -144,14 +216,34 @@ def pad(self, file, size=6):
 
 
 def _payload_size_within_limit(payload, size):
+    """
+
+    Args:
+        payload:
+        size:
+
+    Returns:
+
+    """
     size_in_bytes = size * 1024 * 1024
     if size == 0:
         return True
     else:
-        print('size_of_payload: %s > %s' % (sys.getsizeof(payload), size_in_bytes))
         return sys.getsizeof(payload) < size_in_bytes
 
+
 def _validate_payload_size(payload, size):
-        if not _payload_size_within_limit(payload, size):
-            raise RuntimeError('Record is larger than %sMB. Please increase your max_payload' % size)
-        return True
+    """Check if a payload is within the size in MB threshold. Raise an exception otherwise.
+
+    Args:
+        payload: data that will be checked
+        size (int): max size in MB
+
+    Returns (bool): True if within bounds. if size=0 it will always return True
+    Raises:
+        RuntimeError: If the payload is larger a runtime error is thrown.
+    """
+
+    if not _payload_size_within_limit(payload, size):
+        raise RuntimeError('Record is larger than %sMB. Please increase your max_payload' % size)
+    return True

src/sagemaker/local/entities.py

@@ -18,10 +18,10 @@
 import tempfile
 import time
 import urllib3
-from six.moves.urllib.parse import urlparse
 
 from sagemaker.local.data import BatchStrategyFactory, DataSourceFactory, SplitterFactory
 from sagemaker.local.image import _SageMakerContainer
+from sagemaker.local.utils import copy_directory_structure, move_to_destination
 from sagemaker.utils import get_config_value
 
 logger = logging.getLogger(__name__)
@@ -117,12 +117,13 @@ def start(self, input_data, output_data, **kwargs):
     def _batch_inference(self, input_data, output_data, **kwargs):
         # TODO - Figure if we should pass FileDataSource here instead. Ideally not but the semantics
         # are just weird.
+        print(output_data)
         input_path = input_data['DataSource']['S3DataSource']['S3Uri']
 
         # Transform the input data to feed the serving container. We need to first gather the files
         # from S3 or Local FileSystem. Split them as required (Line, RecordIO, None) and finally batch them
         # according to the batch strategy and limit the request size.
-        data_source = DataSourceFactory.get_instance(input_path)
+        data_source = DataSourceFactory.get_instance(input_path, self.local_session)
         split_type = input_data['SplitType'] if 'SplitType' in input_data else None
         splitter = SplitterFactory.get_instance(split_type)
 
@@ -131,32 +132,33 @@ def _batch_inference(self, input_data, output_data, **kwargs):
         if 'BatchStrategy' in kwargs:
             batch_strategy = kwargs['BatchStrategy']
 
-
         max_payload = 6
         if 'MaxPayloadInMB' in kwargs:
             max_payload = int(kwargs['MaxPayloadInMB'])
 
         final_data = BatchStrategyFactory.get_instance(batch_strategy, splitter)
 
-
         # Output settings
         accept = output_data['Accept'] if 'Accept' in output_data else None
         # TODO - add a warning that we don't support KMS in Local Mode.
 
-
         # Root dir to use for intermediate data location. To make things simple we will write here regardless
         # of the final destination. At the end the files will either be moved or uploaded to S3 and deleted.
         root_dir = get_config_value('local.container_root', self.local_session.config)
         if root_dir:
             root_dir = os.path.abspath(root_dir)
 
+        working_dir = tempfile.mkdtemp(dir=root_dir)
+        dataset_dir = data_source.get_root_dir()
 
-        out_fd, out_path = tempfile.mkstemp(dir=root_dir)
+        for file in data_source.get_file_list():
 
-        output_files = {}
-        with os.fdopen(out_fd, 'w') as f:
-            for file in data_source.get_file_list():
+            relative_path = os.path.dirname(os.path.relpath(file, dataset_dir))
+            filename = os.path.basename(file)
+            copy_directory_structure(working_dir, relative_path)
+            destination_path = os.path.join(working_dir, relative_path, filename + '.out')
 
+            with open(destination_path, 'w') as f:
                 for item in final_data.pad(file, max_payload):
                     # call the container and add the result to inference.
                     response = self.local_session.sagemaker_runtime_client.invoke_endpoint(
@@ -166,12 +168,16 @@ def _batch_inference(self, input_data, output_data, **kwargs):
                     data = response_body.read()
                     response_body.close()
                     print('data: %s' % data)
-                    # TODO - AssembleWith determines if we add a new line or not.
                     f.write(data)
+                    if 'AssembleWith' in output_data and output_data['AssembleWith'] == 'Line':
+                        f.write('\n')
 
-        print(out_path)
+        print(working_dir)
+        move_to_destination(working_dir, output_data['S3OutputPath'], self.local_session)
+        print(output_data['S3OutputPath'])
         self.container.stop_serving()
 
+
 class _LocalModel(object):
 
     def __init__(self, model_name, primary_container):
@@ -244,7 +250,6 @@ def serve(self):
         # the container is running and it passed the healthcheck status is now InService
         self.state = _LocalEndpoint._IN_SERVICE
 
-
     def stop(self):
         if self.container:
             self.container.stop_serving()
@@ -269,7 +274,7 @@ def _wait_for_serving_container(serving_port):
     while True:
         i += 1
         if i >= HEALTH_CHECK_TIMEOUT_LIMIT:
-            raise RuntimeError('Giving up, endpoint: %s didn\'t launch correctly' % self.name)
+            raise RuntimeError('Giving up, endpoint didn\'t launch correctly')
 
         logger.info('Checking if serving container is up, attempt: %s' % i)
         try: