Skip to content

fix(serve): pass instance_type to _deploy once when creating an inference component - #6203

Open
Om-singhaI wants to merge 3 commits into
aws:masterfrom
Om-singhaI:fix/deploy-for-ic-duplicate-kwargs
Open

fix(serve): pass instance_type to _deploy once when creating an inference component#6203
Om-singhaI wants to merge 3 commits into
aws:masterfrom
Om-singhaI:fix/deploy-for-ic-duplicate-kwargs

Conversation

@Om-singhaI

Copy link
Copy Markdown

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.

…ence component

Deploying a CustomOrchestrator in InferenceComponent mode through
ModelBuilder.deploy() always raised

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

In _deploy_for_ic the branch that creates a new inference component read
instance_type and initial_instance_count out of kwargs with kwargs.get,
passed them to self._deploy as explicit keyword arguments, and then spread
the unchanged **kwargs into the same call. kwargs.get leaves the key in
place, so Python saw each of those arguments twice. The only caller of
this branch, ModelBuilder.deploy, always supplies both keys, so the
failure was deterministic.

Consume the two keys with kwargs.pop instead so each is forwarded exactly
once. The defaults (self.instance_type and 1) are unchanged and kwargs is
a local dict, so the caller is unaffected. _deploy_core_endpoint reads the
values back through kwargs.get and receives the same data as before.

Add unit tests covering the create branch of _deploy_for_ic with explicit
instance arguments (fails with the TypeError before this change) and with
the default fallbacks.

Fixes aws#6199
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.

_deploy_for_ic passes instance_type twice to _deploy() causing TypeError when deploying CustomOrchestrator as Inference Component

1 participant