Как DeepCure распространяет модели молекулярных свойств

Для компании, занимающейся поиском лекарств с помощью искусственного интеллекта, данные, предоставляемые нашими моделями машинного обучения, имеют решающее значение. Таким образом, службы, которые мы используем для развертывания и запроса этих моделей, должны быть надежными, производительными и масштабируемыми. В первые дни существования компании мы разработали решение, которое проверяло достаточное количество этих требований, чтобы обойтись. Однако по мере появления новых моделей и увеличения нагрузки на запросы требовались новые решения. Этот пост в блоге будет охватывать только некоторые аспекты этой новой архитектуры. В частности, в этой статье основное внимание будет уделено структуре данных, используемой для распределения обязанностей по прогнозированию для десятков моделей в кластере контейнеров.

Проблема, которую нужно решить

Прежде чем описывать то, что мы построили, нужно понять проблему. Важнейшим компонентом нашего процесса разработки лекарств является прогнозирование свойств молекул-кандидатов ADME-Tox. ADME-Tox — это аббревиатура, обозначающая различные аспекты взаимодействия соединений в организме человека.

Абсорбция — насколько хорошо препарат всасывается в ткани.
Распределение — насколько хорошо лекарство перемещается по организму.
Метаболизм — как и где расщепляется лекарство
Экскреция — сколько времени требуется, чтобы покинуть тело
Токс — Имеет ли препарат какие-либо токсические эффекты?

Было подсчитано, что 40-50% лекарств терпят неудачу из-за низкой эффективности (биодоступность, нежелательная метаболическая стабильность и т.д.)¹. Также было подсчитано, что 30% лекарств терпят неудачу из-за их токсичности¹. В DeepCure мы обучили более 30 моделей прогнозированию свойств по этим измерениям. Это позволяет нам быстро выявлять и отсеивать кандидатов с нежелательными свойствами.

Эти модели генерируют прогнозы в двух разных контекстах: в пакетном режиме и по требованию. Мы используем Apache Spark для пакетного прогнозирования, чтобы запустить произвольное количество моделей для миллионов молекул. Готового решения для предсказания по требованию для одной или нескольких десятков молекул не существовало. При обдумывании того, как обслуживать эти десятки моделей, у нас был список требований, которыми мы руководствовались:

Все 30+ моделей необходимо запрашивать в одном запросе.
Время отклика должно быть менее 25 секунд.
Новые модели или версии существующих моделей необходимо подбирать без изменения кода.
Автоматическое распределение нагрузки.

Как (не)распространять модели

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

Единая служба

Это самое простое решение для рассмотрения: один сервис (или однородный кластер), который будет обслуживать прогнозы из любого набора моделей. Если развернуто всего несколько моделей, это решение может сработать. Однако модели должны быть загружены в память при поступлении запроса, чтобы сократить время отклика. Если запрашивается слишком много моделей, это либо загружает их все в память и рискует получить OOM, либо убивает время отклика, переключаясь на диск и обратно.

АВС Лямбда

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

K8s государство

В предыдущем сообщении в блоге мы обсуждали использование Kubernetes в DeepCure для всех наших микросервисов. Развернув службу с помощью K8s, мы можем полагаться на EKS для большей части управления инфраструктурой, имея при этом свободу управления планированием и базовым оборудованием. Это также позволяет нам назначать подмножество моделей каждому контейнеру, распределяя нагрузку модели по всему кластеру. Недостатком является отсутствие автоматического распределения нагрузки или перераспределения в случае масштабирования кластера или потери контейнера. Это также требует ручного обслуживания и обновления назначений подмножества моделей по мере добавления новых моделей.

Войдите в кольцо

Обслуживание моделей через микросервис, управляемый Kubernetes, больше всего соответствовало нашим требованиям. Но нам по-прежнему нужен был способ автоматического распределения ответственности за модель между узлами в кластере. Решение заключалось в использовании согласованного кольца хеширования, при этом ответственность модели назначалась на основе относительного расположения сущностей в кольце.

Согласованное хеширование описывает распределение данных с использованием алгоритма хеширования для определения местоположения каждого объекта. Кольцо относится к единичному кругу, при этом хэш каждой сущности сопоставляется с местоположением в круге. В приведенном ниже фрагменте кода показано, как можно использовать хэш md5 объекта для размещения его в кольце.

import hashlib

