SageMaker Feature Storeに一括書き込みとレコード列挙APIが登場

発表の概要

Amazon SageMaker Feature Storeは、機械学習モデル向けに特徴量(フィーチャー)を保存・共有・管理するためのフルマネージドサービスです。オンライン推論向けの低レイテンシなオンラインストアと、履歴保存や学習データ向けのオフラインストアを備え、バッチとストリーミングの両方の取り込みパターンをサポートしています。今回、このFeature Storeに対して2つの新しいAPI、BatchWriteRecordとListRecordsが追加されました。

1つ目のBatchWriteRecordは、複数レコードをまとめて書き込むためのAPIで、従来のPutRecordをループで呼び出す際のオーバーヘッドを解消します。2つ目のListRecordsは、フィーチャーグループ内のレコード識別子を列挙するAPIで、特にIn-Memoryストレージ層では、識別子を失った場合の復旧手段がなかったことへの対策でもあります。本記事では、各APIの動作とコード例を基に、既存の運用とどう変わるかを整理します。

BatchWriteRecordの特徴

従来のPutRecordは、1回のAPI呼び出しで1つのフィーチャーグループに1レコードだけを書き込む仕様でした。EventTimeに基づく条件付き書き込みにより、オンラインストアでは最新バージョンが維持されますが、スループットを上げるにはN×M回の呼び出しが必要になります。例えば不正検知パイプラインが毎秒10,000レコードを5つのフィーチャーグループに取り込む場合、毎秒50,000回のAPI呼び出しが必要になるという課題がありました。

BatchWriteRecordは、1回のリクエストで最大25エントリを複数のフィーチャーグループに書き込めます。各レコードは独立して成功または失敗し、一部が失敗してもリクエスト全体は失敗しません。また、PutRecordと同じEventTimeベースの順序保証があり、古いレコードが新しいレコードを上書きすることはありません。失敗したレコードはレスポンスで返され、リトライは失敗分だけを行えばよい設計です。

TTLはレコード単位・リクエスト単位・フィーチャーグループ単位の3段階で指定でき、優先順位が定義されています。ただし、リクエスト全体のエントリ数は最大25件で、この上限は複数フィーチャーグループにまたがる場合も総数で適用されます。TargetStoresを指定すると、オンラインストアとオフラインストアのどちらに書き込むかを各エントリで制御できます。

ListRecordsの位置づけ

Feature StoreにはGetRecordやDeleteRecordがありましたが、いずれも正確なレコード識別子を知っている必要がありました。標準層(Standard tier)ではオフラインストアをAmazon Athenaで検索する回避策があったものの、コストと時間がかかります。In-Memory層ではデフォルトでオフラインストアがなく、識別子が失われるとレコードを復元する手段は事実上ありませんでした。

ListRecordsは、ページネーションを使ってフィーチャーグループ内のレコード識別子を列挙します。標準層(Amazon DynamoDBバックエンド)では最新バージョンを対象に、論理削除済みや期限切れのレコードを除外して返します。In-Memory層(Redisバックエンド)ではキーをスキャンし、ソフト削除されたレコードや内部システムキーを除外して識別子を抽出します。このAPIにより、識別子の喪失によるデータ復旧不能のリスクを軽減できます。

既存APIとの比較と注意点

PutRecordとBatchWriteRecordの主な違いは、1回のリクエストで扱えるレコード数と対象フィーチャーグループ数です。記事では、BatchWriteRecordは従来のPutRecordを置き換えるものではなく、高スループットが必要なワークロード向けの新しい選択肢として説明されています。次の表に主な違いをまとめます。

利用にあたっては、BatchWriteRecordにはPutRecordと同様のIAM権限が必要です。具体的には、対象フィーチャーグループのARNに対するsagemaker:BatchWriteRecordとsagemaker:PutRecordの許可が必要になります。また、部分成功APIであるため、失敗したレコードだけを指数バックオフで再送信する設計が推奨されています。TargetStoresでオンラインストアとオフラインストアの書き込み先も制御できます。

一方、記事ではListRecordsのページネーションの詳細や、各APIの制限値(例えば1リクエストあたりの最大バイト数やクォータなど)については明らかにされていません。また、In-Memory層でのスキャンがパフォーマンスに与える影響も、実際の運用で確認する必要があるとみられます。これらの点は、本番環境へ適用する前に検証すべき項目といえるでしょう。