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,7 +23,7 @@
 from random import random
 from typing import Any, Hashable, Iterable, Optional, TypeAlias
 
-from mip import BINARY, Model, maximize, minimize, xsum
+import pulp
 
 from modelopt.torch._compress.mip.utils import (
     InfeasibleError,
@@ -57,13 +57,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"]:
@@ -78,30 +80,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=True, 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()
 
@@ -111,7 +112,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

@@ -107,7 +107,6 @@
         "hydra-core==1.3.2",
         "immutabledict",
         "lru-dict",
-        "mip",
         "omegaconf==2.3.0",
         "pandas",
         "typeguard",