# Copyright 2018-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""Placeholder docstring"""
from __future__ import absolute_import
import os
import platform
import sys
import tempfile
from abc import ABCMeta
from abc import abstractmethod
from six import with_metaclass
from six.moves.urllib.parse import urlparse
import sagemaker.amazon.common
import sagemaker.local.utils
import sagemaker.utils
def get_data_source_instance(data_source, sagemaker_session):
"""Return an Instance of :class:`sagemaker.local.data.DataSource`.
The instance can handle the provided data_source URI.
data_source can be either file:// or s3://
Args:
data_source (str): a valid URI that points to a data source.
sagemaker_session (:class:`sagemaker.session.Session`): a SageMaker Session to
interact with S3 if required.
Returns:
sagemaker.local.data.DataSource: an Instance of a Data Source
Raises:
ValueError: If parsed_uri scheme is neither `file` nor `s3` , raise an
error.
"""
parsed_uri = urlparse(data_source)
if parsed_uri.scheme == "file":
return LocalFileDataSource(parsed_uri.netloc + parsed_uri.path)
if parsed_uri.scheme == "s3":
return S3DataSource(parsed_uri.netloc, parsed_uri.path, sagemaker_session)
raise ValueError(
"data_source must be either file or s3. parsed_uri.scheme: {}".format(parsed_uri.scheme)
)
def get_splitter_instance(split_type):
"""Return an Instance of :class:`sagemaker.local.data.Splitter`.
The instance returned is according to the specified `split_type`.
Args:
split_type (str): either 'Line' or 'RecordIO'. Can be left as None to
signal no data split will happen.
Returns
:class:`sagemaker.local.data.Splitter`: an Instance of a Splitter
"""
if split_type is None:
return NoneSplitter()
if split_type == "Line":
return LineSplitter()
if split_type == "RecordIO":
return RecordIOSplitter()
raise ValueError("Invalid Split Type: %s" % split_type)
def get_batch_strategy_instance(strategy, splitter):
"""Return an Instance of :class:`sagemaker.local.data.BatchStrategy` according to `strategy`
Args:
strategy (str): Either 'SingleRecord' or 'MultiRecord'
splitter (:class:`sagemaker.local.data.Splitter): splitter to get the data from.
Returns
:class:`sagemaker.local.data.BatchStrategy`: an Instance of a BatchStrategy
"""
if strategy == "SingleRecord":
return SingleRecordStrategy(splitter)
if strategy == "MultiRecord":
return MultiRecordStrategy(splitter)
raise ValueError('Invalid Batch Strategy: %s - Valid Strategies: "SingleRecord", "MultiRecord"')
class DataSource(with_metaclass(ABCMeta, object)):
"""Placeholder docstring"""
@abstractmethod
def get_file_list(self):
"""Retrieve the list of absolute paths to all the files in this data source.
Returns:
List[str]: List of absolute paths.
"""
@abstractmethod
def get_root_dir(self):
"""Retrieve the absolute path to the root directory of this data source.
Returns:
str: absolute path to the root directory of this data source.
"""
class LocalFileDataSource(DataSource):
"""Represents a data source within the local filesystem."""
def __init__(self, root_path):
super(LocalFileDataSource, self).__init__()
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)
def get_file_list(self):
"""Retrieve the list of absolute paths to all the files in this data source.
Returns:
List[str] List of absolute paths.
"""
if os.path.isdir(self.root_path):
return [
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))
]
return [self.root_path]
def get_root_dir(self):
"""Retrieve the absolute path to the root directory of this data source.
Returns:
str: absolute path to the root directory of this data source.
"""
if os.path.isdir(self.root_path):
return self.root_path
return os.path.dirname(self.root_path)
class S3DataSource(DataSource):
"""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 (:class:`sagemaker.session.Session`): a sagemaker_session with the
desired settings
to talk to S3
"""
super(S3DataSource, self).__init__()
# Create a temporary dir to store the S3 contents
root_dir = sagemaker.utils.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)
# Docker cannot mount Mac OS /var folder properly see
# https://forums.docker.com/t/var-folders-isnt-mounted-properly/9600
# Only apply this workaround if the user didn't provide an alternate storage root dir.
if root_dir is None and platform.system() == "Darwin":
working_dir = "/private{}".format(working_dir)
sagemaker.utils.download_folder(bucket, prefix, working_dir, sagemaker_session)
self.files = LocalFileDataSource(working_dir)
def get_file_list(self):
"""Retrieve the list of absolute paths to all the files in this data source.
Returns:
List[str]: List of absolute paths.
"""
return self.files.get_file_list()
def get_root_dir(self):
"""Retrieve the absolute path to the root directory of this data source.
Returns:
str: absolute path to the root directory of this data source.
"""
return self.files.get_root_dir()
class Splitter(with_metaclass(ABCMeta, object)):
"""Placeholder docstring"""
@abstractmethod
def split(self, file):
"""Split a file into records using a specific strategy
Args:
file (str): path to the file to split
Returns:
generator for the individual records that were split from the file
"""
class NoneSplitter(Splitter):
"""Does not split records, essentially reads the whole file."""
# non-utf8 characters.
_textchars = bytearray({7, 8, 9, 10, 12, 13, 27} | set(range(0x20, 0x100)) - {0x7F})
def split(self, filename):
"""Split a file into records using a specific strategy.
For this NoneSplitter there is no actual split happening and the file
is returned as a whole.
Args:
filename (str): path to the file to split
Returns: generator for the individual records that were split from
the file
"""
with open(filename, "rb") as f:
buf = f.read()
if not self._is_binary(buf):
buf = buf.decode()
yield buf
def _is_binary(self, buf):
"""Check whether `buf` contains binary data.
Returns True if `buf` contains any non-utf-8 characters.
Args:
buf (bytes): data to inspect
Returns:
True if data is binary, otherwise False
"""
return bool(buf.translate(None, self._textchars))
class LineSplitter(Splitter):
"""Split records by new line."""
def split(self, file):
"""Split a file into records using a specific strategy
This LineSplitter splits the file on each line break.
Args:
file (str): path to the file to split
Returns: generator for the individual records that were split from
the 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):
"""Split a file into records using a specific strategy
This RecordIOSplitter splits the data into individual RecordIO
records.
Args:
file (str): path to the file to split
Returns: generator for the individual records that were split from
the file
"""
with open(file, "rb") as f:
for record in sagemaker.amazon.common.read_recordio(f):
yield record
class BatchStrategy(with_metaclass(ABCMeta, object)):
"""Placeholder docstring"""
def __init__(self, splitter):
"""Create a Batch Strategy Instance
Args:
splitter (sagemaker.local.data.Splitter): A Splitter to pre-process
the data before batching.
"""
self.splitter = splitter
@abstractmethod
def pad(self, file, size):
"""Group together as many records as possible to fit in the specified size.
Args:
file (str): file path to read the records from.
size (int): maximum size in MB that each group of records will be
fitted to. passing 0 means unlimited size.
Returns:
generator of records
"""
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 pad(self, file, size=6):
"""Group together as many records as possible to fit in the specified size.
Args:
file (str): file path to read the records from.
size (int): maximum size in MB that each group of records will be
fitted to. passing 0 means unlimited size.
Returns:
generator of records
"""
buffer = ""
for element in self.splitter.split(file):
if _payload_size_within_limit(buffer + element, size):
buffer += element
else:
tmp = buffer
buffer = element
yield tmp
if _validate_payload_size(buffer, size):
yield buffer
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 RuntimeError.
"""
def pad(self, file, size=6):
"""Group together as many records as possible to fit in the specified size.
This SingleRecordStrategy will not group any record and will return
them one by one as long as they are within the maximum size.
Args:
file (str): file path to read the records from.
size (int): maximum size in MB that each group of records will be
fitted to. passing 0 means unlimited size.
Returns:
generator of records
"""
for element in self.splitter.split(file):
if _validate_payload_size(element, size):
yield element
def _payload_size_within_limit(payload, size):
"""Placeholder docstring."""
size_in_bytes = size * 1024 * 1024
if size == 0:
return True
return sys.getsizeof(payload) < size_in_bytes
def _validate_payload_size(payload, size):
"""Check if a payload is within the size in MB threshold.
Raise an exception if the payload is beyond the size in MB threshold.
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 _payload_size_within_limit(payload, size):
return True
raise RuntimeError("Record is larger than %sMB. Please increase your max_payload" % size)