easy-sm 0.1.0__py2.py3-none-any.whl

easy_sm/__main__.py

@@ -0,0 +1,27 @@
+import click
+from easy_sm.commands.local import local
+from easy_sm.commands.initialize import init
+from easy_sm.commands.build import build
+from easy_sm.commands.push import push
+from easy_sm.commands.cloud import cloud
+
+@click.group()
+@click.option(u"-t", u"--docker-tag", default=u"latest", help=u"Specify tag for Docker image")
+@click.pass_context
+def cli(ctx, docker_tag):
+    """
+    easy_sm enables training and deploying machine learning models on AWS SageMaker in a few minutes!
+    """
+    ctx.obj = {'docker_tag': docker_tag}
+
+
+def add_commands(cli):
+    cli.add_command(init)
+    cli.add_command(build)
+    cli.add_command(local)
+    cli.add_command(push)
+    cli.add_command(cloud)
+    # cli.add_command(configure)
+
+
+add_commands(cli)

easy_sm/commands/build.py

@@ -0,0 +1,74 @@
+import sys
+import click
+import os
+import subprocess
+from easy_sm.config.config import ConfigManager
+
+
+def _build(source_dir, requirements_dir, image_name, docker_tag, python_version):
+    """
+    Builds a Docker image that contains code under the given source root directory.
+
+    Assumes that Docker is installed and running locally.
+
+    :param source_dir: [str], source root directory
+    :param requirements_dir: [str], path to requirements.txt
+    :param image_name: [str], The name of the Docker image
+    :param docker_tag: [str], the Docker tag for the image
+    """
+    easy_sm_module_path = os.path.relpath(os.path.join(source_dir, 'easy_sm_base/'))
+
+    build_script_path = os.path.join(easy_sm_module_path, 'build.sh')
+    dockerfile_path = os.path.join(easy_sm_module_path, 'Dockerfile')
+
+    train_file_path = os.path.join(easy_sm_module_path, 'training', 'train')
+    serve_file_path = os.path.join(easy_sm_module_path, 'prediction', 'serve')
+    executor_file_path = os.path.join(easy_sm_module_path, 'executor.sh')
+
+    if not os.path.isfile(build_script_path) or not os.path.isfile(train_file_path) or not \
+            os.path.isfile(serve_file_path):
+        raise ValueError("This is not a easy_sm directory: {}".format(source_dir))
+
+    os.chmod(train_file_path, 0o777)
+    os.chmod(serve_file_path, 0o777)
+    os.chmod(executor_file_path, 0o777)
+
+    target_dir_name = os.path.basename(os.path.normpath(source_dir))
+    print("------------ Using docker buildx to build from Mac for Sagemaker ------------")
+    print("------------ this behaviour is set in src/easy_sm_base/build.sh ------------")
+    output = subprocess.check_output(
+        [
+            "{}".format(build_script_path),
+            "{}".format(os.path.relpath(source_dir)),
+            "{}".format(os.path.relpath(target_dir_name)),
+            "{}".format(dockerfile_path),
+            "{}".format(os.path.relpath(requirements_dir)),
+            docker_tag,
+            image_name,
+            python_version
+        ]
+    )
+    print(output)
+
+
+@click.command()
+@click.pass_obj
+def build(obj):
+    """
+    Command to build SageMaker app
+    """
+    print("Started building SageMaker Docker image. It will take some minutes...\n")
+
+    config_file_path = os.path.join('.easy_sm.json')
+    if not os.path.isfile(config_file_path):
+        raise ValueError("This is not a easy_sm directory: {}".format(os.getcwd()))
+
+    config = ConfigManager(config_file_path).get_config()
+    _build(
+        source_dir=config.easy_sm_module_dir,
+        requirements_dir=config.requirements_dir,
+        docker_tag=obj['docker_tag'],
+        image_name=config.image_name,
+        python_version=config.python_version)
+
+    print("Docker image built successfully!")

easy_sm/commands/cloud.py

