Repository: kubeflow/pipelines

Issue:
allow output artifact store configuration (vs hard coded)
it seems like the output artifacts are always stored in a specific minio service, port, namespace, bucket, secrets, etc (`minio-service.kubeflow:9000`). 

see: https://github.com/kubeflow/pipelines/blob/f40a22a3f4a8e06d20cf3e3f425b5058d5c87e0b/sdk/python/kfp/compiler/_op_to_template.py#L148

it would be great to make it flexible, e.g. allow using S3, or change namespace or bucket names.
i suggest making it configurable, i can do such PR if we agree its needed. 
flexible pipeline service (host) path in client SDK 
when creating an SDK `Client()` the path to `ml-pipeline` API service is loaded from a hard coded value (`ml-pipeline.kubeflow.svc.cluster.local:8888`) which indicate a specific k8s namespace. it can be valuable to load that default value from an env variable, i.e. changing the line in `_client.py` from:

`config.host = host if host else Client.IN_CLUSTER_DNS_NAME`

to:

`config.host = host or os.environ.get('ML_PIPELINE_DNS_NAME',Client.IN_CLUSTER_DNS_NAME)`

also note that when a user provide the `host` parameter, the ipython output points to the API server and not to the UI service (see the logic in `_get_url_prefix()`), it seems like a potential bug

if its acceptable i can submit a PR for the line change above
 


Discussion:
There is already a PR for this. But we'd like to solve this in a different way - by having an in-cluster configMap with the storage configuration. Argo already supports that, but we're hesitating to use that mechanism since KF Pipelines might not be the exclusive user of the Argo controller.

Fortunately, the problem will be solved in Argo 2.4 (per-pipeline configMap reference)
Also https://github.com/kubeflow/pipelines/issues/363 https://github.com/kubeflow/pipelines/issues/1195 https://github.com/kubeflow/pipelines/issues/362
@Ark-kun a mechanism like config map sounds better, it will also make the YAMLs cleaner, currently u get the same 27 lines added to every container, see below, would be nice to hide it:

```yaml
    outputs:
      artifacts:
      - name: mlpipeline-ui-metadata
        path: /mlpipeline-ui-metadata.json
        s3:
          accessKeySecret:
            key: accesskey
            name: mlpipeline-minio-artifact
          bucket: mlpipeline
          endpoint: minio-service.kubeflow:9000
          insecure: true
          key: runs/{{workflow.uid}}/{{pod.name}}/mlpipeline-ui-metadata.tgz
          secretKeySecret:
            key: secretkey
            name: mlpipeline-minio-artifact
      - name: mlpipeline-metrics
        path: /mlpipeline-metrics.json
        s3:
          accessKeySecret:
            key: accesskey
            name: mlpipeline-minio-artifact
          bucket: mlpipeline
          endpoint: minio-service.kubeflow:9000
          insecure: true
          key: runs/{{workflow.uid}}/{{pod.name}}/mlpipeline-metrics.tgz
          secretKeySecret:
            key: secretkey
            name: mlpipeline-minio-artifact
```
Fixed by https://github.com/kubeflow/pipelines/pull/1297
>it can be valuable to load that default value from an env variable

I'm not sure about the user scenarios, but I guess it might be useful for some cases.

>ML_PIPELINE_DNS_NAME

Maybe, `KF_PIPELINES_ENDPOINT`?

>it seems like a potential bug

It probably is. The Client class can definitely use some improvements, especially for out-of-cluster usage.

> if its acceptable i can submit a PR

We always appreciate PRs.
Fixed in https://github.com/kubeflow/pipelines/pull/1293

Resolving patch:
diff --git a/components/aws/sagemaker/common/_utils.py b/components/aws/sagemaker/common/_utils.py
--- a/components/aws/sagemaker/common/_utils.py
+++ b/components/aws/sagemaker/common/_utils.py
@@ -471,10 +471,52 @@ def create_model_request(args):
 
     return request
 
+def endpoint_name_exists(client, endpoint_name):
+  try:
+      endpoint_name = client.describe_endpoint(EndpointName=endpoint_name)['EndpointName']
+      logging.info("Endpoint exists: " + endpoint_name)
+      return True
+  except ClientError as e:
+      logging.debug("Endpoint does not exist")
+  return False
+
+def endpoint_config_name_exists(client, endpoint_config_name):
+  try:
+      config_name = client.describe_endpoint_config(EndpointConfigName=endpoint_config_name)['EndpointConfigName']
+      logging.info("Endpoint Config exists: " + config_name)
+      return True
+  except ClientError as e:
+      logging.info("Endpoint Config does not exist:" + endpoint_config_name)
+  return False
+
+def update_deployed_model(client, args):
+    endpoint_config_name = create_endpoint_config(client, args)
+    endpoint_name = update_endpoint(client, args['region'], args['endpoint_name'], endpoint_config_name)
+    return endpoint_name
+
 def deploy_model(client, args):
