Implement lazy callbacks

src/instance/file.jl

@@ -9,9 +9,13 @@ mutable struct FileInstance <: Instance
     loaded::Union{Nothing,JuMPInstance}
     filename::AbstractString
     h5::PyCall.PyObject
+    lazycb::Union{Nothing,Tuple{Function,Function}}
 
-    function FileInstance(filename::AbstractString)::FileInstance
-        instance = new(nothing, nothing, filename)
+    function FileInstance(
+        filename::AbstractString;
+        lazycb::Union{Nothing,Tuple{Function,Function}} = nothing,
+    )::FileInstance
+        instance = new(nothing, nothing, filename, nothing, lazycb)
         instance.py = PyFileInstance(instance)
         instance.h5 = Hdf5Sample(filename)
         instance.filename = filename
@@ -20,16 +24,27 @@ mutable struct FileInstance <: Instance
 end
 
 to_model(instance::FileInstance) = to_model(instance.loaded)
+
 get_instance_features(instance::FileInstance) = get_instance_features(instance.loaded)
+
 get_variable_features(instance::FileInstance, names) =
     get_variable_features(instance.loaded, names)
+
 get_variable_categories(instance::FileInstance, names) =
     get_variable_categories(instance.loaded, names)
+
 get_constraint_features(instance::FileInstance, names) =
     get_constraint_features(instance.loaded, names)
+
 get_constraint_categories(instance::FileInstance, names) =
     get_constraint_categories(instance.loaded, names)
 
+find_violated_lazy_constraints(instance::FileInstance, solver) =
+    find_violated_lazy_constraints(instance.loaded, solver)
+
+enforce_lazy_constraint(instance::FileInstance, solver, violation) =
+    enforce_lazy_constraint(instance.loaded, solver, violation)
+
 function get_samples(instance::FileInstance)
     return [instance.h5]
 end
@@ -40,7 +55,7 @@ end
 
 function load(instance::FileInstance)
     if instance.loaded === nothing
-        instance.loaded = load_instance(instance.filename)
+        instance.loaded = load_instance(instance.filename, lazycb = instance.lazycb)
     end
 end
 
@@ -72,6 +87,10 @@ function __init_PyFileInstance__()
         load(self) = load(self.jl)
         free(self) = free(self.jl)
         flush(self) = flush(self.jl)
+        find_violated_lazy_constraints(self, solver, _) =
+            find_violated_lazy_constraints(self.jl, solver)
+        enforce_lazy_constraint(self, solver, _, violation) =
+            enforce_lazy_constraint(self.jl, solver, violation)
     end
     copy!(PyFileInstance, Class)
 end

src/instance/jump.jl

@@ -12,7 +12,7 @@ mutable struct JuMPInstance <: Instance
     ext::AbstractDict
     samples::Vector{PyCall.PyObject}
 
-    function JuMPInstance(model::JuMP.Model)
+    function JuMPInstance(model::JuMP.Model)::JuMPInstance
         init_miplearn_ext(model)
         instance = new(nothing, model, nothing, model.ext[:miplearn], [])
         py = PyJuMPInstance(instance)
@@ -84,6 +84,18 @@ function create_sample!(instance::JuMPInstance)
     return sample
 end
 
+function find_violated_lazy_constraints(instance::JuMPInstance, solver)::Vector{String}
+    if "lazy_find_cb" ∈ keys(instance.model.ext[:miplearn])
+        return instance.model.ext[:miplearn]["lazy_find_cb"](instance.model, solver.data)
+    else
+        return []
+    end
+end
+
+function enforce_lazy_constraint(instance::JuMPInstance, solver, violation::String)::Nothing
+    instance.model.ext[:miplearn]["lazy_enforce_cb"](instance.model, solver.data, violation)
+end
+
 function __init_PyJuMPInstance__()
     @pydef mutable struct Class <: miplearn.Instance
         function __init__(self, jl)
@@ -101,6 +113,10 @@ function __init_PyJuMPInstance__()
             to_str_array(get_constraint_categories(self.jl, from_str_array(names)))
         get_samples(self) = get_samples(self.jl)
         create_sample(self) = create_sample!(self.jl)
+        find_violated_lazy_constraints(self, solver, _) =
+            find_violated_lazy_constraints(self.jl, solver)
+        enforce_lazy_constraint(self, solver, _, violation) =
+            enforce_lazy_constraint(self.jl, solver, violation)
     end
     copy!(PyJuMPInstance, Class)
 end
@@ -116,7 +132,11 @@ function save(filename::AbstractString, instance::JuMPInstance)::Nothing
     h5 = Hdf5Sample(filename, mode = "w")
     h5.put_scalar("miplearn_version", "0002")
     h5.put_bytes("mps", mps)
-    h5.put_scalar("jump_ext", JSON.json(model.ext[:miplearn]))
+
+    ext = copy(model.ext[:miplearn])
+    delete!(ext, "lazy_find_cb")
+    delete!(ext, "lazy_enforce_cb")
+    h5.put_scalar("jump_ext", JSON.json(ext))
     return
 end
 
@@ -130,12 +150,19 @@ function _check_miplearn_version(h5)
     )
 end
 
