Sansan Tech Blog

Sansanのものづくりを支えるメンバーの技術やデザイン、プロダクトマネジメントの情報を発信

GCS→S3キャッシュプロキシで大規模ML学習のデータ転送コストを削減した話

AWS Google Cloud MLOps

はじめに

技術本部 研究開発部 Architectグループで1カ月強インターンをしていた金古侑大です。*1

Architectグループでは、研究開発に関連するインフラの設計、開発、運用をしています。今回のインターンでは、マルチクラウド環境でのデータ取得コストを削減するためのHTTPプロキシの開発に取り組みました。

Sansan株式会社では、AWS・Google Cloud双方を利用し、サービスを提供、開発をしています。しかし、クラウド間のデータ転送には大きなコストがかかります。特に、大規模な事前学習では、数十TBのデータを100エポック以上学習します。データセットを事前にコピーする方法もありますが、削除依頼への即時対応などデータガバナンス上の要件があります。また、巨大なデータセットの転送による学習開始までの待ち時間も課題となるため、現状は各エポックでソースから直接データを取得しています。

本記事では、この課題に対してどのようなアプローチを検討し、最終的にどのような解決策を実装したのか、そしてインターンを通じて学んだことについて共有します。

背景と課題

研究開発部では、AWS・Google Cloud双方を活用しています。今回は学習用の環境としてAWSを採用しており、Google Cloudからのデータ取得に転送コストが発生します。通常のプロダクト開発では、クラウド間のデータ転送が発生しても、その量は限定的なことが多いです。しかし、大規模な機械学習では話が変わります。 数十TBのデータを、各エポックで再読み込みして学習するため、クラウド間のデータ転送コストが無視できない規模になります。

深層学習は膨大なデータを何度も繰り返し学習するスタイルです。今回想定されるデータ量は約20TBで、これを100エポック以上学習することを考えると、クラウド間のデータ転送量は膨大になります。

上記の環境で学習をした場合、単純計算で約24万ドルという莫大なコストが発生してしまいます。

20TB × 100エポック × ($0.12/GB) ≒ $240,000 (約3,600万円：1 USD = ¥150)

このコスト問題を、研究員の方々にデータ配置やクラウド選択といった負担をかけずに解決することが今回のタスクでした。また、将来的にはGoogle Cloud上で同様の問題が発生した際にも対応できるよう、拡張性のある設計が求められていました。

解決策の検討

基本方針は「GCSに保存されたデータをS3にキャッシュすることでデータ転送料金を削減する」というもので、これを実現する方法として次の4つを検討しました。

検討した選択肢

  • Syncサービスを使用する方法: AWS DataSyncなどを利用し、事前にデータを同期
  • 独自のAPIサーバーを構築する方法: S3/GCSとは別の独自APIを提供
  • HTTPS MITM Proxyを構築する方法: 透過的にリクエストを仲介し、自動的にキャッシュ
  • HTTP Proxyを構築する方法: 明示的なエンドポイント経由でキャッシュ

初期段階での棄却

最初の2つの選択肢は早い段階で見送りました。

Syncサービスには複数の課題がありました。最初のデータ同期に長時間かかること、prefix構造的に不必要なデータもコピーせざるを得ないこと、さらにML側のコードをGCS用からS3用に書き換える必要があることです。

独自APIサーバーは、研究員がOpenAPIを読んで学習コードを大幅に修正する必要があり、API設計から保守まで開発コストも高いため、限られた期間での実現には向いていないと判断しました。

HTTPS MITM Proxyへの挑戦と方針転換

残る2つの選択肢のうち、当初はHTTPS MITM Proxyを選択しました。

MITM（Man-in-the-Middle）Proxyは、ML JobからのAWS・GCPへのリクエストを仲介します。ProxyがML Jobと、AWS・GCPそれぞれとTLSハンドシェイクを行うことで、通信を透過的に処理できます。

この方式を選択した理由は、ML側のコード変更が全く不要という点でした。環境変数（https_proxy）の設定のみで動作し、Proxy側で自動的にS3へキャッシュする仕組みを実現できます。また、学習用インスタンスへの依存を非常に低くでき、将来的には別のインスタンスに立てることも可能で、他のプロジェクトでも再利用できる汎用性の高い解決策になると考えました。

当初、boto3やgoogle.cloudがOAuthトークンのような認証をしていれば、HEADリクエストで認証を確認できると想定していました。認証が通ればキャッシュ用のbucketからデータを取得できます。こうした大きなメリットを見込んで、技術的な難易度は高いものの、チャレンジする価値があると判断しました。

