How to make modeling tools more portable

1

I recently worked on making EpiSewer a bit more portable and wanted to share what came out of it - this is especially focused on deploying stan models in different environments, but I think much of it also applies to other PPLs (whenever they are not widely available on end-user systems). I will talk about two things, containerization and pipeline orchestrators, and I’m curious to hear about others’ experiences and approaches.

Containerizing your statistical model

Anyone who has tried to distribute their stan model to a wider user base probably came across some kind of problem with setting up rstan or cmdstan. For example, installing the cmdstan cli and C++ toolchain can be an obstacle for users on HPCs or other managed environments, as is common in public health settings.

A potential solution to avoid such problems is containerization. This requires the user to have a software like docker installed, but the (probably justified) assumption is that installing docker is more widely supported than installing cmdstan. So we can create a container image for our modeling tool that can be shipped to basically any system that runs docker.

A question of scope

An important question arising at this point is what to put into the container. There are at least three options:

  1. General-purpose image: there already exist a number of “stan” docker images (e.g. here and here) that ship cmdstan and R, together with packages such as cmdstanr, shinystan and more in a container. However, these images are quite heavy: they often provide a full development environment and are not optimized for deployment of already existing models. Importantly, they might also be missing certain dependencies required by your package, so you typically need to add them on top of the image.
  2. Package image: you place only the dependencies required for your R package + the stan tooling in the container. This can avoid some unnecessary software shipped with the “general purpose” dev images described above. Also, you really have everything needed to run your package in one place. However, this approach also comes with limitations. First, some user may struggle to interface with such a container - e.g. if I’m used to working in RStudio, and suddenly have to either work with the R command line inside the container or learn how to connect to it via vs code. Second, it can be hard to foresee what additional software users require for their modeling work. Even a single additional R package requires extending and rebuilding the docker image, which is something most users will not be willing and able to do.
  3. Model image: you place only your stan model in the container and leave the rest of your package outside. This leads to a setup where your R package runs as normal in the user environment, and the container only serves as a “backend” for model fitting. The advantage of this approach is loose coupling - the user environment can be fully customized (and even the parts of your own R package that don’t influence the model) without adjusting the model image. Moreover, the image can be much leaner in principle, as it only needs your model and the sampler/optimizer. The disadvantage is that it requires a suitable environment for the rest of your package (e.g. user needs to have a compatible version of R) and there is added complexity to provide an interface between your package and the container backend.

My attempt at a “lean” model image

I recently tried to implement the third approach (a model image) for EpiSewer using docker. I didn’t manage to go all the way in terms of minimizing the size of the image (see below), but the general architecture seems to work.

My main requirements were:

  • I don’t need functionality to compile new models (just need a container that can fit the stan model of the corresponding package version)
  • I want my model to run out-of-the-box without recompilation in the container
  • I don’t want dependencies on further pre- and postprocessing tools/packages, as these will vary across users and use cases

You can see the resulting Dockerfile here. I tried to make the image as small as possible. This involves a multi-stage build which first installs cmdstan and compiles the model within the container, then starts completely from scratch, copying only the compiled model and some dummy files into the final image (cmdstan not included anymore). In the ideal world, this would be all we need. However, I still wanted to use cmdstanr as an interface, and unfortunately it is not yet able to call models inside a docker container (this would be awesome!). So I ended up including cmdstanr in the container as well (therefore also requiring R, achieved by building on an image like rocker/r-ver).

This resulted in the following setup (as shown in the EpiSewer docker vignette):

  • user runs preprocessing and model specification in their normal R environment (with the EpiSewer R package and any other packages installed)
  • user specifies docker as a backend via model_stan_opts(use_docker = TRUE)
  • EpiSewer translates input data and model specification into cmdstanr arguments (data, inits, and sampling parameters) and writes them to a temporary job file
  • EpiSewer runs docker container with bind-mounts for the job file and some output directory. The container runs a script that
    • i) reads in cmdstanr arguments from the job file
    • ii) fits the precompiled stan model
    • iii) writes the fitted model to the output directory
  • EpiSewer reads the fitted model in R and runs any postprocessing steps
  • user can inspect and further process results in their normal R environment

