Skip to content

feature: Add HubContentStep and Session.import_hub_content - #6565

Open
rohangujarathi wants to merge 1 commit into
aws:masterfrom
rohangujarathi:feature-hub-content-step
Open

rohangujarathi wants to merge 1 commit into
aws:masterfrom
rohangujarathi:feature-hub-content-step

Conversation

@rohangujarathi

Copy link
Copy Markdown
Member

Issue #, if available:

Description of changes:

Adds a pipeline step that imports content into a SageMaker hub.

  • Session.import_hub_content(...) (sagemaker-core): builds an ImportHubContent request and routes it through _intercept_create_request, so under a PipelineSession it returns step arguments instead of calling the service. hub_content_document can be a string or a dict; a dict is serialized when the request is sent.
  • HubContentStep (sagemaker-mlops): a ConfigurableRetryStep that accepts step_args from import_hub_content only, supports cache_config and retry_policies, and exposes properties shaped like DescribeHubContentResponse. A dict document is encoded in arguments the same way JobStep encodes JobConfigDocument, so pipeline variables inside it are resolved at execution time.
  • StepTypeEnum.HUB_CONTENT = "HubContent" and the sagemaker.mlops.workflow export.

Example:

pipeline_session = PipelineSession()
step_args = pipeline_session.import_hub_content(
    hub_name="my-hub",
    hub_content_name="my-model",
    hub_content_type="Model",
    document_schema_version="2.4.0",
    hub_content_document={...},
)
step = HubContentStep(name="ImportModel", step_args=step_args)
# step.properties.HubContentArn, step.properties.HubContentStatus, ...

When hub_content_version is omitted, the service assigns the next version on each execution.

The HubContent step type is being rolled out in SageMaker Pipelines; please hold the merge until it is available, otherwise CreatePipeline rejects the definition and the new integ tests fail.

Testing done:

  • Unit tests in sagemaker-mlops/tests/unit/workflow/test_hub_content_step.py: capture without a service call, request shape, unset optional fields, dict and string documents, pipeline-variable document encoding, step request with cache and retry config, producer validation, properties, pipeline definition, plain Session call.
  • Unit tests in TestImportHubContent (sagemaker-core/tests/unit/helper/test_session_helper.py): required fields, dict serialization, all optional fields and tags, empty optional values, service error propagation.
  • test_job_step.py enum member set updated; existing test_inference_lineage_steps.py, test_job_step.py and test_session_helper.py pass.
  • Integ tests in sagemaker-mlops/tests/integ/workflow/test_hub_content_step.py: an import with a parameterized document feeding a ConditionStep on HubContentStatus and HubContentType, unpinned reruns importing a new version, and a missing hub failing the step. All 3 pass.

Merge Checklist

General

  • I have read the CONTRIBUTING doc
  • I certify that the changes I am 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)
  • If adding any dependency in requirements.txt files, I have spell checked and ensured they exist in PyPi

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 AI code review

This PR adds Session.import_hub_content(...) in sagemaker-core and a HubContentStep pipeline step in sagemaker-mlops, plus a StepTypeEnum.HUB_CONTENT value and workflow export. The implementation closely follows existing patterns: import_hub_content routes through _intercept_create_request like other producers, dict documents are serialized via the shared convert_job_config_document_to_string (so pipeline variables become a Join), step_args is validated with validate_step_args_input, and properties is shaped from the botocore DescribeHubContentResponse shape. I verified the imports (json, format_tags, _append_project_tags, RequestType, _JobStepArguments, Properties), the base-class signatures, and that arguments reads a shallow copy of step_args.args without mutating the captured request. Unit and integ tests are comprehensive. No correctness, security, or backward-compatibility issues found; the implementation is clean.

Note: the PR body asks maintainers to hold the merge until the HubContent step type is available service-side (otherwise CreatePipeline/integ tests fail). That is operational context for reviewers, not a code issue.

Reviewed commit 909574f. Automated review; verify before acting.

papriwal
papriwal previously approved these changes Nov 7, 2026

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 AI code review

This PR adds Session.import_hub_content() (sagemaker-core) and HubContentStep (sagemaker-mlops), a ConfigurableRetryStep that imports content into a SageMaker hub within a pipeline. The producer routes through _intercept_create_request so it returns step args under a PipelineSession, and a dict HubContentDocument is encoded the same way JobStep encodes JobConfigDocument so pipeline variables resolve at execution time. The implementation closely follows the established JobStep/TrainingStep patterns, the StepTypeEnum addition and workflow export are correct, and the unit tests cover request shape, dict/string documents, pipeline-variable encoding, producer validation, properties, and the plain-Session path. I found no correctness or security issues; the code is consistent with existing conventions. Note: the PR body asks reviewers to hold the merge until the HubContent step type is rolled out in Pipelines (otherwise the integ tests fail) — that is a release-coordination item, not a code concern.

Reviewed commit e5d26ec. Automated review; verify before acting.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 AI code review

This PR adds Session.import_hub_content() in sagemaker-core (routing through _intercept_create_request so it returns step args under a PipelineSession, serializing a dict HubContentDocument at request time) and a HubContentStep in sagemaker-mlops, plus StepTypeEnum.HUB_CONTENT and the workflow export. The implementation closely mirrors the established JobStep/JobConfigDocument pattern: the arguments property shallow-copies the captured request and encodes a dict document via convert_job_config_document_to_string for pipeline-variable resolution, and to_request folds in cache config. Imports (json, format_tags, _append_project_tags, validate_step_args_input) are all present at the base commit, tag/empty-value handling matches existing methods, and unit + integ tests cover capture, serialization, optional fields, producer validation, properties, and pipeline definition. I found no correctness, security, or backward-compatibility issues. Note: the PR body asks reviewers to hold the merge until the HubContent step type is rolled out server-side — a process note, not something that changes the code review.

Reviewed commit e9882e0. Automated review; verify before acting.

This branch was successfully deployed

1 active deployment
auto-approve — e9882e0d Deployed Nov 10, 2026 by rohangujarathi via wait-for-approval #834
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.

3 participants