Самый быстрый фреймворк на Диком Западе: ускоряем Django-rest-framework вместе с Rust

от автора

Всем привет! Меня зовут Роман Кабаев, я инженер в отделе разработки инструментов тестирования компании YADRO. Вместе с коллегами мы разрабатываем собственную тест-менеджмент систему с открытым исходным кодом TestY.

На этапе запуска TestY в качестве фреймворка для разработки мы выбрали Django, так как он позволяет в максимально короткие сроки реализовать MVP. Однако развивать такой продукт — добавлять фичи, наращивать число пользователей и объем хранимых данных в системе — бывает сложно.

Мы действительно быстро запустили MVP, перевезли данные из TestRail с помощью плагинов, и команды тестирования YADRO уже более года пользуются системой. Но есть одно «но»: пользовательские сценарии разных команд сильно отличаются. Так, добавление в систему более полумиллиона тестов привело к просадке скорости работы определенных эндпоинтов, завязанных на древовидных структурах. 

Спойлер: камнем преткновения для нас стали CPU-bound задачи с большим количеством данных, о том, как я это выяснил, расскажу ниже. Изучив, как можно ускорить выполнение таких задач в Python, я протестировал несколько решений и нашел оптимальное. Если вы разрабатываете веб-приложение на Django или Python и так же, как я, хотите ускорить работу сервиса, читайте эту статью.

В качестве альтернативы Rust можно посмотреть и на Cython, и на чистый C. Возможно, стоит подумать о микросервисной архитектуре и отдавать тяжелые с точки зрения CPU задачи микросервису, написанному, например, на Go. Я же последовал модным веяниям и решил изучить, как написать Python-пакет с использованием Rust.

Отправная точка — 10,93 секунды

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

Чтобы построить красивое дерево выбора, все тестовые сьюты для выбора и кейсы внутри них отдаются эндпоинтом сразу. А значит, если у вас 10 000 кейсов в одной из сьют, они подгружаются одновременно. Отсюда и такая огромная просадка по скорости отдачи данных.

Проблему можно исправить с помощью хорошего UX-дизайна и небольшого изменения архитектуры, но так как TestY — не единственный наш проект, на это не было ресурсов. Поэтому, чтобы сделать быстрый фикс без изменений UI, нам нужно было ускорить отдачу древовидных структур именно с серверной части приложения.

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

  • 16 419 тестовых кейсов.

  • 445 тестовых сьют, разных степеней вложенности.

  • Размер payload (полезной нагрузки) с 2 МБ данных.

  • Для работы с базой данных используем Django ORM.

  • Запускаем это все в Docker-контейнере на MacBook Pro (M2) 16 ГБ ОЗУ.

  • Версия Python 3.9.13, операционная система контейнера Alpine.

Изначально у нас есть код с использованием Django и Django-rest-framework (DRF) следующего вида:

@cases_search_schema @action(methods=[_GET], url_path='search', url_name='search', detail=False) def cases_search(self, request):     cases = self.get_queryset()     cases = self.filter_queryset(cases)     suites_selector = TestSuiteSelector()     suites = TestSuiteSelector().suites_by_ids_list(         cases.values_list('pk', flat=True),         field_name='test_cases__pk',     )     suites_depth = suites_selector.get_max_level()     suites = (         suites         .get_ancestors(include_self=True)         .filter(parent=None)         .prefetch_related(             *suites_selector.suites_tree_prefetch_children(suites_depth),             *form_tree_prefetch_objects(                 nested_prefetch_field='child_test_suites',                 prefetch_field='test_cases',                 tree_depth=suites_depth,                 queryset=cases,             ),         ).annotate(**suites_selector.cases_count_annotation())     )     return Response(         data=TestSuiteTreeCasesSerializer(             suites,             context=self.get_serializer_context(),             many=True,         ).data,     )

Производительность не очень впечатляющая, эндпоинт отдал данные за 10,93 секунды. Попробуем ускорить решение, чтобы добиться более высокой скорости работы. 

Поиск узкого места

