AWS Common Runtime で Amazon S3 のスループットを高速化

このブログは 2023 年 11 月 27 日に James Bornholt、Abhinav Goyal、Jonathan Henson、Andrew Kutsy によって執筆された内容を日本語化したものです。原文はこちらを参照してください。

データはあらゆる機械学習パイプラインの中心にあります。基盤モデル (FM) の事前トレーニング、ビジネス固有のデータによる FM のファインチューニング、推論クエリの実行など、機械学習ライフサイクルのあらゆる段階においてコンピューティングリソースを常時稼働させて有用な作業を実行できるようにするには、低コストで高性能なデータストレージが必要です。お客様はトレーニングデータやモデルチェックポイントの保存に Amazon Simple Storage Service (Amazon S3) を使用することが推奨されています。理由として弾力性、複数のテラビット/秒にスケールする性能、そして S3 Intelligent-Tiering のようなストレージクラスによってアクセスパターンが変化した際に自動的にストレージコストを節約できるからです。

Amazon S3 との間のデータ転送を自動的に高速化し、機械学習パイプラインの基盤をさらに強化する AWS Command Line Interface (AWS CLI) と AWS SDK for Python (Boto3) の新しいアップデートを発表しました。AWS CLI と Boto3 は、Amazon S3 との間の高スループットデータ転送を提供するために特別に設計・構築された AWS Common Runtime (CRT) S3 クライアントと統合されました。この統合は現在、Amazon EC2 Trn1、P4d、P5 インスタンスタイプではデフォルトで有効化されており、他のインスタンスタイプではオプトインとして有効化することができます。

AWS Common Runtime とは何か

Amazon S3 は、どの HTTP クライアントからもアクセスできるシンプルな REST API で、お客様から好評をいただいています。ただし、大量のデータを扱うアプリケーションで最高のパフォーマンスを得るには、リクエストの並列化、タイムアウト、再試行、バックオフなどのパフォーマンスのベストプラクティスを実装する必要があります。数年前、私たちはこれらのパターンを各 AWS SDK で再実装していて、お客様も独自のアプリケーションにこれらのパターンを実装しなければならないことに気付きました。私たちは、こうした一般的な設計パターンを再実装することなく、どのアプリケーションからでも簡単に S3 の伸縮自在なパフォーマンスにアクセスできるようにしたいと考えていました。

このポータブルなパフォーマンスを実現するために、AWS Common Runtime (CRT) を構築しました。CRT は C で書かれたオープンソースライブラリの集合体で、高性能 HTTP クライアントや暗号化ライブラリなど、AWS サービスとやり取りするための共通機能を実装しています。CRT ライブラリは連携して AWS サービスに高速で信頼性の高いクライアントエクスペリエンスを提供します。Amazon S3 の場合、CRT にはネイティブ S3 クライアントが含まれており、自動リクエストの並列化、リクエストのタイムアウトと再試行、接続の再利用と管理を実装して、ネットワークインターフェイスの過負荷を回避します。たとえば、S3 から 1 つの巨大なオブジェクトをダウンロードする場合、CRT クライアントは複数のバイト範囲を自動的に並行してダウンロードするため、スループットが向上し、ネットワークインターフェースを上手に活用します。

CRT は Java や C++ を含む複数の AWS SDK で既に本番環境で利用可能であり、以前は AWS CLI の実験的なオプションとして提供されていました。また、オープンソースファイルクライアントである Mountpoint for Amazon S3 の基盤でもあります。 私たちはCRTをTrn1、P4d、P5 EC2 インスタンスタイプの AWS CLI と Boto3 で提供を開始しました。これらのインスタンスタイプは大きな CPU とネットワークインターフェースを持ち、これらのパフォーマンス設計パターンから最も恩恵を受けることができます。 他のインスタンスタイプについては、AWS CLI または Boto3 アプリケーションで CRT を使用するように選択すると、多くの場合、自動的にパフォーマンス向上が得られます。

ML パイプラインのパフォーマンス向上

AWS Common Runtime で実現できる可能性のあるパフォーマンス向上を実証するために、ML ライフサイクルの各ステップを表す 4 つのベンチマークデータセットを収集しました：