-  endpoint_config_name = create_endpoint_config(client, args)
-  endpoint_name = create_endpoint(client, args['region'], args['endpoint_name'], endpoint_config_name, args['endpoint_tags'])
-  return endpoint_name
+    args['update_endpoint'] = False
+    endpoint_config_name = create_endpoint_config(client, args)
+    endpoint_name = create_endpoint(client, args['region'], args['endpoint_name'], endpoint_config_name, args['endpoint_tags'])
+    return endpoint_name
+
+def get_endpoint_config(client, endpoint_name):
+    endpoint_config_name = None
+    try:
+        endpoint_config_name = client.describe_endpoint(EndpointName=endpoint_name)['EndpointConfigName']
+        logging.info("Current Endpoint Config Name: " + endpoint_config_name)
+    except ClientError as e:
+        logging.info("Endpoint Config does not exist")
+        ## This is not an error, end point may not exist
+    return endpoint_config_name
+
+def delete_endpoint_config(client, endpoint_config_name):
+    try:
+        client.delete_endpoint_config(EndpointConfigName=endpoint_config_name)
+        return True
+    except ClientError as e:
+        logging.info("Endpoint config may not exist to be deleted: " + e.response['Error']['Message'])
+    return False
 
 def create_endpoint_config_request(args):
     ### Documentation: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.create_endpoint_config
@@ -485,7 +527,7 @@ def create_endpoint_config_request(args):
         logging.error("Must specify at least one model (model name) to host.")
         raise Exception("Could not create endpoint config.")
 
-    endpoint_config_name = args['endpoint_config_name'] if args['endpoint_config_name'] else 'EndpointConfig' + args['model_name_1'][args['model_name_1'].index('-'):]
+    endpoint_config_name = args['endpoint_config_name']
     request['EndpointConfigName'] = endpoint_config_name
 
     if args['resource_encryption_key']:
@@ -517,7 +559,17 @@ def create_endpoint_config_request(args):
 
     return request
 
+def get_endpoint_config_name(client, args):
+    ## boto3 documentation says to update an endpoint, a new EndPointConfig has to be created
+    ## and the one currently in use should NOT be deleted. Appending a random number to resolve conflict
+    endpoint_config_name = args['endpoint_config_name'] if args['endpoint_config_name'] else 'EndpointConfig' + "-" + id_generator(8)
+    if args['update_endpoint'] and endpoint_config_name_exists(client, endpoint_config_name):
+        endpoint_config_name = endpoint_config_name + "-" + id_generator(8)
+        logging.info("Changed endpoint_config_name to: " + endpoint_config_name)
+    return endpoint_config_name
+
 def create_endpoint_config(client, args):
+    args['endpoint_config_name'] = get_endpoint_config_name(client, args)
     request = create_endpoint_config_request(args)
     try:
         create_endpoint_config_response = client.create_endpoint_config(**request)
@@ -528,6 +580,20 @@ def create_endpoint_config(client, args):
     except ClientError as e:
         raise Exception(e.response['Error']['Message'])
 
+def update_endpoint(client, region, endpoint_name, endpoint_config_name):
+  try:
+      update_endpoint_response = client.update_endpoint(
+          EndpointName=endpoint_name,
+          EndpointConfigName=endpoint_config_name,
+          )
+      logging.info("Updating endpoint with name: " + endpoint_name)
+      logging.info("Endpoint in SageMaker: https://{}.console.aws.amazon.com/sagemaker/home?region={}#/endpoints/{}"
+          .format(region, region, endpoint_name))
+      logging.info("CloudWatch logs: https://{}.console.aws.amazon.com/cloudwatch/home?region={}#logStream:group=/aws/sagemaker/Endpoints/{};streamFilter=typeLogStreamPrefix"
+          .format(region, region, endpoint_name))
+      return endpoint_name
+  except ClientError as e:
+      raise Exception(e.response['Error']['Message'])
 
 def create_endpoint(client, region, endpoint_name, endpoint_config_name, endpoint_tags):
   ### Documentation: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.create_endpoint
@@ -565,11 +631,11 @@ def wait_for_endpoint_creation(client, endpoint_name):
 
     if status != 'InService':
       message = resp['FailureReason']
-      logging.info('Create endpoint failed with the following error: {}'.format(message))
+      logging.info('Create/Update endpoint failed with the following error: {}'.format(message))
       raise Exception('Endpoint creation did not succeed')
 
     logging.info("Endpoint Arn: " + resp['EndpointArn'])
-    logging.info("Create endpoint ended with status: " + status)
+    logging.info("Create/Update endpoint ended with status: " + status)
 
 def create_transform_job_request(args):
     ### Documentation: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.create_transform_job
diff --git a/components/aws/sagemaker/deploy/src/deploy.py b/components/aws/sagemaker/deploy/src/deploy.py
--- a/components/aws/sagemaker/deploy/src/deploy.py
+++ b/components/aws/sagemaker/deploy/src/deploy.py
@@ -20,7 +20,7 @@ def create_parser():
   parser = argparse.ArgumentParser(description='SageMaker Training Job')
   _utils.add_default_client_arguments(parser)
   
-  parser.add_argument('--endpoint_config_name', type=str, required=False, help='The name of the endpoint configuration.', default='')
+  parser.add_argument('--endpoint_config_name', type=str, required=False, help='The name of the endpoint configuration. If an existing endpoint is being updated, a suffix is automatically added if this config name exists.', default='')
   parser.add_argument('--variant_name_1', type=str, required=False, help='The name of the production variant.', default='variant-name-1')
   parser.add_argument('--model_name_1', type=str, required=True, help='The model name used for endpoint deployment.')
   parser.add_argument('--initial_instance_count_1', type=int, required=False, help='Number of instances to launch initially.', default=1)
@@ -43,6 +43,7 @@ def create_parser():
   parser.add_argument('--endpoint_config_tags', type=_utils.yaml_or_json_str, required=False, help='An array of key-value pairs, to categorize AWS resources.', default={})
   parser.add_argument('--endpoint_name', type=str, required=False, help='The name of the endpoint.', default='')
   parser.add_argument('--endpoint_tags', type=_utils.yaml_or_json_str, required=False, help='An array of key-value pairs, to categorize AWS resources.', default={})
