From 3a876a72562b55b7a6013fa7896d8cfc5c99bfc6 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Mon, 8 Dec 2025 03:54:31 -0800
Subject: [PATCH] Replace mip package with pulp

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/gpu_tests.yml               |  2 -
 .../mip/mip_with_multi_layer_replacements.py  | 37 ++++++++++---------
 setup.py                                      |  1 -
 3 files changed, 19 insertions(+), 21 deletions(-)

diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index b7baea8b4..be7e88dfc 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.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:
diff --git a/modelopt/torch/_compress/mip/mip_with_multi_layer_replacements.py b/modelopt/torch/_compress/mip/mip_with_multi_layer_replacements.py
index 438db3312..e49b640c1 100644
--- a/modelopt/torch/_compress/mip/mip_with_multi_layer_replacements.py
+++ b/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,7 +80,7 @@ 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():
@@ -86,22 +88,21 @@ def run_mip(
         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)
diff --git a/setup.py b/setup.py
index 20a271fe1..b1ddada62 100644
--- a/setup.py
+++ b/setup.py
@@ -107,7 +107,6 @@
         "hydra-core==1.3.2",
         "immutabledict",
         "lru-dict",
-        "mip",
         "omegaconf==2.3.0",
         "pandas",
         "typeguard",