From 31e76fd27b60d0dfb441c6b0dbfc559d074eb2c6 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Wed, 31 Dec 2025 12:44:22 +0530 Subject: [PATCH 1/2] temp commit: add logging to debug long ttfx in CI --- lib/OptimizationMOI/src/OptimizationMOI.jl | 21 ++++++++++++++++++++- lib/OptimizationMOI/src/moi.jl | 4 ++++ lib/OptimizationMOI/src/nlp.jl | 9 +++++++++ lib/OptimizationMOI/test/runtests.jl | 5 +++++ 4 files changed, 38 insertions(+), 1 deletion(-) diff --git a/lib/OptimizationMOI/src/OptimizationMOI.jl b/lib/OptimizationMOI/src/OptimizationMOI.jl index 7f24d708f..cbe0987e0 100644 --- a/lib/OptimizationMOI/src/OptimizationMOI.jl +++ b/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,7 +330,7 @@ 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 @@ -326,11 +338,15 @@ 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 @@ -376,6 +393,7 @@ function SciMLBase.__init(prob::OptimizationProblem, mtkize = false, kwargs...) cache = if MOI.supports(_create_new_optimizer(opt), MOI.NLPBlock()) + @info "CACHE1" MOIOptimizationNLPCache(prob, opt; maxiters, @@ -385,6 +403,7 @@ function SciMLBase.__init(prob::OptimizationProblem, mtkize, kwargs...) else + @info "CACHE2" MOIOptimizationCache(prob, opt; maxiters, maxtime, abstol, reltol, kwargs...) end return cache diff --git a/lib/OptimizationMOI/src/moi.jl b/lib/OptimizationMOI/src/moi.jl index db298ec56..c2b97305f 100644 --- a/lib/OptimizationMOI/src/moi.jl +++ b/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 diff --git a/lib/OptimizationMOI/src/nlp.jl b/lib/OptimizationMOI/src/nlp.jl index 74e552ae3..ce1cf31af 100644 --- a/lib/OptimizationMOI/src/nlp.jl +++ b/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) diff --git a/lib/OptimizationMOI/test/runtests.jl b/lib/OptimizationMOI/test/runtests.jl index c9f5113ac..7c4fcd330 100644 --- a/lib/OptimizationMOI/test/runtests.jl +++ b/lib/OptimizationMOI/test/runtests.jl @@ -36,10 +36,13 @@ 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 @@ -47,12 +50,14 @@ end @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 From dd5ab024e994b3ed82ffd3b5f20ab64dece9a1af Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Wed, 31 Dec 2025 14:15:53 +0530 Subject: [PATCH 2/2] more logging --- lib/OptimizationBase/src/solve.jl | 14 ++++++++++++++ lib/OptimizationMOI/src/OptimizationMOI.jl | 7 ++++++- 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/lib/OptimizationBase/src/solve.jl b/lib/OptimizationBase/src/solve.jl index 777ca1d8f..42f8e91e9 100644 --- a/lib/OptimizationBase/src/solve.jl +++ b/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 diff --git a/lib/OptimizationMOI/src/OptimizationMOI.jl b/lib/OptimizationMOI/src/OptimizationMOI.jl index cbe0987e0..714a265b6 100644 --- a/lib/OptimizationMOI/src/OptimizationMOI.jl +++ b/lib/OptimizationMOI/src/OptimizationMOI.jl @@ -392,7 +392,12 @@ 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;