From 3aed24eca3d13ab287112514468571a463cfcf7c Mon Sep 17 00:00:00 2001 From: Romik Amipara Date: Tue, 8 Sep 2026 16:23:43 -0700 Subject: [PATCH] feat(feature-store): add UpdateRecord API and Standard_V2 storage type Adds feature-level writes to SageMaker Feature Store: sagemaker-core (generated layer, regenerated from updated models): - Botocore models: UpdateRecord op (PATCH /FeatureGroup/{name}/Record), UpdateRecordRequest/Features shapes, ConflictException/ThrottlingException, and Standard_V2 value on the sagemaker StorageType enum. - FeatureGroup.update_record resource method + shape_dag serialization entries; registered in additional_operations.json. sagemaker-mlops (convenience layer): - OnlineStoreStorageTypeEnum.STANDARD_V2. - feature_utils.update_record(): partial-write helper delegating to core, with client-side validation (<=100 features, no duplicates, offline-only rejected, ttl requires EventTime) and dict/FeatureValue coercion. - Exported update_record from feature_store package. Tests & docs: - Core plumbing tests, mlops unit tests, integ tests, and a usage doc. Task: P495083396 --- docs/ml_ops/index.rst | 67 +++++++++ .../2020-07-01/service-2.json | 58 ++++++++ .../sagemaker/2017-07-24/service-2.json | 5 + .../src/sagemaker/core/resources.py | 58 ++++++++ .../src/sagemaker/core/shapes/shapes.py | 2 + .../core/tools/additional_operations.json | 8 ++ .../core/utils/code_injection/shape_dag.py | 15 +- .../test_feature_store_operations.py | 78 ++++++++++- .../sagemaker/mlops/feature_store/__init__.py | 2 + .../mlops/feature_store/feature_utils.py | 114 +++++++++++++++- .../sagemaker/mlops/feature_store/inputs.py | 1 + .../integ/test_feature_store_update_record.py | 129 ++++++++++++++++++ .../mlops/feature_store/test_feature_utils.py | 126 +++++++++++++++++ .../mlops/feature_store/test_inputs.py | 3 + 14 files changed, 661 insertions(+), 5 deletions(-) create mode 100644 sagemaker-mlops/tests/integ/test_feature_store_update_record.py diff --git a/docs/ml_ops/index.rst b/docs/ml_ops/index.rst index 31923939c7..b3c167d8cb 100644 --- a/docs/ml_ops/index.rst +++ b/docs/ml_ops/index.rst @@ -641,6 +641,73 @@ To include soft-deleted records in the listing: region="us-west-2", ) +**Feature-level writes with UpdateRecord (Standard_V2):** + +``UpdateRecord`` performs a partial write to a record in a feature group whose online store uses +the ``Standard_V2`` or ``InMemory`` storage type. Only the features you supply are written; features +you do not list are preserved. This avoids the ``GetRecord`` -> merge -> ``PutRecord`` round trip and +prevents lost writes when independent pipelines own different features on the same record. The record +must already exist in the online store (use ``PutRecord`` to create it). + +Create the feature group with ``Standard_V2`` storage (feature-level writes require ``Standard_V2`` +or ``InMemory``; they are not supported on the default ``Standard`` tier): + +.. code-block:: python + + from sagemaker.mlops.feature_store import FeatureGroupManager, OnlineStoreStorageTypeEnum + from sagemaker.core.shapes import OnlineStoreConfig + + feature_group = FeatureGroupManager.create( + feature_group_name="customer-features", + record_identifier_feature_name="customer_id", + event_time_feature_name="event_time", + feature_definitions=feature_definitions, + online_store_config=OnlineStoreConfig( + enable_online_store=True, + storage_type=OnlineStoreStorageTypeEnum.STANDARD_V2.value, + ), + role_arn=role, + ) + +You can migrate an existing ``Standard`` feature group to ``Standard_V2`` with ``UpdateFeatureGroup``. +This migration is one-way and cannot be reversed: + +.. code-block:: python + + from sagemaker.core.resources import FeatureGroup + from sagemaker.core.shapes import OnlineStoreConfigUpdate + + feature_group = FeatureGroup.get(feature_group_name="customer-features") + feature_group.update( + online_store_config=OnlineStoreConfigUpdate(storage_type="Standard_V2"), + ) + +Use ``update_record`` to write only the features that changed. Pass ``EventTime`` as a feature +(not a top-level parameter); features you do not include are preserved: + +.. code-block:: python + + from sagemaker.mlops.feature_store import update_record + + update_record( + feature_group_name="customer-features", + record_identifier_value_as_string="cust-1", + features=[ + {"feature_name": "purchase_count", "value_as_string": "11"}, + {"feature_name": "event_time", "value_as_string": "2026-01-02T00:00:00Z"}, + ], + region="us-west-2", + ) + +Notes: + +* Supply at most 100 features per call. If the supplied ``EventTime`` is not greater than the + record's current ``EventTime``, the update is rejected with a ``ConflictException``. +* ``ttl_duration`` requires the record's event-time feature to be present in ``features``. + ``target_stores`` defaults to all stores on the feature group; a value resolving to the + ``OfflineStore`` only is rejected. +* ``UpdateRecord`` is not supported on ``Standard`` (V1) feature groups. + Migration from V2 diff --git a/sagemaker-core/sample/sagemaker-featurestore-runtime/2020-07-01/service-2.json b/sagemaker-core/sample/sagemaker-featurestore-runtime/2020-07-01/service-2.json index 9819e70f98..afd735cfb8 100644 --- a/sagemaker-core/sample/sagemaker-featurestore-runtime/2020-07-01/service-2.json +++ b/sagemaker-core/sample/sagemaker-featurestore-runtime/2020-07-01/service-2.json @@ -111,6 +111,23 @@ {"shape":"AccessForbidden"} ], "documentation":"

