Профилирование памяти в Python

Если вам нужно отладить ошибку в коде, это может оказаться сложной или увлекательной задачей. Выяснить, почему ваша программа/скрипт/инструмент потребляет такой объем памяти, может быть сложно. При анализе нескольких сетевых захватов я столкнулся с несколькими вариантами, работающими по-разному. Я хотел выяснить несколько различий между функциями. Здесь я делюсь несколькими тестами, которые я сделал, и рассказываю о том, как легко выполнить профилирование памяти в Python.

Написание инструмента для разбора сетевых захватов на Python с помощью Scapy Я наткнулся на несколько альтернатив. Имея дело с проблемой производительности для больших файлов, я решил сравнить несколько альтернатив. Это связано с тем, как работают rdpcap() и PcapReader().



Scapy: Способы чтения pcap
Чтение файлов сетевых захватов может быть обычной задачей для извлечения информации, образцов или другого сетевого трафика…medium.com



Некоторые различия были отмечены в предыдущем посте. В поисках разных альтернатив я нашел довольно крутую ветку в StackOverflow с несколькими вариантами, которые можно попробовать. Я выбрал Memory Profiler, в нем есть ключ, который мне нужен.

  • Легко интегрируется
  • Легко запустить
  • Сюжеты

Учебник в рамках проекта GitHub довольно хорош, и функциональность была довольно удобной для того, что мне нужно было сделать. Для меня это было так же просто, как добавить декоратор @profile в функции скриптов (здесь и здесь).

Код выглядит так:

Memory Profiler можно запустить как модуль Python или с помощью команды mprof, и он сгенерирует файл со всеми временными метками. Это довольно удобно и просто в исполнении.

Сначала нам нужно выполнить mprof run <script>, чтобы сгенерировать файл со всеми временными метками, а затем мы можем построить это с помощью mprof plot.

Запустить этот инструмент с файлом Pcap, который я скачал из Анализ трафика вредоносного ПО, было так просто:

mprof run -o sample_run_rdpcap.dat pcap_parsing_rdpcap.py --pcap pcaps/2021-06-04-part-02-after-reboot-Qakbot-with-Cobalt-Strike-and-spambot-activity.pcap --output urls_testing_rdpcap.txt
mprof: Sampling memory every 0.1s
running new process
running as a Python program...
Filename: pcap_parsing_rdpcap.py
Line #    Mem usage    Increment  Occurences   Line Contents
============================================================
    21    100.0 MiB    100.0 MiB           1   @profile
    22                                         def parse_pcap(pcap_path, urls_file):
    23    441.6 MiB    341.6 MiB           1       pcap_flow = rdpcap(pcap_path)
    24    442.7 MiB      1.2 MiB           1       sessions = pcap_flow.sessions()
    25    442.7 MiB      0.0 MiB           1       urls_output = open(urls_file, "wb")
    26    442.7 MiB      0.0 MiB        1775       for session in sessions:
    27    442.7 MiB      0.0 MiB       49194           for packet in sessions[session]:
    28    442.7 MiB      0.0 MiB       47420               try:
    29    442.7 MiB      0.0 MiB       47420                   if packet[TCP].dport == 80:
    30                                                             payload = bytes(packet[TCP].payload)
    31                                                             url = get_url_from_payload(payload)
    32                                                             urls_output.write(url.encode())
    33    442.7 MiB      0.0 MiB        1300               except Exception as e:
    34    442.7 MiB      0.0 MiB        1300                   pass
    35    442.7 MiB      0.0 MiB           1       urls_output.close()

Файл sample_run_rdpcap.dat содержит временные метки и данные о потреблении памяти, измеряемые каждую секунду. Теперь мы можем построить это и сохранить в изображение:

Теперь у нас есть хорошее представление о том, сколько памяти может занять rdpcap при чтении файла. В данном конкретном случае файл pcap имел размер 28 МБ. После загрузки в память он потреблял около 450 МБ, а скрипт выполнялся более 200 секунд.

С PcapReader() мы получаем другие результаты, с меньшим потреблением памяти и значительно меньшим временем выполнения скрипта:

Наконец, мы можем сравнить оба выполнения и отобразить их на одном графике, что является хорошим способом сравнить поведение сценариев и попытаться решить, какой из них более эффективен для той цели, которую я преследовал:

Различия между сценариями говорят сами за себя. Например, rdpcap загружает и анализирует весь файл в памяти, что дает больше информации и возможностей во время чтения сетевого захвата. Тем не менее, pcapreader читает файл последовательно, что быстрее, но менее гибко.

На графике для rdpcap видно, что потребление памяти увеличивается, а затем выходит на плато. Как только весь файл будет прочитан, именно тогда он будет обрабатывать пакеты.

Из любопытства я запустил сценарии для нескольких файлов pcap и попытался посмотреть, как их можно идентифицировать. График выглядит следующим образом:

Все образцы, проанализированные с помощью rdpcap, имеют разное потребление памяти и имеют наклон в начале, когда он загружает файл. PcapReader пиковое потребление памяти остается постоянным и составляет около 100 МБ.

В чем здесь смысл?

Тестируйте, тестируйте и тестируйте. Иногда полезно протестировать, исследовать и выяснить, как лучше всего написать фрагмент кода для инструмента или цели. Я никогда раньше не пробовал профилирование памятив Python, и это только для поверхностного изучения.

Тем не менее, иметь возможность протестировать, сравнить и решить, какой подход лучше всего подходит для решения проблемы, это ключевой момент при попытке понять, как что-то работает. Я еще не закончил со скриптом/инструментом, который я создаю, и у меня определенно есть еще несколько вопросов о профилировании в python.

Этот простой способ протестировать вещи может быть полезен кому-то другому. Есть еще несколько трюков, которые могут быть реализованы и полезны для многих других случаев использования Python.