しかし、実装を進める中で課題が明らかになりました。まず、AWS SigV4認証が想定以上に厳格だったことです。HEADリクエストで認証確認してからGETリクエストを送る方式を想定していましたが、GETの認証情報からHEADリクエストを生成して認証のみ確認することは技術的に困難でした。

さらに、独自証明書の用意と管理が必要になり、Sidecar構成やRoute53での実現も検討しましたが、いずれも運用負荷が高いという問題がありました。

こうした課題により、当初想定していた最大のメリットである「完全な透過性」の実現が困難だと明確になりました。それに伴い、TLS証明書管理や開発コストといったデメリットの比重が相対的に大きくなってしまいました。

上長からのフィードバックもあり、「まずはやりやすい方法でやってみて、それができてからより良いものを作る」 という方針の重要性を痛感しました。特に今回は時間も限られていたため、確実に価値を提供できるシンプルな方針を最初に採用すべきでした。ある程度実装できた後で状況を考慮して方針転換や継続を判断する道のほうが盤石であり、エンジニアが取るべき選択だったと思います。

HTTP Proxyの採択

方針転換に当たり、HTTP Proxyを選択しました。

ML JobはProxyを明示的に指定してHTTPリクエストを送信し、ProxyがS3・GCSのエンドポイントをミラーリングしたAPIサーバーとして動作します。

この方式を選択した理由は次の通りです。

まず、HTTPS MITM Proxyで書いたコードの多くを流用できることが大きな理由でした。既に実施した調査・検証を無駄にすることなく、より現実的なアプローチに切り替えられます。

次に、証明書管理が不要で、実装がHTTPS MITM Proxyより大幅にシンプルという点です。同じインスタンス内に構築することで権限管理がシンプルになり、証明書管理などの複雑な運用も不要になります。開発コストと運用コストのバランスが非常に良いという点も重要でした。

そして何より、ML側の変更がエンドポイント設定のみと最小限で済むにもかかわらず、コスト削減効果は変わらないという点が決定打となりました。この軽微なコード修正は、研究員の方々にとって許容範囲内であることを確認できました。

実装の詳細

技術スタック

  • 言語・フレームワーク: Go, Gin
  • クラウドSDK: aws-sdk-go-v2, google.cloud
  • インフラ: AWS EC2 (ParallelCluster), S3, GCS
  • 認証: EC2 IMDS, Workload Identity Federation
  • デプロイ: Systemd, Nginx
  • IaC: Terraform

技術的な工夫

1. Go + Ginの選択

Proxyに必要な高いパフォーマンスを確保しながら、開発コストを削減するためにGoとGinフレームワークを選択しました。

Goを選んだ最大の理由は、goroutineによる並行処理の優秀さです。大量の同時リクエストを効率的に処理できるだけでなく、大きなファイルの転送でもメモリ使用量を抑制できます。また、AWS SDKやGCP SDKといった豊富なライブラリが揃っており、クラウド連携も容易に実現できました。

他にもPython、C#、Rustといった言語も候補に挙がりましたが、PythonはProxyとしての性能面で劣る可能性が高く、C#とRustは開発経験がなかったため見送りました。

実装では、S3互換のAPIエンドポイントを作成しました。リクエストを受け取ったらまずS3キャッシュの存在を確認し、キャッシュヒットの場合はS3から直接返します。キャッシュミスの場合はGCSから取得しながら、非同期でS3にキャッシュする仕組みを構築しました。

2. ストリーミング対応

大きなファイルでもOOM（Out of Memory）にならないよう、ストリーミング転送を実装しました。今回使用するインスタンスには膨大なメモリが割り当てられており、基本的にOOMの心配はありません。ただし機械学習側でメモリを多く消費する可能性が高く、消費量を抑えたかったため実装に踏み切りました。

この課題に対して、GCSからS3へのコピー、HTTPレスポンスへの返却、すべてストリーミングで処理する方式を採用しました。バッファを使用し、読み込みと書き込みを同時に実行することで、AWS S3 ManagerとGCS Readerを活用してチャンク単位で転送しています。

これにより、ファイルサイズに関わらず一定のメモリ使用量で動作するようになりました。

3. Systemdによるプロセス管理とNginxによる負荷分散

当初はDocker Swarmを使用する予定でしたが、EC2のIMDSv2の制約により、Systemdによるプロセス管理に変更しました。

