fw_utils.py

# Copyright 2017-2018 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.
from __future__ import absolute_import

import os
import re
import shutil
import tempfile
from collections import namedtuple
from six.moves.urllib.parse import urlparse

import sagemaker.utils

_TAR_SOURCE_FILENAME = 'source.tar.gz'

UploadedCode = namedtuple('UserCode', ['s3_prefix', 'script_name'])
"""sagemaker.fw_utils.UserCode: An object containing the S3 prefix and script name.

This is for the source code used for the entry point with an ``Estimator``. It can be
instantiated with positional or keyword arguments.
"""

EMPTY_FRAMEWORK_VERSION_WARNING = 'No framework_version specified, defaulting to version {}.'
LATER_FRAMEWORK_VERSION_WARNING = 'This is not the latest supported version. ' \
                                  'If you would like to use version {latest}, ' \
                                  'please add framework_version={latest} to your constructor.'

EMPTY_FRAMEWORK_VERSION_ERROR = 'framework_version is required for script mode estimator. ' \
                                'Please add framework_version={} to your constructor to avoid this error.'

VALID_PY_VERSIONS = ['py2', 'py3']


def create_image_uri(region, framework, instance_type, framework_version, py_version=None,
                     account='520713654638', optimized_families=None):

    """Return the ECR URI of an image.

    Args:
        region (str): AWS region where the image is uploaded.
        framework (str): framework used by the image.
        instance_type (str): SageMaker instance type. Used to determine device type (cpu/gpu/family-specific optimized).
        framework_version (str): The version of the framework.
        py_version (str): Optional. Python version. If specified, should be one of 'py2' or 'py3'.
            If not specified, image uri will not include a python component.
        account (str): AWS account that contains the image. (default: '520713654638')
        optimized_families (str): Instance families for which there exist specific optimized images.

    Returns:
        str: The appropriate image URI based on the given parameters.
    """
    optimized_families = optimized_families or []

    if py_version and py_version not in VALID_PY_VERSIONS:
        raise ValueError('invalid py_version argument: {}'.format(py_version))

    # Handle Account Number for Gov Cloud
    if region == 'us-gov-west-1':
        account = '246785580436'

    # Handle Local Mode
    if instance_type.startswith('local'):
        device_type = 'cpu' if instance_type == 'local' else 'gpu'
    elif not instance_type.startswith('ml.'):
        raise ValueError('{} is not a valid SageMaker instance type. See: '
                         'https://aws.amazon.com/sagemaker/pricing/instance-types/'.format(instance_type))
    else:
        family = instance_type.split('.')[1]

        # For some frameworks, we have optimized images for specific families, e.g c5 or p3. In those cases,
        # we use the family name in the image tag. In other cases, we use 'cpu' or 'gpu'.
        if family in optimized_families:
            device_type = family
        elif family[0] in ['g', 'p']:
            device_type = 'gpu'
        else:
            device_type = 'cpu'

    if py_version:
        tag = "{}-{}-{}".format(framework_version, device_type, py_version)
    else:
        tag = "{}-{}".format(framework_version, device_type)
    return "{}.dkr.ecr.{}.amazonaws.com/sagemaker-{}:{}" \
        .format(account, region, framework, tag)


def validate_source_dir(script, directory):
    """Validate that the source directory exists and it contains the user script

    Args:
        script (str):  Script filename.
        directory (str): Directory containing the source file.

    Raises:
        ValueError: If ``directory`` does not exist, is not a directory, or does not contain ``script``.
    """
    if directory:
        if not os.path.isfile(os.path.join(directory, script)):
            raise ValueError('No file named "{}" was found in directory "{}".'.format(script, directory))

    return True


def tar_and_upload_dir(session, bucket, s3_key_prefix, script, directory, dependencies=None):
    """Pack and upload source files to S3 only if directory is empty or local.

    Note:
        If the directory points to S3 no action is taken.

    Args:
        session (boto3.Session): Boto session used to access S3.
        bucket (str): S3 bucket to which the compressed file is uploaded.
        s3_key_prefix (str): Prefix for the S3 key.
        script (str): Script filename.
        directory (str or None): Directory containing the source file. If it starts with
                                    "s3://", no action is taken.
        dependencies (List[str]): A list of paths to directories (absolute or relative)
                                containing additional libraries that will be copied into
                                /opt/ml/lib

    Returns:
        sagemaker.fw_utils.UserCode: An object with the S3 bucket and key (S3 prefix) and script name.
    """
    if directory and directory.lower().startswith('s3://'):
        return UploadedCode(s3_prefix=directory, script_name=os.path.basename(script))

    script_name = script if directory else os.path.basename(script)
    dependencies = dependencies or []
    key = '%s/sourcedir.tar.gz' % s3_key_prefix
    tmp = tempfile.mkdtemp()

    try:
        source_files = _list_files_to_compress(script, directory) + dependencies
        tar_file = sagemaker.utils.create_tar_file(source_files, os.path.join(tmp, _TAR_SOURCE_FILENAME))

        session.resource('s3').Object(bucket, key).upload_file(tar_file)
    finally:
        shutil.rmtree(tmp)

    return UploadedCode(s3_prefix='s3://%s/%s' % (bucket, key), script_name=script_name)