The PutRecord API is used to ingest a list of Records into your feature group.

If a new record’s EventTime is greater, the new record is written to both the OnlineStore and OfflineStore. Otherwise, the record is a historic record and it is written only to the OfflineStore.

You can specify the ingestion to be applied to the OnlineStore, OfflineStore, or both by using the TargetStores request parameter.

You can set the ingested record to expire at a given time to live (TTL) duration after the record’s event time, ExpiresAt = EventTime + TtlDuration, by specifying the TtlDuration parameter. A record level TtlDuration is set when specifying the TtlDuration parameter using the PutRecord API call. If the input TtlDuration is null or unspecified, TtlDuration is set to the default feature group level TtlDuration. A record level TtlDuration supersedes the group level TtlDuration.

" + }, + "UpdateRecord":{ + "name":"UpdateRecord", + "http":{ + "method":"POST", + "requestUri":"/FeatureGroup/{FeatureGroupName}/Record" + }, + "input":{"shape":"UpdateRecordRequest"}, + "errors":[ + {"shape":"ValidationError"}, + {"shape":"InternalFailure"}, + {"shape":"ServiceUnavailable"}, + {"shape":"AccessForbidden"}, + {"shape":"ResourceNotFound"}, + {"shape":"ConflictException"} + ], + "documentation":"

Updates one or more feature values for an existing record in the specified feature group. Features that you do not include in the request remain unchanged. You can update up to 100 features per call. This operation requires the online store and is available only for feature groups that use the Standard_V2 or InMemory online store type.

The record must already exist in the online store; if it does not exist or has been soft-deleted, the operation returns a ResourceNotFound error. To create a record, use PutRecord.

Pass EventTime as a feature in the Features list rather than as a top-level parameter. If you provide an EventTime that is older than the record's current EventTime, the update is rejected with a ConflictException; if it is equal to or newer, the update is applied; if you omit it, the record's existing EventTime is kept. If you specify TtlDuration you must also provide an EventTime, otherwise a ValidationError is returned. TargetStores must include the OnlineStore; a value that resolves to the OfflineStore only is rejected.

" } }, "shapes":{ @@ -124,6 +141,15 @@ "exception":true, "synthetic":true }, + "ConflictException":{ + "type":"structure", + "members":{ + "Message":{"shape":"Message"} + }, + "documentation":"

The request conflicts with the current state of the record. This is returned by UpdateRecord when the supplied EventTime is not greater than the record's current EventTime.

", + "error":{"httpStatusCode":409}, + "exception":true + }, "BatchGetRecordError":{ "type":"structure", "required":[ @@ -662,6 +688,38 @@ "member":{"shape":"BatchGetRecordIdentifier"}, "min":0 }, + "UpdateRecordRequest":{ + "type":"structure", + "required":[ + "FeatureGroupName", + "RecordIdentifierValueAsString", + "Features" + ], + "members":{ + "FeatureGroupName":{ + "shape":"FeatureGroupNameOrArn", + "documentation":"

The name or Amazon Resource Name (ARN) of the feature group that contains the record you want to update. The feature group must use an OnlineStoreConfig StorageType of Standard_V2 or InMemory.

", + "location":"uri", + "locationName":"FeatureGroupName" + }, + "RecordIdentifierValueAsString":{ + "shape":"ValueAsString", + "documentation":"

The value for the RecordIdentifier that uniquely identifies the record to update, in string format.

" + }, + "Features":{ + "shape":"Record", + "documentation":"

The feature values to write to the record. Only the features included here are written; features that are not listed are preserved. Pass EventTime as a feature in this list. You can update up to 100 features per call.

" + }, + "TargetStores":{ + "shape":"TargetStores", + "documentation":"

