pavel.pink

Managing Agent Skills with Your Package Manager

pavelzw@gmail.com (Pavel Zwerschke) — Mon, 09 Feb 2026 00:00:00 GMT

AI coding agents like Claude Code can be extended with skills, markdown files that contain domain-specific instructions for how to approach certain tasks. A skill for polars, for example, might teach the agent to prefer lazy evaluation, use collect sparingly, and follow specific patterns for joins.

Skills are just text files. But as your team grows and you accumulate more of them, managing these text files becomes a real problem.

The state of the art

One of the current approaches is skills.sh, a community registry that lets you fetch skills with a single command:

npx skills add https://github.com/davila7/claude-code-templates --skill polars

This downloads the skill markdown file into your project and you’re done. Simple and effective for getting started.

What’s missing

As soon as skills become part of real workflows, across projects, teams, or corporate environments, a few problems emerge:

Markdown sprawl in git. Skills are vendored into your repository. Every project gets its own copy. Updates mean committing diffs of markdown files you didn’t write. Your git history fills up with skill updates that have nothing to do with your project.
No version pinning. npx skills add fetches from main. There’s no lockfile, no version constraint. A skill that worked yesterday might behave differently today because someone pushed an update upstream. If a skill is updated to target a new major version of a library, your project breaks silently and the agent just starts writing code for an API you’re not using.
Enterprise environments. In locked-down corporate networks, access to github.com may be restricted or unavailable. Getting skills into these environments requires manual workarounds.
Supply chain security. Skills are instructions an AI agent follows. A malicious skill could subtly steer an agent toward specific commercial services, leak information through generated code, or introduce unsafe patterns. You might install a polars skill that injects adversarial context: the K-Dense-AI/claude-scientific-skills polars skill does just that! Skills can also spread hallucinated package names: an LLM-generated skill referenced a non-existent npm package, got copy-pasted across 200+ repositories, and agents started trying to install it, creating a supply chain attack vector. A mature packaging ecosystem is a possible way to mitigate some of these problems. It might offer well-defined release processes, auditability and change history to be able to track changes and ensure integrity.

These are not hypothetical problems. These are exactly the problems software package management solved decades ago.

Skills as packages

Turns out: skills are text files. Software packages exist to distribute text files. Package managers already handle versioning, locking, distribution, and security. So why not ship skills as packages?

We use pixi for dependency management across our projects. Since pixi is built on the conda ecosystem and is language-agnostic, it can manage any kind of package, including ones that just contain markdown files.

Creating a skill package

A conda package for a skill is straightforward. Here’s a recipe.yaml for a polars skill:

schema_version: 1

package:
  name: agent-skill-polars
  version: "0.2.0"

build:
  noarch: generic
  script:
    - mkdir -p $PREFIX/share/agent-skills/polars
    - cp $RECIPE_DIR/SKILL.md $PREFIX/share/agent-skills/polars/
    - cp -R $RECIPE_DIR/references $PREFIX/share/agent-skills/polars/

requirements:
  run_constraints:
    - polars >=1.38.0,<2.0

about:
  summary: Agent skill for writing polars code
  license: MIT

$CONDA_PREFIX/share/agent-skills/polars
├── references
│   ├── best_practices.md
│   ├── ...
│   └── transformations
└── SKILL.md

The noarch: generic build means this package is platform-independent — it’s just markdown files. The skill files get installed to share/agent-skills/ in the conda prefix.

The interesting part is run_constraints. This tells the solver: if polars is installed in this environment, it must be >=1.38.0,<2.0. The skill package doesn’t pull in polars as a dependency, but it ensures that the skill matches the version of the library you’re actually using. If you’re on polars 0.x, the solver will refuse to install this skill: you’d need agent-skill-polars of a different version targeting the older API instead.

Do it yourself

You can ship your own skills! Here is a prompt for your favorite agent that can be used to create a skill.

I would like to create a skill for scikit-learn. Do it in the same way as described in https://github.com/pavelzw/skill-forge/blob/main/AGENTS.md.

