From ceaa05485573773fdc958f1a3e7015c3f9651005 Mon Sep 17 00:00:00 2001 From: papriwal Date: Fri, 11 Sep 2026 13:34:52 -0700 Subject: [PATCH] fix: restore model customization reuse state Restore deployment metadata when ModelBuilder reuses a package-backed Model but cannot reuse its endpoint. Preserve explicit resource requirements and recover recipe compute or LoRA adapter state only when creating the replacement endpoint. Make scheduled reuse coverage deterministic and guard against duplicate Model and endpoint writes. --- .../src/sagemaker/serve/model_builder.py | 199 +++++++--- .../test_model_customization_deployment.py | 359 ++++++++++-------- .../tests/unit/test_model_builder.py | 347 +++++++++++++++++ 3 files changed, 687 insertions(+), 218 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index cff687849e..651587c1e4 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -457,6 +457,8 @@ def __post_init__(self) -> None: """Initialize ModelBuilder after instantiation.""" import warnings + self._built_model_was_reused = False + if self.sagemaker_session is None: self.sagemaker_session = self._create_session_with_region() @@ -3227,6 +3229,73 @@ def fetch_endpoint_names_for_base_model(self) -> Set[str]: return endpoint_names + def _resolve_lora_adapter_s3_uri(self, model_package: ModelPackage) -> str: + """Resolve the LoRA adapter URI for a supported model source.""" + if isinstance(self.model, TrainingJob): + model_artifacts = getattr(self.model, "model_artifacts", None) + s3_uri = getattr(model_artifacts, "s3_model_artifacts", None) + suffix = "/checkpoints/hf/" + elif isinstance(self.model, ModelTrainer): + training_job = getattr(self.model, "_latest_training_job", None) + model_artifacts = getattr(training_job, "model_artifacts", None) + s3_uri = getattr(model_artifacts, "s3_model_artifacts", None) + suffix = "/checkpoints/hf/" + elif isinstance(self.model, (AgentRFTJob, ModelPackage)): + try: + s3_uri = ( + model_package.inference_specification.containers[ + 0 + ].model_data_source.s3_data_source.s3_uri + ) + except (AttributeError, IndexError, TypeError): + s3_uri = None + suffix = ( + "/model/checkpoints/hf/" + if isinstance(self.model, ModelPackage) + else "/checkpoints/hf/" + ) + else: + raise ValueError( + "Cannot resolve a LoRA adapter artifact URI from model source type " + f"'{type(self.model).__name__}'. Use a TrainingJob, ModelTrainer, " + "AgentRFTJob, or ModelPackage source." + ) + + if not isinstance(s3_uri, str) or not s3_uri: + raise ValueError( + "Cannot resolve the LoRA adapter artifact URI from the model source. " + "Ensure the training job or model package contains model artifacts." + ) + + if isinstance(self.model, (AgentRFTJob, ModelPackage)): + return s3_uri.rstrip("/") + suffix + return f"{s3_uri}{suffix}" + + def _prepare_reused_model_customization_deployment_state( + self, + model_package: Optional[ModelPackage], + peft_type: Optional[str], + inference_config: Optional[ResourceRequirements], + ) -> None: + """Restore deployment state skipped when a customization Model was reused.""" + if not getattr(self, "_built_model_was_reused", False): + return + if model_package is None or self._is_nova_model(): + return + + if inference_config is None and getattr( + self, "_cached_compute_requirements", None + ) is None: + self._fetch_and_cache_recipe_config() + if getattr(self, "_cached_compute_requirements", None) is None: + raise ValueError( + "Cannot resolve compute requirements for the reused model. " + "Provide ResourceRequirements explicitly." + ) + + if peft_type == "LORA" and not getattr(self, "_adapter_s3_uri", None): + self._adapter_s3_uri = self._resolve_lora_adapter_s3_uri(model_package) + def _build_single_modelbuilder( self, mode: Optional[Mode] = None, @@ -3398,25 +3467,7 @@ def _build_single_modelbuilder( }, ) # Store adapter path for use during deploy - if isinstance(self.model, TrainingJob): - self._adapter_s3_uri = ( - f"{self.model.model_artifacts.s3_model_artifacts}/checkpoints/hf/" - ) - elif isinstance(self.model, ModelTrainer): - self._adapter_s3_uri = ( - f"{self.model._latest_training_job.model_artifacts.s3_model_artifacts}" - "/checkpoints/hf/" - ) - elif isinstance(self.model, AgentRFTJob): - s3_uri = model_package.inference_specification.containers[ - 0 - ].model_data_source.s3_data_source.s3_uri - self._adapter_s3_uri = s3_uri.rstrip("/") + "/checkpoints/hf/" - elif isinstance(self.model, ModelPackage): - s3_uri = model_package.inference_specification.containers[ - 0 - ].model_data_source.s3_data_source.s3_uri - self._adapter_s3_uri = s3_uri.rstrip("/") + "/model/checkpoints/hf/" + self._adapter_s3_uri = self._resolve_lora_adapter_s3_uri(model_package) else: # Non-LORA: Model points at training output from sagemaker.serve.utils.model_package_utils import get_s3_uri_from_inference_spec @@ -4240,6 +4291,10 @@ def _reset_build_state(self): # Core build state self.built_model = None self.secret_key = "" + self._built_model_was_reused = False + for attr in ["_cached_compute_requirements", "_adapter_s3_uri"]: + if hasattr(self, attr): + delattr(self, attr) # JumpStart preparation flags for attr in ["prepared_for_djl", "prepared_for_tgi", "prepared_for_mms"]: @@ -4328,11 +4383,11 @@ def build( configuration chain. (Default: None). region (str, optional): The AWS region for deployment. If specified and different from the current region, a new session will be created. (Default: None). - reuse_resources (bool, optional): If True, checks for an existing endpoint built - from the same model source (with matching deployment configuration) before - creating anything. On a match, build() creates no new resources and sets - ``built_model`` to the existing Model backing that endpoint; the subsequent - deploy() returns the existing endpoint. (Default: False). + reuse_resources (bool, optional): If True, checks for an existing Model built + from the same model source before creating one. On a match, build() creates + no new Model and sets ``built_model`` to the existing Model; on a miss, it + creates a new Model. Endpoint reuse is handled independently by passing + ``reuse_resources=True`` to deploy(). (Default: False). Returns: Union[Model, ModelBuilder, None]: A ``sagemaker.core.resources.Model`` resource @@ -4345,6 +4400,8 @@ def build( >>> endpoint = model_builder.deploy() # Creates Endpoint resource >>> result = endpoint.invoke(data=input_data) """ + self._built_model_was_reused = False + if hasattr(self, "built_model") and self.built_model is not None: logger.warning( "ModelBuilder.build() has already been called. " @@ -4409,6 +4466,7 @@ def build( reused_model.model_name, ) self.built_model = reused_model + self._built_model_was_reused = True return self.built_model deployables = {} @@ -6279,9 +6337,62 @@ def _deploy_model_customization( endpoint.wait_for_status("InService") return endpoint + # Package-backed Nova models with explicit requirements continue through + # the single-IC path without generic LoRA or recipe preparation. + peft_type = ( + self._fetch_peft() if model_package is not None and not is_nova else None + ) + base_model_recipe_name = None + if peft_type == "LORA": + container = model_package.inference_specification.containers[0] + base_model_recipe_name = getattr( + getattr(container, "base_model", None), "recipe_name", None + ) + if not base_model_recipe_name: + raise ValueError( + "Cannot resolve the base model recipe for LoRA deployment. " + "Ensure the model package contains base model metadata." + ) + if not endpoint_name: endpoint_name = f"endpoint-{uuid.uuid4().hex[:8]}" + self._prepare_reused_model_customization_deployment_state( + model_package=model_package, + peft_type=peft_type, + inference_config=inference_config, + ) + + if inference_config is not None: + compute_requirements = InferenceComponentComputeResourceRequirements( + min_memory_required_in_mb=inference_config.min_memory, + max_memory_required_in_mb=inference_config.max_memory, + number_of_cpu_cores_required=inference_config.num_cpus, + number_of_accelerator_devices_required=inference_config.num_accelerators, + ) + copy_count = inference_config.copy_count + else: + compute_requirements = getattr(self, "_cached_compute_requirements", None) + if compute_requirements is None: + raise ValueError( + "Cannot resolve compute requirements for model customization deployment. " + "Provide ResourceRequirements explicitly." + ) + copy_count = 1 + + adapter_s3_uri = None + if peft_type == "LORA": + adapter_s3_uri = getattr(self, "_adapter_s3_uri", None) + if not adapter_s3_uri and isinstance( + self.model, (TrainingJob, ModelTrainer, AgentRFTJob, ModelPackage) + ): + adapter_s3_uri = self._resolve_lora_adapter_s3_uri(model_package) + self._adapter_s3_uri = adapter_s3_uri + if not adapter_s3_uri and getattr(self, "_built_model_was_reused", False): + raise ValueError( + "Cannot resolve the LoRA adapter artifact URI from the model source." + ) + # The endpoint config's network isolation must match the built Model, or # CreateInferenceComponent rejects the mismatch. Nova models are always # created with network isolation enabled; for other models honor the @@ -6326,16 +6437,6 @@ def _deploy_model_customization( else: endpoint = Endpoint.get(endpoint_name=endpoint_name) - # Without a model package (e.g. a Nova CPTTrainer or raw-S3 checkpoint) - # there is no PEFT/recipe metadata, so the deployment follows the - # single-IC path below. - peft_type = self._fetch_peft() if model_package is not None else None - base_model_recipe_name = ( - model_package.inference_specification.containers[0].base_model.recipe_name - if model_package is not None - else None - ) - if peft_type == "LORA": # LORA deployment: base IC + adapter IC @@ -6357,25 +6458,15 @@ def _deploy_model_customization( base_ic_spec = InferenceComponentSpecification( model_name=self.built_model.model_name, + compute_resource_requirements=compute_requirements, ) - if inference_config is not None: - base_ic_spec.compute_resource_requirements = ( - InferenceComponentComputeResourceRequirements( - min_memory_required_in_mb=inference_config.min_memory, - max_memory_required_in_mb=inference_config.max_memory, - number_of_cpu_cores_required=inference_config.num_cpus, - number_of_accelerator_devices_required=inference_config.num_accelerators, - ) - ) - else: - base_ic_spec.compute_resource_requirements = self._cached_compute_requirements InferenceComponent.create( inference_component_name=base_ic_name, endpoint_name=endpoint_name, variant_name=endpoint_name, specification=base_ic_spec, - runtime_config=InferenceComponentRuntimeConfig(copy_count=1), + runtime_config=InferenceComponentRuntimeConfig(copy_count=copy_count), tags=[{"key": "Base", "value": base_model_recipe_name}], ) logger.info("Created base model InferenceComponent: '%s'", base_ic_name) @@ -6389,7 +6480,6 @@ def _deploy_model_customization( # Deploy adapter IC adapter_ic_name = inference_component_name or f"{endpoint_name}-adapter" - adapter_s3_uri = getattr(self, "_adapter_s3_uri", None) adapter_ic_spec = InferenceComponentSpecification( base_inference_component_name=base_ic_name, @@ -6412,26 +6502,15 @@ def _deploy_model_customization( ic_spec = InferenceComponentSpecification( model_name=self.built_model.model_name, + compute_resource_requirements=compute_requirements, ) - if inference_config is not None: - ic_spec.compute_resource_requirements = ( - InferenceComponentComputeResourceRequirements( - min_memory_required_in_mb=inference_config.min_memory, - max_memory_required_in_mb=inference_config.max_memory, - number_of_cpu_cores_required=inference_config.num_cpus, - number_of_accelerator_devices_required=inference_config.num_accelerators, - ) - ) - else: - ic_spec.compute_resource_requirements = self._cached_compute_requirements - InferenceComponent.create( inference_component_name=inference_component_name, endpoint_name=endpoint_name, variant_name=endpoint_name, specification=ic_spec, - runtime_config=InferenceComponentRuntimeConfig(copy_count=1), + runtime_config=InferenceComponentRuntimeConfig(copy_count=copy_count), ) # Create lineage tracking for new endpoints. Lineage is keyed off the diff --git a/sagemaker-serve/tests/integ/test_model_customization_deployment.py b/sagemaker-serve/tests/integ/test_model_customization_deployment.py index 694273eb70..91f8ee0187 100644 --- a/sagemaker-serve/tests/integ/test_model_customization_deployment.py +++ b/sagemaker-serve/tests/integ/test_model_customization_deployment.py @@ -17,9 +17,11 @@ import json import boto3 import time +import uuid import pytest import random import logging +from unittest.mock import patch from botocore.config import Config from botocore.exceptions import ClientError from datetime import datetime, timezone, timedelta @@ -28,7 +30,13 @@ logger = logging.getLogger(__name__) from sagemaker.core.helper.session_helper import Session, get_execution_role -from sagemaker.core.resources import TrainingJob, ModelPackage, InferenceComponent, Endpoint +from sagemaker.core.resources import ( + Endpoint, + EndpointConfig, + InferenceComponent, + ModelPackage, + TrainingJob, +) from sagemaker.core.utils.exceptions import FailedStatusError from sagemaker.serve import ModelBuilder from sagemaker.serve.bedrock_model_builder import BedrockModelBuilder @@ -108,92 +116,208 @@ def test_build_from_training_job(self, training_job_name, sagemaker_session): assert model_builder.instance_type is not None @pytest.mark.skip_in_pr_check - def test_deploy_from_training_job(self, training_job_name, endpoint_name, cleanup_endpoints, sagemaker_session): - """Test deploying model from training job. - - For LORA models, this verifies the two-step deployment: - base IC + adapter IC are both created on the same endpoint. - """ - - - training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) - model_builder = ModelBuilder(model=training_job, instance_type="ml.g5.4xlarge", sagemaker_session=sagemaker_session) - model_builder.accept_eula = True - model_builder.build(model_name=f"test-model-{int(time.time())}-{random.randint(100, 10000)}", region=AWS_REGION) - - peft_type = model_builder._fetch_peft() - adapter_name = f"{endpoint_name}-adapter" + def test_deploy_from_training_job(self, training_job_name, sagemaker_session): + """Deploy, reuse, invoke, and clean up one training-job endpoint.""" + test_id = uuid.uuid4().hex + source_identity = f"model-customization-reuse-{test_id}" + model_name = f"reuse-model-{test_id[:12]}" + endpoint_name = f"reuse-endpoint-{test_id[:12]}" + base_ic_name = f"{endpoint_name}-inference-component" + adapter_ic_name = f"{endpoint_name}-adapter" + + model = None + endpoint = None + base_ic = None + adapter_ic = None + training_job = TrainingJob.get( + training_job_name=training_job_name, region=AWS_REGION + ) + first_builder = ModelBuilder( + model=training_job, + instance_type="ml.g5.4xlarge", + sagemaker_session=sagemaker_session, + ) + first_builder.accept_eula = True try: - endpoint = model_builder.deploy( - endpoint_name=endpoint_name, - inference_component_name=adapter_name if peft_type == "LORA" else None, - ) - except (FailedStatusError, ClientError) as e: - # xfail on environmental capacity/quota limits rather than fail the build. - msg = str(e) - if "InsufficientInstanceCapacity" in msg or "ResourceLimitExceeded" in msg: - cleanup_endpoints.append(endpoint_name) - pytest.xfail( - f"Environmental capacity/quota limit for ml.g5.4xlarge in {AWS_REGION}: {e}" - ) - raise - - cleanup_endpoints.append(endpoint_name) - - assert endpoint is not None - assert endpoint.endpoint_arn is not None - assert endpoint.endpoint_status == "InService" - - # Verify model-source tag is present on the endpoint for reuse discovery. - sm_client = boto3.client("sagemaker", region_name=AWS_REGION) - endpoint_tags = sm_client.list_tags(ResourceArn=endpoint.endpoint_arn).get("Tags", []) - assert MODEL_SOURCE_TAG_KEY in {t["Key"] for t in endpoint_tags}, ( - f"Endpoint {endpoint.endpoint_arn} missing model-source tag for reuse" - ) + with patch.object( + first_builder, "_resolve_model_source_id", return_value=source_identity + ): + model = first_builder.build(model_name=model_name, region=AWS_REGION) + peft_type = first_builder._fetch_peft() + try: + endpoint = first_builder.deploy( + endpoint_name=endpoint_name, + inference_component_name=( + adapter_ic_name if peft_type == "LORA" else None + ), + ) + except (FailedStatusError, ClientError) as error: + message = str(error) + if ( + "InsufficientInstanceCapacity" in message + or "ResourceLimitExceeded" in message + ): + pytest.xfail( + "Environmental capacity or quota limit prevented deployment" + ) + raise + + assert model.model_name == model_name + assert endpoint.endpoint_name == endpoint_name + assert endpoint.endpoint_status == "InService" - if peft_type == "LORA": - # Verify base IC was created - base_ic_name = f"{endpoint_name}-inference-component" - base_ic = InferenceComponent.get(inference_component_name=base_ic_name, region=AWS_REGION) - assert base_ic is not None - assert base_ic.inference_component_status == "InService" + sm_client = boto3.client("sagemaker", region_name=AWS_REGION) + model_tags = sm_client.list_tags(ResourceArn=model.model_arn).get("Tags", []) + endpoint_tags = sm_client.list_tags(ResourceArn=endpoint.endpoint_arn).get( + "Tags", [] + ) + assert any( + tag["Key"] == MODEL_SOURCE_TAG_KEY + and tag["Value"] == source_identity + for tag in model_tags + ) + assert any( + tag["Key"] == MODEL_SOURCE_TAG_KEY + and tag["Value"] == source_identity + for tag in endpoint_tags + ) - # Verify adapter IC was created - adapter_ic = InferenceComponent.get(inference_component_name=adapter_name, region=AWS_REGION) - assert adapter_ic is not None + if peft_type == "LORA": + base_ic = InferenceComponent.get( + inference_component_name=base_ic_name, region=AWS_REGION + ) + adapter_ic = InferenceComponent.get( + inference_component_name=adapter_ic_name, region=AWS_REGION + ) + assert base_ic.inference_component_status == "InService" + assert adapter_ic.inference_component_status == "InService" - # Invoke verification - time.sleep(10) # brief buffer for IC readiness + second_builder = ModelBuilder( + model=training_job, + instance_type="ml.g5.4xlarge", + sagemaker_session=sagemaker_session, + ) + second_builder.accept_eula = True + + with patch.object( + second_builder, "_resolve_model_source_id", return_value=source_identity + ): + deadline = time.monotonic() + 180 + while True: + discovered_model = second_builder._find_reusable_model() + discovered_endpoint = second_builder._find_reusable_endpoint( + instance_type="ml.g5.4xlarge" + ) + if ( + discovered_model is not None + and discovered_model.model_name == model_name + and discovered_endpoint == endpoint_name + ): + break + if time.monotonic() >= deadline: + raise AssertionError( + "Timed out waiting for the exact test resources to become discoverable" + ) + time.sleep(5) + + with ( + patch.object( + second_builder, + "_build_single_modelbuilder", + side_effect=AssertionError( + "Model discovery missed and attempted to create another Model" + ), + ), + patch.object( + second_builder, + "_deploy_model_customization", + side_effect=AssertionError( + "Endpoint discovery missed and attempted another deployment" + ), + ), + ): + reused_model = second_builder.build( + region=AWS_REGION, reuse_resources=True + ) + reused_endpoint = second_builder.deploy( + endpoint_name=endpoint_name, reuse_resources=True + ) - invoke_ic_name = adapter_name if peft_type == "LORA" else f"{endpoint_name}-inference-component" + assert reused_model.model_name == model.model_name + assert reused_model.model_arn == model.model_arn + assert reused_endpoint.endpoint_name == endpoint.endpoint_name + assert reused_endpoint.endpoint_arn == endpoint.endpoint_arn - test_payload = { - "inputs": "What is machine learning?", - "parameters": {"max_new_tokens": 32}, - } + time.sleep(10) + invoke_ic_name = ( + adapter_ic_name if peft_type == "LORA" else base_ic_name + ) + invoke_response = reused_endpoint.invoke( + body=json.dumps( + { + "inputs": "What is machine learning?", + "parameters": {"max_new_tokens": 32}, + } + ), + content_type="application/json", + accept="application/json", + inference_component_name=invoke_ic_name, + ) + response_body = json.loads(invoke_response.body.read()) + assert response_body is not None + if isinstance(response_body, list): + assert response_body + assert ( + "generated_text" in response_body[0] + or "generation" in response_body[0] + ) + elif isinstance(response_body, dict): + assert any( + key in response_body + for key in ("generated_text", "generation", "outputs") + ) + finally: + for component, name in ( + (adapter_ic, adapter_ic_name), + (base_ic, base_ic_name), + ): + if component is None: + try: + component = InferenceComponent.get( + inference_component_name=name, region=AWS_REGION + ) + except Exception: + continue + try: + component.delete() + component.wait_for_delete(timeout=300) + except Exception as error: + logger.warning("Failed to clean up inference component %s: %s", name, error) - invoke_response = endpoint.invoke( - body=json.dumps(test_payload), - content_type="application/json", - accept="application/json", - inference_component_name=invoke_ic_name, - ) + if endpoint is not None: + try: + endpoint.delete() + endpoint.wait_for_delete(timeout=300) + except Exception as error: + logger.warning("Failed to clean up endpoint %s: %s", endpoint_name, error) - response_body = json.loads(invoke_response.body.read()) - - # Validate response structure - assert response_body is not None, f"Empty response from invoke on {invoke_ic_name}" - if isinstance(response_body, list): - assert len(response_body) > 0 - assert "generated_text" in response_body[0] or "generation" in response_body[0] - elif isinstance(response_body, dict): - assert ( - "generated_text" in response_body - or "generation" in response_body - or "outputs" in response_body - ) + try: + EndpointConfig.get( + endpoint_config_name=endpoint_name, region=AWS_REGION + ).delete() + except Exception as error: + logger.warning( + "Failed to clean up endpoint configuration %s: %s", + endpoint_name, + error, + ) + if model is not None: + try: + model.delete() + except Exception as error: + logger.warning("Failed to clean up Model %s: %s", model_name, error) def test_fetch_endpoint_names_for_base_model(self, training_job_name, sagemaker_session): """Test fetching endpoint names for base model.""" @@ -204,87 +328,6 @@ def test_fetch_endpoint_names_for_base_model(self, training_job_name, sagemaker_ assert isinstance(endpoint_names, set) - def test_deploy_reuse_returns_existing_endpoint(self, training_job_name, endpoint_name, cleanup_endpoints, sagemaker_session): - """deploy(reuse_resources=True) finds and returns an existing tagged endpoint. - - Verifies that the reuse mechanism finds an endpoint with the - correct model-source tag. Because prior test runs may have left - behind an endpoint with the same tag, we accept any InService - endpoint carrying the tag as a valid reuse hit. - """ - - training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) - - # Try reuse first — if a tagged endpoint already exists from a prior - # run, reuse should find it without needing to create a new one. - builder = ModelBuilder(model=training_job, instance_type="ml.g5.4xlarge", sagemaker_session=sagemaker_session) - builder.accept_eula = True - builder.build(region=AWS_REGION, reuse_resources=True) - endpoint = builder.deploy(reuse_resources=True) - - if endpoint is not None and hasattr(endpoint, "endpoint_arn") and endpoint.endpoint_arn: - # Reuse found an existing endpoint — verify it has the tag - sm_client = boto3.client("sagemaker", region_name=AWS_REGION) - tags = sm_client.list_tags(ResourceArn=endpoint.endpoint_arn).get("Tags", []) - assert MODEL_SOURCE_TAG_KEY in {t["Key"] for t in tags}, ( - f"Reused endpoint {endpoint.endpoint_arn} missing expected tag {MODEL_SOURCE_TAG_KEY}" - ) - return - - # No reusable endpoint found — create one and verify the tag is applied - model_builder = ModelBuilder(model=training_job, instance_type="ml.g5.4xlarge", sagemaker_session=sagemaker_session) - model_builder.accept_eula = True - model_builder.build(model_name=f"test-model-{int(time.time())}-{random.randint(100, 10000)}", region=AWS_REGION) - - try: - endpoint = model_builder.deploy(endpoint_name=endpoint_name) - except (FailedStatusError, ClientError) as e: - msg = str(e) - if "InsufficientInstanceCapacity" in msg or "ResourceLimitExceeded" in msg: - cleanup_endpoints.append(endpoint_name) - pytest.xfail(f"Capacity/quota limit: {e}") - raise - - cleanup_endpoints.append(endpoint_name) - assert endpoint is not None - assert endpoint.endpoint_status == "InService" - - # Verify model-source tag was applied - sm_client = boto3.client("sagemaker", region_name=AWS_REGION) - tags = sm_client.list_tags(ResourceArn=endpoint.endpoint_arn).get("Tags", []) - assert MODEL_SOURCE_TAG_KEY in {t["Key"] for t in tags}, ( - f"Endpoint {endpoint.endpoint_arn} missing expected tag {MODEL_SOURCE_TAG_KEY}" - ) - - def test_build_reuse_skips_model_creation(self, training_job_name, sagemaker_session): - """build(reuse_resources=True) reuses an existing tagged Model. - - Verifies that the reuse lookup finds a model with the correct - model-source tag. Because prior test runs may have left behind a - model with the same tag, we accept any model carrying the tag as - a valid reuse hit (not only the one created in *this* test run). - """ - - training_job = TrainingJob.get(training_job_name=training_job_name, region=AWS_REGION) - - # Build with reuse — this should find any existing model tagged - # with this training job's source ID. It doesn't matter whether - # the model was created in this test run or a prior one; what - # matters is that the reuse mechanism works. - builder = ModelBuilder(model=training_job, instance_type="ml.g5.4xlarge", sagemaker_session=sagemaker_session) - builder.accept_eula = True - model = builder.build(region=AWS_REGION, reuse_resources=True) - - assert model is not None - assert model.model_arn is not None - - # Verify the model-source tag is actually present on the reused model - sm_client = boto3.client("sagemaker", region_name=AWS_REGION) - tags = sm_client.list_tags(ResourceArn=model.model_arn).get("Tags", []) - assert MODEL_SOURCE_TAG_KEY in {t["Key"] for t in tags}, ( - f"Reused model {model.model_arn} missing expected tag {MODEL_SOURCE_TAG_KEY}" - ) - class TestModelCustomizationFromModelPackage: diff --git a/sagemaker-serve/tests/unit/test_model_builder.py b/sagemaker-serve/tests/unit/test_model_builder.py index c23f79c6e0..f4efb17bc8 100644 --- a/sagemaker-serve/tests/unit/test_model_builder.py +++ b/sagemaker-serve/tests/unit/test_model_builder.py @@ -8,6 +8,7 @@ import json import unittest +from contextlib import ExitStack from unittest.mock import Mock, patch, MagicMock from botocore.exceptions import ClientError @@ -929,6 +930,11 @@ def _patch_lora_deps(self, mb, hosting_uri="s3://bucket/hosting/", patch.object(mb, "_fetch_and_cache_recipe_config"), patch.object(mb, "_is_nova_model", return_value=False), patch.object(mb, "_fetch_peft", return_value="LORA"), + patch.object( + mb, + "_resolve_lora_adapter_s3_uri", + return_value="s3://test-bucket/adapter/checkpoints/hf/", + ), patch.object(mb, "_fetch_hub_document_for_custom_model", return_value=hub_document), ] @@ -1314,6 +1320,347 @@ def test_deploy_reuses_cached_endpoint_when_instance_type_matches( mock_deploy.assert_not_called() assert result == mock_endpoint + def _make_customization_builder(self): + from sagemaker.core.resources import TrainingJob + + training_job = Mock(spec=TrainingJob) + training_job.model_artifacts = Mock( + s3_model_artifacts="s3://test-bucket/training-output" + ) + builder = self._make_builder(model=training_job, instance_type="ml.g5.4xlarge") + return builder + + @staticmethod + def _make_customization_package(recipe_name="test-lora"): + package = Mock() + package.model_package_arn = ( + "arn:aws:sagemaker:us-west-2:123456789012:model-package/test-package/1" + ) + container = Mock() + container.base_model.recipe_name = recipe_name + container.model_data_source.s3_data_source.s3_uri = "s3://test-bucket/model" + package.inference_specification.containers = [container] + return package + + def test_reused_model_marker_lifecycle_and_reset(self): + from sagemaker.core.shapes import InferenceComponentComputeResourceRequirements + + reused_model = Mock(model_name="reused-model", model_arn="reused-model-arn") + builder = self._make_customization_builder() + assert builder._built_model_was_reused is False + + with ( + patch.object(builder, "_get_serve_setting", return_value=Mock()), + patch.object(builder, "_find_reusable_model", return_value=reused_model), + ): + assert builder.build(reuse_resources=True) is reused_model + + assert builder._built_model_was_reused is True + builder._cached_compute_requirements = InferenceComponentComputeResourceRequirements( + min_memory_required_in_mb=1024 + ) + builder._adapter_s3_uri = "s3://test-bucket/adapter" + builder._reset_build_state() + assert builder._built_model_was_reused is False + assert not hasattr(builder, "_cached_compute_requirements") + assert not hasattr(builder, "_adapter_s3_uri") + + created_model = Mock(model_name="created-model", model_arn="created-model-arn") + + def create_model(**_kwargs): + builder.built_model = created_model + return created_model + + with ( + patch.object(builder, "_get_serve_setting", return_value=Mock()), + patch.object(builder, "_find_reusable_model", return_value=None), + patch.object(builder, "_build_single_modelbuilder", side_effect=create_model), + ): + assert builder.build(reuse_resources=True) is created_model + + assert builder._built_model_was_reused is False + + @patch("sagemaker.serve.model_builder.InferenceComponent.create") + @patch("sagemaker.serve.model_builder.EndpointConfig.create") + @patch("sagemaker.serve.model_builder.Endpoint.create") + @patch("sagemaker.serve.model_builder.Model.create") + def test_reuse_hit_returns_exact_model_and_endpoint_without_preparation( + self, mock_model_create, mock_endpoint_create, mock_config_create, mock_ic_create + ): + reused_model = Mock(model_name="reused-model", model_arn="reused-model-arn") + reused_endpoint = Mock(endpoint_name="reused-endpoint", endpoint_arn="reused-endpoint-arn") + builder = self._make_customization_builder() + + with ( + patch.object(builder, "_get_serve_setting", return_value=Mock()), + patch.object(builder, "_find_reusable_model", return_value=reused_model), + patch.object(builder, "_find_reusable_endpoint", return_value="reused-endpoint"), + patch.object(builder, "_deploy_model_customization") as mock_custom_deploy, + patch.object( + builder, "_prepare_reused_model_customization_deployment_state" + ) as mock_prepare, + patch("sagemaker.serve.model_builder.Endpoint.get", return_value=reused_endpoint), + ): + assert builder.build(reuse_resources=True) is reused_model + assert builder.deploy(reuse_resources=True) is reused_endpoint + + assert builder.built_model is reused_model + mock_prepare.assert_not_called() + mock_custom_deploy.assert_not_called() + mock_model_create.assert_not_called() + mock_config_create.assert_not_called() + mock_endpoint_create.assert_not_called() + mock_ic_create.assert_not_called() + + def test_lora_model_reuse_endpoint_miss_restores_state_before_writes(self): + from sagemaker.core.resources import ( + Action, + Artifact, + Association, + Endpoint, + EndpointConfig, + InferenceComponent, + Model, + Tag, + ) + from sagemaker.core.shapes import InferenceComponentComputeResourceRequirements + + package = self._make_customization_package() + reused_model = Mock(model_name="reused-model", model_arn="reused-model-arn") + created_endpoint = Mock(endpoint_name="new-endpoint", endpoint_arn="new-endpoint-arn") + created_endpoint.wait_for_status = Mock() + base_component = Mock(inference_component_arn="base-component-arn") + base_component.wait_for_status = Mock() + compute_requirements = InferenceComponentComputeResourceRequirements( + min_memory_required_in_mb=2048, + number_of_cpu_cores_required=4, + number_of_accelerator_devices_required=1, + ) + adapter_uri = "s3://test-bucket/training-output/checkpoints/hf/" + builder = self._make_customization_builder() + create_calls = [] + + def prepare_recipe(): + builder._cached_compute_requirements = compute_requirements + + def create_config(**kwargs): + assert builder._cached_compute_requirements is compute_requirements + assert builder._adapter_s3_uri == adapter_uri + create_calls.append(("config", kwargs)) + return Mock() + + with ExitStack() as stack: + stack.enter_context( + patch.object(builder, "_get_serve_setting", return_value=Mock()) + ) + stack.enter_context( + patch.object(builder, "_find_reusable_model", return_value=reused_model) + ) + stack.enter_context( + patch.object(builder, "_find_reusable_endpoint", return_value=None) + ) + stack.enter_context( + patch.object(builder, "_resolve_model_source_id", return_value="test-source") + ) + stack.enter_context(patch.object(builder, "add_tags")) + stack.enter_context( + patch.object(builder, "_is_model_customization", return_value=True) + ) + stack.enter_context(patch.object(builder, "_is_nova_model", return_value=False)) + stack.enter_context( + patch.object(builder, "_fetch_model_package", return_value=package) + ) + stack.enter_context(patch.object(builder, "_fetch_peft", return_value="LORA")) + stack.enter_context( + patch.object( + builder, "_fetch_and_cache_recipe_config", side_effect=prepare_recipe + ) + ) + stack.enter_context( + patch.object( + builder, "_resolve_lora_adapter_s3_uri", return_value=adapter_uri + ) + ) + stack.enter_context( + patch.object(builder, "_does_endpoint_exist", return_value=False) + ) + stack.enter_context(patch.object(EndpointConfig, "create", side_effect=create_config)) + stack.enter_context(patch.object(Endpoint, "create", return_value=created_endpoint)) + stack.enter_context(patch.object(InferenceComponent, "get_all", return_value=[])) + mock_ic_create = stack.enter_context( + patch.object(InferenceComponent, "create", return_value=Mock()) + ) + stack.enter_context( + patch.object(InferenceComponent, "get", return_value=base_component) + ) + stack.enter_context(patch.object(Tag, "get_all", return_value=[])) + stack.enter_context( + patch.object(Action, "create", side_effect=Exception("skip lineage")) + ) + stack.enter_context(patch.object(Artifact, "get_all", return_value=[])) + stack.enter_context(patch.object(Association, "add")) + mock_model_create = stack.enter_context(patch.object(Model, "create")) + + assert builder.build(reuse_resources=True) is reused_model + result = builder.deploy(endpoint_name="new-endpoint", reuse_resources=True) + + assert result is created_endpoint + assert builder.built_model is reused_model + assert create_calls[0][0] == "config" + assert mock_ic_create.call_count == 2 + base_spec = mock_ic_create.call_args_list[0].kwargs["specification"] + adapter_spec = mock_ic_create.call_args_list[1].kwargs["specification"] + assert base_spec.compute_resource_requirements is compute_requirements + assert adapter_spec.container.artifact_url == adapter_uri + mock_model_create.assert_not_called() + + def test_reused_non_lora_restores_compute_without_creating_adapter(self): + from sagemaker.core.resources import Endpoint, EndpointConfig, InferenceComponent, Model + from sagemaker.core.shapes import InferenceComponentComputeResourceRequirements + + package = self._make_customization_package(recipe_name="full-finetuning") + endpoint = Mock() + endpoint.wait_for_status = Mock() + compute_requirements = InferenceComponentComputeResourceRequirements( + min_memory_required_in_mb=1024 + ) + builder = self._make_customization_builder() + builder.built_model = Mock(model_name="reused-model") + builder._built_model_was_reused = True + + def prepare_recipe(): + builder._cached_compute_requirements = compute_requirements + + with ( + patch.object(builder, "_is_nova_model", return_value=False), + patch.object(builder, "_fetch_model_package", return_value=package), + patch.object(builder, "_fetch_peft", return_value=None), + patch.object(builder, "_fetch_and_cache_recipe_config", side_effect=prepare_recipe), + patch.object(builder, "_does_endpoint_exist", return_value=False), + patch.object(EndpointConfig, "create"), + patch.object(Endpoint, "create", return_value=endpoint), + patch.object(InferenceComponent, "create", return_value=Mock()) as mock_ic_create, + patch.object(Model, "create") as mock_model_create, + ): + result = builder._deploy_model_customization(endpoint_name="new-endpoint") + + assert result is endpoint + assert mock_ic_create.call_count == 1 + assert ( + mock_ic_create.call_args.kwargs["specification"].compute_resource_requirements + is compute_requirements + ) + mock_model_create.assert_not_called() + + def test_reused_lora_explicit_requirements_preserve_values_and_restore_adapter(self): + from sagemaker.core.inference_config import ResourceRequirements + from sagemaker.core.resources import Endpoint, EndpointConfig, InferenceComponent, Model, Tag + + package = self._make_customization_package() + endpoint = Mock() + endpoint.wait_for_status = Mock() + base_component = Mock(inference_component_arn="base-component-arn") + base_component.wait_for_status = Mock() + requirements = ResourceRequirements( + requests={"num_cpus": 8, "memory": 49152, "num_accelerators": 4, "copies": 3}, + limits={"memory": 98304}, + ) + adapter_uri = "s3://test-bucket/training-output/checkpoints/hf/" + builder = self._make_customization_builder() + builder.built_model = Mock(model_name="reused-model") + builder._built_model_was_reused = True + + with ( + patch.object(builder, "_is_nova_model", return_value=False), + patch.object(builder, "_fetch_model_package", return_value=package), + patch.object(builder, "_fetch_peft", return_value="LORA"), + patch.object(builder, "_fetch_and_cache_recipe_config") as mock_recipe, + patch.object(builder, "_resolve_lora_adapter_s3_uri", return_value=adapter_uri), + patch.object(builder, "_does_endpoint_exist", return_value=False), + patch.object(EndpointConfig, "create"), + patch.object(Endpoint, "create", return_value=endpoint), + patch.object(InferenceComponent, "get_all", return_value=[]), + patch.object(InferenceComponent, "create", return_value=Mock()) as mock_ic_create, + patch.object(InferenceComponent, "get", return_value=base_component), + patch.object(Tag, "get_all", return_value=[]), + patch.object(Model, "create") as mock_model_create, + ): + builder._deploy_model_customization( + endpoint_name="new-endpoint", inference_config=requirements + ) + + mock_recipe.assert_not_called() + base_call = mock_ic_create.call_args_list[0] + compute = base_call.kwargs["specification"].compute_resource_requirements + assert compute.number_of_cpu_cores_required == 8 + assert compute.min_memory_required_in_mb == 49152 + assert compute.max_memory_required_in_mb == 98304 + assert compute.number_of_accelerator_devices_required == 4 + assert base_call.kwargs["runtime_config"].copy_count == 3 + assert mock_ic_create.call_args_list[1].kwargs["specification"].container.artifact_url == adapter_uri + assert requirements.copy_count == 3 + mock_model_create.assert_not_called() + + def test_lora_adapter_uri_resolver_supported_and_rejected_sources(self): + from sagemaker.core.resources import ModelPackage, TrainingJob + from sagemaker.train.agent_rft_job import AgentRFTJob + from sagemaker.train.base_trainer import BaseTrainer + from sagemaker.train.model_trainer import ModelTrainer + + package = self._make_customization_package() + cases = [] + + training_job = Mock(spec=TrainingJob) + training_job.model_artifacts = Mock(s3_model_artifacts="s3://test-bucket/training") + cases.append((training_job, "s3://test-bucket/training/checkpoints/hf/")) + + model_trainer = Mock(spec=ModelTrainer) + model_trainer._latest_training_job = Mock( + model_artifacts=Mock(s3_model_artifacts="s3://test-bucket/trainer") + ) + cases.append((model_trainer, "s3://test-bucket/trainer/checkpoints/hf/")) + + agent_job = Mock(spec=AgentRFTJob) + cases.append((agent_job, "s3://test-bucket/model/checkpoints/hf/")) + + model_package = Mock(spec=ModelPackage) + cases.append((model_package, "s3://test-bucket/model/model/checkpoints/hf/")) + + for model, expected in cases: + with self.subTest(model_type=type(model).__name__): + builder = self._make_builder(model=model) + assert builder._resolve_lora_adapter_s3_uri(package) == expected + + unsupported = Mock(spec=BaseTrainer) + builder = self._make_builder(model=unsupported) + with self.assertRaisesRegex(ValueError, "Use a TrainingJob"): + builder._resolve_lora_adapter_s3_uri(package) + + from sagemaker.core.inference_config import ResourceRequirements + from sagemaker.core.resources import Endpoint, EndpointConfig, InferenceComponent + + builder.built_model = Mock(model_name="reused-model") + builder._built_model_was_reused = True + requirements = ResourceRequirements(requests={"memory": 4096}) + with ( + patch.object(builder, "_is_nova_model", return_value=False), + patch.object(builder, "_fetch_model_package", return_value=package), + patch.object(builder, "_fetch_peft", return_value="LORA"), + patch.object(builder, "_does_endpoint_exist") as mock_endpoint_exists, + patch.object(EndpointConfig, "create") as mock_config_create, + patch.object(Endpoint, "create") as mock_endpoint_create, + patch.object(InferenceComponent, "create") as mock_ic_create, + ): + with self.assertRaisesRegex(ValueError, "Use a TrainingJob"): + builder._deploy_model_customization( + endpoint_name="new-endpoint", inference_config=requirements + ) + + mock_endpoint_exists.assert_not_called() + mock_config_create.assert_not_called() + mock_endpoint_create.assert_not_called() + mock_ic_create.assert_not_called() + class TestReusedEndpointMatchesConfig(unittest.TestCase): """Tests for ModelBuilder._reused_endpoint_matches_config."""