YJITの性能を最大限引き出す方法

RubyのJITコンパイラYJITを開発している弊社Shopifyでは、社内で最もトラフィックが多いストアフロントのアプリにRuby 3.3 (master) をデプロイして平均レスポンスタイムが16%高速化、社内で最も大きなアプリであるモノリスにRuby 3.2をデプロイして平均レスポンスタイムが9%高速化している。他の会社でも、YJITを本番で有効にしたら高速化したという事例をちらほら目にした。

一方で必ずしも良い報告ばかりではなく、YJITを有効化したらメモリを使い切ってしまったりだとか、遅くなったみたいな報告も目に入ることがある。こういった問題は我々も多かれ少なかれ経験しており、それぞれ適切に対処することで解決できたため、その知見を共有する。*1

メモリを使い切ってしまった時

zenn.dev

YJITを有効化すると、YJITが生成する機械語に加えて、それに関するメタデータもメモリを消費する。機械語の最大サイズは --yjit-exec-mem-size (デフォルト 64MiB) で制限されるが、メタデータは特にリミットがない。ただし、メタデータサイズは生成コードのサイズに比例する傾向にあり、かつRuby 3.2の時点ではメタデータは生成コードの3~4倍程度メモリを使うと見積っておくと良い*2。従って、デフォルトではメモリは最大で256~320MiB使われることになる*3。

ここで注意しなければならないのは、この値はあくまで各プロセスあたりのメモリ消費量であること。UnicornやPumaで複数プロセスを走らせる場合、ワーカーがforkする時点で存在しているメモリのページのうち、その後更新がされないものは複数プロセス間で共有される*4が、YJITのコードやメタデータに関しては基本的にワーカーのfork後に生成されるため、メモリの共有は期待できない。そのため、例えばUnicornのプロセスが16ある場合は、最悪の場合 16 x 256~320MiB = 4096~5120MiB 使うことを覚悟しなければならない。

--yjit-call-threshold を大きくする

一番簡単に試せるのはこれ。デフォルトでは --yjit-call-threshold は30で、つまり30回呼ばれたメソッドからコンパイルを開始するようになっているのだが、アプリの初期化のロジック次第では、この閾値では初期化時にしか使われないコードがコンパイルされてしまうということが有り得る。それらのメモリ消費は後で無駄になるので、この値を大きくするとメモリ使用量が大幅に改善することもある。

Shopifyのストアフロントではこれを30よりどれだけ大きくしてもそれほどメモリ使用量は変わらなかったが、20から30に上げた時はメモリ使用量が大幅に減った。あなたのアプリでは30よりもう少し大きくしないとその変化は訪れないかもしれない。なお、これを大きくすればするほどウォームアップが遅くなってしまうため、各ワーカープロセスが処理したリクエスト数合計などのメトリクスと見比べながら、程々の大きさに留める必要がある。

--yjit-exec-mem-size を小さくする

Ruby 3.2のYJITにはCode GCという機構が入った*5。これは生成した機械語のサイズが --yjit-exec-mem-size に達したら全てのコードを開放し、その後必要になったコードだけコンパイルして省サイズ化を目指すものだが、ついでにメタデータの方も開放されるため、これを小さくすればするほどメモリ消費量は抑えられることになる。

これを小さくしすぎるとCode GCが頻繁に行なわれるようになってしまうため、Code GCがどれくらい頻繁に行なわれているかモニタリングする必要がある。 RubyVM::YJIT.runtime_stats[:code_gc_count] の現在の値をRackミドルウェアなどで定期的に記録しておくと良い。Datadogとかだと DataDog/dd-trace-rb#2711 が使えるかもしれない。目安としては、これが0の場合は多少サイズを小さくする余地があり、1で安定する場合や1時間おきに1回走る程度が理想、それより頻繁に値が増える場合はサイズを大きくした方が良い。

肥大化したプロセスをkillする

ワーカー1つあたりに許容するメモリ使用量をあらかじめ決めておき、それを越えたプロセスを止めるようにするという手も有効である。YJITを使っているかに関わらず、本番でメモリリークが発生してしまった時の供えとしても役に立つ。我々は独自の実装を使っているが、OSSのものではunicorn-worker-killerやpuma_worker_killerなど使えばよさそう。一方で、killの頻度が高いと速度的には悪影響であるため、killの回数もモニタリングしておくのが望ましい。