@@ -0,0 +1,266 @@
+import os
+import sys
+import click
+from easy_sm.config.config import ConfigManager
+from easy_sm.sagemaker import sagemaker
+
+
+def _config():
+    if not os.path.isfile('.easy_sm.json'):
+        raise ValueError("This is not a easy_sm directory: {}".format(os.getcwd()))
+    else:
+        return ConfigManager('.easy_sm.json').get_config()
+
+
+@click.group()
+def cloud():
+    """
+    Commands for AWS operations: upload data, train and deploy
+    """
+    pass
+
+
+@click.command(name='upload-data')
+@click.option(u"-i", u"--input-dir", required=True, help="Path to data input directory")
+@click.option(
+    u"-s", u"--s3-dir",
+    required=True,
+    help="s3 location to upload data",
+    type=click.Path()
+)
+@click.option(
+    u"-r",
+    u"--iam-role-arn",
+    required=False,
+    help="The AWS role to use for the upload command"
+)
+def upload_data(input_dir, s3_dir, iam_role_arn):
+    """
+    Command to upload data to S3
+    """
+    print("Started uploading data to S3...\n")
+    config = _config()
+    sage_maker_client = sagemaker.SageMakerClient(config.aws_profile, config.aws_region, iam_role_arn)
+    s3_path = sage_maker_client.upload_data(input_dir, s3_dir)
+    print("Data uploaded to {} successfully".format(s3_path))
+
+
+@click.command(name='train')
+@click.option(
+    u"-i", u"--input-s3-dir",
+    required=True,
+    help="s3 location to input data",
+    type=click.Path()
+)
+@click.option(
+    u"-o", u"--output-s3-dir",
+    required=True,
+    help="s3 location to save output (models, etc)",
+    type=click.Path()
+)
+@click.option(u"-e", u"--ec2-type", required=True, help="ec2 instance type")
+@click.option(
+    u"-r",
+    u"--iam-role-arn",
+    required=True,
+    help="The AWS role to use for the train command"
+)
+@click.option(
+    u"-n",
+    u"--base-job-name",
+    required=False,
+    help="Optional prefix for the SageMaker training job."
+    "If not specified, the estimator generates a default job name, based on the training image name and current timestamp."
+)
+@click.pass_obj
+def train(
+        obj,
+        input_s3_dir,
+        output_s3_dir,
+        ec2_type,
+        iam_role_arn,
+        base_job_name
+):
+    """
+    Command to train ML model(s) on SageMaker
+    """
+
+    print("Started training on SageMaker...\n")
+    config = _config()
+    sage_maker_client = sagemaker.SageMakerClient(config.aws_profile, config.aws_region, iam_role_arn)
+
+    image_name = config.image_name+':'+obj['docker_tag']
+
+    s3_model_location = sage_maker_client.train(
+        image_name=image_name,
+        input_s3_data_location=input_s3_dir,
+        train_instance_type=ec2_type,
+        output_path=output_s3_dir,
+        base_job_name=base_job_name
+    )
+
+    print("Training on SageMaker succeeded")
+    print("Model S3 location: {}".format(s3_model_location))
+    return s3_model_location  # To pipe into other commands
+
+
+@click.command(name='deploy_serverless')
+@click.option(
+    u"-m", u"--s3-model-location",
+    required=True,
+    help="s3 location to model tar.gz",
+    type=click.Path()
+)
+@click.option(u"-s",
+              u"--memory-size-in-mb",
+              required=True,
+              type=click.INT,
+              help="memory size in MB for serverless endpoint")
+@click.option(
+    u"-r",
+    u"--iam-role-arn",
+    required=False,
+    help="The AWS role to use for the deploy command"
+)
+@click.option(
+u"-n",
+    u"--endpoint-name",
+    required=True,
+    default=None,
+    help="Name for the SageMaker endpoint"
+)
+@click.pass_obj
+def deploy_serverless(
+        obj,
+        s3_model_location,
+        memory_size_in_mb,
+        iam_role_arn,
+        endpoint_name
+):
+    """
+    Command to deploy ML model(s) on SageMaker
+    """
+
+    print("Started deployment on SageMaker ...\n")
+    config = _config()
+    image_name = config.image_name+':'+obj['docker_tag']
+
+    sage_maker_client = sagemaker.SageMakerClient(config.aws_profile, config.aws_region, iam_role_arn)
+    endpoint_name = sage_maker_client.deploy_serverless(
+        image_name=image_name,
+        s3_model_location=s3_model_location,
+        memory_size_in_mb=memory_size_in_mb,
+        endpoint_name=endpoint_name
+    )
+
+    print("Endpoint name: {}".format(endpoint_name))
+
+
+@click.command(name='batch-transform')
+@click.option(
+    u"-m", u"--s3-model-location",
+    required=True,
+    help="s3 location to model tar.gz",
+    type=click.Path()
+)
+@click.option(
+    u"-i", u"--s3-input-location",
+    required=True,
+    help="s3 input data location",
+    type=click.Path()
+)
+@click.option(
+    u"-o", u"--s3-output-location",
+    required=True,
+    help="s3 location to save predictions",
+    type=click.Path()
+)
+@click.option(u"-n", u"--num-instances", required=True, type=int, help="Number of ec2 instances")
+@click.option(u"-e", u"--ec2-type", required=True, help="ec2 instance type")
+@click.option(
+    u"-r",
+    u"--iam-role-arn",
+    required=False,
+    help="The AWS role to use for batch transform command"
+)
+@click.option(
+    u"-w",
+    u"--wait",
+    default=False,
+    is_flag=True,
+    help="Wait until Batch Transform is finished. "
+         "Default: don't wait"
+)
+@click.option(
+    u"--job-name",
+    required=True,
+    default=None,
+    help="Name for the SageMaker batch transform job."
+)
+@click.pass_obj
+def batch_transform(
+        obj,
+        s3_model_location,
+        s3_input_location,
+        s3_output_location,
+        num_instances,
+        ec2_type,
+        iam_role_arn,
+        wait,
+        job_name
+):
+    """
+    Command to execute a batch transform job given a trained ML model on SageMaker
+    """
+    print("Started configuration of batch transform on SageMaker ...\n")
+
+    config = _config()
+    image_name = config.image_name+':'+obj['docker_tag']
+
+    sage_maker_client = sagemaker.SageMakerClient(config.aws_profile, config.aws_region, iam_role_arn)
+    status = sage_maker_client.batch_transform(
+        image_name=image_name,
+        s3_model_location=s3_model_location,
+        s3_input_location=s3_input_location,
+        s3_output_location=s3_output_location,
+        transform_instance_count=num_instances,
+        transform_instance_type=ec2_type,
+        wait=wait,
+        job_name=job_name
+    )
+
+    if wait:
+        print("Batch transform on SageMaker finished with status: {}".format(status))
+        if status == "Failed":
+            sys.exit(1)
+    else:
+        print("Started batch transform on SageMaker successfully")
+
+
+@click.command(name='delete_endpoint')
+@click.option(
+u"-n",
+    u"--endpoint-name",
+    required=True,
+    default=None,
+    help="Name of the SageMaker endpoint"
+)
+@click.option(
+    u"-r",
+    u"--iam-role-arn",
+    required=False,
+    help="The AWS role to use for delete command"
+)
+@click.pass_obj
+def delete_endpoint(obj, endpoint_name, iam_role_arn):
+    config = _config()
+    sage_maker_client = sagemaker.SageMakerClient(config.aws_profile, config.aws_region, iam_role_arn)
+    sage_maker_client.shutdown_endpoint(endpoint_name)
+    print(f"Endpoint {endpoint_name} has been deleted")
+
+
+cloud.add_command(upload_data)
+cloud.add_command(train)
+cloud.add_command(deploy_serverless)
+cloud.add_command(batch_transform)
+cloud.add_command(delete_endpoint)

