Feature store lakeformation

[BassemHalim]

Description

This PR adds Lake Formation integration to SageMaker Feature Store, enabling customers to govern access to their offline store data through AWS Lake Formation instead of relying solely on IAM policies.

This simplifies the manual process described in this blog https://aws.amazon.com/blogs/machine-learning/control-access-to-amazon-sagemaker-feature-store-offline-using-aws-lake-formation/

New Features

class LakeFormationConfig(Base):
    """Configuration for Lake Formation governance on Feature Group offline stores.

    enabled: bool = False
    use_service_linked_role: bool = True
    registration_role_arn: Optional[str] = None
    show_s3_policy: bool = False
    disable_hybrid_access_mode: bool = True

FeatureGroup.create() - added a new lake_formation_config parameter

• Enables Lake Formation governance at Feature Group creation time
• Automatically waits for Feature Group to reach "Created" status before configuring Lake Formation

FeatureGroup.enable_lake_formation() method

• Enables Lake Formation on existing Feature Groups
• Three-phase setup:
  1. Registers S3 location with Lake Formation
  2. Grants permissions to the offline store role on the Glue table
  3. Revokes IAMAllowedPrincipal permissions from the Glue table
• Fail-fast behavior with clear error reporting at each phase
• Optional show_s3_policy parameter prints recommended S3 deny policy

Usage

Enable at creation:

from sagemaker.mlops.feature_store import FeatureGroup, LakeFormationConfig

lf_config = LakeFormationConfig()
lf_config.enabled = True

fg = FeatureGroup.create(
    feature_group_name="my-feature-group",
    # ... other params ...
    lake_formation_config=lf_config,
)

Enable on existing Feature Group:

fg = FeatureGroup.get("my-feature-group")
fg.enable_lake_formation(show_s3_policy=True)

Testing

• Unit tests: comprehensive coverage of all new methods and validation logic
• Integration tests: end-to-end tests for both creation workflows and negative scenarios

Notes

• S3 deny policy is provided as a recommendation (not applied automatically) to avoid breaking existing workflows

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

[nargokul]

Please fix integ tests

[nargokul]

This would be confusing to the users to choose between sagemaker.core.FeatureGroup and sagemaker.mlops.FeatureGroup

Can this be utils or called something else ?

[BassemHalim]

The v3 sagemaker.mlops version already re-exports the FeatureGroup from sagemaker-core ref

This class is meant as an extension to add an extra method enable_lake_formation and update the create method to enable lakeformation during creation

[nargokul]

re-exporting is fine.
The class is a sagemaker-core resource class which does basic crud operations exactly mimicing boto.
If we introfuce this change , then there is a difference in the behavior between core and mlops .

Why does this class need to be called FeatureGroup ?

[BassemHalim]

The idea is that it should be a drop-in replacement for the one in sagemaker-core just with extra functionality so that the user can enable lakeformation during feature group creation instead of creating then enabling. I was thinking this might be better UX. What do you suggest renaming the class or moving that to a utility function? I'm open to both but not sure what it can be renamed to

[nargokul]

I think we can go with something like GovernedFeatureGroup or ManagedFeatureGroup.

We can discuss internally .

[nargokul]

Is this the string that gets created by default ?

Can this functionality be gotten from an AWS library instead of hardcoding ?

If this changes , these changes would break

[BassemHalim]

Yes it gets created by default by Lakeformation ref
I didn't find any functionality in the botocore sdk to get that role instead of hardcoding

I'm thinking it's provided in the lakeformation public docs so it's unlikely to change.