From 9192bb78ff703ee7342272b85ec61a114be23b74 Mon Sep 17 00:00:00 2001 From: Liam Neal Reilly Date: Tue, 10 Nov 2026 21:45:32 +0000 Subject: [PATCH 1/7] feature: support capacity_reservation_config in ModelBuilder.deploy Add a capacity_reservation_config parameter to ModelBuilder.deploy() so endpoints can launch on reserved capacity: a SageMaker training plan purchased for the "endpoint" target resource (Flexible Training Plans), or an ODCR-backed ML reservation. It accepts a ProductionVariantCapacityReservationConfig shape or a dict with the CreateEndpointConfig keys, and is applied to the production variant on the model-based, inference-component, model-customization, Nova and recommendation deploy paths. session_helper.production_variant() gains the same parameter. Also fix ModelBuilder.deploy(update_endpoint=True), which called Session.create_endpoint_config() with keyword arguments that method does not accept and raised TypeError. It now builds the endpoint config from the same production variant used for new endpoints, so variant settings (including the capacity reservation) carry over to the update. --- .../src/sagemaker/core/common_utils.py | 52 ++++ .../sagemaker/core/helper/session_helper.py | 12 + .../tests/unit/helper/test_session_helper.py | 26 ++ .../tests/unit/test_common_utils.py | 55 ++++ .../src/sagemaker/serve/model_builder.py | 91 ++++++- .../unit/test_deploy_capacity_reservation.py | 242 ++++++++++++++++++ 6 files changed, 467 insertions(+), 11 deletions(-) create mode 100644 sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py diff --git a/sagemaker-core/src/sagemaker/core/common_utils.py b/sagemaker-core/src/sagemaker/core/common_utils.py index b21c05ca16..21f2d6d494 100644 --- a/sagemaker-core/src/sagemaker/core/common_utils.py +++ b/sagemaker-core/src/sagemaker/core/common_utils.py @@ -2129,6 +2129,58 @@ def _resolve_routing_config(routing_config: Optional[Dict[str, Any]]) -> Optiona return None +_CAPACITY_RESERVATION_CONFIG_KEYS = ("CapacityReservationPreference", "MlReservationArn") + + +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. + + 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 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 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__}." + ) + + @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..8fd6a14c2c 100644 --- a/sagemaker-core/tests/unit/test_common_utils.py +++ b/sagemaker-core/tests/unit/test_common_utils.py @@ -2271,6 +2271,61 @@ 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" + + 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_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.PLAN_ARN) + assert _resolve_capacity_reservation_config(config) == {"MlReservationArn": self.PLAN_ARN} + + def test_dict_is_passed_through_without_none_values(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} + + 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..abb148fe98 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. @@ -3984,6 +4001,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 +4068,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( @@ -4168,6 +4195,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 +4282,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,23 +4315,16 @@ 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) else: @@ -5828,6 +5850,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 +6042,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( @@ -6070,6 +6098,9 @@ def deploy( custom_orchestrator_instance_type: str = None, custom_orchestrator_initial_instance_count: int = None, reuse_resources: bool = False, + capacity_reservation_config: Optional[ + Union[Dict[str, Any], ProductionVariantCapacityReservationConfig] + ] = None, # Recommendation-mode kwargs. These take effect only on the # recommendation deploy path (a recommendation job is attached and # use_recommendation is not False); they are ignored when deploying a @@ -6133,6 +6164,19 @@ def deploy( ``ResourceRequirements``, or a ``modelbuilder_list`` build) are not intercepted by this flag — they manage their own reuse by inference component name (create vs. in-place update). + 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). use_recommendation (bool, optional): Controls the recommendation deploy path. None (default) deploys the recommendation when a recommendation job is attached, else the built model. False forces the built-model path even if a @@ -6175,7 +6219,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 +6293,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"): @@ -6475,6 +6536,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 +6551,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 +6574,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 +6660,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, @@ -6800,6 +6867,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 +6887,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..027750c19d --- /dev/null +++ b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py @@ -0,0 +1,242 @@ +"""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 unittest +from unittest.mock import Mock, patch + +import pytest + +from sagemaker.core.enums import EndpointType +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" +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="ftp-endpoint", + capacity_reservation_config={"MlReservationArn": PLAN_ARN}, + ) + + assert self._deployed_variant()["CapacityReservationConfig"] == { + "MlReservationArn": PLAN_ARN + } + + 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" + ) + self.mock_session.endpoint_from_production_variants.assert_not_called() + + 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, "_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() + + +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 From 86e0357cfa5d897239ffb12597003ae7a17bf81f Mon Sep 17 00:00:00 2001 From: Liam Neal Reilly Date: Tue, 10 Nov 2026 22:01:25 +0000 Subject: [PATCH 2/7] fix: forward capacity_reservation_config on the modelbuilder_list inference component path --- .../src/sagemaker/serve/model_builder.py | 9 ++++- .../unit/test_deploy_capacity_reservation.py | 38 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index abb148fe98..b5f31a6df5 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -6453,8 +6453,15 @@ def deploy( if hasattr(self, "_deployables") and self._deployables: endpoints = [] + # Only the capacity reservation is forwarded here; the endpoint created for the + # first inference component must launch on the reserved capacity. + ic_kwargs = {} + if kwargs.get("capacity_reservation_config") is not None: + ic_kwargs["capacity_reservation_config"] = kwargs["capacity_reservation_config"] for ic in self._deployables.get("InferenceComponents", []): - endpoints.append(self._deploy_for_ic(ic_data=ic, endpoint_name=endpoint_name)) + endpoints.append( + self._deploy_for_ic(ic_data=ic, endpoint_name=endpoint_name, **ic_kwargs) + ) # Handle custom orchestrator if present if self._deployables.get("CustomOrchestrator"): diff --git a/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py index 027750c19d..ca0c5c0982 100644 --- a/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py +++ b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py @@ -196,6 +196,44 @@ def test_deploy_forwards_config_to_deploy_path(self, mock_deploy, *_): 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, "_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, + ) + + assert mock_deploy_for_ic.call_count == 2 + for call in mock_deploy_for_ic.call_args_list: + assert 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, "_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, "_deploy") def test_deploy_rejects_invalid_config_up_front(self, mock_deploy): with pytest.raises(TypeError): From 5d6888b75642e0e71bbc13a93317414f24900926 Mon Sep 17 00:00:00 2001 From: Liam Neal Reilly Date: Wed, 11 Nov 2026 20:28:02 +0000 Subject: [PATCH 3/7] change: warn when capacity_reservation_config cannot apply; add update_endpoint regression test Log a warning when capacity_reservation_config is passed but the reservation cannot take effect because the endpoint already exists: when updating an existing inference component, or when adding an inference component to an endpoint that is already InService. The reservation is bound to the endpoint's production variant, so neither path can apply or change it. Add a regression test for the update_endpoint=True fix that binds the recorded create_endpoint_config and update_endpoint calls to the real Session method signatures, so a plain Mock can no longer hide a signature mismatch. --- .../src/sagemaker/serve/model_builder.py | 20 +++++- .../unit/test_deploy_capacity_reservation.py | 64 +++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index b5f31a6df5..971106656c 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -4185,7 +4185,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, @@ -4444,6 +4452,16 @@ def _deploy_for_ic(self, ic_data: Dict[str, Any], endpoint_name: str, **kwargs) built_model = ic_data.get("Model") if self._does_ic_exist(ic_name): + 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 diff --git a/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py index ca0c5c0982..db5cdf7173 100644 --- a/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py +++ b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py @@ -5,12 +5,15 @@ 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 @@ -149,6 +152,47 @@ def test_update_endpoint_creates_config_from_variant(self): ) self.mock_session.endpoint_from_production_variants.assert_not_called() + 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( @@ -278,3 +322,23 @@ def test_nova_variant_gets_capacity_reservation_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() From 09417d103ab4948d1aa36b48e7229d031ebf0be5 Mon Sep 17 00:00:00 2001 From: Liam Neal Reilly Date: Wed, 11 Nov 2026 21:08:59 +0000 Subject: [PATCH 4/7] fix: address review on capacity_reservation_config deploy paths - modelbuilder_list: pass capacity_reservation_config only to the deploy call that creates the endpoint, so later inference components and the custom orchestrator no longer log a false "ignored" warning. - update_endpoint=True: pass wait through to Session.update_endpoint so deploy(update_endpoint=True, wait=False) does not block. - reuse_resources: _reused_endpoint_matches_config compares the variant's MlReservationArn and CapacityReservationPreference, so an On-Demand or differently reserved endpoint is never reused for a reservation request. - model customization: warn when the endpoint already exists and the reservation cannot apply, matching the other deploy paths. - deploy(): move capacity_reservation_config after endpoint_config_name so existing positional arguments keep their positions. --- .../src/sagemaker/serve/model_builder.py | 119 +++++++++---- .../unit/test_deploy_capacity_reservation.py | 166 +++++++++++++++++- 2 files changed, 249 insertions(+), 36 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 971106656c..6a029f1772 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -2536,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. @@ -2561,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 @@ -2573,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 only if its variant uses the same + reservation ARN and preference. Returns: True if the candidate matches (or the config cannot be read and reuse @@ -2609,6 +2628,19 @@ 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 {} + if any( + existing_reservation.get(key) != requested_reservation.get(key) + for key in ("MlReservationArn", "CapacityReservationPreference") + ): + return False model_name = variant.get("ModelName") if not model_name: # IC-based endpoint: no ModelName to validate container config @@ -4334,7 +4366,9 @@ def _deploy_core_endpoint(self, **kwargs): explainer_config_dict=explainer_config_dict, async_inference_config_dict=async_inference_config_dict, ) - 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, @@ -6116,9 +6150,6 @@ def deploy( custom_orchestrator_instance_type: str = None, custom_orchestrator_initial_instance_count: int = None, reuse_resources: bool = False, - capacity_reservation_config: Optional[ - Union[Dict[str, Any], ProductionVariantCapacityReservationConfig] - ] = None, # Recommendation-mode kwargs. These take effect only on the # recommendation deploy path (a recommendation job is attached and # use_recommendation is not False); they are ignored when deploying a @@ -6133,6 +6164,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``. @@ -6182,19 +6216,6 @@ def deploy( ``ResourceRequirements``, or a ``modelbuilder_list`` build) are not intercepted by this flag — they manage their own reuse by inference component name (create vs. in-place update). - 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). use_recommendation (bool, optional): Controls the recommendation deploy path. None (default) deploys the recommendation when a recommendation job is attached, else the built model. False forces the built-model path even if a @@ -6221,6 +6242,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, @@ -6341,7 +6375,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( @@ -6471,15 +6508,26 @@ def deploy( if hasattr(self, "_deployables") and self._deployables: endpoints = [] - # Only the capacity reservation is forwarded here; the endpoint created for the - # first inference component must launch on the reserved capacity. - ic_kwargs = {} + # The reservation belongs to the endpoint, which the first deploy call creates. + # Later calls add components to that endpoint, so they don't get the config + # (they would otherwise warn that it is ignored). + pending_capacity_kwargs = {} if kwargs.get("capacity_reservation_config") is not None: - ic_kwargs["capacity_reservation_config"] = kwargs["capacity_reservation_config"] + 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, **ic_kwargs) + self._deploy_for_ic( + ic_data=ic, endpoint_name=endpoint_name, **pending_capacity_kwargs + ) ) + pending_capacity_kwargs = {} + # The orchestrator only creates the endpoint when no component came first. + orchestrator_kwargs.update(pending_capacity_kwargs) # Handle custom orchestrator if present if self._deployables.get("CustomOrchestrator"): @@ -6506,7 +6554,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( @@ -6521,7 +6569,7 @@ def deploy( initial_instance_count=custom_orchestrator_initial_instance_count or initial_instance_count, endpoint_name=endpoint_name, - **kwargs, + **orchestrator_kwargs, ) ) @@ -6707,6 +6755,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": diff --git a/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py index db5cdf7173..3a7c143f30 100644 --- a/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py +++ b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py @@ -148,10 +148,24 @@ def test_update_endpoint_creates_config_from_variant(self): 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" + "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_calls_match_real_session_signatures(self): """Regression: the update path once passed kwargs create_endpoint_config doesn't take. @@ -259,9 +273,53 @@ def test_deployables_ic_path_forwards_config(self, mock_deploy_for_ic, *_): capacity_reservation_config=config, ) - assert mock_deploy_for_ic.call_count == 2 - for call in mock_deploy_for_ic.call_args_list: - assert call.kwargs["capacity_reservation_config"] is 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, "_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()}, + } + + 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 @patch.object(ModelBuilder, "_resolve_model_source_id", return_value=None) @patch.object(ModelBuilder, "_is_model_customization", return_value=False) @@ -287,6 +345,80 @@ def test_deploy_rejects_invalid_config_up_front(self, mock_deploy): 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_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)) @@ -342,3 +474,29 @@ def test_updating_existing_ic_warns_capacity_reservation_ignored(self, _): 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() From a2020118a78a199dac713ef194ddb46227054421 Mon Sep 17 00:00:00 2001 From: Liam Neal Reilly Date: Thu, 12 Nov 2026 01:10:50 +0000 Subject: [PATCH 5/7] fix: address second review on capacity_reservation_config - Reuse matching compares only the reservation keys the request sets, so an ARN-only request reuses an endpoint on that reservation regardless of its CapacityReservationPreference. - modelbuilder_list: keep the reservation pending until a component that doesn't exist yet is deployed (only that call can create the endpoint), instead of always handing it to the first component. If every component already exists, warn that the reservation is ignored. - Tests through deploy() for endpoint reuse (On-Demand candidate is not reused; same-reservation candidate is), and per-path tests asserting the CapacityReservationConfig on the variant passed to EndpointConfig.create for recommendation, both model customization variants, and the Nova hand-off. --- .../src/sagemaker/serve/model_builder.py | 31 ++- .../unit/test_deploy_capacity_reservation.py | 226 ++++++++++++++++++ 2 files changed, 249 insertions(+), 8 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 6a029f1772..8aa1b1ac8a 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -2636,9 +2636,11 @@ def _reused_endpoint_matches_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) != requested_reservation.get(key) - for key in ("MlReservationArn", "CapacityReservationPreference") + existing_reservation.get(key) != value + for key, value in requested_reservation.items() ): return False model_name = variant.get("ModelName") @@ -6508,9 +6510,9 @@ def deploy( if hasattr(self, "_deployables") and self._deployables: endpoints = [] - # The reservation belongs to the endpoint, which the first deploy call creates. - # Later calls add components to that endpoint, so they don't get the config - # (they would otherwise warn that it is ignored). + # 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[ @@ -6520,13 +6522,26 @@ def deploy( k: v for k, v in kwargs.items() if k != "capacity_reservation_config" } for ic in self._deployables.get("InferenceComponents", []): + creates_component = bool(pending_capacity_kwargs) and not self._does_ic_exist( + ic.get("Name") + ) endpoints.append( self._deploy_for_ic( - ic_data=ic, endpoint_name=endpoint_name, **pending_capacity_kwargs + ic_data=ic, + endpoint_name=endpoint_name, + **(pending_capacity_kwargs if creates_component else {}), ) ) - pending_capacity_kwargs = {} - # The orchestrator only creates the endpoint when no component came first. + 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 diff --git a/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py index 3a7c143f30..a2ab9cfcbb 100644 --- a/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py +++ b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py @@ -256,6 +256,7 @@ def test_deploy_forwards_config_to_deploy_path(self, mock_deploy, *_): @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() @@ -280,6 +281,7 @@ def test_deployables_ic_path_forwards_config(self, mock_deploy_for_ic, *_): @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() @@ -336,6 +338,105 @@ def test_deployables_ic_path_without_config_is_unchanged(self, mock_deploy_for_i 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) + + 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, *_): + self._reuse_candidate( + { + "ModelName": "m", + "InstanceType": "ml.p5.48xlarge", + "CapacityReservationConfig": EXPECTED_REQUEST, + } + ) + + self._builder().deploy( + instance_type="ml.p5.48xlarge", + reuse_resources=True, + capacity_reservation_config={"MlReservationArn": PLAN_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): @@ -408,6 +509,23 @@ def test_endpoint_on_same_reservation_is_reused(self): is True ) + def test_arn_only_request_matches_endpoint_with_a_preference(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", + capacity_reservation_config={"MlReservationArn": PLAN_ARN}, + ) + is True + ) + def test_no_reservation_requested_keeps_existing_behavior(self): builder = self._builder( { @@ -500,3 +618,111 @@ def test_existing_endpoint_warns_capacity_reservation_ignored(self, mock_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 From e9056e8baf0e26fb5d39327b0931affba5218206 Mon Sep 17 00:00:00 2001 From: Liam Neal Reilly Date: Thu, 12 Nov 2026 18:57:05 +0000 Subject: [PATCH 6/7] fix: default training plan preference; address third review - _resolve_capacity_reservation_config sets CapacityReservationPreference="capacity-reservations-only" when a training-plan ARN is given without a preference, the only value SageMaker accepts for training plans. ML reservation ARNs and explicit preferences are passed through as given. - Multi-component deploy describes each component once: the existence check computed for the reservation is passed to _deploy_for_ic. - Test that an orchestrator that creates the endpoint gets the reservation without the "ignored" warning, and that update_endpoint keeps serverless and accelerator variant settings. - Correct the _reused_endpoint_matches_config docstring: a candidate matches when its variant has each value the request sets. --- .../src/sagemaker/core/common_utils.py | 23 ++++- .../tests/unit/test_common_utils.py | 41 ++++++++- .../src/sagemaker/serve/model_builder.py | 30 ++++-- .../unit/test_deploy_capacity_reservation.py | 92 ++++++++++++++++--- 4 files changed, 157 insertions(+), 29 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/common_utils.py b/sagemaker-core/src/sagemaker/core/common_utils.py index 21f2d6d494..4360bfc5e5 100644 --- a/sagemaker-core/src/sagemaker/core/common_utils.py +++ b/sagemaker-core/src/sagemaker/core/common_utils.py @@ -2130,6 +2130,7 @@ def _resolve_routing_config(routing_config: Optional[Dict[str, Any]]) -> Optiona _CAPACITY_RESERVATION_CONFIG_KEYS = ("CapacityReservationPreference", "MlReservationArn") +_CAPACITY_RESERVATIONS_ONLY = "capacity-reservations-only" def _resolve_capacity_reservation_config( @@ -2146,7 +2147,10 @@ def _resolve_capacity_reservation_config( Returns: Optional[Dict[str, Any]]: The ``CapacityReservationConfig`` request dict, or None if - no config was given or it has no fields set. + 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 @@ -2162,7 +2166,7 @@ def _resolve_capacity_reservation_config( from sagemaker.core.utils.utils import serialize if isinstance(capacity_reservation_config, ProductionVariantCapacityReservationConfig): - return serialize(capacity_reservation_config) or None + 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) @@ -2172,7 +2176,7 @@ def _resolve_capacity_reservation_config( 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 resolved or None + return _default_training_plan_preference(resolved or None) raise TypeError( "capacity_reservation_config must be a dict or a " @@ -2181,6 +2185,19 @@ def _resolve_capacity_reservation_config( ) +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/tests/unit/test_common_utils.py b/sagemaker-core/tests/unit/test_common_utils.py index 8fd6a14c2c..0ae61b321e 100644 --- a/sagemaker-core/tests/unit/test_common_utils.py +++ b/sagemaker-core/tests/unit/test_common_utils.py @@ -2275,6 +2275,7 @@ 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 @@ -2289,18 +2290,50 @@ def test_shape_is_serialized_to_request_keys(self): "MlReservationArn": self.PLAN_ARN, } - def test_shape_with_arn_only_omits_preference(self): + 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.PLAN_ARN) - assert _resolve_capacity_reservation_config(config) == {"MlReservationArn": self.PLAN_ARN} + 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} + 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 diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 8aa1b1ac8a..b8e6367cd5 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -2603,8 +2603,8 @@ def _reused_endpoint_matches_config( 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 only if its variant uses the same - reservation ARN and preference. + 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 @@ -4481,13 +4481,25 @@ 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. @@ -6522,13 +6534,15 @@ def deploy( k: v for k, v in kwargs.items() if k != "capacity_reservation_config" } for ic in self._deployables.get("InferenceComponents", []): - creates_component = bool(pending_capacity_kwargs) and not self._does_ic_exist( - ic.get("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 {}), ) ) diff --git a/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py index a2ab9cfcbb..28d4e22f6b 100644 --- a/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py +++ b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py @@ -23,6 +23,7 @@ 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" EXPECTED_REQUEST = { "CapacityReservationPreference": "capacity-reservations-only", "MlReservationArn": PLAN_ARN, @@ -93,14 +94,24 @@ def test_dict_config_is_accepted(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}, + endpoint_name="ml-reservation-endpoint", + capacity_reservation_config={"MlReservationArn": ML_RESERVATION_ARN}, ) assert self._deployed_variant()["CapacityReservationConfig"] == { - "MlReservationArn": PLAN_ARN + "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", @@ -166,6 +177,30 @@ def test_update_endpoint_honors_wait_false(self): "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. @@ -313,15 +348,19 @@ def test_orchestrator_that_creates_endpoint_gets_config(self, mock_deploy, *_): "CustomOrchestrator": {"Mode": "Endpoint", "Name": "orch", "Model": Mock()}, } - 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, - ) + 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) @@ -386,6 +425,31 @@ def test_all_components_existing_warns_config_ignored(self, mock_deploy_for_ic, 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"} @@ -715,10 +779,10 @@ def test_recommendation_variant_gets_config(self, _): "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() + 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() From ffcb9f4bbb7ef65d7a6dc36e01ef4d22336678cb Mon Sep 17 00:00:00 2001 From: Liam Neal Reilly Date: Thu, 12 Nov 2026 20:09:41 +0000 Subject: [PATCH 7/7] test: cover partial reservation matching with an ML reservation ARN The training plan default adds CapacityReservationPreference before reuse matching, so the ARN-only reuse tests no longer exercised "compare only the keys the request sets". Use an ML reservation ARN, which keeps an ARN-only request, so the tests fail if the comparison reads keys the request didn't set. --- X-AI-Prompt: Fix the reviewer's finding that the reuse-matching tests no longer cover partial matching X-AI-Tool: Claude Code (claude-opus-5-5) --- .../tests/unit/test_deploy_capacity_reservation.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py index 28d4e22f6b..e5e9cc2e4f 100644 --- a/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py +++ b/sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py @@ -24,6 +24,10 @@ 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, @@ -485,18 +489,19 @@ def test_deploy_reuse_skips_on_demand_endpoint_for_reservation_request(self, moc @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": EXPECTED_REQUEST, + "CapacityReservationConfig": ML_RESERVATION_WITH_PREFERENCE, } ) self._builder().deploy( instance_type="ml.p5.48xlarge", reuse_resources=True, - capacity_reservation_config={"MlReservationArn": PLAN_ARN}, + capacity_reservation_config={"MlReservationArn": ML_RESERVATION_ARN}, ) mock_deploy.assert_not_called() @@ -574,18 +579,19 @@ def test_endpoint_on_same_reservation_is_reused(self): ) 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": EXPECTED_REQUEST, + "CapacityReservationConfig": ML_RESERVATION_WITH_PREFERENCE, } ) assert ( builder._reused_endpoint_matches_config( "ep", instance_type="ml.p5.48xlarge", - capacity_reservation_config={"MlReservationArn": PLAN_ARN}, + capacity_reservation_config={"MlReservationArn": ML_RESERVATION_ARN}, ) is True )