+  parser.add_argument('--update_endpoint', type=_utils.str_to_bool, required=False, help='If True, update endpoint if it exists else creates one', default=False)
   parser.add_argument('--endpoint_name_output_path', type=str, default='/tmp/endpoint-name', help='Local output path for the file containing the name of the created endpoint.')
 
   return parser
@@ -53,15 +54,31 @@ def main(argv=None):
 
   logging.getLogger().setLevel(logging.INFO)
   client = _utils.get_sagemaker_client(args.region, args.endpoint_url, assume_role_arn=args.assume_role)
-  logging.info('Submitting Endpoint request to SageMaker...')
-  endpoint_name = _utils.deploy_model(client, vars(args))
-  logging.info('Endpoint creation request submitted. Waiting for completion...')
+  endpoint_name = None
+  old_endpoint_config_name = None
+  if (args.update_endpoint and _utils.endpoint_name_exists(client, args.endpoint_name)):
+    ## Get the old endpoint config to cleanup later
+    old_endpoint_config_name = _utils.get_endpoint_config(client, args.endpoint_name)
+    logging.info('Submitting Update Endpoint request to SageMaker...')
+    endpoint_name = _utils.update_deployed_model(client, vars(args))
+  else:
+    logging.info('Submitting Create Endpoint request to SageMaker...')
+    endpoint_name = _utils.deploy_model(client, vars(args))
+
+  logging.info('Endpoint creation/update request submitted. Waiting for completion...')
   _utils.wait_for_endpoint_creation(client, endpoint_name)
 
-  _utils.write_output(args.endpoint_name_output_path, endpoint_name)
+  ## If updating existing endpoint, cleanup old endpoint config
+  if old_endpoint_config_name:
+      logging.info("Deleting old endpoint config: " + old_endpoint_config_name)
+      if _utils.delete_endpoint_config(client, old_endpoint_config_name):
+        logging.info("Deleted old endpoint config: " + old_endpoint_config_name)
+      else:
+        logging.info("Unable to delete old endpoint config: " + old_endpoint_config_name)
 
-  logging.info('Endpoint creation completed.')
+  _utils.write_output(args.endpoint_name_output_path, endpoint_name)
 
+  logging.info('Endpoint creation/update completed.')
 
 if __name__== "__main__":
   main(sys.argv[1:])


Test patch:
diff --git a/components/aws/sagemaker/tests/integration_tests/component_tests/test_batch_transform_component.py b/components/aws/sagemaker/tests/integration_tests/component_tests/test_batch_transform_component.py
--- a/components/aws/sagemaker/tests/integration_tests/component_tests/test_batch_transform_component.py
+++ b/components/aws/sagemaker/tests/integration_tests/component_tests/test_batch_transform_component.py
@@ -83,8 +83,9 @@ def test_transform_job(
     )
     assert s3_utils.check_object_exists(s3_client, s3_data_bucket, file_key)
 
-    assert not argo_utils.error_in_cw_logs(workflow_json["metadata"]["name"]), \
-        ('Found the CloudWatch error message in the log output. Check SageMaker to see if the job has failed.')
+    assert not argo_utils.error_in_cw_logs(
+        workflow_json["metadata"]["name"]
+    ), "Found the CloudWatch error message in the log output. Check SageMaker to see if the job has failed."
 
     utils.remove_dir(download_dir)
 
@@ -104,9 +105,7 @@ def test_terminate_transformJob(kfp_client, experiment_id, region, sagemaker_cli
     # Generate random prefix for model, job name to avoid errors if resources with same name exists
     test_params["Arguments"]["model_name"] = test_params["Arguments"][
         "job_name"
-    ] = input_job_name = (
-        utils.generate_random_string(4) + "-terminate-job"
-    )
+    ] = input_job_name = (utils.generate_random_string(4) + "-terminate-job")
 
     run_id, _, workflow_json = kfp_client_utils.compile_run_monitor_pipeline(
         kfp_client,
diff --git a/components/aws/sagemaker/tests/integration_tests/component_tests/test_deploy_component.py b/components/aws/sagemaker/tests/integration_tests/component_tests/test_deploy_component.py
--- a/components/aws/sagemaker/tests/integration_tests/component_tests/test_deploy_component.py
+++ b/components/aws/sagemaker/tests/integration_tests/component_tests/test_deploy_component.py
@@ -35,9 +35,7 @@ def np2csv(arr):
     payload = np2csv(train_set[0][30:31])
 
     response = runtime.invoke_endpoint(
-        EndpointName=endpoint_name,
-        ContentType="text/csv",
-        Body=payload,
+        EndpointName=endpoint_name, ContentType="text/csv", Body=payload,
     )
     return json.loads(response["Body"].read().decode())
 
@@ -45,9 +43,10 @@ def np2csv(arr):
 @pytest.mark.parametrize(
     "test_file_dir",
     [
+        pytest.param("resources/config/kmeans-mnist-update-endpoint"),
         pytest.param(
             "resources/config/kmeans-mnist-endpoint", marks=pytest.mark.canary_test
-        )
+        ),
     ],
 )
 def test_create_endpoint(
@@ -104,6 +103,19 @@ def test_create_endpoint(
             ]
             == "InService"
         )