Installing skills

Adding skills to a project is now just a pixi dependency:

[workspace]
channels = ["conda-forge", "https://prefix.dev/skill-forge"]
platforms = ["linux-64", "osx-arm64", "win-64"]

[dependencies]
polars = ">=1,<2"
agent-skill-polars = "*"

pixi install

The skill version is pinned in pixi.lock alongside all your other dependencies. It is reproducible across machines and across time.

We maintain a set of pre-packaged skills on the conda channel prefix.dev/skill-forge. As mentioned above, skills can lead to arbitrary code execution, so please always read through them before using them.

`pixi-skills`: connecting skills to your agent

Installed packages end up in .pixi/envs/default/. My friend Moritz and I wrote a small tool (pavelzw/pixi-skills) that can symlink the skill files from there to wherever your agent expects them:

pixi exec pixi-skills manage

This scans the installed packages for files under share/agent-skills/ and symlinks them into your project’s skill directory. The tool is trivial because the heavy lifting is done by the package manager.

What this gets us

No markdown bloat in git. Skills are installed dependencies, not vendored files. Your .gitignore covers the .pixi/ directory and skill symlinks, so you stop committing other people’s markdown.
Versions are pinned. pixi.lock pins the exact skill version, just like any other dependency. Updates are intentional — you run pixi update agent-skill-polars when you’re ready. And with run_constraints, a skill can express compatibility with the library it targets.
Enterprise distribution works. You probably already have infrastructure for distributing packages internally, like Artifactory, an S3-backed conda channel, or a conda-mirror. Skills move through the same channels as your other packages, with no special cases.
Your existing supply chain security applies. Package signing¹, review workflows, vulnerability scanning, access controls — whatever you’ve built for managing software packages now applies to skills too. You can audit a skill before publishing it to your internal channel, and restrict which skill packages are approved for use.

Drawbacks

Since your skills are no longer tracked in Git, you cannot manually edit them and expect that these changes persist across devices. If you want to edit a skill, you would need to update the package underneath it. With a solid packaging flow, this shouldn’t be a problem, though.

The skills only become available after cloning once you run pixi install. For local development (for example via the Claude Code CLI), this is not an issue since you are running your code anyway. For remote development, it depends on the specific setup and infrastructure you are using. At the time of writing, in the Codex desktop app, you can specify custom setup scripts where you could add pixi install. When using Codex with multiple worktrees locally, the skills get picked up automatically. Strangely, this is not the case with the Codex web version (and mobile app). There, it doesn’t find the skills after running pixi install.

The Claude Code web application (and mobile app) currently do not support this workflow. There, you can create environments like in Codex but unfortunately, it’s only possible to create environment variables and not run custom setup scripts. You could instruct your agent via an AGENTS.md file to run pixi install and look for skills in those locations.

Outlook

Everything described here is pixi-specific, but the core idea applies to any package manager. pip or uv could install skill packages into .venv/share/agent-skills/. bun or npm could write them to node_modules/. The details change, the principle does not.

Skills are text files. Managing text files at scale is a solved problem. The practices we built over decades for shipping software like versioning, lockfiles, registries, signing, and reviews, map directly onto distributing agent instructions.

Package signing is not yet fully supported by the Conda ecosystem but in the works. ↩

Deploying JetBrains Gateway in a Corporate World

pavelzw@gmail.com (Pavel Zwerschke) — Sun, 10 Sep 2023 00:00:00 GMT

When you have worked in a corporate environment, you know that deploying stuff can be challenging. During my work, I deployed JetBrains Gateway in a corporate environment. In this article, I will describe how I did it and what I had to consider.

What is JetBrains Gateway?

JetBrains Gateway is a piece of software that allows you to remotely access your JetBrains IDEs, similarly to how you use VSCode Remote Development. With this, you can run your favorite JetBrains IDEs like IntelliJ IDEA or PyCharm in ✨ ☁️ the cloud ☁️ ✨.

How it usually works