EC2では認証に関わる情報をIMDS（Instance Metadata Service）を通じて取得しますが、Docker SwarmのOverlay Networkを使用すると問題が発生します。IMDSv2はセキュリティのためホップ数制限（TTL=1）を設定しており、Overlay Network経由では169.254.169.254へのリクエストが到達できません。結果として、AWS認証情報の取得ができなくなってしまいます。

ホップ数の変更自体は可能ですが、今回はParallelClusterに任せて作っているインスタンスであり、IaCとして反映しづらい変更を加えたくありませんでした。そのため、Systemdでコンテナを使わずに直接プロセスとしてデプロイする方針に変更しました。

Nginxは引き続きリバースプロキシとして使用し、複数のProxyプロセスへロードバランシングを行います。least_connアルゴリズムで最も接続数の少ないプロセスに振り分け、HTTP/1.1の永続的接続を維持してオーバーヘッドを削減します。

4. ストラテジーパターンによる拡張性

キャッシュシステムをストラテジーパターンで実装することで、将来的な拡張を容易にしました。設計では、キャッシュ層とデータソース層を共通のインターフェースとして定義し、S3キャッシュとGCSソースをそれぞれ独立した実装として切り離しています。

この設計により、将来的にGCP側でも同様の仕組みを実装する際に既存コードの再利用が容易になります。また、新しいキャッシュストレージ（例: ElastiCache）への切り替えも簡単で、テスト時にモックへの差し替えも容易です。

動作フローとしては、まずリクエストを受け取ったらキャッシュするクラウドへのリクエストか確認します。キャッシュするクラウドであれば、キャッシュ層で取得を試みます。キャッシュヒット時は直接返却し、キャッシュミス時はソース層から返却しつつディスクに保存します。ディスクに保存後、非同期でキャッシュ先のクラウドに保存するという流れです。

成果

【従来】
数十TB × (GCSデータ転送料金 + NAT Gatewayコスト) × 100 エポック

【改善後】
数十TB × (GCSデータ転送料金 + NAT Gatewayコスト) × 1回

最終的に、データ転送料金を従来の1/100以下に削減できました。

今回は100エポックでの学習を想定していますが、今後さらに多くのエポック数で学習をする場合、より大きなコスト削減効果が期待できます。

また、パフォーマンス面では、チューニング前はProxyを通さない場合に比べて約50%の遅延がありました。しかし、ストリーミング処理の最適化などを行った結果、チューニング後には10〜20%程度の遅延まで抑えることができました。

最適化はNginxレイヤで実施しました。バックエンドProxyとの接続ではkeepaliveを有効化し、コネクションプールとリクエスト数・アイドル時間の上限を調整しました。これによりTCP接続確立のオーバーヘッドを削減しつつ高スループットを維持できます。また、大容量オブジェクトの転送に合わせてレスポンスバッファや一時ファイルサイズをチューニングし、メモリとディスクI/Oのバランスを取りながらストリーミング性能を引き上げています。さらに、クラウドストレージ特有のレイテンシを考慮して各種タイムアウト値を適切に設定し、無駄な再送や切断を避けています。バックエンド障害時にはproxy_next_upstreamで自動的に別インスタンスへフェイルオーバーし、安定性とレイテンシの両面で性能を向上させました。

研究員の方への影響として、従来は毎回プロダクト側に転送コストを共有する必要があり、研究員の方々はコストの懸念により大規模実験を気軽に踏み出せずにいました。

しかし、このProxy導入でデータ転送コストが大幅に削減され、共有不要なレベルまで下がりました。その結果、コスト問題で今までできなかった大規模事前学習へ取り組めるようになったという成果が得られました。

また、Proxyは事前のデータコピーが不要で、Systemdサービスを起動し、エンドポイント設定を変更するだけで利用開始できます。この手軽さも、研究員の方々に受け入れられた大きな要因だと思います。

おわりに

1カ月強という期間でしたが、非常に濃密なインターンでした。技術的なチャレンジはもちろん、方針転換という苦い経験も含めて、エンジニアとして大きく成長できたと感じています。

特に印象的だったのは、ほぼ毎日行われたインターンランチです。多様なバックグラウンドを持つ方々とのお話は毎回新鮮で、サービスの成り立ちや苦労話を聞く中で、会社への理解が深まりました。また、マネジャー職にも興味を持つようになり、まずはTech Leadを目指そうという新たな目標もできました。

最後に、メンターとして丁寧に指導してくださった方、チームメンバーの皆さま、そして貴重なフィードバックをくださった上長の方々に心から感謝申し上げます。今回学んだことを活かし、今後も成長していきたいと思います。

*1:本記事は、執筆者本人の了承を得て、代理で投稿しています。

© Sansan, Inc.