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
+
+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
+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

test/Project.toml

@@ -1,14 +1,14 @@
 [deps]
-CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 Cbc = "9961bab8-2fa3-5c5a-9d89-47fab24efd76"
 Clp = "e2554f3b-3117-50c0-817c-e040a3ddf72d"
 Conda = "8f4d0f93-b110-5947-807f-2305c1781a2d"
+CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Gurobi = "2e9cd046-0924-5485-92f1-d5272153d98b"
 JSON2 = "2535ab7d-5cd8-5a07-80ac-9b1792aadce3"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
-MIPLearn = "2b1277c3-b477-4c49-a15e-7ba350325c68"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
+MIPLearn = "2b1277c3-b477-4c49-a15e-7ba350325c68"
 PackageCompiler = "9b87118b-4619-50d2-8e1e-99f35a4d4d9d"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"

test/instance/jump_test.jl

@@ -2,29 +2,62 @@
 #  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
+using Cbc
+using JuMP
+using MathOptInterface
 using MIPLearn
+const MOI = MathOptInterface
 
+function find_lazy(model::Model, cb_data)::Vector{String}
+    x = variable_by_name(model, "x")
+    y = variable_by_name(model, "y")
+    x_val = value(cb_data, x)
+    y_val = value(cb_data, y)
+    if x_val + y_val > 1 + 1e-6
+        return ["con"]
+    end
+    return []
+end
+
+function enforce_lazy(model::Model, cb_data, violation::String)::Nothing
+    if violation == "con"
+        x = variable_by_name(model, "x")
+        y = variable_by_name(model, "y")
+        con = @build_constraint(x + y <= 1)
+        submit(cb_data, con)
+    end
+    return
+end
 
+function build_model()
+    model = Model()
+    @variable(model, x, Bin)
+    @variable(model, y, Bin)
+    @objective(model, Max, 2 * x + y)
+    @constraint(model, c1, x + y <= 2)
+    @lazycb(model, find_lazy, enforce_lazy)
+    return model
+end
+
+@testset "Lazy callback" begin
     @testset "JuMPInstance" begin
-    @testset "Save and load" begin
-        # Build instance and solve
-        model = model = build_knapsack_model()
+        model = build_model()
         instance = JuMPInstance(model)
-        solver = LearningSolver(Gurobi.Optimizer)
-        stats = solve!(solver, instance)
-        @test length(instance.py.samples) == 1
+        solver = LearningSolver(Cbc.Optimizer)
+        solve!(solver, instance)
+        @test value(model[:x]) == 1.0
+        @test value(model[:y]) == 0.0
+    end
 
-        # Save model to file
+    @testset "FileInstance" begin
+        model = build_model()
+        instance = JuMPInstance(model)
         filename = tempname()
         save(filename, instance)
-        @test isfile(filename)
-
-        # Read model from file
-        loaded = load_instance(filename)
-        x = variable_by_name(loaded.model, "x")
-        @test loaded.model.ext[:miplearn][:variable_features][x] == [1.0]
-        @test loaded.model.ext[:miplearn][:variable_categories][x] == "cat1"
-        @test loaded.model.ext[:miplearn][:instance_features] == [5.0]
-        @test length(loaded.py.samples) == 1
+        file_instance = FileInstance(filename, lazycb = (find_lazy, enforce_lazy))
+        solver = LearningSolver(Cbc.Optimizer)
+        solve!(solver, file_instance)
+        h5 = MIPLearn.Hdf5Sample(filename)
+        @test h5.get_array("mip_var_values") == [1.0, 0.0]
     end
 end

test/runtests.jl

@@ -10,6 +10,7 @@ MIPLearn.setup_logger()
 @testset "MIPLearn" begin
     include("fixtures/knapsack.jl")
     include("instance/file_test.jl")
+    include("instance/jump_test.jl")
     include("solvers/jump_test.jl")
     include("solvers/learning_test.jl")
     include("utils/benchmark_test.jl")