easy_sm/commands/initialize.py

@@ -0,0 +1,169 @@
+import os
+import sys
+
+import boto3
+import click
+from click import BadParameter
+import os
+from pathlib import Path
+from easy_sm.config.config import ConfigManager
+from distutils.dir_util import copy_tree
+
+_FILE_DIR_PATH = os.path.dirname(os.path.realpath(__file__))
+
+
+def _template_creation(app_name, aws_profile, aws_region, python_version, output_dir, requirements_dir):
+    easy_sm_module_name = 'easy_sm_base'
+
+    easy_sm_exists = os.path.exists(os.path.join(output_dir, easy_sm_module_name))
+    if easy_sm_exists:
+        raise ValueError(
+            "There is a easy_sm directory/module already. "
+            "Please, rename it in order to use easy_sm."
+        )
+
+    Path(output_dir).mkdir(exist_ok=True)
+    Path(os.path.join(output_dir, '__init__.py')).touch()
+
+    # Set 'easy_sm module' directory up
+    copy_tree(os.path.join(_FILE_DIR_PATH, '../template'), output_dir)
+
+    # Set configuration file up
+    config_manager = ConfigManager(os.path.join('.easy_sm.json'))
+    config = config_manager.get_config()
+
+    config.image_name = app_name
+    config.aws_region = aws_region
+    config.aws_profile = aws_profile
+    config.easy_sm_module_dir = output_dir
+    config.python_version = python_version
+    config.requirements_dir = requirements_dir
+    config_manager.set_config(config)
+
+
+def _get_local_aws_profiles():
+    return boto3.Session().available_profiles
+
+
+def ask_for_app_name():
+    return click.prompt(
+        text="Type in a name for your SageMaker app (Only alphanumeric characters and - are allowed))",
+        type=str
+    )
+
+
+def ask_if_existing_project_exists():
+    return click.confirm(text="Are you starting a new project?")
+
+
+def ask_for_root_dir():
+    return click.prompt(text="Type in the directory where your code lives. Example: src", type=str).strip('/')
+
+
+def ask_for_python_version():
+    print("Select Python interpreter:")
+    print('{}'.format('\n'.join(['1 - Python39', '2 - Python310', '3 - Python311'])))
+
+    def _validate_python_option(input_value):
+        if int(input_value) not in {1, 2, 3}:
+            raise BadParameter(
+                message="invalid choice: {}. (choose from 1, 2, 3)".format(str(input_value))
+            )
+
+        return int(input_value)
+
+    chosen_python_index = click.prompt(
+        text="Choose from 1, 2, 3",
+        default=1,
+        value_proc=lambda x: _validate_python_option(x)
+    )
+
+    _index_to_version = {
+        1: '3.9',
+        2: '3.10',
+        3: '3.11'
+    }
+
+    return _index_to_version[chosen_python_index]
+
+
+def ask_for_aws_details():
+    available_profiles = _get_local_aws_profiles()
+
+    if len(available_profiles) == 0:
+        print("aws cli is not configured!")
+        return
+
+    valid_positions = list(range(1, len(available_profiles) + 1))
+    print("Select AWS profile:")
+    print('{}'.format(
+            '\n'.join(
+                [
+                    '{} - {}'.format(pos, profile)
+                    for pos, profile in zip(valid_positions, available_profiles)
+                ]
+            )
+        )
+    )
+
+    def _validate_profile_position(input_pos):
+        if int(input_pos) not in valid_positions:
+            raise BadParameter(
+                message="invalid choice: {}. (choose from {})".format(
+                    input_pos,
+                    ', '.join([str(pos) for pos in valid_positions])
+                )
+            )
+        return int(input_pos)
+
+    chosen_profile_index = click.prompt(
+        text="Choose from {}".format(', '.join([str(pos) for pos in valid_positions])),
+        default=1,
+        value_proc=lambda x: _validate_profile_position(x)
+    )
+    chosen_profile_index -= 1
+
+    chosen_profile = available_profiles[chosen_profile_index]
+
+    chosen_region = click.prompt(
+        text="Type in your preferred AWS region name",
+        default='eu-west-1',
+        type=str
+    )
+
+    return chosen_profile, chosen_region
+
+
+def ask_for_requirements_dir():
+    return click.prompt(text="Type in the path to requirements.txt. Example: requirements.txt", type=str).strip('/')
+
+
+@click.command()
+def init():
+    """
+    Command to initialize SageMaker template
+    """
+
+    easy_sm_app_name = ask_for_app_name()
+
+    is_new_project = ask_if_existing_project_exists()
+
+    root_dir = None
+    if not is_new_project:
+        root_dir = ask_for_root_dir()
+
+    python_version = ask_for_python_version()
+
+    aws_profile, aws_region = ask_for_aws_details()
+
+    requirements_dir = ask_for_requirements_dir()
+    _template_creation(
+        app_name=easy_sm_app_name,
+        aws_profile=aws_profile,
+        aws_region=aws_region,
+        python_version=python_version,
+        output_dir=root_dir if root_dir else 'src',
+        requirements_dir=requirements_dir
+    )
+
+    print("\neasy_sm module is created! ヽ(´▽`)/")