In an open environment that you have control over, the installation of JetBrains Gateway is pretty straightforward.

install the Gateway Client on your local machine
add an SSH connection to your server that should run the IDE
click on Check Connection and Continue
select the IDE you want to install as well as the remote project directory
click Download IDE and Connect
and you’re done! 🎉

Taking a closer look

What happens under the hood when you click on Download IDE and Connect is that the Gateway application on your local machine tries to connect via SSH to the remote server and then downloads your selected IDE from JetBrains’ servers. This IDE is actually the exact same version that you would download as an end-user from JetBrains’ website when you want to install PyCharm¹ or IntelliJ IDEA² on your local machine.

If you look closely into the downloaded file, you will see that there is a bin directory that contains a remote-dev-server.sh script. This script is (more or less) what’s executed on the remote server after your IDE is downloaded.

Corporate setup

In a corporate environment, you usually have to deal with proxies and firewalls. In my case, I had the situation that in the environment to deploy Gateway in was a docker container which I only had SSH access to. The docker container itself had no internet access, but I could access the internet from the host machine.

So I had to download the IDE from my local machine and then copy it into the remote server.

scp pycharm-professional-2023.2.1.tar.gz remote-server:~/
ssh remote-server
tar xzf pycharm-professional-2023.2.1.tar.gz -C /opt

Additionally, the only directory that was persisted in the docker container was my home directory which was mounted from a network share. PyCharm (and probably the other JetBrains IDEs too) requires a persistent directory to store its configuration and caches. The configuration is stored in in ~/.config/JetBrains by default and the caches are stored in ~/.cache/JetBrains. Unfortunately, PyCharm does not like working on a network share at all³, so I had to find a way to store the configuration and caches in a local directory instead. Creating a symlink from the network share to a local directory fixed the issue.

ln -s /tmp/jetbrains-cache ~/.cache/JetBrains

Since these are only cache files, we also won’t lose any meaningful data if the container is restarted.

So when you run this script, you will see some logging and after some time, you will see the following message:

$ REMOTE_DEV_NON_INTERACTIVE=1 /opt/pycharm-2023.2.1/bin/remote-dev-server.sh run \
    "${project_dir}"
    --ssh-link-port 2222 \
    --ssh-link-host remote-server \
    --ssh-link-user remote-user \
    -l 0.0.0.0

...
*********************************************************

Join link: tcp://0.0.0.0:5990#jt=...

Http link: https://code-with-me.jetbrains.com/remoteDev#idePath=%2Fopt%2Fpycharm-2023.2.1&projectPath=%2Fhome%2Fuser%2Fproject&host=remote-server&port=22&user=user&type=ssh&deploy=false&newUi=true

Gateway link: jetbrains-gateway://connect#idePath=%2Fopt%2Fpycharm-2023.2.1&projectPath=%2Fhome%2Fuser%2Fproject&host=remote-server&port=22&user=user&type=ssh&deploy=false&newUi=true

*********************************************************
...

You can now either open the http link in your browser which will open the JetBrains Gateway app for you or you can use the Gateway link and paste it into the Gateway app directly⁴.

And you’re done! 🚀

Plugins and Drivers

What makes the JetBrains IDEs so powerful is the huge ecosystem of drivers and plugins. These drivers and plugins are usually downloaded from the internet when you want to install them. In my case, I had no internet access, so I had to download them manually and copy them into the container.

Drivers

Drivers are usually project-specific and go into ~/.config/JetBrains. For example, when I want to download the Snowflake driver for a project in ~/project, the snowflake-jdbc-3.13.27.jar driver would need to go into

~/.config/JetBrains/RemoteDev-PY/_home_user_project/net/snowflake/snowflake-jdbc/3.13.27/snowflake-jdbc-3.13.27.jar

Which jar exactly is needed depends on the driver you want to install. You can find the exact name of the jar from the IDE error logs. These drivers are usually downloaded from maven central. If your corporate environment has a maven mirror like Nexus or Artifactory, you could maybe even download the driver from there instead. You can also provide your IDE a custom JAR by clicking on the plus sign in the “Driver Files” section.