ワーカーをreforkする (上級者向け)

同僚の@byrootがUnicornのフォークであるPitchforkというのを開発した。これはUnicornと比べてレガシーな依存が一つ外れているモダンなUnicornとしても使うことができるが、それに加えて "Reforking" と呼ばれている機能が追加されている。通常、UnicornやPumaのfork時にはアプリのコードがコンパイル済みでないワーカープロセスが作られるが、Reforkingというのはアプリのコードが既にコンパイルされているプロセスを後から定期的にforkし直すことでそのメモリを複数プロセス間で共有することを目指すというもの。YJITの運用でメモリの使用量を最適化しようとしたら、これが一番効果があると思われる。

デメリットとしては、スレッドを扱うgemなどがfork-safeでないことがあり、アプリが使っているコードが全てfork-safeであることをどうにか保証しておく必要がある。具体的にはgrpc.gemがこの問題を抱えており、byrootたちがGoogleとコミュニケーションを取ってこれに対応している。それから、PitchforkはUnicornのフォークであるため、Pumaのように各プロセスでマルチスレッドワーカーを立てることはできない。

一部ワーカーのみ有効化する (Ruby 3.3以降)

Ruby 3.3では新たに --yjit-disable というフラグと RubyVM::YJIT.enable というメソッドが追加される。この2つを使うと、Ruby起動時にはJITコンパイルを無効にしておき、アプリの初期化が終わった後に手動でコンパイルを開始するということができる。それが主な想定用途なのだが、byrootの発案で、モノリスではUnicorn (Pitchfork) のワーカープロセスのうち、ワーカー番号の若い一部のみ有効化することによって、メモリ使用量を抑えている。これは、Unicornがリクエストを捌く際、全てのワーカーが均等にリクエストを処理するわけではなく、キャパシティに余裕がある場合はワーカー番号が若いものにリクエストが偏るという性質を利用している。そのため、ワークロード次第ではYJITを有効化しているプロセスの割合以上のリクエストがYJIT有効のワーカーに処理されうることになる。

Ruby 3.2にはこの機能はないのでRuby 3.3を待っていただく必要があるが、我々は独自のRuby 3.2フォークを持っており、これにはこの機能がバックポートされている。

遅くなってしまった時

残念ながら、YJITを本番で有効化したら遅くなったという話もちらほら聞く。

--yjit-exec-mem-size を大きくする

YJITを有効化したら遅くなった、と聞いたときに僕が真っ先に疑うのは --yjit-exec-mem-size が小さすぎるというもの。これが小さいとCode GCが頻繁に走りすぎて遅いというリスクが高くなる。その症状になっているかを確認するには、RubyVM::YJIT.runtime_stats[:code_gc_count] をモニタリングし、この値が頻繁に増加していないかを確認すると良い。これが0か1で安定するところまで上げれば、速度的には影響がない状態にできる。

弊社ストアフロントでは --yjit-exec-mem-size=64 を使っていて、これは十分すぎるサイズなのだが、一方弊社モノリスでは --yjit-exec-mem-size=256 を使っていて、これを大きくしてきたときに大幅に速度が改善された。これに関連して、Ruby 3.3 (master)ではこのオプションのデフォルトが128に変更されている。

ワーカープロセスをなるべく長く走らせる

僕が次に疑うのは、ワーカープロセスが定期的にkillされているような環境で、そのkillがあまりにも頻繁すぎるというもの。プロセスが長く走ればコンパイル済みのコードを再利用する機会が増え、速度的には良い影響が期待できる*6が、頻繁にkillされるとコンパイルのオーバーヘッドがかかり続けるということが有り得る。各ワーカープロセスが過去に処理したリクエスト数合計などをモニタリングしておき、その値が大きくなる前にプロセスがkillされてしまっている場合は、killの閾値の見直したり、OOMの影響を確認して割り当てるメモリを増やしたりする必要がある。