easy_sm/commands/local.py

@@ -0,0 +1,80 @@
+import os
+import sys
+import click
+import subprocess
+
+from easy_sm.config.config import ConfigManager
+
+
+@click.group()
+def local():
+    """
+    Commands for local operations: train and deploy
+    """
+    pass
+
+
+@click.command()
+@click.pass_obj
+def train(obj):
+    """
+    Command to train ML model(s) locally
+    """
+    print("Started local training...\n")
+    config = ConfigManager(os.path.join('.easy_sm.json')).get_config()
+    dir = config.easy_sm_module_dir
+    docker_tag = obj['docker_tag']
+    image_name = config.image_name
+    easy_sm_module_path = os.path.join(dir, 'easy_sm_base')
+    local_train_script_path = os.path.join(easy_sm_module_path, 'local_test', 'train_local.sh')
+    test_path = os.path.join(easy_sm_module_path, 'local_test', 'test_dir')
+
+    if not os.path.isdir(test_path):
+        raise ValueError("This is not a easy_sm directory: {}".format(dir))
+
+    output = subprocess.check_output(
+        [
+            "{}".format(local_train_script_path),
+            "{}".format(os.path.abspath(test_path)),
+            docker_tag,
+            image_name
+        ]
+    )
+    print(output)
+
+    print("Local training completed successfully!")
+
+
+@click.command()
+@click.pass_obj
+def deploy(obj):
+    """
+    Command to deploy ML model(s) locally
+    """
+
+    config = ConfigManager(os.path.join('.easy_sm.json')).get_config()
+    dir = config.easy_sm_module_dir
+    docker_tag = obj['docker_tag']
+    image_name = config.image_name
+
+    easy_sm_module_path = os.path.join(dir, 'easy_sm_base')
+    local_deploy_script_path = os.path.join(easy_sm_module_path, 'local_test', 'deploy_local.sh')
+    test_path = os.path.join(easy_sm_module_path, 'local_test', 'test_dir')
+
+    if not os.path.isdir(test_path):
+        raise ValueError("This is not a easy_sm directory: {}".format(dir))
+
+    print("Started local deployment at localhost:8080 ...\n")
+    output = subprocess.check_output(
+        [
+            "{}".format(local_deploy_script_path),
+            "{}".format(os.path.abspath(test_path)),
+            docker_tag,
+            image_name
+        ]
+    )
+    print(output)
+
+
+local.add_command(train)
+local.add_command(deploy)