key = "my-model-id" # identifier of the entity
md5_d = hashlib.md5(key.encode()).digest() # md5 hash of the entity's identifier
hash_val = (md5_d[3] << 24) | (md5_d[2] << 16) | (md5_d[1] << 8) | md5_d[0] # take the first 32 bits
location = (hash_val % 36000) / 100 # entity's location in the ring in degrees, taken to two decimal places

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

Чтобы узнать, за какие модели отвечает узел, найдите его местоположение в кольце и пролистайте против часовой стрелки, пока не будет найдена запись другого узла. С другой стороны, чтобы найти узел, отвечающий за модель, найдите его местоположение и проведите по часовой стрелке, пока не найдете узел.

Определение кольца

Каждый объект в кольце хранится в виде строки в DynamoDB, управляемой AWS базе данных NoSQL. Каждая строка содержит местоположение объекта в кольце и важную информацию о взаимодействии с этими объектами. Для моделей эта информация включает в себя, откуда модель может быть извлечена, например. ведро S3. Для узлов информация включает частный IP-адрес и отметку времени, которые обновляются процессом пульса узла. Узел и запись модели показаны ниже соответственно.

{
    "name": "mol-property-service-746976497d-w4fzh",
    "replica": 1,
    "short_hash": 2938500087,
    "entity_type": "node",
    "location": 0.87,
    "value": "10.0.5.115:8000",
    "ttl": 1667846949
}

{
    "name": "logd_uncertainty",
    "replica": 1,
    "short_hash": 1658134121,
    "entity_type": "pipeline",
    "location": 101.21,
    "value": "s3://foo/bar/logd_uncertainty.tgz"
}

Развертывание новых моделей

Процесс развертывания новой модели очень прост: добавьте новую сущность в кольцо. Когда добавляется новый объект узла, любой запрос прогнозирования для этой модели будет перенаправлен на узел, ответственный за него, в зависимости от их относительного положения. Каждый узел в кластере регулярно консультируется с кольцом на предмет изменений (каждые пять секунд). При добавлении новой модели ответственный узел загрузит ее и добавит в свой кэш в памяти. Если запрос поступает до того, как произойдет эта кольцевая проверка, модель будет загружена для удовлетворения запроса и сохранена.

Обработка отказа

Один из первых вопросов об архитектуре службы — как она будет реагировать на сбои. Каждый узел в нашем сервисе запускает процесс пульса, который обновляет свою запись в кольце. Это обновление установит значение «ttl» в точку в будущем, например. 25 секунд. Если текущая метка времени превышает пороговое значение «ttl», узел будет считаться неисправным. В случае сбоя любые запросы на модели этого узла будут приниматься следующим узлом справа в кольце. На рисунке используется базовый пример кольца, чтобы проиллюстрировать сбой.

Масштабирование горячих моделей

В нашем простом кольце каждая модель принадлежит одному узлу. Это может вызвать проблемы, если подмножество моделей запрашивается чаще, чем другие, или имеет более длительное время прогнозирования. Решение состоит в том, чтобы позволить нескольким узлам обслуживать запросы для одной и той же модели. Это можно сделать путем масштабирования горячей модели (моделей) в кольце, что делается путем добавления дополнительных записей в кольцо. На рисунке ниже показано, как масштабирование горячей модели выглядит в ринге.

Масштабирование узлов

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

Распространение запросов прогнозирования

После распределения моделей по кластеру они должны быть доступны для создания прогнозов. В DeepCure каждый из наших микросервисов находится за балансировщиком нагрузки приложений (ALB), распределяющим объем запросов по узлам кластера, и сервис модели свойств ничем не отличается. Запрос может быть отправлен на любой узел кластера, когда он достигает ALB. Этот узел просматривает запрошенные модели, консультируется с кольцом и выдает прямые запросы соответствующим узлам. Ответы на эти запросы объединяются в один ответ и возвращаются клиенту. На рисунке ниже показано, как запрос на выполнение прогнозов для четырех свойств распределяется по кластеру. Обратите внимание на цвета свойств в запросе, соответствующие цветам узла и записи кольца.

С нетерпением жду

Использование согласованного кольца хеширования для управления распределением моделей очень хорошо удовлетворило наши потребности в DeepCure. В другом сообщении блога мы поднимемся на уровень абстракции и рассмотрим более крупную архитектуру, частью которой является этот сервис, и то, как он обеспечивает быстрое развертывание модели и управление версиями.

Подпишитесь на Блог DeepCure, чтобы узнать больше историй и ресурсов для разработчиков, а также на DeepCure на LinkedIn, чтобы узнать последние новости о компании.