diff --git a/test/e2e/__init__.py b/test/e2e/__init__.py index 7fa31798..0443ca70 100644 --- a/test/e2e/__init__.py +++ b/test/e2e/__init__.py @@ -30,6 +30,7 @@ bootstrap_directory = Path(__file__).parent resource_directory = Path(__file__).parent / "resources" + def sagemaker_client(): return boto3.client("sagemaker") @@ -136,9 +137,7 @@ def wait_resource_endpoint_status( ) -def assert_endpoint_status_in_sync( - endpoint_name, reference, expected_status -): +def assert_endpoint_status_in_sync(endpoint_name, reference, expected_status): assert ( wait_sagemaker_endpoint_status(endpoint_name, expected_status) == wait_resource_endpoint_status(reference, expected_status, 2) diff --git a/test/e2e/common/fixtures.py b/test/e2e/common/fixtures.py index b8bd11c1..f8fc2c3b 100644 --- a/test/e2e/common/fixtures.py +++ b/test/e2e/common/fixtures.py @@ -74,9 +74,7 @@ def xgboost_churn_endpoint(sagemaker_client): ) assert endpoint_resource is not None assert k8s.get_resource_arn(endpoint_resource) is not None - wait_sagemaker_endpoint_status( - replacements["ENDPOINT_NAME"], "InService" - ) + wait_sagemaker_endpoint_status(replacements["ENDPOINT_NAME"], "InService") yield endpoint_spec diff --git a/test/e2e/service_bootstrap.py b/test/e2e/service_bootstrap.py index c2063b9e..b18ab588 100644 --- a/test/e2e/service_bootstrap.py +++ b/test/e2e/service_bootstrap.py @@ -93,9 +93,20 @@ def create_data_bucket() -> str: # duplicate_bucket_contents(source_bucket, destination_bucket) # workaround to copy if buckets are across regions # TODO: check if there is a better way and merge to test-infra - subprocess.call(['mkdir',f'{temp_dir}']) - subprocess.call(['aws', 's3', 'sync', f's3://{SAGEMAKER_SOURCE_DATA_BUCKET}', f'./{temp_dir}/', '--quiet']) - subprocess.call(['aws', 's3', 'sync', f'./{temp_dir}/', f's3://{bucket_name}', '--quiet']) + subprocess.call(["mkdir", f"{temp_dir}"]) + subprocess.call( + [ + "aws", + "s3", + "sync", + f"s3://{SAGEMAKER_SOURCE_DATA_BUCKET}", + f"./{temp_dir}/", + "--quiet", + ] + ) + subprocess.call( + ["aws", "s3", "sync", f"./{temp_dir}/", f"s3://{bucket_name}", "--quiet"] + ) logging.info(f"Synced data bucket") diff --git a/test/e2e/tests/test_adopt_endpoint.py b/test/e2e/tests/test_adopt_endpoint.py index dd4eefb5..e749779a 100644 --- a/test/e2e/tests/test_adopt_endpoint.py +++ b/test/e2e/tests/test_adopt_endpoint.py @@ -30,15 +30,17 @@ create_adopted_resource, wait_sagemaker_endpoint_status, assert_endpoint_status_in_sync, - sagemaker_client + sagemaker_client, ) from e2e.replacement_values import REPLACEMENT_VALUES from e2e.common import config as cfg + @pytest.fixture(scope="module") def name_suffix(): return random_suffix_name("sdk-endpoint", 32) + def sdk_make_model(model_name): data_bucket = REPLACEMENT_VALUES["SAGEMAKER_DATA_BUCKET"] model_input = { @@ -56,6 +58,7 @@ def sdk_make_model(model_name): assert model_response.get("ModelArn", None) is not None return model_input, model_response + def sdk_make_endpoint_config(model_name, endpoint_config_name): endpoint_config_input = { "EndpointConfigName": endpoint_config_name, @@ -75,6 +78,7 @@ def sdk_make_endpoint_config(model_name, endpoint_config_name): assert endpoint_config_response.get("EndpointConfigArn", None) is not None return endpoint_config_input, endpoint_config_response + def sdk_make_endpoint(endpoint_name, endpoint_config_name): endpoint_input = { "EndpointName": endpoint_name, @@ -85,6 +89,7 @@ def sdk_make_endpoint(endpoint_name, endpoint_config_name): return endpoint_input, endpoint_response + @pytest.fixture(scope="module") def sdk_endpoint(name_suffix): model_name = name_suffix + "-model" @@ -92,8 +97,12 @@ def sdk_endpoint(name_suffix): endpoint_name = name_suffix model_input, model_response = sdk_make_model(model_name) - endpoint_config_input, endpoint_config_response = sdk_make_endpoint_config(model_name, endpoint_config_name) - endpoint_input, endpoint_response = sdk_make_endpoint(endpoint_name, endpoint_config_name) + endpoint_config_input, endpoint_config_response = sdk_make_endpoint_config( + model_name, endpoint_config_name + ) + endpoint_input, endpoint_response = sdk_make_endpoint( + endpoint_name, endpoint_config_name + ) yield ( model_input, @@ -103,9 +112,7 @@ def sdk_endpoint(name_suffix): endpoint_input, endpoint_response, ) - wait_sagemaker_endpoint_status( - endpoint_name, cfg.ENDPOINT_STATUS_INSERVICE - ) + 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) @@ -234,7 +241,11 @@ def test_smoke(self, sdk_endpoint, adopted_endpoint): ) endpoint_reference = k8s.create_reference( - CRD_GROUP, CRD_VERSION, cfg.ENDPOINT_RESOURCE_PLURAL, endpoint_name, namespace + CRD_GROUP, + CRD_VERSION, + cfg.ENDPOINT_RESOURCE_PLURAL, + endpoint_name, + namespace, ) endpoint_resource = k8s.wait_resource_consumed_by_controller(endpoint_reference) assert endpoint_resource is not None diff --git a/test/e2e/tests/test_endpoint_config.py b/test/e2e/tests/test_endpoint_config.py index fd12f6fe..06239c3f 100644 --- a/test/e2e/tests/test_endpoint_config.py +++ b/test/e2e/tests/test_endpoint_config.py @@ -22,11 +22,7 @@ from acktest.resources import random_suffix_name from acktest.k8s import resource as k8s -from e2e import ( - service_marker, - create_sagemaker_resource, - sagemaker_client -) +from e2e import service_marker, create_sagemaker_resource, sagemaker_client from e2e.replacement_values import REPLACEMENT_VALUES from e2e.common import config as cfg @@ -65,6 +61,7 @@ def single_variant_config(): _, deleted = k8s.delete_custom_resource(config_reference, 3, 10) assert deleted + def get_sagemaker_endpoint_config(config_name: str): try: return sagemaker_client().describe_endpoint_config( @@ -86,11 +83,13 @@ def test_create_endpoint_config(self, single_variant_config): config_name = resource["spec"].get("endpointConfigName", None) - assert k8s.get_resource_arn(resource) == get_sagemaker_endpoint_config(config_name)["EndpointConfigArn"] + assert ( + k8s.get_resource_arn(resource) + == get_sagemaker_endpoint_config(config_name)["EndpointConfigArn"] + ) # Delete the k8s resource. _, deleted = k8s.delete_custom_resource(reference, 3, 10) assert deleted assert get_sagemaker_endpoint_config(config_name) is None - diff --git a/test/e2e/tests/test_model.py b/test/e2e/tests/test_model.py index f9bedb92..ebd0cae0 100644 --- a/test/e2e/tests/test_model.py +++ b/test/e2e/tests/test_model.py @@ -51,6 +51,7 @@ def xgboost_model(): _, deleted = k8s.delete_custom_resource(reference, 3, 10) assert deleted + def get_sagemaker_model(model_name: str): try: return sagemaker_client().describe_model(ModelName=model_name) @@ -59,6 +60,8 @@ def get_sagemaker_model(model_name: str): f"SageMaker could not find a model with the name {model_name}. Error {error}" ) return None + + @service_marker @pytest.mark.canary class TestModel: @@ -68,11 +71,13 @@ def test_create_model(self, xgboost_model): model_name = resource["spec"].get("modelName", None) - assert k8s.get_resource_arn(resource) == get_sagemaker_model(model_name)["ModelArn"] + assert ( + k8s.get_resource_arn(resource) + == get_sagemaker_model(model_name)["ModelArn"] + ) # Delete the k8s resource. _, deleted = k8s.delete_custom_resource(reference, 3, 10) assert deleted assert get_sagemaker_model(model_name) is None -