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."""