Skip to content

カスタムスクリプトでのmemory_profilerの使用

カスタムスクリプトを使用すると、Treasure ワークフロー内からコンテナ化されたPythonスクリプトを実行でき、カスタムロジックの柔軟性が向上します。

Pythonスクリプトが実行されるコンテナは、アカウント設定に応じてリソース制約を受けます。リソース制約の1つは、カスタムスクリプトが使用できるメモリ量です。カスタムスクリプトが利用可能なメモリを超えた場合、スクリプトは以下のようなエラーメッセージで失敗します:

Error messages from CommandExecutor: OutOfMemoryError: Container killed due to memory usage

メモリの問題でスクリプトが失敗した場合、メモリ使用量を削減するためにスクリプトをトラブルシューティングする必要があります。

memory profilerは、スクリプトの各部分によるメモリ消費量を確認できるPythonが提供するオプションの1つです。これは、プロセスのメモリ消費量とPythonプログラムの行ごとのメモリ消費量分析を監視するためのPythonモジュールです。以下は、memory profilerを使用するカスタムスクリプトの例です。

examples.dig
+task1:
  py>: my_script.call_mem
  docker:
    image: "<image_name>:<version>"
examples.py
import digdag
import sys,os

os.system(f"{sys.executable} -m pip install memory_profiler")

from memory_profiler import profile

def call_mem():
mem_test_1()
mem_test_2()

@profile
def mem_test_1():
a = []
for i in range(10000):
    a.append(i * i)
return a

@profile
def mem_test_2():
a = [1] * (10 ** 6)
b = [2] * (2 * 10 ** 7)
del b
return a

サンプル出力:

Collecting memory_profiler
  Downloading https://files.pythonhosted.org/packages/8f/fd/d92b3295657f8837e0177e7b48b32d6651436f0293af42b76d134c3bb489/memory_profiler-0.58.0.tar.gz
Collecting psutil (from memory_profiler)
  Downloading https://files.pythonhosted.org/packages/84/da/f7efdcf012b51506938553dbe302aecc22f3f43abd5cffa8320e8e0588d5/psutil-5.8.0-cp37-cp37m-manylinux2010_x86_64.whl (296kB)
Building wheels for collected packages: memory-profiler
  Building wheel for memory-profiler (setup.py): started
  Building wheel for memory-profiler (setup.py): finished with status 'done'
  Stored in directory: /root/.cache/pip/wheels/02/e4/0b/aaab481fc5dd2a4ea59e78bc7231bb6aae7635ca7ee79f8ae5
Successfully built memory-profiler
Installing collected packages: psutil, memory-profiler
Successfully installed memory-profiler-0.58.0 psutil-5.8.0
WARNING: You are using pip version 19.1.1, however version 20.3.3 is available.
You should consider upgrading via the 'pip install --upgrade pip' command.

Filename: /my_script.py
Line #    Mem usage    Increment  Occurences   Line Contents
============================================================
    12     17.3 MiB     17.3 MiB           1   @profile
    13                                         def mem_test_1():
    14     17.3 MiB      0.0 MiB           1       a = []
    15     17.3 MiB      0.0 MiB       10001       for i in range(10000):
    16     17.3 MiB      0.0 MiB       10000           a.append(i * i)
    17     17.3 MiB      0.0 MiB           1       return a

Filename: /my_script.py
Line #    Mem usage    Increment  Occurences   Line Contents
============================================================
    19     17.3 MiB     17.3 MiB           1   @profile
    20                                         def mem_test_2():
    21     24.8 MiB      7.5 MiB           1       a = [1] * (10 ** 6)
    22    177.4 MiB    152.6 MiB           1       b = [2] * (2 * 10 ** 7)
    23     24.9 MiB   -152.5 MiB           1       del b
    24     24.9 MiB      0.0 MiB           1       return a