Replace mip package with pulp

Signed-off-by: Keval Morabia <[email protected]>

Author: kevalmorabia97

.github/workflows/gpu_tests.yml

@@ -73,8 +73,6 @@ jobs:
       - name: Setup environment variables
         run: |
           echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu" >> $GITHUB_ENV
-      - name: Install dependencies for mip
-        run: apt-get update && apt-get install -y libffi-dev
       - name: Run gpu tests
         run: pip install tox-current-env && tox -e py312-cuda12-gpu --current-env
   gpu-tests-non-pr:

modelopt/torch/_compress/mip/mip_with_multi_layer_replacements.py

@@ -23,9 +23,9 @@
 from random import random
 from typing import Any, Hashable, Iterable, Optional, TypeAlias
 
-from mip import BINARY, Model, maximize, minimize, xsum
+import pulp
 
-from .utils import InfeasibleError, consecutive_ngrams, get_nested_key, sort_replacements
+from modelopt.torch._compress.mip.utils import consecutive_ngrams, get_nested_key, sort_replacements
 
 ReplacementID: TypeAlias = Hashable
 Replacement: TypeAlias = dict[str, Any]
@@ -38,6 +38,7 @@ def run_mip(
     constraints: dict[str, float],
     bigger_is_better: bool,
     max_seconds_per_solution: Optional[float] = None,
+    verbose: bool = True,
 ) -> tuple[ChosenReplacements, float, dict[str, float]]:
     orig_num_replacements = len(replacements)
     replacements = {
@@ -52,13 +53,15 @@ def run_mip(
         )
         print("\n\n\n")
 
-    mip_model = Model()
+    # Create pulp problem with appropriate sense (minimize or maximize)
+    sense = pulp.LpMaximize if bigger_is_better else pulp.LpMinimize
+    problem = pulp.LpProblem(name="multi_layer_replacement", sense=sense)
 
     objective_vars = []
     constraint_vars = {constraint_key: [] for constraint_key in constraints.keys()}
     choice_indicators_by_layer = defaultdict(list)
-    for replacement_id, replacement in replacements.items():
-        is_chosen = mip_model.add_var(var_type=BINARY)
+    for i, (replacement_id, replacement) in enumerate(replacements.items()):
+        is_chosen = pulp.LpVariable(f"choice_{i}", cat=pulp.LpBinary)
         replacement["is_chosen"] = is_chosen
 
         for parent_layer_idx in replacement["parent_layer_indices"]:
@@ -73,30 +76,29 @@ def run_mip(
 
     # MIP constraints: each parent layer must come from exactly one chosen replacement
     for parent_layer_idx, curr_choice_indicators in choice_indicators_by_layer.items():
-        mip_model += xsum(curr_choice_indicators) == 1
+        problem += pulp.lpSum(curr_choice_indicators) == 1
 
     # MIP constraints: the sum of chosen replacement costs must be lower than the max cost
     for constraint_key, max_cost in constraints.items():
         min_cost = None
         if isinstance(max_cost, Iterable):
             min_cost, max_cost = max_cost
 
-        if max_cost is not None:
-            mip_model += xsum(constraint_vars[constraint_key]) <= max_cost
-        if min_cost is not None:
-            mip_model += xsum(constraint_vars[constraint_key]) >= min_cost
+        # PuLP is stricter than mip - it doesn't allow NaN/inf in constraints
+        if max_cost is not None and math.isfinite(max_cost):
+            problem += pulp.lpSum(constraint_vars[constraint_key]) <= max_cost
+        if min_cost is not None and math.isfinite(min_cost):
+            problem += pulp.lpSum(constraint_vars[constraint_key]) >= min_cost
 
     # MIP objective
-    mip_model.objective = (
-        maximize(xsum(objective_vars)) if bigger_is_better else minimize(xsum(objective_vars))
-    )
-
-    if max_seconds_per_solution is not None:
-        mip_model.max_seconds = max_seconds_per_solution
+    problem += (pulp.lpSum(objective_vars), "objective")
 
-    mip_model.optimize()
+    # Configure and run solver
+    solver = pulp.PULP_CBC_CMD(msg=verbose, timeLimit=max_seconds_per_solution)
+    problem.solve(solver)
 
-    if is_chosen.x is None:
+    # Check if solution is feasible
+    if problem.status != pulp.LpStatusOptimal:
         return []
         # raise InfeasibleError()
 
@@ -106,7 +108,7 @@ def run_mip(
     chosen_replacements: ChosenReplacements = []
     chosen_layers = []
     for replacement_id, replacement in replacements.items():
-        is_chosen = replacement["is_chosen"].x >= 0.99
+        is_chosen = replacement["is_chosen"].varValue >= 0.99
         if is_chosen:
             assert replacement not in chosen_replacements
             chosen_replacements.append(replacement)

setup.py

@@ -111,7 +111,6 @@
         "typeguard",
         "pandas",
         "immutabledict",
-        "mip",
     ],
 }