diff --git a/sagemaker-core/src/sagemaker/core/common_utils.py b/sagemaker-core/src/sagemaker/core/common_utils.py index b21c05ca16..4360bfc5e5 100644 --- a/sagemaker-core/src/sagemaker/core/common_utils.py +++ b/sagemaker-core/src/sagemaker/core/common_utils.py @@ -2129,6 +2129,75 @@ def _resolve_routing_config(routing_config: Optional[Dict[str, Any]]) -> Optiona return None +_CAPACITY_RESERVATION_CONFIG_KEYS = ("CapacityReservationPreference", "MlReservationArn") +_CAPACITY_RESERVATIONS_ONLY = "capacity-reservations-only" + + +def _resolve_capacity_reservation_config( + capacity_reservation_config: Optional[Any], +) -> Optional[Dict[str, Any]]: + """Resolve a production variant capacity reservation config to its request dict. + + Args: + capacity_reservation_config (Optional[Union[dict, + sagemaker.core.shapes.ProductionVariantCapacityReservationConfig]]): + Either a ``ProductionVariantCapacityReservationConfig`` shape or a dict using the + ``CreateEndpointConfig`` request keys (``CapacityReservationPreference``, + ``MlReservationArn``). + + Returns: + Optional[Dict[str, Any]]: The ``CapacityReservationConfig`` request dict, or None if + no config was given or it has no fields set. A training plan ARN + (``...:training-plan/``) given without a preference gets + ``CapacityReservationPreference="capacity-reservations-only"``, the only value + SageMaker accepts for training plans. Other ARNs are returned as given. + + Raises: + ValueError: If a dict contains keys other than ``CapacityReservationPreference`` and + ``MlReservationArn``. + TypeError: If the config is neither a dict nor a + ``ProductionVariantCapacityReservationConfig``. + """ + if capacity_reservation_config is None: + return None + + # Imported lazily: sagemaker.core.shapes imports from this module. + from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig + from sagemaker.core.utils.utils import serialize + + if isinstance(capacity_reservation_config, ProductionVariantCapacityReservationConfig): + return _default_training_plan_preference(serialize(capacity_reservation_config) or None) + + if isinstance(capacity_reservation_config, dict): + unknown_keys = set(capacity_reservation_config) - set(_CAPACITY_RESERVATION_CONFIG_KEYS) + if unknown_keys: + raise ValueError( + f"Unsupported capacity_reservation_config keys: {sorted(unknown_keys)}. " + f"Supported keys are: {list(_CAPACITY_RESERVATION_CONFIG_KEYS)}." + ) + resolved = {k: v for k, v in capacity_reservation_config.items() if v is not None} + return _default_training_plan_preference(resolved or None) + + raise TypeError( + "capacity_reservation_config must be a dict or a " + "sagemaker.core.shapes.ProductionVariantCapacityReservationConfig, got " + f"{type(capacity_reservation_config).__name__}." + ) + + +def _default_training_plan_preference( + resolved: Optional[Dict[str, Any]], +) -> Optional[Dict[str, Any]]: + """Set the preference a training plan requires when the caller left it unset.""" + if ( + resolved + and ":training-plan/" in str(resolved.get("MlReservationArn", "")) + and "CapacityReservationPreference" not in resolved + ): + return {**resolved, "CapacityReservationPreference": _CAPACITY_RESERVATIONS_ONLY} + return resolved + + @lru_cache def get_instance_rate_per_hour( instance_type: str, diff --git a/sagemaker-core/src/sagemaker/core/helper/session_helper.py b/sagemaker-core/src/sagemaker/core/helper/session_helper.py index 4b6e0f1487..e80309af8a 100644 --- a/sagemaker-core/src/sagemaker/core/helper/session_helper.py +++ b/sagemaker-core/src/sagemaker/core/helper/session_helper.py @@ -49,6 +49,7 @@ instance_supports_kms, create_paginator_config, validate_path_within_directory, + _resolve_capacity_reservation_config, ) from sagemaker.core.config.config_utils import _log_sagemaker_config_merge @@ -2940,6 +2941,7 @@ def production_variant( managed_instance_scaling=None, routing_config=None, inference_ami_version=None, + capacity_reservation_config=None, ): """Create a production variant description suitable for use in a ``ProductionVariant`` list. @@ -2973,6 +2975,13 @@ def production_variant( https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-inference-code.html #your-algorithms -inference-algo-ping-requests + capacity_reservation_config (Union[dict, + sagemaker.core.shapes.ProductionVariantCapacityReservationConfig]): Runs the + variant's instances on reserved capacity, either a SageMaker training plan + created for the ``endpoint`` target resource or an ML reservation backed by + On-Demand Capacity Reservations (ODCRs). A dict must use the + ``CreateEndpointConfig`` request keys ``CapacityReservationPreference`` and + ``MlReservationArn``. Ignored for serverless variants (default: None). Returns: dict[str, str]: An SageMaker ``ProductionVariant`` description @@ -3002,6 +3011,9 @@ def production_variant( ModelDataDownloadTimeoutInSeconds=model_data_download_timeout, ContainerStartupHealthCheckTimeoutInSeconds=container_startup_health_check_timeout, RoutingConfig=routing_config, + CapacityReservationConfig=_resolve_capacity_reservation_config( + capacity_reservation_config + ), ) if inference_ami_version: diff --git a/sagemaker-core/tests/unit/helper/test_session_helper.py b/sagemaker-core/tests/unit/helper/test_session_helper.py index 1b05aa1491..77022a5803 100644 --- a/sagemaker-core/tests/unit/helper/test_session_helper.py +++ b/sagemaker-core/tests/unit/helper/test_session_helper.py @@ -1539,6 +1539,32 @@ def test_production_variant_serverless(self): assert result["ServerlessConfig"] == serverless_config assert "InstanceType" not in result + def test_production_variant_with_capacity_reservation_config(self): + """Test production variant bound to a training plan.""" + from sagemaker.core.helper.session_helper import production_variant + from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig + + plan_arn = "arn:aws:sagemaker:us-west-2:123456789012:training-plan/my-plan" + result = production_variant( + model_name="my-model", + instance_type="ml.p5.48xlarge", + capacity_reservation_config=ProductionVariantCapacityReservationConfig( + capacity_reservation_preference="capacity-reservations-only", + ml_reservation_arn=plan_arn, + ), + ) + assert result["CapacityReservationConfig"] == { + "CapacityReservationPreference": "capacity-reservations-only", + "MlReservationArn": plan_arn, + } + + def test_production_variant_without_capacity_reservation_config(self): + """Test production variant omits CapacityReservationConfig by default.""" + from sagemaker.core.helper.session_helper import production_variant + + result = production_variant(model_name="my-model", instance_type="ml.m5.xlarge") + assert "CapacityReservationConfig" not in result + def test_production_variant_with_volume_size(self): """Test production variant with volume size.""" from sagemaker.core.helper.session_helper import production_variant diff --git a/sagemaker-core/tests/unit/test_common_utils.py b/sagemaker-core/tests/unit/test_common_utils.py index 8edc9dcdf7..0ae61b321e 100644 --- a/sagemaker-core/tests/unit/test_common_utils.py +++ b/sagemaker-core/tests/unit/test_common_utils.py @@ -2271,6 +2271,94 @@ def test_resolve_routing_config_none(self): assert result is None +class TestResolveCapacityReservationConfig: + """Test _resolve_capacity_reservation_config function.""" + + PLAN_ARN = "arn:aws:sagemaker:us-west-2:123456789012:training-plan/my-plan" + ML_ARN = "arn:aws:sagemaker:us-west-2:123456789012:ml-reservation/my-reservation" + + def test_shape_is_serialized_to_request_keys(self): + from sagemaker.core.common_utils import _resolve_capacity_reservation_config + from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig + + config = ProductionVariantCapacityReservationConfig( + capacity_reservation_preference="capacity-reservations-only", + ml_reservation_arn=self.PLAN_ARN, + ) + assert _resolve_capacity_reservation_config(config) == { + "CapacityReservationPreference": "capacity-reservations-only", + "MlReservationArn": self.PLAN_ARN, + } + + def test_shape_with_ml_reservation_arn_only_omits_preference(self): + from sagemaker.core.common_utils import _resolve_capacity_reservation_config + from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig + + config = ProductionVariantCapacityReservationConfig(ml_reservation_arn=self.ML_ARN) + assert _resolve_capacity_reservation_config(config) == {"MlReservationArn": self.ML_ARN} + + def test_dict_is_passed_through_without_none_values(self): + from sagemaker.core.common_utils import _resolve_capacity_reservation_config + + config = {"MlReservationArn": self.ML_ARN, "CapacityReservationPreference": None} + assert _resolve_capacity_reservation_config(config) == {"MlReservationArn": self.ML_ARN} + + def test_shape_with_training_plan_arn_only_gets_reservations_only(self): + from sagemaker.core.common_utils import _resolve_capacity_reservation_config + from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig + + config = ProductionVariantCapacityReservationConfig(ml_reservation_arn=self.PLAN_ARN) + assert _resolve_capacity_reservation_config(config) == { + "MlReservationArn": self.PLAN_ARN, + "CapacityReservationPreference": "capacity-reservations-only", + } + + def test_dict_with_training_plan_arn_only_gets_reservations_only(self): + from sagemaker.core.common_utils import _resolve_capacity_reservation_config + + config = {"MlReservationArn": self.PLAN_ARN, "CapacityReservationPreference": None} + assert _resolve_capacity_reservation_config(config) == { + "MlReservationArn": self.PLAN_ARN, + "CapacityReservationPreference": "capacity-reservations-only", + } + + def test_training_plan_dict_input_is_not_mutated(self): + from sagemaker.core.common_utils import _resolve_capacity_reservation_config + + config = {"MlReservationArn": self.PLAN_ARN} + _resolve_capacity_reservation_config(config) + assert config == {"MlReservationArn": self.PLAN_ARN} + + def test_explicit_preference_is_kept(self): + from sagemaker.core.common_utils import _resolve_capacity_reservation_config + + config = {"MlReservationArn": self.PLAN_ARN, "CapacityReservationPreference": "other"} + assert _resolve_capacity_reservation_config(config) == config + + def test_empty_config_resolves_to_none(self): + from sagemaker.core.common_utils import _resolve_capacity_reservation_config + from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig + + assert _resolve_capacity_reservation_config(None) is None + assert _resolve_capacity_reservation_config({}) is None + assert ( + _resolve_capacity_reservation_config(ProductionVariantCapacityReservationConfig()) + is None + ) + + def test_dict_with_unknown_key_raises(self): + from sagemaker.core.common_utils import _resolve_capacity_reservation_config + + with pytest.raises(ValueError, match="ml_reservation_arn"): + _resolve_capacity_reservation_config({"ml_reservation_arn": self.PLAN_ARN}) + + def test_unsupported_type_raises(self): + from sagemaker.core.common_utils import _resolve_capacity_reservation_config + + with pytest.raises(TypeError): + _resolve_capacity_reservation_config(self.PLAN_ARN) + + class TestWaitUntil: """Test _wait_until function.""" diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 07e7c18268..b8e6367cd5 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -50,6 +50,7 @@ InferenceComponentContainerSpecification, InferenceComponentRuntimeConfig, ProductionVariant, + ProductionVariantCapacityReservationConfig, ) from sagemaker.core.resources import ( ModelPackage, @@ -120,6 +121,7 @@ from sagemaker.core.enums import EndpointType from sagemaker.core.common_utils import ( Tags, + _resolve_capacity_reservation_config, _resolve_routing_config, format_tags, resolve_value_from_config, @@ -182,6 +184,21 @@ SAGEMAKER_OUTPUT_LOCATION = "sagemaker_s3_output" +def _capacity_reservation_config_shape( + capacity_reservation_config: Optional[ + Union[Dict[str, Any], ProductionVariantCapacityReservationConfig] + ], +) -> Optional[ProductionVariantCapacityReservationConfig]: + """Normalize a dict or shape capacity reservation config to the core shape.""" + resolved = _resolve_capacity_reservation_config(capacity_reservation_config) + if not resolved: + return None + return ProductionVariantCapacityReservationConfig( + capacity_reservation_preference=resolved.get("CapacityReservationPreference"), + ml_reservation_arn=resolved.get("MlReservationArn"), + ) + + @dataclass class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuilderUtils): """Unified interface for building and deploying machine learning models. @@ -2519,14 +2536,23 @@ def _find_reusable_model(self) -> Optional["Model"]: return None - def _find_reusable_endpoint(self, instance_type: Optional[str] = None) -> Optional[str]: + def _find_reusable_endpoint( + self, + instance_type: Optional[str] = None, + capacity_reservation_config: Optional[ + Union[Dict[str, Any], ProductionVariantCapacityReservationConfig] + ] = None, + ) -> Optional[str]: """Return the name of an existing endpoint that can be reused, if any. A candidate must carry the same model-source tag and match the requested - deployment configuration (env vars, image URI, instance type). + deployment configuration (env vars, image URI, instance type, capacity + reservation). Args: instance_type: The instance type requested for this deploy, if any. + capacity_reservation_config: The capacity reservation requested for this + deploy, if any. Returns: The reusable endpoint name, or None if there is no suitable match. @@ -2544,7 +2570,9 @@ def _find_reusable_endpoint(self, instance_type: Optional[str] = None) -> Option existing_name = existing_arn.rsplit("/", 1)[-1] if self._reused_endpoint_matches_config( - existing_name, instance_type=instance_type or self.instance_type + existing_name, + instance_type=instance_type or self.instance_type, + capacity_reservation_config=capacity_reservation_config, ): return existing_name @@ -2556,19 +2584,27 @@ def _find_reusable_endpoint(self, instance_type: Optional[str] = None) -> Option return None def _reused_endpoint_matches_config( - self, endpoint_name: str, instance_type: Optional[str] = None + self, + endpoint_name: str, + instance_type: Optional[str] = None, + capacity_reservation_config: Optional[ + Union[Dict[str, Any], ProductionVariantCapacityReservationConfig] + ] = None, ) -> bool: """Check that a reuse candidate endpoint matches the requested deploy config. A source-tag match only proves the endpoint was built from the same model artifacts. Before reusing it, confirm the runtime configuration the caller - requested (container environment variables, instance type, and image URI) - also matches, so a differently-configured request does not silently get an - endpoint that contradicts it. + requested (container environment variables, instance type, image URI, and + capacity reservation) also matches, so a differently-configured request does + not silently get an endpoint that contradicts it. Args: endpoint_name: Name of the candidate endpoint to inspect. instance_type: The instance type requested for this deploy, if any. + capacity_reservation_config: The capacity reservation requested for this + deploy, if any. A candidate matches when its variant has each value the + request sets; keys the request leaves unset are not compared. Returns: True if the candidate matches (or the config cannot be read and reuse @@ -2592,6 +2628,21 @@ def _reused_endpoint_matches_config( and instance_type != variant["InstanceType"] ): return False + # Like InstanceType, the reservation lives on the variant: check it before + # the IC early-return so an On-Demand endpoint is never reused for a + # reserved-capacity request. + requested_reservation = _resolve_capacity_reservation_config( + capacity_reservation_config + ) + if requested_reservation: + existing_reservation = variant.get("CapacityReservationConfig") or {} + # Compare only the keys the request sets: an ARN-only request matches an + # endpoint on that reservation whatever its preference. + if any( + existing_reservation.get(key) != value + for key, value in requested_reservation.items() + ): + return False model_name = variant.get("ModelName") if not model_name: # IC-based endpoint: no ModelName to validate container config @@ -3984,6 +4035,10 @@ def _deploy_core_endpoint(self, **kwargs): inference_ami_version = kwargs.get( "inference_ami_version", getattr(self, "inference_ami_version", None) ) + # Resolve (and validate) before any resources are created. + capacity_reservation_config = _resolve_capacity_reservation_config( + kwargs.get("capacity_reservation_config", None) + ) serializer = kwargs.get("serializer", None) deserializer = kwargs.get("deserializer", None) @@ -4047,6 +4102,12 @@ def _deploy_core_endpoint(self, **kwargs): "serverless_inference_config needs to be a ServerlessInferenceConfig object" ) + if is_serverless and capacity_reservation_config: + raise ValueError( + "capacity_reservation_config is not supported for serverless inference. " + "Deploy to an instance-based endpoint to use reserved capacity." + ) + if self._is_sharded_model: if endpoint_type != EndpointType.INFERENCE_COMPONENT_BASED: logger.warning( @@ -4158,7 +4219,15 @@ def _deploy_core_endpoint(self, **kwargs): else: managed_instance_scaling_config["MinInstanceCount"] = initial_instance_count - if not self.sagemaker_session.endpoint_in_service_or_not(self.endpoint_name): + endpoint_exists = self.sagemaker_session.endpoint_in_service_or_not(self.endpoint_name) + if endpoint_exists and capacity_reservation_config: + logger.warning( + "capacity_reservation_config is ignored because endpoint %s already " + "exists; the inference component is added to the endpoint's existing " + "capacity.", + self.endpoint_name, + ) + if not endpoint_exists: production_variant = session_helper.production_variant( instance_type=instance_type, initial_instance_count=initial_instance_count, @@ -4168,6 +4237,7 @@ def _deploy_core_endpoint(self, **kwargs): managed_instance_scaling=managed_instance_scaling_config, routing_config=routing_config, inference_ami_version=inference_ami_version, + capacity_reservation_config=capacity_reservation_config, ) self.sagemaker_session.endpoint_from_production_variants( @@ -4254,6 +4324,7 @@ def _deploy_core_endpoint(self, **kwargs): container_startup_health_check_timeout=container_startup_health_check_timeout, routing_config=routing_config, inference_ami_version=inference_ami_version, + capacity_reservation_config=capacity_reservation_config, ) if endpoint_name: self.endpoint_name = endpoint_name @@ -4286,25 +4357,20 @@ def _deploy_core_endpoint(self, **kwargs): explainer_config_dict = explainer_config._to_request_dict() if update_endpoint: + # Reuse the variant built above so every variant setting (including + # capacity_reservation_config) carries over to the new endpoint config. endpoint_config_name = self.sagemaker_session.create_endpoint_config( - name=self.model_name, - model_name=self.model_name, - initial_instance_count=initial_instance_count, - instance_type=instance_type, - accelerator_type=accelerator_type, + name=name_from_base(base_from_name(self.model_name)), + production_variants=[production_variant], tags=tags, kms_key=kms_key, data_capture_config_dict=data_capture_config_dict, - volume_size=volume_size, - model_data_download_timeout=model_data_download_timeout, - container_startup_health_check_timeout=container_startup_health_check_timeout, explainer_config_dict=explainer_config_dict, async_inference_config_dict=async_inference_config_dict, - serverless_inference_config_dict=serverless_inference_config_dict, - routing_config=routing_config, - inference_ami_version=inference_ami_version, ) - self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name) + self.sagemaker_session.update_endpoint( + self.endpoint_name, endpoint_config_name, wait=wait + ) else: self.sagemaker_session.endpoint_from_production_variants( name=self.endpoint_name, @@ -4415,13 +4481,35 @@ def _update_inference_component( runtime_config=runtime_config, ) - def _deploy_for_ic(self, ic_data: Dict[str, Any], endpoint_name: str, **kwargs) -> Endpoint: - """Deploy/update inference component and return V3 Endpoint.""" + def _deploy_for_ic( + self, + ic_data: Dict[str, Any], + endpoint_name: str, + ic_exists: Optional[bool] = None, + **kwargs, + ) -> Endpoint: + """Deploy/update inference component and return V3 Endpoint. + + ``ic_exists`` lets a caller that already described the component skip a second + DescribeInferenceComponent call; None means look it up here. + """ ic_name = ic_data.get("Name") resource_requirements = ic_data.get("ResourceRequirements") built_model = ic_data.get("Model") - if self._does_ic_exist(ic_name): + if ic_exists is None: + ic_exists = self._does_ic_exist(ic_name) + if ic_exists: + if kwargs.get("capacity_reservation_config") is not None: + # The reservation is bound to the endpoint's production variant, not to the + # inference component, so an in-place IC update cannot apply or change it. + logger.warning( + "capacity_reservation_config is ignored when updating existing inference " + "component %s; capacity is reserved on the endpoint %s, which is not " + "recreated.", + ic_name, + endpoint_name, + ) # Update existing IC self._update_inference_component(ic_name, resource_requirements, **kwargs) # Return existing endpoint @@ -5828,6 +5916,9 @@ def _deploy_recommendation( tags: Optional[List[Dict[str, str]]], wait: bool, auto_approve: bool = False, + capacity_reservation_config: Optional[ + Union[Dict[str, Any], ProductionVariantCapacityReservationConfig] + ] = None, ) -> Endpoint: """Deploy a row from ``self._recommendation_job`` to a live endpoint. @@ -6017,6 +6108,9 @@ def _deploy_recommendation( variant_kwargs["routing_config"] = _ProductionVariantRoutingConfig( routing_strategy="LEAST_OUTSTANDING_REQUESTS" ) + capacity_reservation_shape = _capacity_reservation_config_shape(capacity_reservation_config) + if capacity_reservation_shape is not None: + variant_kwargs["capacity_reservation_config"] = capacity_reservation_shape production_variant = _ProductionVariant(**variant_kwargs) EndpointConfig.create( @@ -6084,6 +6178,9 @@ def deploy( role: Optional[str] = None, model_name: Optional[str] = None, endpoint_config_name: Optional[str] = None, + capacity_reservation_config: Optional[ + Union[Dict[str, Any], ProductionVariantCapacityReservationConfig] + ] = None, **kwargs, ) -> Union[Endpoint, LocalEndpoint, Transformer]: """Deploy the built model to an ``Endpoint``. @@ -6159,6 +6256,19 @@ def deploy( endpoint_config_name (str, optional): Recommendation deploy only. Name for the created EndpointConfig; auto-generated if omitted. Ignored when deploying a normally-built model. + capacity_reservation_config (Union[dict, + sagemaker.core.shapes.ProductionVariantCapacityReservationConfig], optional): + Launches the endpoint's instances on reserved capacity instead of the + On-Demand pool. ``ml_reservation_arn`` is either the ARN of a SageMaker + training plan purchased with target resource ``endpoint`` (Flexible Training + Plans), or the ARN of an ML reservation backed by On-Demand Capacity + Reservations (ODCRs). With ``capacity_reservation_preference= + "capacity-reservations-only"``, instances launch only into the reservation + and the deployment fails if it has no free capacity; this setting is + required for training plans. A dict must use the API keys + ``CapacityReservationPreference`` and ``MlReservationArn``. The + ``instance_type`` must match the reservation's instance type. Not supported + with serverless inference. (Default: None). Returns: Union[Endpoint, LocalEndpoint, Transformer]: A ``sagemaker.core.resources.Endpoint`` resource representing the deployed endpoint, a ``LocalEndpoint`` for local mode, @@ -6175,7 +6285,23 @@ def deploy( >>> model = model_builder.build() # Creates Model resource >>> endpoint = model_builder.deploy(endpoint_name="my-endpoint") # Creates Endpoint resource >>> result = endpoint.invoke(data=input_data) # Make predictions + + Deploy onto a training plan reserved for endpoints: + + >>> from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig + >>> endpoint = model_builder.deploy( + ... instance_type="ml.p5.48xlarge", + ... capacity_reservation_config=ProductionVariantCapacityReservationConfig( + ... capacity_reservation_preference="capacity-reservations-only", + ... ml_reservation_arn=training_plan_arn, + ... ), + ... ) """ + if capacity_reservation_config is not None: + # Validate up front, then carry it to every deploy path through kwargs. + _resolve_capacity_reservation_config(capacity_reservation_config) + kwargs["capacity_reservation_config"] = capacity_reservation_config + if hasattr(self, "_deployed") and self._deployed: logger.warning( "ModelBuilder.deploy() has already been called. " @@ -6233,6 +6359,7 @@ def deploy( tags=kwargs.get("tags"), wait=wait, auto_approve=auto_approve, + capacity_reservation_config=capacity_reservation_config, ) if not hasattr(self, "built_model") and not hasattr(self, "_deployables"): @@ -6262,7 +6389,10 @@ def deploy( # build() does not look for or cache an endpoint. if reuse_resources and not is_inference_component_deploy: requested_instance_type = instance_type or self.instance_type - reusable_endpoint = self._find_reusable_endpoint(instance_type=requested_instance_type) + reusable_endpoint = self._find_reusable_endpoint( + instance_type=requested_instance_type, + capacity_reservation_config=capacity_reservation_config, + ) if reusable_endpoint: if endpoint_name and endpoint_name != reusable_endpoint: logger.warning( @@ -6392,8 +6522,41 @@ def deploy( if hasattr(self, "_deployables") and self._deployables: endpoints = [] + # The reservation belongs to the endpoint, which only a new component's deploy can + # create. Pass the config to the first component that doesn't exist yet; existing + # components are updated in place and later ones join that endpoint. + pending_capacity_kwargs = {} + if kwargs.get("capacity_reservation_config") is not None: + pending_capacity_kwargs["capacity_reservation_config"] = kwargs[ + "capacity_reservation_config" + ] + orchestrator_kwargs = { + k: v for k, v in kwargs.items() if k != "capacity_reservation_config" + } for ic in self._deployables.get("InferenceComponents", []): - endpoints.append(self._deploy_for_ic(ic_data=ic, endpoint_name=endpoint_name)) + # Only look the component up here when a reservation is pending; the result is + # passed on so _deploy_for_ic doesn't describe it again. + ic_exists = self._does_ic_exist(ic.get("Name")) if pending_capacity_kwargs else None + creates_component = bool(pending_capacity_kwargs) and not ic_exists + endpoints.append( + self._deploy_for_ic( + ic_data=ic, + endpoint_name=endpoint_name, + ic_exists=ic_exists, + **(pending_capacity_kwargs if creates_component else {}), + ) + ) + if creates_component: + pending_capacity_kwargs = {} + if pending_capacity_kwargs and not self._deployables.get("CustomOrchestrator"): + logger.warning( + "capacity_reservation_config is ignored because every inference component " + "already exists on endpoint %s; capacity is reserved when the endpoint is " + "created.", + endpoint_name, + ) + # The orchestrator gets the config only if no component created the endpoint. + orchestrator_kwargs.update(pending_capacity_kwargs) # Handle custom orchestrator if present if self._deployables.get("CustomOrchestrator"): @@ -6420,7 +6583,7 @@ def deploy( "endpoint_type": EndpointType.INFERENCE_COMPONENT_BASED, } - deploy_kwargs.update(kwargs) + deploy_kwargs.update(orchestrator_kwargs) endpoints.append(self._deploy(**deploy_kwargs)) elif custom_orchestrator["Mode"] == "InferenceComponent": logger.info( @@ -6435,7 +6598,7 @@ def deploy( initial_instance_count=custom_orchestrator_initial_instance_count or initial_instance_count, endpoint_name=endpoint_name, - **kwargs, + **orchestrator_kwargs, ) ) @@ -6475,6 +6638,9 @@ def _deploy_model_customization( # component deployment; otherwise the model is placed directly on the # production variant. is_ic_deploy = isinstance(inference_config, ResourceRequirements) + capacity_reservation_shape = _capacity_reservation_config_shape( + kwargs.get("capacity_reservation_config") + ) # Nova models without IC resources use the direct model-on-variant path. # Nova models WITH a ResourceRequirements inference_config fall through to @@ -6487,6 +6653,7 @@ def _deploy_model_customization( endpoint_name=endpoint_name, initial_instance_count=initial_instance_count, wait=kwargs.get("wait", True), + capacity_reservation_config=capacity_reservation_shape, ) # The model package may be absent (e.g. a Nova CPTTrainer or raw-S3 @@ -6509,6 +6676,7 @@ def _deploy_model_customization( model_name=self.built_model.model_name, instance_type=self.instance_type, initial_instance_count=initial_instance_count or 1, + capacity_reservation_config=capacity_reservation_shape, ) ], ) @@ -6594,6 +6762,7 @@ def _deploy_model_customization( variant_name=endpoint_name, instance_type=self.instance_type, initial_instance_count=initial_instance_count or 1, + capacity_reservation_config=capacity_reservation_shape, ) ], execution_role_arn=self.role_arn, @@ -6615,6 +6784,13 @@ def _deploy_model_customization( ) endpoint.wait_for_status("InService") else: + if capacity_reservation_shape is not None: + logger.warning( + "capacity_reservation_config is ignored because endpoint %s already " + "exists; the inference component is added to the endpoint's existing " + "capacity.", + endpoint_name, + ) endpoint = Endpoint.get(endpoint_name=endpoint_name) if peft_type == "LORA": @@ -6800,6 +6976,7 @@ def _deploy_nova_model( endpoint_name: str, initial_instance_count: int = 1, wait: bool = True, + capacity_reservation_config: Optional[ProductionVariantCapacityReservationConfig] = None, ) -> Endpoint: """Deploy a Nova model directly to an endpoint without inference components. @@ -6819,6 +6996,7 @@ def _deploy_nova_model( model_name=self.built_model.model_name, instance_type=self.instance_type, initial_instance_count=initial_instance_count, + capacity_reservation_config=capacity_reservation_config, ) ], ) diff --git a/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py new file mode 100644 index 0000000000..e5e9cc2e4f --- /dev/null +++ b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py @@ -0,0 +1,798 @@ +"""Unit tests for deploying ModelBuilder endpoints onto reserved capacity. + +Covers ``capacity_reservation_config`` on ``ModelBuilder.deploy()``, which binds the +endpoint's production variant to a SageMaker training plan (Flexible Training Plans) +or an ODCR-backed ML reservation. +""" + +import inspect +import logging +import unittest +from unittest.mock import Mock, patch + +import pytest + +from sagemaker.core.enums import EndpointType +from sagemaker.core.helper.session_helper import Session +from sagemaker.core.inference_config import ResourceRequirements +from sagemaker.core.resources import Endpoint +from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig +from sagemaker.serve.model_builder import ModelBuilder, _capacity_reservation_config_shape +from sagemaker.serve.mode.function_pointers import Mode +from sagemaker.core.inference_config import ServerlessInferenceConfig + +ROLE_ARN = "arn:aws:iam::123456789012:role/TestRole" +PLAN_ARN = "arn:aws:sagemaker:us-west-2:123456789012:training-plan/p5-for-endpoints" +ML_RESERVATION_ARN = "arn:aws:sagemaker:us-west-2:123456789012:ml-reservation/p5-reserved" +ML_RESERVATION_WITH_PREFERENCE = { + "MlReservationArn": ML_RESERVATION_ARN, + "CapacityReservationPreference": "capacity-reservations-only", +} +EXPECTED_REQUEST = { + "CapacityReservationPreference": "capacity-reservations-only", + "MlReservationArn": PLAN_ARN, +} + + +def _plan_config(): + return ProductionVariantCapacityReservationConfig( + capacity_reservation_preference="capacity-reservations-only", + ml_reservation_arn=PLAN_ARN, + ) + + +class _CoreEndpointTestBase(unittest.TestCase): + def setUp(self): + self.mock_session = Mock() + self.mock_session.boto_region_name = "us-west-2" + self.mock_session.config = {} + self.mock_session.sagemaker_config = {} + self.mock_session.boto_session = Mock() + self.mock_session.boto_session.region_name = "us-west-2" + self.mock_session.endpoint_in_service_or_not = Mock(return_value=False) + self.mock_session.create_endpoint_config = Mock(return_value="new-endpoint-config") + + patches = [ + patch("sagemaker.serve.model_builder.Endpoint.get", return_value=Mock(spec=Endpoint)), + patch.object(ModelBuilder, "_wait_for_endpoint"), + patch.object( + ModelBuilder, "_resolve_inference_ami_version", side_effect=lambda **kw: None + ), + ] + for p in patches: + p.start() + self.addCleanup(p.stop) + + def _builder(self): + builder = ModelBuilder( + model=Mock(), + role_arn=ROLE_ARN, + sagemaker_session=self.mock_session, + mode=Mode.SAGEMAKER_ENDPOINT, + ) + builder.built_model = Mock() + builder.built_model.model_name = "test-model" + builder.model_name = "test-model" + builder.inference_component_name = None # normally set during build() + return builder + + def _deployed_variant(self): + call = self.mock_session.endpoint_from_production_variants.call_args + return call.kwargs["production_variants"][0] + + +class TestDeployCoreEndpointCapacityReservation(_CoreEndpointTestBase): + def test_model_based_endpoint_gets_capacity_reservation_config(self): + self._builder()._deploy_core_endpoint( + instance_type="ml.p5.48xlarge", + initial_instance_count=1, + endpoint_name="ftp-endpoint", + capacity_reservation_config=_plan_config(), + ) + + variant = self._deployed_variant() + assert variant["InstanceType"] == "ml.p5.48xlarge" + assert variant["CapacityReservationConfig"] == EXPECTED_REQUEST + + def test_dict_config_is_accepted(self): + self._builder()._deploy_core_endpoint( + instance_type="ml.p5.48xlarge", + initial_instance_count=1, + endpoint_name="ml-reservation-endpoint", + capacity_reservation_config={"MlReservationArn": ML_RESERVATION_ARN}, + ) + + assert self._deployed_variant()["CapacityReservationConfig"] == { + "MlReservationArn": ML_RESERVATION_ARN + } + + def test_training_plan_without_preference_gets_reservations_only(self): + self._builder()._deploy_core_endpoint( + instance_type="ml.p5.48xlarge", + initial_instance_count=1, + endpoint_name="ftp-endpoint", + capacity_reservation_config={"MlReservationArn": PLAN_ARN}, + ) + + assert self._deployed_variant()["CapacityReservationConfig"] == EXPECTED_REQUEST + + def test_no_config_leaves_variant_unchanged(self): + self._builder()._deploy_core_endpoint( + instance_type="ml.m5.large", + initial_instance_count=1, + endpoint_name="plain-endpoint", + ) + + assert "CapacityReservationConfig" not in self._deployed_variant() + + def test_inference_component_endpoint_gets_capacity_reservation_config(self): + resources = Mock(spec=ResourceRequirements) + resources.num_cpus = None + resources.copy_count = 1 + resources.get_compute_resource_requirements.return_value = { + "NumberOfAcceleratorDevicesRequired": 8 + } + + self._builder()._deploy_core_endpoint( + instance_type="ml.p5.48xlarge", + initial_instance_count=1, + endpoint_name="ftp-ic-endpoint", + endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED, + resources=resources, + capacity_reservation_config=_plan_config(), + ) + + variant = self._deployed_variant() + assert "ModelName" not in variant + assert variant["CapacityReservationConfig"] == EXPECTED_REQUEST + self.mock_session.create_inference_component.assert_called_once() + + def test_update_endpoint_creates_config_from_variant(self): + self._builder()._deploy_core_endpoint( + instance_type="ml.p5.48xlarge", + initial_instance_count=2, + endpoint_name="existing-endpoint", + update_endpoint=True, + capacity_reservation_config=_plan_config(), + ) + + call = self.mock_session.create_endpoint_config.call_args + variant = call.kwargs["production_variants"][0] + assert variant["ModelName"] == "test-model" + assert variant["InitialInstanceCount"] == 2 + assert variant["CapacityReservationConfig"] == EXPECTED_REQUEST + assert call.kwargs["name"].startswith("test-model") + self.mock_session.update_endpoint.assert_called_once_with( + "existing-endpoint", "new-endpoint-config", wait=True + ) + self.mock_session.endpoint_from_production_variants.assert_not_called() + + def test_update_endpoint_honors_wait_false(self): + self._builder()._deploy_core_endpoint( + instance_type="ml.p5.48xlarge", + initial_instance_count=1, + endpoint_name="existing-endpoint", + update_endpoint=True, + wait=False, + capacity_reservation_config=_plan_config(), + ) + + self.mock_session.update_endpoint.assert_called_once_with( + "existing-endpoint", "new-endpoint-config", wait=False + ) + + def test_update_endpoint_keeps_serverless_and_accelerator_settings(self): + self._builder()._deploy_core_endpoint( + serverless_inference_config=ServerlessInferenceConfig(), + endpoint_name="existing-endpoint", + update_endpoint=True, + ) + serverless_variant = self.mock_session.create_endpoint_config.call_args.kwargs[ + "production_variants" + ][0] + assert "ServerlessConfig" in serverless_variant + assert "InstanceType" not in serverless_variant + + self._builder()._deploy_core_endpoint( + instance_type="ml.m5.large", + initial_instance_count=1, + accelerator_type="ml.eia2.medium", + endpoint_name="existing-endpoint", + update_endpoint=True, + ) + accelerator_variant = self.mock_session.create_endpoint_config.call_args.kwargs[ + "production_variants" + ][0] + assert accelerator_variant["AcceleratorType"] == "ml.eia2.medium" + + def test_update_endpoint_calls_match_real_session_signatures(self): + """Regression: the update path once passed kwargs create_endpoint_config doesn't take. + + A plain Mock accepts any arguments, so bind each recorded call to the real + Session method signature to catch drift between the caller and the session API. + """ + self._builder()._deploy_core_endpoint( + instance_type="ml.m5.large", + initial_instance_count=1, + endpoint_name="existing-endpoint", + update_endpoint=True, + ) + + for method_name in ("create_endpoint_config", "update_endpoint"): + call = getattr(self.mock_session, method_name).call_args + assert call is not None, f"{method_name} was not called" + inspect.signature(getattr(Session, method_name)).bind( + self.mock_session, *call.args, **call.kwargs + ) + + def test_existing_endpoint_ic_deploy_warns_capacity_reservation_ignored(self): + self.mock_session.endpoint_in_service_or_not = Mock(return_value=True) + resources = Mock(spec=ResourceRequirements) + resources.num_cpus = None + resources.copy_count = 1 + resources.get_compute_resource_requirements.return_value = {} + + with self.assertLogs("sagemaker", level=logging.WARNING) as logs: + self._builder()._deploy_core_endpoint( + instance_type="ml.p5.48xlarge", + initial_instance_count=1, + endpoint_name="existing-ic-endpoint", + endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED, + resources=resources, + capacity_reservation_config=_plan_config(), + ) + + assert any("already exists" in line for line in logs.output) + self.mock_session.endpoint_from_production_variants.assert_not_called() + self.mock_session.create_inference_component.assert_called_once() + + def test_serverless_with_capacity_reservation_raises(self): + with pytest.raises(ValueError, match="not supported for serverless"): + self._builder()._deploy_core_endpoint( + serverless_inference_config=ServerlessInferenceConfig(), + capacity_reservation_config=_plan_config(), + ) + self.mock_session.endpoint_from_production_variants.assert_not_called() + + def test_invalid_config_raises_before_creating_resources(self): + with pytest.raises(ValueError, match="Unsupported capacity_reservation_config keys"): + self._builder()._deploy_core_endpoint( + instance_type="ml.p5.48xlarge", + initial_instance_count=1, + capacity_reservation_config={"ml_reservation_arn": PLAN_ARN}, + ) + self.mock_session.endpoint_from_production_variants.assert_not_called() + + +class TestDeployForwardsCapacityReservation(unittest.TestCase): + def setUp(self): + self.mock_session = Mock() + self.mock_session.boto_region_name = "us-west-2" + self.mock_session.config = {} + self.mock_session.sagemaker_config = {} + self.mock_session.boto_session = Mock() + self.mock_session.boto_session.region_name = "us-west-2" + + def _builder(self): + builder = ModelBuilder( + model=Mock(), + role_arn=ROLE_ARN, + sagemaker_session=self.mock_session, + mode=Mode.SAGEMAKER_ENDPOINT, + ) + builder.built_model = Mock() + return builder + + @patch.object(ModelBuilder, "_resolve_model_source_id", return_value=None) + @patch.object(ModelBuilder, "_is_model_customization", return_value=False) + @patch.object(ModelBuilder, "_deploy", return_value=Mock(spec=Endpoint)) + def test_deploy_forwards_config_to_deploy_path(self, mock_deploy, *_): + config = _plan_config() + + self._builder().deploy(instance_type="ml.p5.48xlarge", capacity_reservation_config=config) + + assert mock_deploy.call_args.kwargs["capacity_reservation_config"] is config + + @patch.object(ModelBuilder, "_resolve_model_source_id", return_value=None) + @patch.object(ModelBuilder, "_is_model_customization", return_value=False) + @patch.object(ModelBuilder, "_does_ic_exist", return_value=False) + @patch.object(ModelBuilder, "_deploy_for_ic", return_value=Mock(spec=Endpoint)) + def test_deployables_ic_path_forwards_config(self, mock_deploy_for_ic, *_): + config = _plan_config() + builder = self._builder() + builder._deployables = { + "InferenceComponents": [ + {"Name": "ic-1", "ResourceRequirements": Mock(), "Model": Mock()}, + {"Name": "ic-2", "ResourceRequirements": Mock(), "Model": Mock()}, + ] + } + + builder.deploy( + endpoint_name="multi-ic-endpoint", + instance_type="ml.p5.48xlarge", + capacity_reservation_config=config, + ) + + # Only the first call creates the endpoint, so only it gets the reservation. + first_call, second_call = mock_deploy_for_ic.call_args_list + assert first_call.kwargs["capacity_reservation_config"] is config + assert "capacity_reservation_config" not in second_call.kwargs + + @patch.object(ModelBuilder, "_resolve_model_source_id", return_value=None) + @patch.object(ModelBuilder, "_is_model_customization", return_value=False) + @patch.object(ModelBuilder, "_does_ic_exist", return_value=False) + @patch.object(ModelBuilder, "_deploy_for_ic", return_value=Mock(spec=Endpoint)) + def test_orchestrator_after_components_does_not_get_config(self, mock_deploy_for_ic, *_): + config = _plan_config() + builder = self._builder() + builder._deployables = { + "InferenceComponents": [ + {"Name": "ic-1", "ResourceRequirements": Mock(), "Model": Mock()}, + ], + "CustomOrchestrator": {"Mode": "InferenceComponent", "Name": "orch", "Model": Mock()}, + } + + builder.deploy( + endpoint_name="multi-ic-endpoint", + instance_type="ml.p5.48xlarge", + capacity_reservation_config=config, + ) + + component_call, orchestrator_call = mock_deploy_for_ic.call_args_list + assert component_call.kwargs["capacity_reservation_config"] is config + assert "capacity_reservation_config" not in orchestrator_call.kwargs + + @patch.object(ModelBuilder, "_resolve_model_source_id", return_value=None) + @patch.object(ModelBuilder, "_is_model_customization", return_value=False) + @patch.object(ModelBuilder, "_deploy", return_value=Mock(spec=Endpoint)) + def test_orchestrator_that_creates_endpoint_gets_config(self, mock_deploy, *_): + config = _plan_config() + builder = self._builder() + builder._deployables = { + "CustomOrchestrator": {"Mode": "Endpoint", "Name": "orch", "Model": Mock()}, + } + + with patch("sagemaker.serve.model_builder.logger") as mock_logger: + builder.deploy( + endpoint_name="orchestrator-endpoint", + instance_type="ml.p5.48xlarge", + custom_orchestrator_instance_type="ml.p5.48xlarge", + custom_orchestrator_initial_instance_count=1, + capacity_reservation_config=config, + ) + + assert mock_deploy.call_args.kwargs["capacity_reservation_config"] is config + # The orchestrator consumed the config, so no "ignored" warning is logged. + warnings = [str(c.args[0]) for c in mock_logger.warning.call_args_list if c.args] + assert not any("is ignored" in w for w in warnings) + + @patch.object(ModelBuilder, "_resolve_model_source_id", return_value=None) + @patch.object(ModelBuilder, "_is_model_customization", return_value=False) + @patch.object(ModelBuilder, "_deploy_for_ic", return_value=Mock(spec=Endpoint)) + def test_deployables_ic_path_without_config_is_unchanged(self, mock_deploy_for_ic, *_): + builder = self._builder() + builder._deployables = { + "InferenceComponents": [ + {"Name": "ic-1", "ResourceRequirements": Mock(), "Model": Mock()} + ] + } + + builder.deploy(endpoint_name="multi-ic-endpoint", instance_type="ml.p5.48xlarge") + + assert "capacity_reservation_config" not in mock_deploy_for_ic.call_args.kwargs + + @patch.object(ModelBuilder, "_resolve_model_source_id", return_value=None) + @patch.object(ModelBuilder, "_is_model_customization", return_value=False) + @patch.object(ModelBuilder, "_deploy_for_ic", return_value=Mock(spec=Endpoint)) + def test_existing_first_component_does_not_consume_config(self, mock_deploy_for_ic, *_): + config = _plan_config() + builder = self._builder() + builder._deployables = { + "InferenceComponents": [ + {"Name": "existing-ic", "ResourceRequirements": Mock(), "Model": Mock()}, + {"Name": "new-ic", "ResourceRequirements": Mock(), "Model": Mock()}, + ] + } + + with patch.object( + ModelBuilder, "_does_ic_exist", side_effect=lambda name: name == "existing-ic" + ): + builder.deploy( + endpoint_name="multi-ic-endpoint", + instance_type="ml.p5.48xlarge", + capacity_reservation_config=config, + ) + + existing_call, new_call = mock_deploy_for_ic.call_args_list + assert "capacity_reservation_config" not in existing_call.kwargs + assert new_call.kwargs["capacity_reservation_config"] is config + + @patch.object(ModelBuilder, "_resolve_model_source_id", return_value=None) + @patch.object(ModelBuilder, "_is_model_customization", return_value=False) + @patch.object(ModelBuilder, "_does_ic_exist", return_value=True) + @patch.object(ModelBuilder, "_deploy_for_ic", return_value=Mock(spec=Endpoint)) + def test_all_components_existing_warns_config_ignored(self, mock_deploy_for_ic, *_): + builder = self._builder() + builder._deployables = { + "InferenceComponents": [ + {"Name": "existing-ic", "ResourceRequirements": Mock(), "Model": Mock()} + ] + } + + with self.assertLogs("sagemaker", level=logging.WARNING) as logs: + builder.deploy( + endpoint_name="multi-ic-endpoint", + instance_type="ml.p5.48xlarge", + capacity_reservation_config=_plan_config(), + ) + + assert "capacity_reservation_config" not in mock_deploy_for_ic.call_args.kwargs + assert any("every inference component already exists" in line for line in logs.output) + + @patch.object(ModelBuilder, "_resolve_model_source_id", return_value=None) + @patch.object(ModelBuilder, "_is_model_customization", return_value=False) + @patch.object(ModelBuilder, "_update_inference_component") + @patch.object(ModelBuilder, "_deploy", return_value=Mock(spec=Endpoint)) + @patch("sagemaker.serve.model_builder.Endpoint.get", return_value=Mock(spec=Endpoint)) + def test_each_component_is_described_once(self, *_): + builder = self._builder() + builder._deployables = { + "InferenceComponents": [ + {"Name": "existing-ic", "ResourceRequirements": Mock(), "Model": Mock()}, + {"Name": "new-ic", "ResourceRequirements": Mock(), "Model": Mock()}, + ] + } + + with patch.object( + ModelBuilder, "_does_ic_exist", side_effect=lambda name: name == "existing-ic" + ) as mock_exists: + builder.deploy( + endpoint_name="multi-ic-endpoint", + instance_type="ml.p5.48xlarge", + capacity_reservation_config=_plan_config(), + ) + + assert [c.args[0] for c in mock_exists.call_args_list] == ["existing-ic", "new-ic"] + + def _reuse_candidate(self, variant): + client = self.mock_session.sagemaker_client + client.describe_endpoint.return_value = {"EndpointConfigName": "cfg"} + client.describe_endpoint_config.return_value = {"ProductionVariants": [variant]} + client.describe_model.return_value = { + "PrimaryContainer": {"Environment": {}, "Image": "img"} + } + + @patch( + "sagemaker.serve.model_builder.find_existing_sagemaker_endpoint", + return_value="arn:aws:sagemaker:us-west-2:123456789012:endpoint/existing-endpoint", + ) + @patch.object(ModelBuilder, "_resolve_model_source_id", return_value="model-source") + @patch.object(ModelBuilder, "_is_model_customization", return_value=False) + @patch.object(ModelBuilder, "_deploy", return_value=Mock(spec=Endpoint)) + def test_deploy_reuse_skips_on_demand_endpoint_for_reservation_request(self, mock_deploy, *_): + self._reuse_candidate({"ModelName": "m", "InstanceType": "ml.p5.48xlarge"}) + + self._builder().deploy( + instance_type="ml.p5.48xlarge", + reuse_resources=True, + capacity_reservation_config=_plan_config(), + ) + + mock_deploy.assert_called_once() + + @patch("sagemaker.serve.model_builder.Endpoint.get", return_value=Mock(spec=Endpoint)) + @patch( + "sagemaker.serve.model_builder.find_existing_sagemaker_endpoint", + return_value="arn:aws:sagemaker:us-west-2:123456789012:endpoint/existing-endpoint", + ) + @patch.object(ModelBuilder, "_resolve_model_source_id", return_value="model-source") + @patch.object(ModelBuilder, "_is_model_customization", return_value=False) + @patch.object(ModelBuilder, "_deploy", return_value=Mock(spec=Endpoint)) + def test_deploy_reuses_endpoint_on_same_reservation(self, mock_deploy, *_): + # An ML reservation ARN, because a training plan ARN gets a preference added before matching. + self._reuse_candidate( + { + "ModelName": "m", + "InstanceType": "ml.p5.48xlarge", + "CapacityReservationConfig": ML_RESERVATION_WITH_PREFERENCE, + } + ) + + self._builder().deploy( + instance_type="ml.p5.48xlarge", + reuse_resources=True, + capacity_reservation_config={"MlReservationArn": ML_RESERVATION_ARN}, + ) + + mock_deploy.assert_not_called() + + @patch.object(ModelBuilder, "_deploy") + def test_deploy_rejects_invalid_config_up_front(self, mock_deploy): + with pytest.raises(TypeError): + self._builder().deploy( + instance_type="ml.p5.48xlarge", capacity_reservation_config=PLAN_ARN + ) + mock_deploy.assert_not_called() + + +def test_capacity_reservation_config_does_not_shift_positional_parameters(): + params = list(inspect.signature(ModelBuilder.deploy).parameters) + # Added after the last pre-existing named parameter, so positional callers are unaffected. + assert params.index("capacity_reservation_config") == params.index("endpoint_config_name") + 1 + assert params.index("use_recommendation") == params.index("reuse_resources") + 1 + + +class TestReuseMatchesCapacityReservation: + def _builder(self, variant): + session = Mock() + session.boto_session = Mock() + builder = ModelBuilder(model=Mock(), role_arn=ROLE_ARN, sagemaker_session=session) + client = Mock() + client.describe_endpoint.return_value = {"EndpointConfigName": "cfg"} + client.describe_endpoint_config.return_value = {"ProductionVariants": [variant]} + client.describe_model.return_value = { + "PrimaryContainer": {"Environment": {}, "Image": "img"} + } + builder.sagemaker_session.sagemaker_client = client + return builder + + def test_on_demand_endpoint_is_not_reused_for_reservation_request(self): + builder = self._builder({"ModelName": "m", "InstanceType": "ml.p5.48xlarge"}) + assert ( + builder._reused_endpoint_matches_config( + "ep", instance_type="ml.p5.48xlarge", capacity_reservation_config=_plan_config() + ) + is False + ) + + def test_inference_component_endpoint_on_other_reservation_is_not_reused(self): + other = { + "MlReservationArn": PLAN_ARN.replace("p5-for-endpoints", "other-plan"), + "CapacityReservationPreference": "capacity-reservations-only", + } + # No ModelName: inference-component endpoint, checked before the IC early return. + builder = self._builder( + {"InstanceType": "ml.p5.48xlarge", "CapacityReservationConfig": other} + ) + assert ( + builder._reused_endpoint_matches_config( + "ep", instance_type="ml.p5.48xlarge", capacity_reservation_config=_plan_config() + ) + is False + ) + + def test_endpoint_on_same_reservation_is_reused(self): + # DescribeEndpointConfig may return extra utilization fields; only ARN and preference matter. + existing = dict(EXPECTED_REQUEST, TotalInstanceCount=2) + builder = self._builder( + { + "ModelName": "m", + "InstanceType": "ml.p5.48xlarge", + "CapacityReservationConfig": existing, + } + ) + assert ( + builder._reused_endpoint_matches_config( + "ep", instance_type="ml.p5.48xlarge", capacity_reservation_config=_plan_config() + ) + is True + ) + + def test_arn_only_request_matches_endpoint_with_a_preference(self): + # An ML reservation ARN, because a training plan ARN gets a preference added before matching. + builder = self._builder( + { + "ModelName": "m", + "InstanceType": "ml.p5.48xlarge", + "CapacityReservationConfig": ML_RESERVATION_WITH_PREFERENCE, + } + ) + assert ( + builder._reused_endpoint_matches_config( + "ep", + instance_type="ml.p5.48xlarge", + capacity_reservation_config={"MlReservationArn": ML_RESERVATION_ARN}, + ) + is True + ) + + def test_no_reservation_requested_keeps_existing_behavior(self): + builder = self._builder( + { + "ModelName": "m", + "InstanceType": "ml.p5.48xlarge", + "CapacityReservationConfig": EXPECTED_REQUEST, + } + ) + assert builder._reused_endpoint_matches_config("ep", instance_type="ml.p5.48xlarge") is True + + +class TestCapacityReservationShapeHelper: + def test_dict_becomes_shape(self): + shape = _capacity_reservation_config_shape(dict(EXPECTED_REQUEST)) + assert isinstance(shape, ProductionVariantCapacityReservationConfig) + assert shape.capacity_reservation_preference == "capacity-reservations-only" + assert shape.ml_reservation_arn == PLAN_ARN + + def test_none_and_empty_return_none(self): + assert _capacity_reservation_config_shape(None) is None + assert _capacity_reservation_config_shape({}) is None + + +class TestNovaDeployCapacityReservation(unittest.TestCase): + @patch("sagemaker.serve.model_builder.Endpoint.create") + @patch("sagemaker.serve.model_builder.EndpointConfig.create") + def test_nova_variant_gets_capacity_reservation_config( + self, mock_config_create, mock_endpoint_create + ): + builder = ModelBuilder.__new__(ModelBuilder) + builder.built_model = Mock() + builder.built_model.model_name = "nova-model" + builder.instance_type = "ml.p5.48xlarge" + builder._tags = None + builder._base_model_name = Mock(return_value="nova-lite") + builder._fetch_model_package = Mock(return_value=None) + mock_endpoint_create.return_value = Mock(spec=Endpoint) + + builder._deploy_nova_model( + endpoint_name="nova-endpoint", + wait=False, + capacity_reservation_config=_plan_config(), + ) + + variant = mock_config_create.call_args.kwargs["production_variants"][0] + assert variant.capacity_reservation_config.ml_reservation_arn == PLAN_ARN + + +class TestDeployForIcExistingComponent(unittest.TestCase): + @patch("sagemaker.serve.model_builder.Endpoint.get", return_value=Mock(spec=Endpoint)) + def test_updating_existing_ic_warns_capacity_reservation_ignored(self, _): + builder = ModelBuilder.__new__(ModelBuilder) + builder.sagemaker_session = Mock() + builder.region = "us-west-2" + builder._does_ic_exist = Mock(return_value=True) + builder._update_inference_component = Mock() + + with self.assertLogs("sagemaker", level=logging.WARNING) as logs: + builder._deploy_for_ic( + ic_data={"Name": "ic-1", "ResourceRequirements": Mock(), "Model": Mock()}, + endpoint_name="existing-endpoint", + capacity_reservation_config=_plan_config(), + ) + + assert any("existing inference component ic-1" in line for line in logs.output) + builder._update_inference_component.assert_called_once() + + +class TestModelCustomizationExistingEndpoint(unittest.TestCase): + @patch("sagemaker.serve.model_builder.Endpoint.get", side_effect=RuntimeError("stop")) + @patch("sagemaker.serve.model_builder.EndpointConfig.create") + def test_existing_endpoint_warns_capacity_reservation_ignored(self, mock_config_create, _): + builder = ModelBuilder.__new__(ModelBuilder) + builder._is_nova_model = Mock(return_value=False) + builder._fetch_model_package = Mock(return_value=None) + builder._prepare_reused_model_customization_deployment_state = Mock() + builder._cached_compute_requirements = Mock() + builder.built_model = Mock(enable_network_isolation=False) + builder._enable_network_isolation = False + builder._does_endpoint_exist = Mock(return_value=True) + + # Endpoint.get raises right after the existing-endpoint warning, so the + # inference-component deploy logic that follows isn't exercised here. + with self.assertLogs("sagemaker", level=logging.WARNING) as logs: + with pytest.raises(RuntimeError, match="stop"): + builder._deploy_model_customization( + endpoint_name="existing-endpoint", + capacity_reservation_config=_plan_config(), + ) + + assert any("endpoint existing-endpoint already exists" in line for line in logs.output) + mock_config_create.assert_not_called() + + +class TestModelCustomizationVariantsGetCapacityReservation(unittest.TestCase): + def _builder(self, model_package=None): + builder = ModelBuilder.__new__(ModelBuilder) + builder._is_nova_model = Mock(return_value=False) + builder._fetch_model_package = Mock(return_value=model_package) + builder._fetch_peft = Mock(return_value=None) + builder._prepare_reused_model_customization_deployment_state = Mock() + builder._cached_compute_requirements = Mock() + builder.built_model = Mock(enable_network_isolation=False, model_name="custom-model") + builder._enable_network_isolation = False + builder._does_endpoint_exist = Mock(return_value=False) + builder.instance_type = "ml.p5.48xlarge" + builder.role_arn = ROLE_ARN + builder._tags = None + return builder + + @staticmethod + def _variant(mock_config_create): + return mock_config_create.call_args.kwargs["production_variants"][0] + + # Endpoint.create raises "stop" right after EndpointConfig.create, which is what's under test. + @patch("sagemaker.serve.model_builder.Endpoint.create", side_effect=RuntimeError("stop")) + @patch("sagemaker.serve.model_builder.EndpointConfig.create") + def test_new_endpoint_variant_gets_config(self, mock_config_create, _): + with pytest.raises(RuntimeError, match="stop"): + self._builder()._deploy_model_customization( + endpoint_name="custom-endpoint", capacity_reservation_config=_plan_config() + ) + + assert ( + self._variant(mock_config_create).capacity_reservation_config.ml_reservation_arn + == PLAN_ARN + ) + + @patch( + "sagemaker.serve.utils.model_package_utils.is_restricted_model_package", return_value=True + ) + @patch("sagemaker.serve.model_builder.Endpoint.create", side_effect=RuntimeError("stop")) + @patch("sagemaker.serve.model_builder.EndpointConfig.create") + def test_restricted_model_package_variant_gets_config(self, mock_config_create, *_): + with pytest.raises(RuntimeError, match="stop"): + self._builder(model_package=Mock())._deploy_model_customization( + endpoint_name="custom-endpoint", capacity_reservation_config=_plan_config() + ) + + variant = self._variant(mock_config_create) + assert variant.model_name == "custom-model" + assert variant.capacity_reservation_config.ml_reservation_arn == PLAN_ARN + + def test_nova_model_is_handed_the_config(self): + builder = self._builder() + builder._is_nova_model = Mock(return_value=True) + builder._deploy_nova_model = Mock(return_value=Mock(spec=Endpoint)) + + builder._deploy_model_customization( + endpoint_name="nova-endpoint", capacity_reservation_config=_plan_config() + ) + + shape = builder._deploy_nova_model.call_args.kwargs["capacity_reservation_config"] + assert shape.ml_reservation_arn == PLAN_ARN + + +class TestRecommendationDeployCapacityReservation: + @patch( + "sagemaker.serve.model_builder.resolve_and_validate_role", + return_value="arn:aws:iam::123456789012:role/ambient-execution-role", + ) + def test_recommendation_variant_gets_config(self, _): + from types import SimpleNamespace + from unittest.mock import MagicMock + + session = MagicMock() + session.sagemaker_config = {} + builder = ModelBuilder(sagemaker_session=session) + builder.role_arn = ROLE_ARN + builder._recommendation_job = SimpleNamespace( + recommendations=[ + SimpleNamespace( + model_details=SimpleNamespace( + model_package_arn="arn:aws:sagemaker:us-west-2:1:model-package/p/1", + inference_specification_name="A", + ), + deployment_configuration=SimpleNamespace( + instance_type="ml.p5.48xlarge", instance_count=1, copy_count_per_instance=1 + ), + expected_performance=[], + ) + ], + ai_recommendation_job_status="Completed", + ) + session.sagemaker_client.describe_model_package.return_value = { + "ModelApprovalStatus": "Approved", + "InferenceSpecification": {"Containers": [{"Image": "123.dkr.ecr/x:1"}]}, + } + + with ( + patch("sagemaker.core.resources.Model.create"), + patch("sagemaker.serve.model_builder.EndpointConfig.create") as mock_config_create, + patch("sagemaker.serve.model_builder.Endpoint.create", return_value=MagicMock()), + ): + builder.deploy( + endpoint_name="rec-endpoint", wait=False, capacity_reservation_config=_plan_config() + ) + + variant = mock_config_create.call_args.kwargs["production_variants"][0] + assert variant.capacity_reservation_config.ml_reservation_arn == PLAN_ARN