Test cross-region support (#141)

Description of changes:

-  Test ACKs cross region support so that it does not need to be tested manually every release


By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

Author: ryansteakley

test/e2e/__init__.py

@@ -11,6 +11,7 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
+from acktest.aws.identity import get_region
 import pytest
 import logging
 import botocore
@@ -34,8 +35,9 @@
 resource_directory = Path(__file__).parent / "resources"
 
 
-def sagemaker_client():
-    return boto3.client("sagemaker")
+def sagemaker_client(region: str = None):
+    region = region or get_region()
+    return boto3.client("sagemaker", region_name=region)
 
 
 def create_sagemaker_resource(
@@ -325,9 +327,10 @@ def get_sagemaker_endpoint(endpoint_name: str):
         return None
 
 
-def get_sagemaker_model(model_name: str):
+def get_sagemaker_model(model_name: str, sm_client=None):
+    sm_client = sm_client or sagemaker_client()
     try:
-        return sagemaker_client().describe_model(ModelName=model_name)
+        return sm_client.describe_model(ModelName=model_name)
     except botocore.exceptions.ClientError as error:
         logging.error(
             f"SageMaker could not find a model with the name {model_name}. Error {error}"

test/e2e/resources/cross_region_model.yaml

@@ -0,0 +1,12 @@
+apiVersion: sagemaker.services.k8s.aws/v1alpha1
+kind: Model
+metadata:
+  name: $MODEL_NAME
+  annotations:
+    services.k8s.aws/region: $REGION
+spec:
+  modelName: $MODEL_NAME
+  primaryContainer:
+    containerHostname: xgboost
+    image: $XGBOOST_V1_IMAGE_URI
+  executionRoleARN: $SAGEMAKER_EXECUTION_ROLE_ARN

test/e2e/tests/test_cross_region.py

@@ -0,0 +1,85 @@
+# Copyright 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.
+"""Integration test for ACKs Cross Region Support.
+"""
+
+import logging
+from acktest.aws.identity import get_region
+import pytest
+import time
+from typing import Dict
+
+from acktest.resources import random_suffix_name
+from acktest.k8s import resource as k8s
+from e2e import (
+    service_marker,
+    create_sagemaker_resource,
+    get_sagemaker_model,
+    sagemaker_client,
+)
+from e2e.replacement_values import REPLACEMENT_VALUES, XGBOOST_V1_IMAGE_URIS
+from e2e.common import config as cfg
+
+
+@pytest.fixture(scope="module")
+def cross_region_model():
+    resource_name = random_suffix_name("cross-region-model", 32)
+    region = get_cross_region()
+
+    replacements = REPLACEMENT_VALUES.copy()
+    replacements["MODEL_NAME"] = resource_name
+    replacements["REGION"] = region
+    replacements[
+        "XGBOOST_V1_IMAGE_URI"
+    ] = f"{XGBOOST_V1_IMAGE_URIS[region]}/xgboost:latest"
+
+    reference, spec, resource = create_sagemaker_resource(
+        resource_plural=cfg.MODEL_RESOURCE_PLURAL,
+        resource_name=resource_name,
+        spec_file="cross_region_model",
+        replacements=replacements,
+    )
+    assert resource is not None
+
+    yield (reference, resource)
+
+    # Delete the k8s resource if not already deleted by tests
+    if k8s.get_resource_exists(reference):
+        _, deleted = k8s.delete_custom_resource(reference, 3, 10)
+        assert deleted
+
+
+def get_cross_region():
+    region = get_region()
+    if region != "us-west-2":
+        return "us-west-2"
+    return "us-east-1"
+
+
+@service_marker
+class TestCrossRegionModel:
+    def test_create_cross_region_model(self, cross_region_model):
+        (reference, resource) = cross_region_model
+        assert k8s.get_resource_exists(reference)
+
+        sm_client = sagemaker_client(get_cross_region())
+        model_name = resource["spec"].get("modelName", None)
+        model_desc = get_sagemaker_model(model_name, sm_client)
+        cross_region_model_arn = model_desc["ModelArn"]
+        assert k8s.get_resource_arn(resource) == cross_region_model_arn
+
+        # Delete the k8s resource.
+        _, deleted = k8s.delete_custom_resource(reference, 3, 10)
+        assert deleted
+
+        assert get_sagemaker_model(model_name, sm_client) is None