Чтобы решить какую-либо задачу, сначала нужно ее проанализировать и понять корень проблемы. Так как мы говорим про веб-приложение, первое, с чего я начал, — это посмотрел, как быстро работает запрос в базу данных с помощью django-debug-toolbar. Он съедает часть производительности системы, поэтому мы включаем его только единожды, чтобы посмотреть, как много времени у нас занимают запросы в базу данных. Все запросы в БД занимают 371 миллисекунду.

Упрощаем DRF-сериализаторы

Сериализаторы DRF очень тяжеловесны и содержат много скрытой логики и CPU-bound нагрузки. Попробуем их облегчить:

  • избавимся от ModelSerializer в пользу Serializer,

  • сделаем все поля read_only.

Получаем сериализаторы следующего вида:

from rest_framework import serializers   class TestSuiteTestSerializer(serializers.Serializer):     id = serializers.IntegerField(read_only=True)     name = serializers.CharField(read_only=True)     parent_id = serializers.IntegerField(read_only=True)     title = serializers.CharField(read_only=True)     test_cases = GenericIdNameSerializer(many=True, read_only=True)     parent = GenericIdNameSerializer(read_only=True)     children = serializers.SerializerMethodField()      def get_children(self, value):         return type(self)(             value.child_test_suites.all(),             many=True,             context=self.context,         ).data

Получаем результат в 9,12 секунды. Такой медленный ответ пользователя не устроит, продолжим поиски.

Используем стороннюю библиотеку для перевода Python-объектов в JSON

Так как конвертирование Python-словаря в JSON — не самая «легкая» cpu-bound операция из-за рекурсивного обхода вложенных Python-объектов и манипуляций со строками, ищем альтернативу.

Вспоминаем про библиотеку Pydantic, которая часто используется для сериализации объектов в связке с FastAPI. Нам не нужен весь функционал Pydantic, интересует только сериализация в JSON, а для этого Pydantic использует Orjson — библиотеку, написанную на Rust как раз для ускорения сериализации в JSON.

Используем Orjson для сериализации в JSON и получаем результат в 8,32 секунды. Пользователь в ярости от такой нерасторопности системы: возможно, уже обновил страницу кулаком и не один раз повысил нагрузку на сервер. Результат — плохо всем.

Как и большинство веб-проектов на Django, для запуска в production-окружении мы используем application server gunicorn. Так как мы не можем обрабатывать один запрос более двух минут, чтобы не сломать систему и тем более виртуальную машину, на которой у нас развернуто приложение, мы выставляем timeout. Если запрос исполняется больше двух минут, gunicorn worker ликвидируется и выдает неприятную ошибку в системе мониторинга Sentry.

Стадия принятия

Каким бы прекрасным ни был Python, когда разработчики сталкиваются с ограничением производительности, то обращаются к более высокопроизводительным языкам, а затем делают биндинги в Python. Воспользуемся этим. Orjson использует Rust, и мы решили написать на этом языке свой пакет.

Так как сборка древовидных структур происходит по ID его родителя, вся нагрузка ложится на Python. Но мы можем получить наборы данных в плоской структуре — в виде списков словарей, а затем передать их в пакет, написанный на Rust, и обработать данные там.

Что мы делали раньше:

  1. Сформировали запрос в базу данных с помощью django-orm.

  2. Отдали набор данных QuerySet в DRF сериализатор.

  3. Достали из него Python-словарь.

  4. Конвертировали Python словарь в JSON (данный шаг скрыт под капотом у класса Response из библиотеки DRF).

  5. Отдали ответ пользователю.

Как это выглядит теперь:

  1. Получаем QuerySet всех тестовых кейсов, используем метод values (). Он вернет QuerySet, содержащий словари. Обернем все это в список — так мы получим список словарей, с которым проще всего будет работать в Rust. Таким же образом получаем все тестовые кейсы, в values указываем все нужные нам поля, чтобы не тянуть лишние данные из базы.

  2. Ищем тестовые сьюты, к которым относятся эти тестовые кейсы, и получаем их предков (родители родителя и так до 0 уровня). Получаем плоский набор данных по принципу, описанному выше.

  3. Передаем наборы данных в функцию из нашей Rust-библиотеки и производим необходимую обработку.

  4. Возвращаем список словарей. Словари, в свою очередь, представляют иерархическое дерево.

  5. Сериализуем полученный список в JSON с помощью Orjson.

  6. Отдаем ответ пользователю.

