fixes for test flakyness post 0.0.4 release (#97)

Description of changes:
- there is a slight delay sometimes when tags are available on resource post creation
- adopted_endpoint resource was missing deleting the cr which caused false positive error message clutter in logs
- Job stopping sometimes takes a longer time than timeout
- minor linting and clean up


Testing:
PR build

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

Author: surajkota

test/e2e/__init__.py

@@ -313,3 +313,35 @@ def assert_training_status_in_sync(training_job_name, reference, expected_status
         == wait_resource_training_status(reference, expected_status)
         == expected_status
     )
+
+
+def get_sagemaker_endpoint(endpoint_name: str):
+    try:
+        return sagemaker_client().describe_endpoint(EndpointName=endpoint_name)
+    except botocore.exceptions.ClientError as error:
+        logging.error(
+            f"SageMaker could not find a endpoint with the name {endpoint_name}. Error {error}"
+        )
+        return None
+
+
+def get_sagemaker_model(model_name: str):
+    try:
+        return sagemaker_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}"
+        )
+        return None
+
+
+def get_sagemaker_endpoint_config(config_name: str):
+    try:
+        return sagemaker_client().describe_endpoint_config(
+            EndpointConfigName=config_name
+        )
+    except botocore.exceptions.ClientError as error:
+        logging.error(
+            f"SageMaker could not find an endpoint config with the name {config_name}. Error {error}"
+        )
+        return None

test/e2e/common/config.py

@@ -33,5 +33,7 @@
 DELETE_WAIT_PERIOD = 4
 DELETE_WAIT_LENGTH = 30
 
-JOB_DELETE_WAIT_PERIODS = 8
+JOB_DELETE_WAIT_PERIODS = 12
 JOB_DELETE_WAIT_LENGTH = 30
+
+TAG_DELAY_SLEEP = 20

test/e2e/service_bootstrap.py

@@ -18,6 +18,7 @@
 import logging
 import time
 import subprocess
+import random
 
 from acktest import resources
 from acktest.aws.identity import get_region, get_account_id
@@ -48,10 +49,15 @@ def create_execution_role() -> str:
         Description="SageMaker execution role for ACK integration and canary tests",
     )
 
+    # random sleep to prevent throttling
+    time.sleep(random.randrange(1, 3))
     iam.attach_role_policy(
         RoleName=role_name,
         PolicyArn="arn:aws:iam::aws:policy/AmazonSageMakerFullAccess",
     )
+
+    # random sleep to prevent throttling
+    time.sleep(random.randrange(1, 3))
     iam.attach_role_policy(
         RoleName=role_name, PolicyArn="arn:aws:iam::aws:policy/AmazonS3FullAccess"
     )

test/e2e/tests/test_adopt_endpoint.py

@@ -13,13 +13,9 @@
 """Integration tests for the SageMaker Endpoint API.
 """
 
-import boto3
 import pytest
-import logging
-import time
 from typing import Dict
 
-from acktest.aws import s3
 from acktest.resources import random_suffix_name
 from acktest.k8s import resource as k8s
 
@@ -31,6 +27,9 @@
     wait_sagemaker_endpoint_status,
     assert_endpoint_status_in_sync,
     sagemaker_client,
+    get_sagemaker_endpoint,
+    get_sagemaker_endpoint_config,
+    get_sagemaker_model,
 )
 from e2e.replacement_values import REPLACEMENT_VALUES
 from e2e.common import config as cfg
@@ -112,10 +111,16 @@ def sdk_endpoint(name_suffix):
         endpoint_input,
         endpoint_response,
     )
-    wait_sagemaker_endpoint_status(endpoint_name, cfg.ENDPOINT_STATUS_INSERVICE)
-    sagemaker_client().delete_endpoint(EndpointName=endpoint_name)
-    sagemaker_client().delete_endpoint_config(EndpointConfigName=endpoint_config_name)
-    sagemaker_client().delete_model(ModelName=model_name)
+
+    if get_sagemaker_endpoint(endpoint_name) is not None:
+        wait_sagemaker_endpoint_status(endpoint_name, cfg.ENDPOINT_STATUS_INSERVICE)
+        sagemaker_client().delete_endpoint(EndpointName=endpoint_name)
+    if get_sagemaker_endpoint_config(endpoint_config_name) is not None:
+        sagemaker_client().delete_endpoint_config(
+            EndpointConfigName=endpoint_config_name
+        )
+    if get_sagemaker_model(model_name) is not None:
+        sagemaker_client().delete_model(ModelName=model_name)
 
 
 @pytest.fixture(scope="module")
