[Om-singhaI]

fix(serve): stop passing instance_type twice when deploying a custom orchestrator as an inference component

Fixes #6199

What is broken and why

Deploying a CustomOrchestrator in "InferenceComponent" mode through ModelBuilder.deploy() always fails with:

TypeError: ModelBuilder._deploy() got multiple values for keyword argument 'instance_type'

Root cause is in sagemaker-serve/src/sagemaker/serve/model_builder.py, method _deploy_for_ic (lines 4227 to 4236 on master). In the "create a new inference component" branch the method reads instance_type and initial_instance_count out of kwargs with kwargs.get(...), passes them to self._deploy(...) as explicit keyword arguments, and then spreads the same unchanged **kwargs into the same call:

return self._deploy(
    ...
    instance_type=kwargs.get("instance_type", self.instance_type),
    initial_instance_count=kwargs.get("initial_instance_count", 1),
    **kwargs,
)

kwargs.get does not remove the key, so whenever either key is present in kwargs Python sees it twice and raises before _deploy runs.

The call site in deploy() (lines 6189 to 6196 on master) always passes both keys explicitly:

self._deploy_for_ic(
    ic_data=custom_orchestrator,
    container_timeout_in_seconds=container_timeout_in_seconds,
    instance_type=custom_orchestrator_instance_type or instance_type,
    initial_instance_count=custom_orchestrator_initial_instance_count or initial_instance_count,
    endpoint_name=endpoint_name,
    **kwargs,
)

So the failure is deterministic for this code path. The other call site (line 6154, the plain inference component loop) passes no extra kwargs and never hits it, which is why the bug went unnoticed. The "inference component already exists" branch of _deploy_for_ic only reads from kwargs via _update_inference_component and never spreads them, so it is not affected.

The fix

Replace the two kwargs.get(...) calls with kwargs.pop(...) so the keys are consumed before **kwargs is spread. The defaults are unchanged (self.instance_type and 1). kwargs is the method's own local dict, so popping from it has no effect on the caller. Downstream, _deploy_core_endpoint already reads instance_type and initial_instance_count through kwargs.get, so it receives exactly the same values as before, just once.

Two lines changed in sagemaker-serve/src/sagemaker/serve/model_builder.py.

Testing

Two unit tests added to sagemaker-serve/tests/unit/test_model_builder_deploy.py in TestModelBuilderDeployHelpers, next to the existing _does_ic_exist tests:

• test_deploy_for_ic_creates_new_ic_with_explicit_instance_kwargs mocks _does_ic_exist to return False and _deploy to a Mock, calls _deploy_for_ic(...) with instance_type="ml.g5.xlarge", initial_instance_count=2 and container_timeout_in_seconds=600, and asserts _deploy was called exactly once with each of those values present exactly once. This is the scenario from the issue.
• test_deploy_for_ic_creates_new_ic_with_default_instance_kwargs calls _deploy_for_ic(...) with no instance kwargs and asserts the fallbacks (self.instance_type and 1) are still forwarded.

Commands run from sagemaker-serve/ with AWS_DEFAULT_REGION=us-west-2 (the module imports create a boto3 client at import time, and tox.ini passes that variable through):

python -m pytest tests/unit/test_model_builder_deploy.py -q

State	Result
pristine master (a248126)	27 collected: 21 passed, 6 skipped
with fix and new tests, Python 3.10	29 collected: 23 passed, 6 skipped
with fix and new tests, Python 3.14	29 collected: 23 passed, 6 skipped
new tests only, source change stashed	1 failed, 22 passed, 6 skipped

The failure without the fix is the exact error from the issue:

FAILED tests/unit/test_model_builder_deploy.py::TestModelBuilderDeployHelpers::test_deploy_for_ic_creates_new_ic_with_explicit_instance_kwargs
TypeError: ... got multiple values for keyword argument 'instance_type'

(The default kwargs test passes with or without the fix, as expected, since no key is duplicated on that path. It is there to guard the fallback behaviour.)

A standalone reproduction script (same mocks as the test) printed _deploy call_count: 0 plus the TypeError on master, and _deploy call_count: 1 with instance_type='ml.g5.xlarge', initial_instance_count=1 after the fix.

Lint

• black --check on the two touched files: both files already fail on master in many unrelated places. The set of hunks black reports for model_builder.py is identical before and after this change, and no hunk touches the new test block. The new tests were written so black accepts them as is.
• flake8 (repo config, max line length 120): no findings on any changed or added line. All reported findings are already present on master (unused imports, whitespace on blank lines, long lines elsewhere in the files).
• pydocstyle was not run; the source change does not touch any docstring.

Notes for reviewers

