From b218debe88010d783513e6ec970573678fa204b8 Mon Sep 17 00:00:00 2001 From: David Matthews Date: Sat, 16 Jun 2018 19:18:29 +0100 Subject: [PATCH] Added region name to the sagemaker cli. (#24) --- mlflow/sagemaker/cli.py | 8 +++++--- mlflow/sagemaker/deploy.py | 4 ++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/mlflow/sagemaker/cli.py b/mlflow/sagemaker/cli.py index e2514b77e45c1..1c8358aeb2b31 100644 --- a/mlflow/sagemaker/cli.py +++ b/mlflow/sagemaker/cli.py @@ -28,7 +28,8 @@ def commands(): @click.option("--bucket", "-b", help="S3 bucket to store model artifacts", required=True) @click.option("--run_id", "-r", default=None, help="Run id") @click.option("--container", "-c", default="mlflow_sage", help="container name") -def deploy(app_name, model_path, execution_role_arn, bucket, run_id=None, container="mlflow_sage"): # noqa +@click.option("--region-name", default="us-west-2", help="region name") +def deploy(app_name, model_path, execution_role_arn, bucket, run_id=None, container="mlflow_sage", region_name="us-west-2"): # noqa """ Deploy model on sagemaker. :param app_name: Name of the deployed app. @@ -51,7 +52,8 @@ def deploy(app_name, model_path, execution_role_arn, bucket, run_id=None, contai container_name=container, app_name=app_name, model_s3_path=model_s3_path, - run_id=run_id) + run_id=run_id, + region_name=region_name) @commands.command("run-local") @@ -130,4 +132,4 @@ def _check_compatible(path): path = os.path.abspath(path) servable = Model.load(os.path.join(path, "MLmodel")) if pyfunc.FLAVOR_NAME not in servable.flavors: - raise Exception("Currenlty only supports pyfunc format.") \ No newline at end of file + raise Exception("Currenlty only supports pyfunc format.") diff --git a/mlflow/sagemaker/deploy.py b/mlflow/sagemaker/deploy.py index dd270dc69ff66..04dc81d1180da 100644 --- a/mlflow/sagemaker/deploy.py +++ b/mlflow/sagemaker/deploy.py @@ -47,7 +47,7 @@ def _upload_s3(local_model_path, bucket, prefix): shutil.rmtree(tmp_dir) -def _deploy(role, container_name, app_name, model_s3_path, run_id): +def _deploy(role, container_name, app_name, model_s3_path, run_id, region_name): """ Deploy model on sagemaker. :param role: @@ -57,7 +57,7 @@ def _deploy(role, container_name, app_name, model_s3_path, run_id): :param run_id: :return: """ - sage_client = boto3.client('sagemaker', region_name="us-west-2") + sage_client = boto3.client('sagemaker', region_name) ecr_client = boto3.client("ecr") repository_conf = ecr_client.describe_repositories( repositoryNames=[container_name])['repositories'][0]