Pros and cons of the model image approach

As explained above, the “model image” strategy has a couple of nice properties:

  • Everything except the model fitting runs outside the container and can thus be tailored to the user’s needs without creating a new image. Normally, this will be in the user’s typical R environment, but I will discuss a more reproducible setup in the next section.
  • As long as you make no changes to your stan model, your R package can be updated without building a new image. For EpiSewer, I set up a CI workflow to build the image that is only triggered upon changes to one of the stan files.
  • No overhead for compiling the stan model and no compilation conflicts due to mismatches in cmdstan versions etc.

There are also some disadvantages:

  • Fitting models requires the docker daemon to be running on the system
  • One-time overhead for pulling the docker image from github or docker hub
  • Slight overhead when running models in a container (I haven’t quantified this, but it does not seem to be enormous) + I/O overhead for writing temporary files
  • Risk of mismatches between your R package version and the docker image version. This can however be mitigated using tags/digests: For example, the EpiSewer R package will always pull a specific version of the docker image linked to the right stan model version (this can all be managed through CI). In addition, I added a check that compares the hashes of the stan model in the package vs. inside the container.

Using containerized models in inference pipelines

Modern pipeline orchestration software comes with support for containers and package management. This aligns nicely with the above setup. For example, I recently built a Snakemake pipeline for real-time wastewater surveillance in Switzerland that fits EpiSewer models to the newest gene concentration measurements across different pathogens (you can find a blueprint of the core pipeline here).

The pipeline looks like this:

  • Rules for data preprocessing and model specification run in a conda environment specified in a workflow-specific yaml file (including EpiSewer and other packages from the tidyverse).
  • As one of the rule’s output, the resulting EpiSewer job file is written to a temp/scratch directory.
  • There is special rule for model fitting, with a container directive pointing to our stan model image on docker hub. Snakemake automatically pulls the container and runs it with the required inputs and outputs as bind mounts. The rule executes the model fitting script in the container.
  • Rules for postprocessing can use the fitted model file as input and are again executed in a workflow-specific conda environment.

Overall, this gives quite a reproducible and maintainable setup. I have previously been using the targets package in R for pipeline orchestration - which has some great features and is very strong at caching - but the natural support for containers in Snakemake is a considerable advantage in my opinion. Particularly because Snakemake also runs smoothly in HPC environments (using a job scheduler like slurm). Importantly, if conda cannot be installed on the system, there is also the option to use another container for the pre- and postprocessing steps (either using a custom image with all required R packages installed, or by letting snakemake run conda inside a container, see here).

A question of scope, again

For a fixed inference pipeline, I am actually not sure if the “model image” approach I followed for EpiSewer is really the best one. Typically, pipelines are more static and not used in an interactive fashion, so it might be fine to commit to a fixed set of packages that gets updated only once in a while. Moreover, as orchestrators like Snakemake support the use of different environments across steps/rules, the issue of “user-specific” software requirements is practically solved - users can just add further steps/rules with a custom environment or container. So having your full package and model combined in a single docker image should be fine - and is probably more elegant - in this setting.

Final remarks

I would be curious to hear your opinions on model containerization and pipelines. Also, I’m grateful if someone wants to test the EpiSewer vignette with the docker backend or the snakemake pipeline on their system - there are probably failure modes I haven’t foreseen…

Overall, I believe that shifting system requirements towards more generic and widely used frameworks like docker or Snakemake can improve the portability of modeling tools - even if it’s just a question of support (i.e. IT departments more likely helping with docker problems than with installing cmdstan and C++ toolchains). But again, I am curious what others think: (when) is it worth the effort?