Precompilation workload

Project.toml

@@ -17,6 +17,7 @@ JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
@@ -38,6 +39,7 @@ JuMP = "^1.28"
 LinearAlgebra = "1"
 Logging = "1"
 MathOptInterface = "1"
+PrecompileTools = "1"
 PrettyTables = "3.1"
 Random = "^1.10"
 Serialization = "1"

src/InfrastructureOptimizationModels.jl

@@ -115,6 +115,7 @@ function has_service end
 function validate_time_series! end
 
 import TimerOutputs
+import PrecompileTools
 
 # Base Imports
 import Base.getindex
@@ -667,4 +668,9 @@ include("utils/jump_utils.jl")
 include("utils/component_utils.jl")
 include("utils/time_series_utils.jl")
 include("utils/datetime_utils.jl")
+
+# Must remain the LAST include: bindings referenced by the workload resolve at
+# precompile execution time, so anything defined or imported after this line
+# would be invisible to it and fail only during Pkg.precompile.
+include("precompile_workload.jl")
 end

src/precompile_workload.jl

@@ -0,0 +1,78 @@
+# Precompilation workload. Exercises the key-type-independent container layers
+# (Settings/container init, JuMP fill of DenseAxisArray containers, objective
+# assembly) against private mock types so downstream consumers skip that compile.
+# Structs sit at top level because @setup_workload lowers into a `let` block.
+
+mutable struct _WorkloadSystem <: IS.InfrastructureSystemsContainer
+    base_power::Float64
+end
+
+get_base_power(sys::_WorkloadSystem) = sys.base_power
+stores_time_series_in_memory(::_WorkloadSystem) = true
+
+struct _WorkloadComponent <: IS.InfrastructureSystemsComponent end
+struct _WorkloadVariable <: VariableType end
+struct _WorkloadConstraint <: ConstraintType end
+struct _WorkloadExpression <: ExpressionType end
+
+function _run_precompile_workload()
+    sys = _WorkloadSystem(100.0)
+    settings = Settings(
+        sys;
+        horizon = Dates.Hour(24),
+        resolution = Dates.Hour(1),
+        time_series_cache_size = 0,
+    )
+    container = OptimizationContainer(sys, settings, nothing, IS.Deterministic)
+    set_time_steps!(container, 1:24)
+    jump_model = get_jump_model(container)
+    names = ["c1", "c2", "c3"]
+    time_steps = get_time_steps(container)
+    variables = add_variable_container!(
+        container,
+        _WorkloadVariable,
+        _WorkloadComponent,
+        names,
+        time_steps,
+    )
+    for name in names, t in time_steps
+        variables[name, t] = JuMP.@variable(
+            jump_model,
+            base_name = "V_{$(name), $(t)}",
+            lower_bound = 0.0,
+            upper_bound = 10.0
+        )
+    end
+    expressions = add_expression_container!(
+        container,
+        _WorkloadExpression,
+        _WorkloadComponent,
+        names,
+        time_steps,
+    )
+    for name in names, t in time_steps
+        expressions[name, t] = JuMP.AffExpr(0.0)
+        JuMP.add_to_expression!(expressions[name, t], variables[name, t])
+    end
+    constraints = add_constraints_container!(
+        container,
+        _WorkloadConstraint,
+        _WorkloadComponent,
+        names,
+        time_steps,
+    )
+    for name in names, t in time_steps
+        constraints[name, t] = JuMP.@constraint(jump_model, expressions[name, t] <= 5.0)
+    end
+    add_to_objective_invariant_expression!(container, 2.0 * variables["c1", 1])
+    objective_function = get_objective_expression(container)
+    objective = get_objective_expression(objective_function)
+    JuMP.@objective(jump_model, MOI.MIN_SENSE, objective)
+    return
+end
+
+PrecompileTools.@setup_workload begin
+    PrecompileTools.@compile_workload begin
+        _run_precompile_workload()
+    end
+end