Transformador em AppSec on-premise: como integramos um modelo ML para classificação de segredos em um produto sem GPU

CodeScoring 12,66 Classificação 21 Assinantes Assinar nouhadonosor Há 36 minutos Transformador em AppSec on-premise: como integramos um modelo ML para classificação de segredos em um produto sem GPU Médio 9 min 1.1K Blog da empresa CodeScoring Segurança da Informação * Aprendizado de Máquina * Python * Código Aberto * Tutorial TL;DR; Contamos como integramos um modelo de classificação de segredos baseado em CodeBERT em um produto de produção com restrições rígidas de hardware, reduzindo o tempo de inferência de 320 para 90 segundos e o tamanho do modelo de ~600 para ~130 MB - sem aceleradores discretos e dependências pesadas. No artigo anterior, contamos como preparamos e montamos nosso modelo baseado no modelo CodeBERT para classificar os acionamentos dos mecanismos de busca de segredos no código. Falamos sobre conjuntos de dados, tokenizador, métricas e como conseguimos aumentar o PR AUC* de 0,70 para 0,90. *PR AUC - Área sob a Curva Precisão-Recall, literalmente “área sob a curva “precisão-exaustividade””, em palavras simples: a confiança do modelo na avaliação probabilística da classe rara positivamente verdadeira na amostra geral. Sobre mim Eu sou Nathan, techlead do módulo Secrets na CodeScoring. Trabalho com Secrets há quase 2 anos, antes trabalhei em FinTech e EdTech, agora escrevo em Python e Go. Sobre a empresa Na CodeScoring, estamos envolvidos na criação de uma solução para trabalhar com segurança com código aberto, verificar a compatibilidade de licenças, procurar segredos e avaliar a qualidade do código em termos de equipe. Sobre o que é o artigo? Como abordamos a questão da integração de um modelo computacionalmente intensivo em um produto AppSec dentro de uma máquina alvo com capacidade limitada. O artigo será útil principalmente para MLOps que desejam combinar aprendizado de máquina e segurança da informação, bem como para aqueles que integram soluções que consomem muitos recursos em condições de poder computacional limitado sem reduzir o desempenho do produto. Contexto: o que é on-premise e quais são os dados de entrada em geral? Nosso produto é lançado on-premise, ou seja, nas instalações do cliente. O cenário alvo é ajudar a verificar os resultados da varredura de segredos no código dentro dos pipelines CI/CD. Na implementação anterior do modelo, tínhamos um processo refinado e familiar para o usuário, então queríamos encaixar a nova versão nele para não causar atrito desnecessário para o usuário. Transformadores são uma história pesada, e antes de seguir em frente, era preciso responder honestamente: podemos executá-lo em condições de recursos de hardware limitados do cliente e requisitos de segurança com qualidade de UX aceitável? Após várias rodadas de discussão e filtragem do real e do irreal, chegamos aos seguintes requisitos: O modelo funciona estritamente na máquina do cliente. A inferência dentro da análise não deve levar mais de 5 minutos para um projeto hipotético que contenha 1000 segredos, pois o trabalho do modelo em segredos é a parte mais demorada. Na prática, esses volumes quase nunca são encontrados. Também é preciso levar em consideração que Segredos não estão sozinhos e competirão por recursos com outras partes do produto e do sistema. Minimizamos o volume de dependências, isso é necessário para: Minimizar o peso da montagem de produção. Reduzir o número de possíveis vulnerabilidades. Facilitar o processo de montagem a partir do código-fonte, esta é uma etapa obrigatória na certificação da FSTEC da Rússia. Como a máquina alvo do cliente, consideramos um certo servidor Linux, com arquitetura de processador x86 com 4 threads lógicos, uma velocidade de clock de 2 GHz, 16 GB de RAM e a capacidade de executar o Docker. O modelo reside no disco e é executado por um contêiner com um consumidor de tarefas. Ele realiza inferência dentro da tarefa de análise do código-fonte para a presença de segredos, onde primeiro os segredos são encontrados usando mecanismos como gitleaks ou trufflehog, após o que o modelo os avalia (respondendo à pergunta “o segredo é positivamente verdadeiro?”) por uma série de propriedades, onde as principais são o próprio segredo e seu ambiente. O usuário inicia a análise através da interface (manualmente ou por agendamento), após o que a tarefa é executada em segundo plano. Além disso, o usuário pode treinar o modelo em seus próprios dados marcados diretamente na interface do produto, sem transferi-los para fora do contorno - isso é fundamental para nossos clientes com requisitos de segurança de dados. Primeira iteração: PyTorch, pesado e desajeitado O primeiro instinto foi óbvio: pegar o modelo que foi originalmente implementado no PyTorch e integrá-lo ao produto. A estrutura do próprio modelo no PyTorch é a seguinte: Ou seja, o modelo é uma sequência de vários módulos: pré-processamento, tokenizador, transformador, camada de saída. Peso e dependências Árvore para torch v2.12.0 sob Linux x86 junto com dependências CUDA (imagem grande) PyTorch arrasta uma enorme árvore de dependências transitivas. Para um produto no campo do AppSec, é muito importante ser seguro por conta própria. Cada dependência extra é uma superfície de ataque potencial que precisa ser rastreada, corrigida e auditada. Referência Temos nosso próprio conjunto de dados, que é um conjunto de registros, cada um dos quais contém: uma string com a ocorrência de um segredo, a própria ocorrência do segredo, vários atributos relacionados e marcação, que consideramos verdade e que nos diz se o segredo é positivamente verdadeiro. Pegamos 1000 desses registros e os usamos como uma “régua”. Velocidade de inferência As primeiras medições em nossa referência, na imagem Debian, deram 320 segundos. A inferência faz parte do processo de análise, que o usuário inicia de forma interativa ou por horário. Portanto, o tempo máximo permitido para 1000 segredos foi de 4-5 minutos, levando em consideração que as outras partes da análise passarão rápido o suficiente. Alpine Linux e musl libc Nossas imagens de produção são construídas no Alpine Linux, que usa musl libc em vez de glibc. Uma pequena explicação do que é libc e qual é a diferença entre musl e glibc: libc é um conjunto básico de componentes em sistemas Linux que permite que programas em C e C++ funcionem normalmente, a implementação mais popular desta biblioteca é glibc. Ele tem muitas funções úteis que tornam mais fácil escrever e executar vários softwares no glibc Linux, mas como acontece com frequência, tudo de bom precisa ser pago de alguma forma, e aqui também. Os sistemas Linux no glibc costumam conter um grande número de vulnerabilidades e, em geral, têm um volume maior, portanto, se a leveza e a segurança forem muito importantes para você, você deve recorrer a outra implementação do libc. No nosso caso, usamos imagens Alpine, que implementam libc através de seu conjunto musl, é leve e seguro, mas muitas vezes não tem funcionalidade “açucarada”. PyTorch não distribui oficialmente rodas sob musl. Isso significava ou rejeitar o Alpine ou construir o PyTorch a partir do código-fonte. Recusar o Alpine seria bastante difícil para nós, então decidimos tentar construir o PyTorch sob o Alpine. No entanto, isso nos levou ao fato de que essas rodas precisam ser montadas de uma certa maneira, e não havia certeza clara de que rodas específicas funcionariam em plataformas diferentes. Além disso, embora funcionalmente musl e glibc sejam (quase) idênticos, ainda havia preocupações de que o PyTorch funcionasse de forma diferente neles - e não em vão. O principal problema foi a biblioteca com a implementação de threads libpthread em glibc, que tem a função pthread_attr_setaffinity_np() e está ausente na implementação musl. Isso foi tratado desativando o sinalizador USE_DISTRIBUTED ao montar, mas isso matou completamente a velocidade de inferência do modelo. Ele, é claro, funcionou na CPU, mas nossos descendentes provavelmente já poderiam desfrutar dos resultados de seus cálculos. Portanto, esta opção também não nos serviu. Segunda iteração: Wolfi+PyTorch, complexidade injustificada Antes de chegar à solução final, tentamos contornar o problema com musl de outra forma. Como usamos o alpine devido ao seu pequeno footprint e baixo número de vulnerabilidades, a transição para imagens padrão com glibc não era adequada para nós - elas são muito pesadas e costumam conter vulnerabilidades. No entanto, encontramos uma exceção: a imagem base Wolfi, que funciona no glibc, mas ao mesmo tempo em suas propriedades (peso e segurança) é muito próxima do Alpine. Como o modelo é executado em um contêiner com um consumidor de tarefas, a ideia era simples: substituir sua imagem base por Wolfi e, assim, fechar todos os requisitos de assunto e negócios. Chegamos à montagem real da imagem - e até a montamos. Mas duas vezes apareceram imediatamente problemas: Wolfi não permite baixar versões de pacotes gratuitamente, exceto a última. Para um produto de produção, onde a reprodutibilidade da montagem é crítica, isso é inaceitável: você não pode corrigir uma versão específica e ter certeza de que a imagem será montada de forma idêntica em um mês. PyTorch dentro do contêiner Wolfi se comportou de forma imprevisível: rodas com diferentes parâmetros de montagem se comportaram de forma diferente em plataformas diferentes, um colega pode estar bem, mas a inferência pode ficar presa em algum estágio e depurar o PyTorch - esta não é a atividade mais gratificante em termos de tempo e esforço, então algo mais era necessário, algo simples e elegante. A partir daí, surgiu a ideia de que deveríamos evitar completamente o PyTorch no ambiente de produção. Última iteração: ONNX, um experimento bem-sucedido ONNX (Open Neural Network Exchange) - um formato e ambiente de execução para redes neurais, originalmente criado pela Microsoft como uma alternativa independente de plataforma a frameworks como PyTorch e TensorFlow. É muito mais leve, modular e não arrasta uma pilha de dependências. ONNX, ao contrário do PyTorch, cria um gráfico de cálculo estático a partir do modelo, que é o principal segredo de seu desempenho. Além disso, tendo tal gráfico, você pode literalmente olhar para o esqueleto do modelo convertido para o formato ONNX. Parte do gráfico do modelo, criado usando netron. Me ajudou a distinguir os candidatos quando ficaram muitos Durante o desenvolvimento, precisávamos transformar o modelo PyTorch em formato ONNX e, em seguida, fazê-lo funcionar no ambiente do nosso produto, levando em consideração a máquina alvo. O PyTorch tem uma API que permite converter modelos no formato que precisamos, então, através de uma série de testes de diferentes abordagens, chegamos a um determinado script de conversão que nos satisfez, e obtivemos nosso modelo no formato certo. É verdade, havia uma armadilha aqui também: não existem rodas onnxruntime oficiais sob Alpine / musl. Mas, ao contrário do PyTorch, o ONNX é organizado de forma diferente - ele é dividido em pacotes separados, dos quais você pode pegar apenas a funcionalidade necessária. Isso tornou a tarefa de montagem a partir do código-fonte muito mais fácil. Nós bifurcamos o repositório microsoft/onnxruntime em nosso GitLab interno, fizemos os patches necessários para compatibilidade com musl e nossos requisitos, e também configuramos o pipeline CI, que monta as rodas e as publica em nosso repositório interno de pacotes. O resultado final: dois pacotes onnxruntime e onnxruntime-extensions pesando cerca de 25 MB. Ao montar a roda principal, removemos a montagem e execução de testes unitários após a montagem usando --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=OFF e --skip_tests, pois isso exigia dependências adicionais e complicava o pipeline. Também vale a pena montar na configuração Release usando --config Release, e não há necessidade de puxar todos os submódulos do repositório - nós mesmos os colocamos em nosso repositório, então adicionamos --skip_submodule_sync. Além disso, dois locais foram corrigidos ao copiar o repositório: onnxruntime/core/platform/posix/stacktrace.cc, que inclui a resolução do rastreamento da pilha quando a biblioteca falha - não precisamos disso na produção. version.txt - aqui estava apenas a versão errada, que substituímos pela real Você também pode conectar o construtor Ninja e o cache para CMAKE, que juntos devem acelerar a montagem como um todo e sua repetição. O próprio fato da conversão do modelo do formato PyTorch para ONNX deu um bônus imediato: ~15% de aumento na velocidade de inferência sem quaisquer otimizações adicionais. Além disso, para onnx existem bibliotecas em outros idiomas, e ele próprio transforma o modelo em um formato universal .onnx, que é amigável com outras ferramentas e frameworks, por exemplo, OpenVINO. Otimização: de 600 MB e 320 segundos para 130 MB e 90 segundos A conversão para ONNX não foi suficiente - 320 segundos se transformaram em 270, o que ainda é inaceitável, pois deixa uma pequena margem de tempo para todo o resto. Portanto, aplicamos uma série de técnicas adicionais. Quantização do modelo Uma representação aproximada de como a quantização funciona Quantização - esta é uma redução na precisão da representação de dados no modelo para reduzir o tamanho e acelerar os cálculos. Após várias séries de testes de um número bastante grande de candidatos derivados da quantização, paramos na quantização int8 no nível de operações individuais dentro do gráfico do modelo, combinando-a com uma série de otimizações estruturais do próprio modelo e sua embalagem no código. Como resultado, isso nos permitiu: Reduzir o tamanho do modelo de ~500 MB para ~130 MB Reduzir significativamente o tempo de inferência de ~270 segundos para ~90 segundos Metodologia de referência Para todas as medições, usamos uma abordagem: o conjunto de dados proprietário marcado mencionado acima de 1000 segredos, executado em uma máquina com 4 threads de CPU / 2 GHz / 16 GB de RAM. Esta é a nossa definição de “máquina média do cliente”: queríamos ter certeza de que isso funcionava com o cliente, e não apenas localmente com os desenvolvedores. Resultado final: 90 segundos para 1000 segredos. Isso é 3,5 vezes mais rápido que a versão original, embora com uma perda mínima de precisão (cerca de 2%). Dado que o conjunto de dados real do cliente raramente contém 1000 segredos ao mesmo tempo, na prática o tempo de trabalho é significativamente menor. Tabela comparativa dos principais parâmetros de cada abordagem: Opção Peso da biblioteca Tamanho do modelo Tempo de inferência (1000 segredos) PyTorch (linha de base) 1 GB (somente CPU) ~600 MB ~320 seg ONNX sem otimizações 25 MB ~500 MB ~250-230 seg ONNX + quantização + otimizações 25 MB ~130 MB ~90 seg Conclusão Conseguimos integrar um novo modelo ao produto em conformidade com todas as condições, que são a interseção de várias necessidades na forma de: conveniência do usuário, velocidade, economia de recursos e conformidade com os requisitos de segurança. A principal lição Em suma: não tenha medo de experimentar. No momento em que enfrentamos essa tarefa, havia poucas informações sobre a aplicação prática do ONNX em um ambiente de produção - especialmente em relação aos ambientes musl. Parte das soluções teve que ser encontrada por conta própria por tentativa e erro. Mas o resultado justificou o esforço. Falando mais amplamente: a introdução de transformadores em um produto on-premise não é uma tarefa apenas para a equipe de engenheiros de dados. Esta é uma tarefa de engenharia complexa, na qual você precisa não apenas criar uma hipótese sobre o modelo, testá-lo, otimizar a precisão, mas também reunir tudo isso em um módulo funcional do restante do sistema, que, além de ter sua própria pilha e padrões de trabalho, não funciona na nuvem, mas na máquina do cliente (que, geralmente, não terá aceleradores discretos a bordo). Agradecimentos Muito obrigado ao meu colega Nikita Besperstov por ajudar a montar as bibliotecas a partir do código-fonte e abrir caminho na floresta escura das bibliotecas do sistema Linux. Também gostaria de agradecer a Anton Volodchenko por ajudar a formular este artigo e orientar o processo de publicação como um todo. Para os leitores Se você tem experiência em executar transformadores em ambientes on-premise com restrições rígidas - teremos prazer em discutir nos comentários. Assine o Codescoring em Telegram, YouTube ou VK. Tags: appsec onnx mlops código aberto otimização codescoring ml Hubs: Blog da empresa CodeScoring Segurança da Informação Aprendizado de Máquina Python Código Aberto +3 4 0 CodeScoring Site Telegram 3 Carma Nathan Vlasenko @nouhadonosor AppSec TechLead Assinar Habr Carreira Habr Especialistas O fluxo de Segurança da Informação está disponível 24 horas por dia, 7 dias por semana, graças ao apoio dos amigos do Habr Habr Cursos para todos PUBLICIDADE Prática, Hexlet, SkyPro, cursos do autor - reunimos todos e pedimos descontos. Resta escolher! Ir Ir para o fluxo de Segurança da Informação Comentar Melhor do dia Semelhante