Repository navigation
feature: support capacity_reservation_config in ModelBuilder.deploy - #6575
lhnealreilly wants to merge 10 commits into
Conversation
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.
There was a problem hiding this comment.
🤖 AI code review
This PR adds a capacity_reservation_config parameter to ModelBuilder.deploy() (and the underlying production_variant helper plus a new _resolve_capacity_reservation_config util), letting users bind an endpoint's production variant to a SageMaker training plan or ODCR-backed ML reservation. It also fixes a real bug in the model-based update_endpoint=True path, which previously called create_endpoint_config with invalid keyword arguments. The core resolution/serialization logic is correct and well tested, and the parameter is threaded through most deploy paths. The main issue is that the _deployables (modelbuilder_list) inference-component path silently drops the config, which contradicts the PR's "never silently dropped" claim.
[{"path": "sagemaker-serve/src/sagemaker/serve/model_builder.py", "line": 6396, "severity": "medium", "body": "This _deploy_for_ic call does not forward **kwargs, so capacity_reservation_config (placed into kwargs by deploy()) is silently dropped on the _deployables / modelbuilder_list inference-component deploy path. The PR description states the setting "is never silently dropped" and is "applied on every deploy path that builds a production variant," but on this path the variant is created via _deploy(... **kwargs) inside _deploy_for_ic, which never receives the config. The custom-orchestrator IC path at line 6431 does forward **kwargs, so the two IC paths behave inconsistently. Consider passing the relevant kwargs here too, and add a test covering the _deployables path."}, {"path": "sagemaker-serve/src/sagemaker/serve/model_builder.py", "line": 4195, "severity": "low", "body": "capacity_reservation_config is already resolved to a request dict at line ~4004 before being passed here, and production_variant() calls _resolve_capacity_reservation_config on it again. This is harmless (re-resolving a valid dict is idempotent) but does the validation/serialization work twice. Minor; could pass the already-resolved value straight into the variant dict instead."}]
Reviewed commit 9192bb7. Automated review; verify before acting.
…erence component path
|
Addressed the AI review in 86e0357:
|
There was a problem hiding this comment.
🤖 AI code review
This PR adds capacity_reservation_config to ModelBuilder.deploy(), threading a ProductionVariantCapacityReservationConfig (or dict) through every deploy path (_deploy_core_endpoint, model-customization, Nova, recommendation) and adding a _resolve_capacity_reservation_config helper plus a production_variant() parameter in sagemaker-core. It also fixes the previously broken deploy(update_endpoint=True) model-based path, which called create_endpoint_config with keyword args that don't exist in that method's signature. I verified the base signatures, serialization behavior, and that no existing unit tests assert the old (broken) update call. The change is well-structured and thoroughly tested; I only have minor notes.
Findings not attached to a changed line:
sagemaker-serve/src/sagemaker/serve/model_builder.py:4426🟡 Low: When an inference component already exists,_deploy_for_icroutes to_update_inference_component(ic_name, resource_requirements, **kwargs)withcapacity_reservation_confignow present inkwargs._update_inference_componentsilently ignores it (it only reads a few specific keys), so updating an existing IC on a capacity-reserved endpoint does not actually apply/change the reservation. That is arguably acceptable (the reservation lives on the endpoint config, not the IC), but it is a silent drop; consider a brief comment or log so callers are not surprised thatcapacity_reservation_confighas no effect on the existing-IC update path.
Reviewed commit 86e0357. Automated review; verify before acting.
| capacity_reservation_config=capacity_reservation_config, | ||
| ) | ||
| if endpoint_name: | ||
| self.endpoint_name = endpoint_name |
There was a problem hiding this comment.
🟡 Low: This hunk bundles a non-trivial bugfix (the update_endpoint=True model-based path previously always raised TypeError because create_endpoint_config was called with model_name=/instance_type=/initial_instance_count= etc. that its signature (name, production_variants, tags, kms_key, data_capture_config_dict, async_inference_config_dict, explainer_config_dict, vpc_config, enable_network_isolation, role) does not accept) into a feature PR. Prior review precedent flags mixing orthogonal changes (Source: #5764). As you offered in the description, splitting this into a separate fix: PR would make both changes easier to review and bisect. At minimum, add a regression test that exercises the real (unmocked-session) create_endpoint_config call so a future signature drift is caught — the existing test_deploy_with_update_endpoint mocks _deploy and would not have caught the original bug.
…e_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.
|
Addressed the second AI review in 5d6888b:
On the integ failures: |
- 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.
There was a problem hiding this comment.
🤖 AI code review
This PR adds capacity_reservation_config support to ModelBuilder.deploy(), wiring a ProductionVariantCapacityReservationConfig (shape or dict) into every production-variant build path (core endpoint, inference-component, Nova, model-customization, recommendation), and incidentally fixes a fully-broken deploy(update_endpoint=True) path whose create_endpoint_config call used non-existent keyword args. The implementation is thorough, consistently threads the config, validates up front, and has extensive, well-targeted unit tests (including a regression test that binds mocked session calls to the real signatures). I verified the shape exists, serializes to the API keys, imports resolve, and the new production_variant parameter is appended so positional callers are unaffected. Overall this looks correct and well-tested; only minor observations below.
Findings not attached to a changed line:
sagemaker-serve/src/sagemaker/serve/model_builder.py:4307🟡 Low: Theupdate_endpoint=Truepath now passeswait=waittoSession.update_endpoint(previously it was called with nowait, defaulting toTrue). This is a behavioral change for existingupdate_endpointcallers who passwait=False— they will now get a non-blocking update where previously the call blocked. It is arguably the more correct behavior, but it is a change bundled with the capacity-reservation feature and worth calling out in the PR description / changelog so consumers relying on the old blocking behavior are aware.
Reviewed commit 09417d1. Automated review; verify before acting.
There was a problem hiding this comment.
🤖 AI code review
This PR adds capacity_reservation_config support to ModelBuilder.deploy(), wiring a new ProductionVariantCapacityReservationConfig (or dict) through every deploy path (model-based, inference-component, model customization, Nova, recommendation, multi-IC deployables), and adding _resolve_capacity_reservation_config / production_variant(capacity_reservation_config=...) helpers. It also fixes a real pre-existing bug where the update_endpoint=True model path called Session.create_endpoint_config() with kwargs that method never accepted (model_name, initial_instance_count, etc.), which raised TypeError.
The implementation is coherent and very well tested (serialization keys, serverless rejection, reuse matching, update-path signature binding, multi-IC forwarding, warnings for already-existing endpoints/ICs). I verified the session signatures, update_args None-dropping, and the serverless gating. I found no correctness/security bugs. One low-severity edge case is noted below. The update_endpoint config-name change from deterministic to name_from_base affects all update callers, but the old path was broken anyway, so this is a net fix.
Reviewed commit f19ebf4. Automated review; verify before acting.
| orchestrator_kwargs = { | ||
| k: v for k, v in kwargs.items() if k != "capacity_reservation_config" | ||
| } | ||
| for ic in self._deployables.get("InferenceComponents", []): |
There was a problem hiding this comment.
🟡 Low: In the multi-IC deployables loop, pending_capacity_kwargs is unconditionally reset to {} after the first _deploy_for_ic call, on the assumption that the first component always creates the endpoint. If the first inference component already exists (_does_ic_exist returns True), _deploy_for_ic only logs the "ignored" warning and does not create the endpoint — but the reservation kwargs are still consumed and dropped. A subsequent component that actually creates the endpoint will then launch on On-Demand capacity with the reservation silently lost. Consider carrying pending_capacity_kwargs forward until a component actually creates the endpoint (e.g. only clear it on the call that creates the endpoint), or at least documenting this limitation.
There was a problem hiding this comment.
Fixed in a202011. The reservation stays pending until a component that doesn't exist yet is deployed, because only that call can create the endpoint. The _does_ic_exist check runs only when a reservation was requested. If every component already exists, one warning says the reservation is ignored. The custom orchestrator gets the config only if no component consumed it. Added test_existing_first_component_does_not_consume_config and test_all_components_existing_warns_config_ignored.
- 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.
There was a problem hiding this comment.
🤖 AI code review
This PR adds capacity_reservation_config support to ModelBuilder.deploy(), threading a ProductionVariantCapacityReservationConfig (or dict) through every deploy path (model-based, inference-component, model customization, Nova, recommendation, and multi-IC/orchestrator), plus a new _resolve_capacity_reservation_config helper and a production_variant() parameter in session_helper. It also fixes the previously-broken model-based update_endpoint=True path, which was calling Session.create_endpoint_config() with kwargs that signature doesn't accept.
I verified the key claims against the base commit: the old create_endpoint_config(name=, model_name=, instance_type=, ...) call did not match the real signature (name, production_variants, tags, kms_key, ...), so the fix is correct; production_variant is in scope in the update branch; the core shape and serialization drop explicit None fields as the shape helper relies on. The change is well-covered by unit tests. I found no high-severity issues; only minor observations below.
Reviewed commit 0b4b90c. Automated review; verify before acting.
| ic.get("Name") | ||
| ) | ||
| endpoints.append( | ||
| self._deploy_for_ic( |
There was a problem hiding this comment.
🟡 Low: _does_ic_exist(ic.get("Name")) is evaluated here to compute creates_component, and then _deploy_for_ic calls _does_ic_exist again internally for the same component. Each call issues a describe_inference_component API request, so every component is now described twice. Minor, but you could pass the already-computed result down or cache it.
There was a problem hiding this comment.
Fixed in e9056e8. The loop passes the existence result it already computed to _deploy_for_ic(ic_exists=...), so each component is described once. test_each_component_is_described_once covers it.
| @@ -6392,8 +6510,39 @@ def deploy( | |||
|
|
|||
There was a problem hiding this comment.
🟡 Low: The update_endpoint=True fix changes observable behavior beyond capacity reservation: the generated endpoint-config name moves from a fixed self.model_name to name_from_base(...), and update_endpoint is now called with wait=wait. These are reasonable fixes, but since the old path raised TypeError for all callers it's worth a changelog/release note so consumers who were catching or working around that failure are aware the update path now actually executes.
There was a problem hiding this comment.
Added to the PR description under "Also fixes update_endpoint=True". This repo's CHANGELOG.md is only edited by release commits, so I didn't add an entry.
| _CAPACITY_RESERVATION_CONFIG_KEYS = ("CapacityReservationPreference", "MlReservationArn") | ||
|
|
||
|
|
||
| def _resolve_capacity_reservation_config( |
There was a problem hiding this comment.
One thing from my live testing for FTP/ODCR change in Sherpa: CreateEndpointConfig rejects a training-plan ARN unless CapacityReservationPreference is capacity-reservations-only. The docstring says this, but a caller who passes only the ARN won't find out until the API call fails. Would it make sense to default the preference when the ARN is a training-plan/ ARN, or raise a clear client-side error? botocore doesn't validate the ARN pattern or the enum either, so this would be the first place that catches it.
There was a problem hiding this comment.
Done in e9056e8. A training-plan/ ARN without a preference now gets CapacityReservationPreference="capacity-reservations-only" in _resolve_capacity_reservation_config. Every deploy path goes through that function. ML reservation ARNs and explicit preferences are passed as given, and the service still decides whether they're valid. Tests cover the shape and dict inputs, check that the caller's dict isn't mutated, and check the request built by deploy.
| initial_instance_count=initial_instance_count, | ||
| instance_type=instance_type, | ||
| accelerator_type=accelerator_type, | ||
| name=name_from_base(base_from_name(self.model_name)), |
There was a problem hiding this comment.
Small question: this branch now builds the config from the prebuilt variant, which is nice. It also changes the config name from self.model_name to name_from_base(...), adds wait=wait and drops the serverless/accelerator kwargs. Is that intended as part of this PR? If it is, could the description mention it? I wasn't sure whether the serverless update path still behaves the same.
There was a problem hiding this comment.
Yes, it's intended: the old call raised TypeError on every update. Serverless still behaves the same. The prebuilt variant already gets serverless_inference_config, accelerator_type, volume_size and the timeouts, so the separate arguments were duplicates. In e9056e8, test_update_endpoint_keeps_serverless_and_accelerator_settings checks that the update config keeps ServerlessConfig (with no InstanceType) and AcceleratorType. The description now has an "Also fixes update_endpoint=True" section that lists the config name, wait=wait, and the variant-based config.
…city-reservation-config
- _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.
There was a problem hiding this comment.
🤖 AI code review
This PR adds a capacity_reservation_config parameter to ModelBuilder.deploy() (and the underlying session_helper.production_variant()), threading it through the model-based, inference-component, model-customization, Nova and recommendation deploy paths, and also fixes a pre-existing TypeError in the update_endpoint=True model path by reusing the prebuilt production variant and calling create_endpoint_config with its real signature. I verified the production_variant variable is in scope in the update branch, the create_endpoint_config/update_endpoint signatures match the new call sites, the ProductionVariant/ProductionVariantCapacityReservationConfig shapes carry the expected fields, CapacityReservationConfig is correctly confined to the instance (non-serverless) branch, and the server wrappers pass kwargs through to _deploy_core_endpoint. The change is careful about edge cases (existing endpoints/ICs are warned rather than silently dropped, reuse matching compares reservation config) and is backed by thorough unit tests, including a regression test that binds recorded calls to the real Session signatures. I found no correctness or compatibility issues worth flagging.
[]
Reviewed commit e9056e8. Automated review; verify before acting.
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)
There was a problem hiding this comment.
🤖 AI code review
This PR adds capacity_reservation_config support to ModelBuilder.deploy(), threading a ProductionVariantCapacityReservationConfig (or dict) through all deploy paths (core endpoint, inference-component, model-customization, Nova, recommendation, and multi-IC orchestration), with validation that happens up front. It also fixes a pre-existing bug where deploy(update_endpoint=True) on the model-based path called Session.create_endpoint_config() with the wrong keyword arguments and raised TypeError. The implementation is careful about where reserved capacity can and cannot apply (warning rather than silently dropping on existing endpoints/ICs), and it is backed by an extensive, well-targeted unit test suite. Overall the change is correct and thorough; I found only one minor documentation inconsistency.
Reviewed commit ffcb9f4. Automated review; verify before acting.
| 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). |
There was a problem hiding this comment.
🟡 Low: This docstring says capacity_reservation_config is "Ignored for serverless variants", but the deploy() / _deploy_core_endpoint() path now raises ValueError ("capacity_reservation_config is not supported for serverless inference") when combined with serverless. production_variant() itself silently drops it for serverless (it's only read inside the non-serverless else branch), so the two layers behave differently. Consider aligning the wording to avoid confusing callers who read this helper's doc directly.
Issue
ModelBuilder.deploy()has no way to launch an endpoint on reserved capacity. TheCreateEndpointConfigAPI supportsProductionVariant.CapacityReservationConfig, which binds a variant to either:endpoint(Flexible Training Plans for inference endpoints), orWithout SDK support, users who hit insufficient On-Demand GPU capacity (common for large open-weight models on p5/p5en/p6) have to drop to raw
boto3create_endpoint_configcalls, and lose the image, model and resource resolution thatModelBuilderdoes.Description of changes
New:
capacity_reservation_configonModelBuilder.deploy()ProductionVariantCapacityReservationConfigshape or a dict with the API keys (CapacityReservationPreference,MlReservationArn).deploy()before any resources are created. Unknown dict keys raiseValueError, wrong types raiseTypeError, and combining it with serverless inference raisesValueError._deploy_core_endpoint(which the JumpStart, DJL, TGI, TEI and MMS wrappers route through), model customization, Nova, recommendation, andmodelbuilder_listinference-component deploys.CapacityReservationPreference="capacity-reservations-only", the only value SageMaker accepts for training plans. ML reservation ARNs and explicit preferences are passed as given.sagemaker.core.helper.session_helper.production_variant()gains the samecapacity_reservation_configparameter. The newsagemaker.core.common_utils._resolve_capacity_reservation_config()normalizes it.Also fixes
deploy(update_endpoint=True), which raisedTypeErrorThe model-based update path called
Session.create_endpoint_config()withmodel_name=,instance_type=,initial_instance_count=and similar keyword arguments. That method's signature is(name, production_variants, tags, kms_key, ...), so every call failed:Unit tests mock the session, so this wasn't caught. Kept in this PR because moving an existing endpoint onto or off a training plan needs the update path to work. Observable changes on that path:
production_variantdict as a new endpoint, so every variant setting carries over: instance type and count, serverless config, accelerator type, volume size, timeouts, routing, AMI version, and the capacity reservation. The serverless and accelerator arguments that the old call listed separately are already on that variant (covered bytest_update_endpoint_keeps_serverless_and_accelerator_settings).name_from_base(...)instead of a fixedself.model_name, so repeated updates don't collide.waitis honored.update_endpointis called withwait=wait, sodeploy(update_endpoint=True, wait=False)returns without waiting.A regression test binds the recorded
create_endpoint_configandupdate_endpointcalls to the realSessionsignatures, so a plainMockcan no longer hide a mismatch.Testing done
sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py(38 tests):update_endpointvariants carryCapacityReservationConfigupdate_endpointcalls bind to the realSession.create_endpoint_config/Session.update_endpointsignatures (regression test for theTypeError)deploy()forwards the config, including on themodelbuilder_listinference-component pathProductionVariant_resolve_capacity_reservation_config(10, including the training-plan default) andproduction_variant(2).sagemaker-coretests/unit/test_common_utils.pyandtests/unit/helper/test_session_helper.py: 422 passed.sagemaker-serveunit suite: no new failures. The same 49 tests fail with and without this change because the local environment has no AWS credentials (JumpStart S3 reads) and no TensorFlow.Merge Checklist
General
Tests
By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.