[DO NOT MERGE] debug long ttfx in OptimizationMOI CI

lib/OptimizationBase/src/solve.jl

@@ -110,33 +110,44 @@ function solve(
 end
 
 function _check_opt_alg(prob::SciMLBase.OptimizationProblem, alg; kwargs...)
+    @info "CHECK" 1
     !allowsbounds(alg) && (!isnothing(prob.lb) || !isnothing(prob.ub)) &&
         throw(IncompatibleOptimizerError("The algorithm $(typeof(alg)) does not support box constraints. Either remove the `lb` or `ub` bounds passed to `OptimizationProblem` or use a different algorithm."))
+    @info "CHECK" 2
     requiresbounds(alg) && isnothing(prob.lb) &&
         throw(IncompatibleOptimizerError("The algorithm $(typeof(alg)) requires box constraints. Either pass `lb` and `ub` bounds to `OptimizationProblem` or use a different algorithm."))
+    @info "CHECK" 3
     !allowsconstraints(alg) && !isnothing(prob.f.cons) &&
         throw(IncompatibleOptimizerError("The algorithm $(typeof(alg)) does not support constraints. Either remove the `cons` function passed to `OptimizationFunction` or use a different algorithm."))
+    @info "CHECK" 4
     requiresconstraints(alg) && isnothing(prob.f.cons) &&
         throw(IncompatibleOptimizerError("The algorithm $(typeof(alg)) requires constraints, pass them with the `cons` kwarg in `OptimizationFunction`."))
+    @info "CHECK" 5
     # Check that if constraints are present and the algorithm supports constraints, both lcons and ucons are provided
     allowsconstraints(alg) && !isnothing(prob.f.cons) &&
         (isnothing(prob.lcons) || isnothing(prob.ucons)) &&
         throw(ArgumentError("Constrained optimization problem requires both `lcons` and `ucons` to be provided to OptimizationProblem. " *
                             "Example: OptimizationProblem(optf, u0, p; lcons=[-Inf], ucons=[0.0])"))
+    @info "CHECK" 5
     !allowscallback(alg) && !(get(kwargs, :callback, DEFAULT_CALLBACK) isa NullCallback) &&
         throw(IncompatibleOptimizerError("The algorithm $(typeof(alg)) does not support callbacks, remove the `callback` keyword argument from the `solve` call."))
+    @info "CHECK" 6
     requiresgradient(alg) &&
         !(prob.f isa SciMLBase.AbstractOptimizationFunction) &&
         throw(IncompatibleOptimizerError("The algorithm $(typeof(alg)) requires gradients, hence use `OptimizationFunction` to generate them with an automatic differentiation backend e.g. `OptimizationFunction(f, AutoForwardDiff())` or pass it in with `grad` kwarg."))
+    @info "CHECK" 7
     requireshessian(alg) &&
         !(prob.f isa SciMLBase.AbstractOptimizationFunction) &&
         throw(IncompatibleOptimizerError("The algorithm $(typeof(alg)) requires hessians, hence use `OptimizationFunction` to generate them with an automatic differentiation backend e.g. `OptimizationFunction(f, AutoFiniteDiff(); kwargs...)` or pass them in with `hess` kwarg."))
+    @info "CHECK" 8
     requiresconsjac(alg) &&
         !(prob.f isa SciMLBase.AbstractOptimizationFunction) &&
         throw(IncompatibleOptimizerError("The algorithm $(typeof(alg)) requires constraint jacobians, hence use `OptimizationFunction` to generate them with an automatic differentiation backend e.g. `OptimizationFunction(f, AutoFiniteDiff(); kwargs...)` or pass them in with `cons` kwarg."))
+    @info "CHECK" 9
     requiresconshess(alg) &&
         !(prob.f isa SciMLBase.AbstractOptimizationFunction) &&
         throw(IncompatibleOptimizerError("The algorithm $(typeof(alg)) requires constraint hessians, hence use `OptimizationFunction` to generate them with an automatic differentiation backend e.g. `OptimizationFunction(f, AutoFiniteDiff(), AutoFiniteDiff(hess=true); kwargs...)` or pass them in with `cons` kwarg."))
+    @info "CHECK" 10
     return
 end
 
@@ -184,10 +195,13 @@ See also [`solve(prob::OptimizationProblem, alg, args...; kwargs...)`](@ref)
 """
 function init(prob::SciMLBase.OptimizationProblem, alg, args...;
         kwargs...)::SciMLBase.AbstractOptimizationCache
+    @info "Base.init"
     if prob.u0 !== nothing && !isconcretetype(eltype(prob.u0))
         throw(SciMLBase.NonConcreteEltypeError(eltype(prob.u0)))
     end
+    @info "Base check alg" 
     _check_opt_alg(prob::SciMLBase.OptimizationProblem, alg; kwargs...)
+    @info "Base __init"
     cache = __init(prob, alg, args...; prob.kwargs..., kwargs...)
     return cache
 end

lib/OptimizationMOI/src/OptimizationMOI.jl

@@ -177,10 +177,12 @@ function rep_pars_vals!(expr::Expr, expr_map)
             isequal(f, expr) && return n
         end
     end
+    @info "rep_pars_vals - threaded `for`"
     Threads.@sync for i in eachindex(expr.args)
         i == 1 && expr.head == :call && continue # first arg is the operator
         Threads.@spawn expr.args[i] = rep_pars_vals!(expr.args[i], expr_map)
     end
+    @info "rep_pars_vals - threaded `for` end"
     return expr
 end
 
@@ -213,14 +215,19 @@ parameters with `x[i]` and `p[i]`.
 function convert_to_expr(eq, expr_map; expand_expr = false)
     if expand_expr
         eq = try
+            @info "cte - expand"
             Symbolics.expand(eq) # PolyForm sometimes errors
         catch e
+            @info "cte - expand again" e
             Symbolics.expand(eq)
         end
     end
+    @info "cte - toexpr"
     expr = ModelingToolkitBase.toexpr(eq)
 
+    @info "cte - rep_pars_vals"
     expr = rep_pars_vals!(expr, expr_map)
+    @info "cte - symbolify"
     expr = symbolify!(expr)
     return expr
 end
@@ -294,18 +301,23 @@ function generate_exprs(prob::OptimizationProblem)
     Unsupported parameter object type $(typeof(pobj)) for expression construction.
     """
     @variables x[1:length(prob.u0)] p[1:length(pobj)]
+    @info "gen_exprs - trace"
     obj = prob.f.f(collect(x), collect(p))
+    @info "gen_exprs - toexpr"
     obj_expr = SU.Code.toexpr(SU.expand(SU.unwrap(obj)))
+    @info "gen_exprs - symbolify!"
     symbolify!(obj_expr)
     if prob.lcons === nothing && prob.ucons === nothing
         return SciMLBase.remake(f; expr = obj_expr)
     end
+    @info "gen_exprs - trace cons"
     if SciMLBase.isinplace(prob)
         cons_expr = zeros(Num, length(prob.lcons))
         prob.f.cons(cons_expr, collect(x), collect(p))
     else
         cons_expr = prob.f.cons(collect(x), collect(p))
     end
+    @info "gen_exprs - toexpr cons"
     cons_expr = SU.Code.toexpr.(SU.expand.(SU.unwrap.(cons_expr)))::Vector{Expr}
     for i in eachindex(cons_expr)
         cons_expr[i] = if prob.lcons[i] == prob.ucons[i]
@@ -318,19 +330,23 @@ function generate_exprs(prob::OptimizationProblem)
             Expr(:comparison, prob.lcons[i], :(<=), cons_expr[i], :(<=), prob.ucons[i])
         end
     end
-    symbolify!(obj_expr)
+    @info "gen_exprs - symbolify! cons"
     symbolify!.(cons_expr)
     newf = SciMLBase.remake(f; expr = obj_expr, cons_expr)
     return newf
 end
 
 function process_system_exprs(prob::OptimizationProblem, f::OptimizationFunction)
     @assert f.sys !== nothing
+    @info "pse - get_expr_map"
     expr_map = get_expr_map(prob.f.sys)
+    @info "pse - convert_to_expr"
     expr = convert_to_expr(f.expr, expr_map; expand_expr = false)
+    @info "pse - repl_getindex"
     expr = repl_getindex!(expr)
     cons = MTK.constraints(f.sys)
     cons_expr = Vector{Expr}(undef, length(cons))
+    @info "pse - threaded for"
     Threads.@sync for i in eachindex(cons)
         Threads.@spawn if prob.lcons[i] == prob.ucons[i] == 0
             cons_expr[i] = Expr(:call, :(==),
@@ -345,6 +361,7 @@ function process_system_exprs(prob::OptimizationProblem, f::OptimizationFunction
             expand_expr = false)), 0)
         end
     end