-function load_instance(filename::AbstractString)::JuMPInstance
+function load_instance(
+    filename::AbstractString;
+    lazycb::Union{Nothing,Tuple{Function,Function}} = nothing,
+)::JuMPInstance
     h5 = Hdf5Sample(filename)
     _check_miplearn_version(h5)
     mps = h5.get_bytes("mps")
-    ext = h5.get_scalar("jump_ext")
-    instance = JuMPInstance(Vector{UInt8}(mps), JSON.parse(ext))
+    ext = JSON.parse(h5.get_scalar("jump_ext"))
+    if lazycb !== nothing
+        ext["lazy_find_cb"] = lazycb[1]
+        ext["lazy_enforce_cb"] = lazycb[2]
+    end
+    instance = JuMPInstance(Vector{UInt8}(mps), ext)
     return instance
 end
 

src/solvers/jump.jl

@@ -9,6 +9,8 @@ using MathOptInterface
 using TimerOutputs
 const MOI = MathOptInterface
 
+import JuMP: value
+
 mutable struct JuMPSolverData
     optimizer_factory::Any
     varname_to_var::Dict{String,VariableRef}
@@ -19,6 +21,7 @@ mutable struct JuMPSolverData
     solution::Dict{JuMP.VariableRef,Float64}
     reduced_costs::Vector{Float64}
     dual_values::Dict{JuMP.ConstraintRef,Float64}
+    cb_data::Any
 end
 
 
@@ -175,10 +178,25 @@ function remove_constraints(data::JuMPSolverData, names::Vector{String})::Nothin
 end
 
 
-function solve(data::JuMPSolverData; tee::Bool = false, iteration_cb = nothing)
+function solve(
+    data::JuMPSolverData;
+    tee::Bool = false,
+    iteration_cb = nothing,
+    lazy_cb = nothing,
+)
     model = data.model
     wallclock_time = 0
     log = ""
+
+    if lazy_cb !== nothing
+        function lazy_cb_wrapper(cb_data)
+            data.cb_data = cb_data
+            lazy_cb(nothing, nothing)
+            data.cb_data = nothing
+        end
+        MOI.set(model, MOI.LazyConstraintCallback(), lazy_cb_wrapper)
+    end
+
     while true
         wallclock_time += @elapsed begin
             log *= _optimize_and_capture_output!(model, tee = tee)
@@ -189,6 +207,7 @@ function solve(data::JuMPSolverData; tee::Bool = false, iteration_cb = nothing)
             break
         end
     end
+
     if is_infeasible(data)
         data.solution = Dict()
         primal_bound = nothing
@@ -452,6 +471,7 @@ function __init_JuMPSolver__()
                 Dict(),  # solution
                 [],  # reduced_costs
                 Dict(),  # dual_values
+                nothing,  # cb_data
             )
         end
 
@@ -555,12 +575,27 @@ function __init_JuMPSolver__()
             iteration_cb = nothing,
             lazy_cb = nothing,
             user_cut_cb = nothing,
-        ) = solve(self.data, tee = tee, iteration_cb = iteration_cb)
+        ) = solve(self.data, tee = tee, iteration_cb = iteration_cb, lazy_cb = lazy_cb)
 
         solve_lp(self; tee = false) = solve_lp(self.data, tee = tee)
     end
     copy!(JuMPSolver, Class)
 end
 
+function value(solver::JuMPSolverData, var::VariableRef)
+    if solver.cb_data !== nothing
+        return JuMP.callback_value(solver.cb_data, var)
+    else
+        return JuMP.value(var)
+    end
+end
 
-export JuMPSolver
+function submit(solver::JuMPSolverData, con::AbstractConstraint, name::String = "")
+    if solver.cb_data !== nothing
+        MOI.submit(solver.model, MOI.LazyConstraint(solver.cb_data), con)
+    else
+        JuMP.add_constraint(solver.model, con, name)
+    end
+end
+
+export JuMPSolver, submit

src/solvers/macros.jl

@@ -53,6 +53,13 @@ function set_category!(c::ConstraintRef, category::String)::Nothing
     return
 end
 
+function set_lazy_callback!(model::Model, find_cb::Function, enforce_cb::Function)::Nothing
+    ext = init_miplearn_ext(model)
+    ext["lazy_find_cb"] = find_cb
+    ext["lazy_enforce_cb"] = enforce_cb
+    return
+end
+
 
 macro feature(obj, features)
     quote
@@ -67,6 +74,12 @@ macro category(obj, category)
     end
 end
 
+macro lazycb(obj, find_cb, enforce_cb)
+    quote
+        set_lazy_callback!($(esc(obj)), $(esc(find_cb)), $(esc(enforce_cb)))
+    end
+end
+
 function _get_and_check_name(obj)
     n = name(obj)
     length(n) > 0 || error(
@@ -77,4 +90,4 @@ function _get_and_check_name(obj)
 end
 
 
-export @feature, @category
+export @feature, @category, @lazycb

src/utils/sysimage.jl

@@ -4,10 +4,10 @@
 
 using PackageCompiler
 
-using CSV
 using Cbc
 using Clp
 using Conda
+using CSV
 using DataFrames
 using Distributed
 using JLD2
@@ -20,10 +20,10 @@ using PyCall
 using TimerOutputs
 
 pkg = [
-    :CSV
     :Cbc
     :Clp
     :Conda
+    :CSV
     :DataFrames
     :Distributed
     :JLD2