Skip to content

Commit 5ee68d9

Browse files
fix: avoid world-age/serialization issues in distributed workflows
Co-authored-by: Claude <noreply@anthropic.com>
1 parent 6d0dd38 commit 5ee68d9

2 files changed

Lines changed: 51 additions & 2 deletions

File tree

lib/ModelingToolkitBase/src/systems/codegen_utils.jl

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,16 @@ function eval_or_rgf(expr::Expr; eval_expression = false, eval_module = @__MODUL
137137
if eval_expression
138138
return eval_module.eval(expr)
139139
else
140-
return drop_expr(RuntimeGeneratedFunction(eval_module, eval_module, expr))
140+
# Only function-definition expressions benefit from RuntimeGeneratedFunction (avoids
141+
# world-age issues with new methods). Module-level references such as
142+
# `:(ModelingToolkitBase._oop_unimplemented)` are not function definitions; wrapping
143+
# them in an RGF would fail. Evaluate them directly instead — no new method is
144+
# introduced so there is no world-age concern.
145+
if Meta.isexpr(expr, :function) || Meta.isexpr(expr, :->)
146+
return drop_expr(RuntimeGeneratedFunction(eval_module, eval_module, expr))
147+
else
148+
return eval_module.eval(expr)
149+
end
141150
end
142151
end
143152

@@ -595,12 +604,36 @@ Base.@nospecializeinfer function build_function_wrapper(
595604
end
596605

597606
optimize = resolve_optimize_option(optimize)
598-
return Symbolics.codegen_function(ir, expr, args; wrap_code, similarto, cse, optimize, kwargs...)
607+
result = Symbolics.codegen_function(ir, expr, args; wrap_code, similarto, cse, optimize, kwargs...)
608+
# When iip_config disables one side, Symbolics generates an anonymous `unimplemented`
609+
# function expression. Replace it here — where we know which side is disabled — with a
610+
# reference to a named module-level function so all downstream paths (expression=Val{true},
611+
# eval_expression=true, distributed serialization) get a stable, serializable callable.
612+
iip_config = get(kwargs, :iip_config, (true, true))
613+
if result isa NTuple{2, Expr}
614+
oop_expr, iip_expr = result
615+
if !iip_config[1]
616+
oop_expr = OOP_UNIMPLEMENTED_EXPR
617+
end
618+
if !iip_config[2]
619+
iip_expr = IIP_UNIMPLEMENTED_EXPR
620+
end
621+
result = (oop_expr, iip_expr)
622+
end
623+
return result
599624
end
600625

601626
resolve_optimize_option(x) = x
602627
resolve_optimize_option(::Nothing) = nothing
603628

629+
# Module-level fallback functions for the disabled side of an `iip_config` pair.
630+
# Using named module-level functions ensures correct serialization across all codegen paths
631+
# (expression=Val{true} evaluated by the user, eval_expression=true, distributed workers).
632+
_oop_unimplemented(args...) = throw(Symbolics.FunctionUnimplementedError("out-of-place"))
633+
_iip_unimplemented(args...) = throw(Symbolics.FunctionUnimplementedError("in-place"))
634+
const OOP_UNIMPLEMENTED_EXPR = :($ModelingToolkitBase._oop_unimplemented)
635+
const IIP_UNIMPLEMENTED_EXPR = :($ModelingToolkitBase._iip_unimplemented)
636+
604637
"""
605638
$(TYPEDEF)
606639

lib/ModelingToolkitBase/test/odesystem.jl

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,22 @@ f.f(du, u, p, 0.1)
110110
@test du == [4, 0, -16]
111111
@test_throws Symbolics.FunctionUnimplementedError f.f(u, p, 0.1)
112112

113+
# check that iip_config with expression = Val{true} produces serializable functions
114+
# (regression test for https://github.com/SciML/ModelingToolkit.jl/issues/4464)
115+
f_iip_expr = ODEFunction(de; iip_config = (false, true), expression = Val{true})
116+
f_iip_from_expr = eval(f_iip_expr)
117+
# verify the OOP stub is the stable module-level function (not a Main closure)
118+
@test f_iip_from_expr.f.f_oop === ModelingToolkitBase._oop_unimplemented
119+
# verify round-trip serialization works (simulates distributed usage)
120+
@testset "Issue#4464" begin
121+
using Serialization
122+
buf = IOBuffer()
123+
serialize(buf, f_iip_from_expr)
124+
seekstart(buf)
125+
f2 = deserialize(buf)
126+
@test f2.f.f_oop === ModelingToolkitBase._oop_unimplemented
127+
end
128+
113129
#check iip
114130
f = eval(ODEFunction(de; expression = Val{true}))
115131
f2 = ODEFunction(de)

0 commit comments

Comments
 (0)