Plugins

Plugins can be downloaded manually from the JetBrains Website. For example, if you want to install the GitToolBox plugin, you can just search for it on the website and download it for your IDE version. The downloaded zip needs to be extracted to ~/.config/JetBrains/RemoteDev-PY/_home_user_project/plugins/gittoolbox/.

You need to differentiate between host plugins and client plugins. Host plugins are installed on the remote server and client plugins can be installed on the local machine (mostly theme-related plugins). For example, GitToolBox is a host plugin, but IdeaVim is a client plugin.

Putting it all together

I’d like to end with a short summary of the steps that are needed to deploy JetBrains Gateway in a corporate world.

install the Gateway Client on your local machine
download your ide from JetBrains’ website
scp <ide>-<version>.tar.gz remote-server:~/
extract the tarball to your installation directory
ln -s /tmp/jetbrains-cache ~/.cache/JetBrains (if $HOME is a network share)
add drivers and plugins to ~/.config/JetBrains/RemoteDev-PY/_home_user_project/⁵ (if needed)
if you need a proxy on your local machine, configure it in the Gateway app
run the remote-dev-server.sh script on the remote server
???
profit!!! 💸

jetbrains.com/pycharm/download?section=linux or download.jetbrains.com/product?code=PY&latest&distribution=linux ↩
jetbrains.com/idea/download/?section=linux or download.jetbrains.com/product?code=IU&latest&distribution=linux ↩
See issue GTW-5697 on YouTrack. ↩
In corporate environments, it is usually the case that you cannot access the internet from your local machine directly but through a proxy. You may need to configure the proxy settings in the Gateway app. px-proxy might come in handy if you use NTLM authentication since JetBrains products cannot handle this natively. ↩
You need to update the paths accordingly to your project directory. Also, RemoteDev-PY might be different for you depending on the IDE you want to use. ↩

Automatically Transform Complex Python Methods to Polars Expressions

pavelzw@gmail.com (Pavel Zwerschke) — Fri, 25 Aug 2023 00:00:00 GMT

Update 28.09.2023: My colleague Bela held a Lightning Talk at the Big PyData BBQ. You can check it out on YouTube!

In this post, we introduce polarIFy, a Python function decorator that gives you a simpler way to write logical statements for Polars. With polarIFy, you can use Python’s language structures like if / elif / else statements and transform them into pl.when(..).then(..).otherwise(..) statements. This makes your code more readable and less cumbersome to write.

At QuantCo, we frequently deal with insurance contracts where the calculation behavior is contingent upon the exact tariff. For instance, we might calculate certain properties differently for older contracts compared to newer ones.

This introduces a degree of conditional business logic into our data processing pipelines. When working with Polars, these conditional operations often translate into nested pl.when(..).then(..).otherwise(..) statements. While this structure is powerful and allows for a great deal of flexibility, it can quickly become complex and difficult to read, especially when dealing with multiple conditions or branches.

This is where polarIFy comes in. By allowing us to write these conditional statements in a more Pythonic way using if / elif / else structures, polarIFy greatly simplifies our code. It makes our Polars pipelines cleaner, more readable, and less prone to errors, especially when dealing with the intricate conditional logic often required in our work with insurance contracts.

Using polarIFy

polarIFy can automatically transform Python functions using if / elif / else statements into Polars expressions. Here’s an example:

@polarify
def func(x: pl.Expr) -> pl.Expr:
    s = 1
    if x > 10:
        return s + 10
    else:
        t = 2

    if x > 0:
        return t
    else:
        return -s

This gets transformed into:

def func_polarified(x: pl.Expr) -> pl.Expr:
    return (
        pl.when(x > 10)
        .then(1 + 10)
        .otherwise(
            pl.when(x > 0)
                .then(2)
                .otherwise(-1)
        )
    )

polarIFy can also handle multiple statements and nested statements, making it a versatile tool for simplifying your Polars code.

