Daskで使用メモリ容量に制限をかける

pandasで処理しきれないデータを扱う際には、pandasライクにout-of-core処理がかけるDaskがよく利用されます。今回の記事はそんなDaskをサーバ１台の環境やローカルのマシンで動かす際にメモリ使用制限をかけたい、という方向けのお話です。
めちゃくちゃニッチな話かもしれませんね。。。需要ないかも。
（少し拡張すれば分散環境にも適用できるので、そちらに興味のある方にも響くことを期待してます。）

概要

メモリに使用制限をかけるにはDaskの分散処理モード？(distributed)のLocalClusterを利用します。LocalClusterはローカルのマシンに擬似的な分散処理環境を作ります。擬似的な分散処理環境のため、schedulerとworkerが必要になります。この時、実際に処理をさばくworkerのメモリ使用量に上限を設定することで、目的のメモリ使用量に制限をかけることを実現します。

方法

ドキュメントでは下記でschedulerとworkerが立ち上がるとの記載があります。この時、Clientの引数memory_limitを指定することでメモリーの使用制限がかけられます。

from dask.distributed import Client, LocalCluster
cluster = LocalCluster()
client = Client(cluster, memory_limit='~GB')

はずでした。。。これが、環境によってできる、できないがあるようです。下記のようなエラーをはいて、workerが起動しない場合があります。詳細についてはこちらをみていただければと思います。

・・・・
AttributeError: 'NoneType' object has no attribute 'terminate'
distributed.utils - ERROR - Worker failed to start

この場合、明示的にschedulerとworkerをローカルの環境で起動してあげる必要があります。そしてworkerの起動時にmemory-limitを記載し使用メモリの上限を定めます。

$ dask-scheduler
$ dask-worker localhost:8786  --memory-limit '~GB' # this will have to be in a different terminal

そしてスクリプト側からは下記の記載で接続が可能になります。

from dask.distributed import Client
client = Client('localhost:8786')

あとはDaskのAPIにしたがって書くだけです。簡単ですね。

これで、あなたもDaskにメモリ使用上限をかけられるようになりましたね。１台のサーバのメモリを複数人でとり合う等の状況になってしまったら、是非使ってみてください!

参考