A list of stores to which the update is applied. By default, Feature Store applies the update to all of the stores that you're using for the FeatureGroup. A value that resolves to the OfflineStore only is rejected.

" + }, + "TtlDuration":{ + "shape":"TtlDuration", + "documentation":"

Time to live duration, where the record is hard deleted after the expiration time is reached; ExpiresAt = EventTime + TtlDuration. Specifying TtlDuration requires EventTime to be present in Features.

" + } + } + }, "ValidationError":{ "type":"structure", "members":{ diff --git a/sagemaker-core/sample/sagemaker/2017-07-24/service-2.json b/sagemaker-core/sample/sagemaker/2017-07-24/service-2.json index 6b551c5fd6..5a21eb4eee 100644 --- a/sagemaker-core/sample/sagemaker/2017-07-24/service-2.json +++ b/sagemaker-core/sample/sagemaker/2017-07-24/service-2.json @@ -38648,6 +38648,10 @@ "TtlDuration":{ "shape":"TtlDuration", "documentation":"

Time to live duration, where the record is hard deleted after the expiration time is reached; ExpiresAt = EventTime + TtlDuration. For information on HardDelete, see the DeleteRecord API in the Amazon SageMaker API Reference guide.

" + }, + "StorageType":{ + "shape":"StorageType", + "documentation":"

The online store storage type to migrate the feature group to. Use this parameter to migrate an existing feature group from Standard to Standard_V2 storage format, enabling support for the UpdateRecord operation. Migration is a one-way operation and cannot be reversed.

" } }, "documentation":"

Updates the feature group online store configuration.