How it works

Let’s take a look at the above example in a control flow graph:

In order to transform this control flow into a Polars expression, we need to keep track of all possible assignments to variables to determine when to return what. We do that by creating a dictionary assignments for each node that maps variable names to their values.

At the beginning, we haven’t assigned any values to any variables, so we start with an empty dictionary: assignments = {}.

In the first step, we assign the value 1 to the variable s, so we update the dictionary: assignments = {'s': 1}.

Going further, we evaluate if x > 10. This translates to the polars expression pl.when(x > 10).then(..).otherwise(..). We need to keep track recursively of all possible assignments to variables and the return value of the then and else branches of our control flow.

Since the then branch directly returns s + 10, we can just put a s + 10 in our polars expression: pl.when(x > 10).then(assignments['s'] + 10).otherwise(..).

Now, let’s look at the else branch. We assign the value 2 to the variable t, so we create a new dictionary from the old one with our new assignment: assignments = {'s': 1, 't': 2}.

We now have another if statement which gets translated into another pl.when(x > 0).then(..).otherwise(..) expression.

The then branch returns the value of t, which we can get from our assignments dictionary. Thus, we can put assignments['t'] into our polars expression: pl.when(x > 0).then(assignments['t']).otherwise(..). The else branch returns -s, which we can also get from our assignments dictionary: pl.when(x > 0).then(..).otherwise(-assignments['s']).

All in all, we land at the following polars expression:

(
    pl.when(x > 10)
    .then(1 + 10)
    .otherwise(
        pl.when(x > 0)
            .then(2)
            .otherwise(-1)
    )
)

Limitations

Since polarIFy transforms Python functions into Polars expressions, it can only handle logic that Polars itself can also do. For instance, it can’t handle for and while loops, since it is not possible to do loops for a single element with Polars expressions. Also, functions with side effects (like print, raise or pl.write_csv) are not supported since these don’t make sense in a Polars expression context.

Since polarIFy is a relatively new project, it might not work for all use cases yet. As of August 2023, there are some known limitations:

match ... case statements are not supported
the walrus operator := is not supported

Conclusion

polarIFy is a powerful tool that simplifies the process of writing conditional statements in Polars. By allowing you to write conditions in a more Pythonic way, it makes your code cleaner and easier to understand. We’re excited about its potential to make working with Polars even more efficient and enjoyable.

We’re always looking for feedback and contributions, so feel free to check out the polarIFy GitHub repository and let us know what you think!

This is a cross-post from the QuantCo blog. Check out the other posts there!

Optimize Pickling Disk Space for Deploying scikit-learn Trees to Production

pavelzw@gmail.com (Pavel Zwerschke) — Thu, 23 Jun 2022 00:00:00 GMT

Update 04/2023: We have created a package that implements the ideas presented in this article and presented it at PyData Berlin 2023. Feel free to check slim-trees out on GitHub.

At QuantCo, we use ML models such as random forests to create value from data for our customers. To ship these models to production, we have to include them in the application, e.g., bake them into a Docker container or fetch them from a model registry. To save these models, we pickle them.

If we want to improve the model, we sometimes increase the number of parameters. This comes at the cost of disk space because the model is getting more complex. In our case, we were not able to deploy a larger model because of a file size constraint. Nevertheless, in general it makes sense to strive for smaller Docker images to achieve quicker transfer and deployment.

One way to solve this problem would be to use a different format altogether, like ONNX. Instead, we challenged the scikit-learn pickle format to see if there are any gains to be had. We found a 4x improvement over the native scikit-learn format.

Here are the numbers:

dtype compression	compression algorithm	size	dump time	load time
no	no	496 MB	1.4 s	0.7 s
yes	no	128 MB	1.6 s	0.9 s
no	lzma	144 MB	217 s	8.4 s
yes	lzma	32 MB	53 s	2.8 s

In this article, we present how we did it.

How a `Tree` is saved

