From fe71f18b9cbdd2ead13c4675baf263128afec428 Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Tue, 21 Jul 2026 14:31:59 -0700 Subject: [PATCH] fix(serve/integ): delete endpoints by name so cleanup survives deploy failures The sagemaker-serve integration tests each build a model and deploy a real SageMaker endpoint, then delete it in a finally block. Cleanup was guarded on the handle returned by deploy(): core_model = core_endpoint = None try: core_model, core_endpoint = build_and_deploy() # may fail / hang ... finally: if core_model and core_endpoint: # skipped on failure cleanup_resources(core_model, core_endpoint) When deploy() raises (endpoint -> Failed) or the build is killed mid-call (CodeBuild timeout), core_endpoint is never assigned, the guard is falsy, and the endpoint SageMaker already created is never torn down. Deleting by handle also can't clean up a resource whose creating call never returned. This is the root cause of the endpoint accumulation in the CI account (thousands of Failed and InService endpoints). Fix: add tests/integ/cleanup_helpers.py with delete_quietly() (best-effort, per-resource) and cleanup_by_name() (deletes endpoint -> endpoint config -> model by name). Each test now generates its resource names up front and always calls cleanup_by_name() in finally, so a Failed or half-created endpoint is still deleted, and one failed delete no longer aborts the rest. - Applies to: jumpstart, optimize, tgi, tei, huggingface, triton, train_inference_e2e, and both ai_inference_recommender tests. - ai_inference_recommender keeps its job / workload-config / model-package cleanup, now via the shared delete_quietly(); its local duplicate helpers are removed. - in_process is unchanged (deploy_local creates no cloud endpoint). This only changes test cleanup; no SDK source or test assertions are affected. --- .../tests/integ/cleanup_helpers.py | 104 ++++++++++++++++++ ...st_ai_inference_recommender_integration.py | 99 ++++++++--------- .../integ/test_huggingface_integration.py | 54 ++++----- .../tests/integ/test_jumpstart_integration.py | 55 ++++----- .../tests/integ/test_optimize_integration.py | 60 +++++----- .../tests/integ/test_tei_integration.py | 60 +++++----- .../tests/integ/test_tgi_integration.py | 56 ++++------ .../test_train_inference_e2e_integration.py | 71 ++++++------ .../tests/integ/test_triton_integration.py | 58 +++++----- 9 files changed, 326 insertions(+), 291 deletions(-) create mode 100644 sagemaker-serve/tests/integ/cleanup_helpers.py diff --git a/sagemaker-serve/tests/integ/cleanup_helpers.py b/sagemaker-serve/tests/integ/cleanup_helpers.py new file mode 100644 index 0000000000..b0acb6c445 --- /dev/null +++ b/sagemaker-serve/tests/integ/cleanup_helpers.py @@ -0,0 +1,104 @@ +# 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. +"""Shared, resilient resource cleanup for sagemaker-serve integration tests. + +These tests each ``build()`` a model and ``deploy()`` a real SageMaker endpoint, +then delete it in a ``finally`` block. The historical pattern guarded cleanup on +the object handle returned by ``deploy()``:: + + core_model = core_endpoint = None + try: + core_model, core_endpoint = build_and_deploy() # deploy() may fail/hang + ... + finally: + if core_model and core_endpoint: # <-- skipped on failure + cleanup_resources(core_model, core_endpoint) + +That guard is the root cause of the endpoint accumulation in the CI account: when +``deploy()`` raises (endpoint goes to ``Failed``) or is killed mid-call (build +timeout), ``core_endpoint`` is never assigned, so the ``if`` is falsy and the +endpoint SageMaker already created is never deleted. Deleting by the *handle* also +can't clean up a resource whose creation call never returned. + +The helpers here delete by **name** instead, so cleanup no longer depends on +``deploy()`` returning. A test captures the names it is about to use up front and +always calls :func:`cleanup_by_name` in ``finally``. Each delete is best-effort +(:func:`delete_quietly`): a missing resource or a transient error on one delete +never prevents the others from running. +""" +from __future__ import absolute_import + +import logging + +from sagemaker.core.resources import Endpoint, EndpointConfig, Model + +logger = logging.getLogger(__name__) + + +def delete_quietly(delete_fn, label): + """Run a single best-effort delete, logging and swallowing any failure. + + ``delete_fn`` is a zero-arg callable that fetches and deletes one resource + (e.g. ``lambda: Endpoint.get(endpoint_name=name).delete()``). It is expected + to raise when the resource does not exist (already cleaned up / never + created) or on a transient API error; either way we log and continue so a + later delete in the same cleanup sequence still runs. + """ + try: + delete_fn() + logger.info("Deleted %s", label) + except Exception as exc: # noqa: E722 - best-effort cleanup must never raise + # Debug, not warning: a "does not exist" here is the common, expected + # case (e.g. the endpoint config was never created because deploy() + # failed early), and we don't want to red the log for normal cleanup. + logger.debug("Skipping delete of %s: %s", label, exc) + + +def cleanup_by_name( + endpoint_name=None, + endpoint_config_name=None, + model_name=None, +): + """Best-effort delete of an endpoint, its endpoint config, and its model by name. + + Every argument is optional and independent; pass whatever names the test + intended to create. Deletion order mirrors reverse order of creation + (endpoint -> endpoint config -> model) so a live endpoint is torn down before + the config it references. Because deletes are keyed on name (not on the + objects ``deploy()`` returns), this cleans up resources even when the + creating call failed or never returned. + + ``endpoint_config_name`` defaults to ``endpoint_name`` when omitted, matching + ``ModelBuilder.deploy()``'s default of reusing the endpoint name for its + auto-created endpoint config. + """ + if endpoint_name and endpoint_config_name is None: + endpoint_config_name = endpoint_name + + if endpoint_name: + delete_quietly( + lambda: Endpoint.get(endpoint_name=endpoint_name).delete(), + f"Endpoint {endpoint_name}", + ) + if endpoint_config_name: + delete_quietly( + lambda: EndpointConfig.get( + endpoint_config_name=endpoint_config_name + ).delete(), + f"EndpointConfig {endpoint_config_name}", + ) + if model_name: + delete_quietly( + lambda: Model.get(model_name=model_name).delete(), + f"Model {model_name}", + ) diff --git a/sagemaker-serve/tests/integ/test_ai_inference_recommender_integration.py b/sagemaker-serve/tests/integ/test_ai_inference_recommender_integration.py index 5b0d5fd1fc..da83603bee 100644 --- a/sagemaker-serve/tests/integ/test_ai_inference_recommender_integration.py +++ b/sagemaker-serve/tests/integ/test_ai_inference_recommender_integration.py @@ -25,7 +25,6 @@ AIBenchmarkJob, AIRecommendationJob, AIWorkloadConfig, - EndpointConfig, Model, ModelPackage, ) @@ -36,6 +35,8 @@ from sagemaker.serve.model_builder import ModelBuilder from sagemaker.train.configs import Compute +from cleanup_helpers import cleanup_by_name, delete_quietly + logger = logging.getLogger(__name__) # Qwen2-1.5B-Instruct on the JumpStart DJL-LMI-22 0.36 container natively @@ -78,16 +79,18 @@ def test_benchmark_workflow_end_to_end(): role = get_execution_role(sagemaker_session=Session()) bench_job_name = f"air-bench-job-{unique_id}" workload_config_name = f"air-bench-wl-{unique_id}" - core_model = None - core_endpoint = None + # Names are generated up front so cleanup can run by name even if build() + # or deploy() raises or is killed before returning a resource handle. + model_name = f"air-bench-model-{unique_id}" + endpoint_name = f"air-bench-ep-{unique_id}" try: model_builder = _build_jumpstart_model_builder(role_arn=role) - core_model = model_builder.build(model_name=f"air-bench-model-{unique_id}") + core_model = model_builder.build(model_name=model_name) logger.info(f"Model created: {core_model.model_name}") core_endpoint = model_builder.deploy( - endpoint_name=f"air-bench-ep-{unique_id}" + endpoint_name=endpoint_name ) logger.info(f"Endpoint InService: {core_endpoint.endpoint_name}") @@ -117,14 +120,17 @@ def test_benchmark_workflow_end_to_end(): logger.info(f"Parsed {len(result.metrics.all_metrics)} benchmark metrics") finally: - if core_endpoint and core_model: - cleanup_resources(core_model, core_endpoint) - _delete_quietly( - lambda: AIBenchmarkJob.get(ai_benchmark_job_name=bench_job_name), + # Best-effort cleanup by name; runs even if deploy() failed/hung so a + # Failed or half-created endpoint is still torn down. + cleanup_by_name(endpoint_name=endpoint_name, model_name=model_name) + delete_quietly( + lambda: AIBenchmarkJob.get(ai_benchmark_job_name=bench_job_name).delete(), f"AIBenchmarkJob {bench_job_name}", ) - _delete_quietly( - lambda: AIWorkloadConfig.get(ai_workload_config_name=workload_config_name), + delete_quietly( + lambda: AIWorkloadConfig.get( + ai_workload_config_name=workload_config_name + ).delete(), f"AIWorkloadConfig {workload_config_name}", ) @@ -139,16 +145,15 @@ def test_recommendation_workflow_end_to_end(): rec_job_name = f"air-rec-job-{unique_id}" workload_config_name = f"air-rec-wl-{unique_id}" rec_endpoint_config_name = f"air-rec-cfg-{unique_id}" + rec_endpoint_name = f"air-rec-ep-{unique_id}" rec_model_name = f"air-rec-model-{unique_id}" + source_model_name = f"air-rec-source-{unique_id}" - source_model = None - rec_model = None - rec_endpoint = None rec_model_package_arn = None try: model_builder = _build_jumpstart_model_builder(role_arn=role) - source_model = model_builder.build(model_name=f"air-rec-source-{unique_id}") + model_builder.build(model_name=source_model_name) recommendation_job = model_builder.generate_deployment_recommendations( workload=_build_synthetic_workload(), @@ -192,7 +197,7 @@ def test_recommendation_workflow_end_to_end(): # The recommendation's ModelPackage is unapproved; opt in to approving # it as part of deploy for this end-to-end test. rec_endpoint = model_builder.deploy( - endpoint_name=f"air-rec-ep-{unique_id}", + endpoint_name=rec_endpoint_name, model_name=rec_model_name, endpoint_config_name=rec_endpoint_config_name, role=role, @@ -203,51 +208,39 @@ def test_recommendation_workflow_end_to_end(): assert rec_endpoint.endpoint_status == "InService", ( f"Endpoint did not reach InService: {rec_endpoint.endpoint_status}" ) - rec_model = Model.get(model_name=rec_model_name) finally: - if rec_endpoint and rec_model: - cleanup_resources(rec_model, rec_endpoint, rec_endpoint_config_name) - if source_model: - source_model.delete() - _delete_quietly( - lambda: AIRecommendationJob.get(ai_recommendation_job_name=rec_job_name), + # Best-effort cleanup by name; runs even if deploy() failed/hung so a + # Failed or half-created endpoint is still torn down. The + # recommendation path uses a distinct endpoint-config name, so it is + # passed explicitly rather than defaulting to the endpoint name. + cleanup_by_name( + endpoint_name=rec_endpoint_name, + endpoint_config_name=rec_endpoint_config_name, + model_name=rec_model_name, + ) + delete_quietly( + lambda: Model.get(model_name=source_model_name).delete(), + f"Model {source_model_name}", + ) + delete_quietly( + lambda: AIRecommendationJob.get( + ai_recommendation_job_name=rec_job_name + ).delete(), f"AIRecommendationJob {rec_job_name}", ) - _delete_quietly( - lambda: AIWorkloadConfig.get(ai_workload_config_name=workload_config_name), + delete_quietly( + lambda: AIWorkloadConfig.get( + ai_workload_config_name=workload_config_name + ).delete(), f"AIWorkloadConfig {workload_config_name}", ) # The recommendation job publishes (and this test approves) a # ModelPackage; delete it so repeated runs don't accumulate packages. if rec_model_package_arn: - _delete_quietly( - lambda: ModelPackage.get(model_package_name=rec_model_package_arn), + delete_quietly( + lambda: ModelPackage.get( + model_package_name=rec_model_package_arn + ).delete(), f"ModelPackage {rec_model_package_arn}", ) - - -def cleanup_resources(core_model, core_endpoint, endpoint_config_name=None): - """Delete the model, endpoint, and endpoint config in reverse order of creation. - - ``endpoint_config_name`` defaults to ``core_endpoint.endpoint_name`` for the - benchmark-test path where ``ModelBuilder.deploy()`` uses a single name for - both. Must be passed explicitly for the recommendation-test path, where - ``ModelBuilder.deploy(recommendation_index=N)`` generates a distinct - endpoint-config name. - """ - config_name = endpoint_config_name or core_endpoint.endpoint_name - _delete_quietly(lambda: core_model, f"Model {core_model.model_name}") - _delete_quietly(lambda: core_endpoint, f"Endpoint {core_endpoint.endpoint_name}") - _delete_quietly( - lambda: EndpointConfig.get(endpoint_config_name=config_name), - f"EndpointConfig {config_name}", - ) - - -def _delete_quietly(resource_factory, label): - """Best-effort delete; log and continue on any failure.""" - try: - resource_factory().delete() - except Exception as exc: - logger.warning("Failed to delete %s: %s", label, exc) diff --git a/sagemaker-serve/tests/integ/test_huggingface_integration.py b/sagemaker-serve/tests/integ/test_huggingface_integration.py index 250d2884ec..1a82db3327 100644 --- a/sagemaker-serve/tests/integ/test_huggingface_integration.py +++ b/sagemaker-serve/tests/integ/test_huggingface_integration.py @@ -20,7 +20,8 @@ from sagemaker.serve.model_builder import ModelBuilder from sagemaker.serve.utils.types import ModelServer from sagemaker.core.training.configs import Compute -from sagemaker.core.resources import EndpointConfig + +from cleanup_helpers import cleanup_by_name logger = logging.getLogger(__name__) @@ -34,30 +35,33 @@ def test_huggingface_build_deploy_invoke_cleanup(): """Integration test for HuggingFace model build, deploy, invoke, and cleanup workflow""" logger.info("Starting HuggingFace integration test...") - - core_model = None - core_endpoint = None - + + # Names are generated up front so cleanup can run by name even if build() + # or deploy() raises or is killed before returning a resource handle. + unique_id = str(uuid.uuid4())[:8] + model_name = f"{MODEL_NAME_PREFIX}-{unique_id}" + endpoint_name = f"{ENDPOINT_NAME_PREFIX}-{unique_id}" + try: # Build and deploy logger.info("Building and deploying HuggingFace model...") - core_model, core_endpoint = build_and_deploy() - + core_endpoint = build_and_deploy(model_name, endpoint_name) + # Make prediction logger.info("Making prediction...") make_prediction(core_endpoint) - + # Test passed successfully logger.info("HuggingFace integration test completed successfully") - + except Exception as e: logger.error(f"HuggingFace integration test failed: {str(e)}") raise finally: - # Cleanup resources - if core_model and core_endpoint: - logger.info("Cleaning up resources...") - cleanup_resources(core_model, core_endpoint) + # Best-effort cleanup by name; runs even if deploy() failed/hung so a + # Failed or half-created endpoint is still torn down. + logger.info("Cleaning up resources...") + cleanup_by_name(endpoint_name=endpoint_name, model_name=model_name) def create_schema_builder(): @@ -70,12 +74,11 @@ def create_schema_builder(): return SchemaBuilder(sample_input, sample_output) -def build_and_deploy(): +def build_and_deploy(model_name, endpoint_name): """Build and deploy HuggingFace model - preserving exact logic from manual test""" hf_model_id = MODEL_ID - + schema_builder = create_schema_builder() - unique_id = str(uuid.uuid4())[:8] compute = Compute( instance_type="ml.g5.xlarge", @@ -95,13 +98,13 @@ def build_and_deploy(): ) # Build and deploy your model. Returns SageMaker Core Model and Endpoint objects - core_model = model_builder.build(model_name=f"{MODEL_NAME_PREFIX}-{unique_id}") + core_model = model_builder.build(model_name=model_name) logger.info(f"Model Successfully Created: {core_model.model_name}") - core_endpoint = model_builder.deploy(endpoint_name=f"{ENDPOINT_NAME_PREFIX}-{unique_id}") + core_endpoint = model_builder.deploy(endpoint_name=endpoint_name) logger.info(f"Endpoint Successfully Created: {core_endpoint.endpoint_name}") - - return core_model, core_endpoint + + return core_endpoint def make_prediction(core_endpoint): @@ -120,14 +123,3 @@ def make_prediction(core_endpoint): # Decode the output of the invocation and print the result prediction = json.loads(result.body.read().decode('utf-8')) logger.info(f"Result of invoking endpoint: {prediction}") - - -def cleanup_resources(core_model, core_endpoint): - """Fully clean up model and endpoint creation - preserving exact logic from manual test""" - core_endpoint_config = EndpointConfig.get(endpoint_config_name=core_endpoint.endpoint_name) - - core_model.delete() - core_endpoint.delete() - core_endpoint_config.delete() - - logger.info("Model and Endpoint Successfully Deleted!") \ No newline at end of file diff --git a/sagemaker-serve/tests/integ/test_jumpstart_integration.py b/sagemaker-serve/tests/integ/test_jumpstart_integration.py index 40532d2a37..c35804d196 100644 --- a/sagemaker-serve/tests/integ/test_jumpstart_integration.py +++ b/sagemaker-serve/tests/integ/test_jumpstart_integration.py @@ -19,9 +19,10 @@ from sagemaker.serve.model_builder import ModelBuilder from sagemaker.core.jumpstart.configs import JumpStartConfig -from sagemaker.core.resources import EndpointConfig from sagemaker.train.configs import Compute +from cleanup_helpers import cleanup_by_name + logger = logging.getLogger(__name__) # Configuration - easily customizable @@ -36,49 +37,50 @@ def test_jumpstart_build_deploy_invoke_cleanup(): """Integration test for JumpStart model build, deploy, invoke, and cleanup workflow""" logger.info("Starting JumpStart integration test...") - - core_model = None - core_endpoint = None - core_endpoint_config = None - + + # Names are generated up front so cleanup can run by name even if build() + # or deploy() raises or is killed before returning a resource handle. + unique_id = str(uuid.uuid4())[:8] + model_name = f"{MODEL_NAME_PREFIX}-{unique_id}" + endpoint_name = f"{ENDPOINT_NAME_PREFIX}-{unique_id}" + try: # Build and deploy logger.info("Building and deploying JumpStart model...") - core_model, core_endpoint = build_and_deploy() - + core_endpoint = build_and_deploy(model_name, endpoint_name) + # Make prediction logger.info("Making prediction...") make_prediction(core_endpoint) - + # Test passed successfully logger.info("JumpStart integration test completed successfully") - + except Exception as e: logger.error(f"JumpStart integration test failed: {str(e)}") raise finally: - # Cleanup resources - if core_model and core_endpoint: - logger.info("Cleaning up resources...") - cleanup_resources(core_model, core_endpoint) + # Best-effort cleanup by name; runs even if deploy() failed/hung so a + # Failed or half-created endpoint is still torn down. + logger.info("Cleaning up resources...") + cleanup_by_name(endpoint_name=endpoint_name, model_name=model_name) -def build_and_deploy(): +def build_and_deploy(model_name, endpoint_name): """Build and deploy JumpStart model - preserving exact logic from manual test""" # Initialize model_builder object with JumpStart configuration compute = Compute(instance_type="ml.g5.2xlarge") jumpstart_config = JumpStartConfig(model_id=MODEL_ID) model_builder = ModelBuilder.from_jumpstart_config(jumpstart_config=jumpstart_config, compute=compute) - unique_id = str(uuid.uuid4())[:8] - + # Build and deploy your model. Returns SageMaker Core Model and Endpoint objects - core_model = model_builder.build(model_name=f"{MODEL_NAME_PREFIX}-{unique_id}") + core_model = model_builder.build(model_name=model_name) logger.info(f"Model Successfully Created: {core_model.model_name}") - core_endpoint = model_builder.deploy(endpoint_name=f"{ENDPOINT_NAME_PREFIX}-{unique_id}") + core_endpoint = model_builder.deploy(endpoint_name=endpoint_name) logger.info(f"Endpoint Successfully Created: {core_endpoint.endpoint_name}") - - return core_model, core_endpoint + + return core_endpoint def make_prediction(core_endpoint): @@ -93,14 +95,3 @@ def make_prediction(core_endpoint): # Decode the output of the invocation and print the result prediction = json.loads(result.body.read().decode('utf-8')) logger.info(f"Result of invoking endpoint: {prediction}") - - -def cleanup_resources(core_model, core_endpoint): - """Fully clean up model and endpoint creation - preserving exact logic from manual test""" - core_endpoint_config = EndpointConfig.get(endpoint_config_name=core_endpoint.endpoint_name) - - core_model.delete() - core_endpoint.delete() - core_endpoint_config.delete() - - logger.info("Model and Endpoint Successfully Deleted!") \ No newline at end of file diff --git a/sagemaker-serve/tests/integ/test_optimize_integration.py b/sagemaker-serve/tests/integ/test_optimize_integration.py index 9ee222ea9e..8c15c8950a 100644 --- a/sagemaker-serve/tests/integ/test_optimize_integration.py +++ b/sagemaker-serve/tests/integ/test_optimize_integration.py @@ -20,10 +20,11 @@ import boto3 from sagemaker.serve.model_builder import ModelBuilder -from sagemaker.core.resources import EndpointConfig from sagemaker.core.helper.session_helper import Session from sagemaker.core import image_uris +from cleanup_helpers import cleanup_by_name + logger = logging.getLogger(__name__) # Configuration - easily customizable @@ -42,30 +43,33 @@ def test_optimize_build_deploy_invoke_cleanup(): """Integration test for Optimize workflow""" logger.info("Starting Optimize integration test...") - - optimized_model = None - core_endpoint = None - + + # Names are generated up front so cleanup can run by name even if optimize() + # or deploy() raises or is killed before returning a resource handle. + unique_id = str(uuid.uuid4())[:8] + model_name = f"{MODEL_NAME_PREFIX}-{unique_id}" + endpoint_name = f"{ENDPOINT_NAME_PREFIX}-{unique_id}" + try: # Build and deploy logger.info("Optimizing and deploying model...") - optimized_model, core_endpoint = build_and_deploy() - + core_endpoint = build_and_deploy(model_name, endpoint_name, unique_id) + # Make prediction logger.info("Making prediction...") make_prediction(core_endpoint) - + # Test passed successfully logger.info("Optimize integration test completed successfully") - + except Exception as e: logger.error(f"Optimize integration test failed: {str(e)}") raise finally: - # Cleanup resources - if optimized_model and core_endpoint: - logger.info("Cleaning up resources...") - cleanup_resources(optimized_model, core_endpoint) + # Best-effort cleanup by name; runs even if deploy() failed/hung so a + # Failed or half-created endpoint is still torn down. + logger.info("Cleaning up resources...") + cleanup_by_name(endpoint_name=endpoint_name, model_name=model_name) def create_schema_builder(): @@ -78,19 +82,18 @@ def create_schema_builder(): return SchemaBuilder(sample_input, sample_output) -def build_and_deploy(): +def build_and_deploy(model_name, endpoint_name, unique_id): """Optimize and deploy JumpStart model - exact logic from optimize test.""" schema_builder = create_schema_builder() boto_session = boto3.Session(region_name=AWS_REGION) sagemaker_session = Session(boto_session=boto_session) - unique_id = str(uuid.uuid4())[:8] - + model_builder = ModelBuilder( model=MODEL_ID, schema_builder=schema_builder, sagemaker_session=sagemaker_session, ) - + # Optimize the model logger.info("Optimizing JumpStart model...") default_bucket = sagemaker_session.default_bucket() @@ -100,9 +103,9 @@ def build_and_deploy(): version=DJL_LMI_VERSION, ) logger.info(f"Resolved DJL LMI image URI: {djl_lmi_image_uri}") - + optimized_model = model_builder.optimize( - model_name=f"{MODEL_NAME_PREFIX}-{unique_id}", + model_name=model_name, instance_type="ml.g5.2xlarge", output_path=f"s3://{default_bucket}/optimize-output/jumpstart-{unique_id}/", quantization_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}}, @@ -112,17 +115,17 @@ def build_and_deploy(): region=AWS_REGION ) logger.info(f"Model Successfully Optimized: {optimized_model.model_name}") - + # Deploy the optimized model logger.info("Deploying optimized model to endpoint...") core_endpoint = model_builder.deploy( - endpoint_name=f"{ENDPOINT_NAME_PREFIX}-{unique_id}", + endpoint_name=endpoint_name, initial_instance_count=1, instance_type="ml.g5.2xlarge" ) logger.info(f"Endpoint Successfully Created: {core_endpoint.endpoint_name}") - - return optimized_model, core_endpoint + + return core_endpoint def make_prediction(core_endpoint): @@ -140,14 +143,3 @@ def make_prediction(core_endpoint): response_body = result.body.read().decode('utf-8') prediction = json.loads(response_body) logger.info(f"Result of invoking optimized endpoint: {prediction}") - - -def cleanup_resources(optimized_model, core_endpoint): - """Clean up optimized model and endpoint - preserving exact logic from manual test""" - core_endpoint_config = EndpointConfig.get(endpoint_config_name=core_endpoint.endpoint_name) - - optimized_model.delete() - core_endpoint.delete() - core_endpoint_config.delete() - - logger.info("Optimized model and endpoint successfully deleted!") \ No newline at end of file diff --git a/sagemaker-serve/tests/integ/test_tei_integration.py b/sagemaker-serve/tests/integ/test_tei_integration.py index 0847d20fbe..a4b0530074 100644 --- a/sagemaker-serve/tests/integ/test_tei_integration.py +++ b/sagemaker-serve/tests/integ/test_tei_integration.py @@ -16,13 +16,12 @@ import uuid import pytest import logging -import boto3 from sagemaker.serve.model_builder import ModelBuilder from sagemaker.serve.utils.types import ModelServer from sagemaker.train.configs import Compute -from sagemaker.core.resources import EndpointConfig -from sagemaker.core.helper.session_helper import Session + +from cleanup_helpers import cleanup_by_name logger = logging.getLogger(__name__) @@ -36,30 +35,33 @@ def test_tei_build_deploy_invoke_cleanup(): """Integration test for TEI model build, deploy, invoke, and cleanup workflow""" logger.info("Starting TEI integration test...") - - core_model = None - core_endpoint = None - + + # Names are generated up front so cleanup can run by name even if build() + # or deploy() raises or is killed before returning a resource handle. + unique_id = str(uuid.uuid4())[:8] + model_name = f"{MODEL_NAME_PREFIX}-{unique_id}" + endpoint_name = f"{ENDPOINT_NAME_PREFIX}-{unique_id}" + try: # Build and deploy logger.info("Building and deploying TEI model...") - core_model, core_endpoint = build_and_deploy() - + core_endpoint = build_and_deploy(model_name, endpoint_name) + # Make prediction logger.info("Making prediction...") make_prediction(core_endpoint) - + # Test passed successfully logger.info("TEI integration test completed successfully") - + except Exception as e: logger.error(f"TEI integration test failed: {str(e)}") raise finally: - # Cleanup resources - if core_model and core_endpoint: - logger.info("Cleaning up resources...") - cleanup_resources(core_model, core_endpoint) + # Best-effort cleanup by name; runs even if deploy() failed/hung so a + # Failed or half-created endpoint is still torn down. + logger.info("Cleaning up resources...") + cleanup_by_name(endpoint_name=endpoint_name, model_name=model_name) def create_schema_builder(): @@ -72,37 +74,36 @@ def create_schema_builder(): return SchemaBuilder(sample_input, sample_output) -def build_and_deploy(): +def build_and_deploy(model_name, endpoint_name): """Build and deploy TEI model - exact logic from backup file.""" # Use HuggingFace model string for TEI (text embeddings) hf_model_id = MODEL_ID - + schema_builder = create_schema_builder() - unique_id = str(uuid.uuid4())[:8] compute = Compute( instance_type="ml.g5.xlarge", instance_count=1, ) - + model_builder = ModelBuilder( model=hf_model_id, # Use HuggingFace model string model_server=ModelServer.TEI, schema_builder=schema_builder, compute=compute, ) - + # Build and deploy your model. Returns SageMaker Core Model and Endpoint objects - core_model = model_builder.build(model_name=f"{MODEL_NAME_PREFIX}-{unique_id}") + core_model = model_builder.build(model_name=model_name) logger.info(f"Model Successfully Created: {core_model.model_name}") core_endpoint = model_builder.deploy( - endpoint_name=f"{ENDPOINT_NAME_PREFIX}-{unique_id}", + endpoint_name=endpoint_name, initial_instance_count=1, ) logger.info(f"Endpoint Successfully Created: {core_endpoint.endpoint_name}") - - return core_model, core_endpoint + + return core_endpoint def make_prediction(core_endpoint): @@ -117,14 +118,3 @@ def make_prediction(core_endpoint): # Decode the output of the invocation and print the result prediction = json.loads(result.body.read().decode('utf-8')) logger.info(f"Result of invoking endpoint: {prediction}") - - -def cleanup_resources(core_model, core_endpoint): - """Fully clean up model and endpoint creation - preserving exact logic from manual test""" - core_endpoint_config = EndpointConfig.get(endpoint_config_name=core_endpoint.endpoint_name) - - core_model.delete() - core_endpoint.delete() - core_endpoint_config.delete() - - logger.info("Model and Endpoint Successfully Deleted!") \ No newline at end of file diff --git a/sagemaker-serve/tests/integ/test_tgi_integration.py b/sagemaker-serve/tests/integ/test_tgi_integration.py index 63fac89be3..874c775625 100644 --- a/sagemaker-serve/tests/integ/test_tgi_integration.py +++ b/sagemaker-serve/tests/integ/test_tgi_integration.py @@ -16,13 +16,12 @@ import uuid import pytest import logging -import boto3 from sagemaker.serve.model_builder import ModelBuilder from sagemaker.serve.utils.types import ModelServer from sagemaker.train.configs import Compute -from sagemaker.core.resources import EndpointConfig -from sagemaker.core.helper.session_helper import Session + +from cleanup_helpers import cleanup_by_name logger = logging.getLogger(__name__) @@ -36,30 +35,33 @@ def test_tgi_build_deploy_invoke_cleanup(): """Integration test for TGI model build, deploy, invoke, and cleanup workflow""" logger.info("Starting TGI integration test...") - - core_model = None - core_endpoint = None - + + # Names are generated up front so cleanup can run by name even if build() + # or deploy() raises or is killed before returning a resource handle. + unique_id = str(uuid.uuid4())[:8] + model_name = f"{MODEL_NAME_PREFIX}-{unique_id}" + endpoint_name = f"{ENDPOINT_NAME_PREFIX}-{unique_id}" + try: # Build and deploy logger.info("Building and deploying TGI model...") - core_model, core_endpoint = build_and_deploy() - + core_endpoint = build_and_deploy(model_name, endpoint_name) + # Make prediction logger.info("Making prediction...") make_prediction(core_endpoint) - + # Test passed successfully logger.info("TGI integration test completed successfully") - + except Exception as e: logger.error(f"TGI integration test failed: {str(e)}") raise finally: - # Cleanup resources - if core_model and core_endpoint: - logger.info("Cleaning up resources...") - cleanup_resources(core_model, core_endpoint) + # Best-effort cleanup by name; runs even if deploy() failed/hung so a + # Failed or half-created endpoint is still torn down. + logger.info("Cleaning up resources...") + cleanup_by_name(endpoint_name=endpoint_name, model_name=model_name) def create_schema_builder(): @@ -72,13 +74,12 @@ def create_schema_builder(): return SchemaBuilder(sample_input, sample_output) -def build_and_deploy(): +def build_and_deploy(model_name, endpoint_name): """Build and deploy TGI model - exact logic from backup file.""" # Use HuggingFace model string for TGI (no local artifacts needed) hf_model_id = MODEL_ID - + schema_builder = create_schema_builder() - unique_id = str(uuid.uuid4())[:8] compute = Compute( instance_type="ml.g5.xlarge", @@ -101,16 +102,16 @@ def build_and_deploy(): ) # Build and deploy your model. Returns SageMaker Core Model and Endpoint objects - core_model = model_builder.build(model_name=f"{MODEL_NAME_PREFIX}-{unique_id}") + core_model = model_builder.build(model_name=model_name) logger.info(f"Model Successfully Created: {core_model.model_name}") core_endpoint = model_builder.deploy( - endpoint_name=f"{ENDPOINT_NAME_PREFIX}-{unique_id}", + endpoint_name=endpoint_name, initial_instance_count=1, ) logger.info(f"Endpoint Successfully Created: {core_endpoint.endpoint_name}") - - return core_model, core_endpoint + + return core_endpoint def make_prediction(core_endpoint): @@ -128,14 +129,3 @@ def make_prediction(core_endpoint): # Decode the output of the invocation and print the result prediction = json.loads(result.body.read().decode('utf-8')) logger.info(f"Result of invoking endpoint: {prediction}") - - -def cleanup_resources(core_model, core_endpoint): - """Fully clean up model and endpoint creation - preserving exact logic from manual test""" - core_endpoint_config = EndpointConfig.get(endpoint_config_name=core_endpoint.endpoint_name) - - core_model.delete() - core_endpoint.delete() - core_endpoint_config.delete() - - logger.info("Model and Endpoint Successfully Deleted!") \ No newline at end of file diff --git a/sagemaker-serve/tests/integ/test_train_inference_e2e_integration.py b/sagemaker-serve/tests/integ/test_train_inference_e2e_integration.py index 892163b3d6..6959cccfb5 100644 --- a/sagemaker-serve/tests/integ/test_train_inference_e2e_integration.py +++ b/sagemaker-serve/tests/integ/test_train_inference_e2e_integration.py @@ -24,7 +24,8 @@ from sagemaker.serve.utils.types import ModelServer from sagemaker.train.model_trainer import ModelTrainer from sagemaker.train.configs import SourceCode -from sagemaker.core.resources import EndpointConfig + +from cleanup_helpers import cleanup_by_name logger = logging.getLogger(__name__) @@ -38,35 +39,37 @@ def test_train_inference_e2e_build_deploy_invoke_cleanup(): """Integration test for Train-Inference E2E workflow""" logger.info("Starting Train-Inference E2E integration test...") - - model_trainer = None - core_model = None - core_endpoint = None - + + # Names are generated up front so cleanup can run by name even if train(), + # build() or deploy() raises or is killed before returning a resource handle. + unique_id = str(uuid.uuid4())[:8] + model_name = f"{MODEL_NAME_PREFIX}-{unique_id}" + endpoint_name = f"{ENDPOINT_NAME_PREFIX}-{unique_id}" + try: # Step 1: Train model logger.info("Training model...") - model_trainer, unique_id = train_model() - + model_trainer = train_model(unique_id) + # Step 2: Build and deploy logger.info("Building and deploying model...") - core_model, core_endpoint = build_and_deploy(model_trainer, unique_id) - + core_endpoint = build_and_deploy(model_trainer, model_name, endpoint_name) + # Step 3: Test inference logger.info("Making prediction...") make_prediction(core_endpoint) - + # Test passed successfully logger.info("Train-Inference E2E integration test completed successfully") - + except Exception as e: logger.error(f"Train-Inference E2E integration test failed: {str(e)}") raise finally: - # Cleanup resources - if core_model and core_endpoint: - logger.info("Cleaning up resources...") - cleanup_resources(core_model, core_endpoint) + # Best-effort cleanup by name; runs even if deploy() failed/hung so a + # Failed or half-created endpoint is still torn down. + logger.info("Cleaning up resources...") + cleanup_by_name(endpoint_name=endpoint_name, model_name=model_name) def create_pytorch_training_code(): @@ -137,18 +140,17 @@ def create_schema_builder(): return SchemaBuilder(sample_input, sample_output) -def train_model(): +def train_model(unique_id): """Train model using ModelTrainer.""" from sagemaker.core import image_uris from sagemaker.core.helper.session_helper import Session - + # Get the current region from a session session = Session() region = session.boto_region_name - + training_code_dir = create_pytorch_training_code() - unique_id = str(uuid.uuid4())[:8] - + # Get training image for the current region training_image = image_uris.retrieve( framework="pytorch", @@ -171,11 +173,11 @@ def train_model(): model_trainer.train() logger.info("Model Training Completed!") - - return model_trainer, unique_id + + return model_trainer -def build_and_deploy(model_trainer, unique_id): +def build_and_deploy(model_trainer, model_name, endpoint_name): """Build and deploy model using ModelBuilder.""" from sagemaker.serve.spec.inference_spec import InferenceSpec from sagemaker.core import image_uris @@ -215,16 +217,16 @@ def invoke(self, input_object, model): dependencies={"auto": False}, ) - core_model = model_builder.build(model_name=f"{MODEL_NAME_PREFIX}-{unique_id}") + core_model = model_builder.build(model_name=model_name) logger.info(f"Model Successfully Created: {core_model.model_name}") - + core_endpoint = model_builder.deploy( - endpoint_name=f"{ENDPOINT_NAME_PREFIX}-{unique_id}", + endpoint_name=endpoint_name, initial_instance_count=1 ) logger.info(f"Endpoint Successfully Created: {core_endpoint.endpoint_name}") - - return core_model, core_endpoint + + return core_endpoint def make_prediction(core_endpoint): @@ -238,14 +240,3 @@ def make_prediction(core_endpoint): prediction = json.loads(result.body.read().decode('utf-8')) logger.info(f"Result of invoking endpoint: {prediction}") - - -def cleanup_resources(core_model, core_endpoint): - """Fully clean up model and endpoint creation - preserving exact logic from manual test""" - core_endpoint_config = EndpointConfig.get(endpoint_config_name=core_endpoint.endpoint_name) - - core_model.delete() - core_endpoint.delete() - core_endpoint_config.delete() - - logger.info("Model and Endpoint Successfully Deleted!") \ No newline at end of file diff --git a/sagemaker-serve/tests/integ/test_triton_integration.py b/sagemaker-serve/tests/integ/test_triton_integration.py index 5b2c3b2e3a..ea970fdc80 100644 --- a/sagemaker-serve/tests/integ/test_triton_integration.py +++ b/sagemaker-serve/tests/integ/test_triton_integration.py @@ -21,12 +21,13 @@ from sagemaker.serve.model_builder import ModelBuilder from sagemaker.serve.utils.types import ModelServer -from sagemaker.core.resources import EndpointConfig # PyTorch Imports import torch import torch.nn as nn +from cleanup_helpers import cleanup_by_name + logger = logging.getLogger(__name__) # Configuration - easily customizable @@ -48,30 +49,33 @@ def forward(self, x): def test_triton_build_deploy_invoke_cleanup(): """Integration test for Triton model build, deploy, invoke, and cleanup workflow""" logger.info("Starting Triton integration test...") - - core_model = None - core_endpoint = None - + + # Names are generated up front so cleanup can run by name even if build() + # or deploy() raises or is killed before returning a resource handle. + unique_id = str(uuid.uuid4())[:8] + model_name = f"{MODEL_NAME_PREFIX}-{unique_id}" + endpoint_name = f"{ENDPOINT_NAME_PREFIX}-{unique_id}" + try: # Build and deploy logger.info("Building and deploying Triton model...") - core_model, core_endpoint = build_and_deploy() - + core_endpoint = build_and_deploy(model_name, endpoint_name) + # Make prediction logger.info("Making prediction...") make_prediction(core_endpoint) - + # Test passed successfully logger.info("Triton integration test completed successfully") - + except Exception as e: logger.error(f"Triton integration test failed: {str(e)}") raise finally: - # Cleanup resources - if core_model and core_endpoint: - logger.info("Cleaning up resources...") - cleanup_resources(core_model, core_endpoint) + # Best-effort cleanup by name; runs even if deploy() failed/hung so a + # Failed or half-created endpoint is still torn down. + logger.info("Cleaning up resources...") + cleanup_by_name(endpoint_name=endpoint_name, model_name=model_name) def create_schema_builder(): @@ -86,31 +90,30 @@ def create_schema_builder(): return SchemaBuilder(sample_input, sample_output) -def build_and_deploy(): +def build_and_deploy(model_name, endpoint_name): """Build and deploy Triton model - preserving exact logic from manual test""" # Create model and save state dictionary (assume model is pre-trained!) pytorch_model = SimpleModel() model_path = tempfile.mkdtemp() torch.save(pytorch_model.state_dict(), os.path.join(model_path, "model.pth")) - + schema_builder = create_schema_builder() - + model_builder = ModelBuilder( model=pytorch_model, model_path=model_path, model_server=ModelServer.TRITON, schema_builder=schema_builder ) - - unique_id = str(uuid.uuid4())[:8] + # Build and deploy your model. Returns SageMaker Core Model and Endpoint objects - core_model = model_builder.build(model_name=f"{MODEL_NAME_PREFIX}-{unique_id}") + core_model = model_builder.build(model_name=model_name) logger.info(f"Model Successfully Created: {core_model.model_name}") - core_endpoint = model_builder.deploy(endpoint_name=f"{ENDPOINT_NAME_PREFIX}-{unique_id}") + core_endpoint = model_builder.deploy(endpoint_name=endpoint_name) logger.info(f"Endpoint Successfully Created: {core_endpoint.endpoint_name}") - - return core_model, core_endpoint + + return core_endpoint def make_prediction(core_endpoint): @@ -135,14 +138,3 @@ def make_prediction(core_endpoint): # Decode the output of the invocation and print the result prediction = json.loads(result.body.read().decode('utf-8')) logger.info(f"Result of invoking endpoint: {prediction}") - - -def cleanup_resources(core_model, core_endpoint): - """Fully clean up model and endpoint creation - preserving exact logic from manual test""" - core_endpoint_config = EndpointConfig.get(endpoint_config_name=core_endpoint.endpoint_name) - - core_model.delete() - core_endpoint.delete() - core_endpoint_config.delete() - - logger.info("Model and Endpoint Successfully Deleted!") \ No newline at end of file