Add a precompile statement to PlutoRunner and some minor refactorings

[rikhuijzer]

Some small changes to make the PlutoRunner slightly more snappy.

Changes:

• Drop Base.@kwdef to make the code simpler and avoid having to compile a kwarg method.
• Fix some type inferability problems by adding type assertions. This is tested against Julia 1.8-beta3.
• Narrow down the type of old_logger a bit. This doesn't help in performance, but probably helps in readability.
• Add a precompile statement for run_expression.

Time to first X

On Linux with Julia 1.8-beta3.

Before (master)

$ julia --project -ie 'using Pluto'
[...]

julia> module Foo end
Main.Foo

julia> @time @eval Pluto.PlutoRunner.run_expression(Foo, :(1 + 1), Pluto.uuid1(), nothing);
[...]
  1.653029 seconds (2.94 M allocations: 152.834 MiB, 13.22% gc time, 99.80% compilation time)

After (this PR)

$ julia --project -ie 'using Pluto'
[...]

julia> module Foo end
Main.Foo

julia> @time @eval Pluto.PlutoRunner.run_expression(Foo, :(1 + 1), Pluto.uuid1(), nothing);
[...]
  1.287273 seconds (1.60 M allocations: 82.489 MiB, 2.78% gc time, 99.70% compilation time)

[github-actions]

Try this Pull Request!

Open Julia and type:

julia> import Pkg
julia> Pkg.activate(temp=true)
julia> Pkg.add(url="https://github.com/fonsp/Pluto.jl", rev="rh/precompile-plutorunner")
julia> using Pluto

[Pangoraw]

@rikhuijzer you should be interested by tooking a look at #1881 where this improvements could be more beneficial. PlutoRunner is currently included() in each notebook process in Loader.jl and I don't think precompile happens in this case?

[Pangoraw]

In which way does adding type assertion help? any should be easy to infer.

[rikhuijzer]

Normally, yes. But expr.args is a Vector{Any}:

julia> ex = :(1 + 1)
:(1 + 1)

julia> f(ex) = ex.args;

julia> @code_warntype f(ex)
MethodInstance for f(::Expr)
  from f(ex) in Main at REPL[15]:1
Arguments
  #self#::Core.Const(f)
  ex::Expr
Body::Vector{Any}
1 ─ %1 = Base.getproperty(ex, :args)::Vector{Any}
└──      return %1

which results in a Union{Missing, Bool} instead of Bool:

julia> g() = any(==("foo"), ex.args)

julia> @code_warntype g()
MethodInstance for g()
  from g() in Main at REPL[17]:1
Arguments
  #self#::Core.Const(g)
Body::Union{Missing, Bool}
1 ─ %1 = (==)("foo")::Base.Fix2{typeof(==), String}
│   %2 = Base.getproperty(Main.ex, :args)::Any
│   %3 = Main.any(%1, %2)::Union{Missing, Bool}
└──      return %3

In general, you're right indeed.

julia h() = any(==("foo"), ["bar"]);

has return type Bool.

[fonsp]

But does a return type help with inference? I thought that f(x) = z(x)::Bool helps, but f(x)::Bool = z(x) does not? Not sure why i think that

[fonsp]

@rikhuijzer

[rikhuijzer]

Sorry for late reply. I didn't see this one yet.

But does a return type help with inference? I thought that f(x) = z(x)::Bool helps, but f(x)::Bool = z(x) does not? Not sure why i think that

Depends on where you care about the inference. In this case, the assertion on the method for has_hook_style_pluto_properties_in_expr doesn't help to improve the running time of this method because the any returns an Union{Missing,Bool} and the type annotation on the function sets that to Bool. So, no real gains there indeed. Instead the method might take slightly longer to run.

However, this may help a lot for places where the method is used, because thanks to the annotation, the type can be inferred which, in turn, may allow more code to be hit during precompilation.

In this case, I just added it because has_hook_style_pluto_properties_in_expr must be a boolean at line 281. So, adding the annotation improves:

1. code readability
2. code inferability
3. debugability because it fails at an earlier point. Namely, when an incorrect type is returned instead of when the incorrect type is passed to the CachedMacroExpansion struct.

[rikhuijzer]

Thanks a lot for the review Paul. Much appreciated ❤️

@rikhuijzer you should be interested by tooking a look at #1881 where this improvements could be more beneficial.

Yeah, I've waited now for a few weeks for that PR to make progress, but I figured that it's not going forward. Performance gains to PlutoRunner are very important for the user, I think. So, I figured that I should start working on PlutoRunner

[rikhuijzer]

I've added PlutoRunner.run_expression to the benchmark now so that we can track it. Output on Julia 1.8-beta3 on my system:

Before (master)

               Total measured:                   33.1s            2.41GiB

 Section                            ncalls     time    %tot     alloc    %tot
 ────────────────────────────────────────────────────────────────────────────
 import Pluto                            1    673ms    2.1%   76.2MiB    3.1%
 compiletimes.jl                         1    31.1s   97.9%   2.29GiB   96.9%
   Pluto.Cell                            1   45.4ms    0.1%   62.3KiB    0.0%
   Pluto.Notebook                        1    600ms    1.9%   87.0MiB    3.6%
   SessionActions.open                   1    13.5s   42.5%    674MiB   27.8%
   SessionActions.shutdown               1    459ms    1.4%   42.9MiB    1.8%
   Configuration.from_flat_kwargs        1   14.7ms    0.0%   13.6KiB    0.0%
   Pluto.run                             1    5.55s   17.4%    558MiB   23.0%
 ──────────────────────────────────────────────────────────────────────────── 

After (this PR)

                                            ───────────────   ───────────────
               Total measured:                   29.9s            2.48GiB

 Section                            ncalls     time    %tot     alloc    %tot
 ────────────────────────────────────────────────────────────────────────────
 import Pluto                            1    561ms    1.9%   77.6MiB    3.1%
 compiletimes.jl                         1    28.3s   98.1%   2.36GiB   96.9%
   Pluto.Cell                            1   39.7ms    0.1%   62.3KiB    0.0%
   Pluto.Notebook                        1    540ms    1.9%   87.0MiB    3.5%
   PlutoRunner.run_expression            1    1.51s    5.2%   86.8MiB    3.5%
   SessionActions.open                   1    10.4s   36.1%    563MiB   22.6%
   SessionActions.shutdown               1    415ms    1.4%   43.0MiB    1.7%
   Configuration.from_flat_kwargs        1   13.1ms    0.0%   13.5KiB    0.0%
   Pluto.run                             1    4.41s   15.3%    558MiB   22.4%
 ────────────────────────────────────────────────────────────────────────────

Conclusion

More progress is needed. PlutoRunner.run_expression has still 86MiB allocations. Reducing this will reduce the time for SessionActions.open.

[fonsp]

Could you also take a look at PlutoRunner.format_output((a=[1,2],c=Dict("d"=>(5,6,true))))? That's a big source of compile times right now: every time you output a new "tree viewable" type like dict, vector, tuple, it recompiles all of our show methods for the new type

[fonsp]

🎉 Awesome!! Did this really reduce SessionActions.open from 13s to 10s? That's a lot!

[rikhuijzer]

Nope. Sorry voor de onduidelijke communicatie. Dat kan beter. In dit geval maar 20 MiB reducties helaas. Een gedeelte van de compile time is simpelweg verplaatst naar PlutoRunner.run_expression omdat die nu ook in de test zit maar ik snap de verwarring.