Usually, when pickling an object, the methods __reduce__ and __setstate__ get called to serialize the object. The following code illustrates some of what is happening when calling pickle.dump(obj).

def pickle_obj(obj):
    cls, init_args, state = obj.__reduce__()
    return unpickle_obj, (cls, init_args, state)


def unpickle_obj(cls, init_args, state):
    obj = cls(*init_args)
    obj.__setstate__(state)
    return obj


myobj = ...
unpickle_func, myobj_data = pickle_obj(myobj)
save_to_file("myobj.pkl", unpickle_func, myobj_data)
...
unpickle_func, myobj_data = load_from_file("myobj.pkl")
myobj_unpickled = unpickle_func(*myobj_data)

Let’s look at the structure of state when we call tree.__reduce__().

{
    'max_depth': 19,
    'node_count': 469,
    'nodes': np.array(...,
        shape=(469,),
        dtype=np.dtype([
            ('left_child', '<i8'),
            ('right_child', '<i8'),
            ('feature', '<i8'),
            ('threshold', '<f8'),
            ('impurity', '<f8'),
            ('n_node_samples', '<i8'),
            ('weighted_n_node_samples', '<f8')])),
    'values': np.array(...,
        shape=(469, 1, 1), dtype=float64)
}

Since max_depth and node_count are only single integers, they can be neglected when trying to minimize the disk space. nodes is a NumPy data type object, each element of the array is actually a tuple containing the indices of the children, the used feature, etc. per node. This has a lot of potential for savings. value is a separate array because it could be a multidimensional output too which doesn’t fit in the data type object structure of nodes.

Returning to the state of our Tree, we can see that all the elements of nodes are either of type <i8 or <f8, i.e., int64 or float64. This takes up a lot of disk space. To mitigate this problem, we could simply change how these parameters of the model are stored to make the serialization more efficient.

How can we implement changing all float64 and int64 numbers to more efficient datatypes? To do this, we can customize the behaviour of scikit-learn when pickling a Tree.

Custom picklers

We will be creating a custom pickler whose behavior varies when pickling a Tree object. Thus, every instance of a Tree in our model structure will get pickled in a custom manner while the rest of the pickling behavior stays the same.

To achieve this, we need to customize the dispatch table of said pickler.

def dump_dtype_reduction(model: Any, file: BinaryIO):
    p = pickle.Pickler(file)
    p.dispatch_table = copyreg.dispatch_table.copy()
    p.dispatch_table[Tree] = compressed_tree_pickle
    p.dump(model)

The method compressed_tree_pickle is supposed to replace the method pickle_obj from above.

Since we want to reduce the size of a Tree’s state, we want to do something like this.

def compressed_tree_pickle(tree):
    assert isinstance(tree, Tree)
    cls, init_args, state = tree.__reduce__()
    compressed_state = compress_tree_state(state)
    return compressed_tree_unpickle, (cls, init_args, compressed_state)


def compressed_tree_unpickle(cls, init_args, state):
    tree = cls(*init_args)
    decompressed_state = decompress_tree_state(state)
    tree.__setstate__(decompressed_state)
    return tree

Now, all that’s left is to implement compress_tree_state and decompress_tree_state to a more efficient format.

Compressing a `Tree` state

We can see that left_child and right_child are stored as <i8, i.e., int64 numbers. But since there are not that many leaves in a tree, the possible values of the array containing the children of each node are relatively small. Thus, we can change the dtype of the array to int16 for example. The same can be done with the feature array. This attribute stores the index of the features to be used in the node. The value array contains the actual output values of the leaves as float64numbers. With a bit of loss of precision, we could use float32 here.

There are now four remaining attributes of state: threshold, impurity, n_node_samples, and weighted_node_samples. The last three are all purely analytical and not used during inference. Thus, we can just leave them out in our custom pickling solution.

Now let’s look at threshold, i.e., the splitting points at each node, whether to go to the left or right child in the tree. We could simply do the same as with the value array and call it a day. But unlike the value array, we need to have in mind that error propagation can occur. If a “wrong” turn is made in the beginning, we could land at a completely different output node than we should.

