feat: add throughput management support for feature group

doc/api/prep_data/feature_store.rst

@@ -75,6 +75,14 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfig
+ :members:
+ :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfigUpdate
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.OnlineStoreConfig
  :members:
  :show-inheritance:
@@ -99,6 +107,10 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputModeEnum
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.ResourceEnum
  :members:
  :show-inheritance:

src/sagemaker/feature_store/feature_group.py

@@ -64,6 +64,8 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputConfigUpdate,
 )
 from sagemaker.utils import resolve_value_from_config, format_tags, Tags
 
@@ -541,6 +543,7 @@ def create(
  tags: Optional[Tags] = None,
  table_format: TableFormatEnum = None,
  online_store_storage_type: OnlineStoreStorageTypeEnum = None,
+ throughput_config: ThroughputConfig = None,
  ) -> Dict[str, Any]:
  """Create a SageMaker FeatureStore FeatureGroup.
 
@@ -570,6 +573,8 @@ def create(
  table_format (TableFormatEnum): format of the offline store table (default: None).
  online_store_storage_type (OnlineStoreStorageTypeEnum): storage type for the
  online store (default: None).
+ throughput_config (ThroughputConfig): throughput configuration of the
+ feature group (default: None).
 
  Returns:
  Response dict from service.
@@ -618,6 +623,9 @@ def create(
  )
  create_feature_store_args.update({"online_store_config": online_store_config.to_dict()})
 
+ if throughput_config:
+ create_feature_store_args.update({"throughput_config": throughput_config.to_dict()})
+
  # offline store configuration
  if s3_uri:
  s3_storage_config = S3StorageConfig(s3_uri=s3_uri)
@@ -656,17 +664,17 @@ def update(
  self,
  feature_additions: Sequence[FeatureDefinition] = None,
  online_store_config: OnlineStoreConfigUpdate = None,
+ throughput_config: ThroughputConfigUpdate = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup and add new features from the given feature definitions.
 
  Args:
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
  online_store_config (OnlineStoreConfigUpdate): online store config to be updated.
-
+ throughput_config (ThroughputConfigUpdate): target throughput configuration
  Returns:
  Response dict from service.
  """
-
  if feature_additions is None:
  feature_additions_parameter = None
  else:
@@ -679,10 +687,15 @@ def update(
  else:
  online_store_config_parameter = online_store_config.to_dict()
 
+ throughput_config_parameter = (
+ None if throughput_config is None else throughput_config.to_dict()
+ )
+
  return self.sagemaker_session.update_feature_group(
  feature_group_name=self.name,
  feature_additions=feature_additions_parameter,
  online_store_config=online_store_config_parameter,
+ throughput_config=throughput_config_parameter,
  )
 
  def update_feature_metadata(

src/sagemaker/feature_store/inputs.py

@@ -453,3 +453,79 @@ class ExpirationTimeResponseEnum(Enum):
 
  DISABLED = "Disabled"
  ENABLED = "Enabled"
+
+
+class ThroughputModeEnum(Enum):
+ """Enum of throughput modes supported by feature group.
+
+ Throughput mode of feature group can be ON_DEMAND or PROVISIONED.
+ """
+
+ ON_DEMAND = "OnDemand"
+ PROVISIONED = "Provisioned"
+
+
+@attr.s
+class ThroughputConfig(Config):
+ """Throughput configuration of the feature group.
+
+ Throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )
+
+
+@attr.s
+class ThroughputConfigUpdate(Config):
+ """Target throughput configuration for the feature group.
+
+ Target throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Target throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )

src/sagemaker/session.py

@@ -5679,6 +5679,7 @@ def create_feature_group(
  role_arn: str = None,
  online_store_config: Dict[str, str] = None,
  offline_store_config: Dict[str, str] = None,
+ throughput_config: Dict[str, Any] = None,
  description: str = None,
  tags: Optional[Tags] = None,
  ) -> Dict[str, Any]:
@@ -5694,6 +5695,8 @@ def create_feature_group(
  feature online store.
  offline_store_config (Dict[str, str]): dict contains configuration of the
  feature offline store.
+ throughput_config (Dict[str, str]): dict contains throughput configuration
+ for the feature group.
  description (str): description of the FeatureGroup.
  tags (Optional[Tags]): tags for labeling a FeatureGroup.
 
@@ -5729,6 +5732,7 @@ def create_feature_group(
  kwargs,
  OnlineStoreConfig=inferred_online_store_from_config,
  OfflineStoreConfig=inferred_offline_store_from_config,
+ ThroughputConfig=throughput_config,
  Description=description,
  Tags=tags,
  )
@@ -5757,28 +5761,32 @@ def update_feature_group(
  feature_group_name: str,
  feature_additions: Sequence[Dict[str, str]] = None,
  online_store_config: Dict[str, any] = None,
+ throughput_config: Dict[str, Any] = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup
 
-either adding new features from the given feature definitions
-or updating online store config
+Supports modifications like adding new features from the given feature definitions,
+ updating online store and throughput configurations.
 
  Args:
  feature_group_name (str): name of the FeatureGroup to update.
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
+ online_store_config (Dict[str, Any]): updates to online store config
+ throughput_config (Dict[str, Any]): target throughput configuration of the feature group
  Returns:
  Response dict from service.
  """
+ update_req = {"FeatureGroupName": feature_group_name}
+ if online_store_config is not None:
+ update_req["OnlineStoreConfig"] = online_store_config
 
- if feature_additions is None:
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name,
- OnlineStoreConfig=online_store_config,
- )
+ if throughput_config is not None:
+ update_req["ThroughputConfig"] = throughput_config
 
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name, FeatureAdditions=feature_additions
- )
+ if feature_additions is not None:
+ update_req["FeatureAdditions"] = feature_additions
+
+ return self.sagemaker_client.update_feature_group(**update_req)
 
  def list_feature_groups(
  self,

tests/integ/test_feature_store.py

@@ -43,6 +43,9 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputModeEnum,
+ ThroughputConfigUpdate,
 )
 from sagemaker.feature_store.dataset_builder import (
  JoinTypeEnum,
@@ -410,6 +413,78 @@ def test_create_feature_group_standard_storage_type(
  assert storage_type == "Standard"
 
 
+def test_throughput_create_as_provisioned_and_update_to_ondemand(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.PROVISIONED, 4000, 4000),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 4000
+ assert wcu == 4000
+
+ feature_group.update(throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.ON_DEMAND))
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+
+def test_throughput_create_as_ondemand_and_update_to_provisioned(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.ON_DEMAND),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+ feature_group.update(
+ throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.PROVISIONED, 100, 200)
+ )
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 100
+ assert wcu == 200
+
+
 def test_ttl_duration(
  feature_store_session,
  role,

doc/api/prep_data/feature_store.rst

@@ -75,6 +75,14 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfig
+ :members:
+ :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfigUpdate
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.OnlineStoreConfig
  :members:
  :show-inheritance:
@@ -99,6 +107,10 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputModeEnum
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.ResourceEnum
  :members:
  :show-inheritance:

src/sagemaker/feature_store/feature_group.py

@@ -64,6 +64,8 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputConfigUpdate,
 )
 from sagemaker.utils import resolve_value_from_config, format_tags, Tags
 
@@ -541,6 +543,7 @@ def create(
  tags: Optional[Tags] = None,
  table_format: TableFormatEnum = None,
  online_store_storage_type: OnlineStoreStorageTypeEnum = None,
+ throughput_config: ThroughputConfig = None,
  ) -> Dict[str, Any]:
  """Create a SageMaker FeatureStore FeatureGroup.
 
@@ -570,6 +573,8 @@ def create(
  table_format (TableFormatEnum): format of the offline store table (default: None).
  online_store_storage_type (OnlineStoreStorageTypeEnum): storage type for the
  online store (default: None).
+ throughput_config (ThroughputConfig): throughput configuration of the
+ feature group (default: None).
 
  Returns:
  Response dict from service.
@@ -618,6 +623,9 @@ def create(
  )
  create_feature_store_args.update({"online_store_config": online_store_config.to_dict()})
 
+ if throughput_config:
+ create_feature_store_args.update({"throughput_config": throughput_config.to_dict()})
+
  # offline store configuration
  if s3_uri:
  s3_storage_config = S3StorageConfig(s3_uri=s3_uri)
@@ -656,17 +664,17 @@ def update(
  self,
  feature_additions: Sequence[FeatureDefinition] = None,
  online_store_config: OnlineStoreConfigUpdate = None,
+ throughput_config: ThroughputConfigUpdate = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup and add new features from the given feature definitions.
 
  Args:
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
  online_store_config (OnlineStoreConfigUpdate): online store config to be updated.
-
+ throughput_config (ThroughputConfigUpdate): target throughput configuration
  Returns:
  Response dict from service.
  """
-
  if feature_additions is None:
  feature_additions_parameter = None
  else:
@@ -679,10 +687,15 @@ def update(
  else:
  online_store_config_parameter = online_store_config.to_dict()
 
+ throughput_config_parameter = (
+ None if throughput_config is None else throughput_config.to_dict()
+ )
+
  return self.sagemaker_session.update_feature_group(
  feature_group_name=self.name,
  feature_additions=feature_additions_parameter,
  online_store_config=online_store_config_parameter,
+ throughput_config=throughput_config_parameter,
  )
 
  def update_feature_metadata(

src/sagemaker/feature_store/inputs.py

@@ -453,3 +453,79 @@ class ExpirationTimeResponseEnum(Enum):
 
  DISABLED = "Disabled"
  ENABLED = "Enabled"
+
+
+class ThroughputModeEnum(Enum):
+ """Enum of throughput modes supported by feature group.
+
+ Throughput mode of feature group can be ON_DEMAND or PROVISIONED.
+ """
+
+ ON_DEMAND = "OnDemand"
+ PROVISIONED = "Provisioned"
+
+
+@attr.s
+class ThroughputConfig(Config):
+ """Throughput configuration of the feature group.
+
+ Throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )
+
+
+@attr.s
+class ThroughputConfigUpdate(Config):
+ """Target throughput configuration for the feature group.
+
+ Target throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Target throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )

src/sagemaker/session.py

@@ -5679,6 +5679,7 @@ def create_feature_group(
  role_arn: str = None,
  online_store_config: Dict[str, str] = None,
  offline_store_config: Dict[str, str] = None,
+ throughput_config: Dict[str, Any] = None,
  description: str = None,
  tags: Optional[Tags] = None,
  ) -> Dict[str, Any]:
@@ -5694,6 +5695,8 @@ def create_feature_group(
  feature online store.
  offline_store_config (Dict[str, str]): dict contains configuration of the
  feature offline store.
+ throughput_config (Dict[str, str]): dict contains throughput configuration
+ for the feature group.
  description (str): description of the FeatureGroup.
  tags (Optional[Tags]): tags for labeling a FeatureGroup.
 
@@ -5729,6 +5732,7 @@ def create_feature_group(
  kwargs,
  OnlineStoreConfig=inferred_online_store_from_config,
  OfflineStoreConfig=inferred_offline_store_from_config,
+ ThroughputConfig=throughput_config,
  Description=description,
  Tags=tags,
  )
@@ -5757,28 +5761,32 @@ def update_feature_group(
  feature_group_name: str,
  feature_additions: Sequence[Dict[str, str]] = None,
  online_store_config: Dict[str, any] = None,
+ throughput_config: Dict[str, Any] = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup
 
-either adding new features from the given feature definitions
-or updating online store config
+Supports modifications like adding new features from the given feature definitions,
+ updating online store and throughput configurations.
 
  Args:
  feature_group_name (str): name of the FeatureGroup to update.
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
+ online_store_config (Dict[str, Any]): updates to online store config
+ throughput_config (Dict[str, Any]): target throughput configuration of the feature group
  Returns:
  Response dict from service.
  """
+ update_req = {"FeatureGroupName": feature_group_name}
+ if online_store_config is not None:
+ update_req["OnlineStoreConfig"] = online_store_config
 
- if feature_additions is None:
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name,
- OnlineStoreConfig=online_store_config,
- )
+ if throughput_config is not None:
+ update_req["ThroughputConfig"] = throughput_config
 
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name, FeatureAdditions=feature_additions
- )
+ if feature_additions is not None:
+ update_req["FeatureAdditions"] = feature_additions
+
+ return self.sagemaker_client.update_feature_group(**update_req)
 
  def list_feature_groups(
  self,

tests/integ/test_feature_store.py

@@ -43,6 +43,9 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputModeEnum,
+ ThroughputConfigUpdate,
 )
 from sagemaker.feature_store.dataset_builder import (
  JoinTypeEnum,
@@ -410,6 +413,78 @@ def test_create_feature_group_standard_storage_type(
  assert storage_type == "Standard"
 
 
+def test_throughput_create_as_provisioned_and_update_to_ondemand(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.PROVISIONED, 4000, 4000),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 4000
+ assert wcu == 4000
+
+ feature_group.update(throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.ON_DEMAND))
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+
+def test_throughput_create_as_ondemand_and_update_to_provisioned(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.ON_DEMAND),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+ feature_group.update(
+ throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.PROVISIONED, 100, 200)
+ )
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 100
+ assert wcu == 200
+
+
 def test_ttl_duration(
  feature_store_session,
  role,

doc/api/prep_data/feature_store.rst

@@ -75,6 +75,14 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfig
+ :members:
+ :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfigUpdate
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.OnlineStoreConfig
  :members:
  :show-inheritance:
@@ -99,6 +107,10 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputModeEnum
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.ResourceEnum
  :members:
  :show-inheritance:

src/sagemaker/feature_store/feature_group.py

@@ -64,6 +64,8 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputConfigUpdate,
 )
 from sagemaker.utils import resolve_value_from_config, format_tags, Tags
 
@@ -541,6 +543,7 @@ def create(
  tags: Optional[Tags] = None,
  table_format: TableFormatEnum = None,
  online_store_storage_type: OnlineStoreStorageTypeEnum = None,
+ throughput_config: ThroughputConfig = None,
  ) -> Dict[str, Any]:
  """Create a SageMaker FeatureStore FeatureGroup.
 
@@ -570,6 +573,8 @@ def create(
  table_format (TableFormatEnum): format of the offline store table (default: None).
  online_store_storage_type (OnlineStoreStorageTypeEnum): storage type for the
  online store (default: None).
+ throughput_config (ThroughputConfig): throughput configuration of the
+ feature group (default: None).
 
  Returns:
  Response dict from service.
@@ -618,6 +623,9 @@ def create(
  )
  create_feature_store_args.update({"online_store_config": online_store_config.to_dict()})
 
+ if throughput_config:
+ create_feature_store_args.update({"throughput_config": throughput_config.to_dict()})
+
  # offline store configuration
  if s3_uri:
  s3_storage_config = S3StorageConfig(s3_uri=s3_uri)
@@ -656,17 +664,17 @@ def update(
  self,
  feature_additions: Sequence[FeatureDefinition] = None,
  online_store_config: OnlineStoreConfigUpdate = None,
+ throughput_config: ThroughputConfigUpdate = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup and add new features from the given feature definitions.
 
  Args:
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
  online_store_config (OnlineStoreConfigUpdate): online store config to be updated.
-
+ throughput_config (ThroughputConfigUpdate): target throughput configuration
  Returns:
  Response dict from service.
  """
-
  if feature_additions is None:
  feature_additions_parameter = None
  else:
@@ -679,10 +687,15 @@ def update(
  else:
  online_store_config_parameter = online_store_config.to_dict()
 
+ throughput_config_parameter = (
+ None if throughput_config is None else throughput_config.to_dict()
+ )
+
  return self.sagemaker_session.update_feature_group(
  feature_group_name=self.name,
  feature_additions=feature_additions_parameter,
  online_store_config=online_store_config_parameter,
+ throughput_config=throughput_config_parameter,
  )
 
  def update_feature_metadata(

src/sagemaker/feature_store/inputs.py

@@ -453,3 +453,79 @@ class ExpirationTimeResponseEnum(Enum):
 
  DISABLED = "Disabled"
  ENABLED = "Enabled"
+
+
+class ThroughputModeEnum(Enum):
+ """Enum of throughput modes supported by feature group.
+
+ Throughput mode of feature group can be ON_DEMAND or PROVISIONED.
+ """
+
+ ON_DEMAND = "OnDemand"
+ PROVISIONED = "Provisioned"
+
+
+@attr.s
+class ThroughputConfig(Config):
+ """Throughput configuration of the feature group.
+
+ Throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )
+
+
+@attr.s
+class ThroughputConfigUpdate(Config):
+ """Target throughput configuration for the feature group.
+
+ Target throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Target throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )

src/sagemaker/session.py

@@ -5679,6 +5679,7 @@ def create_feature_group(
  role_arn: str = None,
  online_store_config: Dict[str, str] = None,
  offline_store_config: Dict[str, str] = None,
+ throughput_config: Dict[str, Any] = None,
  description: str = None,
  tags: Optional[Tags] = None,
  ) -> Dict[str, Any]:
@@ -5694,6 +5695,8 @@ def create_feature_group(
  feature online store.
  offline_store_config (Dict[str, str]): dict contains configuration of the
  feature offline store.
+ throughput_config (Dict[str, str]): dict contains throughput configuration
+ for the feature group.
  description (str): description of the FeatureGroup.
  tags (Optional[Tags]): tags for labeling a FeatureGroup.
 
@@ -5729,6 +5732,7 @@ def create_feature_group(
  kwargs,
  OnlineStoreConfig=inferred_online_store_from_config,
  OfflineStoreConfig=inferred_offline_store_from_config,
+ ThroughputConfig=throughput_config,
  Description=description,
  Tags=tags,
  )
@@ -5757,28 +5761,32 @@ def update_feature_group(
  feature_group_name: str,
  feature_additions: Sequence[Dict[str, str]] = None,
  online_store_config: Dict[str, any] = None,
+ throughput_config: Dict[str, Any] = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup
 
-either adding new features from the given feature definitions
-or updating online store config
+Supports modifications like adding new features from the given feature definitions,
+ updating online store and throughput configurations.
 
  Args:
  feature_group_name (str): name of the FeatureGroup to update.
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
+ online_store_config (Dict[str, Any]): updates to online store config
+ throughput_config (Dict[str, Any]): target throughput configuration of the feature group
  Returns:
  Response dict from service.
  """
+ update_req = {"FeatureGroupName": feature_group_name}
+ if online_store_config is not None:
+ update_req["OnlineStoreConfig"] = online_store_config
 
- if feature_additions is None:
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name,
- OnlineStoreConfig=online_store_config,
- )
+ if throughput_config is not None:
+ update_req["ThroughputConfig"] = throughput_config
 
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name, FeatureAdditions=feature_additions
- )
+ if feature_additions is not None:
+ update_req["FeatureAdditions"] = feature_additions
+
+ return self.sagemaker_client.update_feature_group(**update_req)
 
  def list_feature_groups(
  self,

tests/integ/test_feature_store.py

@@ -43,6 +43,9 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputModeEnum,
+ ThroughputConfigUpdate,
 )
 from sagemaker.feature_store.dataset_builder import (
  JoinTypeEnum,
@@ -410,6 +413,78 @@ def test_create_feature_group_standard_storage_type(
  assert storage_type == "Standard"
 
 
+def test_throughput_create_as_provisioned_and_update_to_ondemand(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.PROVISIONED, 4000, 4000),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 4000
+ assert wcu == 4000
+
+ feature_group.update(throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.ON_DEMAND))
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+
+def test_throughput_create_as_ondemand_and_update_to_provisioned(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.ON_DEMAND),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+ feature_group.update(
+ throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.PROVISIONED, 100, 200)
+ )
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 100
+ assert wcu == 200
+
+
 def test_ttl_duration(
  feature_store_session,
  role,

doc/api/prep_data/feature_store.rst

@@ -75,6 +75,14 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfig
+ :members:
+ :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfigUpdate
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.OnlineStoreConfig
  :members:
  :show-inheritance:
@@ -99,6 +107,10 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputModeEnum
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.ResourceEnum
  :members:
  :show-inheritance:

src/sagemaker/feature_store/feature_group.py

@@ -64,6 +64,8 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputConfigUpdate,
 )
 from sagemaker.utils import resolve_value_from_config, format_tags, Tags
 
@@ -541,6 +543,7 @@ def create(
  tags: Optional[Tags] = None,
  table_format: TableFormatEnum = None,
  online_store_storage_type: OnlineStoreStorageTypeEnum = None,
+ throughput_config: ThroughputConfig = None,
  ) -> Dict[str, Any]:
  """Create a SageMaker FeatureStore FeatureGroup.
 
@@ -570,6 +573,8 @@ def create(
  table_format (TableFormatEnum): format of the offline store table (default: None).
  online_store_storage_type (OnlineStoreStorageTypeEnum): storage type for the
  online store (default: None).
+ throughput_config (ThroughputConfig): throughput configuration of the
+ feature group (default: None).
 
  Returns:
  Response dict from service.
@@ -618,6 +623,9 @@ def create(
  )
  create_feature_store_args.update({"online_store_config": online_store_config.to_dict()})
 
+ if throughput_config:
+ create_feature_store_args.update({"throughput_config": throughput_config.to_dict()})
+
  # offline store configuration
  if s3_uri:
  s3_storage_config = S3StorageConfig(s3_uri=s3_uri)
@@ -656,17 +664,17 @@ def update(
  self,
  feature_additions: Sequence[FeatureDefinition] = None,
  online_store_config: OnlineStoreConfigUpdate = None,
+ throughput_config: ThroughputConfigUpdate = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup and add new features from the given feature definitions.
 
  Args:
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
  online_store_config (OnlineStoreConfigUpdate): online store config to be updated.
-
+ throughput_config (ThroughputConfigUpdate): target throughput configuration
  Returns:
  Response dict from service.
  """
-
  if feature_additions is None:
  feature_additions_parameter = None
  else:
@@ -679,10 +687,15 @@ def update(
  else:
  online_store_config_parameter = online_store_config.to_dict()
 
+ throughput_config_parameter = (
+ None if throughput_config is None else throughput_config.to_dict()
+ )
+
  return self.sagemaker_session.update_feature_group(
  feature_group_name=self.name,
  feature_additions=feature_additions_parameter,
  online_store_config=online_store_config_parameter,
+ throughput_config=throughput_config_parameter,
  )
 
  def update_feature_metadata(

src/sagemaker/feature_store/inputs.py

@@ -453,3 +453,79 @@ class ExpirationTimeResponseEnum(Enum):
 
  DISABLED = "Disabled"
  ENABLED = "Enabled"
+
+
+class ThroughputModeEnum(Enum):
+ """Enum of throughput modes supported by feature group.
+
+ Throughput mode of feature group can be ON_DEMAND or PROVISIONED.
+ """
+
+ ON_DEMAND = "OnDemand"
+ PROVISIONED = "Provisioned"
+
+
+@attr.s
+class ThroughputConfig(Config):
+ """Throughput configuration of the feature group.
+
+ Throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )
+
+
+@attr.s
+class ThroughputConfigUpdate(Config):
+ """Target throughput configuration for the feature group.
+
+ Target throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Target throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )

src/sagemaker/session.py

@@ -5679,6 +5679,7 @@ def create_feature_group(
  role_arn: str = None,
  online_store_config: Dict[str, str] = None,
  offline_store_config: Dict[str, str] = None,
+ throughput_config: Dict[str, Any] = None,
  description: str = None,
  tags: Optional[Tags] = None,
  ) -> Dict[str, Any]:
@@ -5694,6 +5695,8 @@ def create_feature_group(
  feature online store.
  offline_store_config (Dict[str, str]): dict contains configuration of the
  feature offline store.
+ throughput_config (Dict[str, str]): dict contains throughput configuration
+ for the feature group.
  description (str): description of the FeatureGroup.
  tags (Optional[Tags]): tags for labeling a FeatureGroup.
 
@@ -5729,6 +5732,7 @@ def create_feature_group(
  kwargs,
  OnlineStoreConfig=inferred_online_store_from_config,
  OfflineStoreConfig=inferred_offline_store_from_config,
+ ThroughputConfig=throughput_config,
  Description=description,
  Tags=tags,
  )
@@ -5757,28 +5761,32 @@ def update_feature_group(
  feature_group_name: str,
  feature_additions: Sequence[Dict[str, str]] = None,
  online_store_config: Dict[str, any] = None,
+ throughput_config: Dict[str, Any] = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup
 
-either adding new features from the given feature definitions
-or updating online store config
+Supports modifications like adding new features from the given feature definitions,
+ updating online store and throughput configurations.
 
  Args:
  feature_group_name (str): name of the FeatureGroup to update.
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
+ online_store_config (Dict[str, Any]): updates to online store config
+ throughput_config (Dict[str, Any]): target throughput configuration of the feature group
  Returns:
  Response dict from service.
  """
+ update_req = {"FeatureGroupName": feature_group_name}
+ if online_store_config is not None:
+ update_req["OnlineStoreConfig"] = online_store_config
 
- if feature_additions is None:
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name,
- OnlineStoreConfig=online_store_config,
- )
+ if throughput_config is not None:
+ update_req["ThroughputConfig"] = throughput_config
 
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name, FeatureAdditions=feature_additions
- )
+ if feature_additions is not None:
+ update_req["FeatureAdditions"] = feature_additions
+
+ return self.sagemaker_client.update_feature_group(**update_req)
 
  def list_feature_groups(
  self,

tests/integ/test_feature_store.py

@@ -43,6 +43,9 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputModeEnum,
+ ThroughputConfigUpdate,
 )
 from sagemaker.feature_store.dataset_builder import (
  JoinTypeEnum,
@@ -410,6 +413,78 @@ def test_create_feature_group_standard_storage_type(
  assert storage_type == "Standard"
 
 
+def test_throughput_create_as_provisioned_and_update_to_ondemand(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.PROVISIONED, 4000, 4000),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 4000
+ assert wcu == 4000
+
+ feature_group.update(throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.ON_DEMAND))
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+
+def test_throughput_create_as_ondemand_and_update_to_provisioned(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.ON_DEMAND),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+ feature_group.update(
+ throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.PROVISIONED, 100, 200)
+ )
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 100
+ assert wcu == 200
+
+
 def test_ttl_duration(
  feature_store_session,
  role,

doc/api/prep_data/feature_store.rst

@@ -75,6 +75,14 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfig
+ :members:
+ :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfigUpdate
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.OnlineStoreConfig
  :members:
  :show-inheritance:
@@ -99,6 +107,10 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputModeEnum
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.ResourceEnum
  :members:
  :show-inheritance:

src/sagemaker/feature_store/feature_group.py

@@ -64,6 +64,8 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputConfigUpdate,
 )
 from sagemaker.utils import resolve_value_from_config, format_tags, Tags
 
@@ -541,6 +543,7 @@ def create(
  tags: Optional[Tags] = None,
  table_format: TableFormatEnum = None,
  online_store_storage_type: OnlineStoreStorageTypeEnum = None,
+ throughput_config: ThroughputConfig = None,
  ) -> Dict[str, Any]:
  """Create a SageMaker FeatureStore FeatureGroup.
 
@@ -570,6 +573,8 @@ def create(
  table_format (TableFormatEnum): format of the offline store table (default: None).
  online_store_storage_type (OnlineStoreStorageTypeEnum): storage type for the
  online store (default: None).
+ throughput_config (ThroughputConfig): throughput configuration of the
+ feature group (default: None).
 
  Returns:
  Response dict from service.
@@ -618,6 +623,9 @@ def create(
  )
  create_feature_store_args.update({"online_store_config": online_store_config.to_dict()})
 
+ if throughput_config:
+ create_feature_store_args.update({"throughput_config": throughput_config.to_dict()})
+
  # offline store configuration
  if s3_uri:
  s3_storage_config = S3StorageConfig(s3_uri=s3_uri)
@@ -656,17 +664,17 @@ def update(
  self,
  feature_additions: Sequence[FeatureDefinition] = None,
  online_store_config: OnlineStoreConfigUpdate = None,
+ throughput_config: ThroughputConfigUpdate = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup and add new features from the given feature definitions.
 
  Args:
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
  online_store_config (OnlineStoreConfigUpdate): online store config to be updated.
-
+ throughput_config (ThroughputConfigUpdate): target throughput configuration
  Returns:
  Response dict from service.
  """
-
  if feature_additions is None:
  feature_additions_parameter = None
  else:
@@ -679,10 +687,15 @@ def update(
  else:
  online_store_config_parameter = online_store_config.to_dict()
 
+ throughput_config_parameter = (
+ None if throughput_config is None else throughput_config.to_dict()
+ )
+
  return self.sagemaker_session.update_feature_group(
  feature_group_name=self.name,
  feature_additions=feature_additions_parameter,
  online_store_config=online_store_config_parameter,
+ throughput_config=throughput_config_parameter,
  )
 
  def update_feature_metadata(

src/sagemaker/feature_store/inputs.py

@@ -453,3 +453,79 @@ class ExpirationTimeResponseEnum(Enum):
 
  DISABLED = "Disabled"
  ENABLED = "Enabled"
+
+
+class ThroughputModeEnum(Enum):
+ """Enum of throughput modes supported by feature group.
+
+ Throughput mode of feature group can be ON_DEMAND or PROVISIONED.
+ """
+
+ ON_DEMAND = "OnDemand"
+ PROVISIONED = "Provisioned"
+
+
+@attr.s
+class ThroughputConfig(Config):
+ """Throughput configuration of the feature group.
+
+ Throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )
+
+
+@attr.s
+class ThroughputConfigUpdate(Config):
+ """Target throughput configuration for the feature group.
+
+ Target throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Target throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )

src/sagemaker/session.py

@@ -5679,6 +5679,7 @@ def create_feature_group(
  role_arn: str = None,
  online_store_config: Dict[str, str] = None,
  offline_store_config: Dict[str, str] = None,
+ throughput_config: Dict[str, Any] = None,
  description: str = None,
  tags: Optional[Tags] = None,
  ) -> Dict[str, Any]:
@@ -5694,6 +5695,8 @@ def create_feature_group(
  feature online store.
  offline_store_config (Dict[str, str]): dict contains configuration of the
  feature offline store.
+ throughput_config (Dict[str, str]): dict contains throughput configuration
+ for the feature group.
  description (str): description of the FeatureGroup.
  tags (Optional[Tags]): tags for labeling a FeatureGroup.
 
@@ -5729,6 +5732,7 @@ def create_feature_group(
  kwargs,
  OnlineStoreConfig=inferred_online_store_from_config,
  OfflineStoreConfig=inferred_offline_store_from_config,
+ ThroughputConfig=throughput_config,
  Description=description,
  Tags=tags,
  )
@@ -5757,28 +5761,32 @@ def update_feature_group(
  feature_group_name: str,
  feature_additions: Sequence[Dict[str, str]] = None,
  online_store_config: Dict[str, any] = None,
+ throughput_config: Dict[str, Any] = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup
 
-either adding new features from the given feature definitions
-or updating online store config
+Supports modifications like adding new features from the given feature definitions,
+ updating online store and throughput configurations.
 
  Args:
  feature_group_name (str): name of the FeatureGroup to update.
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
+ online_store_config (Dict[str, Any]): updates to online store config
+ throughput_config (Dict[str, Any]): target throughput configuration of the feature group
  Returns:
  Response dict from service.
  """
+ update_req = {"FeatureGroupName": feature_group_name}
+ if online_store_config is not None:
+ update_req["OnlineStoreConfig"] = online_store_config
 
- if feature_additions is None:
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name,
- OnlineStoreConfig=online_store_config,
- )
+ if throughput_config is not None:
+ update_req["ThroughputConfig"] = throughput_config
 
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name, FeatureAdditions=feature_additions
- )
+ if feature_additions is not None:
+ update_req["FeatureAdditions"] = feature_additions
+
+ return self.sagemaker_client.update_feature_group(**update_req)
 
  def list_feature_groups(
  self,

tests/integ/test_feature_store.py

@@ -43,6 +43,9 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputModeEnum,
+ ThroughputConfigUpdate,
 )
 from sagemaker.feature_store.dataset_builder import (
  JoinTypeEnum,
@@ -410,6 +413,78 @@ def test_create_feature_group_standard_storage_type(
  assert storage_type == "Standard"
 
 
+def test_throughput_create_as_provisioned_and_update_to_ondemand(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.PROVISIONED, 4000, 4000),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 4000
+ assert wcu == 4000
+
+ feature_group.update(throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.ON_DEMAND))
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+
+def test_throughput_create_as_ondemand_and_update_to_provisioned(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.ON_DEMAND),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+ feature_group.update(
+ throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.PROVISIONED, 100, 200)
+ )
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 100
+ assert wcu == 200
+
+
 def test_ttl_duration(
  feature_store_session,
  role,

doc/api/prep_data/feature_store.rst

@@ -75,6 +75,14 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfig
+ :members:
+ :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfigUpdate
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.OnlineStoreConfig
  :members:
  :show-inheritance:
@@ -99,6 +107,10 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputModeEnum
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.ResourceEnum
  :members:
  :show-inheritance:

src/sagemaker/feature_store/feature_group.py

@@ -64,6 +64,8 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputConfigUpdate,
 )
 from sagemaker.utils import resolve_value_from_config, format_tags, Tags
 
@@ -541,6 +543,7 @@ def create(
  tags: Optional[Tags] = None,
  table_format: TableFormatEnum = None,
  online_store_storage_type: OnlineStoreStorageTypeEnum = None,
+ throughput_config: ThroughputConfig = None,
  ) -> Dict[str, Any]:
  """Create a SageMaker FeatureStore FeatureGroup.
 
@@ -570,6 +573,8 @@ def create(
  table_format (TableFormatEnum): format of the offline store table (default: None).
  online_store_storage_type (OnlineStoreStorageTypeEnum): storage type for the
  online store (default: None).
+ throughput_config (ThroughputConfig): throughput configuration of the
+ feature group (default: None).
 
  Returns:
  Response dict from service.
@@ -618,6 +623,9 @@ def create(
  )
  create_feature_store_args.update({"online_store_config": online_store_config.to_dict()})
 
+ if throughput_config:
+ create_feature_store_args.update({"throughput_config": throughput_config.to_dict()})
+
  # offline store configuration
  if s3_uri:
  s3_storage_config = S3StorageConfig(s3_uri=s3_uri)
@@ -656,17 +664,17 @@ def update(
  self,
  feature_additions: Sequence[FeatureDefinition] = None,
  online_store_config: OnlineStoreConfigUpdate = None,
+ throughput_config: ThroughputConfigUpdate = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup and add new features from the given feature definitions.
 
  Args:
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
  online_store_config (OnlineStoreConfigUpdate): online store config to be updated.
-
+ throughput_config (ThroughputConfigUpdate): target throughput configuration
  Returns:
  Response dict from service.
  """
-
  if feature_additions is None:
  feature_additions_parameter = None
  else:
@@ -679,10 +687,15 @@ def update(
  else:
  online_store_config_parameter = online_store_config.to_dict()
 
+ throughput_config_parameter = (
+ None if throughput_config is None else throughput_config.to_dict()
+ )
+
  return self.sagemaker_session.update_feature_group(
  feature_group_name=self.name,
  feature_additions=feature_additions_parameter,
  online_store_config=online_store_config_parameter,
+ throughput_config=throughput_config_parameter,
  )
 
  def update_feature_metadata(

src/sagemaker/feature_store/inputs.py

@@ -453,3 +453,79 @@ class ExpirationTimeResponseEnum(Enum):
 
  DISABLED = "Disabled"
  ENABLED = "Enabled"
+
+
+class ThroughputModeEnum(Enum):
+ """Enum of throughput modes supported by feature group.
+
+ Throughput mode of feature group can be ON_DEMAND or PROVISIONED.
+ """
+
+ ON_DEMAND = "OnDemand"
+ PROVISIONED = "Provisioned"
+
+
+@attr.s
+class ThroughputConfig(Config):
+ """Throughput configuration of the feature group.
+
+ Throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )
+
+
+@attr.s
+class ThroughputConfigUpdate(Config):
+ """Target throughput configuration for the feature group.
+
+ Target throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Target throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )

src/sagemaker/session.py

@@ -5679,6 +5679,7 @@ def create_feature_group(
  role_arn: str = None,
  online_store_config: Dict[str, str] = None,
  offline_store_config: Dict[str, str] = None,
+ throughput_config: Dict[str, Any] = None,
  description: str = None,
  tags: Optional[Tags] = None,
  ) -> Dict[str, Any]:
@@ -5694,6 +5695,8 @@ def create_feature_group(
  feature online store.
  offline_store_config (Dict[str, str]): dict contains configuration of the
  feature offline store.
+ throughput_config (Dict[str, str]): dict contains throughput configuration
+ for the feature group.
  description (str): description of the FeatureGroup.
  tags (Optional[Tags]): tags for labeling a FeatureGroup.
 
@@ -5729,6 +5732,7 @@ def create_feature_group(
  kwargs,
  OnlineStoreConfig=inferred_online_store_from_config,
  OfflineStoreConfig=inferred_offline_store_from_config,
+ ThroughputConfig=throughput_config,
  Description=description,
  Tags=tags,
  )
@@ -5757,28 +5761,32 @@ def update_feature_group(
  feature_group_name: str,
  feature_additions: Sequence[Dict[str, str]] = None,
  online_store_config: Dict[str, any] = None,
+ throughput_config: Dict[str, Any] = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup
 
-either adding new features from the given feature definitions
-or updating online store config
+Supports modifications like adding new features from the given feature definitions,
+ updating online store and throughput configurations.
 
  Args:
  feature_group_name (str): name of the FeatureGroup to update.
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
+ online_store_config (Dict[str, Any]): updates to online store config
+ throughput_config (Dict[str, Any]): target throughput configuration of the feature group
  Returns:
  Response dict from service.
  """
+ update_req = {"FeatureGroupName": feature_group_name}
+ if online_store_config is not None:
+ update_req["OnlineStoreConfig"] = online_store_config
 
- if feature_additions is None:
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name,
- OnlineStoreConfig=online_store_config,
- )
+ if throughput_config is not None:
+ update_req["ThroughputConfig"] = throughput_config
 
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name, FeatureAdditions=feature_additions
- )
+ if feature_additions is not None:
+ update_req["FeatureAdditions"] = feature_additions
+
+ return self.sagemaker_client.update_feature_group(**update_req)
 
  def list_feature_groups(
  self,

tests/integ/test_feature_store.py

@@ -43,6 +43,9 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputModeEnum,
+ ThroughputConfigUpdate,
 )
 from sagemaker.feature_store.dataset_builder import (
  JoinTypeEnum,
@@ -410,6 +413,78 @@ def test_create_feature_group_standard_storage_type(
  assert storage_type == "Standard"
 
 
+def test_throughput_create_as_provisioned_and_update_to_ondemand(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.PROVISIONED, 4000, 4000),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 4000
+ assert wcu == 4000
+
+ feature_group.update(throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.ON_DEMAND))
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+
+def test_throughput_create_as_ondemand_and_update_to_provisioned(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.ON_DEMAND),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+ feature_group.update(
+ throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.PROVISIONED, 100, 200)
+ )
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 100
+ assert wcu == 200
+
+
 def test_ttl_duration(
  feature_store_session,
  role,

doc/api/prep_data/feature_store.rst

@@ -75,6 +75,14 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfig
+ :members:
+ :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfigUpdate
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.OnlineStoreConfig
  :members:
  :show-inheritance:
@@ -99,6 +107,10 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputModeEnum
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.ResourceEnum
  :members:
  :show-inheritance:

src/sagemaker/feature_store/feature_group.py

@@ -64,6 +64,8 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputConfigUpdate,
 )
 from sagemaker.utils import resolve_value_from_config, format_tags, Tags
 
@@ -541,6 +543,7 @@ def create(
  tags: Optional[Tags] = None,
  table_format: TableFormatEnum = None,
  online_store_storage_type: OnlineStoreStorageTypeEnum = None,
+ throughput_config: ThroughputConfig = None,
  ) -> Dict[str, Any]:
  """Create a SageMaker FeatureStore FeatureGroup.
 
@@ -570,6 +573,8 @@ def create(
  table_format (TableFormatEnum): format of the offline store table (default: None).
  online_store_storage_type (OnlineStoreStorageTypeEnum): storage type for the
  online store (default: None).
+ throughput_config (ThroughputConfig): throughput configuration of the
+ feature group (default: None).
 
  Returns:
  Response dict from service.
@@ -618,6 +623,9 @@ def create(
  )
  create_feature_store_args.update({"online_store_config": online_store_config.to_dict()})
 
+ if throughput_config:
+ create_feature_store_args.update({"throughput_config": throughput_config.to_dict()})
+
  # offline store configuration
  if s3_uri:
  s3_storage_config = S3StorageConfig(s3_uri=s3_uri)
@@ -656,17 +664,17 @@ def update(
  self,
  feature_additions: Sequence[FeatureDefinition] = None,
  online_store_config: OnlineStoreConfigUpdate = None,
+ throughput_config: ThroughputConfigUpdate = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup and add new features from the given feature definitions.
 
  Args:
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
  online_store_config (OnlineStoreConfigUpdate): online store config to be updated.
-
+ throughput_config (ThroughputConfigUpdate): target throughput configuration
  Returns:
  Response dict from service.
  """
-
  if feature_additions is None:
  feature_additions_parameter = None
  else:
@@ -679,10 +687,15 @@ def update(
  else:
  online_store_config_parameter = online_store_config.to_dict()
 
+ throughput_config_parameter = (
+ None if throughput_config is None else throughput_config.to_dict()
+ )
+
  return self.sagemaker_session.update_feature_group(
  feature_group_name=self.name,
  feature_additions=feature_additions_parameter,
  online_store_config=online_store_config_parameter,
+ throughput_config=throughput_config_parameter,
  )
 
  def update_feature_metadata(

src/sagemaker/feature_store/inputs.py

@@ -453,3 +453,79 @@ class ExpirationTimeResponseEnum(Enum):
 
  DISABLED = "Disabled"
  ENABLED = "Enabled"
+
+
+class ThroughputModeEnum(Enum):
+ """Enum of throughput modes supported by feature group.
+
+ Throughput mode of feature group can be ON_DEMAND or PROVISIONED.
+ """
+
+ ON_DEMAND = "OnDemand"
+ PROVISIONED = "Provisioned"
+
+
+@attr.s
+class ThroughputConfig(Config):
+ """Throughput configuration of the feature group.
+
+ Throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )
+
+
+@attr.s
+class ThroughputConfigUpdate(Config):
+ """Target throughput configuration for the feature group.
+
+ Target throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Target throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )

src/sagemaker/session.py

@@ -5679,6 +5679,7 @@ def create_feature_group(
  role_arn: str = None,
  online_store_config: Dict[str, str] = None,
  offline_store_config: Dict[str, str] = None,
+ throughput_config: Dict[str, Any] = None,
  description: str = None,
  tags: Optional[Tags] = None,
  ) -> Dict[str, Any]:
@@ -5694,6 +5695,8 @@ def create_feature_group(
  feature online store.
  offline_store_config (Dict[str, str]): dict contains configuration of the
  feature offline store.
+ throughput_config (Dict[str, str]): dict contains throughput configuration
+ for the feature group.
  description (str): description of the FeatureGroup.
  tags (Optional[Tags]): tags for labeling a FeatureGroup.
 
@@ -5729,6 +5732,7 @@ def create_feature_group(
  kwargs,
  OnlineStoreConfig=inferred_online_store_from_config,
  OfflineStoreConfig=inferred_offline_store_from_config,
+ ThroughputConfig=throughput_config,
  Description=description,
  Tags=tags,
  )
@@ -5757,28 +5761,32 @@ def update_feature_group(
  feature_group_name: str,
  feature_additions: Sequence[Dict[str, str]] = None,
  online_store_config: Dict[str, any] = None,
+ throughput_config: Dict[str, Any] = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup
 
-either adding new features from the given feature definitions
-or updating online store config
+Supports modifications like adding new features from the given feature definitions,
+ updating online store and throughput configurations.
 
  Args:
  feature_group_name (str): name of the FeatureGroup to update.
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
+ online_store_config (Dict[str, Any]): updates to online store config
+ throughput_config (Dict[str, Any]): target throughput configuration of the feature group
  Returns:
  Response dict from service.
  """
+ update_req = {"FeatureGroupName": feature_group_name}
+ if online_store_config is not None:
+ update_req["OnlineStoreConfig"] = online_store_config
 
- if feature_additions is None:
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name,
- OnlineStoreConfig=online_store_config,
- )
+ if throughput_config is not None:
+ update_req["ThroughputConfig"] = throughput_config
 
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name, FeatureAdditions=feature_additions
- )
+ if feature_additions is not None:
+ update_req["FeatureAdditions"] = feature_additions
+
+ return self.sagemaker_client.update_feature_group(**update_req)
 
  def list_feature_groups(
  self,

tests/integ/test_feature_store.py

@@ -43,6 +43,9 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputModeEnum,
+ ThroughputConfigUpdate,
 )
 from sagemaker.feature_store.dataset_builder import (
  JoinTypeEnum,
@@ -410,6 +413,78 @@ def test_create_feature_group_standard_storage_type(
  assert storage_type == "Standard"
 
 
+def test_throughput_create_as_provisioned_and_update_to_ondemand(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.PROVISIONED, 4000, 4000),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 4000
+ assert wcu == 4000
+
+ feature_group.update(throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.ON_DEMAND))
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+
+def test_throughput_create_as_ondemand_and_update_to_provisioned(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.ON_DEMAND),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+ feature_group.update(
+ throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.PROVISIONED, 100, 200)
+ )
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 100
+ assert wcu == 200
+
+
 def test_ttl_duration(
  feature_store_session,
  role,

doc/api/prep_data/feature_store.rst

@@ -75,6 +75,14 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfig
+ :members:
+ :show-inheritance:
+
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputConfigUpdate
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.OnlineStoreConfig
  :members:
  :show-inheritance:
@@ -99,6 +107,10 @@ Inputs
  :members:
  :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.ThroughputModeEnum
+ :members:
+ :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.ResourceEnum
  :members:
  :show-inheritance:

src/sagemaker/feature_store/feature_group.py

@@ -64,6 +64,8 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputConfigUpdate,
 )
 from sagemaker.utils import resolve_value_from_config, format_tags, Tags
 
@@ -541,6 +543,7 @@ def create(
  tags: Optional[Tags] = None,
  table_format: TableFormatEnum = None,
  online_store_storage_type: OnlineStoreStorageTypeEnum = None,
+ throughput_config: ThroughputConfig = None,
  ) -> Dict[str, Any]:
  """Create a SageMaker FeatureStore FeatureGroup.
 
@@ -570,6 +573,8 @@ def create(
  table_format (TableFormatEnum): format of the offline store table (default: None).
  online_store_storage_type (OnlineStoreStorageTypeEnum): storage type for the
  online store (default: None).
+ throughput_config (ThroughputConfig): throughput configuration of the
+ feature group (default: None).
 
  Returns:
  Response dict from service.
@@ -618,6 +623,9 @@ def create(
  )
  create_feature_store_args.update({"online_store_config": online_store_config.to_dict()})
 
+ if throughput_config:
+ create_feature_store_args.update({"throughput_config": throughput_config.to_dict()})
+
  # offline store configuration
  if s3_uri:
  s3_storage_config = S3StorageConfig(s3_uri=s3_uri)
@@ -656,17 +664,17 @@ def update(
  self,
  feature_additions: Sequence[FeatureDefinition] = None,
  online_store_config: OnlineStoreConfigUpdate = None,
+ throughput_config: ThroughputConfigUpdate = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup and add new features from the given feature definitions.
 
  Args:
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
  online_store_config (OnlineStoreConfigUpdate): online store config to be updated.
-
+ throughput_config (ThroughputConfigUpdate): target throughput configuration
  Returns:
  Response dict from service.
  """
-
  if feature_additions is None:
  feature_additions_parameter = None
  else:
@@ -679,10 +687,15 @@ def update(
  else:
  online_store_config_parameter = online_store_config.to_dict()
 
+ throughput_config_parameter = (
+ None if throughput_config is None else throughput_config.to_dict()
+ )
+
  return self.sagemaker_session.update_feature_group(
  feature_group_name=self.name,
  feature_additions=feature_additions_parameter,
  online_store_config=online_store_config_parameter,
+ throughput_config=throughput_config_parameter,
  )
 
  def update_feature_metadata(

src/sagemaker/feature_store/inputs.py

@@ -453,3 +453,79 @@ class ExpirationTimeResponseEnum(Enum):
 
  DISABLED = "Disabled"
  ENABLED = "Enabled"
+
+
+class ThroughputModeEnum(Enum):
+ """Enum of throughput modes supported by feature group.
+
+ Throughput mode of feature group can be ON_DEMAND or PROVISIONED.
+ """
+
+ ON_DEMAND = "OnDemand"
+ PROVISIONED = "Provisioned"
+
+
+@attr.s
+class ThroughputConfig(Config):
+ """Throughput configuration of the feature group.
+
+ Throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )
+
+
+@attr.s
+class ThroughputConfigUpdate(Config):
+ """Target throughput configuration for the feature group.
+
+ Target throughput configuration can be ON_DEMAND, or PROVISIONED with valid values for
+ read and write capacity units. ON_DEMAND works best for less predictable traffic,
+ while PROVISIONED works best for consistent and predictable traffic.
+
+ Attributes:
+ mode (ThroughputModeEnum): Target throughput mode
+ provisioned_read_capacity_units (int): For provisioned feature groups, this indicates
+ the read throughput you are billed for and can consume without throttling.
+ provisioned_write_capacity_units (int): For provisioned feature groups, this indicates
+ the write throughput you are billed for and can consume without throttling.
+ """
+
+ mode: ThroughputModeEnum = attr.ib(default=None)
+ provisioned_read_capacity_units: int = attr.ib(default=None)
+ provisioned_write_capacity_units: int = attr.ib(default=None)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Construct a dictionary based on the attributes provided.
+
+ Returns:
+ dict represents the attributes.
+ """
+ return Config.construct_dict(
+ ThroughputMode=self.mode.value if self.mode else None,
+ ProvisionedReadCapacityUnits=self.provisioned_read_capacity_units,
+ ProvisionedWriteCapacityUnits=self.provisioned_write_capacity_units,
+ )

src/sagemaker/session.py

@@ -5679,6 +5679,7 @@ def create_feature_group(
  role_arn: str = None,
  online_store_config: Dict[str, str] = None,
  offline_store_config: Dict[str, str] = None,
+ throughput_config: Dict[str, Any] = None,
  description: str = None,
  tags: Optional[Tags] = None,
  ) -> Dict[str, Any]:
@@ -5694,6 +5695,8 @@ def create_feature_group(
  feature online store.
  offline_store_config (Dict[str, str]): dict contains configuration of the
  feature offline store.
+ throughput_config (Dict[str, str]): dict contains throughput configuration
+ for the feature group.
  description (str): description of the FeatureGroup.
  tags (Optional[Tags]): tags for labeling a FeatureGroup.
 
@@ -5729,6 +5732,7 @@ def create_feature_group(
  kwargs,
  OnlineStoreConfig=inferred_online_store_from_config,
  OfflineStoreConfig=inferred_offline_store_from_config,
+ ThroughputConfig=throughput_config,
  Description=description,
  Tags=tags,
  )
@@ -5757,28 +5761,32 @@ def update_feature_group(
  feature_group_name: str,
  feature_additions: Sequence[Dict[str, str]] = None,
  online_store_config: Dict[str, any] = None,
+ throughput_config: Dict[str, Any] = None,
  ) -> Dict[str, Any]:
  """Update a FeatureGroup
 
-either adding new features from the given feature definitions
-or updating online store config
+Supports modifications like adding new features from the given feature definitions,
+ updating online store and throughput configurations.
 
  Args:
  feature_group_name (str): name of the FeatureGroup to update.
  feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
+ online_store_config (Dict[str, Any]): updates to online store config
+ throughput_config (Dict[str, Any]): target throughput configuration of the feature group
  Returns:
  Response dict from service.
  """
+ update_req = {"FeatureGroupName": feature_group_name}
+ if online_store_config is not None:
+ update_req["OnlineStoreConfig"] = online_store_config
 
- if feature_additions is None:
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name,
- OnlineStoreConfig=online_store_config,
- )
+ if throughput_config is not None:
+ update_req["ThroughputConfig"] = throughput_config
 
- return self.sagemaker_client.update_feature_group(
- FeatureGroupName=feature_group_name, FeatureAdditions=feature_additions
- )
+ if feature_additions is not None:
+ update_req["FeatureAdditions"] = feature_additions
+
+ return self.sagemaker_client.update_feature_group(**update_req)
 
  def list_feature_groups(
  self,

tests/integ/test_feature_store.py

@@ -43,6 +43,9 @@
  TtlDuration,
  OnlineStoreConfigUpdate,
  OnlineStoreStorageTypeEnum,
+ ThroughputConfig,
+ ThroughputModeEnum,
+ ThroughputConfigUpdate,
 )
 from sagemaker.feature_store.dataset_builder import (
  JoinTypeEnum,
@@ -410,6 +413,78 @@ def test_create_feature_group_standard_storage_type(
  assert storage_type == "Standard"
 
 
+def test_throughput_create_as_provisioned_and_update_to_ondemand(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.PROVISIONED, 4000, 4000),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 4000
+ assert wcu == 4000
+
+ feature_group.update(throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.ON_DEMAND))
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+
+def test_throughput_create_as_ondemand_and_update_to_provisioned(
+ feature_store_session,
+ role,
+ feature_group_name,
+ pandas_data_frame,
+):
+ feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+ feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+ with cleanup_feature_group(feature_group):
+ feature_group.create(
+ s3_uri=False,
+ record_identifier_name="feature1",
+ event_time_feature_name="feature3",
+ role_arn=role,
+ enable_online_store=True,
+ throughput_config=ThroughputConfig(ThroughputModeEnum.ON_DEMAND),
+ )
+ _wait_for_feature_group_create(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ assert mode == ThroughputModeEnum.ON_DEMAND.value
+
+ feature_group.update(
+ throughput_config=ThroughputConfigUpdate(ThroughputModeEnum.PROVISIONED, 100, 200)
+ )
+ _wait_for_feature_group_update(feature_group)
+
+ tp_config = feature_group.describe().get("ThroughputConfig")
+ mode = tp_config.get("ThroughputMode")
+ rcu = tp_config.get("ProvisionedReadCapacityUnits")
+ wcu = tp_config.get("ProvisionedWriteCapacityUnits")
+ assert mode == ThroughputModeEnum.PROVISIONED.value
+ assert rcu == 100
+ assert wcu == 200
+
+
 def test_ttl_duration(
  feature_store_session,
  role,