Для написания пакета нам нужен сам Rust, с инструкцией по его установке можно ознакомится на официальном сайте. Также понадобится Python-пакет maturin.

Maturin создает проект с необходимым содержимым для того, чтобы мы начали писать свой Python-пакет на Rust:

  • Cargo.toml — конфигурационный файл пакетного  менеджера Rust Cargo 

  • pyproject.toml — конфигурация нашего Python-пакета,

  • lib.rs — точка входа Rust-кода.

Сначала напишем функцию, которую будет вызывать Python-код. Это функция с декоратором #[pyfunction]:

// src/serialization.rs #[pyfunction] pub fn serialize_tree(    py: Python,    data_set_object: TreeObject,    prefetch_objects: Vec<Prefetch>,    is_tree: bool, ) -> Vec<PythonInstance> {    let mut dict_map = BTreeMap::new();    let mut root_refs: Vec<PythonInstanceRef> = Vec::new();    for mut dict in data_set_object.instances {        for prefetch_object in &prefetch_objects {            dict.related_objects.insert(prefetch_object.group_key.clone(), vec![]);        }        dict_map.insert(dict.id, Rc::new(RefCell::new(dict)));    }    add_related_objects(py, prefetch_objects, &mut dict_map);    if !is_tree {        return dict_map.values().map(|obj| obj.borrow().clone()).collect();    }    for dict in dict_map.values() {        let borrowed_dict = dict.borrow();          if let Some(parent_id) = borrowed_dict.parent_id {            if let Some(parent) = dict_map.get(&parent_id) {                map_parent(parent, parent_id, dict, borrowed_dict, py)            }        } else {            root_refs.push(Rc::clone(dict))        }    }    root_refs.iter().map(|root| root.borrow().clone()).collect() }   fn map_parent(    parent: &Rc<RefCell<PythonInstance>>,    parent_id: u32,    dict: &Rc<RefCell<PythonInstance>>,    borrowed_dict: Ref<PythonInstance>,    py: Python, ) {    let py_dict = borrowed_dict        .dict        .downcast::<PyDict>(py)        .expect("Could not downcast PyObject to PyDict");    let mut borrowed_parent = parent.borrow_mut();    borrowed_parent.children.push(Rc::clone(dict));    let name: String = borrowed_parent        .dict        .downcast::<PyDict>(py)        .expect("Could not cast PyAny to PyDict")        .get_item("name")        .expect("name was not found in PyDict")        .expect("unwrapping")        .extract()        .expect("Could not extract String from PyObject");    py_dict.set_item(        "parent",        Parent {            id: parent_id,            name,        }.into_py(py),    ).expect("Could not set item in dict"); }   fn add_related_objects(    py: Python,    prefetch_objects: Vec<Prefetch>,    dict_map: &mut BTreeMap<u32, PythonInstanceRef>, ) {    for prefetch_object in prefetch_objects {        for instance in prefetch_object.instances {            let related_dict = instance                .downcast::<PyDict>(py)                .expect("Could not cast PyObject to PyDict");            let fk: u32 = related_dict                .get_item(&prefetch_object.fk_key)                .expect("Could not find fk by provided key from dict")                .expect("fk was not found")                .extract()                .expect("Could not convert PyObject to u32");            let relation_map = &mut dict_map.get_mut(&fk).expect("Could not get map").borrow_mut().related_objects;            if let Some(related_objects) = relation_map.get_mut(&prefetch_object.group_key) {                related_objects.push(instance);            }        }    } }

Регистрируем ее в наш модуль в файле lib.rs:

// src/lib.rs  use pyo3::prelude::*;  mod serialization; mod sql; #[pymodule] fn rusty(_py: Python, m: &PyModule) -> PyResult<()> {     m.add_function(wrap_pyfunction!(serialization::serialize_tree, m)?)?;     m.add_function(wrap_pyfunction!(sql::cases_search, m)?)?;     Ok(()) }

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