There are three ways to continue:

accept the extra loss and continue;
use float64 again in the threshold array and don’t gain disk space in this array;
examine the structure of the threshold array and try to find a way to do lossless compression.

Let’s try out the latter option.

Compressing half ints

The following only works if you have many categorical columns in your dataset. If you don’t, you may want to go with 1. or 2. or look into other methods to do lossless compression.

When we looked at the structure of the threshold array, we realized that there are a lot of values that are of the form x.5, i.e., 3.5, 0.5, 9.5, etc. Let’s call them half ints. This is because our dataset has a lot of categorical values which are converted to int and then cast to float64 for fitting by scikit-learn. Since every point still lies on an integer, the optimal splitting point is in the middle of two integers. These optimal splitting points are exactly what the sklearn trees learn. Now what can we do with this information?

We could compress these half ints by multiplying them with 2 and casting them to int8.

Because of floating point precision, we need to make sure not only to catch the exact integers but an epsilon neighborhood of the integer values.

a2 = np.array([1.9999999999999, 2.0000000000001])
np.minimum(np.abs(a2 % 1 - 1), a2 % 1)
# array([9.99200722e-14, 9.99200722e-14])

def compress_half_int_float_array(a, compression_dtype=np.int8):
    info = np.iinfo(compression_dtype)
    a2 = 2. * a
    is_compressible = (np.minimum(np.abs(a2 % 1 - 1), a2 % 1) < 1e-12) & \
                      (a2 >= info.min) & (a2 <= info.max)
    not_compressible = np.logical_not(is_compressible)

    a2_compressible = a2[is_compressible].astype(compression_dtype)
    a_incompressible = a[not_compressible]

    state = {
        "is_compressible": is_compressible,
        "a2_compressible": a2_compressible,
        "a_incompressible": a_incompressible,
    }

    return state


def decompress_half_int_float_array(state):
    is_compressible = state["is_compressible"]
    a = np.zeros(len(is_compressible), dtype="float64")
    a[is_compressible] = state["a2_compressible"] / 2.
    a[np.logical_not(is_compressible)] = state["a_incompressible"]
    return a

This format is very efficient if you have a lot of these half int values in your array. state["is_compressible"] is a bool array which takes up only one bit per entry, state["a2_compressible"] is an int8 array which only takes up eight bits per entry and state["a_incompressible"] takes up the full 64 bits per entry.

In our productive example, the threshold array contained 4000 elements with 65-70% of them being half integers. This leads to a theoretical compression of

compared to the float32 compression of . So better theoretical compression performance while staying lossless.

Further optimizations

All the arrays that we have talked about have the same length, namely the number of nodes in the tree. But when thinking about it, the leaves don’t really need a left/right child, a corresponding feature or threshold. It only needs the value. That’s why the values in the children_left and children_right fields are -1 there and the value of features and thresholds are -2.

We can just delete these elements and save some space again by creating a bool array is_leaf that is used as a mask for getting and setting the corresponding values.

def compress_tree_state(state: dict):
    ...
    is_leaf = nodes['left_child'] == -1
    is_not_leaf = np.logical_not(is_leaf)

    # feature, threshold and children are irrelevant for leaves
    children_left = nodes['left_child'][is_not_leaf].astype(dtype_child)
    children_right = nodes['right_child'][is_not_leaf].astype(dtype_child)
    features = nodes['feature'][is_not_leaf].astype(dtype_feature)
    # do lossless compression for thresholds by downcasting
    # half ints (e.g. 5.5, 10.5, ...) to dtype_halfint
    thresholds = nodes['threshold'][is_not_leaf].astype(dtype_threshold)
    thresholds = compress_half_int_float_array(thresholds, dtype_halfint)

    # value is irrelevant when node is not a leaf
    values = state['values'][is_leaf].astype(dtype_value)
    ...


