Julia and compilation

[hartytp]

#102 introduces an OBE solver, using Julia as a back end.

An issue with Julia is that it is a bit slow to startup due to not being a compiled language. On my laptop, running the bloch equations test takes ~20s the first time I run it, and about 50ms subsequently. The difference here is essentially all compilation. Even just importing JuliaCall takes around 5s on my laptop.

I've explored using PackageCompiler.jl to speed up startup by making a system image. I'm using this issue as a place to dump thoughts while I decide what to do.

If this turns out to be useful, I imagine doing something like adding a poe trace_jl task to automate generating a trace.

Precompilation misses

To get better visibility about where recompilation misses are happening, I added an additional option into JuliaCall's __init__.py CONFIG['opt_trace'] = path_option('trace_compile', envkey='PYTHON_JULIACALL_TRACE')[0]. NB there is an upstream issue to add support for this, but it hasn't been implemented yet, so I just hacked this locally for now. If this turns out to be useful I'll submit a PR upstream.

With this I then run

import os
os.environ['PYTHON_JULIACALL_TRACE'] = 'bloch_trace.jl'
from atomic_physics.tests import test_bloch

To generate a trace.

Precompiling QuantumOptics.jl

With the trace we can then create a system image as follows:

from juliacall import Main as jl
jl.seval("using PackageCompiler")
jl.seval("using QuantumOptics")
jl.seval('PackageCompiler.create_sysimage(["QuantumOptics"]; sysimage_path="bloch.so", precompile_statements_file="bloch_trace.jl")')

We can then use the system image as follows:

import os
os.environ['PYTHON_JULIACALL_SYSIMAGE'] = 'bloch.so'

from atomic_physics.tests import test_bloch

Using the system image roughly halves the startup time for the test from ~20s to ~10s. The remaining time is essentially all associated with importing julicall or the various bits of wrapper code it uses.

Precompiling JuliaCall

This does not seem to be supported. See for example JuliaPy/PythonCall.jl#436

As a minimal reproduction:

from julicall import Main as jl
jl.seval("using PythonCall")
jl.seval("using PackageCompiler")
jl.seval('PackageCompiler.create_sysimage(["PythonCall"]; sysimage_path="python_call.so")')

Then running

import os
os.environ['PYTHON_JULIACALL_SYSIMAGE'] = 'python_call.so'

from julicall import Main as jl

That crashes out with

fatal: error thrown and no exception handler available.
InitError(mod=:C, error=ArgumentError(msg="Package CondaPkg does not have pixi_jll in its dependencies:\n- You may have a partially installed environment. Try `Pkg.instantiate()`\n  to ensure all packages in the environment are installed.\n- Or, if you have CondaPkg checked out for development and have\n  added pixi_jll as a dependency but haven't updated your primary\n  environment's manifest file, try `Pkg.resolve()`.\n- Otherwise you may need to report an issue with CondaPkg"))
macro expansion at .\loading.jl:2299 [inlined]
macro expansion at .\lock.jl:273 [inlined]
__require at .\loading.jl:2271
...

That's a shame. I'm not sure if there is an easy way around this.