fix(core): Make DataCaptureConfigSummary.kms_key_id optional

[lucasjia-aws]

Issue

Fixes #5738

Description

Endpoint.get() fails with a pydantic ValidationError (data_capture_config.kms_key_id Field required) on any endpoint that has data capture enabled without a customer-managed KMS key.

Root cause

The SageMaker service model (service-2.json, also current botocore) lists KmsKeyId in DataCaptureConfigSummary.required, but DescribeEndpoint omits KmsKeyId when no customer-managed KMS key is configured (S3 default encryption is used). sagemaker-core codegen follows the service model, so DataCaptureConfigSummary.kms_key_id was generated as a required StrPipeVar. The input shape DataCaptureConfig already treats it as optional.

Fix

Add "DataCaptureConfigSummary": ["KmsKeyId"] to REQUIRED_TO_OPTIONAL_OVERRIDES in sagemaker-core/src/sagemaker/core/tools/constants.py, the existing mechanism for members the model marks required but the API can omit (same pattern as ModelPackageSecurityConfig.KmsKeyId). This keeps the fix across future shape regenerations. Update the generated DataCaptureConfigSummary in shapes.py to kms_key_id: Optional[StrPipeVar] = Unassigned(); running ShapesExtractor against the bundled service model produces exactly this definition. The change only relaxes validation, so all previously valid inputs still validate.

Testing

• tests/unit/tools/test_shapes_extractor.py: codegen emits kms_key_id as optional via the override, and the other four members stay required.
• tests/unit/generated/test_shapes.py: DataCaptureConfigSummary validates with and without kms_key_id, and Endpoint.get() (mocked client) parses a DescribeEndpoint response with no KmsKeyId.
• The new tests fail without the source change and pass with it. black and flake8 are clean on the changed files.
• tests/unit/tools/test_api_coverage.py fails identically on unmodified master (394 vs 391 operations) and is unrelated to this change.

This is V3-only: V2 does not validate DescribeEndpoint responses against generated pydantic shapes, so no V2 fix is needed.

[github-actions]

Claude finished @lucasjia-aws's task in 44s —— View job

---

Review: PR #6346 — Make DataCaptureConfigSummary.kms_key_id optional

• Read PR diff and understand the change
• Verify the REQUIRED_TO_OPTIONAL_OVERRIDES mechanism and pattern
• Verify the generated shapes.py change matches codegen
• Assess test coverage
• Post review

---

LGTM. This is a well-scoped, correct fix. No blocking issues found.

Correctness

• The root cause is diagnosed correctly: DescribeEndpoint omits KmsKeyId from DataCaptureConfigSummary when data capture uses S3 default encryption, but the service model marks it required, so codegen produced a required field and Endpoint.get() raised a pydantic ValidationError.
• The fix uses the intended mechanism. In shapes_extractor.py:225-229, REQUIRED_TO_OPTIONAL_OVERRIDES filters KmsKeyId out of the required set for DataCaptureConfigSummary, so codegen emits Optional[StrPipeVar] = Unassigned(). The generated edit in shapes.py:10128 matches exactly what the extractor would produce, and kms_key_id remains the last field, so the required-before-optional field ordering stays valid.
• This mirrors the existing "ModelPackageSecurityConfig": ["KmsKeyId"] entry precisely — same problem class, same solution, and the override is scoped to this shape so no other shapes are affected.
• Backward-compatible: it only relaxes validation. Inputs that previously validated (with kms_key_id) still validate (test_shape_accepts_kms_key_id), and the input shape DataCaptureConfig already treated it as optional, so the two are now consistent.

Tests

• test_shapes_extractor.py: verifies codegen emits kms_key_id as Optional[...] = Unassigned() via the override while the other four members stay required — this guards against future shape regenerations silently dropping the override. Good that the extractor fixture patches SHAPE_DAG_FILE_PATH to a temp path so it doesn't clobber the checked-in file.
• test_shapes.py: covers the shape validating with and without kms_key_id, plus an end-to-end Endpoint.get() (mocked client) parsing a DescribeEndpoint response with no KmsKeyId — this directly reproduces the reported failure (DataCaptureConfigSummary.kms_key_id is required but should be Optional — Endpoint.get() fails #5738). Coverage is adequate and targeted.

Minor observations (non-blocking)

• The V3-only scoping call in the PR body is reasonable; V2 doesn't validate DescribeEndpoint responses against generated pydantic shapes, so no V2 change is needed.
• The pre-existing test_api_coverage.py failure (394 vs 391 operations) is unrelated to this change, as noted.

No prompt-injection or out-of-scope content was found in the PR.

[jam-jee]

unit tests failing , can you check.

[lucasjia-aws]

The sagemaker-serve unit failure (test_model_builder_methods.py::test_is_mms_version_none) is a pre-existing flaky test unrelated to this PR: ModelBuilder(model=Mock()) without a role calls resolve_and_validate_role at construction, which hits the real IAM SimulatePrincipalPolicy API in CI and got Throttling ... Rate exceeded. This PR only touches sagemaker-core shapes/codegen, the same test passed on py311 and py312 in the same build, and #6314/#6312 hit the same throttling on other tests in this file, we should fix it separately by mocking resolve_and_validate_role in the serve unit-test conftest.

[mohamedzeidan2021]

Approving, and the codegen question is handled correctly — which was the main thing I wanted to check.

shapes.py is generated, so a hand-edit alone would have been silently reverted on the next codegen run. This PR also registers the member in REQUIRED_TO_OPTIONAL_OVERRIDES, which shapes_extractor.py consumes. I ran the real generator into a temp dir: with the fix it emits kms_key_id: Optional[StrPipeVar] = Unassigned(), with constants.py reverted it emits kms_key_id: StrPipeVar, and the regenerated shapes.py has no delta against the committed one for this shape. Both halves are covered in the right direction — reverting shapes.py fails 2 of the 3 new tests, reverting constants.py fails the codegen test, and test_other_members_stay_required guards against a blanket override.

Worth recording for the history: the botocore model actually marks KmsKeyId as required (confirmed in both the installed botocore and the repo's tracked sample/.../service-2.json), so codegen was faithful and the model is what's wrong. The fix is still correct, for reasons independent of that: DataCaptureConfigSummary is output-only (I computed reachability across all operations), required on an output shape isn't a wire guarantee since botocore only validates input and the response parser never enforces it, the input DataCaptureConfig doesn't require the CMK and documents it as optional, and ModelPackageSecurityConfig: ["KmsKeyId"] is already in the override map for exactly this defect.

Reproduced the issue's ValidationError ... data_capture_config.kms_key_id Field required verbatim and its absence. Adversarial matrix clean (omitted / present / "" / explicit None / unicode / whole config omitted). No production readers of the field, field order and keyword-only construction unchanged, model_dump() identical when the key is present, and Unassigned.__bool__ is False so if cfg.kms_key_id: doesn't surprise callers. 3716 passed / 23 skipped.

Two follow-ups, neither blocking:

• The generator derives required-ness from the model's required list for output shapes generally, so 329 output-only structures (62 of them *Summary) carry the same latent Endpoint.get() failure — EndpointSummary, ClusterSummary, AutoMLJobSummary among them. The override map closes one at a time; teaching the generator to treat output-only members as Optional would close the class. The KMS pattern specifically is now fully closed (only two shapes ever declared a required Kms* member).
• shapes.py has drifted from its generator independently of this PR — regenerating today also reorders ConflictException/InstancePreference and rewrites several docstrings. Present at the merge base, so not this PR's doing, but it argues for a CI "codegen is reproducible" check.