def decompress_tree_state(state: dict):
    ...
    n_edges = len(is_leaf)

    children_left = np.empty(n_edges, dtype=np.int64)
    children_right = np.empty(n_edges, dtype=np.int64)
    features = np.empty(n_edges, dtype=np.int64)
    thresholds = np.empty(n_edges, dtype=np.float64)

    # same shape as values but with all edges instead of only the leaves
    values = np.zeros((n_edges, *state['values'].shape[1:]), dtype=np.float64)

    children_left[is_not_leaf] = state['children_left']
    children_left[is_leaf] = -1  # child of leaves is -1
    children_right[is_not_leaf] = state['children_right']
    children_right[is_leaf] = -1  # child of leaves is -1
    features[is_not_leaf] = state['features']
    features[is_leaf] = -2  # feature of leaves is -2
    thresholds[is_not_leaf] = decompress_half_int_float_array(state['thresholds'])
    thresholds[is_leaf] = -2  # threshold of leaves is -2
    values[is_leaf] = state['values']
    ...

Reconstructing the state of the tree

To be able to reconstruct a Tree, we need the exact same state format as in the original, i.e., a dict with max_depth, node_count, nodes and values, all in the same format as before. nodes is actually a NumPy data type object as mentioned before, so we need to specify the structure manually.

def decompress_tree_state(state: dict):
    ...
    dtype = np.dtype([('left_child', '<i8'), ('right_child', '<i8'),
                      ('feature', '<i8'), ('threshold', '<f8'),
                      ('impurity', '<f8'), ('n_node_samples', '<i8'),
                      ('weighted_n_node_samples', '<f8')])
    nodes = np.zeros(n_edges, dtype=dtype)
    nodes['left_child'] = children_left
    nodes['right_child'] = children_right
    nodes['feature'] = features
    nodes['threshold'] = thresholds

    return {'max_depth': state['max_depth'],
            'node_count': state['node_count'],
            'nodes': nodes,
            'values': values}

Compression errors

In our production example, we tested the predictions on a dataset with 3 500 000 entries. The average compression error was at 5.59e-8, there were 8 / 3 500 000 elements with a prediction error > 1e-4 with the largest error being 2.87e-4. When we cumulate all prediction errors, we get an absolute error of 0.19. For our use case, these errors are completely fine.

Discussion

In our production example, we could shrink down the model size from 496 MB to 128 MB, i.e., compress the model to 26% of its original size.

One can also use a compression algorithm like gzip or lzma to compress the model even further. We tried out zstd, gzip, lzma, bz2 and zlib. We opted for lzma because it is part of the Python stdlib and it yielded the best compression results. We managed to go down to 32 MB (without the extra dtype reduction steps and only lzma compression we got to 144 MB, so we still get the model down to 22% of its original “lzma size”).

These extra compression algorithms do bring additional startup time with them, though. You can find the numbers at the beginning of this blog post.

This is although only the size on the disk. When you actually load the model, these steps don’t help on reducing the size of the model in memory. It still has its use cases, for example when providing a docker image that is as small as possible.

# dump compressed
with lzma.open("/path/to/model.pickle.lzma", mode="wb") as f:
    dump_dtype_reduction(model, file)

# load compressed
with lzma.open("/path/to/model.pickle.lzma", mode="rb") as f:
    model_dtype_reduction = pickle.load(file)

Pickling has several disadvantages though like arbitrary code execution and no version compatibility. Uwe Korn did a blog post about the implications of pickling ML models and Ned Batchelder wrote one about pickle’s nine flaws.

Another alternative could be using ONNX to compile the model. This also could shrink down the model and improve the inference speed and the simplification of deployment. ONNX brings its own challenges though: you need create workflows to transform your fitted model to ONNX which is not trivial for more complicated models. Also, you would need to make sure that the converted ONNX model predicts the same things as your original model.

The code for this blog post can be found on GitHub Gist. The whole procedure was tested on scikit-learn=0.23.2, the internal structure of Tree may however change in upcoming versions.

This is a cross-post from the QuantCo blog. Check out the other posts there!