• Caltech-256：画像データセットで、平均 40 kB サイズの 30,607 の小さな画像ファイルを含み、総データセットサイズは 1.1 GB です。
• Caltech-256-WebDataset：同じ Caltech 256 の画像データセットですが、WebDataset 形式を使用して保存されており、複数の画像を 100 MB のシャードオブジェクトにまとめています。シャード化されたデータセットは、ML トレーニングで Amazon S3 を使用する際に、より高いパフォーマンスを達成できることがよくあります。
• C4-en：Common Crawl corpusに基づく C4 データセットの英語サブセットで、320 MB の 1,024 ファイルを含みます。
• Checkpoint：単一の 7.6 GB PyTorch チェックポイントファイルで、大規模 ML モデルのシャード化されたチェックポイントを代表しています。

AWS CLIを使用して、これらの各データセットを trn1n.32xlarge EC2 インスタンスからアップロードおよびダウンロードしました。AWS CRTが有効になっていない場合と有効になっている場合の両方です。結果は次のとおりです：

CRT は、AWS CLI の最新リリースへの更新以外の追加作業なしで、これらのベンチマーク全体で 2 ～ 6 倍のパフォーマンス向上を実現します。CRT を有効にして Boto3 を使用する Python アプリケーションでも、同様のパフォーマンス向上が見られました。

アプリケーションで CRT を使い始める

AWS CLI で CRT を使用するには、まず AWS CLI の最新バージョンをインストールしてください。まだ AWS CLI のバージョン 2 にアップデートしていない場合は絶好の機会です。CRT との統合はバージョン 2 でのみ利用可能です。Trn1、P4d、または P5 EC2 インスタンスで実行している場合はこれだけで準備完了です — aws s3 sync のような CLI コマンドを使用する際に、CRT はデフォルトで有効になります。その他のインスタンスタイプでは、次のコマンドを実行して CRT を有効にすることができます：

aws configure set s3.preferred_transfer_client crt

Boto3 で CRT を使用するには、まず追加の crt 機能を含む Boto3 をインストールしていることを確認してください。たとえば、pip を使用してインストールする場合は、以下を実行します：

pip install boto3[crt]

Trn1、P4d、および P5 インスタンスでは、Boto3 が crt 機能付きでインストールされると、upload_file と download_file の呼び出しに自動的に CRT が使用されます。たとえば、CRT を使用してファイルを S3 にアップロードするには、以下を実行します：

import boto3
s3 = boto3.client('s3')
s3.upload_file('/tmp/hello.txt', 'doc-example-bucket', 'hello.txt')

また、s3transfer パッケージを使用して Boto3 で CRT にアクセスできますが、このパッケージはまだ一般提供を開始されておらず、将来変更される可能性があります。

パフォーマンスチューニング

CRT は S3 を使用するアプリケーションのパフォーマンスを自動的に最適化します。デフォルト設定では多くの状況で速度が向上します。これらのデフォルトでは、CPU トポロジー、メモリ量、Elastic Network Adapter (ENA) インターフェイスの数とレイアウトなど、実行中のインスタンスタイプの仕様に基づいて CRT が自動的に設定されます。これらの詳細に基づいて、 CRTはS3リクエストの並列化戦略（並列接続の数、各リクエストのサイズ、S3 IPアドレスごとのリクエスト数など）を選択します。

CRT 転送で使用されるネットワーク帯域幅の量を制限する場合など、状況によってはこれらのデフォルトをオーバーライドしたい場合があります。AWS CLI で CRT を使用する場合は、target_bandwidth パラメーターを設定することでデフォルトをオーバーライドできます。たとえば、転送を 5 ギガビット/秒に制限するには、以下を実行します：

aws configure set s3.target_bandwidth 5Gb/s

注意事項とオプトアウト

CLI と Boto3 用の CRT の今回のリリースでは、多くの ML アプリケーションのパフォーマンスが向上しますが、注意すべき点が 3 つあります。

マルチプロセス実行

CRT は、S3 リクエストを複数のスレッドで並行して行うことで、高スループットのデータ転送を実現します。これは、一度に 1 つの S3 クライアントしか使用しないアプリケーションに最適です。これらのスレッドはインスタンスの vCPU に分散する可能性があるためです。ただし、それぞれが独自の S3 クライアントを作成する複数のプロセスを使用する場合、これらのスレッドは互いに競合して S3 のパフォーマンスを低下させる可能性があります。また、これらの複数のクライアントがネットワーク帯域幅をめぐって競合し、輻輳が発生してパフォーマンスが低下する場合もあります。

