Load time of MTK models

General Usage Performance
, , , ,
1

I am not happy with the TTFX of our MTK models (see: RamAirKites.jl). We already cache our simplified model after building it and use a custom system image. But nevertheless, the time to load a model from our cache is still about 35s, even though a second run in the same REPL takes only 5.5s.

This is what AI explains:

  • ModelingToolkit.build_function doesn’t produce an ordinary Julia method. It produces a RuntimeGeneratedFunction{argnames, cache_tag, context_tag, id, B} — a wrapper whose body (an Expr) is stashed in a Dict (_cache_body, line ~226) that lives inside a dynamically-created module, and whose actual calls go through a @generated function generated_callfunc_body(...) indirection (line 194).
  • The first time that generated function is called with a given id, Julia has to run type inference + LLVM codegen on it, same as any other freshly-seen method — that’s real JIT work, tied to this specific model’s RHS/Jacobian/getter expressions (129 pts, 108 segs, 6 twist groups, 4 winches).
  • Once compiled, the resulting native code sits in the process’s ordinary method-instance cache for the rest of that Julia session. That’s why re-running init! (even on a fresh sam, since the disk deserialize repopulates the same ids) in the same REPL only costs 5.5 s — you’re paying deserialize + object reconstruction, but the generated_callfunc_body methods for those ids are already compiled in this process, so no JIT.
  • A sysimage is built by PackageCompiler walking the precompile caches of the packages passed to create_sysimage (:StableAWEModels, :VortexStepMethod, …) — i.e. methods that exist statically in those packages’ module method tables at using time. The RuntimeGeneratedFunction methods created by eval while test_for_precompile.jl ran during the sysimage build live in a transient module/id created in that build process. They were real, compiled native code — but only inside that now-dead build process. PackageCompiler’s snapshot doesn’t walk “every method that got JIT-compiled while the precompile script executed,” only the target packages’ own precompiled method tables, so this per-model, eval-minted code never makes it into the .so.
    So every fresh julia -J bin/kps-image-1.12.so process starts with an empty method-instance cache for this model’s RuntimeGeneratedFunction ids and has to redo that JIT work once — that’s your 38 s. Inside that same process, it’s free from then on — that’s your 5.5 s. This is a known SciML/ModelingToolkit + PackageCompiler limitation (the whole point of the RuntimeGeneratedFunction indirection is to avoid invalidating a huge sysimage’s method cache every time a user builds a new symbolic model — the tradeoff is that the model-specific code can never be pre-baked into that sysimage).

Is there no workaround to include this compiled code in the system image for production use?

2

No idea?

Where should I create an issue:

  • in MTK
  • in Julia
  • in PackageCompiler

?

This is an issue that could technically be solved, but I don’t know where. The issue is that generated functions are not saved in the system image generated by package compiler.

3

This is something we’ve been working on. But examples of cases where this is going bad are always welcome. I believe MTK is the right place

4

I will create an issue and provide an example.

5

I am now wanting to use a package compiler on one of my MTK models so user can use it without startup time etc.

Gemini is claiming it can be done with the eval_expression=true kwargs.
I don’t know enough about how the whole code generation part of MTK works to say whether this is hallucination or not – any opinions?

edit: it seems to be based off this page of MTK documentation

"

When combining PackageCompiler.jl with ModelingToolkit.jl (MTK) models, you must account for how MTK generates functions at runtime. [1, 2]

Because MTK relies heavily on runtime evaluation and RuntimeGeneratedFunctions to safely build expressions without invalidating caches, simply bundling MTK into a basic system image will not eliminate your first-solve compilation overhead. [1, 2]

To successfully precompile and deploy MTK models using PackageCompiler, you must explicitly opt out of runtime expressions and capture the fully native code. [1, 2]

1. Configure the MTK Model for Precompilation

By default, ODEProblem creates code using eval statements that live purely in short-lived processes. To bypass this and make your model ahead-of-time (AOT) compileable, you must set eval_expression = true and target your specific module. [1, 2]

Wrap your model definition inside a dedicated module:

julia

# ModelModule.jl
module ModelModule
    using ModelingToolkit
    using OrdinaryDiffEq

    # Define a simple system
    @variables x(ModelingToolkit.t_nounits)
    @named sys = System([ModelingToolkit.D_nounits(x) ~ -x + 1], ModelingToolkit.t_nounits)
    
    # CRITICAL: eval_expression=true and eval_module=@__MODULE__
    # This forces MTK to output static code that can be cached.
    const compiled_sys = mtkcompile(sys)
    const prob = ODEProblem(compiled_sys, [x => 30.0], (0.0, 100.0), [], 
                            eval_expression = true, 
                            eval_module = @__MODULE__)

    # Precompile trigger function
    function run_model()
        sol = solve(prob, Tsit5())
        return sol[x][end]
    end
end

Use code with caution.

2. Create a Precompilation Script

PackageCompiler needs a script to trace execution and record the native machine instructions into the new system image (sysimage). [1, 2, 3]

julia

# precompile_script.jl
using ModelingToolkit
using OrdinaryDiffEq
include("ModelModule.jl")

# Force execution so Julia saves the compiled native methods
ModelModule.run_model()

Use code with caution.

3. Generate the Custom Sysimage

Run the compilation task from the standard Julia REPL. This bundles the core packages and your staticized model into a single shared library binary (.so, .dll, or .dylib). [1, 2]

julia

using PackageCompiler

create_sysimage(
    [:ModelingToolkit, :OrdinaryDiffEq]; 
    sysimage_path="MtkSysimage.so", 
    precompile_execution_file="precompile_script.jl"
)

Use code with caution.

4. Run with the Compiled Sysimage

To run your model with near-zero startup and first-solve latency, execute your production scripts by pointing Julia to the compiled binary bundle: [1]

bash

julia -J MtkSysimage.so -e 'include("ModelModule.jl"); println(ModelModule.run_model())'

"