Package: dirba
DirbaНабор утилит для работы с ML моделями и их запуском в kafka.СодержаниеУстановкаРабота с kafkaЧто такое kafkaConsumerКонфигурацияProducerRunner для моделейModelAbstractBaseKafkaRunnerStrictRunnerAbstractKafkaRunnerЗапуск моделей через APIРабота с каталогамиАсинхронный APIСинхронный APIРабота с источниками данныхВалидация моделейУстановкаДанный модуль подразумевает несколько вариантов установкиСтандартныйpip install dirba. Подразумевает установку базовых библиотек. Позволяет пользоваться
каталогами и работать с источниками данныхДля работы с kafka. В таком случаеpip install dirba[kafka]установит все необходимые зависимостиДля работы с валидацией потребуется выполнитьpip install dirba[validation]Работа с kafkaЧто такое kafkaApache Kafka — распределённый программный брокер сообщений. Основное его назначение – организация асинхронного
взаимодействия между различными сервисами с гарантиями доставки.Если проще, то это промежуточное приложение, для общения различных приложений. Вполне логичный вопрос, который
может возникнуть – "зачем нужно это промежуточное приложение ?"
На деле, при реализации взаимодействия между различными приложениями, встаёт ряд проблем:что делать, если одно из приложений будет недоступнокак распределить нагрузку между несколькими экземплярами приложения (для увеличения производительности)как хранить историю сообщений (чтобы можно было что-то отладить/посмотреть)как гарантировать обработку сообщенияПри прямой схеме взаимодействия (например, по HTTP протоколу) решение вышеперечисленных проблем потребует
реализации дополнительного кода для каждого из приложений, что замедляет разработку и увеличивает вероятность
ошибок.Kafka же, использует асинхронный подход во взаимодействии различных сервисов. Она принимает сообщение и хранит
его до тех пор, пока другое приложение не заберёт предназначенное для него сообщение.В такой схеме работы есть несколько ключевых понятий, которые стоит уточнить:Producer– внешнее приложение,отправляющеесообщениевkafkaConsumer– внешнее приложение,получающеесообщениеизkafkaBroker– экземпляр (инстанс) приложенияApache Kafka(обычно их запускают от 3х штук, на разных
серверах, чтобы обеспечить бесперебойную работу (ведь без kafka приложение не смогут общаться)). В
совокупности превращаются вKafka Cluster, с которым, фактически, и общаются ваши приложения (любое
приложение может переключиться на активный инстанс или работать сразу с несколькими)Topic– "тема", определённый ключ, с которым ассоциируются отправляемые в kafka сообщения. Этот же ключ
используется для получения сообщенийPartition– "часть топика", логическая единица, для записи информации изтопика. Кол-во партиций
связано с максимальным кол-вом приложений-consumer'ов (в одномтопикене могут параллельно писать **
консьюмеров** больше, чем кол-вопартиций). Как было сказано, партиция – логическая единица, которая могут
быть расположены в произвольном порядке на экземплярах приложенияApache KafkaИтого мы имеем:В рамках нашей системы у нас есть некоторыйKafka Cluster, к которому мы можем подключитьсяУ нас есть наше приложение-producer, отправляющее сообщения с данными в конкретныйTopicна этом
кластере.Другое приложение-consumer считывает сообщения из этогоTopicи обрабатывает их согласно своей
бизнес-логике. При обработке сообщения приложением, оно подтверждает факт обработки, отправкой специального
commit-сигнала (как правило реализуется драйвером для kafka, но можно выполнять и самостоятельно).Как consumer понимает, какое сообщение братьЕсли вы внимательно прочитали главу выше, то вы должны были подметить, что там ни слова не сказано о том, как
консьюмер понимает, какие данные ему брать из kafka. Да, есть топик, которые разграничивает разные по виду
сообщения, но что если их в kafka накопилось несколько сотен или даже тысяч ? Для решения этой проблемы в
kafka существует механизмгрупп(group) исдвигов(offset).Группа – это обобщение для определённого приложения. Подразумевает выдачу некоторого общего идентификатора.
Такой идентификатор называютgroup_id.Сдвиг – обозначение № сообщения, на котором находится та или иная группа.Давайте рассмотрим следующий пример.Мы организовываем работу в некотором банке. Понятно, что почти все приложения банка – высоконагруженные, а
потеря данных в них недопустима. Поэтому мы само собой вспоминаем про kafka и радостно её берём.Для обработки всех транзакций (информации о платежах), мы выставляем 3 экземпляра написанного нами
приложенияacquiring. Скорее всего это приложение представляет собой REST или gRPC API, которое
взаимодействует с клиентским приложением, принимает транзакцию и тут же отправляет её в kafka,а это значит
что приложение acquiring – Producer. В принципе, нам не важно, как приложение будет это делать, главное
чтобы приложения имели одну структуру, в независимости от экземпляра. Но на всякий случай мывыдадим всем
экземплярам одинаковый group_id, чтобы явно идентифицировать приложение.После того как данные оказались в kafka, можно выдохнуть. Как минимум, мы уже их не потеряем, если что-то
навернётся. Значит, можно начать обрабатывать наши транзакции. Обработка транзакций – сложный процесс с кучей
логики, поэтому для них у нас поднято несколько экземпляров (на рисунке 2, а в реальности может быть и сильно
больше). Экземпляры приложенияbusiness logicявляется Consumer'ами, т.к. достают данные из kafka, мы их
тоже пометим отдельнымgroup_id. В данном случае, одинаковыйgroup_idобязателендля экземпляров **
одного и того же приложения**. Так мы гарантируем, что одно и то же сообщение не будет обработано более 1
раза.Ну и осталось приложениеlogger, которое такжеявляется Consumer'ом. Для него у нас всего-лишь 1
экземпляр, т.к. логирование нам не к спеху и пригодится только при инцидентах. Ему мывыставим
другойgroup_id, отличный отbusiness logic. Тем самым, мы и логгеру и бизнес логике позволим прочитать
одни и те же сообщения (ведь мы хотим все сообщения и обработать и залоггировать).ConsumerДля того чтобы реализовать свой Consumer, вам необходимо унаследовать абстрактный
классrunners.kafka_runner.AbstractKafkaConsumer.Для этого вам необходимо:указать формат входной модели (подразумевается, что все сообщения в кафке хранятся в форматеjson).
Входная модель –pydanticмодель, описывающая перечень полей вjsonсообщении. за входную модель
отвечает поле классаInputModelреализовать методon_startupреализовать методprocess??????PROFIT!!!Давайте посмотрим на пример такого сервисаimportosimportaiomiscimportdirba.logging_utilsfromdirba.runnersimportkafkaaskafka_runnerclassExampleConsumer(kafka_runner.AbstractKafkaConsumer):InputModel=kafka_runner.topic_schemas.LoaderMessageasyncdefon_startup(self):print('started up')asyncdefprocess(self,message:InputModel):print('message',message.uid_loaded_data,'consumed')if__name__=='__main__':config=kafka_runner.KafkaConfig(input_topic='loaded_data',group_id='example_dirba_consumer',bootstrap_servers=os.environ['BOOTSTRAP_SERVERS'])consumer=ExampleConsumer(config,from_topic_begin=False)withaiomisc.entrypoint(consumer)asloop:dirba.logging_utils.set_logging(sentry_url=os.environ['SENTRY_URL'])loop.run_forever()В данном случае у нас получится echo сервис, который будет выводить содержимое полученного из kafka сообщения.КонфигурацияДля подключения к kafka вам необходимо заполнитьKafkaConfig. Он содержит информацию о:топикеinput_topic, из которого будут считаны сообщенияидентификаторе приложенияgroup_id(должен быть уникален для каждого приложения)адресах подключенияbootstrap_servers(<ip1>:<port1>,<ip2>:<port2>, т.е. адреса для подключения к
инстансам, через запятую)Также, для отслеживания работы сервиса, необходимо сконфигурировать логирование с помощьюdirba.logging_utils.set_logging(sentry_url=os.environ['SENTRY_URL'])Где,os.environ['SENTRY_URL']– url для отправки информации в sentry сервис (для PRODUCTION – обязательно,
для локальной разработки – по желанию)Особенности сервисаЧто под капотом делает сервис:сериализует данные из kafkaвалидирует данные по указанной схемеобеспечивает работоспособность сервиса при ошибках.отправляет информацию об ошибках в sentryотсылает информацию о производительности в sentryЧтоне гарантирует сервис:commitпосле каждого сообщения (атомарность для обработки сообщений)повторную обработку сообщения при ошибкеОсновной кейс использования – потоковая обработка большого кол-ва сообщений, с допущением на частичную потерю
данных (в наших реалиях это менее 0,0001% от всех данных, но вы всё ещё можете повторить обработку, отследив
сообщения через sentry).ProducerПо реализации, очень похож на Consumer~~, однако, как правило, сам по себе бесполезен~~.Для запуска вам необходимо:указать формат выходной (output) модели (подразумевается, что все сообщения в кафке хранятся в формате json)
. Входная модель – pydantic модель, описывающая перечень полей в json сообщении. за выходную модель отвечает
поле класса OutputModelреализовать методon_startupреализовать методpackдля сериализации данных в выходную модельпереопределить методstart, добавив необходимую логикуimportosfromtypingimportAnyimportaiomiscimportpydanticimportdirba.logging_utilsfromdirba.runnersimportkafkaaskafka_runnerclassExampleMessageSchema(pydantic.BaseModel):id:intmessage:strclassExampleProducer(kafka_runner.AbstractKafkaProducer):OutputModel=ExampleMessageSchemadefpack(self,data:Any)->OutputModel:returndataasyncdefon_startup(self):print('started up')asyncdefstart(self,trigger_start=True):awaitsuper(ExampleProducer,self).start(trigger_start=True)foriinrange(10):message=ExampleMessageSchema(id=i,message=f'test{i}')awaitself.send_message(message)print('sent',message)if__name__=='__main__':config=kafka_runner.KafkaConfig(None,group_id='example_dirba_producer',bootstrap_servers=os.environ['BOOTSTRAP_SERVERS'],output_topic='test__dirba')producer=ExampleProducer(config)withaiomisc.entrypoint(producer)asloop:dirba.logging_utils.set_logging(sentry_url=os.environ['SENTRY_URL'])loop.run_forever()В данном примере мы просто немного поспамим сообщениями в kafka.На самом деле классAbstractKafkaProducerявляется скорее промежуточным (поэтому тут нужно переопределять
методstart) и предполагается, что он будет использоваться в множественном наследовании или будет запущен в
параллель с другим сервисом (смотри/examples/example_producer.pyв репозитории).Данный сервис лишь гарантирует robust'ное (устойчивое к сбоям) соединение с kafka, однако никоим образом не
обрабатывает ошибки, которые могут возникнуть в вашей логике.Конфиг заполняется так же, как и вproducer'e, но вместоinput_topic, необходимо
указатьoutput_topic, в который будут отправляться сообщения.Runner для моделейЕсливы реализуетеMLмодель,то вам сюдаRunner'ы – семейство классов, обеспечивающие абстракцию различного уровня для потоковой работы с kafka (принял
сообщение, создал новое сообщение).Все Runner'ы работают с базовой абстракцией –ModelModelModel – сущность обеспечивающая обработку входных данных, с получением вероятностной оценки о принадлежности к
кому-либо классу. Иначе говоря – обёртка для классификаторов.К сожалению, Model, имеет достаточно строгий API и приспособить его для другой задачи будет довольно
проблематично. Поэтому давайте посмотрим на то, как его реализовать.AbstractProhibitedModelAbstractProhibitedModel– частный случайAbstractModel. Единственное отличие – явная привязка к некоторому
каталогу. Подразумевается, что любая выдаваемая оценка от такой модели, должна быть длякатегории,
существующей вкаталоге.classDamboInput(pydantic.BaseModel):text:strclassDamboModel(AbstractProhibitedModel):def__init__(self,category_catalog:CategoryCatalog):super().__init__(category_catalog)self.category_catalog.load_catalog_sync()defpredict(self,features:DamboInput)->List[Predict]:return[Predict(score=1,category=i)foriinself.category_catalog.catalog_values]defpreprocess(self,features:LoaderMessage)->DamboInput:returnDamboInput(text='dambooooooooooooo')defauthor(self)->Author:returnAuthor(name='dambo',version='0.0.1')Модель подразумевает отдельную функцию для предварительной обработки входного сообщения, и функцию по
получению вероятностей принадлежности к классу.Советыесли модель не отнесла входной материал к какой-либо категории, то нужно возвращатьпустой списоквалидацию данных лучше оставить вне логики моделилучше избегать многопроцессорной обработки. Т.к. предполагается запуск через kafka, гораздо проще будет
увеличить кол-во инстансов, чем мучиться с последствиями многопоточности.AbstractBaseKafkaRunnerБазовый класс для обработки данных из kafka моделью. Подразумевает следующий порядок работы:получение сообщения из kafkaвалидация по входной моделипередача данных моделиотправка результирующего сообщения в kafkaПо своей сути – небольшая обёртка надConsumerиProducer, с добавлением вызова
моделиНе стоит использовать данный класс для деплоя модели, если не хотите писать много кодаimportosimportuuidfromtypingimportOptionalimportaiomiscfromdirba.models.abcimportPredictfromdirba.models.damboimportDamboModelfromdirba.runners.kafka_runnersimportAbstractBaseKafkaRunner,KafkaConfigfromdirba.runners.kafka_runners.topic_schemasimportLoaderMessage,AnalysisMessage,AnalysisResult,ModelOutputfromdirba.utils.catalogsimportCategoryCatalogclassDamboBaseKafkaRunner(AbstractBaseKafkaRunner):InputModel=LoaderMessageOutputModel=AnalysisMessagedefpack_model_answer(self,message:InputModel,predict:Predict)->Optional[OutputModel]:input_data=message.dict(include={"uid_query","query_id","driver_id","category_id","type_id","uid_search","uid_filter_link","uid_loader","uid_loaded_data","uid_analysis",})model_answer=ModelOutput(category=predict.category,estimate=predict.score)result=AnalysisResult(content_ref=message.result,model=model_answer,type_content=message.type_content)packed=self.OutputModel(uid_analysis=uuid.uuid4(),author=self.model.author(),result=result,**input_data)returnpackeddefis_adorable(self,input_message:InputModel)->bool:returninput_message.type_content=='text'asyncdefon_startup(self):print('started_up')if__name__=='__main__':config=KafkaConfig(input_topic='loaded_data',output_topic='dirba_simple_runner_test',group_id='example_dirba_consumer',bootstrap_servers=os.environ['BOOTSTRAP_SERVERS'])category_catalog=CategoryCatalog(os.environ['CATALOG_URL'])model=DamboModel(category_catalog)runner=DamboBaseKafkaRunner(model,config,from_topic_begin=True)withaiomisc.entrypoint(runner)asloop:loop.run_forever()StrictRunnerТот жеAbstractRunner, однако привязан к определённому формату данных. При получении сообщения делает
проверку на наличие категории в переданном каталоге.Не стоит использовать данный класс для деплоя модели, если не хотите писать много кодаimportloggingimportosfromrandomimportrandomfromtypingimportListimportaiomiscfromdirba.models.abcimportPredictfromdirba.models.damboimportDamboModel,DamboInputfromdirba.runners.kafka_runnersimportAbstractStrictBaseKafkaRunner,KafkaConfigfromdirba.utils.catalogsimportCategoryCatalogclassDumbModel(DamboModel):defpredict(self,features:DamboInput)->List[Predict]:ifrandom()>0.5:return[Predict(score=1,category=i)foriinself.category_catalog.catalog_values]else:return[Predict(score=1,category=-20)]classDambKafkaRunner(AbstractStrictBaseKafkaRunner):defis_adorable(self,input_message:AbstractStrictBaseKafkaRunner.InputModel)->bool:returninput_message.type_content=='text'asyncdefon_startup(self):print('started_up')if__name__=='__main__':config=KafkaConfig(input_topic='loaded_data',output_topic='dirba_simple_runner_test',group_id='example_dirba_consumer',bootstrap_servers=os.environ['BOOTSTRAP_SERVERS'])category_catalog=CategoryCatalog(os.environ['CATALOG_URL'])model=DumbModel(category_catalog)runner=DambKafkaRunner(model,config,from_topic_begin=True,category_catalog=category_catalog,produce_incorrect_categories=True)withaiomisc.entrypoint(runner,log_level=logging.DEBUG)asloop:loop.run_forever()AbstractKafkaRunnerОсновной класс для использованияИз коробки содержит:мониторинг через sentryвыгрузка метрик в prometheus (роут/metrics/)Вот пример его использованияimportosimporttimefromrandomimportrandomfromtypingimportIterableimportaiomiscimportdirbafromdirba.models.abcimportPredictfromdirba.models.damboimportDamboModel,DamboInputfromdirba.runners.kafka_runnersimportKafkaConfigfromdirba.runners.kafka_runners.runnerimportAbstractKafkaRunnerfromdirba.utils.catalogsimportCategoryCatalogfromdirba.utils.metrics.prometheus_exporterimportPrometheusExporterclassExampleRunner(AbstractKafkaRunner):InputModel=AbstractKafkaRunner.InputModelOutputModel=AbstractKafkaRunner.OutputModeldefis_adorable(self,input_message:InputModel)->bool:returninput_message.type_content=='text'asyncdefon_startup(self):print('started up')classDumbModel(DamboModel):defpredict(self,features:DamboInput)->Iterable[Predict]:# для использования дополнительных метрик они должны быть сконфигурированыself.metric.add_label_values(ora='jojo',muda='dio')time.sleep(1)# в тех ситуациях, когда нужно прокинуть разные значения метрик для каждого predict'a,# можно возвращать их с помощью генератораforiinself.category_catalog.catalog_values:ifrandom()>0.3:self.metric.add_label_values(ora=f'jojo_{i}',muda='dio')yieldPredict(score=1,category=i)if__name__=='__main__':config=KafkaConfig(input_topic='loaded_data',output_topic='dirba_simple_runner_test',group_id='example_dirba_runner',bootstrap_servers=os.environ['BOOTSTRAP_SERVERS'])# т.к. это "строгий" runner (как и модель), они должны взаимодействовать с каталогомcategory_catalog=CategoryCatalog(os.environ['CATALOG_URL'])model=DumbModel(category_catalog)# from_topic_begin - отладочный вариант запуска. В проде он должен быть выставлен в Falserunner=ExampleRunner(model,config,from_topic_begin=True,category_catalog=category_catalog)# для прокидывания метрик также необходимо запустить экспортёрexporter=PrometheusExporter(port=int(os.environ['PORT']),address='0.0.0.0')# после инциализации runner'a можно сконфигурировать доп. набор метрик, при необходимостиrunner.material_metric.register_labels('ora','muda')# конфигурация логированияdirba.logging_utils.set_logging(sentry_url=os.environ['SENTRY_URL'])withaiomisc.entrypoint(runner,exporter)asloop:loop.run_forever()Из отличий – добавился ещё один сервисPrometheusExporter, который отвечает за выгрузку метрик. Метрики
можно получить по адресу, который был передан экспортёру, и роуту/metrics/. И да, экспортёр можно и не
запускать, но в таком случае и метрики выведены не будут. Однако это никоим образом не сломает раннер.Также в примере продемонстрирована возможность добавления label'ов к метрике prometheus.Запуск моделей через APIWIPРабота с каталогамиТ.к. все модели ориентированы на получение вероятности принадлежности к категории, то мы где-то должны эти
категории брать. Хардкод – дело неблагодарное, поэтому у нас есть общий сервис, в котором содержится
информация о всех категориях. Раз это сервис, то, наверное, хотелось бы иметь обвязку для обращения к нему, да
ещё и с кэшом желательно.Для этого у нас естьutils.catalogs.abc.AbstractCatalog. Сам класс представляет собой простенькую обвязку
для обращения к стороннему HTTP сервису.Если вы захотите поработать именно сAbstractCatalog, то вам необходимо:Определитьpydanticмодель данных для данных каталогаПереопределить методparse_catalog_responseдля сериализации данных каталога в словарь, для последующей
обработкиПример вы можете найти в обвязке вокруг конкретного каталогакатегорийвutils.catalogs.category.У каталога есть 2 варианта API – синхронный и асинхронный.Асинхронный APIimportasyncioimportosimporttimeimportloggingfromdirba.utils.catalogsimportCategoryCataloglogging.basicConfig(level=logging.DEBUG)catalog=CategoryCatalog(os.environ['CATALOG_URL'])asyncdefasync_example():# демонстрация работы кэшаstart_time=time.time()count=20foriinrange(count):val=awaitcatalog.get_value(catalog_id=14)print(count,'requests ended up in',time.time()-start_time,'seconds')print(val)awaitcatalog.get_catalog_data()# в том числе и для несуществующих значенийforiinrange(800,810):val=awaitcatalog.get_value(catalog_id=i)# все последующие обращения берутся из кэшаval=awaitcatalog.get_value(catalog_id=i)val=awaitcatalog.get_value(catalog_id=i)if__name__=='__main__':asyncio.run(async_example())Основная функция для взаимодействия –get_value. Под капотом, она самостоятельно обновляет кэш значений из
каталога (если вы запрашиваете несуществующее значение, то каталоги будут обновлены, а дальше вступит в силу
кэш). Можно работать с результирующим словарём каталога напрямуюкрайне не рекомендуется.И в синхронном варианте это будет выглядеть как-то такСинхронный APIimportosimporttimeimportloggingfromdirba.utils.catalogsimportCategoryCataloglogging.basicConfig(level=logging.DEBUG)catalog=CategoryCatalog(os.environ['CATALOG_URL'])defsync_example():# демонстрация работы кэшаstart_time=time.time()count=20foriinrange(count):val=catalog.get_value_sync(catalog_id=14)print(count,'sync requests ended up in',time.time()-start_time,'seconds')print(val)# в том числе и для несуществующих значенийforiinrange(800,810):val=catalog.get_value_sync(catalog_id=i)# все последующие обращения берутся из кэшаval=catalog.get_value_sync(catalog_id=i)val=catalog.get_value_sync(catalog_id=i)if__name__=='__main__':sync_example()Разницы по скорости между ними нет, однако в зависимости от вашего контекста исполнения, вам могут
понадобиться синхронный или асинхронный варианты.Более подробную информацию вы можете найти в документации методов и самого классаРабота с источниками данныхЕсли вы разрабатываете ML модель, то у вас скорее всего встанет ещё один вопрос, как достать данные для их
прогона через модель, ведь вам приходит лишь идентификатор.Для работы с данными есть набор классов вutils.data_loader. Останавливаться на подробностях
имплементацииAbstractDataLoaderсмысла не вижу (т.к. все реализации уже готовы). Скажу лишь, что из коробки
там прописанаretryполитика для HTTP запросов, чтобы вам поменьше страдать.Используются классы для данных максимально просто:выбираете нужный по названию сущностиуказываете URL до сервисаполучаете данные по идентификатору сущности, который вы получилиimportosfromdirba.utils.data_loaderimportTextLoader,HtmlLoader,ImageLoaderif__name__=='__main__':service_url=os.environ['SERVICE_URL']text_loader=TextLoader(service_url)print(text_loader.get_content(entity_id=1))html_loader=HtmlLoader(service_url)print(html_loader.get_content(entity_id=1))image_loader=ImageLoader(service_url)image=image_loader.get_content(entity_id=1)withopen('./tempo.jpg','wb')asf:f.write(image[1].content)Валидация моделейWIP