fix: accept core shapes and omitted content types in ModelBuilder.register

[mohamedzeidan2021]

Issue #, if available: Fixes #6479, Fixes #6480

Description

Two independent AttributeError defects in ModelBuilder.register(). They sit on the same method and the same customer workflow (register / deploy-from-package), so they are fixed together.

1. register() rejected the ModelMetrics type its own signature declares (#6479)

ModelBuilder.register() types model_metrics as sagemaker.core.shapes.ModelMetrics, but get_model_package_args() converted it with model_metrics._to_request_dict(). The sagemaker.core.shapes classes are pydantic models and do not implement that method, so passing the documented type raised:

AttributeError: 'ModelMetrics' object has no attribute '_to_request_dict'

Only the older sagemaker.core.model_metrics.ModelMetrics compatibility class implements _to_request_dict(), and that class is not the type in the public signature.

The same mismatch applies to the three sibling arguments converted the same way in the same function — metadata_properties, drift_check_baselines and model_life_cycle are all declared in register() as sagemaker.core.shapes types, and none of those shapes implements _to_request_dict() either. Fixing only model_metrics would have left the customer hitting the identical error on the next argument, so all four now go through one helper.

get_model_package_args() gained a module-private _to_request_dict(value) that:

• calls value._to_request_dict() when the value implements it — the compatibility classes keep their existing behaviour unchanged, so they remain supported for the rest of their deprecation period;
• otherwise falls back to the core serializer sagemaker.core.utils.utils.serialize(), which is what handles the sagemaker.core.shapes types;
• passes a plain dict that is already in request form straight through.

Why not the patch suggested on the issue. The issue proposes replacing the call with serialize(model_metrics) unconditionally. I verified that this fixes the broken path but regresses the one that works today. The compatibility class carries flat fields (model_statistics, model_constraints, …) and serialize() renders them literally:

# serialize(sagemaker.core.model_metrics.ModelMetrics(model_statistics=...))
{"ModelStatistics": {"ContentType": "application/json", "S3Uri": "s3://bucket/evaluation.json"}}

That is not a valid CreateModelPackage ModelMetrics shape — the API expects the nested ModelQuality.Statistics block that the compatibility class's own _to_request_dict() builds. So serialize() is the right direction for the core shapes only, and the dispatch above is required rather than a straight substitution.

2. register() failed when content and response types were omitted (#6480)

content_types and response_types are declared optional with None defaults, but register() only assigned self.content_types / self.response_types when the argument was not None, and then read those attributes unconditionally:

if content_types is not None:
    self.content_types = content_types
...
model_pkg_args = get_model_package_args(
    self.content_types,      # AttributeError on a first call that omits them
    self.response_types,

A first registration that omits them raised:

AttributeError: 'ModelBuilder' object has no attribute 'content_types'

The read now uses getattr(self, "content_types", None). This is deliberately a two-line change at the read site rather than the alternative suggested on the issue (initialising both in ModelBuilder.__init__):

• it preserves the existing behaviour where a later register() that omits the types reuses values supplied to an earlier call, which the issue explicitly asks for;
• ModelBuilder is a dataclass, so initialising in __init__ would add two new public constructor fields next to the unrelated and already-public content_type field — Python's own suggestion on the traceback is Did you mean: 'content_type'?, so the near-duplicate names would be a lasting source of confusion;
• it does not invent attributes on the builder as a side effect of calling register().

When the types are absent they are now simply left out of the request, as the issue expects.

Testing

All validation was done offline. There were no usable AWS credentials in this environment, so there is no live CreateModelPackage round-trip in this evidence — request correctness is established against the service contract, not against the live service.

Reproduction. Both defects reproduce against released sagemaker==3.25.0 with the exact errors quoted in the issues (model_registry.py:79 for #6479, model_builder.py:6989 for #6480), and both are present on master at 8e07c37.

Request correctness. ModelBuilder.register() was driven with a mocked SageMaker client, and the resulting wire request read out of create_model_package.call_args.kwargs was validated with botocore.validate.ParamValidator against botocore.session.get_session().get_service_model("sagemaker").operation_model("CreateModelPackage").input_shape. With the fix, the declared core shape produces exactly the structure the issue asks for, and the request validates against the contract:

"ModelMetrics": {"ModelQuality": {"Statistics": {
  "ContentType": "application/json", "S3Uri": "s3://bucket/evaluation.json"}}}

The compatibility-class request and the content-types-supplied request are byte-identical before and after the change, confirming the working paths are untouched. Reuse of types across repeated register() calls was checked explicitly and still works, and register() no longer leaves content_types / response_types attributes behind when they were never supplied.

New tests — 13 total.

• sagemaker-core/tests/unit/test_model_registry.py — 4 added: core shapes.ModelMetrics, the model_metrics compatibility class, core shapes.DriftCheckBaselines, and a request-form dict.
• sagemaker-serve/tests/unit/test_model_builder_register.py — 9 added: register() with the types omitted, with them supplied, reusing them from an earlier call, not inventing attributes, with both ModelMetrics types, and the serialized request shape for both.

Negative check. Each fix was reverted independently to confirm the new tests actually guard the behaviour:

State	test_model_builder_register.py	test_model_registry.py
Both fixes applied	9 passed	25 passed
Both reverted	7 failed, 2 passed	3 failed, 22 passed
Only the #6479 fix reverted	1 failed, 8 passed	3 failed, 22 passed
Only the #6480 fix reverted	7 failed, 2 passed	25 passed

Regression runs. Full tests/unit of both touched packages, with the fixes applied and again with them reverted:

• sagemaker-core: 2 failed, 3963 passed, 21 skipped. Reverted: 5 failed, 3960 passed — the same 2 failures plus exactly the 3 new tests that guard ModelBuilder.register() rejects its declared core ModelMetrics #6479. No failure appears with the fixes that was not there without them.
• sagemaker-serve: 70 failed, 2079 passed, 25 skipped — the same 70 failures, identical set, with the fixes reverted (2070 passed; the 9-test delta is exactly the new file).

Those pre-existing failures are environmental in this sandbox, not caused by this change: 65 of the 70 sagemaker-serve failures are botocore.exceptions.ClientError: ExpiredToken from GetCallerIdentity (the tests construct ModelBuilder, which resolves a serving role via STS), and the other 5 are ModuleNotFoundError: No module named 'tensorflow'. The 2 sagemaker-core failures come from running repo tests against the installed 3.25.0 wheel (a djl-neuronx region added by the preceding commit, and a removed-module guidance message not present in the wheel).

Lint. Run as CI does, from inside each package: black --line-length 100 --check clean, flake8 --max-line-length 120 clean, pydocstyle clean on both changed src files. pylint --rcfile=../.pylintrc introduced no new findings — model_registry.py 9.71 → 9.72, model_builder.py 9.84 → 9.84; the only message text change is the pre-existing wrong-import-order finding now also enumerating the new import.

Merge Checklist

Put an x in the boxes that apply.

General

• I have read the CONTRIBUTING doc
• I certify that the changes I'm 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 passed the region in to all S3 and STS clients that I've initialized as part of this change
• I have updated any necessary documentation, including READMEs and API docs (if appropriate)

Tests

• I have added tests that prove my fix is effective or that my feature works (if appropriate)
• I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes
• I have checked that my tests are not configured for a specific region or account (if appropriate)
• I have used unique_name_from_base to create resource names in integ tests (if appropriate)

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

🤖 AI code review

The PR fixes two AttributeError defects in ModelBuilder.register(): (1) get_model_package_args() now dispatches via a new module-private _to_request_dict() helper that calls value._to_request_dict() for the legacy compatibility classes and falls back to the core serializer for sagemaker.core.shapes pydantic types (and plain dicts); (2) the read of content_types/response_types now uses getattr(..., None) so a first registration that omits them no longer raises. I verified the serializer produces the expected PascalCase request shape, the compatibility classes implement _to_request_dict() so hasattr dispatch is correct, and the attribute-assignment behavior at the register site. The implementation is correct and well-tested; only minor coverage gaps noted.

Findings not attached to a changed line:

• sagemaker-core/tests/unit/test_model_registry.py:83 🟡 Low: The PR extends the _to_request_dict dispatch to four arguments (model_metrics, drift_check_baselines, metadata_properties, model_life_cycle), but only model_metrics and drift_check_baselines get a core-shape test here. metadata_properties and model_life_cycle core-shape serialization are untested. Since they share the helper the risk is low, but a shape like ModelLifeCycle may serialize to a structure worth asserting (e.g. to confirm the serializer output matches what CreateModelPackage expects, as noted in the PR body for ModelMetrics). Consider adding at least one assertion for each.

_{Reviewed commit 8eddfc4. Automated review; verify before acting.}