Поговорим немного о преобразовании Python-объектов в Rust-объекты. Нам нужна какая-то структура, которая помогла бы преобразовывать динамически типизированные Python-объекты в статически типизированные Rust-объекты, поэтому напишем следующие структуры:

// src/serialization.rs  use std::{cell::RefCell, collections::BTreeMap, rc::Rc};  use pyo3::{pyfunction, types::PyDict, FromPyObject, IntoPy, PyAny, PyObject, PyResult, Python};  type PythonInstanceRef = Rc<RefCell<PythonInstance>>;  #[derive(Clone, Debug)] pub struct PythonInstance {     id: u32,     parent_id: Option<u32>,     dict: PyObject,     children: Vec<PythonInstanceRef>,     related_objects: BTreeMap<String, Vec<PyObject>>, }  #[derive(Clone, Debug, FromPyObject)] pub struct Prefetch {     group_key: String,     fk_key: String,     instances: Vec<PyObject>, }  struct Parent {     id: u32,     name: String, }  pub struct TreeObject {     instances: Vec<PythonInstance>, }

Built-in типы PyO3 может преобразовывать и сам. Например, list[int]) преобразуется в Vec без явного вмешательства с нашей стороны. Однако нам нужно передать определенную метаинформацию, чтобы потенциально переиспользовать функцию из Python-Rust-пакета для обработки других структур, отличных от тестовых сьют с кейсами. А еще мы бы хотели работать с объектами Python как с не строго типизированными, потому что мы можем получить разное количество ключей в словарях или передать словари с другими полями. Поэтому у нас есть Python-объекты (дата-классы), которые несут определенную метаинформацию.

Объекты, которые мы хотим обработать, оставляем в формате PyObject. В процессе формирования иерархической структуры мы достаем из них только ключи, необходимые нам для формирования древовидных структур. Чтобы добиться этого, мы реализуем Rust traits из PyO3, такие как FromPyObject и IntoPy.

// src/serialization.rs  impl<'source> FromPyObject<'source> for TreeObject {     fn extract(obj: &'source PyAny) -> PyResult<TreeObject> {         let parent_key: Option<String> = obj             .getattr("parent_key")             .expect("Could not get parent key from python object")             .extract()?;         let pk_key: String = obj             .getattr("pk_key")             .expect("Could not_get pk key from python object")             .extract()?;         let mut instances = Vec::new();         let py_instances: Vec<PyObject> = obj             .getattr("instances")             .unwrap()             .extract()?;         for py_instance in py_instances {             let py_dict = py_instance.downcast::<PyDict>(obj.py())?;             let id = py_dict                 .get_item(&pk_key)                 .unwrap()                 .unwrap()                 .extract()?;             let mut parent_id = None;             if let Some(parent_id_from_py) = py_dict                 .get_item(&parent_key)                 .expect("Could not get element from python dict")             {                 parent_id = parent_id_from_py                     .extract()                     .expect("Could not parse parent id value as u32");             }             instances.push(PythonInstance {                 id,                 dict: py_instance,                 parent_id,                 related_objects: BTreeMap::new(),                 children: Vec::new(),             });         }         Ok(TreeObject { instances })     } }  impl IntoPy<PyObject> for Parent {     fn into_py(self, py: Python) -> PyObject {         let dict = PyDict::new(py);         dict.set_item("name", self.name)             .expect("Could not set name for parent");         dict.set_item("id", self.id)             .expect("Could not set id for parent");         dict.into()     } } 

FromPyObject конвертирует Python-объект в Rust-структуру, IntoPy — это обратная операция. Чтобы с модулем было удобнее работать из основного приложения, нужно добавить тайп-хинты и объекты группировки. Последние подскажут пользователю, как работать с нашим модулем.

Так как мы хотим использовать и Python, и Rust в одном проекте, сделаем папку pysrc. Выкладка нашего модуля будет выглядеть так:

В pyproject.toml, который сгенерировал пакет maturin, обозначим, что Python-исходник — это пункт python-source:

# Cargo.toml [build-system] requires = ["maturin>=1.0,<2.0"] build-backend = "maturin"  [tool.maturin] # "extension-module" tells pyo3 we want to build an extension module (skips linking against libpython.so) features = ["pyo3/extension-module"] python-source = "pysrc"  [project] name = "TestY-rust" requires-python = ">=3.8" classifiers = [     "Programming Language :: Rust",     "Programming Language :: Python :: Implementation :: CPython",     "Programming Language :: Python :: Implementation :: PyPy", ] dynamic = ["version"]

Важный момент: в init.py файле в pysrc/rusty нужно будет импортировать все из пакета rusty (тот пакет, который мы обозначили в lib.rs), иначе в пакете мы не найдем наши функции.

# pysrc/rusty/__init__.py from typing import Any  from .constants import * from .rusty import * from .types import *  __all__ = (     "CASES_SEARCH_SUITE_FIELDS",     "CASES_SEARCH_CASE_FIELDS",     "CaseSearchQueryParams",     "Prefetch",     "DataSetObject",     "serialize_tree",     "cases_search", )

В pysrc/rusty/types сделаем пару простых дата-классов, которые будут содержать необходимую для форматирования данных метаинформацию. Prefetch — это объекты, которые мы будем группировать на другие объекты, в нашем случае тестовые кейсы на сьюты. Нам нужен ключ, по которому мы будем складывать наши тестовые кейсы в сьюты. fk_key — это имя ключа родительского элемента, по нему будем искать нужный объект, instances — это наши тестовые кейсы.

DataSetObject — это наши cьюты, pk_key — это ключ первичного ключа, parent_key — ключ родительского элемента, instances — это сами сьюты:

# pysrc/rusty/types.py from dataclasses import dataclass from typing import Any, Optional  @dataclass class Prefetch:     group_key: str     fk_key: str     instances: list[dict[str, Any]]   @dataclass class DataSetObject:     parent_key: Optional[str]     pk_key: str     instances: list[dict[str, Any]]

И чтобы однозначно обозначить, что наши функции принимают, и чтобы IDE подсказывала нам, что эти функции существуют, сделаем .pyi-файл и положим его рядом с кодом на Python в pysrc/rusty:

# pysrc/rusty/__init__.pyi  def serialize_tree(     data_set_object: DataSetObject,     prefetch_objects: list[Prefetch],     is_tree: bool, ) -> list[dict[str, Any]]:     """     Retrieve python objects retrieved from queries and build dict for json.      Args:         data_set_object: description type of provided data         prefetch_objects: description type of data that should be "prefetched" on main object         is_tree: boolean flag to build the main data set as a tree.      Returns:         list of dictionaries ready to be given to the final user.     """

Так мы приходим к следующему коду в нашем view поиска тестовых кейсов:

@cases_search_schema @action(methods=[_GET], url_path='search', url_name='search', detail=False) def cases_search(self, request):     cases = self.get_queryset()     cases = self.filter_queryset(cases)     suites = list(         TestSuite.objects.filter(             id__in=cases.values_list('suite_id', flat=True),         ).get_ancestors(             include_self=True,         ).annotate(             title=F(_NAME),         ).values(_ID, _NAME, 'title', 'parent_id').order_by('name'),     )     data = rusty.serialize_tree(         rusty.DataSetObject('parent_id', _ID, instances=suites),         prefetch_objects=[             rusty.Prefetch(                 'test_cases',                 'suite_id',                 instances=list(cases.values(_ID, _NAME, 'suite_id', 'labels').order_by('name')),             ),         ],         is_tree=True,     )     data = orjson.dumps(data)     return HttpResponse(content=data, content_type='application/json')

Результатом работы становится ответ сервера в 300 мс — это отлично, если сравнивать с исходным результатом в 10,93 секунды

Но мы пойдем дальше и посмотрим, что будет, если сделать этот же запрос на самый нагруженный проект в системе. В проекте 221 203 тестовых кейса и 101 тестовый сьют. В результате получим ответ от сервера за 722 мс с пейлоадом в 22.3 МБ. Это более чем удовлетворительный результат, подходящий пользователям.

