Transformer in on-premise AppSec: how we integrated an ML model for classifying secrets into a product without GPU

In our previous article, we talked about how we trained and assembled our model based on CodeBERT for classifying alerts from secret detection engines in code. We covered datasets, tokenizer, metrics, and how we managed to boost PR AUC* from 0.70 to 0.90.

*PR AUC — Area Under the Precision-Recall Curve, literally “area under the precision-recall curve”, in simple terms: the model's confidence when probabilistically evaluating a true positive rare class in the overall sample.

About Me

I'm Nathan, tech lead of the Secrets module at CodeScoring. I've been working on Secrets for nearly 2 years; I previously worked in FinTech and EdTech, and now I write code in Python and Go.

About the Company

At CodeScoring, we develop a solution for secure open source work, license compatibility checks, secret detection, and team-level code quality assessment.

What Is This Article About?

How we approached integrating a computationally heavy model into an AppSec product within the constraints of a target machine with limited compute capacity.

This article will be most useful for MLOps specialists who want to combine machine learning and information security, as well as for those who integrate resource-intensive solutions in environments with limited compute capacity without sacrificing product performance.

Context: What Is On-Premise and What Are the Overall Initial Conditions?

Our product runs on-premise, i.e. on the client's infrastructure. The target use case is assisting with verifying secret scanning results in code as part of CI/CD pipelines.

In the previous model implementation, we had a polished, user-familiar process, so we wanted to integrate the new version into it to avoid causing unnecessary friction for users.

Transformers are a complex undertaking, and before moving forward, we had to honestly answer: can we launch this under the conditions of limited client hardware resources and security requirements with acceptable UX quality?

After several rounds of discussions and filtering out what is feasible and what is not, we arrived at the following requirements:

1. The model runs exclusively on the client machine.

2. Inference during analysis should take no more than 5 minutes for a hypothetical project containing 1000 secrets, since model processing of secrets is the most time-consuming part. In practice, such volumes are almost never encountered. It is also worth considering that Secrets do not operate in isolation and will compete for resources with other parts of the product and the system.

3. We minimize the volume of dependencies, this is necessary for:

   Read also:

   • Which motherboard to choose for Intel Core i5-12400 and Intel Core i5-13400

   1. Minimizing the size of the production build.

   2. Reducing the number of potential vulnerabilities.

   3. Simplifying the source build process, which is a mandatory step when certifying with the Russian FSTEC (Federal Service for Technical and Export Control).

We consider the target client machine to be a Linux server with an x86 processor architecture, 4 logical threads, 2 GHz clock speed, 16 GB of RAM, and the ability to run Docker.

The model is stored on disk and launched via a container with the task consumer. It performs inference as part of the source code analysis task for secrets, where secrets are first detected using engines like gitleaks or trufflehog, after which the model evaluates them (answering the question "is the secret a true positive?") based on a number of properties, with the key ones being the secret itself and its surrounding context. The user launches the analysis via the interface (manually or on a schedule), after which the task runs in the background.

In addition, users can fine-tune the model on their own labeled data directly in the product interface, without transferring this data outside the perimeter — this is critical for our clients with data security requirements.

First iteration: PyTorch, heavy and unwieldy

The first instinct was obvious: take the model that was originally implemented in PyTorch and embed it in the product.

That is, the model is a sequence of several modules: preprocessing, tokenizer, transformer, output layer.

Weight and dependencies

PyTorch pulls a huge tree of transitive dependencies. For a product in the AppSec field, it is very important to be secure itself. Every extra dependency is a potential attack surface that needs to be monitored, patched, and audited.

Benchmark

We have our own dataset, which is a set of records, each containing: a string with a secret occurrence, the secret occurrence itself, several related attributes, and markup that we consider to be the truth and that tells us whether the secret is truly positive. We took 1000 such records and used them as a "ruler".

Inference speed

The first measurements on our benchmark, on a Debian image, gave 320 seconds. Inference is part of the analysis process that the user runs interactively or on a schedule. Therefore, we took 4-5 minutes as the maximum allowed time for 1000 secrets, taking into account that the other parts of the analysis will be quite fast.

Alpine Linux and musl libc

Our production images are built on Alpine Linux, which uses musl libc instead of glibc.

A brief explanation of what libc is and the difference between musl and glibc: libc is a basic set of components in Linux systems that allows C and C++ programs to run properly, the most popular implementation of this library is glibc. It has a large number of useful functions that make it easier to write and run various software on glibc Linux, but as is often the case, you have to pay for all the good things, and it is exactly the same here. glibc-based Linux systems often contain a large number of vulnerabilities and are generally larger in size, so if ease of use and security are very important to you, you should turn to alternative libc implementations. In our case, we use Alpine images, which implement libc via their own musl set, it is lightweight and secure, but often lacks "sugary" functionality.

PyTorch does not officially distribute wheels for musl. This meant either giving up on Alpine or building PyTorch from source. Giving up Alpine would be quite difficult for us, so we decided to try building PyTorch for Alpine.

However, this led us to the fact that these wheels had to be built in a specific way, and there was no clear certainty that specific wheels would work on different platforms. In addition, although musl and glibc are (almost) functionally identical, there were still concerns that PyTorch would run differently on them — and these concerns were not unfounded.

The main problem was the libpthread library, which implements threads in glibc, and has the pthread_attr_setaffinity_np() function that is missing in the musl implementation. This was fixed by disabling the USE_DISTRIBUTED flag during the build, however this completely killed the model inference speed. It, of course, worked on the CPU, but it would most likely be our descendants who would get to enjoy the results of its calculations. So this option was not suitable for us either.

Second iteration: Wolfi + PyTorch, unjustified complexity

