YJIT を有効化してパフォーマンス向上した話

はじめに

初めまして、delyでサーバーサイドエンジニア兼 PM している宇野です。
今回はクラシルリワードで Ruby の YJIT を有効化するまでにやったこと・結果を紹介します。

YJIT とは

まず初めに、YJIT について簡単に説明します。
YJITは、Shopify が開発した JIT(Just In Time) コンパイルで、バージョン 3.1.0 から実験的に導入され、バージョン 3.2.0 で正式にサポートされました。
Ruby 3.2.0 のリリースノートには、以下の画像が掲載されておりパフォーマンス改善されることが期待できます。
Ruby 3.2.0 の時点ではメモリ使用量が多く懸念がありましたが、3.3.0 では大きく改善されているみたいです。(※1)

(※1)

• 3.2 と 3.3 の比較
• 3.3 のリリースノート

クラシルリワードでの YJIT 有効化の流れ

当時クラシルリワードでは、Ruby 3.1.4, Rails 7.0.2 を使用していました。
公式のリリースノートや note さんの投稿を Slack で話すなど YJIT を有効化する機運が高まっていました。

YJIT 有効化のためには Ruby のバージョンを上げるだけで良いのですが、Rails が一部 Ruby 3.2.0 で削除されたメソッドを利用していました😭(3.2.0 のリリースノート)
そのため、Ruby だけではなく Rails もバージョンアップをすることにしました。
Rails と Ruby、さらに YJIT 有効化を同時に行うと不具合があった時の特定が難しいので、以下のように段階的に YJIT 有効化を進めました。

1. Rails を 7.1 系にバージョンアップ。1週間ほど様子見して問題がないことを確認
2. Ruby を 3.3.0 にあげる。1週間ほど様子見して問題がないことを確認
3. YJIT の有効化

Rails と Ruby のバージョンアップは、テストカバレッジが高くほとんどのエラーはテストで見つけることができました。
ただ、あるバッチ処理だけ Rails のバージョンを上げることでメモリ使用量がスパイクする問題が発生しました。

このバッチ処理では以下のように自作の find_in_batches のようなメソッドを使っており、ここで大量のメモリが使用されていました。
find_in_batches メソッドは id でソートするのですが、ここでは token でソートしつつ分割してレコードを取得したいという意図がありメソッドを自作していました。
これは pluck を select にすることで、メモリ使用量を抑えて解決できました。

# 一部コードを抜粋
def self.find_in_batches_order_by_token(batch_size:)
  token_offset = 0
  loop do
    tokens = where('token > ?', token_offset).unique_by_token.order(token: :asc).limit(batch_size)
    break if tokens.blank?

    yield(tokens)
    token_offset = tokens.last.token
  end
end

# pluck を select にすることで、メモリ使用量のスパイクが抑えられた
User.active.where(id: user_ids).pluck(:id).find_in_batches_order_by_token(batch_size: 10_000) do |block|
  ...
end

本題のYJIT の有効化は、以下でできます。簡単ですね😄

config/initializers/enable_yjit.rb

if defined? RubyVM::YJIT.enable
  Rails.application.config.after_initialize do
    RubyVM::YJIT.enable
  end
end

YJIT は様々なパラメーターがあるのですが、最適な値の検証が難しいのでリリースしてから調整すれば良いと考えデフォルトでリリースをしました。

パフォーマンスへの影響

• レスポンスタイム

デプロイしたタイミングで跳ねていますが、そこからのレスポンスタイムのベースラインが下がっています。

• メモリ使用率

YJIT は仕組み上 JIT コードやメタデータをメモリ上で保持をするのでメモリ使用率が上がることは把握していました。リリース前は懸念していましたが、許容範囲に収まっていて問題なさそうでした。

• ECS タスク数の削減

これはリリース前は考慮していなかったのですが、ECS の1タスクで捌けるリクエストが増えたことでタスク数が減りコスト削減にも繋がりました🎉
ただし、これはクラシルリワードのインフラ設定に依存する部分もあるので一般的に効果があるものではないと思います。

ここまでは嬉しいことなのですが、YJIT を有効化したことによりデプロイするとそのタイミングだけレスポンスタイムがスパイクするようになってしまいました💦

当時 クラシルリワードは Blue/Gree デプロイを採用しており、デプロイ時に一括でトラフィックを新ターゲットに変えるようにしていました。この結果、デプロイ直後に短い期間で JIT のコンパイルが多発するようになりグラフのような挙動になりました。ちょうど同じタイミングでカナリアリリースの導入を検討していたので、段階的に新ターゲットに切り替わるようになり解決しました。

おわりに

この記事では、YJIT 有効化までの流れとパフォーマンスへの影響を紹介しました。

YJIT は統計情報を取得することができ、それを元に更にパフォーマンスが良くなるようにチューニングができます。
クラシルリワードはまだデフォルトの設定で YJIT を有効化しているだけですが、今後は統計情報を元にチューニングを行いさらなるパフォーマンス向上を目指していきます。

皆さんも YJIT を有効化してみて、パフォーマンス向上を実感してみてください。