feat(feature-store): Add lineage registration to DatasetBuilder#6014
Open
Vishakha263 wants to merge 2 commits into
Open
feat(feature-store): Add lineage registration to DatasetBuilder#6014Vishakha263 wants to merge 2 commits into
Vishakha263 wants to merge 2 commits into
Conversation
Add register_as_dataset parameter to DatasetBuilder.create() that registers the built dataset as a HubContent with source feature group ARNs in ContentMetadata. This enables auto-lineage between Feature Groups and Training Jobs via SM Dataset (HubContent) as bridge. Changes: - Add register_as_dataset flag to DatasetBuilder.create() - Collect source FG ARNs from base and merged feature groups - Call HubContent.import_hub_content() with sourceFeatureGroups, extractionMethod, and athenaQueryExecutionId metadata - Graceful fallback on AccessDenied or generic errors (logs warning) - Add 10 unit tests covering opt-in, ARN collection, dedup, DataFrame path, HubContent invocation, error handling
Vishakha263
had a problem deploying
to
manual-approval
GitHub Actions
Error
Vishakha263
had a problem deploying
to
manual-approval
GitHub Actions
Error
Vishakha263
requested a deployment
to
manual-approval
GitHub Actions
Waiting
Vishakha263
requested a deployment
to
manual-approval
GitHub Actions
Waiting
huiqunli
reviewed
Jul 10, 2026
| try: | ||
| from sagemaker.core.resources import HubContent | ||
|
|
||
| HubContent.import_hub_content( |
Contributor
There was a problem hiding this comment.
Instead of calling HubContent.import_hub_content directly, might consider to use DataSet.create from sagemaker.ai_registry.dataset. It handles document construction and schema compliance internally. We'd need to extend DataSet.create to accept a metadata param for passing sourceFeatureGroups though.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Add optional lineage registration to DatasetBuilder that creates a SM Dataset (HubContent) after extracting features from Feature Groups. This enables automatic lineage tracking between Feature Groups and Training Jobs.
Motivation
SageMaker's auto-lineage currently has a gap: there's no connection between Feature Groups and Training Jobs. When a customer extracts features via DatasetBuilder.to_csv_file() and uses the resulting CSV for training, that relationship is lost. This PR fills that gap by registering the extraction as a SM Dataset with source Feature Group metadata.
Implementation details:
Files changed:
Testing
Added 10 unit tests in TestDatasetBuilderRegisterAsDataset:
Permissions required (for customers who opt in)
sagemaker:ImportHubContent on the execution role. If missing, a warning is logged and the CSV is still returned.
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.