AWS CLI と Boto3 の CRT 統合では、複数のプロセスが CRT ベースの S3 クライアントを作成していることを自動的に検出し、このような場合は非 CRT ベースの S3 クライアントにフォールバックします。このフォールバックにより、CRT クライアントが 1 つだけ存在するようになるため、システム上で競合が発生するリスクが軽減されますが、その結果、他のクライアントのパフォーマンスが低下する可能性があります。この制限は複数の S3 クライアントにのみ影響します。1 つの S3 クライアントを同じプロセス内の複数のスレッドで共有することも、同じ AWS CLI 呼び出し内の多数の S3 転送で共有することもできます。

アプリケーションが複数のプロセスで独自の S3 クライアントを作成してしまうパターンは主に 2 つあります。1 つ目は、AWS CLI の呼び出しを複数同時に実行すると、各 CLI プロセスに独自の S3 クライアントが存在することになります。たとえば、以前に AWS CLI を parallel や xargs -P ユーティリティを実行してパフォーマンスを向上させたことがある場合は、複数の AWS CLI プロセスを使用し、それぞれに独自の S3 クライアントを用意することになります。新しい CRT 統合では、CLI プロセスを 1 つだけ使用し、CLI に転送の並列処理を任せることをお勧めします。2 つ目に、PyTorch のような ML フレームワークで Boto3 を使用する場合、データをロードするために複数のワーカープロセス （たとえば、PyTorch の DataLoader の num_workers 引数） を使用することになります。

マルチリージョンおよびクロスリージョンアクセス

AWS CLI と Boto3 の CRT 統合は、現在 S3 バケットの自動リージョン検出をサポートしていません。つまり、アプリケーションがインスタンスが実行されているリージョンとは異なるリージョンの S3 バケットにアクセスする場合、ターゲットリージョンを手動で指定する必要があります。これを行うには、AWS CLI の –region argument を使用するか、AWS CLI と Boto3 の両方に AWS_REGION 環境変数を設定します。Boto3 の場合、リージョンはクライアントの作成時に設定されるため、この制限により、1 つの S3 クライアントは 1 つのリージョンのバケットにしかアクセスできないことにもなります。複数のリージョンのバケットにアクセスする必要がある場合は、複数のクライアントを作成する必要があります。

Transfer コンフィグレーション

Boto3 の CRT 統合は、クライアントを転送ごとに設定するための TransferConfig API をサポートしていません。代わりに、CRT はネットワーク帯域幅を最大化するようにクライアントを自動的に設定し、同じプロセスで同時に実行される S3 リクエストすべてでその帯域幅を共有します。

CRT からのオプトアウト

これらの制限のいずれかを回避する必要がある場合は、CRT をオプトアウトできます。AWS CLI の CRT 統合を無効にするには、以下を実行します：

aws configure set s3.preferred_transfer_client classic

同様に、Boto3 の CRT S3 統合を無効にするには、boto3 転送コールで使用するときに TransferConfig の preferred_transfer_client を classic に設定してください。

from boto3.s3.transfer import TransferConfig
config = TransferConfig(preferred_transfer_client='classic')
client = boto3.client('s3', region_name='us-west-2')
client.upload_file('/tmp/file', Bucket='doc-example-bucket',
Key='test_file', Config=config)

まとめと今後の改善点

Amazon S3 は伸縮自在でパフォーマンスが高いため、ML トレーニングデータやモデルチェックポイントを保存するのに最適です。AWS CLI と Boto3 の改善により、ML パイプラインで S3 にアクセスする際のパフォーマンスをさらに簡単に最適化できるようになり、ジョブをより早く完了し、コストを削減できるようになりました。将来的には、AWS Common Runtime をデフォルトでより多くのインスタンスタイプで有効にし、よりきめ細かな調整機能を公開して、ワークロードのパフォーマンスをさらに最適化できるようにする予定です。AWS CLI、Boto3、AWS Common Runtime はすべてオープンソースプロジェクトであり、それぞれの GitHub リポジトリに関するフィードバックをお待ちしています。