+    @info "pse - threaded for end"
     return expr, cons_expr
 end
 
@@ -375,7 +392,13 @@ function SciMLBase.__init(prob::OptimizationProblem,
         reltol::Union{Number, Nothing} = nothing,
         mtkize = false,
         kwargs...)
-    cache = if MOI.supports(_create_new_optimizer(opt), MOI.NLPBlock())
+    @info "create new opt"
+    _opt = _create_new_optimizer(opt)
+    @info "blk"
+    blk = MOI.NLPBlock()
+    @info "supports"
+    cache = if MOI.supports(_opt, blk)
+        @info "CACHE1"
         MOIOptimizationNLPCache(prob,
             opt;
             maxiters,
@@ -385,6 +408,7 @@ function SciMLBase.__init(prob::OptimizationProblem,
             mtkize,
             kwargs...)
     else
+        @info "CACHE2"
         MOIOptimizationCache(prob, opt; maxiters, maxtime, abstol, reltol, kwargs...)
     end
     return cache

lib/OptimizationMOI/src/moi.jl

@@ -17,8 +17,11 @@ function MOIOptimizationCache(prob::OptimizationProblem, opt; kwargs...)
     reinit_cache = OptimizationBase.ReInitCache(prob.u0, prob.p)
     if isnothing(f.sys)
         if f.adtype isa OptimizationBase.AutoSymbolics
+            @info "AutoSymbolics"
             num_cons = prob.ucons === nothing ? 0 : length(prob.ucons)
+            @info "generate_exprs"
             f = generate_exprs(prob)
+            @info "instantiate_fn"
             f = OptimizationBase.instantiate_function(f,
                 reinit_cache,
                 prob.f.adtype,
@@ -30,6 +33,7 @@ function MOIOptimizationCache(prob::OptimizationProblem, opt; kwargs...)
 
     # TODO: check if the problem is at most bilinear, i.e. affine and or quadratic terms in two variables
     if f.sys !== nothing
+        @info "process_system_exprs"
         expr, cons_expr = process_system_exprs(prob, f)
         f = remake(f; expr, cons_expr)
     end

lib/OptimizationMOI/src/nlp.jl

@@ -108,16 +108,21 @@ function MOIOptimizationNLPCache(prob::OptimizationProblem,
         mtkize = false,
         callback = nothing,
         kwargs...)
+    @info "NLPCACHE"
     reinit_cache = OptimizationBase.ReInitCache(prob.u0, prob.p) # everything that can be changed via `reinit`
 
     num_cons = prob.ucons === nothing ? 0 : length(prob.ucons)
     if prob.f.adtype isa ADTypes.AutoSymbolics || (prob.f.adtype isa ADTypes.AutoSparse &&
         prob.f.adtype.dense_ad isa ADTypes.AutoSymbolics)
+        @info "AutoSymbolics"
+        @info "generate_exprs"
         f = generate_exprs(prob)
+        @info "instantiate_function"
         f = OptimizationBase.instantiate_function(
             f, reinit_cache, prob.f.adtype, num_cons;
             g = true, h = true, cons_j = true, cons_h = true)
     else
+        @info "Non-symbolic instantiate_function"
         f = OptimizationBase.instantiate_function(
             prob.f, reinit_cache, prob.f.adtype, num_cons;
             g = true, h = true, cons_j = true, cons_vjp = true, lag_h = true)
@@ -149,19 +154,23 @@ function MOIOptimizationNLPCache(prob::OptimizationProblem,
     ucons = prob.ucons === nothing ? fill(T(Inf), num_cons) : prob.ucons
 
     if f.sys isa SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing} && mtkize
+        @info "MTKIZE"
         try
+            @info "do mtkize"
             sys = MTK.modelingtoolkitize(prob)
         catch err
             throw(ArgumentError("Automatic symbolic expression generation with ModelingToolkit failed with error: $err.
             Try by setting `mtkize = false` instead if the solver doesn't require symbolic expressions."))
         end
         if !isnothing(prob.p) && !(prob.p isa SciMLBase.NullParameters)
+            @info "sysprob"
             unames = variable_symbols(sys)
             pnames = parameter_symbols(sys)
             us = [unames[i] => prob.u0[i] for i in 1:length(prob.u0)]
             ps = [pnames[i] => prob.p[i] for i in 1:length(prob.p)]
             sysprob = OptimizationProblem(sys, us, ps)
         else
+            @info "sysprob2"
             unames = variable_symbols(sys)
             us = [unames[i] => prob.u0[i] for i in 1:length(prob.u0)]
             sysprob = OptimizationProblem(sys, us)

lib/OptimizationMOI/test/runtests.jl

@@ -36,23 +36,28 @@ end
     x0 = zeros(2)
     _p = [1.0, 100.0]
     cons_circ = (res, x, p) -> res .= [x[1]^2 + x[2]^2]
+    @info "OPTFN"
     optprob = OptimizationFunction(
         rosenbrock, AutoZygote();
         cons = cons_circ)
+    @info "OPTPROB"
     prob = OptimizationProblem(optprob, x0, _p, ucons = [Inf], lcons = [0.0])
+    @info "INIT"
     evaluator = init(prob, Ipopt.Optimizer()).evaluator
 
     x = prob.u0
     # vector-constraint jacobian product
     @test (evaluator.f.cons_j !== nothing) || (evaluator.f.cons_jvp !== nothing)
     y = zeros(1)
     w = ones(2)
+    @info "EVAL CJAC PROD"
     @test MathOptInterface.eval_constraint_jacobian_product(evaluator, y, x, w) === nothing
 
     # constraint jacobian-vector product
     @test (evaluator.f.cons_j !== nothing) || (evaluator.f.cons_vjp !== nothing)
     y = zeros(2)
     w = ones(1)
+    @info "EVAL CJAC TRANSPOSE PROD"
     @test MathOptInterface.eval_constraint_jacobian_transpose_product(
         evaluator, y, x, w) === nothing
 end