@@ -262,3 +267,9 @@ def test_smoke(self, sdk_endpoint, adopted_endpoint):
             endpoint_name, endpoint_reference, cfg.ENDPOINT_STATUS_INSERVICE,
         )
         assert k8s.wait_on_condition(endpoint_reference, "ACK.ResourceSynced", "True")
+
+        for cr in (model_reference, config_reference, endpoint_reference):
+            _, deleted = k8s.delete_custom_resource(
+                cr, cfg.DELETE_WAIT_PERIOD, cfg.DELETE_WAIT_LENGTH
+            )
+            assert deleted

test/e2e/tests/test_adopt_model_package.py

@@ -242,16 +242,8 @@ def test_smoke(self, sdk_model_package, adopted_model_package):
             model_package_reference, "ACK.ResourceSynced", "True"
         )
 
-        _, deleted = k8s.delete_custom_resource(
-            model_package_reference,
-            cfg.JOB_DELETE_WAIT_PERIODS,
-            cfg.JOB_DELETE_WAIT_LENGTH,
-        )
-        assert deleted is True
-
-        _, deleted = k8s.delete_custom_resource(
-            model_package_group_reference,
-            cfg.JOB_DELETE_WAIT_PERIODS,
-            cfg.JOB_DELETE_WAIT_LENGTH,
-        )
-        assert deleted is True
+        for cr in (model_package_reference, model_package_group_reference):
+            _, deleted = k8s.delete_custom_resource(
+                cr, cfg.JOB_DELETE_WAIT_PERIODS, cfg.JOB_DELETE_WAIT_LENGTH
+            )
+            assert deleted

test/e2e/tests/test_data_quality_job_definition.py

@@ -16,12 +16,14 @@
 import pytest
 import logging
 import botocore
+import time
 
 from e2e import service_marker, assert_tags_in_sync
 from e2e.common.fixtures import (
     xgboost_churn_data_quality_job_definition,
     xgboost_churn_endpoint,
 )
+from e2e.common import config as cfg
 from acktest.k8s import resource as k8s
 
 # Access variable so it is loaded as a fixture
@@ -56,6 +58,8 @@ def test_smoke(self, sagemaker_client, xgboost_churn_data_quality_job_definition
         job_definition_arn = job_definition_desc["JobDefinitionArn"]
         assert k8s.get_resource_arn(resource) == job_definition_arn
 
+        # random sleep before we check for tags to reduce test flakyness
+        time.sleep(cfg.TAG_DELAY_SLEEP)
         resource_tags = resource["spec"].get("tags", None)
         assert_tags_in_sync(job_definition_arn, resource_tags)
         # Delete the k8s resource.

test/e2e/tests/test_endpoint.py

@@ -13,11 +13,9 @@
 """Integration tests for the SageMaker Endpoint API.
 """
 
-import botocore
-from botocore import endpoint
 import pytest
 import logging
-import time
+
 from typing import Dict
 
 from acktest.aws import s3
@@ -29,13 +27,17 @@
     create_sagemaker_resource,
     assert_endpoint_status_in_sync,
     assert_tags_in_sync,
+    get_sagemaker_endpoint,
 )
 from e2e.replacement_values import REPLACEMENT_VALUES
 from e2e.common import config as cfg
 
 FAIL_UPDATE_ERROR_MESSAGE = "EndpointUpdateError: unable to update endpoint. check FailureReason. latest EndpointConfigName is "
 # annontation key for last endpoint config name used for update
-LAST_ENDPOINTCONFIG_UPDATE_ANNOTATION = "sagemaker.services.k8s.aws/last-endpoint-config-for-update"
+LAST_ENDPOINTCONFIG_UPDATE_ANNOTATION = (
+    "sagemaker.services.k8s.aws/last-endpoint-config-for-update"
+)
+
 
 @pytest.fixture(scope="module")
 def name_suffix():