--yjit-call-threshold を調整する

関連した問題として、--yjit-call-thresholdはワーカーのリクエスト処理数に応じて調整する必要がある。--yjit-call-threshold が小さすぎると、起動時に多くのコードがコンパイルされ、ウォームアップのオーバーヘッドが一気にかかりがちになる。その一方で --yjit-call-threshold を例えば1000まで上げた時、ワーカーが頻繁に再起動されていてプロセスあたり1000リクエスト足らずでkillされてしまっている場合、1000回目のリクエストで初めてコンパイルされるパスが有り得ることになる。その場合は閾値を100くらいまで下げると、より早く、一方で早すぎずウォームアップが行なわれ、速度が改善したりする。

ratio_in_yjit を確認する

これはRuby 3.2では少し手間がかかるのだが、Rubyのconfigure時に --enable-yjit=stats をつけてビルドしておき、Rubyの起動時に --yjit-stats をつけ、RubyVM::YJIT.runtime_stats[:ratio_in_yjit] を確認すると、実行されているVM命令のうち何%がYJITで実行されているか確認することができる。ワーカープロセスがピーク性能に達した状態でこれが90%くらいあれば性能が改善していることが期待できるが、これが例えば80%を下回るなどしている場合、アプリのコードかYJITの実装のどちらかに問題があることになる。例えばTracePointがどこかで有効になっていると、ものによってはそれがVM命令を全てtrace命令に書き換えることがあり、その場合YJITはコンパイルを諦めてしまう。いずれにせよ、例えばワーカープロセスが10000リクエスト処理した後もこの値が小さい場合は、--yjit-stats を使用した状態での RubyVM::YJIT.runtime_stats の中身を丸ごとYJITチームに共有*7していただけると、よしなに対応できると思う。

Ruby 3.3ではデフォルトのビルドで ratio_in_yjit が確認できるようになっている。YJITが使える全てのRubyのビルドで、起動時に --yjit-stats をつけておけば、RubyVM::YJIT.runtime_stats[:ratio_in_yjit] が利用できる。一方で、Ruby 3.3には ratio_in_yjit を改善する様々な実装が入っており、以前は92%だったのが現在では97%まで改善していたりするので、そもそもこれを確認しなくても速くなるようになっているかもしれない。

まとめ

簡単にできるアクションのまとめとしては、RubyVM::YJIT.runtime_stats[:code_gc_count] と各ワーカープロセスが処理したリクエスト数をモニタリングし、それらや速度、メモリ使用量に応じて --yjit-call-threshold や --yjit-exec-mem-size を調整したり、割り当てられているメモリのサイズやワーカーのkillの設定などを見直していただくのが良いと思われる。

これを読んで「YJITを使うのは面倒くさい」と思った人もいるかもしれないが、これらの知見を可能な限りデフォルトのパラメータにフィードバックしており、基本的にはデフォルトの設定で有効化しただけで特に苦労なく速くなる状態を目指して開発されている。また、より多くのVM命令の対応が日々コミットされ続けているので、Ruby 3.2で遅くても、Ruby 3.3にしたら速くなった、ということも有り得ると思う。

参考資料

1月に書いた(のを6月にやっと公開した)記事なので若干情報は古いが、会社のブログに書いた以下の記事も参考になるかもしれない。

*1:余談だが、こういった話のプロポーザルをRubyKaigi 2023に出していたのだが、RJITの方の話が勝ったため話しそびれていた。

*2:なお、Ruby 3.3ではメタデータのサイズがより小さくなっており、生成コードの2~3倍程度になる。

*3:メモリは仮想ページごとに必要に応じて割り当てが行なわれるため、アプリのサイズ次第ではメモリ消費量はもっと小さくなる。

*4:Copy on Writeの話をしている

*5:僕が書いた

*6:一方で長く走らせすぎるとメモリが断片化していき若干性能が減衰する問題もあるが、ここで話している問題とは別なので置いておく

*7:https://bugs.ruby-lang.org/ にチケットを起票するか、https://github.com/Shopify/yjit にissueを書くか、gistとかに貼ってX (Twitter) で僕にリプライなど