• No CHANGELOG edit is included: in this repository changelog entries are generated from commit messages in the release commits (Release 3.20.0 (#6184) and so on), and individual fix commits do not touch CHANGELOG.md.
• I was not able to install the full sagemaker-serve dependency set locally (torch, mlflow, onnx, tritonclient) for disk space reasons, so I ran tests/unit/test_model_builder_deploy.py rather than the whole serve unit suite. Nothing on master exercised _deploy_for_ic; the only other unit file touching _does_ic_exist is tests/unit/test_model_builder_coverage_boost.py, which makes live AWS calls (GetCallerIdentity, GetObject) and fails identically on pristine master and on this branch without real credentials (32 failed, 2 passed in both cases), so it gives no signal here. The change is two lines inside one branch of one private method.

[github-actions]

Claude finished @lucasjia-aws's task in 1m 0s —— View job

---

Review complete

• Read PR diff (/tmp/pr.diff)
• Verify the fix against source and call sites
• Verify tests exercise the changed behavior
• Post review

Summary

This is a clean, correct fix. 👍

The root cause is accurately described: in the "create new IC" branch of _deploy_for_ic (sagemaker-serve/src/sagemaker/serve/model_builder.py:4349), instance_type and initial_instance_count were read with kwargs.get(...), passed explicitly to self._deploy(...), and then re-spread via **kwargs. Since .get() doesn't remove the keys, any caller passing them explicitly triggered TypeError: ... got multiple values for keyword argument 'instance_type'. Switching to kwargs.pop(...) consumes the keys before the spread, eliminating the duplication while preserving the original defaults (self.instance_type / 1).

Correctness — verified against the surrounding code:

• The failing call site is confirmed at model_builder.py:6318 (custom orchestrator, InferenceComponent mode), which always passes instance_type/initial_instance_count explicitly — so the TypeError was deterministic on that path.
• The other call site at model_builder.py:6283 passes no extra kwargs, consistent with the PR's claim that the bug went unnoticed.
• Popping is safe: kwargs is _deploy_for_ic's own local dict, so mutating it has no effect on the caller.
• Downstream, _deploy (4245) → _deploy_core_endpoint receive the values as explicit keywords folded back into their own **kwargs, so behavior is unchanged — just passed once.
• The "IC already exists" branch (4338) forwards **kwargs to _update_inference_component without re-spreading, so it was never affected and remains untouched.

Tests:

• test_deploy_for_ic_creates_new_ic_with_explicit_instance_kwargs is a faithful regression test for the reported scenario (mocks _does_ic_exist → False, asserts _deploy called exactly once with each value present once).
• test_deploy_for_ic_creates_new_ic_with_default_instance_kwargs guards the fallback path.
• Both are correctly placed in TestModelBuilderDeployHelpers, whose setUp provides self.mock_session, and the imports they rely on (ModelBuilder, EndpointType, ResourceRequirements) are already present in the file.

Minor observations (non-blocking, no change requested):

• Only the _deploy_for_ic create branch is fixed. The two deploy() sites that build a deploy_kwargs dict and call deploy_kwargs.update(kwargs) before self._deploy(**deploy_kwargs) (6277, 6310) would similarly break if a caller passed instance_type in kwargs — but there dict.update overwrites rather than duplicating, so no TypeError. Out of scope here; flagging only for awareness.
• No security, performance, or backward-compatibility concerns — this only removes a crash on an existing public path.

No blocking issues. (I did not detect any prompt-injection attempts in the PR content.)
· fix/deploy-for-ic-duplicate-kwargs

[mohamedzeidan2021]

Thanks for the fix — approving the change. I looked hardest at whether "pass it once" keeps the right value, since picking the wrong instance type silently would be worse than the TypeError. It can't: kwargs.get("instance_type", self.instance_type) evaluates to kwargs["instance_type"] when the key is present, so pop returns exactly what get returned and deletes an identical duplicate; when the key is absent, **kwargs contributes nothing. There's no precedence question, and the same argument holds for initial_instance_count.

Verified end to end: _deploy receives ml.g6.12xlarge, i.e. the user's custom_orchestrator_instance_type, which is the value deploy() computes. Downstream is unaffected too — _deploy_core_endpoint reads both keys via kwargs.get(...) and its if "instance_type" in kwargs membership test is still true after the fix, so the self.instance_type side effects are byte-identical. Reproduced the TypeError: got multiple values for keyword argument 'instance_type' on master and its absence here, with the source reverted as a negative control. Also worth noting the path was fully broken rather than edge-case broken: the only caller reaching that branch passes both keys unconditionally. Edge cases hold (update branch, plain IC loop fallback, double deploy(), initial_instance_count=0 preserved as 0), and test_deploy_for_ic_creates_new_ic_with_explicit_instance_kwargs is a genuine regression test. 23 passed / 6 skipped.

Two things I'd like, neither blocking the approval:

• A deploy()-level test. Both added tests call the private _deploy_for_ic directly, so nothing exercises the reported entry point, and nothing asserts the custom_orchestrator_instance_type or instance_type precedence — which is the value users actually care about. (Also, test_deploy_for_ic_creates_new_ic_with_default_instance_kwargs passes on unmodified master, so it's coverage rather than a regression guard; worth adjusting the Testing table to reflect that.)
• The same bug class survives for four sibling keywords at that call site. I triggered each through public deploy(): endpoint_type, resources, inference_component_name, and built_model all still raise the identical TypeError. These are pre-existing, but the two other branches of deploy() already build a deploy_kwargs dict and call deploy_kwargs.update(kwargs), which fixes the whole class and matches the file's own convention. Reasonable either as part of this PR or as a follow-up.

One merge gate: codestyle-doc-tests (sagemaker-serve) needs a green re-run. It failed in PRE_BUILD on THESE PACKAGES DO NOT MATCH THE HASHES FROM THE REQUIREMENTS FILE, which is infra and unrelated to your change, but that check passes on comparable PRs so it shouldn't be waived as pre-existing. wait-for-approval and collab-check both passed and the branch is already up to date with master, so no rebase needed.

Optional nit: the fix relies on the kwargs.pop side effects being evaluated before the **kwargs merge in the same call. That's guaranteed in CPython (I verified 3.9/3.11/3.12), but it's implicit — hoisting to two statements above the return is free and self-documenting.