@@ -212,35 +214,17 @@ def faulty_config(name_suffix, single_container_model):
 @service_marker
 @pytest.mark.canary
 class TestEndpoint:
-    def _get_resource_endpoint_arn(self, resource: Dict):
-        assert (
-            "ackResourceMetadata" in resource["status"]
-            and "arn" in resource["status"]["ackResourceMetadata"]
-        )
-        return resource["status"]["ackResourceMetadata"]["arn"]
-
-    def _describe_sagemaker_endpoint(self, sagemaker_client, endpoint_name: str):
-        try:
-            return sagemaker_client.describe_endpoint(EndpointName=endpoint_name)
-        except botocore.exceptions.ClientError as error:
-            logging.error(
-                f"SageMaker could not find a endpoint with the name {endpoint_name}. Error {error}"
-            )
-            return None
-
-    def create_endpoint_test(self, sagemaker_client, xgboost_endpoint):
+    def create_endpoint_test(self, xgboost_endpoint):
         (reference, resource, _) = xgboost_endpoint
         assert k8s.get_resource_exists(reference)
 
         # endpoint has correct arn and status
         endpoint_name = resource["spec"].get("endpointName", None)
         assert endpoint_name is not None
 
-        endpoint_desc = self._describe_sagemaker_endpoint(
-            sagemaker_client, endpoint_name
-        )
+        endpoint_desc = get_sagemaker_endpoint(endpoint_name)
         endpoint_arn = endpoint_desc["EndpointArn"]
