Магнит: прогноз спроса для 33 000 магазинов с релизами через Airflow

Команда MLOps FnR (Forecast & Replenishment) в Magnit Tech прогнозирует спрос и считает пополнение для магазинов сети «Магнит». Прогноз покрывает 33 000 магазинов и сотни тысяч SKU. Изначально расчёты жили в ноутбуках Zeppelin и запускались вручную, затем их перенесли в Airflow, но по-прежнему правили код на горячую: версий, отката и стейджа не было. Автор приводит пример: проверка гипотезы о сезонности продаж блинной муки перед Масленицей заняла день, а вывод изменения в прод через дата-инженеров потребовал бы три-четыре недели. Команда поставила задачу сделать релиз атомарным, версионированным, с тестами перед продом и откатом.

Система состоит из offline-пайплайна из 45 задач (проверки качества данных источников, сезонность, эластичность, фактические цены, итоговый датасет), инференса, генерации плана прогноза, интеграции и проверок качества данных, это пять репозиториев. Расчёты написаны на PySpark в Kubernetes, таблицы хранятся в Delta Lake поверх S3, оркестрирует Airflow. Для запуска задач команда написала собственный фреймворк PySpark Toolkit и оператор для Airflow. Релиз оформляется как git-тег, а данные версионируются вместе с кодом. Окружения dev, stg, uat и prod реализованы как схемы в Hive Metastore, которые ссылаются на каталог версии в S3.

Выкат устроен как Airflow DAG с формой параметров: дежурный data scientist выбирает теги репозиториев, пишет комментарий и запускает процесс. Дальше автоматика при необходимости копирует данные предыдущего тега, прогоняет offline-пайплайн, запускает тестовый инференс на стейджинговых схемах и перелинковывает таблицы в боевые схемы. Регулярные DAG подхватывают новый тег из Airflow Variable. Один релиз обрабатывает 200 ТБ данных, релизов бывает от одного до шести в месяц. Откат сводится к возврату переменной на предыдущий тег: так команда откатила релиз, когда обновление интеграции команды пополнения разошлось по времени с релизом прогноза. Ручными остаются выбор тегов, запуск и разбор ночных падений утром.