AWSを活用した機械翻訳のためのGPU並列処理環境の構築

Tags
AWSを活用した機械翻訳のためのGPU並列処理環境の構築
Page content

はじめに

こんにちは、ストックマークでエンジニアをしている麻生です。ストックマークでは、「Anews」というウェブサービスを提供しています。この度、Anewsで新機能導入のために日次バッチの大規模なインフラ変更を行い、GPU並列処理環境を構築しましたのでご紹介します。

組織の自律化を支援するナレッジプラットフォーム「Anews」

Anewsは国内外30,000メディアのニュースを毎日収集し、最先端の自然言語処理で個人や組織のミッションに即したニュースをレコメンドします。コメント機能で簡単にチームにアイデアを共有でき、社内の知見者から学ぶことでチームの情報感度が底上げされます。

エンタープライズを中心に、累計1500社以上のお客様にご利用いただいているサービスです。

Anewsのプロダクトイメージ
Anewsのプロダクトイメージ

英語記事をレコメンドする上での課題

Anewsでは、記事への行動履歴からユーザーや組織の好みを学習し、記事をレコメンドしています。ユーザーが記事に対してシェアやコメントをするほど行動履歴が蓄積され、記事のレコメンド精度が向上する仕組みです。

レコメンドについてはこちらの記事に詳しい解説があります。

アクションを取らなければ記事のレコメンド精度の向上が見込めないのですが、そこで問題になるのが英語記事です。日本語を母国語とする以上、英語記事を読むのは負荷のかかる作業で、どうしても記事へのアクションが少なくなってしまいます。そのため、英語記事のレコメンド精度は上がりにくく、精度が低いことでレコメンドに繋がる行動履歴も蓄積されない、という悪循環に陥っていました。

お客様によっては海外の情報が重要なケースもあり、英語記事のレコメンド精度をいかに上げるかが課題でした。

日本語記事への行動履歴活用によるレコメンド精度向上

そこで今回持ち上がったのが、日本語記事への行動履歴を元に、英語記事をレコメンドするというアイデアです。英語記事を読まないユーザーも日本語記事への行動履歴は蓄積されているため、それが使えれば英語記事のレコメンド精度を向上させられるはずです。

単純なアイデアですが、英語と日本語では文章の性質が全く異なるため、日本語記事への行動履歴をそのまま英語記事に活用できません。一度、日本語記事を英語記事に翻訳することで、あたかも英語記事へのアクションがあったかのようなデータを作成する必要があります。

レコメンドの流れ
レコメンドの流れ

社内の機械学習チームの検証結果を踏まえ、transformer-align というモデルによる翻訳処理をプロダクトに組み込むことにしました。

翻訳処理導入に向けたインフラの見直し

今回の施策によってワークフローが変更になります。新しいワークフローでは、英語記事のレコメンドの前に翻訳処理が入ります。

変更後のワークフロー
変更後のワークフロー

この変更により、EC2インスタンス1台体制だったインフラを大幅に見直す必要が出てきました。

ポイントは以下の3点です。


翻訳処理用のGPU並列処理環境を用意する

翻訳処理をCPUで実行すると時間がかかりすぎて、ユーザーがアプリを使うまでに処理が終わりません。GPUが利用可能、かつ、記事量に応じて並列数を増やしていけるような環境が必要でした。


既存処理環境(CPU環境)と翻訳処理環境(GPU環境)を分離する

GPUを必要としない既存処理もGPU環境で実行してしまうと、必要以上にコストが高くなってしまいます。そのため、既存処理と翻訳処理の環境を別々のサーバに分けることにしました。


ワークフローエンジンを導入する

CPU環境とGPU環境を分離すると、サーバを跨いだワークフロー管理が必要になります。EC2インスタンス1台のときには不要だったワークフローエンジンを導入します。

翻訳機能導入後のインフラ

変更前後のインフラ構成は下図の通りです。選定理由を簡単にまとめていますが、ここに至るまでに動作確認とインフラの再構築を繰り返したり、AWSの方にアドバイスいただいたりと試行錯誤しながら進めました。

変更前後のインフラ
変更前後のインフラ構成


翻訳処理(GPU処理)の環境は AWS Batch + Amazon EC2 と AWS Lambda

翻訳処理には、GPUを利用可能、かつ、並列処理を実装しやすい AWS Batch + Amazon EC2 を採用しました。

並列処理は配列ジョブで実装しています。翻訳する記事数に応じて並列数が決まるように、一度 AWS Lambda で記事数をカウントして並列数を決定した後、ジョブキューにジョブを送っています。配列ジョブでは、各ジョブがAWS_BATCH_JOB_ARRAY_INDEXという環境変数を割り当てられます。各ジョブはこの環境変数によって自身が処理する記事を識別します。