-        assert self._get_resource_endpoint_arn(resource) == endpoint_arn
+        assert k8s.get_resource_arn(resource) == endpoint_arn
 
         # endpoint transitions Creating -> InService state
         assert_endpoint_status_in_sync(
@@ -257,7 +241,7 @@ def create_endpoint_test(self, sagemaker_client, xgboost_endpoint):
         assert_tags_in_sync(endpoint_arn, resource_tags)
 
     def update_endpoint_failed_test(
-        self, sagemaker_client, single_variant_config, faulty_config, xgboost_endpoint
+        self, single_variant_config, faulty_config, xgboost_endpoint
     ):
         (endpoint_reference, _, endpoint_spec) = xgboost_endpoint
         (_, faulty_config_resource) = faulty_config
@@ -284,25 +268,28 @@ def update_endpoint_failed_test(
         )
 
         assert k8s.wait_on_condition(endpoint_reference, "ACK.ResourceSynced", "False")
-        
+
         (_, old_config_resource) = single_variant_config
-        current_config_name = old_config_resource["spec"].get("endpointConfigName", None)
+        current_config_name = old_config_resource["spec"].get(
+            "endpointConfigName", None
+        )
         assert k8s.assert_condition_state_message(
-            endpoint_reference, "ACK.Terminal", "True", FAIL_UPDATE_ERROR_MESSAGE+current_config_name,
+            endpoint_reference,
+            "ACK.Terminal",
+            "True",
+            FAIL_UPDATE_ERROR_MESSAGE + current_config_name,
         )
 
         endpoint_resource = k8s.get_resource(endpoint_reference)
         assert endpoint_resource["status"].get("failureReason", None) is not None
 
-    def update_endpoint_successful_test(
-        self, sagemaker_client, multi_variant_config, xgboost_endpoint
-    ):
+    def update_endpoint_successful_test(self, multi_variant_config, xgboost_endpoint):
         (endpoint_reference, endpoint_resource, endpoint_spec) = xgboost_endpoint
 
         endpoint_name = endpoint_resource["spec"].get("endpointName", None)
-        production_variants = self._describe_sagemaker_endpoint(
-            sagemaker_client, endpoint_name
-        )["ProductionVariants"]
+        production_variants = get_sagemaker_endpoint(endpoint_name)[
+            "ProductionVariants"
+        ]
         old_variant_instance_count = production_variants[0]["CurrentInstanceCount"]
         old_variant_name = production_variants[0]["VariantName"]
 
@@ -338,9 +325,9 @@ def update_endpoint_successful_test(
         assert endpoint_resource["status"].get("failureReason", None) is None
 
         # RetainAllVariantProperties - variant properties were retained + is a multi-variant endpoint
-        new_production_variants = self._describe_sagemaker_endpoint(
-            sagemaker_client, endpoint_name
-        )["ProductionVariants"]
+        new_production_variants = get_sagemaker_endpoint(endpoint_name)[
+            "ProductionVariants"
+        ]
         assert len(new_production_variants) > 1
         new_variant_instance_count = None
         for variant in new_production_variants:
@@ -349,14 +336,16 @@ def update_endpoint_successful_test(
 
         assert new_variant_instance_count == old_variant_instance_count
 
-    def delete_endpoint_test(self, sagemaker_client, xgboost_endpoint):
+    def delete_endpoint_test(self, xgboost_endpoint):
         (reference, resource, _) = xgboost_endpoint
         endpoint_name = resource["spec"].get("endpointName", None)
 
-        _, deleted = k8s.delete_custom_resource(reference, cfg.DELETE_WAIT_PERIOD, cfg.DELETE_WAIT_LENGTH)
+        _, deleted = k8s.delete_custom_resource(
+            reference, cfg.DELETE_WAIT_PERIOD, cfg.DELETE_WAIT_LENGTH
+        )
         assert deleted
 
-        assert self._describe_sagemaker_endpoint(sagemaker_client, endpoint_name) is None
+        assert get_sagemaker_endpoint(endpoint_name) is None
 
     def test_driver(
         self,
@@ -366,13 +355,11 @@ def test_driver(
         multi_variant_config,
         xgboost_endpoint,
     ):
-        self.create_endpoint_test(sagemaker_client, xgboost_endpoint)
+        self.create_endpoint_test(xgboost_endpoint)
         self.update_endpoint_failed_test(
-            sagemaker_client, single_variant_config, faulty_config, xgboost_endpoint
+            single_variant_config, faulty_config, xgboost_endpoint
         )
         # Note: the test has been intentionally ordered to run a successful update after a failed update
         # check that controller updates the endpoint, removes the terminal condition and clears the failure reason
-        self.update_endpoint_successful_test(
-            sagemaker_client, multi_variant_config, xgboost_endpoint
-        )
-        self.delete_endpoint_test(sagemaker_client, xgboost_endpoint)
+        self.update_endpoint_successful_test(multi_variant_config, xgboost_endpoint)
+        self.delete_endpoint_test(xgboost_endpoint)

test/e2e/tests/test_endpoint_config.py

@@ -13,19 +13,19 @@
 """Integration tests for the SageMaker EndpointConfig API.
 """
 
-import botocore
 import pytest
 import logging
 from typing import Dict
+import time
 
 from acktest.resources import random_suffix_name
 from acktest.k8s import resource as k8s
 
 from e2e import (
     service_marker,
     create_sagemaker_resource,
-    sagemaker_client,
     assert_tags_in_sync,
+    get_sagemaker_endpoint_config,
 )
 from e2e.replacement_values import REPLACEMENT_VALUES
 from e2e.common import config as cfg
@@ -70,18 +70,6 @@ def single_variant_config():
         assert deleted
 
 
-def get_sagemaker_endpoint_config(config_name: str):
-    try:
-        return sagemaker_client().describe_endpoint_config(
-            EndpointConfigName=config_name
-        )
-    except botocore.exceptions.ClientError as error:
-        logging.error(
-            f"SageMaker could not find a config with the name {config_name}. Error {error}"
-        )
-        return None
-
-
 @service_marker
 @pytest.mark.canary
 class TestEndpointConfig:
@@ -94,6 +82,8 @@ def test_create_endpoint_config(self, single_variant_config):
         endpoint_arn = endpoint_config_desc["EndpointConfigArn"]
         assert k8s.get_resource_arn(resource) == endpoint_arn
 
+        # random sleep before we check for tags to reduce test flakyness
+        time.sleep(cfg.TAG_DELAY_SLEEP)
         resource_tags = resource["spec"].get("tags", None)
         assert_tags_in_sync(endpoint_arn, resource_tags)
         # Delete the k8s resource.