Skip to content

test(serve): add ModelBuilder.register framework attribute default tests - #6117

Open
aviruthen wants to merge 1 commit into
aws:masterfrom
aviruthen:test/modelbuilder-hf-env-and-register-tests
Open

test(serve): add ModelBuilder.register framework attribute default tests#6117
aviruthen wants to merge 1 commit into
aws:masterfrom
aviruthen:test/modelbuilder-hf-env-and-register-tests

Conversation

@aviruthen

Copy link
Copy Markdown
Collaborator

Issue #, if available:

Description of changes:

Add regression tests verifying a freshly constructed ModelBuilder exposes
framework and framework_version defaulted to None, so register() does not raise
AttributeError when the user never sets them.

Closes #5391

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

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown

Claude finished @aviruthen's task in 2m 25s —— View job


Code Review — test(serve): ModelBuilder.register framework defaults

  • Read PR diff and the referenced source (sagemaker-serve/src/sagemaker/serve/model_builder.py)
  • Verify test targets match the real register() implementation
  • Assess correctness, conventions, and coverage

Summary

This is a tests-only PR and it looks clean. It adds a focused regression test for #5391, locking in that a freshly constructed ModelBuilder exposes framework/framework_version defaulted to None so register() doesn't raise AttributeError.

I verified the tests against the source and they exercise the intended behavior accurately:

  • Root-cause guard_initialize_defaults sets both attributes to None at construction (model_builder.py:542-545). test_framework_attributes_default_to_none_on_construction asserts exactly this.
  • End-to-end path — with model_package_group_name set and the default (non-PROPRIETARY) model_type, register() builds a single-container container_def and passes self.framework / self.framework_version into update_container_with_inference_params (model_builder.py:6178-6185). test_register_does_not_require_user_set_framework_attributes captures both values there, confirming they were read (as None), not missing.
  • Decorator handling is anticipated correctly — the test pins session.sagemaker_config = {} (needed by the @_telemetry_emitter wrapper) and uses a plain MagicMock session so the @runnable_by_pipeline decorator takes the non-pipeline branch and returns the ARN (model_builder.py:6232-6242). Good attention to detail.
  • Patching resolve_and_validate_role keeps construction offline (matches model_builder.py:530), consistent with how the SDK avoids IAM work when a role isn't supplied.

Minor observations (non-blocking)

  • The tests depend on several internal patch targets (create_model_package_from_containers, get_model_package_args, update_container_with_inference_params, _prepare_container_def). That's expected for a unit test at this layer, but it does couple the test to register()'s internal structure — a future refactor of that method could require updating these patches. The docstrings already explain the intent well, which mitigates this.
  • Consider a brief comment noting that the empty get_model_package_args return ({}) is only safe because create_model_package_from_containers is also mocked — otherwise a reader might expect real package args to flow through.

Neither is worth blocking on. No correctness, security, or performance concerns — the PR adds no production code and the tests are well-scoped and accurate.

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.

ModelBuilder cannot call register() method Without Setting framework and framework_version Attributes

1 participant