翻訳処理の実行例がこちらで、10並列でジョブが走っています(キューのメッセージ取得、インスタンス起動、コンテナ起動とオーバーヘッドがあるため完全に同じタイミングにはなっていません)。最大32並列で実行可能な設定にしています。

並列処理実行例
並列処理実行例


既存処理(CPU処理)の環境は Amazon ECS + AWS Fargate

既存処理はEC2インスタンス1台で運用してきましたが、EC2 は他の AWS リソースとの連携が困難なため、翻訳処理と連携が取れるように環境を変更することにしました。 候補には AWS Lambda、Amazon ECS、AWS Batch が挙げられるのですが、既に一部のタスクを Amazon ECS+Fargate で実行できるようにしていたため、それベースに他のタスクに拡張することとしました。

EC2インスタンス1台でやりくりしていた処理をコンテナレベルで分散することで、処理速度の向上というメリットも得られました。翻訳処理の追加による時間の増分を含めても、全体の処理時間は短くなっています。


ワークフローエンジンは AWS Step Functions

AWS のマネージドなワークフローエンジンである AWS Step Functions を採用しました。ワークフローエンジンには他にも Airflow, Digdag, Argo とあるのですが、社内での利用実績と AWS リソースとの連携のしやすさから採用しました。

Step Functions はワークフローの可視化が可能で、以下のように実行状況を確認できます。

Step Functionis で可視化されたワークフロー
Step Functionis で可視化されたワークフロー

インフラ再構築のポイント

今回の開発を振り返ると、いかに効率的に試行錯誤を繰り返せるかがポイントでした。特に、利用経験のないAWSサービスを組み込んでインフラを変更する場合、1回目の変更で全て順調に進むことはないと思います。大抵、想定とは違った挙動をしていたり、設定値の変更が必要だったりします。 その点、開発チームで導入を進めてきた Infrastracture as Code (IaC) と CI/CD に助けられました。インフラ変更、デプロイの手順はシンプルで、ミスを抑えて高速で試行錯誤を繰り返すことができました。


IaC

IaC ではコードで書いてある通りにインフラの追加、削除、変更がされるので、大幅なインフラ変更もコマンド数行で一気に実行できます。開発環境では細かい修正を繰り返しやすいですし、本番環境には開発環境で確認済みのコードを適用するだけなので精神的負担がかなり抑えられます。

また、コード自体がインフラの設計書として機能するので、現行インフラの引き継ぎをする際にも全体像を理解する助けになりました。

使用しているのは Terraform と Serverless Framework です。Lambda と Step Functions は Serverless Framework、それ以外は Terraform というような使い分けをしています。統一した方がいいという考え方もあるのですが、一方にしかない機能があるので併用しています(例えば、Lambda に関数をデプロイする際にzipが必要なのですが、Serverless Framework は zipを自動作成してくれます)。


CI/CD

バッチでは以下のような CI/CD パイプラインを構築しています。

  • CI機構:GitHub でプルリクエストを作成すると CodeBuild へ Webhook イベントが届き、自動テスト用のワークフローが実行される
  • CD機構:GitHub で特定のタグを付与すると CodeBuild へ Webhook イベントが届き、自動デプロイ用のワークフローが実行される

今回特に役に立ったのはCD機構です。インフラ変更により、以下の手順が必要になるのですが、全て CodeBuild 上で自動実行されるため、開発者からみると GitHub でタグを付与するだけです。

  • Dockerイメージをビルドする
  • Dockerイメージを ECR へプッシュする
  • 最新のDockerイメージを使うように ECS のタスク定義を更新する
  • 最新のDockerイメージを使うように Batch のジョブ定義を更新する

IaC と同様に、開発環境での試行錯誤のしやすさ、本番環境リリース時のミス防止、精神的負担の軽減といったメリットがありました。

自動デプロイフロー
自動デプロイフロー

おわりに

今回、Anewsへの新機能追加と、それに伴うインフラ変更についてご紹介しました。 ストックマークでは、お客様の新価値創造を支援するため、機械学習チームとエンジニアが連携してプロダクトをつくり上げています。 Wantedlyに募集職種を載せておりますので、興味のある方はぜひご覧ください!

告知

今回ご紹介した内容に加えて、ブログには書ききれなかった内容もお話するイベントを実施いたします。オンラインで実施いたしますのでご興味をお持ちの方はぜひ下記URLよりエントリーお願いいたします!

参考