+        # Verify that the update was successful by checking that InstanceType changed
+        if "ExpectedInstanceType" in test_params.keys():
+            new_endpoint_config_name = sagemaker_utils.describe_endpoint(
+                sagemaker_client, input_endpoint_name
+            )["EndpointConfigName"]
+            response = sagemaker_utils.describe_endpoint_config(
+                sagemaker_client, new_endpoint_config_name
+            )
+            prod_variant = response["ProductionVariants"][0]
+            print(f"Production Variant item: {prod_variant}")
+            instance_type = prod_variant["InstanceType"]
+            print(f"Production Variant item InstanceType: {instance_type}")
+            assert instance_type == test_params["ExpectedInstanceType"]
 
         # Validate the model for use by running a prediction
         result = run_predict_mnist(boto3_session, input_endpoint_name, download_dir)
diff --git a/components/aws/sagemaker/tests/integration_tests/component_tests/test_groundtruth_component.py b/components/aws/sagemaker/tests/integration_tests/component_tests/test_groundtruth_component.py
--- a/components/aws/sagemaker/tests/integration_tests/component_tests/test_groundtruth_component.py
+++ b/components/aws/sagemaker/tests/integration_tests/component_tests/test_groundtruth_component.py
@@ -22,12 +22,7 @@ def create_initial_workteam(
 
     # First create a workteam using a separate pipeline and get the name, arn of the workteam created.
     create_workteamjob(
-        kfp_client,
-        test_params,
-        experiment_id,
-        region,
-        sagemaker_client,
-        download_dir,
+        kfp_client, test_params, experiment_id, region, sagemaker_client, download_dir,
     )
 
     workteam_arn = sagemaker_utils.get_workteam_arn(sagemaker_client, workteam_name)
@@ -74,7 +69,9 @@ def test_groundtruth_labeling_job(
         test_params["Arguments"][
             "ground_truth_train_job_name"
         ] = ground_truth_train_job_name = (
-            test_params["Arguments"]["ground_truth_train_job_name"] + "-by-" + workteam_name
+            test_params["Arguments"]["ground_truth_train_job_name"]
+            + "-by-"
+            + workteam_name
         )
 
         run_id, _, _ = kfp_client_utils.compile_run_monitor_pipeline(
@@ -115,7 +112,9 @@ def test_groundtruth_labeling_job(
     finally:
         # Cleanup the SageMaker Resources
         if ground_truth_train_job_name:
-            sagemaker_utils.stop_labeling_job(sagemaker_client, ground_truth_train_job_name)
+            sagemaker_utils.stop_labeling_job(
+                sagemaker_client, ground_truth_train_job_name
+            )
         if workteam_name:
             sagemaker_utils.delete_workteam(sagemaker_client, workteam_name)
 
diff --git a/components/aws/sagemaker/tests/integration_tests/component_tests/test_hpo_component.py b/components/aws/sagemaker/tests/integration_tests/component_tests/test_hpo_component.py
--- a/components/aws/sagemaker/tests/integration_tests/component_tests/test_hpo_component.py
+++ b/components/aws/sagemaker/tests/integration_tests/component_tests/test_hpo_component.py
@@ -10,11 +10,7 @@
 
 @pytest.mark.parametrize(
     "test_file_dir",
-    [
-        pytest.param(
-            "resources/config/kmeans-mnist-hpo", marks=pytest.mark.canary_test
-        )
-    ],
+    [pytest.param("resources/config/kmeans-mnist-hpo", marks=pytest.mark.canary_test)],
 )
 def test_hyperparameter_tuning(
     kfp_client, experiment_id, region, sagemaker_client, test_file_dir
diff --git a/components/aws/sagemaker/tests/integration_tests/component_tests/test_processing_component.py b/components/aws/sagemaker/tests/integration_tests/component_tests/test_processing_component.py
--- a/components/aws/sagemaker/tests/integration_tests/component_tests/test_processing_component.py
+++ b/components/aws/sagemaker/tests/integration_tests/component_tests/test_processing_component.py
@@ -15,7 +15,7 @@
             "resources/config/kmeans-algo-mnist-processing",
             marks=pytest.mark.canary_test,
         ),
-        "resources/config/assume-role-processing"
+        "resources/config/assume-role-processing",
     ],
 )
 def test_processingjob(
@@ -80,8 +80,9 @@ def test_processingjob(
     for output in process_response["ProcessingOutputConfig"]["Outputs"]:
         assert processing_outputs[output["OutputName"]] == output["S3Output"]["S3Uri"]
 
-    assert not argo_utils.error_in_cw_logs(workflow_json["metadata"]["name"]), \
-        ('Found the CloudWatch error message in the log output. Check SageMaker to see if the job has failed.')
+    assert not argo_utils.error_in_cw_logs(
+        workflow_json["metadata"]["name"]
+    ), "Found the CloudWatch error message in the log output. Check SageMaker to see if the job has failed."
 
     utils.remove_dir(download_dir)
 
diff --git a/components/aws/sagemaker/tests/integration_tests/component_tests/test_train_component.py b/components/aws/sagemaker/tests/integration_tests/component_tests/test_train_component.py
--- a/components/aws/sagemaker/tests/integration_tests/component_tests/test_train_component.py
+++ b/components/aws/sagemaker/tests/integration_tests/component_tests/test_train_component.py
@@ -13,13 +13,10 @@
         pytest.param(
             "resources/config/simple-mnist-training", marks=pytest.mark.canary_test
         ),
-        pytest.param(
-            "resources/config/fsx-mnist-training",
-            marks=pytest.mark.fsx_test
-        ),
+        pytest.param("resources/config/fsx-mnist-training", marks=pytest.mark.fsx_test),
         "resources/config/spot-sample-pipeline-training",
         "resources/config/assume-role-training",
-        "resources/config/xgboost-mnist-trainingjob-debugger"
+        "resources/config/xgboost-mnist-trainingjob-debugger",
     ],
 )
 def test_trainingjob(
diff --git a/components/aws/sagemaker/tests/integration_tests/component_tests/test_workteam_component.py b/components/aws/sagemaker/tests/integration_tests/component_tests/test_workteam_component.py
--- a/components/aws/sagemaker/tests/integration_tests/component_tests/test_workteam_component.py
+++ b/components/aws/sagemaker/tests/integration_tests/component_tests/test_workteam_component.py
@@ -31,12 +31,7 @@ def create_workteamjob(
 
 @pytest.mark.parametrize(
     "test_file_dir",
-    [
-        pytest.param(
-            "resources/config/create-workteam",
-            marks=pytest.mark.canary_test
-        )
-    ],
+    [pytest.param("resources/config/create-workteam", marks=pytest.mark.canary_test)],
 )
 def test_workteamjob(
     kfp_client, experiment_id, region, sagemaker_client, test_file_dir
@@ -58,7 +53,12 @@ def test_workteamjob(
 
     try:
         workflow_json = create_workteamjob(
-            kfp_client, test_params, experiment_id, region, sagemaker_client, download_dir
+            kfp_client,
+            test_params,
+            experiment_id,
+            region,
+            sagemaker_client,
+            download_dir,
         )
 
         outputs = {"sagemaker-private-workforce": ["workteam_arn"]}
@@ -80,14 +80,11 @@ def test_workteamjob(
         assert response["Workteam"]["WorkteamArn"] == workteam_arn
 
     finally:
-        workteams = sagemaker_utils.list_workteams(
-            sagemaker_client
-        )["Workteams"]
+        workteams = sagemaker_utils.list_workteams(sagemaker_client)["Workteams"]
         workteam_names = list(map((lambda x: x["WorkteamName"]), workteams))
         # Check workteam was successfully created
         if workteam_name in workteam_names:
             sagemaker_utils.delete_workteam(sagemaker_client, workteam_name)
 
-
     # Delete generated files only if the test is successful
     utils.remove_dir(download_dir)
diff --git a/components/aws/sagemaker/tests/integration_tests/conftest.py b/components/aws/sagemaker/tests/integration_tests/conftest.py
--- a/components/aws/sagemaker/tests/integration_tests/conftest.py
+++ b/components/aws/sagemaker/tests/integration_tests/conftest.py
@@ -146,7 +146,7 @@ def get_experiment_id(kfp_client):
 
 @pytest.fixture(scope="session")
 def experiment_id(kfp_client, tmp_path_factory, worker_id):
-    if not worker_id:
+    if worker_id == "master":
         return get_experiment_id(kfp_client)
 
     # Locking taking as an example from
diff --git a/components/aws/sagemaker/tests/integration_tests/environment.yml b/components/aws/sagemaker/tests/integration_tests/environment.yml
--- a/components/aws/sagemaker/tests/integration_tests/environment.yml
+++ b/components/aws/sagemaker/tests/integration_tests/environment.yml
@@ -5,8 +5,8 @@ channels:
 dependencies:
   - python=3.7.*
   - pip=20.0.*
-  - awscli=1.18.*
-  - boto3=1.12.*
+  - awscli=1.18.141
+  - boto3=1.15.0
   - pytest-xdist=2.1.0
   - pyyaml=5.3.*
   - flake8=3.7.*
diff --git a/components/aws/sagemaker/tests/integration_tests/resources/config/kmeans-mnist-update-endpoint/config.yaml b/components/aws/sagemaker/tests/integration_tests/resources/config/kmeans-mnist-update-endpoint/config.yaml
new file mode 100644
--- /dev/null
+++ b/components/aws/sagemaker/tests/integration_tests/resources/config/kmeans-mnist-update-endpoint/config.yaml
@@ -0,0 +1,24 @@
+PipelineDefinition: resources/definition/update_endpoint_pipeline.py
+TestName: kmeans-update-endpoint-test
+Timeout: 7200
+ExpectedInstanceType: ml.m5.xlarge
+ExpectedPrediction:
+  predictions:
+    - distance_to_cluster: 7.448746204376221
+      closest_cluster: 2.0
+Arguments:
+  region: ((REGION))
+  model_name: kmeans-mnist-model
+  endpoint_config_name: kmeans-mnist-deploy-model
+  endpoint_name: kmeans-mnist-deploy-model
+  image: ((KMEANS_REGISTRY)).dkr.ecr.((REGION)).amazonaws.com/kmeans:1
+  model_artifact_url: s3://((DATA_BUCKET))/mnist_kmeans_example/model/kmeans-mnist-model/model.tar.gz
+  variant_name_1: variant-1
+  initial_variant_weight_1: 1.0
+  instance_type_1: ml.m4.xlarge
+  instance_type_2: ml.m5.xlarge
+  initial_instance_count_1: 1
+  network_isolation: "True"
+  role: ((ROLE_ARN))
+  update_endpoint: "True"
+  
diff --git a/components/aws/sagemaker/tests/integration_tests/resources/definition/create_endpoint_pipeline.py b/components/aws/sagemaker/tests/integration_tests/resources/definition/create_endpoint_pipeline.py
--- a/components/aws/sagemaker/tests/integration_tests/resources/definition/create_endpoint_pipeline.py
+++ b/components/aws/sagemaker/tests/integration_tests/resources/definition/create_endpoint_pipeline.py
@@ -47,6 +47,8 @@ def create_endpoint_pipeline(
         initial_variant_weight_1=initial_variant_weight_1,
     )
 
+    dsl.get_pipeline_conf().set_image_pull_policy(policy="Always")
+
 
 if __name__ == "__main__":
     kfp.compiler.Compiler().compile(
diff --git a/components/aws/sagemaker/tests/integration_tests/resources/definition/update_endpoint_pipeline.py b/components/aws/sagemaker/tests/integration_tests/resources/definition/update_endpoint_pipeline.py
new file mode 100644
--- /dev/null
+++ b/components/aws/sagemaker/tests/integration_tests/resources/definition/update_endpoint_pipeline.py
@@ -0,0 +1,71 @@
+import kfp
+from kfp import components
+from kfp import dsl
+
+sagemaker_model_op = components.load_component_from_file("../../model/component.yaml")
+sagemaker_deploy_op = components.load_component_from_file("../../deploy/component.yaml")
+
+
+@dsl.pipeline(
+    name="Create Hosting Endpoint in SageMaker",
+    description="SageMaker deploy component test",
+)
+def update_endpoint_pipeline(
+    region="",
+    endpoint_url="",
+    image="",
+    model_name="",
+    endpoint_config_name="",
+    endpoint_name="",
+    model_artifact_url="",
+    variant_name_1="",
+    instance_type_1="",
+    instance_type_2="",
+    initial_instance_count_1="",
+    initial_variant_weight_1="",
+    network_isolation="",
+    role="",
+    update_endpoint="",
+):
+    create_model = sagemaker_model_op(
+        region=region,
+        endpoint_url=endpoint_url,
+        model_name=model_name,
+        image=image,
+        model_artifact_url=model_artifact_url,
+        network_isolation=network_isolation,
+        role=role,
+    )
+
+    deploy_model = sagemaker_deploy_op(
+        region=region,
+        endpoint_url=endpoint_url,
+        endpoint_config_name=endpoint_config_name,
+        endpoint_name=endpoint_name,
+        model_name_1=create_model.output,
+        variant_name_1=variant_name_1,
+        instance_type_1=instance_type_1,
+        initial_instance_count_1=initial_instance_count_1,
+        initial_variant_weight_1=initial_variant_weight_1,
+    )
+
+    update_model = sagemaker_deploy_op(
+        region=region,
+        endpoint_url=endpoint_url,
+        endpoint_config_name=endpoint_config_name,
+        endpoint_name=deploy_model.output,
+        model_name_1=create_model.output,
+        variant_name_1=variant_name_1,
+        instance_type_1=instance_type_2,
+        initial_instance_count_1=initial_instance_count_1,
+        initial_variant_weight_1=initial_variant_weight_1,
+        update_endpoint=update_endpoint,
+    )
+
+    dsl.get_pipeline_conf().set_image_pull_policy(policy="Always")
+
+
+if __name__ == "__main__":
+    kfp.compiler.Compiler().compile(
+        create_endpoint_pipeline, "SageMaker_update_hosting_pipeline" + ".yaml"
+    )
diff --git a/components/aws/sagemaker/tests/integration_tests/scripts/kmeans_preprocessing.py b/components/aws/sagemaker/tests/integration_tests/scripts/kmeans_preprocessing.py
--- a/components/aws/sagemaker/tests/integration_tests/scripts/kmeans_preprocessing.py
+++ b/components/aws/sagemaker/tests/integration_tests/scripts/kmeans_preprocessing.py
@@ -6,19 +6,24 @@
 
 print("Extracting MNIST data set")
 # Load the dataset
-with gzip.open('/opt/ml/processing/input/mnist.pkl.gz', 'rb') as f:
-    train_set, valid_set, test_set = pickle.load(f, encoding='latin1')
+with gzip.open("/opt/ml/processing/input/mnist.pkl.gz", "rb") as f:
+    train_set, valid_set, test_set = pickle.load(f, encoding="latin1")
 
 # process the data
 # Convert the training data into the format required by the SageMaker KMeans algorithm
 print("Writing training data")
-with open('/opt/ml/processing/output_train/train_data', 'wb') as train_file:
+with open("/opt/ml/processing/output_train/train_data", "wb") as train_file:
     write_numpy_to_dense_tensor(train_file, train_set[0], train_set[1])
 
 print("Writing test data")
-with open('/opt/ml/processing/output_test/test_data', 'wb') as test_file:
+with open("/opt/ml/processing/output_test/test_data", "wb") as test_file:
     write_numpy_to_dense_tensor(test_file, test_set[0], test_set[1])
 
 print("Writing validation data")
 # Convert the valid data into the format required by the SageMaker KMeans algorithm
-numpy.savetxt('/opt/ml/processing/output_valid/valid-data.csv', valid_set[0], delimiter=',', fmt='%g')
+numpy.savetxt(
+    "/opt/ml/processing/output_valid/valid-data.csv",
+    valid_set[0],
+    delimiter=",",
+    fmt="%g",
+)
diff --git a/components/aws/sagemaker/tests/integration_tests/scripts/run_integration_tests b/components/aws/sagemaker/tests/integration_tests/scripts/run_integration_tests
--- a/components/aws/sagemaker/tests/integration_tests/scripts/run_integration_tests
+++ b/components/aws/sagemaker/tests/integration_tests/scripts/run_integration_tests
@@ -256,4 +256,4 @@ fi
 
 DIR_THIS_FILE="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 
-cd $DIR_THIS_FILE/../ && python -m pytest "${pytest_args[@]}" --junitxml ./integration_tests.log -n auto
+cd $DIR_THIS_FILE/../ && python -m pytest "${pytest_args[@]}" --junitxml ./integration_tests.log -n 10
diff --git a/components/aws/sagemaker/tests/integration_tests/utils/__init__.py b/components/aws/sagemaker/tests/integration_tests/utils/__init__.py
--- a/components/aws/sagemaker/tests/integration_tests/utils/__init__.py
+++ b/components/aws/sagemaker/tests/integration_tests/utils/__init__.py
@@ -85,12 +85,14 @@ def replace_placeholders(input_filename, output_filename):
         "((ROLE_ARN))": get_role_arn(),
         "((DATA_BUCKET))": get_s3_data_bucket(),
         "((KMEANS_REGISTRY))": get_algorithm_image_registry("kmeans", region, "1"),
-        "((XGBOOST_REGISTRY))": get_algorithm_image_registry("xgboost", region, "1.0-1"),
+        "((XGBOOST_REGISTRY))": get_algorithm_image_registry(
+            "xgboost", region, "1.0-1"
+        ),
         "((BUILTIN_RULE_IMAGE))": get_algorithm_image_registry("debugger", region),
         "((FSX_ID))": get_fsx_id(),
         "((FSX_SUBNET))": get_fsx_subnet(),
         "((FSX_SECURITY_GROUP))": get_fsx_security_group(),
-        "((ASSUME_ROLE_ARN))": get_assume_role_arn()
+        "((ASSUME_ROLE_ARN))": get_assume_role_arn(),
     }
 
     filedata = ""
diff --git a/components/aws/sagemaker/tests/integration_tests/utils/argo_utils.py b/components/aws/sagemaker/tests/integration_tests/utils/argo_utils.py
--- a/components/aws/sagemaker/tests/integration_tests/utils/argo_utils.py
+++ b/components/aws/sagemaker/tests/integration_tests/utils/argo_utils.py
@@ -5,13 +5,18 @@ def print_workflow_logs(workflow_name):
     output = get_workflow_logs(workflow_name)
     print(f"workflow logs:\n", output.decode())
 
+
 def find_in_logs(workflow_name, sub_str):
     logs = get_workflow_logs(workflow_name).decode()
     return logs.find(sub_str) >= 0
 
+
 def get_workflow_logs(workflow_name):
-    return utils.run_command(f"argo logs {workflow_name} -n {utils.get_kfp_namespace()}")
+    return utils.run_command(
+        f"argo logs {workflow_name} -n {utils.get_kfp_namespace()}"
+    )
+
 
 def error_in_cw_logs(workflow_name):
-    ERROR_MESSAGE = 'Error in fetching CloudWatch logs for SageMaker job'
-    return find_in_logs(workflow_name, ERROR_MESSAGE)
\ No newline at end of file
+    ERROR_MESSAGE = "Error in fetching CloudWatch logs for SageMaker job"
+    return find_in_logs(workflow_name, ERROR_MESSAGE)
diff --git a/components/aws/sagemaker/tests/integration_tests/utils/sagemaker_utils.py b/components/aws/sagemaker/tests/integration_tests/utils/sagemaker_utils.py
--- a/components/aws/sagemaker/tests/integration_tests/utils/sagemaker_utils.py
+++ b/components/aws/sagemaker/tests/integration_tests/utils/sagemaker_utils.py
@@ -9,9 +9,15 @@ def describe_model(client, model_name):
 def describe_endpoint(client, endpoint_name):
     return client.describe_endpoint(EndpointName=endpoint_name)
 
+
 def list_endpoints(client, name_contains):
     return client.list_endpoints(NameContains=name_contains)
 
+
+def describe_endpoint_config(client, endpoint_config_name):
+    return client.describe_endpoint_config(EndpointConfigName=endpoint_config_name)
+
+
 def delete_endpoint(client, endpoint_name):
     client.delete_endpoint(EndpointName=endpoint_name)
     waiter = client.get_waiter("endpoint_deleted")
diff --git a/components/aws/sagemaker/tests/unit_tests/tests/test_deploy.py b/components/aws/sagemaker/tests/unit_tests/tests/test_deploy.py
--- a/components/aws/sagemaker/tests/unit_tests/tests/test_deploy.py
+++ b/components/aws/sagemaker/tests/unit_tests/tests/test_deploy.py
@@ -41,6 +41,50 @@ def test_main(self):
       call('/tmp/output', 'test-endpoint-name')
     ])
 
+  def test_main_update_no_endpoint(self):
+    # Mock out all of utils except parser
+    deploy._utils = MagicMock()
+    deploy._utils.add_default_client_arguments = _utils.add_default_client_arguments
+
+    # Update is set to true, but endpoint does not exist
+    # Endpoint should be created in this case
+    # Set some static returns
+    update_endpoint_args = required_args + ['--update_endpoint', 'True']
+    deploy._utils.endpoint_name_exists.return_value = False
+    deploy._utils.deploy_model.return_value = 'test-endpoint-name'
+    deploy.main(update_endpoint_args)
+
+    # Check if correct requests were created and triggered
+    deploy._utils.deploy_model.assert_called()
+    deploy._utils.wait_for_endpoint_creation.assert_called()
+
+    # Check the file outputs
+    deploy._utils.write_output.assert_has_calls([
+      call('/tmp/output', 'test-endpoint-name')
+    ])
+
+  def test_main_update_endpoint(self):
+    # Update is set to true and endpoint exits
+    # Mock out all of utils except parser
+    deploy._utils = MagicMock()
+    deploy._utils.add_default_client_arguments = _utils.add_default_client_arguments
+
+    # Set some static returns
+    update_endpoint_args = required_args + ['--update_endpoint', 'True']
+    deploy._utils.endpoint_name_exists.return_value = True
+    deploy._utils.update_deployed_model.return_value = 'test-endpoint-name'
+    deploy.main(update_endpoint_args)
+
+    # Check if correct requests were created and triggered
+    deploy._utils.update_deployed_model.assert_called()
+    deploy._utils.wait_for_endpoint_creation.assert_called()
+    deploy._utils.delete_endpoint_config.assert_called()
+
+    # Check the file outputs
+    deploy._utils.write_output.assert_has_calls([
+      call('/tmp/output', 'test-endpoint-name')
+    ])
+
   def test_main_assumes_role(self):
     # Mock out all of utils except parser
     deploy._utils = MagicMock()
@@ -71,6 +115,22 @@ def test_deploy_model(self):
 
     self.assertEqual(response, 'test-endpoint-name')
 
+  def test_update_deployed_model(self):
+    mock_client = MagicMock()
+    mock_args = self.parser.parse_args(required_args + ['--endpoint_name', 'test-endpoint-name', '--endpoint_config_name', 'test-endpoint-config-name'])
+    response = _utils.update_deployed_model(mock_client, vars(mock_args))
+
+    mock_client.create_endpoint_config.assert_called_once_with(
+      EndpointConfigName='test-endpoint-config-name',
+      ProductionVariants=[
+        {'VariantName': 'variant-name-1', 'ModelName': 'model-test', 'InitialInstanceCount': 1,
+         'InstanceType': 'ml.m4.xlarge', 'InitialVariantWeight': 1.0}
+      ],
+      Tags=[]
+    )
+
+    self.assertEqual(response, 'test-endpoint-name')
+
   def test_sagemaker_exception_in_deploy_model(self):
     mock_client = MagicMock()
     mock_exception = ClientError({"Error": {"Message": "SageMaker broke"}}, "deploy_model")
@@ -80,6 +140,15 @@ def test_sagemaker_exception_in_deploy_model(self):
     with self.assertRaises(Exception):
       _utils.deploy_model(mock_client, vars(mock_args))
 
+  def test_sagemaker_exception_in_update_deployed_model(self):
+    mock_client = MagicMock()
+    mock_exception = ClientError({"Error": {"Message": "SageMaker broke"}}, "update_deployed_model")
+    mock_client.create_endpoint_config.side_effect = mock_exception
+    mock_args = self.parser.parse_args(required_args)
+
+    with self.assertRaises(Exception):
+      _utils.update_deployed_model(mock_client, vars(mock_args))
+
   def test_model_name_exception(self):
     mock_client = MagicMock()
     mock_args = vars(self.parser.parse_args(required_args))
@@ -96,6 +165,38 @@ def test_create_endpoint_exception(self):
     with self.assertRaises(Exception):
       _utils.create_endpoint(mock_client, 'us-east-1', 'fake-endpoint', 'fake-endpoint-config', {})
 
+  def test_get_endpoint_config_exception(self):
+    mock_client = MagicMock()
+    mock_exception = ClientError({"Error": {"Message": "SageMaker broke"}}, "describe_endpoint")
+    mock_client.describe_endpoint.side_effect = mock_exception
+
+    response = _utils.get_endpoint_config(mock_client, 'fake-endpoint')
+    self.assertEqual(response, None)
+
+  def test_delete_endpoint_config_exception(self):
+    mock_client = MagicMock()
+    mock_exception = ClientError({"Error": {"Message": "SageMaker broke"}}, "delete_endpoint")
+    mock_client.delete_endpoint_config.side_effect = mock_exception
+
+    response = _utils.delete_endpoint_config(mock_client, 'fake-endpoint-config')
+    self.assertEqual(response, False)
+
+  def test_endpoint_name_exists(self):
+    mock_client = MagicMock()
+    mock_exception = ClientError({"Error": {"Message": "SageMaker broke"}}, "delete_endpoint")
+    mock_client.describe_endpoint.side_effect = mock_exception
+
+    response = _utils.endpoint_name_exists(mock_client, 'fake-endpoint')
+    self.assertEqual(response, False)
+
+  def test_endpoint_config_name_exists_exception(self):
+    mock_client = MagicMock()
+    mock_exception = ClientError({"Error": {"Message": "SageMaker broke"}}, "describe_endpoint_config")
+    mock_client.describe_endpoint_config.side_effect = mock_exception
+
+    response = _utils.endpoint_config_name_exists(mock_client, 'fake-endpoint-config')
+    self.assertEqual(response, False)
+
   def test_wait_for_endpoint_creation(self):
     mock_client = MagicMock()
     mock_client.describe_endpoint.side_effect = [
@@ -123,13 +224,14 @@ def test_wait_for_failed_job(self):
   def test_get_endpoint_name_from_job(self):
     mock_client = MagicMock()
 
-    # if we don't pass --endpoint_name argument then endpoint name is constructed using --model_name_1
-    self.assertEqual(_utils.deploy_model(mock_client, vars(self.parser.parse_args(required_args))), 'Endpoint-test')
+    # if we don't pass --endpoint_name argument then endpoint name is constructed
+    self.assertTrue('Endpoint-' in _utils.deploy_model(mock_client, vars(self.parser.parse_args(required_args))))
 
   def test_pass_most_args(self):
     arguments =  [
       '--region', 'us-west-2',
       '--endpoint_url', 'fake-url',
+      '--endpoint_config_name','EndpointConfig-test-1',
       '--model_name_1', 'model-test-1',
       '--accelerator_type_1', 'ml.eia1.medium',
       '--model_name_2', 'model-test-2',