Before arriving at the final decision, we tried to work around the musl issue differently. Since we use Alpine because of its small footprint and low number of vulnerabilities, switching to standard glibc-based images was not suitable for us — they are too heavy and often contain vulnerabilities. However, we found an exception: the base image Wolfi, which runs on glibc, but at the same time in terms of its properties (weight and security) is very close to Alpine. Since the model runs in a container with the task consumer, the idea was simple: replace its base image with Wolfi and thus meet all functional and business requirements. We got as far as actually building the image — and even built it. But two problems immediately came up:

1. Wolfi does not allow you to freely pull package versions other than the latest. For a production product, where build reproducibility is critical, this is unacceptable: you cannot fix a specific version and be sure that the image will be assembled identically a month later.

2. PyTorch inside the Wolfi container behaved unpredictably: wheels with different build parameters behaved differently on different platforms, a colleague might have had everything work fine, while my inference could get stuck at one of the stages, and debugging PyTorch is not the most rewarding activity in terms of time and effort, so we needed something else, something simple and elegant.

This is where the idea came from that we should completely avoid PyTorch in production.

Last iteration: ONNX, a successful experiment

ONNX (Open Neural Network Exchange) is a format and runtime environment for neural networks, originally created by Microsoft as a platform-neutral alternative to frameworks like PyTorch and TensorFlow. It is much lighter, more modular, and does not drag along a whole pile of dependencies.

Unlike PyTorch, ONNX creates a static computation graph from the model, which is the main secret of its performance. Also, having such a graph, you can literally look at the skeleton of the model converted to ONNX format.

During development, we needed to convert the PyTorch model to ONNX format and then get it to work in our product environment taking into account the target machine. PyTorch has an API that allows converting models to the format we needed, so after a series of tests of different approaches, we arrived at a specific conversion script that met our requirements, and we got our model in the required format.

However, there was a catch here too: official onnxruntime wheels for Alpine/musl also do not exist. But unlike PyTorch, ONNX is structured differently: it is split into separate packages, from which you can take only the functionality you need. This made building from source much easier.

We forked the microsoft/onnxruntime repository to our internal GitLab, added the necessary patches for compatibility with musl and our requirements, and also set up a CI pipeline that builds the wheels and publishes them to our internal package repository. Final result: two packages, onnxruntime and onnxruntime-extensions, weighing about 25 MB.

When building the main wheel, we disabled building and running unit tests after the build using --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=OFF and --skip_tests, since this required additional dependencies and complicated the pipeline.

It is also worth building in Release configuration using --config Release, and there is no need to pull all submodules from the repository: we add them to our own repository ourselves, so we add --skip_submodule_sync.

In addition, two places were patched when copying the repository:

1. onnxruntime/core/platform/posix/stacktrace.cc, which enables stacktrace resolution when the library crashes — we do not need this in production.

2. version.txt — there was simply an incorrect version here, which we replaced with the actual one.

You can also integrate the Ninja build system and a cache for CMake, which together should speed up the build in general and its repetition.

The very fact of converting the model from PyTorch to ONNX format gave an immediate bonus: ~15% increase in inference speed without any additional optimizations.

In addition, there are libraries for ONNX in other programming languages, and ONNX itself converts a model into the universal .onnx format, which is compatible with other tools and frameworks, for example OpenVINO.

Optimization: from 600 MB and 320 seconds to 130 MB and 90 seconds

Converting to ONNX was not enough — the 320 seconds turned into 270, which is still unacceptable, as it leaves very little time margin for everything else. So we applied a number of additional techniques.

Model quantization

Quantization is the reduction of data representation precision in a model to decrease its size and speed up computations. After several rounds of testing a fairly large number of candidates produced by quantization, we settled on int8 quantization at the level of individual operations inside the model graph, combining it with a number of structural optimizations of the model itself and its code wrapper. As a result, this allowed us to:

• Reduce the model size from ~500 MB to ~130 MB

• Significantly reduce inference time from ~270 seconds to ~90 seconds

Benchmarking methodology

For all measurements, we used the same approach: the aforementioned proprietary labeled dataset of 1000 secrets, run on a machine with 4 CPU threads / 2 GHz / 16 GB RAM. This is our definition of an "average client machine": we wanted to make sure that it works for actual clients, not just locally for developers.

Final result: 90 seconds for 1000 secrets. This is 3.5 times faster than the initial version, albeit with minimal accuracy loss (around 2%). Given that a real client dataset rarely contains 1000 secrets at once, the actual runtime is significantly lower in practice. Comparative table of key parameters for each approach:

Summary

We successfully integrated the new model into the product while meeting all requirements that represent the intersection of several key needs: user convenience, performance, resource efficiency, and compliance with security requirements.

Key takeaway

In short: don't be afraid to experiment. At the time we encountered this task, there was very little information on the practical use of ONNX in production environments, especially for musl environments. We had to figure out some solutions on our own through trial and error. But the result was worth the effort.

More broadly: integrating transformers into an on-premise product is a task not only for the data engineering team. It is a comprehensive engineering task that requires not only developing a model hypothesis, testing it, and optimizing accuracy, but also assembling all of this into a functional module of the rest of the system, which not only has its own tech stack and operational patterns, but also runs not in the cloud, but on the client's machine (which typically does not have discrete accelerators on board).

Acknowledgments

Special thanks to my colleague Nikita Besperstov for help with building libraries from source and navigating the dark forest of Linux system libraries.

I also want to thank Anton Volodchenko for help with drafting this article and guidance on the overall publication process.

To readers

If you have experience running transformers in highly constrained on-premise environments, we would be happy to discuss it in the comments.