" @@ -45368,6 +45372,7 @@ "type":"string", "enum":[ "Standard", + "Standard_V2", "InMemory" ] }, diff --git a/sagemaker-core/src/sagemaker/core/resources.py b/sagemaker-core/src/sagemaker/core/resources.py index 322ff2f356..398ea76fe2 100644 --- a/sagemaker-core/src/sagemaker/core/resources.py +++ b/sagemaker-core/src/sagemaker/core/resources.py @@ -12282,6 +12282,64 @@ def put_record( response = client.put_record(**operation_input_args) logger.debug(f"Response: {response}") + @Base.add_validate_call + def update_record( + self, + record_identifier_value_as_string: StrPipeVar, + features: List[FeatureValue], + target_stores: Optional[List[StrPipeVar]] = Unassigned(), + ttl_duration: Optional[TtlDuration] = Unassigned(), + session: Optional[Session] = None, + region: Optional[str] = None, + ) -> None: + """ + The UpdateRecord API performs a feature-level write to a Record in a feature group whose OnlineStoreConfig StorageType is Standard_V2 or InMemory. Only the supplied Features are written; features not included are preserved. The record must already exist in the online store. + + Parameters: + record_identifier_value_as_string: The value for the RecordIdentifier that uniquely identifies the record to update, in string format. + features: The list of FeatureValues to update. Only the features included here are written; features that are not listed are preserved. Pass EventTime as a feature in this list. A maximum of 100 features can be updated in a single request. + target_stores: A list of stores to which the update is applied. By default, Feature Store applies the update to all of the stores that you're using for the FeatureGroup. A value that resolves to the OfflineStore only is rejected. + ttl_duration: Time to live duration, where the record is hard deleted after the expiration time is reached; ExpiresAt = EventTime + TtlDuration. Specifying TtlDuration requires EventTime to be present in Features. + session: Boto3 session. + region: Region name. + + Raises: + botocore.exceptions.ClientError: This exception is raised for AWS service related errors. + The error message and error code can be parsed from the exception as follows: + ``` + try: + # AWS service call here + except botocore.exceptions.ClientError as e: + error_message = e.response['Error']['Message'] + error_code = e.response['Error']['Code'] + ``` + AccessForbidden: You do not have permission to perform an action. + ConflictException: There was a conflict when you attempted to modify a record; the supplied EventTime was not greater than the record's current EventTime. + InternalFailure: An internal failure occurred. Try your request again. If the problem persists, contact Amazon Web Services customer support. + ResourceNotFound: A resource that is required to perform an action was not found. + ServiceUnavailable: The service is currently unavailable. + ValidationError: There was an error validating your request. + """ + + operation_input_args = { + "FeatureGroupName": self.feature_group_name, + "RecordIdentifierValueAsString": record_identifier_value_as_string, + "Features": features, + "TargetStores": target_stores, + "TtlDuration": ttl_duration, + } + # serialize the input request + operation_input_args = serialize(operation_input_args) + logger.debug(f"Serialized input request: {operation_input_args}") + + client = Base.get_sagemaker_client( + session=session, region_name=region, service_name="sagemaker-featurestore-runtime" + ) + + logger.debug(f"Calling update_record API") + response = client.update_record(**operation_input_args) + logger.debug(f"Response: {response}") + @Base.add_validate_call def delete_record( self, diff --git a/sagemaker-core/src/sagemaker/core/shapes/shapes.py b/sagemaker-core/src/sagemaker/core/shapes/shapes.py index e0e79a5751..a52715c21c 100644 --- a/sagemaker-core/src/sagemaker/core/shapes/shapes.py +++ b/sagemaker-core/src/sagemaker/core/shapes/shapes.py @@ -14626,9 +14626,11 @@ class OnlineStoreConfigUpdate(Base): Attributes ---------------------- ttl_duration: Time to live duration, where the record is hard deleted after the expiration time is reached; ExpiresAt = EventTime + TtlDuration. For information on HardDelete, see the DeleteRecord API in the Amazon SageMaker API Reference guide. + storage_type: The online store storage type to migrate the feature group to. Use this parameter to migrate an existing feature group from Standard to Standard_V2 storage format, enabling support for the UpdateRecord operation. Migration is a one-way operation and cannot be reversed. """ ttl_duration: Optional[TtlDuration] = Unassigned() + storage_type: Optional[StrPipeVar] = Unassigned() class Parent(Base): diff --git a/sagemaker-core/src/sagemaker/core/tools/additional_operations.json b/sagemaker-core/src/sagemaker/core/tools/additional_operations.json index 25ba6b0fce..adf5696a5d 100644 --- a/sagemaker-core/src/sagemaker/core/tools/additional_operations.json +++ b/sagemaker-core/src/sagemaker/core/tools/additional_operations.json @@ -459,6 +459,14 @@ "method_type": "object", "service_name": "sagemaker-featurestore-runtime" }, + "UpdateRecord": { + "operation_name": "UpdateRecord", + "resource_name": "FeatureGroup", + "method_name": "update_record", + "return_type": "None", + "method_type": "object", + "service_name": "sagemaker-featurestore-runtime" + }, "DeleteRecord": { "operation_name": "DeleteRecord", "resource_name": "FeatureGroup", diff --git a/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py b/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py index 439fef722d..e29f2c0f02 100644 --- a/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py +++ b/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py @@ -14100,7 +14100,10 @@ "type": "structure", }, "OnlineStoreConfigUpdate": { - "members": [{"name": "TtlDuration", "shape": "TtlDuration", "type": "structure"}], + "members": [ + {"name": "TtlDuration", "shape": "TtlDuration", "type": "structure"}, + {"name": "StorageType", "shape": "StorageType", "type": "string"}, + ], "type": "structure", }, "OnlineStoreSecurityConfig": { @@ -18687,6 +18690,16 @@ ], "type": "structure", }, + "UpdateRecordRequest": { + "members": [ + {"name": "FeatureGroupName", "shape": "FeatureGroupNameOrArn", "type": "string"}, + {"name": "RecordIdentifierValueAsString", "shape": "ValueAsString", "type": "string"}, + {"name": "Features", "shape": "Record", "type": "list"}, + {"name": "TargetStores", "shape": "TargetStores", "type": "list"}, + {"name": "TtlDuration", "shape": "TtlDuration", "type": "structure"}, + ], + "type": "structure", + }, "UpdateTrialComponentResponse": { "members": [{"name": "TrialComponentArn", "shape": "TrialComponentArn", "type": "string"}], "type": "structure", diff --git a/sagemaker-core/tests/unit/generated/test_feature_store_operations.py b/sagemaker-core/tests/unit/generated/test_feature_store_operations.py index 9b54acb34d..f42a37f23e 100644 --- a/sagemaker-core/tests/unit/generated/test_feature_store_operations.py +++ b/sagemaker-core/tests/unit/generated/test_feature_store_operations.py @@ -10,7 +10,7 @@ # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. -"""Unit tests for FeatureGroup batch_write_record and list_records methods.""" +"""Unit tests for FeatureGroup batch_write_record, list_records, and update_record methods.""" from __future__ import absolute_import import pytest @@ -22,8 +22,10 @@ BatchWriteRecordResponse, FeatureValue, ListRecordsResponse, + OnlineStoreConfigUpdate, TtlDuration, ) +from sagemaker.core.utils.utils import serialize @pytest.fixture @@ -322,3 +324,77 @@ def test_list_records_accepts_next_token_parameter( call_kwargs = mock_client.list_records.call_args[1] # The explicitly passed next_token should be used, not self.next_token assert call_kwargs["NextToken"] == "list-records-page-2-token" + + +class TestUpdateRecord: + """Tests for FeatureGroup.update_record method (feature-level writes on Standard_V2).""" + + @patch("sagemaker.core.resources.Base.get_sagemaker_client") + def test_update_record_success(self, mock_get_client, mock_feature_group): + """Test that update_record calls the client with FeatureGroupName and Features.""" + mock_client = MagicMock() + mock_get_client.return_value = mock_client + + mock_feature_group.update_record( + record_identifier_value_as_string="r1", + features=[ + FeatureValue(feature_name="city", value_as_string="seattle"), + FeatureValue(feature_name="EventTime", value_as_string="1700000000"), + ], + ) + + mock_get_client.assert_called_once_with( + session=None, region_name=None, service_name="sagemaker-featurestore-runtime" + ) + mock_client.update_record.assert_called_once() + call_kwargs = mock_client.update_record.call_args[1] + assert call_kwargs["FeatureGroupName"] == "test-feature-group" + assert call_kwargs["RecordIdentifierValueAsString"] == "r1" + assert [f["FeatureName"] for f in call_kwargs["Features"]] == ["city", "EventTime"] + + @patch("sagemaker.core.resources.Base.get_sagemaker_client") + def test_update_record_serializes_target_stores_and_ttl( + self, mock_get_client, mock_feature_group + ): + """Test update_record serializes optional target_stores and ttl_duration.""" + mock_client = MagicMock() + mock_get_client.return_value = mock_client + + mock_feature_group.update_record( + record_identifier_value_as_string="r1", + features=[FeatureValue(feature_name="EventTime", value_as_string="1700000000")], + target_stores=["OnlineStore"], + ttl_duration=TtlDuration(unit="Days", value=7), + ) + + call_kwargs = mock_client.update_record.call_args[1] + assert call_kwargs["TargetStores"] == ["OnlineStore"] + assert call_kwargs["TtlDuration"] == {"Unit": "Days", "Value": 7} + + @patch("sagemaker.core.resources.Base.get_sagemaker_client") + def test_update_record_does_not_pass_next_token(self, mock_get_client, mock_feature_group): + """update_record has no pagination; self.next_token must not leak into the call.""" + mock_client = MagicMock() + mock_get_client.return_value = mock_client + + mock_feature_group.update_record( + record_identifier_value_as_string="r1", + features=[FeatureValue(feature_name="city", value_as_string="seattle")], + ) + + call_kwargs = mock_client.update_record.call_args[1] + assert "NextToken" not in call_kwargs + + +class TestOnlineStoreConfigUpdateStorageType: + """UpdateFeatureGroup -> Standard_V2 migration via OnlineStoreConfigUpdate.storage_type.""" + + def test_storage_type_field_present(self): + cfg = OnlineStoreConfigUpdate(storage_type="Standard_V2") + assert cfg.storage_type == "Standard_V2" + + def test_storage_type_serializes_to_pascal_case(self): + # Confirms the value reaches UpdateFeatureGroup on the wire as StorageType. + assert serialize(OnlineStoreConfigUpdate(storage_type="Standard_V2")) == { + "StorageType": "Standard_V2" + } diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/__init__.py index 34ff18e2bb..f789f53c49 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/__init__.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/__init__.py @@ -56,6 +56,7 @@ ingest_dataframe, list_records, load_feature_definitions_from_dataframe, + update_record, ) # Classes (local) @@ -119,6 +120,7 @@ "ingest_dataframe", "list_records", "load_feature_definitions_from_dataframe", + "update_record", # Classes "AthenaQuery", "DatasetBuilder", diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py index 8c6d9b2615..d5f206f77e 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py @@ -4,14 +4,16 @@ import logging import os import time +from collections import Counter from pathlib import Path import re -from typing import Any, Dict, Sequence, Union +from typing import Any, Dict, Optional, Sequence, Union import boto3 import pandas import pandas as pd from pandas import DataFrame, Series, read_csv +from sagemaker.core.shapes import FeatureValue, TtlDuration from sagemaker.core.utils.utils import Unassigned from sagemaker.mlops.feature_store import FeatureGroup as CoreFeatureGroup, FeatureGroup from sagemaker.core.helper.session_helper import Session @@ -25,6 +27,7 @@ StringFeatureDefinition, ) from sagemaker.mlops.feature_store.ingestion_manager_pandas import IngestionManagerPandas +from sagemaker.mlops.feature_store.inputs import TargetStoreEnum from sagemaker.core.utils import unique_name_from_base @@ -94,8 +97,12 @@ "SELECT, DESCRIBE, and ALTER permissions on the table in Lake Formation, " "in addition to IAM permissions.\n" "If this feature group uses IAM governance, ensure your role has " - "glue:GetTable and glue:UpdateTable permissions on the feature group's Glue table." - ) + "glue:GetTable and glue:UpdateTable permissions on the feature group's Glue table." + ) + +# UpdateRecord supports at most 100 features per call. +MAX_UPDATE_RECORD_FEATURES = 100 + def _get_athena_client(session: Session): """Get Athena client from session.""" @@ -578,6 +585,107 @@ def list_records( return fg.list_records(**kwargs) +def _to_feature_value(feature: Union[FeatureValue, Dict[str, Any]]) -> FeatureValue: + """Coerce a dict or FeatureValue into a core FeatureValue. + + Args: + feature: A ``FeatureValue`` or a dict with ``feature_name`` (or ``FeatureName``) + and exactly one of ``value_as_string``/``ValueAsString`` or + ``value_as_string_list``/``ValueAsStringList``. + + Returns: + A ``FeatureValue`` instance. + """ + if isinstance(feature, FeatureValue): + return feature + if isinstance(feature, dict): + name = feature.get("feature_name", feature.get("FeatureName")) + value = feature.get("value_as_string", feature.get("ValueAsString")) + value_list = feature.get("value_as_string_list", feature.get("ValueAsStringList")) + kwargs: Dict[str, Any] = {"feature_name": name} + if value is not None: + kwargs["value_as_string"] = value + if value_list is not None: + kwargs["value_as_string_list"] = value_list + return FeatureValue(**kwargs) + raise TypeError(f"Unsupported feature type: {type(feature)}. Expected FeatureValue or dict.") + + +@_telemetry_emitter(Feature.FEATURE_STORE, "update_record") +def update_record( + feature_group_name: str, + record_identifier_value_as_string: str, + features: Sequence[Union[FeatureValue, Dict[str, Any]]], + target_stores: Optional[Sequence[str]] = None, + ttl_duration: Optional[TtlDuration] = None, + region: str = None, +) -> None: + """Perform a feature-level (partial) write to a record via the UpdateRecord API. + + ``UpdateRecord`` is supported only for feature groups whose ``OnlineStoreConfig`` + ``StorageType`` is ``Standard_V2`` or ``InMemory``. Unlike ``PutRecord``, which overwrites + the whole record, only the features supplied in ``features`` are written; any feature not + included is preserved. This avoids the ``GetRecord`` -> merge -> ``PutRecord`` round + trip and prevents lost writes when independent pipelines own different features on the + same record. The record must already exist in the online store (use ``PutRecord`` to + create it); otherwise the service returns ``ResourceNotFound``. + + Args: + feature_group_name: Name or ARN of the FeatureGroup to update (``Standard_V2`` or + ``InMemory`` online store). + record_identifier_value_as_string: The record identifier value, in string format. + features: The features to update (up to 100). Each entry is a ``FeatureValue`` or a + dict. Features not listed here are preserved. Pass ``EventTime`` as a feature + in this list, not as a top-level parameter. + target_stores: Stores to apply the update to. Defaults to all stores configured on + the FeatureGroup. A value that resolves to the ``OfflineStore`` only is rejected. + ttl_duration: Time to live for the record; ``ExpiresAt = EventTime + TtlDuration``. + The service requires the record's event-time feature to be present in ``features``. + region: Region name. + + Raises: + ValueError: If ``features`` is empty, exceeds 100 entries, contains duplicate + feature names, or ``target_stores`` resolves to the ``OfflineStore`` only. + """ + if not features: + raise ValueError("features must contain at least one feature to update.") + if len(features) > MAX_UPDATE_RECORD_FEATURES: + raise ValueError( + f"features may contain at most {MAX_UPDATE_RECORD_FEATURES} entries, " + f"got {len(features)}." + ) + + feature_values = [_to_feature_value(f) for f in features] + + feature_names = [fv.feature_name for fv in feature_values] + duplicates = sorted(name for name, count in Counter(feature_names).items() if count > 1) + if duplicates: + raise ValueError(f"Duplicate feature names are not allowed: {duplicates}.") + + resolved_target_stores = list(target_stores) if target_stores is not None else None + if resolved_target_stores is not None and set(resolved_target_stores) == { + TargetStoreEnum.OFFLINE_STORE.value + }: + raise ValueError( + "UpdateRecord cannot target the OfflineStore only; include the OnlineStore." + ) + + fg = CoreFeatureGroup.get(feature_group_name=feature_group_name, region=region) + + kwargs: Dict[str, Any] = { + "record_identifier_value_as_string": record_identifier_value_as_string, + "features": feature_values, + } + if resolved_target_stores is not None: + kwargs["target_stores"] = resolved_target_stores + if ttl_duration is not None: + kwargs["ttl_duration"] = ttl_duration + if region is not None: + kwargs["region"] = region + + fg.update_record(**kwargs) + + @_telemetry_emitter(Feature.FEATURE_STORE, "get_feature_group_as_dataframe") def get_feature_group_as_dataframe( feature_group_name: str, diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/inputs.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/inputs.py index f264059eb3..470741854c 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/inputs.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/inputs.py @@ -10,6 +10,7 @@ class OnlineStoreStorageTypeEnum(Enum): """Storage types for online store.""" STANDARD = "Standard" IN_MEMORY = "InMemory" + STANDARD_V2 = "Standard_V2" class TableFormatEnum(Enum): """Offline store table formats.""" diff --git a/sagemaker-mlops/tests/integ/test_feature_store_update_record.py b/sagemaker-mlops/tests/integ/test_feature_store_update_record.py new file mode 100644 index 0000000000..a0a07ca490 --- /dev/null +++ b/sagemaker-mlops/tests/integ/test_feature_store_update_record.py @@ -0,0 +1,129 @@ +"""Integration tests for UpdateRecord (feature-level writes) on Standard_V2 feature groups. + +These tests require an updated boto3/botocore that ships the UpdateRecord operation and a +region where Feature Store Standard_V2 storage is available. +""" +import time +import pytest +import pandas as pd + +from sagemaker.core.helper.session_helper import Session, get_execution_role +from sagemaker.mlops.feature_store import ( + FeatureGroup, + OnlineStoreConfig, + OnlineStoreStorageTypeEnum, +) +from sagemaker.mlops.feature_store.feature_utils import ( + load_feature_definitions_from_dataframe, + ingest_dataframe, + update_record, +) +from sagemaker.core.utils import unique_name_from_base + + +@pytest.fixture(scope="module") +def sagemaker_session(): + return Session() + + +@pytest.fixture(scope="module") +def role(): + return get_execution_role() + + +@pytest.fixture +def feature_group_name(): + return unique_name_from_base("integ-test-updaterecord-fg") + + +@pytest.fixture +def sample_dataframe(): + current_time = int(time.time()) + return pd.DataFrame( + { + "record_id": [f"id-{i}" for i in range(3)], + "city": ["seattle", "portland", "denver"], + "temperature": [float(i) for i in range(3)], + "event_time": [float(current_time + i) for i in range(3)], + } + ) + + +def cleanup_feature_group(feature_group_name): + try: + fg = FeatureGroup.get(feature_group_name=feature_group_name) + fg.delete() + time.sleep(2) + except Exception: + pass + + +def _create_standard_v2_group(feature_group_name, sample_dataframe, role): + feature_definitions = load_feature_definitions_from_dataframe(sample_dataframe) + fg = FeatureGroup.create( + feature_group_name=feature_group_name, + record_identifier_feature_name="record_id", + event_time_feature_name="event_time", + feature_definitions=feature_definitions, + role_arn=role, + online_store_config=OnlineStoreConfig( + enable_online_store=True, + storage_type=OnlineStoreStorageTypeEnum.STANDARD_V2.value, + ), + ) + fg.wait_for_status("Created") + return fg + + +def test_update_record_preserves_unlisted_features( + feature_group_name, sample_dataframe, role +): + """UpdateRecord writes only the supplied features; others are preserved.""" + try: + fg = _create_standard_v2_group(feature_group_name, sample_dataframe, role) + ingest_dataframe(feature_group_name=feature_group_name, data_frame=sample_dataframe) + time.sleep(15) + + new_event_time = float(int(time.time()) + 100) + update_record( + feature_group_name=feature_group_name, + record_identifier_value_as_string="id-0", + features=[ + {"feature_name": "city", "value_as_string": "tacoma"}, + {"feature_name": "event_time", "value_as_string": str(new_event_time)}, + ], + ) + time.sleep(10) + + record = fg.get_record(record_identifier_value_as_string="id-0") + values = {fv.feature_name: fv.value_as_string for fv in record.record} + assert values["city"] == "tacoma" # updated + assert values["temperature"] == "0.0" # preserved (not in the update) + finally: + cleanup_feature_group(feature_group_name) + + +def test_update_record_stale_event_time_conflict( + feature_group_name, sample_dataframe, role +): + """An EventTime not greater than the current one is rejected with a conflict.""" + from botocore.exceptions import ClientError + + try: + _create_standard_v2_group(feature_group_name, sample_dataframe, role) + ingest_dataframe(feature_group_name=feature_group_name, data_frame=sample_dataframe) + time.sleep(15) + + stale_event_time = "1.0" + with pytest.raises(ClientError) as exc: + update_record( + feature_group_name=feature_group_name, + record_identifier_value_as_string="id-0", + features=[ + {"feature_name": "city", "value_as_string": "tacoma"}, + {"feature_name": "event_time", "value_as_string": stale_event_time}, + ], + ) + assert exc.value.response["Error"]["Code"] in ("ConflictException", "ValidationError") + finally: + cleanup_feature_group(feature_group_name) diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py index 93311b2284..07cb6a3b8f 100644 --- a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py +++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py @@ -870,3 +870,129 @@ def test_calls_load_feature_definitions(self, mock_fg_class): ) mock_fg.load_feature_definitions.assert_called_once() + + +class TestUpdateRecord: + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_updates_subset_of_features(self, mock_fg_class): + mock_fg = MagicMock() + mock_fg_class.get.return_value = mock_fg + + from sagemaker.mlops.feature_store.feature_utils import update_record + + update_record( + feature_group_name="test-fg", + record_identifier_value_as_string="r1", + features=[ + {"feature_name": "city", "value_as_string": "seattle"}, + {"feature_name": "EventTime", "value_as_string": "1700000000"}, + ], + ) + + mock_fg_class.get.assert_called_once_with(feature_group_name="test-fg", region=None) + mock_fg.update_record.assert_called_once() + call = mock_fg.update_record.call_args.kwargs + assert call["record_identifier_value_as_string"] == "r1" + assert [fv.feature_name for fv in call["features"]] == ["city", "EventTime"] + assert "target_stores" not in call + assert "ttl_duration" not in call + + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_accepts_feature_value_objects(self, mock_fg_class): + mock_fg = MagicMock() + mock_fg_class.get.return_value = mock_fg + + from sagemaker.mlops.feature_store.feature_utils import update_record + from sagemaker.core.shapes import FeatureValue + + update_record( + feature_group_name="test-fg", + record_identifier_value_as_string="r1", + features=[FeatureValue(feature_name="city", value_as_string="seattle")], + ) + call = mock_fg.update_record.call_args.kwargs + assert call["features"][0].feature_name == "city" + + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_passes_target_stores_and_ttl(self, mock_fg_class): + mock_fg = MagicMock() + mock_fg_class.get.return_value = mock_fg + + from sagemaker.mlops.feature_store.feature_utils import update_record + from sagemaker.core.shapes import TtlDuration + + ttl = TtlDuration(unit="Days", value=7) + update_record( + feature_group_name="test-fg", + record_identifier_value_as_string="r1", + features=[ + {"feature_name": "city", "value_as_string": "seattle"}, + {"feature_name": "EventTime", "value_as_string": "1700000000"}, + ], + target_stores=["OnlineStore", "OfflineStore"], + ttl_duration=ttl, + ) + call = mock_fg.update_record.call_args.kwargs + assert call["target_stores"] == ["OnlineStore", "OfflineStore"] + assert call["ttl_duration"] is ttl + + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_empty_features_raises(self, mock_fg_class): + from sagemaker.mlops.feature_store.feature_utils import update_record + + with pytest.raises(ValueError, match="at least one feature"): + update_record("test-fg", "r1", features=[]) + mock_fg_class.get.assert_not_called() + + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_too_many_features_raises(self, mock_fg_class): + from sagemaker.mlops.feature_store.feature_utils import update_record + + features = [{"feature_name": f"f{i}", "value_as_string": "v"} for i in range(101)] + with pytest.raises(ValueError, match="at most 100"): + update_record("test-fg", "r1", features=features) + + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_duplicate_features_raises(self, mock_fg_class): + from sagemaker.mlops.feature_store.feature_utils import update_record + + with pytest.raises(ValueError, match="Duplicate feature names"): + update_record( + "test-fg", + "r1", + features=[ + {"feature_name": "city", "value_as_string": "a"}, + {"feature_name": "city", "value_as_string": "b"}, + ], + ) + + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_offline_only_target_store_raises(self, mock_fg_class): + from sagemaker.mlops.feature_store.feature_utils import update_record + + with pytest.raises(ValueError, match="OfflineStore only"): + update_record( + "test-fg", + "r1", + features=[{"feature_name": "city", "value_as_string": "a"}], + target_stores=["OfflineStore"], + ) + + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_ttl_without_event_time_is_forwarded_to_service(self, mock_fg_class): + """The SDK does not client-side enforce the ttl/event-time rule (feature name is + FG-defined); it forwards to the service, which validates.""" + mock_fg = MagicMock() + mock_fg_class.get.return_value = mock_fg + + from sagemaker.mlops.feature_store.feature_utils import update_record + from sagemaker.core.shapes import TtlDuration + + ttl = TtlDuration(unit="Days", value=7) + update_record( + "test-fg", + "r1", + features=[{"feature_name": "city", "value_as_string": "a"}], + ttl_duration=ttl, + ) + assert mock_fg.update_record.call_args.kwargs["ttl_duration"] is ttl diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_inputs.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_inputs.py index 44e3ec6085..5290e96d92 100644 --- a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_inputs.py +++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_inputs.py @@ -32,6 +32,9 @@ def test_standard(self): def test_in_memory(self): assert OnlineStoreStorageTypeEnum.IN_MEMORY.value == "InMemory" + def test_standard_v2(self): + assert OnlineStoreStorageTypeEnum.STANDARD_V2.value == "Standard_V2" + class TestTableFormatEnum: def test_glue(self):