Zenn
📚

機械学習バッチジョブの実行時間を4時間から1時間にした話

2022/10/05に公開
Python
機械学習
tech

こんにちはnasaです。

最近、機械学習バッチジョブのパフォーマンス改善に取り組み実行時間を4時間30分から1時間まで改善できたのでやったことを残しておこうと思います。

やったこと

取り組みは次の４つです。それぞれ説明していこうと思います。

  • マシンのスケールアップ
  • GCSを経由してBigQueryからデータを取得する
  • 並列処理
  • dataframeのCSV出力をpickle出力にする

マシンのスケールアップ

金の弾丸と呼ばれるやつですね。

インフラコストとの相談ですが手っ取り早く高速化したいならお金の力でぶん殴るのは良い手段だと思います。

今回実行時間の短縮に取り組んだバッチジョブはEC2インスタンス上で動いていました(k8s上で動いているけどこの辺は省略します)

EC2 r5系列のマシンで実行していましたが新しい世代がいくつか追加されておりCPUのパフォーマンスが上がっていることも確認できました。

1世代上のr6系列で実行してみたところ実行時間が4時間30分から4時間になっていたのでええやん！ということでインスタンスタイプを変更しました。

EC2のインスタンス一覧を見るとr6系列よりコンピューティングに優れたインスタンスがあったのですが、ARMベースのプロセッサになっておりバッチジョブを動かすのにかなり苦労するのが目に見えていたので今回はインテルのプロセッサを使用しているものの中から最新のものを選びました。

(余談ですが、GPUを使いたい！という思いは大いにありけりですがまだ取り組めていない。。。)

GCSを経由してBigQueryからデータを取得する

今回扱っているバッチジョブは10GBほどのデータをBigQueryから取ってくるんですがこれに1時間30分かかっていました。
弊社のBigQueryはUSにあり、EC2はtokyoにあるのでデータ転送にかなりの時間がかかっているのだと思います。

実験したところTokyoリージョンのGCSからのデータ取得は大して時間がかからないこと、BigQuery USからTokyo GCSへのデータ転送は爆速なことが分かったので、この手法でデータを取得することにしました。

BigQueryのEXPORTを使って欲しいデータを圧縮しGCSに書き出します。これを取得し解凍することで1時間30分かかっていた処理が数分になりました。

EXPORT DATA OPTIONS(
      uri="gs://hoge*.csv.gz",
      format='CSV',
      compression="GZIP",
      header=TRUE,
      overwrite=TRUE
  ) AS
    select * from <欲しいデータ>

上記のスケールアップとGCS経由で実行速度は4時間半から2時間になっています。

並列処理

バッチジョブはpythonで書かれており、一部の処理はasync/awaitで並行処理していました。
async/awaitはマルチスレッドで並行処理を実現していますがpythonにはGILがありCPUバウンドな処理はCPU数に対してスケールしません。(16コアマシンで実行していたのですがこのとき悲しきhtopを見ることが出来ました。。。)

機械学習バッチジョブはCPUでゴリゴリ計算を行っているためこの部分の実行速度を改善するために複数のCPUコアを活用したいと考えました。

スレッドが駄目ならマルチプロセスにするしか無いかなーということで単一プロセスで動いていたバッチジョブをマルチプロセス化することにしました。
(余談ですが、僕はwebアプリケーションのバックエンドエンジニアをメインでやっているので時代に逆行してるなーと思いつつマルチプロセスベースの並列処理を試していました)

pythonで並列処理を書くのは初めてだったのですがこの記事のおかげで凄く助かりました

https://zenn.dev/forrep/articles/5da59f28a7732b

ここまでの改善で実行速度は4時間半から1時間20分になっています。

dataframeのCSV出力をpickle出力にする

これはあんまり話すことがないですね。dataframeの書き出しをCSVでなくpickle形式に変更しました。

df.read_csv()
df.to_csv()

↓

df.read_pickle()
df.to_pickle()

10分かかるCSV出力がいくつかありそこそこ時間かかっていたのですがpickleに変更した結果バッチジョブ全体で20分ほど速くなりました。

まとめ

この4つの改善を行い4時間半から1時間まで実行時間を短縮することが出来ました

  • マシンのスケールアップ
  • GCSを経由してBigQueryからデータを取得する
  • 並列処理
  • dataframeのCSV出力をpickle出力にする

機械学習バッチジョブの実行時間に困っている人の参考になれば嬉しいです〜

GitHubで編集を提案

Discussion