Skip to content

Add BatchWriteRecord and ListRecod apis to Sagemaker Feature Store Runtime #5983

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
BassemHalim wants to merge 1 commit into
base: master
Choose a base branch
from
+800 −12
Open

Add BatchWriteRecord and ListRecod apis to Sagemaker Feature Store Runtime #5983

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
190 changes: 188 additions & 2 deletions sagemaker-core/sample/sagemaker-featurestore-runtime/2020-07-01/service-2.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,13 @@
"endpointPrefix":"featurestore-runtime.sagemaker",
"jsonVersion":"1.1",
"protocol":"rest-json",
"protocols":["rest-json"],
"serviceFullName":"Amazon SageMaker Feature Store Runtime",
"serviceId":"SageMaker FeatureStore Runtime",
"signatureVersion":"v4",
"signingName":"sagemaker",
"uid":"sagemaker-featurestore-runtime-2020-07-01"
"uid":"sagemaker-featurestore-runtime-2020-07-01",
"auth":["aws.auth#sigv4"]
},
"operations":{
"BatchGetRecord":{
Expand All @@ -28,6 +30,23 @@
],
"documentation":"<p>Retrieves a batch of <code>Records</code> from a <code>FeatureGroup</code>.</p>"
},
"BatchWriteRecord":{
"name":"BatchWriteRecord",
"http":{
"method":"POST",
"requestUri":"/BatchWriteRecord"
},
"input":{"shape":"BatchWriteRecordRequest"},
"output":{"shape":"BatchWriteRecordResponse"},
"errors":[
{"shape":"ValidationError"},
{"shape":"ResourceNotFound"},
{"shape":"InternalFailure"},
{"shape":"ServiceUnavailable"},
{"shape":"AccessForbidden"}
],
"documentation":"<p>Writes a batch of <code>Records</code> to one or more <code>FeatureGroup</code>s. Use this API for bulk ingestion of records into the <code>OnlineStore</code> and <code>OfflineStore</code>.</p> <p>You can set the ingested records to expire at a given time to live (TTL) duration after the record's event time by specifying the <code>TtlDuration</code> parameter. A request level <code>TtlDuration</code> applies to all entries that do not specify their own <code>TtlDuration</code>.</p>"
},
"DeleteRecord":{
"name":"DeleteRecord",
"http":{
Expand All @@ -41,7 +60,7 @@
{"shape":"ServiceUnavailable"},
{"shape":"AccessForbidden"}
],
"documentation":"<p>Deletes a <code>Record</code> from a <code>FeatureGroup</code> in the <code>OnlineStore</code>. Feature Store supports both <code>SoftDelete</code> and <code>HardDelete</code>. For <code>SoftDelete</code> (default), feature columns are set to <code>null</code> and the record is no longer retrievable by <code>GetRecord</code> or <code>BatchGetRecord</code>. For <code>HardDelete</code>, the complete <code>Record</code> is removed from the <code>OnlineStore</code>. In both cases, Feature Store appends the deleted record marker to the <code>OfflineStore</code>. The deleted record marker is a record with the same <code>RecordIdentifer</code> as the original, but with <code>is_deleted</code> value set to <code>True</code>, <code>EventTime</code> set to the delete input <code>EventTime</code>, and other feature values set to <code>null</code>.</p> <p>Note that the <code>EventTime</code> specified in <code>DeleteRecord</code> should be set later than the <code>EventTime</code> of the existing record in the <code>OnlineStore</code> for that <code>RecordIdentifer</code>. If it is not, the deletion does not occur:</p> <ul> <li> <p>For <code>SoftDelete</code>, the existing (not deleted) record remains in the <code>OnlineStore</code>, though the delete record marker is still written to the <code>OfflineStore</code>.</p> </li> <li> <p> <code>HardDelete</code> returns <code>EventTime</code>: <code>400 ValidationException</code> to indicate that the delete operation failed. No delete record marker is written to the <code>OfflineStore</code>.</p> </li> </ul> <p>When a record is deleted from the <code>OnlineStore</code>, the deleted record marker is appended to the <code>OfflineStore</code>. If you have the Iceberg table format enabled for your <code>OfflineStore</code>, you can remove all history of a record from the <code>OfflineStore</code> using Amazon Athena or Apache Spark. For information on how to hard delete a record from the <code>OfflineStore</code> with the Iceberg table format enabled, see <a href=\"https://docs.aws.amazon.com/sagemaker/latest/dg/feature-store-delete-records-offline-store.html#feature-store-delete-records-offline-store\">Delete records from the offline store</a>.</p>"
"documentation":"<p>Deletes a <code>Record</code> from a <code>FeatureGroup</code> in the <code>OnlineStore</code>. Feature Store supports both <code>SoftDelete</code> and <code>HardDelete</code>. For <code>SoftDelete</code> (default), feature columns are set to <code>null</code> and the record is no longer retrievable by <code>GetRecord</code> or <code>BatchGetRecord</code>. For <code>HardDelete</code>, the complete <code>Record</code> is removed from the <code>OnlineStore</code>. In both cases, Feature Store appends the deleted record marker to the <code>OfflineStore</code>. The deleted record marker is a record with the same <code>RecordIdentifer</code> as the original, but with <code>is_deleted</code> value set to <code>True</code>, <code>EventTime</code> set to the delete input <code>EventTime</code>, and other feature values set to <code>null</code>.</p> <p>Note that the <code>EventTime</code> specified in <code>DeleteRecord</code> should be set later than the <code>EventTime</code> of the existing record in the <code>OnlineStore</code> for that <code>RecordIdentifer</code>. If it is not, the deletion does not occur:</p> <ul> <li> <p>For <code>SoftDelete</code>, the existing (not deleted) record remains in the <code>OnlineStore</code>, though the delete record marker is still written to the <code>OfflineStore</code>.</p> </li> <li> <p> <code>HardDelete</code> returns <code>EventTime</code>: <code>400 ValidationException</code> to indicate that the delete operation failed. No delete record marker is written to the <code>OfflineStore</code>.</p> </li> </ul> <p>When a record is deleted from the <code>OnlineStore</code>, the deleted record marker is appended to the <code>OfflineStore</code>. If you have the Iceberg table format enabled for your <code>OfflineStore</code>, you can remove all history of a record from the <code>OfflineStore</code> using Amazon Athena or Apache Spark. For information on how to hard delete a record from the <code>OfflineStore</code> with the Iceberg table format enabled, see <a href=\"https://docs.aws.amazon.com/sagemaker/latest/dg/feature-store-delete-records.html#feature-store-delete-records-offline-store\">Delete records from the offline store</a>.</p>"
},
"GetRecord":{
"name":"GetRecord",
Expand All @@ -60,6 +79,24 @@
],
"documentation":"<p>Use for <code>OnlineStore</code> serving from a <code>FeatureStore</code>. Only the latest records stored in the <code>OnlineStore</code> can be retrieved. If no Record with <code>RecordIdentifierValue</code> is found, then an empty result is returned. </p>"
},
"ListRecords":{
"name":"ListRecords",
"http":{
"method":"POST",
"requestUri":"/FeatureGroup/{FeatureGroupName}/ListRecords"
},
"input":{"shape":"ListRecordsRequest"},
"output":{"shape":"ListRecordsResponse"},
"errors":[
{"shape":"ValidationError"},
{"shape":"ResourceNotFound"},
{"shape":"InternalFailure"},
{"shape":"ServiceUnavailable"},
{"shape":"AccessForbidden"}
],
"documentation":"<p>Lists the <code>RecordIdentifier</code> values of all records stored in a <code>FeatureGroup</code>'s <code>OnlineStore</code>. This enables you to discover which records exist without retrieving the full record data.</p>",
"readonly":true
},
"PutRecord":{
"name":"PutRecord",
"http":{
Expand Down Expand Up @@ -216,6 +253,98 @@
"member":{"shape":"BatchGetRecordResultDetail"},
"min":0
},
"BatchWriteRecordEntries":{
"type":"list",
"member":{"shape":"BatchWriteRecordEntry"},
"max":25,
"min":1
},
"BatchWriteRecordEntry":{
"type":"structure",
"required":[
"FeatureGroupName",
"Record"
],
"members":{
"FeatureGroupName":{
"shape":"FeatureGroupNameOrArn",
"documentation":"<p>The name or Amazon Resource Name (ARN) of the <code>FeatureGroup</code> to write the record to.</p>"
},
"Record":{
"shape":"Record",
"documentation":"<p>List of FeatureValues to be inserted. This will be a full over-write.</p>"
},
"TargetStores":{
"shape":"TargetStores",
"documentation":"<p>A list of stores to which you're adding the record. By default, Feature Store adds the record to all of the stores that you're using for the <code>FeatureGroup</code>.</p>"
},
"TtlDuration":{
"shape":"TtlDuration",
"documentation":"<p>Time to live duration for this entry, where the record is hard deleted after the expiration time is reached; <code>ExpiresAt</code> = <code>EventTime</code> + <code>TtlDuration</code>. This overrides the request level <code>TtlDuration</code>.</p>"
}
},
"documentation":"<p>An entry to write as part of a <code>BatchWriteRecord</code> request.</p>"
},
"BatchWriteRecordError":{
"type":"structure",
"required":[
"Entry",
"ErrorCode",
"ErrorMessage"
],
"members":{
"Entry":{
"shape":"BatchWriteRecordEntry",
"documentation":"<p>The entry that failed to be written.</p>"
},
"ErrorCode":{
"shape":"ValueAsString",
"documentation":"<p>The error code for the failed record write.</p>"
},
"ErrorMessage":{
"shape":"Message",
"documentation":"<p>The error message for the failed record write.</p>"
}
},
"documentation":"<p>The error that has occurred when attempting to write a record in a batch.</p>"
},
"BatchWriteRecordErrors":{
"type":"list",
"member":{"shape":"BatchWriteRecordError"},
"min":0
},
"BatchWriteRecordRequest":{
"type":"structure",
"required":["Entries"],
"members":{
"Entries":{
"shape":"BatchWriteRecordEntries",
"documentation":"<p>A list of records to write. Each entry specifies the <code>FeatureGroup</code>, the record data, and optionally target stores and a TTL duration.</p>"
},
"TtlDuration":{
"shape":"TtlDuration",
"documentation":"<p>Time to live duration applied to all entries in the batch that do not specify their own <code>TtlDuration</code>; <code>ExpiresAt</code> = <code>EventTime</code> + <code>TtlDuration</code>. For information on HardDelete, see the <a href=\"https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_feature_store_DeleteRecord.html\">DeleteRecord</a> API in the Amazon SageMaker API Reference guide.</p>"
}
}
},
"BatchWriteRecordResponse":{
"type":"structure",
"required":[
"Errors",
"UnprocessedEntries"
],
"members":{
"Errors":{
"shape":"BatchWriteRecordErrors",
"documentation":"<p>A list of errors that occurred when writing records in the batch.</p>"
},
"UnprocessedEntries":{
"shape":"UnprocessedBatchWriteRecordEntries",
"documentation":"<p>A list of entries that were not processed. These entries can be retried.</p>"
}
}
},
"Boolean":{"type":"boolean"},
"DeleteRecordRequest":{
"type":"structure",
"required":[
Expand Down Expand Up @@ -364,6 +493,54 @@
"fault":true,
"synthetic":true
},
"ListRecordsMaxResults":{
"type":"integer",
"max":100,
"min":1
},
"ListRecordsNextToken":{
"type":"string",
"max":8192,
"min":1
},
"ListRecordsRequest":{
"type":"structure",
"required":["FeatureGroupName"],
"members":{
"FeatureGroupName":{
"shape":"FeatureGroupNameOrArn",
"documentation":"<p>The name or Amazon Resource Name (ARN) of the feature group to list records from.</p>",
"location":"uri",
"locationName":"FeatureGroupName"
},
"MaxResults":{
"shape":"ListRecordsMaxResults",
"documentation":"<p>The maximum number of record identifiers to return in a single page of results. For the <code>InMemory</code> tier, this value is a hint and not a strict requirement. The response may contain more or fewer results than the specified <code>MaxResults</code>.</p>"
},
"NextToken":{
"shape":"ListRecordsNextToken",
"documentation":"<p>A token to resume pagination of <code>ListRecords</code> results.</p>"
},
"IncludeSoftDeletedRecords":{
"shape":"Boolean",
"documentation":"<p>If set to <code>true</code>, the result includes records that have been soft deleted.</p>"
}
}
},
"ListRecordsResponse":{
"type":"structure",
"required":["RecordIdentifiers"],
"members":{
"RecordIdentifiers":{
"shape":"RecordIdentifierList",
"documentation":"<p>A list of record identifier values for the records stored in the <code>OnlineStore</code>.</p>"
},
"NextToken":{
"shape":"ListRecordsNextToken",
"documentation":"<p>A token to resume pagination if the response includes more record identifiers than <code>MaxResults</code>.</p>"
}
}
},
"Message":{
"type":"string",
"max":2048
Expand Down Expand Up @@ -400,6 +577,10 @@
"member":{"shape":"FeatureValue"},
"min":1
},
"RecordIdentifierList":{
"type":"list",
"member":{"shape":"ValueAsString"}
},
"RecordIdentifiers":{
"type":"list",
"member":{"shape":"ValueAsString"},
Expand Down Expand Up @@ -471,6 +652,11 @@
"type":"integer",
"min":1
},
"UnprocessedBatchWriteRecordEntries":{
"type":"list",
"member":{"shape":"BatchWriteRecordEntry"},
"min":0
},
"UnprocessedIdentifiers":{
"type":"list",
"member":{"shape":"BatchGetRecordIdentifier"},
Expand Down
116 changes: 116 additions & 0 deletions sagemaker-core/src/sagemaker/core/resources.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12394,6 +12394,122 @@ def batch_get_record(
transformed_response = transform(response, "BatchGetRecordResponse")
return BatchGetRecordResponse(**transformed_response)

@Base.add_validate_call
def batch_write_record(
self,
entries: List[BatchWriteRecordEntry],
ttl_duration: Optional[TtlDuration] = Unassigned(),
session: Optional[Session] = None,
region: Optional[str] = None,
) -> Optional[BatchWriteRecordResponse]:
"""
Writes a batch of Records to one or more FeatureGroups.

Parameters:
entries: A list of records to write. Each entry specifies the FeatureGroup, the record data, and optionally target stores and a TTL duration.
ttl_duration: Time to live duration applied to all entries in the batch that do not specify their own TtlDuration; ExpiresAt = EventTime + TtlDuration. For information on HardDelete, see the DeleteRecord API in the Amazon SageMaker API Reference guide.
session: Boto3 session.
region: Region name.

Returns:
BatchWriteRecordResponse

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.
InternalFailure: An internal failure occurred. Try your request again. If the problem persists, contact Amazon Web Services customer support.
ResourceNotFound: Resource being access is not found.
ServiceUnavailable: The service is currently unavailable.
ValidationError: There was an error validating your request.
"""

operation_input_args = {
"Entries": entries,
"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 batch_write_record API")
response = client.batch_write_record(**operation_input_args)
logger.debug(f"Response: {response}")

transformed_response = transform(response, "BatchWriteRecordResponse")
return BatchWriteRecordResponse(**transformed_response)

@Base.add_validate_call
def list_records(
self,
max_results: Optional[int] = Unassigned(),
next_token: Optional[StrPipeVar] = Unassigned(),
include_soft_deleted_records: Optional[bool] = Unassigned(),
session: Optional[Session] = None,
region: Optional[str] = None,
) -> Optional[ListRecordsResponse]:
"""
Lists the RecordIdentifier values of all records stored in a FeatureGroup's OnlineStore.

Parameters:
max_results: The maximum number of record identifiers to return in a single page of results. For the InMemory tier, this value is a hint and not a strict requirement. The response may contain more or fewer results than the specified MaxResults.
next_token: A token to resume pagination of ListRecords results.
include_soft_deleted_records: If set to true, the result includes records that have been soft deleted.
session: Boto3 session.
region: Region name.

Returns:
ListRecordsResponse

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.
InternalFailure: An internal failure occurred. Try your request again. If the problem persists, contact Amazon Web Services customer support.
ResourceNotFound: Resource being access is not found.
ServiceUnavailable: The service is currently unavailable.
ValidationError: There was an error validating your request.
"""

operation_input_args = {
"FeatureGroupName": self.feature_group_name,
"MaxResults": max_results,
"NextToken": next_token,
"IncludeSoftDeletedRecords": include_soft_deleted_records,
}
# 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 list_records API")
response = client.list_records(**operation_input_args)
logger.debug(f"Response: {response}")

transformed_response = transform(response, "ListRecordsResponse")
return ListRecordsResponse(**transformed_response)


class FeatureMetadata(Base):
"""
Expand Down
Loading
Loading