def _list_files_to_compress(script, directory):
    if directory is None:
        return [script]

    basedir = directory if directory else os.path.dirname(script)
    return [os.path.join(basedir, name) for name in os.listdir(basedir)]


def framework_name_from_image(image_name):
    """Extract the framework and Python version from the image name.

    Args:
        image_name (str): Image URI, which should be one of the following forms:
            legacy:
            '<account>.dkr.ecr.<region>.amazonaws.com/sagemaker-<fw>-<py_ver>-<device>:<container_version>'
            legacy:
            '<account>.dkr.ecr.<region>.amazonaws.com/sagemaker-<fw>-<py_ver>-<device>:<fw_version>-<device>-<py_ver>'
            current:
            '<account>.dkr.ecr.<region>.amazonaws.com/sagemaker-<fw>:<fw_version>-<device>-<py_ver>'

    Returns:
        tuple: A tuple containing:
            str: The framework name
            str: The Python version
            str: The image tag
    """
    # image name format: <account>.dkr.ecr.<region>.amazonaws.com/sagemaker-<framework>-<py_ver>-<device>:<tag>
    sagemaker_pattern = re.compile(r'^(\d+)(\.)dkr(\.)ecr(\.)(.+)(\.)amazonaws.com(/)(.*:.*)$')
    sagemaker_match = sagemaker_pattern.match(image_name)
    if sagemaker_match is None:
        return None, None, None
    else:
        # extract framework, python version and image tag
        # We must support both the legacy and current image name format.
        name_pattern = re.compile('^sagemaker-(tensorflow|mxnet|chainer|pytorch):(.*?)-(.*?)-(py2|py3)$')
        legacy_name_pattern = re.compile('^sagemaker-(tensorflow|mxnet)-(py2|py3)-(cpu|gpu):(.*)$')
        name_match = name_pattern.match(sagemaker_match.group(8))
        legacy_match = legacy_name_pattern.match(sagemaker_match.group(8))

        if name_match is not None:
            fw, ver, device, py = name_match.group(1), name_match.group(2), name_match.group(3), name_match.group(4)
            return fw, py, '{}-{}-{}'.format(ver, device, py)
        elif legacy_match is not None:
            return legacy_match.group(1), legacy_match.group(2), legacy_match.group(4)
        else:
            return None, None, None


def framework_version_from_tag(image_tag):
    """Extract the framework version from the image tag.

    Args:
        image_tag (str): Image tag, which should take the form '<framework_version>-<device>-<py_version>'

    Returns:
        str: The framework version.
    """
    tag_pattern = re.compile('^(.*)-(cpu|gpu)-(py2|py3)$')
    tag_match = tag_pattern.match(image_tag)
    return None if tag_match is None else tag_match.group(1)


def parse_s3_url(url):
    """Returns an (s3 bucket, key name/prefix) tuple from a url with an s3 scheme

    Args:
        url (str):

    Returns:
        tuple: A tuple containing:
            str: S3 bucket name
            str: S3 key
    """
    parsed_url = urlparse(url)
    if parsed_url.scheme != "s3":
        raise ValueError("Expecting 's3' scheme, got: {} in {}".format(parsed_url.scheme, url))
    return parsed_url.netloc, parsed_url.path.lstrip('/')


def model_code_key_prefix(code_location_key_prefix, model_name, image):
    """Returns the s3 key prefix for uploading code during model deployment

    The location returned is a potential concatenation of 2 parts
        1. code_location_key_prefix if it exists
        2. model_name or a name derived from the image

    Args:
        code_location_key_prefix (str): the s3 key prefix from code_location
        model_name (str): the name of the model
        image (str): the image from which a default name can be extracted

    Returns:
        str: the key prefix to be used in uploading code
    """
    training_job_name = sagemaker.utils.name_from_image(image)
    return '/'.join(filter(None, [code_location_key_prefix, model_name or training_job_name]))


def empty_framework_version_warning(default_version, latest_version):
    msgs = [EMPTY_FRAMEWORK_VERSION_WARNING.format(default_version)]
    if default_version != latest_version:
        msgs.append(LATER_FRAMEWORK_VERSION_WARNING.format(latest=latest_version))
    return ' '.join(msgs)