Раз мы удовлетворены перформансом эндпоинта, поговорим о подводных камнях при установке пакета.

Установка Rust/Python-модуля

Так как мы хотим, чтобы наш пакет работал на большом количестве конфигураций систем, было бы хорошо включить в наш Dockerfile Rust и собирать пакет на лету. Но это увеличивает и так не маленькое время установки зависимостей проекта и запуска контейнеров. Примерно на 4−8 минут, так как сначала нужно установить Rust, а затем уже скомпилировать нашу зависимость. Как разработчиков системы это нас очень печалит.

Вспоминаем, что есть предсобранные файлы wheels или же .whl Python-пакеты, но есть одно «но». Даже у разработчиков в нашей команде разные конфигурации систем MacOS (arm/x86), Debian (x86) и Windows (x86), то есть мы не можем собрать .whl-файл для какой-то конкретной системы, даже для нашего внутреннего использования. Также надо помнить, что TestY — это проект с открытым исходным кодом, а значит, он должен собираться на большом количестве разных систем.

Что мы можем сделать, чтобы уменьшить шансы ошибки сборки проекта? Собрать .whl-файлы для самых часто встречающихся конфигураций машин. Поэтому мы решили автоматизировать сборку .whl-файлов и их загрузку на pypi. И когда пользователь будет делать pip-install, pypi сам будет определять необходимую пользователю зависимость. Для этого мы воспользуемся GitHub Workflows, в частности maturin-action. Как только мы запушим новый tag в наш репозиторий на GitHub, maturin-action соберет необходимые нам .whl-файлы и зальет их на pypi, если не найдет идентичные .whl-файлы.

Пример кода пайплана для сборки .whl-файлов для различных дистрибутивов Linux на разных архитектурах:

# .github/workflows/CI.yml  name: CI  on:   push:     branches:       - main       - master     tags:       - '*'   pull_request:   workflow_dispatch:  permissions:   contents: read  jobs:   linux:     runs-on: ubuntu-latest     strategy:       matrix:         arch: [ i686, x86_64, aarch64 ]         image: [ manylinux2014, musllinux_1_1 ]      steps:       - uses: actions/checkout@v4       - uses: actions/setup-python@v5         with:           python-version: '3.10'       - name: Build wheels - manylinux         uses: PyO3/maturin-action@v1         with:           target: ${{ matrix.arch }}           args: --release --out dist -i python3.9           sccache: 'true'           manylinux: ${{ matrix.image }}       - name: Upload wheels         uses: actions/upload-artifact@v4         with:           name: wheels-linux-${{ matrix.image }}-${{ matrix.arch }}           path: dist

Часть пайплайна, которая загружает собранные .whl-файлы в pypi, опущена, также как и сборка под другие операционные системы. Теперь наш пакет полностью готов к использованию.

Выводы по задаче и альтернативные решения

В ходе разработки я понял, что интеграция Rust-модуля в Python — не самая тривиальная задача. Хоть на эту тему есть много гайдов от других инженеров, всегда возникают подводные камни, которые никто не описывал. Как по мне, документация Rust-пакетов, в частности PyO3, могла бы быть лучше, но все знают, что наличие документации в целом — уже чудо.

Такая оптимизация — это обходной путь для нехватки человеческого ресурса и времени на разработку. Тем не менее, мы всеми силами стараемся сделать систему лучше, не отказываемся от Django и не хотим переписывать всю систему на Rust, так как это сильно замедлит скорость разработки. Мы используем все возможные средства для улучшения пользовательского опыта, не замыкаемся в стандартных решениях для Django или ограничениях языка.

Если тоже хотите поучаствовать в улучшении тест-менеджмент системы TestY, пишите на почту testy@yadro.com. А скачать репозиторий с версией TestY 1.3, о которой мы недавно писали, можно по ссылке.


ссылка на оригинал статьи https://habr.com/ru/articles/830024/


Комментарии

Добавить комментарий

Ваш адрес email не будет опубликован. Обязательные поля помечены *