Skip to content

feature: support capacity_reservation_config in ModelBuilder.deploy - #6575

Open
lhnealreilly wants to merge 10 commits into
aws:masterfrom
lhnealreilly:feature/deploy-capacity-reservation-config
Open

lhnealreilly wants to merge 10 commits into
aws:masterfrom
lhnealreilly:feature/deploy-capacity-reservation-config

Conversation

@lhnealreilly

@lhnealreilly lhnealreilly commented Nov 10, 2026 •

Copy link
Copy Markdown
Contributor

Issue

ModelBuilder.deploy() has no way to launch an endpoint on reserved capacity. The CreateEndpointConfig API supports ProductionVariant.CapacityReservationConfig, which binds a variant to either:

Without 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 boto3 create_endpoint_config calls, and lose the image, model and resource resolution that ModelBuilder does.

Description of changes

New: capacity_reservation_config on ModelBuilder.deploy()

from sagemaker.core.shapes import ProductionVariantCapacityReservationConfig

endpoint = model_builder.deploy(
    instance_type="ml.p5.48xlarge",  # must match the reservation's instance type
    capacity_reservation_config=ProductionVariantCapacityReservationConfig(
        capacity_reservation_preference="capacity-reservations-only",
        ml_reservation_arn=training_plan_arn,
    ),
)
  • Accepts a ProductionVariantCapacityReservationConfig shape or a dict with the API keys (CapacityReservationPreference, MlReservationArn).
  • Validated in deploy() before any resources are created. Unknown dict keys raise ValueError, wrong types raise TypeError, and combining it with serverless inference raises ValueError.
  • Applied on every deploy path that builds a production variant: model-based and inference-component endpoints in _deploy_core_endpoint (which the JumpStart, DJL, TGI, TEI and MMS wrappers route through), model customization, Nova, recommendation, and modelbuilder_list inference-component deploys.
  • The reservation is bound to the endpoint's production variant, so it only takes effect when the endpoint is created. If the endpoint already exists (an inference component added to an InService endpoint, or an existing inference component updated in place), the config can't apply and a warning is logged rather than dropping it silently.
  • A training plan ARN given without a preference gets 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 same capacity_reservation_config parameter. The new sagemaker.core.common_utils._resolve_capacity_reservation_config() normalizes it.

Also fixes deploy(update_endpoint=True), which raised TypeError

The model-based update path called Session.create_endpoint_config() with model_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:

TypeError: missing a required argument: 'production_variants'

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:

  • Config built from the prebuilt variant. The new endpoint config uses the same production_variant dict 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 by test_update_endpoint_keeps_serverless_and_accelerator_settings).
  • Unique config name. The endpoint-config name comes from name_from_base(...) instead of a fixed self.model_name, so repeated updates don't collide.
  • wait is honored. update_endpoint is called with wait=wait, so deploy(update_endpoint=True, wait=False) returns without waiting.

A regression test binds the recorded create_endpoint_config and update_endpoint calls to the real Session signatures, so a plain Mock can no longer hide a mismatch.

Testing done

  • New sagemaker-serve/tests/unit/test_deploy_capacity_reservation.py (38 tests):
    • model-based, inference-component and update_endpoint variants carry CapacityReservationConfig
    • update_endpoint calls bind to the real Session.create_endpoint_config / Session.update_endpoint signatures (regression test for the TypeError)
    • dict input is accepted; omitting the config leaves the variant unchanged
    • serverless and invalid configs are rejected before any resources are created
    • deploy() forwards the config, including on the modelbuilder_list inference-component path
    • a warning is logged when the endpoint or inference component already exists
    • the Nova path sets it on the core ProductionVariant
  • New core tests for _resolve_capacity_reservation_config (10, including the training-plan default) and production_variant (2).
  • sagemaker-core tests/unit/test_common_utils.py and tests/unit/helper/test_session_helper.py: 422 passed.
  • sagemaker-serve unit 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.
  • Not yet verified against a live endpoint on an Active training plan.

Merge Checklist

General

  • I have read the CONTRIBUTING doc
  • I certify that the changes I am introducing will be backward compatible, and I have discussed concerns about this, if any, with the Python SDK team
  • I used the commit message format described in CONTRIBUTING
  • I have updated any necessary documentation, including docstrings

Tests

  • I have added tests that prove my fix is effective or that my feature works
  • I have checked that my tests are not configured for a specific region or account (if appropriate)

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

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.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@lhnealreilly

Copy link
Copy Markdown
Contributor Author

Addressed the AI review in 86e0357:

  • Medium (_deployables IC path drops the config): confirmed. The modelbuilder_list loop now forwards capacity_reservation_config to each _deploy_for_ic call, so the endpoint created for the first inference component launches on the reservation. Only this key is forwarded, so other kwargs on that path behave as before. Added tests for this path, both with and without a config.
  • Low (config resolved twice): leaving as is. production_variant() is a public helper that also takes the shape directly, so it has to resolve its input. ModelBuilder resolves earlier to fail fast before any resources are created. Re-resolving an already-resolved dict returns it unchanged and costs very little.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 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_ic routes to _update_inference_component(ic_name, resource_requirements, **kwargs) with capacity_reservation_config now present in kwargs. _update_inference_component silently 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 that capacity_reservation_config has 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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

evakravi
evakravi previously approved these changes Nov 11, 2026
…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.
@lhnealreilly

Copy link
Copy Markdown
Contributor Author

Addressed the second AI review in 5d6888b:

  • Existing inference component update ignores capacity_reservation_config: agreed. The reservation is bound to the endpoint's production variant, so an in-place inference component update can't apply it. A warning is now logged in that case. The same applies when an inference component is added to an endpoint that's already InService, so that path warns too. Both have tests.
  • Bundled update_endpoint fix: keeping it in this PR, because moving an existing endpoint onto or off a training plan depends on the update path working. Added the suggested regression test: it binds the recorded create_endpoint_config and update_endpoint calls to the real Session method signatures. I confirmed it fails with the original TypeError: missing a required argument: 'production_variants' when run against the old call.

On the integ failures: integ-tests (sagemaker-core) is also failing on #6572, #6571, #6566, #6565 and #6558, and in the nightly ci-health-v3-master run for the last 3 days. integ-tests (sagemaker-train) is intermittent across recent PRs. integ-tests (sagemaker-serve) passes, and that's where the deploy changes are. I couldn't read the CodeBuild output because GitHub's job-log API returns 502 for this repo, so a maintainer with CI log access may want to confirm.

evakravi
evakravi previously approved these changes Nov 11, 2026

@evakravi evakravi left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Findings inline.

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

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 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: The update_endpoint=True path now passes wait=wait to Session.update_endpoint (previously it was called with no wait, defaulting to True). This is a behavioral change for existing update_endpoint callers who pass wait=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.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 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", []):

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@evakravi evakravi left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Findings inline.

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

@evakravi evakravi left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Findings inline.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 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(

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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(

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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

@evakravi evakravi left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Findings inline.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 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)

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 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).

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

This branch was successfully deployed

1 active deployment
auto-approve — ffcb9f4b Deployed Nov 12, 2026 by lhnealreilly via wait-for-approval #2222
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants