feature: support capacity_reservation_config in ModelBuilder.deploy

sagemaker-core/src/sagemaker/core/common_utils.py

@@ -2129,6 +2129,75 @@ def _resolve_routing_config(routing_config: Optional[Dict[str, Any]]) -> Optiona
     return None
 
 
+_CAPACITY_RESERVATION_CONFIG_KEYS = ("CapacityReservationPreference", "MlReservationArn")
+_CAPACITY_RESERVATIONS_ONLY = "capacity-reservations-only"
+
+
+def _resolve_capacity_reservation_config(
+    capacity_reservation_config: Optional[Any],
+) -> Optional[Dict[str, Any]]:
+    """Resolve a production variant capacity reservation config to its request dict.
+
+    Args:
+        capacity_reservation_config (Optional[Union[dict,
+            sagemaker.core.shapes.ProductionVariantCapacityReservationConfig]]):
+            Either a ``ProductionVariantCapacityReservationConfig`` shape or a dict using the
+            ``CreateEndpointConfig`` request keys (``CapacityReservationPreference``,
+            ``MlReservationArn``).
+
+    Returns:
+        Optional[Dict[str, Any]]: The ``CapacityReservationConfig`` request dict, or None if
+            no config was given or it has no fields set. A training plan ARN
+            (``...:training-plan/<name>``) given without a preference gets
+            ``CapacityReservationPreference="capacity-reservations-only"``, the only value
+            SageMaker accepts for training plans. Other ARNs are returned as given.
+
+    Raises:
+        ValueError: If a dict contains keys other than ``CapacityReservationPreference`` and
+            ``MlReservationArn``.
+        TypeError: If the config is neither a dict nor a
+            ``ProductionVariantCapacityReservationConfig``.
+    """
+    if capacity_reservation_config is None:
+        return None
+
+    # Imported lazily: sagemaker.core.shapes imports from this module.
+    from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig
+    from sagemaker.core.utils.utils import serialize
+
+    if isinstance(capacity_reservation_config, ProductionVariantCapacityReservationConfig):
+        return _default_training_plan_preference(serialize(capacity_reservation_config) or None)
+
+    if isinstance(capacity_reservation_config, dict):
+        unknown_keys = set(capacity_reservation_config) - set(_CAPACITY_RESERVATION_CONFIG_KEYS)
+        if unknown_keys:
+            raise ValueError(
+                f"Unsupported capacity_reservation_config keys: {sorted(unknown_keys)}. "
+                f"Supported keys are: {list(_CAPACITY_RESERVATION_CONFIG_KEYS)}."
+            )
+        resolved = {k: v for k, v in capacity_reservation_config.items() if v is not None}
+        return _default_training_plan_preference(resolved or None)
+
+    raise TypeError(
+        "capacity_reservation_config must be a dict or a "
+        "sagemaker.core.shapes.ProductionVariantCapacityReservationConfig, got "
+        f"{type(capacity_reservation_config).__name__}."
+    )
+
+
+def _default_training_plan_preference(
+    resolved: Optional[Dict[str, Any]],
+) -> Optional[Dict[str, Any]]:
+    """Set the preference a training plan requires when the caller left it unset."""
+    if (
+        resolved
+        and ":training-plan/" in str(resolved.get("MlReservationArn", ""))
+        and "CapacityReservationPreference" not in resolved
+    ):
+        return {**resolved, "CapacityReservationPreference": _CAPACITY_RESERVATIONS_ONLY}
+    return resolved
+
+
 @lru_cache
 def get_instance_rate_per_hour(
     instance_type: str,

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:

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

sagemaker-core/tests/unit/test_common_utils.py

@@ -2271,6 +2271,94 @@ def test_resolve_routing_config_none(self):
         assert result is None
 
 
+class TestResolveCapacityReservationConfig:
+    """Test _resolve_capacity_reservation_config function."""
+
+    PLAN_ARN = "arn:aws:sagemaker:us-west-2:123456789012:training-plan/my-plan"
+    ML_ARN = "arn:aws:sagemaker:us-west-2:123456789012:ml-reservation/my-reservation"
+
+    def test_shape_is_serialized_to_request_keys(self):
+        from sagemaker.core.common_utils import _resolve_capacity_reservation_config
+        from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig
+
+        config = ProductionVariantCapacityReservationConfig(
+            capacity_reservation_preference="capacity-reservations-only",
+            ml_reservation_arn=self.PLAN_ARN,
+        )
+        assert _resolve_capacity_reservation_config(config) == {
+            "CapacityReservationPreference": "capacity-reservations-only",
+            "MlReservationArn": self.PLAN_ARN,
+        }
+
+    def test_shape_with_ml_reservation_arn_only_omits_preference(self):
+        from sagemaker.core.common_utils import _resolve_capacity_reservation_config
+        from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig
+
+        config = ProductionVariantCapacityReservationConfig(ml_reservation_arn=self.ML_ARN)
+        assert _resolve_capacity_reservation_config(config) == {"MlReservationArn": self.ML_ARN}
+
+    def test_dict_is_passed_through_without_none_values(self):
+        from sagemaker.core.common_utils import _resolve_capacity_reservation_config
+
+        config = {"MlReservationArn": self.ML_ARN, "CapacityReservationPreference": None}
+        assert _resolve_capacity_reservation_config(config) == {"MlReservationArn": self.ML_ARN}
+
+    def test_shape_with_training_plan_arn_only_gets_reservations_only(self):
+        from sagemaker.core.common_utils import _resolve_capacity_reservation_config
+        from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig
+
+        config = ProductionVariantCapacityReservationConfig(ml_reservation_arn=self.PLAN_ARN)
+        assert _resolve_capacity_reservation_config(config) == {
+            "MlReservationArn": self.PLAN_ARN,
+            "CapacityReservationPreference": "capacity-reservations-only",
+        }
+
+    def test_dict_with_training_plan_arn_only_gets_reservations_only(self):
+        from sagemaker.core.common_utils import _resolve_capacity_reservation_config
+
+        config = {"MlReservationArn": self.PLAN_ARN, "CapacityReservationPreference": None}
+        assert _resolve_capacity_reservation_config(config) == {
+            "MlReservationArn": self.PLAN_ARN,
+            "CapacityReservationPreference": "capacity-reservations-only",
+        }
+
+    def test_training_plan_dict_input_is_not_mutated(self):
+        from sagemaker.core.common_utils import _resolve_capacity_reservation_config
+
+        config = {"MlReservationArn": self.PLAN_ARN}
+        _resolve_capacity_reservation_config(config)
+        assert config == {"MlReservationArn": self.PLAN_ARN}
+
+    def test_explicit_preference_is_kept(self):
+        from sagemaker.core.common_utils import _resolve_capacity_reservation_config
+
+        config = {"MlReservationArn": self.PLAN_ARN, "CapacityReservationPreference": "other"}
+        assert _resolve_capacity_reservation_config(config) == config
+
+    def test_empty_config_resolves_to_none(self):
+        from sagemaker.core.common_utils import _resolve_capacity_reservation_config
+        from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig
+
+        assert _resolve_capacity_reservation_config(None) is None
+        assert _resolve_capacity_reservation_config({}) is None
+        assert (
+            _resolve_capacity_reservation_config(ProductionVariantCapacityReservationConfig())
+            is None
+        )
+
+    def test_dict_with_unknown_key_raises(self):
+        from sagemaker.core.common_utils import _resolve_capacity_reservation_config
+
+        with pytest.raises(ValueError, match="ml_reservation_arn"):
+            _resolve_capacity_reservation_config({"ml_reservation_arn": self.PLAN_ARN})
+
+    def test_unsupported_type_raises(self):
+        from sagemaker.core.common_utils import _resolve_capacity_reservation_config
+
+        with pytest.raises(TypeError):
+            _resolve_capacity_reservation_config(self.PLAN_ARN)
+
+
 class TestWaitUntil:
     """Test _wait_until function."""