Всем привет. На связи Артём Гойлик @ArtoLord и Владислав Волох @Chillintano из команды DataSphere в Yandex Cloud. Мы создаём инфраструктуру для ML‑разработчиков. И сегодня расскажем про одну задачу, которая, как и многие другие, начиналась с болей наших пользователей.
Такая боль
Всё началось с команды, которая обратилась к нам, потому что у неё никак не получались ML‑эксперименты.
Пока её разработчики тестировали код локально, всё шло хорошо. Но как только они переключались на удалённую машину с GPU, чтобы провести тесты на реальных данных, начинались проблемы. Команда вручную переносила окружение на удалённую машину: ставила библиотеки, переносила локальный код. В этот момент что‑то обязательно ломалось, разработчики всё исправляли. А когда после долгих усилий запускали код — он не работал. Приходилось возвращаться к локальной машине и начинать всё заново.
Примерно так:
![](https://habrastorage.org/getpro/habr/upload_files/9ca/4bc/282/9ca4bc282ebd2053639acd327f85ed08.png)
Мы предложили им переносить окружение, используя Docker. Команда собрала огромный Docker со всеми нужными библиотеками и драйверами и вернулась с новыми проблемами:
Каждому разработчику нужны ещё какие‑то свои библиотеки. Команда бесконтрольно обновляет Docker, доламывая то, что работало.
Обновление базового образа усложняет ситуацию, и кто‑то не обновляет базовые образы, а создаёт новые — получается зоопарк.
Локальный код всё равно нужно копировать руками.
Никто не умеет писать Dockerfile.
Всё очень долго.
![](https://habrastorage.org/getpro/habr/upload_files/740/d62/236/740d62236eadddb30a609fa0de6b9fe1.png)
На этом этапе мы поняли, что без разработки решения не обойтись, и написали для страдающей команды простенькую утилиту.
Велосипед MVP. По нашему плану, эта утилита:
1. Локально соберёт окружение через pip freeze
и перенесёт его на удалённую машину.
2. Запакует и скопирует локальный код.
3. Развернёт данные из пунктов 1 и 2 поверх базового образа.
В теории после этого всё должно было заработать. На деле оказалось, что в окружении ML‑разработчиков очень много пакетов: они переносятся и ставятся целую вечность. Pip freeze
при этом не помогает, потому что выдаёт все версии пакетов, которые могут быть несовместимы друг с другом. Потестив это всё, мы поняли, что нам нужен новый план.
Уточнённые условия задачи
Мы поговорили с пользователями и сформулировали требования к пайплайну:
![](https://habrastorage.org/getpro/habr/upload_files/3cc/0a6/675/3cc0a6675b53b9cf711c03ddf16d8daa.png)
Чаще всего команда тестировала код локально, прежде чем запустить его на удалённой машине. Но иногда пропускала тестирование на локалке и переходила сразу к большим данным. Если что‑то ломалось, цикл повторялся. А потом всё ехало в деплой. Но деплой — уже не наша забота, от нас требовалось сделать быстрый перенос окружения. Идеальное решение:
кросс‑платформенное: не зависит от ноутбуков и дистрибутивов;
обладает низким порогом входа: не требует дополнительного обучения;
воспроизводимо между запусками.
Новое решение
Мы выбрали следующий алгоритм:
Динамическая сборка окружения в рантайме + возможность использовать кастомный базовый образ
Сначала запустим import
в коде программ и посмотрим, от каких пакетов он действительно зависит: едва ли в рантайме нужны абсолютно все пакеты, установленные в системе. А дальше — оставим возможность задавать базовый образ, где можно установить то, что не удастся перенести.
Плюсы такого подхода:
Находит только нужные модули.
Видит локальные модули.
Различает типы модулей.
В корнер‑кейсах можно использовать кастомный Docker‑образ.
Минимальный порог входа.
Есть и минусы:
Нужно запускать импорты локально.
Не всегда работает (об этих корнер‑кейсах мы расскажем дальше).
Шаг 1 из 3. Ищем модули для переноса
Первый шаг — найти только те пакеты и модули, от которых зависит рантайм, и сделать это как можно быстрее.
Для этого начнём с теории: что такое неймспейсы в Python.
Неймспейсы — это словари, в которых хранятся все объекты, которые есть в рантайме в Python:
Глобальный неймспейс — словарь, где есть все объекты текущего модуля: импортированные модули, функции, классы, глобальные переменные и прочее. Такой неймспейс есть у каждого модуля.
Локальный неймспейс — словарь, который создаётся во время исполнения функции и содержит все локальные переменные. Удаляется после окончания исполнения.
Как их обойти? Построим граф, где каждое ребро соединяет два зависящих друг от друга модуля:
import functools
a = 42
def foo():
pass
print(globals())
# {
# '__name__': '__main__’,
# ...,
# 'functools': <module 'functools'>,
# 'a': 42,
# 'foo': <function foo at ...>
# }
Определим, что такое зависимость. Зависимость — это:
модули в неймспейсе других модулей;
модули объектов в неймспейсе других модулей.
Всё бы было хорошо, но в Python всегда есть всякие подводные камни.
Корнер-кейсы
Во-первых, некоторые пакеты (например, pkg_resources
) импортируют внутри вообще все установленные модули. Их нужно пропускать, чтобы не обходить всё-всё-всё, что есть в системе.
Во-вторых, локальные неймспейсы. Как мы уже обсудили, они существуют только внутри функций. Но проблема в том, что внутри функции можно сделать import
, и мы его не увидим, потому что не запускаем функции. Это, конечно, антипаттерн, так никому не надо делать, но такие случаи встречаются. Сейчас мы их не поддерживаем, но в будущем можно будет просто составлять функции АСТ-дерева и смотреть, что импортируется на самом деле.
Мы написали такую штуку: первый раз она находит все нужные модули за полсекунды, дальше работает сильно быстрее. И находит вот для такого окружения две с чем-то тысячи модулей:
import torch
import catboost
import pandas
import numpy
timeit(lambda: get_dependencies(globals()), number=1)
# >> 0.49449617799999857
len(get_dependencies(globals()))
# >> 2293
На первый взгляд выглядит очень страшно. Но на самом деле эти модули принадлежат одним пакетам: на следующем шаге мы их объединим.
Шаг 2 из 3. Переносим модули
Итак, мы нашли тысячи нужных модулей. Теперь требуется объединить их по пакетам, классифицировать и собрать артефакты для переноса.
И снова нам нужно немного теории: напомню, что такое distribution.
Distribution — архив, где хранится информация обо всех пакетах и модулях, которые есть на диске: версии пакетов, их местоположение на диске, откуда их установили и так далее.
Информация доступна только для пакетов, установленных с помощью pip. Её предоставляет библиотека importlib-metadata
:
import importlib_metadata
dist = importlib_metadata.distribution("torch”)
dist.version
# '1.13.1’
dist.files
# [
# PackagePath('../../../bin/convert-caffe2 to-onnx’),
# PackagePath('functorch/.dylibs/libiomp5.dylib’),
# PackagePath('functorch/_C.cpython-39.so’),
# PackagePath('functorch/__init__.py’),
# ...
# ]
Именно metadata поможет нам собрать все возможные артефакты.
Теперь классифицируем модули:
![](https://habrastorage.org/getpro/habr/upload_files/f1f/148/838/f1f148838616d1497db5920069e42c52.png)
Чтобы отнести модуль к одной из трёх групп, нужно проверить наличие (или отсутствие) distribution и откуда он установлен: из Pypi или нет.
Модули, которые есть в Pypi, перевозим с помощью списка в conda.yaml вместе с версией Python. Локальные и editable‑модули — архивами через s3.
Ну а бинарные локальные модули — не перевозим, выводим на экран пользователя предупреждения: например, просим их вставить проблемные модули в базовый Docker‑образ, который позже поднимем на удалённой машине.
Корнер‑кейсы
1. Namespace modules
Namespace modules — модули, которые состоят из нескольких пакетов. Например:
files_to_dist["<>/google/api/annotations.proto"]\
.name
# 'googleapis-common-protos’
files_to_dist["<>/google/rpc/status.proto"].name
# 'pylzy'
Для работы с namespace modules нужна специальная логика: кроме того, что найти внутренние модули будет сложнее обычного, такой модуль может состоять наполовину из локального модуля, наполовину — из модуля Pypi.
2. Editable‑установка, которую conda делает плохо
Это такой вид установки пакетов в Python, когда мы не копируем все файлики в директорию с питонячьими модулями, а указываем путь на реальные файлы на диске: создаём ссылку на то, что лежит где‑то в другом месте.
Проблема в том, что conda устанавливает так вообще все свои пакеты: вместо пути в реальной файловой системе она указывает какие‑то свои метаданные, которые никуда не ведут.
И снова, чтобы распознать всё это, понадобится отдельная логика.
3. Файлы одного пакета в разных местах, как в Jupyter
На первый взгляд, нет ничего страшного в том, что файлы одного пакета могут лежать в разных местах. Но пакеты могут класть свои файлы, например, в системные директории. Такие файлы перенести не получится: мы не можем гарантировать, что они заведутся на удалённой машине.
4. Разные схемы Pypi (и то, что он иногда лежит)
Pypi, по идее, должен работать у всех. На деле оказалось, что у многих есть локальные версии Pypi, которые отличаются от стандартной. А сам стандартный Pypi постоянно лежит и возвращает странные коды.
Приготовьтесь обложиться ретраями, чтобы всё это починить.
Итак, нам удалось преодолеть все трудности и не споткнуться о корнер‑кейсы: теперь можно собрать окружение.
timeit(lambda: classify(all_deps), number=1)
# 2.6664085060001526
len(classify(all_deps))
# 27
classify(all_deps)
# [
# PypiDistribution(name='numpy', version='1.24.1’),
# PypiDistribution(name='torch', version='1.13.1’),
# …
# ]
Всё работает довольно быстро: две с половиной секунды, большую часть которых занимают походы в Pypi. Можно сделать быстрее, оптимизировав код: наш пример совсем простой, мы идём в Pypi за каждым пакетом.
Из 2000+ модулей мы получили всего 27 пакетов, объединили их и перенесли на удалённую машину.
Шаг 3 из 3. Разворачиваем окружение
Остался последний шаг: удалённо развернуть все элементы окружения, которые мы перенесли.
Это совсем просто: запускаем базовый образ, через conda устанавливаем conda yaml и распаковываем локальные модули в этом же образе (так их можно будет получить из кода).
![](https://habrastorage.org/getpro/habr/upload_files/1c8/9ef/a3f/1c89efa3f221583346079b0510797b9a.png)
Conda медленная, но именно благодаря ей мы можем вообще не думать про версию Python. Нельзя поддержать все текущие минорные версии языка — они часто бывают несовместимы друг с другом, даже если речь идёт о подверсиях минорной версии, мы такое ловили. Conda умеет самостоятельно ставить нужную версию: ту, которая была локально.
Если нужно ускорить решение, можно кешировать локальные модули по хешу в s3 между исполнениями кода, чтобы сохранять запущенный образ с уже установленным окружением и устанавливать только те модули, которых ещё нет.
Новый инструмент для распределённого исполнения кода, который автоматически собирает окружение, мы назвали Lzy и выпустили его в опенсорс. Приходите ставить звёздочки и контрибьютить.
Но если вы не хотите самостоятельно собирать окружение и переносить его либо модифицировать свой код, чтобы использовать Lzy, попробуйте DataSphere Jobs.
DataSphere Jobs
![](https://habrastorage.org/getpro/habr/upload_files/4ef/1fb/c3f/4ef1fbc3f2b6f516ce018c569b2ce77d.png)
Мы уже вкратце рассказывали, как это работает:
Вы запускаете команду у себя в консоли.
Консольная утилита собирает ваше окружение и вместе с входными данными загружает его в Object Storage.
Мы выделяем вам виртуальную машину с нужной конфигурацией. Если у вас уже есть ресурсы в DataSphere (например, s3 или датасеты), то мы можем подключить их к этой виртуальной машине.
Мы исполняем ваш код. Следите за процессом: логи попадают к вам в консоль.
После исполнения кода мы выгрузим выходные данные в Object Storage и отправим вам ссылку на скачивание.
Но давайте проверим на примере: превратим фотографии котиков в смешной стикерпак.
План такой: находим на фото кошачьи мордочки, вырезаем их, а затем обрабатываем нейросетью, чтобы перерисовать их в мультяшном стиле. Вот такие прекрасные котики:
![](https://habrastorage.org/getpro/habr/upload_files/cd4/fa6/864/cd4fa68641b07ef6574fb3bf46a69dce.png)
И вот такой код:
![](https://habrastorage.org/getpro/habr/upload_files/372/61b/966/37261b96634a086968a05d303a8a5ef1.png)
Выделение объектов работает очень быстро даже на процессоре ноутбука: исполнение этого кода заняло всего пару минут.
![](https://habrastorage.org/getpro/habr/upload_files/ea1/a02/13b/ea1a0213b93c5e16b1f889fd0f887f73.png)
И вот такие прекрасные мордочки получились:
![](https://habrastorage.org/getpro/habr/upload_files/6de/91d/26e/6de91d26e079f88a3de1dce8ba7b1d9f.png)
Угадайте, какой из этих котов принадлежит одному из авторов статьи?
Следующий шаг — обработать мордочки моделью, которая перегенерирует их в стиле рисунков. Когда мы запустили локальное окружение с этой моделью, обнаружили, что одно изображение будет генерироваться около часа. И решили ускорить этот процесс:
![](https://habrastorage.org/getpro/habr/upload_files/58d/ec2/15d/58dec215ddfa07ace1c30e4861093204.png)
Можно воспользоваться облачными ресурсами: получить доступ к машине с более мощной видеокартой. Для этого мы написали простой конфигурационный файл: описал, как запустить мой скрипт, какие у него входные и выходные данные:
![](https://habrastorage.org/getpro/habr/upload_files/cc2/17c/6c0/cc217c6c02fcb7f27e7010544a0190b3.png)
Также мы немного переписали скрипт генерации, чтобы он принимал и отдавал архивы: упростил конфигурационный файл и скачивание данных. В самом конце файла указали конфигурацию вычислительных ресурсов.
![](https://habrastorage.org/getpro/habr/upload_files/541/9a7/7d5/5419a77d5177bfc6a996ebb58d79ab82.png)
Запустили этот скрипт командой DataSphere: около пяти минут ушло на загрузку окружения и входных данных, а затем исполнение запустилось и логи исполнения начали отображаться прямо в консоли. Генерация одного изображения, как видите, стала занимать порядка полутора‑двух минут.
Через десять минут после запуска скрипта мы получили ссылку на скачивание результатов. А повторные запуски будут и того быстрее, мы закешируем окружение.
![](https://habrastorage.org/getpro/habr/upload_files/2bc/5cf/ea1/2bc5cfea1a349f21a25fbf37aef21490.png)
Как вы узнали из примера, можно заниматься локальной разработкой в любимой IDE и пользоваться облачными мощностями, как только вам перестаёт хватать ресурсов на исполнение кода.
DataSphere Jobs — удобный способ батчевого инференса для тяжёлых моделей. Поскольку запуск исполнения происходит из консоли, у конфигурационного файла простая и чёткая структура, с которой легко автоматизировать исполнение кода.
Спасибо, что дошли до конца. Приходите обсудить наши решения в комментариях.