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

LakeFormationConfig — declarative configuration for Lake Formation governance:

class LakeFormationConfig(Base):                                                                                                       
    enabled: bool = False                                                                                                              
    use_service_linked_role: bool = True                                                                                               
    registration_role_arn: Optional[str] = None                                                                                        
    disable_hybrid_access_mode: bool          # Required — no default                                                                  
    acknowledge_risk: Optional[bool] = None   # None=interactive prompt, True=skip, False=abort                                        

FeatureGroupManager.create() — added 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

FeatureGroupManager.enable_lake_formation() — new method

• Enables Lake Formation on existing Feature Groups

• Three-phase setup:

  1. Registers S3 location with Lake Formation
  2. Grants permissions to the execution role on the Glue table
  3. Optionally revokes IAMAllowedPrincipal permissions from the Glue table (controlled by disable_hybrid_access_mode)

• Fail-fast behavior with clear error reporting at each phase

• Interactive confirmation prompts warn users about risks (controllable via acknowledge_risk)

• Logs recommended S3 deny policy as a warning

Usage

Enable at creation:

  from sagemaker.mlops.feature_store import FeatureGroupManager, LakeFormationConfig                                                     
                                                                                                                                         
  lf_config = LakeFormationConfig(                                                                                                       
      enabled=True,                                                                                                                      
      disable_hybrid_access_mode=True,                                                                                                   
      acknowledge_risk=True,  # skip interactive prompt                                                                                  
  )                                                                                                                                      
                                                                                                                                         
  fg = FeatureGroupManager.create(                                                                                                       
      feature_group_name="my-feature-group",                                                                                             
      # ... other params ...                                                                                                             
      lake_formation_config=lf_config,                                                                                                   
  )                                                                                                                                      
                                                                                                                                         
  #Enable on existing Feature Group:                                                                                                      
                                                                                                                                         
  fg = FeatureGroupManager.get("my-feature-group")                                                                                       
  result = fg.enable_lake_formation(                                                                                                     
      disable_hybrid_access_mode=True,                                                                                                   
      acknowledge_risk=True,                                                                                                             
  )                                                                                                                                      

Testing

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

Notes

• S3 deny policy is logged as a recommendation (not applied automatically) to avoid breaking existing workflows
• disable_hybrid_access_mode is a required field with no default — users must explicitly choose whether to revoke IAMAllowedPrincipal permissions
• acknowledge_risk controls interactive confirmation: None (default) prompts via input(), True skips the prompt, False aborts with RuntimeError

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

[mollyheamazon]

input() is not safe for non-interactive environments (CI, pipelines, SageMaker jobs).
Please require acknowledge_risk to be explicitly set to True or False and remove the None interactive path.

[mollyheamazon]

Same comment as above:
input() is not safe for non-interactive environments (CI, pipelines, SageMaker jobs).
Please require acknowledge_risk to be explicitly set to True or False and remove the None interactive path.