easy_sm/commands/push.py

@@ -0,0 +1,79 @@
+import sys
+import click
+import subprocess
+import os
+from easy_sm.config.config import ConfigManager
+
+
+def _push(dir, docker_tag, aws_region, iam_role_arn, aws_profile, external_id, image_name):
+    """
+    Push Docker image to AWS ECR
+
+    :param dir: [str], source root directory
+    :param docker_tag: [str], the Docker tag for the image
+    :param aws_region: [str], the AWS region to push the image to
+    :param iam_role_arn: [str], the AWS role used to push the image to ECR
+    :param aws_profile: [str], the AWS profile used to push the image to ECR
+    :param external_id: [str], Optional external id used when using an IAM role
+    :param image_name: [str], The name of the Docker image
+    """
+
+    easy_sm_module_path = os.path.relpath(os.path.join(dir, 'easy_sm_base/'))
+    push_script_path = os.path.join(easy_sm_module_path, 'push.sh')
+
+    if not os.path.isfile(push_script_path):
+        raise ValueError("This is not a easy_sm directory: {}".format(dir))
+
+    output = subprocess.check_output([
+                                     "{}".format(push_script_path),
+                                     docker_tag,
+                                     aws_region,
+                                     iam_role_arn,
+                                     aws_profile,
+                                     external_id,
+                                     image_name])
+    print(output)
+
+@click.command()
+@click.option(u"-r", u"--aws-region", required=False, help="The AWS region to push the image to")
+@click.option(u"-i", u"--iam-role-arn", required=False, help="The AWS role to use for the push command")
+@click.option(u"-p", u"--aws-profile", required=False, help="The AWS profile to use for the push command")
+@click.option(u"-e", u"--external-id", required=False, help="Optional external id used when using an IAM role")
+@click.pass_obj
+def push(obj, aws_region, iam_role_arn, aws_profile, external_id):
+    """
+    Command to push Docker image to AWS ECR
+    """
+
+    if iam_role_arn is not None and aws_profile is not None:
+        print('Only one of iam-role-arn and aws-profile can be used.')
+        sys.exit(2)
+
+    if iam_role_arn is not None:
+        aws_profile = ''
+
+    config_file_path = os.path.join('.easy_sm.json')
+    if not os.path.isfile(config_file_path):
+        raise ValueError("This is not a easy_sm directory: {}".format(os.getcwd()))
+
+    config = ConfigManager(config_file_path).get_config()
+    image_name = config.image_name
+    aws_region = config.aws_region if aws_region is None else aws_region
+    aws_profile = config.aws_profile if (aws_profile is None and iam_role_arn is None) else aws_profile
+    external_id = "" if external_id is None else external_id
+    iam_role_arn = "" if iam_role_arn is None else iam_role_arn
+
+    print("Started pushing Docker image to AWS ECR. It will take some time. Please, be patient...\n")
+
+    _push(
+        dir=config.easy_sm_module_dir,
+        docker_tag=obj['docker_tag'],
+        aws_region=aws_region,
+        iam_role_arn=iam_role_arn,
+        aws_profile=aws_profile,
+        external_id=external_id,
+        image_name=image_name)
+
+    print("Docker image pushed to ECR successfully!")
+
+