def find_ranking(comparisons, equal_width=0.2, max_rank=-1, verbose=False):
"""
Find the least changes to a set of comparisons so that they are consistent
(transitive), it returns a topological ranking.
comparisons A dictionary with tuple keys in the form of (i, j), values
are scalars indicating the probability of i > j. It is
assumed that comparisons are symmetric. Use 0 for i < j,
0.5 for i == j, and 1 for i > j (and any value in between).
equal_width 0..0.5-equal_width/2 is considered '<=' and 0.5..0.5+equal_width/2
is considered '>='. In between it is considered to be '=='.
max_rank Maximal rank, a low value forces the model to choose more
equal cases.
verbose Whether to print gurobi's progress.
Returns:
A tuple of size two:
0) Ranking derived from topological sort (list of ranks in order of
nodes);
1) Sum of absolute changes to the comparisons.
"""
# remove unnecessary variables
comparisons = {(i, j) if i < j else (j, i): value if i < j else 1 - value
for (i, j), value in comparisons.items()}
nodes = np.unique([i for ij in comparisons.keys() for i in ij])
# define variables
model = Model('comparison')
model.setParam('OutputFlag', verbose)
values = np.fromiter(comparisons.values(), dtype=float)
assert values.max() <= 1 and values.min() >= 0
# variables to encode the error of comparisons
E_ij = model.addVars(comparisons.keys(),
name='e_ij',
vtype=GRB.CONTINUOUS,
ub=1.0 - values,
lb=-values)
# variables to encode hard choice of >=, <=, ==
Ge_ij = model.addVars(comparisons.keys(), name='ge_ij', vtype=GRB.BINARY)
Le_ij = model.addVars(comparisons.keys(), name='le_ij', vtype=GRB.BINARY)
Eq_ij = model.addVars(comparisons.keys(), name='eq_ij', vtype=GRB.BINARY)
# variables to help with transitivity in non-fully connected graphs
if max_rank < 1:
max_rank = len(nodes)
R_i = model.addVars(nodes,
name='r_i',
vtype=GRB.CONTINUOUS,
lb=0,
ub=max_rank)
# variables to emulate abs
T_ij_pos = {}
T_ij_neg = {}
index = (values != 1) & (values != 0)
T_ij_pos = model.addVars(
(ij for ij, value in comparisons.items() if value not in [0.0, 1.0]),
vtype=GRB.CONTINUOUS,
name='T_ij_pos',
lb=0,
ub=1 - values[index])
T_ij_neg = model.addVars(
(ij for ij, value in comparisons.items() if value not in [0.0, 1.0]),
vtype=GRB.CONTINUOUS,
name='T_ij_neg',
lb=0,
ub=values[index])
model.update()
# emulate abs for non-binary comparisons: E_ij = T_ij_pos - T_ij_neg
model.addConstrs((E_ij[ij] == T_ij_pos[ij] - T_ij_neg[ij]
for ij in T_ij_pos), 'E_ij = T_ij_pos - T_ij_neg')
# hard decision of >=, <=, and ==
lower_bound = 0.5 - equal_width / 2.0
upper_bound = 0.5 + equal_width / 2.0
# <=
model.addConstrs((E_ij[ij] + comparisons[ij] - upper_bound <= ge_ij
for ij, ge_ij in Ge_ij.items()), 'ge_ij_lower_bound')
model.addConstrs((E_ij[ij] + comparisons[ij] - upper_bound >= -1 + ge_ij
for ij, ge_ij in Ge_ij.items()), 'ge_ij_upper_bound')
# >=
model.addConstrs((E_ij[ij] + comparisons[ij] - lower_bound >= -le_ij
for ij, le_ij in Le_ij.items()), 'le_ij_lower_bound')
model.addConstrs((E_ij[ij] + comparisons[ij] - lower_bound <= 1 - le_ij
for ij, le_ij in Le_ij.items()), 'le_ij_upper_bound')
# ==
model.addConstrs((
le + eq + ge == 1
for le, eq, ge in zip(Le_ij.values(), Eq_ij.values(), Ge_ij.values())),
'eq_ij')
# transitivity
for (i, j), eq_a in Eq_ij.items():
le_a = Le_ij[i, j]
ge_a = Ge_ij[i, j]
for k in nodes:
j_, k_ = j, k
if j > k:
j_, k_ = k, j
eq_b = Eq_ij.get((j_, k_), None)
if eq_b is None:
continue
else:
le_b = Le_ij[j_, k_]
ge_b = Ge_ij[j_, k_]
if j_ != j:
le_b, ge_b = ge_b, le_b
i_, k_ = i, k
if i > k:
i_, k_ = k, i
eq_c = Eq_ij.get((i_, k_), None)
if eq_c is None:
continue
else:
le_c = Le_ij[i_, k_]
ge_c = Ge_ij[i_, k_]
if i_ != i:
le_c, ge_c = ge_c, le_c
# a <= b and b <= c -> a <= c
model.addLConstr(ge_a + ge_b, GRB.LESS_EQUAL, 1 + ge_c,
f'transitivity_ge_{i},{j},{k}')
# a >= b and b >= c -> a >= c
model.addLConstr(le_a + le_b, GRB.LESS_EQUAL, 1 + le_c,
f'transitivity_le_{i},{j},{k}')
# a <= b and b == c -> a <= c
model.addLConstr(le_a + eq_b, GRB.LESS_EQUAL, 1 + le_c,
f'transitivity_leeq_{i},{j},{k}')
# a == b and b <= c -> a <= c
model.addLConstr(eq_a + le_b, GRB.LESS_EQUAL, 1 + le_c,
f'transitivity_eqle_{i},{j},{k}')
# a >= b and b == c --> a >= c
model.addLConstr(ge_a + eq_b, GRB.LESS_EQUAL, 1 + ge_c,
f'transitivity_geeq_{i},{j},{k}')
# a == b and b >= c --> a >= c
model.addLConstr(eq_a + ge_b, GRB.LESS_EQUAL, 1 + ge_c,
f'transitivity_eqge_{i},{j},{k}')
# a == b and b == c --> a == c
model.addLConstr(eq_a + eq_b, GRB.LESS_EQUAL, 1 + eq_c,
f'transitivity_eq_{i},{j},{k}')
# transitivity helper (for not-fully connected graphs)
# also provides a latent rank
big_m = max_rank
model.addConstrs(((1 - ge_ij) * big_m + R_i[i] >= R_i[j] + 1
for (i, j), ge_ij in Ge_ij.items()),
'rank_transitivity_larger')
model.addConstrs(((1 - le_ij) * big_m + R_i[j] >= R_i[i] + 1
for (i, j), le_ij in Le_ij.items()),
'rank_transitivity_smaller')
model.addConstrs(((1 - eq_ij) * big_m + R_i[j] >= R_i[i]
for (i, j), eq_ij in Eq_ij.items()),
'rank_transitivity_equal1')
model.addConstrs(((1 - eq_ij) * big_m + R_i[i] >= R_i[j]
for (i, j), eq_ij in Eq_ij.items()),
'rank_transitivity_equal2')
# objective function
objective = LinExpr()
for ij, value in comparisons.items():
if value == 1.0:
objective += -E_ij[ij]
elif value == 0.0:
objective += E_ij[ij]
else:
objective += T_ij_pos[ij] + T_ij_neg[ij]
model.setObjective(objective, GRB.MINIMIZE)
# solve
model.optimize()
# verify abs emulation: one T_ij has to be 0
for ij, value in T_ij_pos.items():
assert value.X == 0 or T_ij_neg[ij] == 0, \
f'T_{ij} pos {value.X} neg {T_ij_neg[ij]}'
# find minimal Rs
model_ = Model('comparison')
model_.setParam('OutputFlag', verbose)
R_i = model_.addVars(nodes,
name='r_i',
vtype=GRB.CONTINUOUS,
lb=0,
ub=len(nodes))
for ((i, j), ge_ij), le_ij in zip(Ge_ij.items(), Le_ij.values()):
if ge_ij.x == 1:
model_.addConstr(R_i[i] >= R_i[j] + 1)
elif le_ij.x == 1:
model_.addConstr(R_i[j] >= R_i[i] + 1)
else:
model_.addConstr(R_i[j] == R_i[i])
model_.setObjective(R_i.sum(), GRB.MINIMIZE)
model_.optimize()
return [model_.getVarByName(f'r_i[{i}]').X for i in range(len(nodes))], \
model.objVal
def find_ranking(comparisons, verbose=False):
"""
Find the least changes to a set of comparisons so that they are consistent
(transitive), it returns a topological ranking.
comparisons A dictionary with tuple keys in the form of (i, j), values
are scalars indicating the probability of i >= j. It is
assumed that comparisons are symmetric. The only way this
function it can handle '=' is by setting the probability to
0.5. In this case the MILP can treat it as > or < (but not
as =).
verbose Whether to print gurobi's progress.
Returns:
A tuple of size two:
0) Ranking derived from topological sort (list of ranks in order of nodes);
1) Sum of absolute changes to the comparisons.
"""
# remove unnecessary variables
comparisons = {(i, j) if i < j else (j, i): value if i < j else 1 - value
for (i, j), value in comparisons.items()}
nodes = np.unique([i for ij in comparisons.keys() for i in ij])
# define variables
model = Model('comparison')
model.setParam('OutputFlag', verbose)
values = np.fromiter(comparisons.values(), dtype=float)
assert values.max() <= 1 and values.min() >= 0
# variables to encode the error of comparisons
E_ij = model.addVars(comparisons.keys(),
name='e_ij',
vtype=GRB.CONTINUOUS,
ub=1.0 - values,
lb=-values)
# variable to encode hard choice of >= and <=
Z_ij = model.addVars(comparisons.keys(), name='z_ij', vtype=GRB.BINARY)
# variables to help with transitivity in non-fully connected graphs
R_i = model.addVars(nodes,
name='r_i',
vtype=GRB.CONTINUOUS,
lb=0,
ub=len(nodes))
# variables to emulate abs
T_ij_pos = {}
T_ij_neg = {}
index = (values != 1) & (values != 0)
T_ij_pos = model.addVars(
(ij for ij, value in comparisons.items() if value not in [0.0, 1.0]),
vtype=GRB.CONTINUOUS,
name='T_ij_pos',
lb=0,
ub=1 - values[index])
T_ij_neg = model.addVars(
(ij for ij, value in comparisons.items() if value not in [0.0, 1.0]),
vtype=GRB.CONTINUOUS,
name='T_ij_neg',
lb=0,
ub=values[index])
model.update()
# emulate abs for non-binary comparisons: E_ij = T_ij_pos - T_ij_neg
model.addConstrs((E_ij[ij] == T_ij_pos[ij] - T_ij_neg[ij]
for ij in T_ij_pos), 'E_ij = T_ij_pos - T_ij_neg')
# hard decision of >= and <=: z_ij == 1 <-> i > j
model.addConstrs((E_ij[ij] + comparisons[ij] - 0.5 >= -1 + z_ij
for ij, z_ij in Z_ij.items()), 'z_ij_upper_bound')
model.addConstrs((E_ij[ij] + comparisons[ij] - 0.5 <= z_ij
for ij, z_ij in Z_ij.items()), 'z_ij_lower_bound')
# transitivity
for (i, j), a in Z_ij.items():
for k in nodes:
j_, k_ = j, k
if j > k:
j_, k_ = k, j
b = Z_ij.get((j_, k_), None)
if b is None:
continue
elif j_ != j:
b = 1 - b
i_, k_ = i, k
if i > k:
i_, k_ = k, i
c = Z_ij.get((i_, k_), None)
if c is None:
continue
elif i_ != i:
c = 1 - c
# a <= b and b <= c -> a <= c
model.addLConstr(a + b, GRB.LESS_EQUAL, 1 + c,
f'transitivity_ge_{i},{j},{k}')
# a >= b and b >= c -> a >= c
model.addLConstr(a + b, GRB.GREATER_EQUAL, c,
f'transitivity_le_{i},{j},{k}')
# transitivity helper (for not-fully connected graphs)
# also provides a latent rank
big_m = len(nodes)
model.addConstrs(((1 - z_ij) * big_m + R_i[i] >= R_i[j] + 1
for (i, j), z_ij in Z_ij.items()),
'rank_transitivity_larger')
model.addConstrs(
(z_ij * big_m + R_i[j] >= R_i[i] + 1 for (i, j), z_ij in Z_ij.items()),
'rank_transitivity_smaller')
# objective function
objective = LinExpr()
for ij, value in comparisons.items():
if value == 1.0:
objective += -E_ij[ij]
elif value == 0.0:
objective += E_ij[ij]
else:
objective += T_ij_pos[ij] + T_ij_neg[ij]
model.setObjective(objective, GRB.MINIMIZE)
# solve
model.optimize()
# verify abs emulation: one T_ij has to be 0
for ij, value in T_ij_pos.items():
assert value.X == 0 or T_ij_neg[ij] == 0, \
f'T_{ij} pos {value.X} neg {T_ij_neg[ij]}'
# find minimal Rs
model_ = Model('comparison')
model_.setParam('OutputFlag', verbose)
R_i = model_.addVars(nodes,
name='r_i',
vtype=GRB.CONTINUOUS,
lb=0,
ub=len(nodes))
for (i, j), z_ij in Z_ij.items():
if z_ij.x == 1:
model_.addConstr(R_i[i] >= R_i[j] + 1)
else:
model_.addConstr(R_i[j] >= R_i[i] + 1)
model_.setObjective(R_i.sum(), GRB.MINIMIZE)
model_.optimize()
return [model_.getVarByName(f'r_i[{i}]').X for i in range(len(nodes))], \
model.objVal
class ILPSolver:
def __init__(
self,
g: DFGraph,
budget: int,
eps_noise=None,
seed_s=None,
integral=True,
imposed_schedule: ImposedSchedule = ImposedSchedule.FULL_SCHEDULE,
solve_r=True,
write_model_file: Optional[PathLike] = None,
gurobi_params: Dict[str, Any] = None):
self.GRB_CONSTRAINED_PRESOLVE_TIME_LIMIT = 300 # todo (paras): read this from gurobi_params
self.gurobi_params = gurobi_params
self.num_threads = self.gurobi_params.get('Threads', 1)
self.model_file = write_model_file
self.seed_s = seed_s
self.integral = integral
self.imposed_schedule = imposed_schedule
self.solve_r = solve_r
self.eps_noise = eps_noise
self.budget = budget
self.g: DFGraph = g
self.solve_time = None
if not self.integral:
assert not self.solve_r, "Can't solve for R if producing a fractional solution"
self.init_constraints = [] # used for seeding the model
self.m = Model("checkpointmip_gc_{}_{}".format(self.g.size,
self.budget))
if gurobi_params is not None:
for k, v in gurobi_params.items():
setattr(self.m.Params, k, v)
T = self.g.size
self.ram_gcd = self.g.ram_gcd(self.budget)
if self.integral:
self.R = self.m.addVars(T, T, name="R", vtype=GRB.BINARY)
self.S = self.m.addVars(T, T, name="S", vtype=GRB.BINARY)
self.Free_E = self.m.addVars(T,
len(self.g.edge_list),
name="FREE_E",
vtype=GRB.BINARY)
else:
self.R = self.m.addVars(T,
T,
name="R",
vtype=GRB.CONTINUOUS,
lb=0.0,
ub=1.0)
self.S = self.m.addVars(T,
T,
name="S",
vtype=GRB.CONTINUOUS,
lb=0.0,
ub=1.0)
self.Free_E = self.m.addVars(T,
len(self.g.edge_list),
name="FREE_E",
vtype=GRB.CONTINUOUS,
lb=0.0,
ub=1.0)
self.U = self.m.addVars(T,
T,
name="U",
lb=0.0,
ub=float(budget) / self.ram_gcd)
for x in range(T):
for y in range(T):
self.m.addLConstr(self.U[x, y], GRB.GREATER_EQUAL, 0)
self.m.addLConstr(self.U[x, y], GRB.LESS_EQUAL,
float(budget) / self.ram_gcd)
def build_model(self):
T = self.g.size
dict_val_div = lambda cost_dict, divisor: {
k: v / divisor
for k, v in cost_dict.items()
}
permute_ram = dict_val_div(self.g.cost_ram, self.ram_gcd)
budget = self.budget / self.ram_gcd
permute_eps = lambda cost_dict, eps: {
k: v * (1. + eps * np.random.randn())
for k, v in cost_dict.items()
}
permute_cpu = dict_val_div(self.g.cost_cpu, self.g.cpu_gcd())
if self.eps_noise:
permute_cpu = permute_eps(permute_cpu, self.eps_noise)
with Timer("Gurobi model construction",
extra_data={
'T': str(T),
'budget': str(budget)
}):
with Timer("Objective construction",
extra_data={
'T': str(T),
'budget': str(budget)
}):
# seed solver with a baseline strategy
if self.seed_s is not None:
for x in range(T):
for y in range(T):
if self.seed_s[x, y] < 1:
self.init_constraints.append(
self.m.addLConstr(self.S[x, y], GRB.EQUAL,
0))
self.m.update()
# define objective function
self.m.setObjective(
quicksum(self.R[t, i] * permute_cpu[i] for t in range(T)
for i in range(T)), GRB.MINIMIZE)
with Timer("Variable initialization",
extra_data={
'T': str(T),
'budget': str(budget)
}):
if self.imposed_schedule == ImposedSchedule.FULL_SCHEDULE:
self.m.addLConstr(
quicksum(self.R[t, i] for t in range(T)
for i in range(t + 1, T)), GRB.EQUAL, 0)
self.m.addLConstr(
quicksum(self.S[t, i] for t in range(T)
for i in range(t, T)), GRB.EQUAL, 0)
self.m.addLConstr(quicksum(self.R[t, t] for t in range(T)),
GRB.EQUAL, T)
elif self.imposed_schedule == ImposedSchedule.COVER_ALL_NODES:
self.m.addLConstr(quicksum(self.S[0, i] for i in range(T)),
GRB.EQUAL, 0)
for i in range(T):
self.m.addLConstr(
quicksum(self.R[t, i] for t in range(T)),
GRB.GREATER_EQUAL, 1)
elif self.imposed_schedule == ImposedSchedule.COVER_LAST_NODE:
self.m.addLConstr(quicksum(self.S[0, i] for i in range(T)),
GRB.EQUAL, 0)
# note: the integrality gap is very large as this constraint
# is only applied to the last node (last column of self.R).
self.m.addLConstr(
quicksum(self.R[t, T - 1] for t in range(T)),
GRB.GREATER_EQUAL, 1)
with Timer("Correctness constraints",
extra_data={
'T': str(T),
'budget': str(budget)
}):
# ensure all checkpoints are in memory
for t in range(T - 1):
for i in range(T):
self.m.addLConstr(self.S[t + 1, i], GRB.LESS_EQUAL,
self.S[t, i] + self.R[t, i])
# ensure all computations are possible
for (u, v) in self.g.edge_list:
for t in range(T):
self.m.addLConstr(self.R[t, v], GRB.LESS_EQUAL,
self.R[t, u] + self.S[t, u])
# define memory constraints
def _num_hazards(t, i, k):
if t + 1 < T:
return 1 - self.R[t, k] + self.S[t + 1, i] + quicksum(
self.R[t, j] for j in self.g.successors(i) if j > k)
return 1 - self.R[t, k] + quicksum(
self.R[t, j] for j in self.g.successors(i) if j > k)
def _max_num_hazards(t, i, k):
num_uses_after_k = sum(1 for j in self.g.successors(i)
if j > k)
if t + 1 < T:
return 2 + num_uses_after_k
return 1 + num_uses_after_k
with Timer("Constraint: upper bound for 1 - Free_E",
extra_data={
'T': str(T),
'budget': str(budget)
}):
for t in range(T):
for eidx, (i, k) in enumerate(self.g.edge_list):
self.m.addLConstr(1 - self.Free_E[t, eidx],
GRB.LESS_EQUAL,
_num_hazards(t, i, k))
with Timer("Constraint: lower bound for 1 - Free_E",
extra_data={
'T': str(T),
'budget': str(budget)
}):
for t in range(T):
for eidx, (i, k) in enumerate(self.g.edge_list):
self.m.addLConstr(
_max_num_hazards(t, i, k) *
(1 - self.Free_E[t, eidx]), GRB.GREATER_EQUAL,
_num_hazards(t, i, k))
with Timer(
"Constraint: initialize memory usage (includes spurious checkpoints)",
extra_data={
'T': str(T),
'budget': str(budget)
}):
for t in range(T):
self.m.addLConstr(
self.U[t,
0], GRB.EQUAL, self.R[t, 0] * permute_ram[0] +
quicksum(self.S[t, i] * permute_ram[i]
for i in range(T)))
with Timer("Constraint: memory recurrence",
extra_data={
'T': str(T),
'budget': str(budget)
}):
for t in range(T):
for k in range(T - 1):
mem_freed = quicksum(
permute_ram[i] * self.Free_E[t, eidx]
for (eidx, i) in self.g.predecessors_indexed(k))
self.m.addLConstr(
self.U[t, k + 1], GRB.EQUAL, self.U[t, k] +
self.R[t, k + 1] * permute_ram[k + 1] - mem_freed)
if self.model_file is not None and self.g.size < 200: # skip for big models to save runtime
with Timer("Saving model",
extra_data={
'T': str(T),
'budget': str(budget)
}):
self.m.write(self.model_file)
return None # return value ensures ray remote call can be chained
def solve(self):
T = self.g.size
with Timer('Gurobi model optimization',
extra_data={
'T': str(T),
'budget': str(self.budget)
}):
if self.seed_s is not None:
self.m.Params.TimeLimit = self.GRB_CONSTRAINED_PRESOLVE_TIME_LIMIT
self.m.optimize()
if self.m.status == GRB.INFEASIBLE:
print(f"Infeasible ILP seed at budget {self.budget:.2E}")
self.m.remove(self.init_constraints)
self.m.Params.TimeLimit = self.gurobi_params.get('TimeLimit', 0)
self.m.message("\n\nRestarting solve\n\n")
with Timer("ILPSolve") as solve_ilp:
self.m.optimize()
self.solve_time = solve_ilp.elapsed
infeasible = (self.m.status == GRB.INFEASIBLE)
if infeasible:
raise ValueError(
"Infeasible model, check constraints carefully. Insufficient memory?"
)
if self.m.solCount < 1:
raise ValueError(
f"Model status is {self.m.status} (not infeasible), but solCount is {self.m.solCount}"
)
Rout = np.zeros((T, T),
dtype=remat.core.utils.solver_common.SOLVER_DTYPE
if self.integral else np.float)
Sout = np.zeros((T, T),
dtype=remat.core.utils.solver_common.SOLVER_DTYPE
if self.integral else np.float)
Uout = np.zeros((T, T),
dtype=remat.core.utils.solver_common.SOLVER_DTYPE
if self.integral else np.float)
Free_Eout = np.zeros((T, len(self.g.edge_list)),
dtype=remat.core.utils.solver_common.SOLVER_DTYPE)
solver_dtype_cast = int if self.integral else float
try:
for t in range(T):
for i in range(T):
try:
Rout[t][i] = solver_dtype_cast(self.R[t, i].X)
except (AttributeError, TypeError) as e:
Rout[t][i] = solver_dtype_cast(self.R[t, i])
try:
Sout[t][i] = solver_dtype_cast(self.S[t, i])
except (AttributeError, TypeError) as e:
Sout[t][i] = solver_dtype_cast(self.S[t, i].X)
try:
Uout[t][i] = self.U[t, i].X * self.ram_gcd
except (AttributeError, TypeError) as e:
Uout[t][i] = self.U[t, i] * self.ram_gcd
for e in range(len(self.g.edge_list)):
try:
Free_Eout[t][e] = solver_dtype_cast(self.Free_E[t,
e].X)
except (AttributeError, TypeError) as e:
Free_Eout[t][e] = solver_dtype_cast(self.Free_E[t, e])
except AttributeError as e:
logging.exception(e)
return None, None, None, None
# prune R using closed-form solver
if self.solve_r and self.integral:
Rout = solve_r_opt(self.g, Sout)
return Rout, Sout, Uout, Free_Eout
class MaxBatchILPSolver:
def __init__(
self,
g: DFGraph,
budget: int,
eps_noise=None,
model_file=None,
remote=False,
gurobi_params: Dict[str, Any] = None,
cpu_fwd_factor: int = 2,
):
self.cpu_fwd_factor = cpu_fwd_factor
self.logger = logging.getLogger(__name__)
self.remote = remote
self.profiler = functools.partial(Timer, print_results=True)
self.gurobi_params = gurobi_params
self.num_threads = self.gurobi_params.get("Threads", 1)
self.model_file = model_file
self.eps_noise = eps_noise
self.budget = budget
self.g = g
self.solve_time = None
self.init_constraints = [] # used for seeding the model
self.m = Model("checkpointmip_gc_maxbs_{}_{}".format(self.g.size, self.budget))
if gurobi_params is not None:
for k, v in gurobi_params.items():
setattr(self.m.Params, k, v)
T = self.g.size
CPU_VALS = list(self.g.cost_cpu.values())
RAM_VALS = list(self.g.cost_ram.values())
self.logger.info(
"RAM: [{:.2E}, {:.2E}], {:.2E} +- {:.2E}".format(
np.min(RAM_VALS), np.max(RAM_VALS), np.mean(RAM_VALS), np.std(RAM_VALS)
)
)
self.logger.info(
"CPU: [{:.2E}, {:.2E}], {:.2E} +- {:.2E}".format(
np.min(CPU_VALS), np.max(CPU_VALS), np.mean(CPU_VALS), np.std(CPU_VALS)
)
)
self.ram_gcd = int(max(self.g.ram_gcd(self.budget), 1))
self.cpu_gcd = int(max(self.g.cpu_gcd(), 1))
self.logger.info("ram_gcd = {} cpu_gcd = {}".format(self.ram_gcd, self.cpu_gcd))
self.batch_size = self.m.addVar(lb=1, ub=1024 * 16, name="batch_size")
self.R = self.m.addVars(T, T, name="R", vtype=GRB.BINARY)
self.S = self.m.addVars(T, T, name="S", vtype=GRB.BINARY)
self.Free_E = self.m.addVars(T, len(self.g.edge_list), name="FREE_E", vtype=GRB.BINARY)
self.U = self.m.addVars(T, T, name="U", lb=0.0, ub=self.budget)
for x in range(T):
for y in range(T):
self.m.addLConstr(self.U[x, y], GRB.GREATER_EQUAL, 0)
self.m.addLConstr(self.U[x, y], GRB.LESS_EQUAL, float(budget) / self.ram_gcd)
def build_model(self):
T = self.g.size
dict_val_div = lambda cost_dict, divisor: {k: np.ceil(v / divisor) for k, v in cost_dict.items()}
permute_ram = dict_val_div(self.g.cost_ram, self.ram_gcd)
budget = self.budget / self.ram_gcd
permute_eps = lambda cost_dict, eps: {k: v * (1.0 + eps * np.random.randn()) for k, v in cost_dict.items()}
permute_cpu = dict_val_div(self.g.cost_cpu, self.cpu_gcd)
if self.eps_noise:
permute_cpu = permute_eps(permute_cpu, self.eps_noise)
with self.profiler("Gurobi model construction", extra_data={"T": str(T), "budget": str(budget)}):
with self.profiler("Objective construction", extra_data={"T": str(T), "budget": str(budget)}):
# define objective function
self.m.setObjective(self.batch_size, GRB.MAXIMIZE)
with self.profiler("Variable initialization", extra_data={"T": str(T), "budget": str(budget)}):
self.m.addLConstr(quicksum(self.R[t, i] for t in range(T) for i in range(t + 1, T)), GRB.EQUAL, 0)
self.m.addLConstr(quicksum(self.S[t, i] for t in range(T) for i in range(t, T)), GRB.EQUAL, 0)
self.m.addLConstr(quicksum(self.R[t, t] for t in range(T)), GRB.EQUAL, T)
with self.profiler("Correctness constraints", extra_data={"T": str(T), "budget": str(budget)}):
# ensure all checkpoints are in memory
for t in range(T - 1):
for i in range(T):
self.m.addLConstr(self.S[t + 1, i], GRB.LESS_EQUAL, self.S[t, i] + self.R[t, i])
# ensure all computations are possible
for (u, v) in self.g.edge_list:
for t in range(T):
self.m.addLConstr(self.R[t, v], GRB.LESS_EQUAL, self.R[t, u] + self.S[t, u])
# define memory constraints
def _num_hazards(t, i, k):
if t + 1 < T:
return 1 - self.R[t, k] + self.S[t + 1, i] + quicksum(self.R[t, j] for j in self.g.successors(i) if j > k)
return 1 - self.R[t, k] + quicksum(self.R[t, j] for j in self.g.successors(i) if j > k)
def _max_num_hazards(t, i, k):
num_uses_after_k = sum(1 for j in self.g.successors(i) if j > k)
if t + 1 < T:
return 2 + num_uses_after_k
return 1 + num_uses_after_k
with self.profiler("Constraint: upper bound for 1 - Free_E", extra_data={"T": str(T), "budget": str(budget)}):
for t in range(T):
for eidx, (i, k) in enumerate(self.g.edge_list):
self.m.addLConstr(1 - self.Free_E[t, eidx], GRB.LESS_EQUAL, _num_hazards(t, i, k))
with self.profiler("Constraint: lower bound for 1 - Free_E", extra_data={"T": str(T), "budget": str(budget)}):
for t in range(T):
for eidx, (i, k) in enumerate(self.g.edge_list):
self.m.addLConstr(
_max_num_hazards(t, i, k) * (1 - self.Free_E[t, eidx]), GRB.GREATER_EQUAL, _num_hazards(t, i, k)
)
with self.profiler(
"Constraint: initialize memory usage (includes spurious checkpoints)",
extra_data={"T": str(T), "budget": str(budget)},
):
for t in range(T):
self.m.addConstr(
self.U[t, 0]
== self.batch_size
* (self.R[t, 0] * permute_ram[0] + quicksum(self.S[t, i] * permute_ram[i] for i in range(T))),
name="init_mem",
)
with self.profiler("Constraint: memory recurrence", extra_data={"T": str(T), "budget": str(budget)}):
for t in range(T):
for k in range(T - 1):
mem_freed = quicksum(
permute_ram[i] * self.Free_E[t, eidx] for (eidx, i) in self.g.predecessors_indexed(k)
)
self.m.addConstr(
self.U[t, k + 1]
== self.U[t, k] + self.batch_size * (self.R[t, k + 1] * permute_ram[k + 1] - mem_freed),
name="update_mem",
)
if self.cpu_fwd_factor:
with self.profiler("Constraint: recomputation overhead"):
compute_fwd = sum([permute_cpu[i] for i in self.g.vfwd])
bwd_compute = sum([permute_cpu[i] for i in self.g.v if i not in self.g.vfwd])
max_mem = self.cpu_fwd_factor * compute_fwd + bwd_compute
self.logger.info("Solver using compute overhead ceiling of {}".format(max_mem))
self.m.addConstr(
quicksum(self.R[t, i] * permute_cpu[i] for t in range(T) for i in range(T)) <= max_mem,
name="limit_cpu",
)
else:
self.logger.info("No compute limit")
with self.profiler("Model update"):
self.m.update()
if self.model_file is not None and self.g.size < 200: # skip for big models to save runtime
with self.profiler("Saving model", extra_data={"T": str(T), "budget": str(budget)}):
self.m.write(self.model_file)
return None # return value ensures ray remote call can be chained
def solve(self):
T = self.g.size
with self.profiler("Gurobi model optimization", extra_data={"T": str(T), "budget": str(self.budget)}):
with Timer("ILPSolve") as solve_ilp:
self.m.optimize()
self.solve_time = solve_ilp.elapsed
infeasible = self.m.status == GRB.INFEASIBLE
try:
_ = self.R[0, 0].X
_ = self.S[0, 0].X
_ = self.U[0, 0].X
_ = self.batch_size.X
except AttributeError as e:
infeasible = True
if infeasible:
raise ValueError("Infeasible model, check constraints carefully. Insufficient memory?")
Rout = np.zeros((T, T), dtype=SOLVER_DTYPE)
Sout = np.zeros((T, T), dtype=SOLVER_DTYPE)
Uout = np.zeros((T, T), dtype=SOLVER_DTYPE)
Free_Eout = np.zeros((T, len(self.g.edge_list)), dtype=SOLVER_DTYPE)
batch_size = self.batch_size.X
try:
for t in range(T):
for i in range(T):
Rout[t][i] = int(self.R[t, i].X)
Sout[t][i] = int(self.S[t, i].X)
Uout[t][i] = self.U[t, i].X * self.ram_gcd
for e in range(len(self.g.edge_list)):
Free_Eout[t][e] = int(self.Free_E[t, e].X)
except AttributeError as e:
logging.exception(e)
return None, None, None, None
Rout = solve_r_opt(self.g, Sout) # prune R using optimal recomputation solver
ilp_aux_data = ILPAuxData(
U=Uout,
Free_E=Free_Eout,
ilp_approx=False,
ilp_time_limit=0,
ilp_eps_noise=0,
ilp_num_constraints=self.m.numConstrs,
ilp_num_variables=self.m.numVars,
)
schedule, aux_data = schedule_from_rs(self.g, Rout, Sout)
return (
ScheduledResult(
solve_strategy=SolveStrategy.OPTIMAL_ILP_GC,
solver_budget=self.budget,
feasible=True,
schedule=schedule,
schedule_aux_data=aux_data,
solve_time_s=self.solve_time,
ilp_aux_data=ilp_aux_data,
),
batch_size,
)
class ILPSolver:
def __init__(self, si: SolverInfo, budget, gurobi_params: Dict[str, Any] = None, ablation=False, overhead=False):
self.gurobi_params = gurobi_params
self.num_threads = self.gurobi_params.get('Threads', 1)
self.budget = int(budget * MEM_GCD_MULTIPLIER * GB_TO_KB)
self.si: SolverInfo = si
self.solve_time = None
self.ablation = ablation
self.overhead = overhead
V = self.si.loss + 1
T = len(self.si.nodes) - self.si.loss
Y = 3
budget = self.budget
self.m = Model("monet{}".format(self.budget))
if gurobi_params is not None:
for k, v in gurobi_params.items():
setattr(self.m.Params, k, v)
self.ram = np.array([math.ceil(self.si.nodes[i].mem*MEM_GCD_MULTIPLIER/1024) for i in self.si.nodes]) # Convert to KB
self.cpu = dict(( i, [math.ceil(val*CPU_GCD_MULTIPLIER) for val in self.si.nodes[i].workspace_compute] ) for i in self.si.nodes)
self.cpu_recompute = dict(( i, [math.ceil(val*CPU_GCD_MULTIPLIER) for val in self.si.nodes[i].recompute_workspace_compute] ) for i in self.si.nodes)
self.cpu_inplace = dict(( i, [math.ceil(val*CPU_GCD_MULTIPLIER) for val in self.si.nodes[i].inplace_workspace_compute] ) for i in self.si.nodes)
self.R = self.m.addVars(T, V, name="R", vtype=GRB.BINARY) # Recomputation
self.P = self.m.addVars(T, V, name="P", vtype=GRB.BINARY) # In-memory
self.S = self.m.addVars(T+1, V, name="S", vtype=GRB.BINARY) # Stored
self.M = self.m.addVars(T, Y, name="M", vtype=GRB.BINARY) # Backward operator implementation
self.SM = self.m.addVars(T, V, Y, name="SM", vtype=GRB.BINARY) # Linearization of S * M
if self.si.select_conv_algo:
self.RF = self.m.addVars(T, len(self.si.conv_list), self.si.num_conv_algos, name="RF", vtype=GRB.BINARY) # Conv operator implementations
if self.si.do_inplace:
self.IP = self.m.addVars(T, V, name="IP", vtype=GRB.BINARY) # In-place
def cmem(self, j, recompute=False, inplace=False):
if inplace:
return [math.ceil(val * MEM_GCD_MULTIPLIER * 1024) for val in self.si.nodes[j].inplace_workspace_mem]
else:
if recompute:
return [math.ceil(val * MEM_GCD_MULTIPLIER * 1024) for val in self.si.nodes[j].recompute_workspace_mem]
else:
return [math.ceil(val * MEM_GCD_MULTIPLIER * 1024) for val in self.si.nodes[j].workspace_mem]
def local_memory(self, j):
return math.ceil(self.si.nodes[j].local_memory * MEM_GCD_MULTIPLIER / 1024)
def fixed_ram(self, j):
return math.ceil(self.si.nodes[j].fixed_mem * MEM_GCD_MULTIPLIER / 1024)
def build_model(self):
V = self.si.loss + 1
T = len(self.si.nodes) - self.si.loss
Y = 3
budget = self.budget
# define objective function
if self.si.select_conv_algo:
fwd_compute = quicksum(self.R[0, i] * self.cpu[i][0] for i in range(V) if i not in self.si.conv_list)
fwd_recompute = quicksum(quicksum(self.R[t, i]*self.cpu_recompute[i][0] for t in range(1,T)) for i in range(V) if i not in self.si.conv_list)
conv_fwd_compute = quicksum(quicksum(self.RF[t,self.si.conv_list[i],c] * self.cpu[i][c] for t in range(T)) for i in self.si.conv_list if i<=self.si.loss for c in range(len(self.cpu[i]))) # conv's compute and recompute same
if self.si.do_inplace:
inplace_fwd_compute = quicksum(quicksum(self.IP[t,i] * (self.cpu_inplace[i][0]-self.cpu[i][0]) for t in range(T)) for i in self.si.inplace_list) # relu's compute and recompute same
conv_bwd_compute = quicksum(quicksum(self.RF[t,self.si.conv_list[i],c] * self.cpu[i][c] for t in range(T)) for i in self.si.conv_list if i>self.si.loss for c in range(len(self.cpu[i]))) # conv's compute and recompute same
bwd_compute = quicksum(self.M[t, p] * self.cpu[t+V-1][p] for t in range(T) for p in range(Y) if (p<len(self.cpu[t+V-1]) and (t+V-1) not in self.si.conv_list))
if self.si.do_inplace:
compute_fwd = fwd_compute + fwd_recompute + conv_fwd_compute + inplace_fwd_compute
else:
compute_fwd = fwd_compute + fwd_recompute + conv_fwd_compute
compute_bwd = conv_bwd_compute + bwd_compute
compute_cost = compute_fwd + compute_bwd
else:
fwd_compute = quicksum(self.R[0, i] * self.cpu[i][0] for i in range(V))
fwd_recompute = quicksum(self.R[t, i] * self.cpu_recompute[i][0] for t in range(1,T) for i in range(V))
if self.si.do_inplace:
inplace_fwd_compute = quicksum(quicksum(self.IP[t,i] * (self.cpu_inplace[i][0]-self.cpu[i][0]) for t in range(T)) for i in self.si.inplace_list) # relu's compute and recompute same
bwd_compute = quicksum(self.M[t, p] * self.cpu[t+V-1][p] for t in range(T) for p in range(Y) if p<len(self.cpu[t+V-1]))
if self.si.do_inplace:
compute_fwd = fwd_compute + fwd_recompute + inplace_fwd_compute
else:
compute_fwd = fwd_compute + fwd_recompute
compute_bwd = bwd_compute
compute_cost = compute_fwd + compute_bwd
self.m.setObjective(compute_cost, GRB.MINIMIZE)
self.compute_cost = compute_cost
self.compute_fwd = compute_fwd
self.compute_bwd = compute_bwd
# Add in-place constraints
for t in range(T):
for v in range(V):
if self.si.do_inplace:
self.m.addLConstr(self.P[t,v], GRB.LESS_EQUAL, self.R[t,v])
self.m.addLConstr(self.IP[t,v], GRB.LESS_EQUAL, self.R[t,v])
# print(list(self.si.inplace_list.keys()))
if v not in self.si.inplace_list:
self.m.addLConstr(self.IP[t,v], GRB.EQUAL, 0)
elif not isinstance(self.si.nodes[v], IntNode):
# If IP[v], then P[u] = 0, else P[u] = R[u]
self.m.addLConstr(self.P[t,self.si.inplace_list[v]], GRB.GREATER_EQUAL, self.R[t,self.si.inplace_list[v]] - 2*self.IP[t,v])
self.m.addLConstr(self.P[t,self.si.inplace_list[v]], GRB.LESS_EQUAL, 2 - 2*self.IP[t,v])
self.m.addLConstr(self.S[t+1,self.si.inplace_list[v]], GRB.LESS_EQUAL, 2 - 2*self.IP[t,v])
else:
self.m.addLConstr(self.P[t,v], GRB.EQUAL, self.R[t,v])
# store nothing in the beginning
for i in range(V):
self.m.addLConstr(self.S[0,i] , GRB.EQUAL, 0)
# Recompute full forward
for i in range(V):
if not isinstance(self.si.nodes[i], IntNode):
self.m.addLConstr(self.R[0,i] , GRB.EQUAL, 1)
# All M which have no paths are set as 0
self.m.addLConstr(self.M[0,0], GRB.EQUAL, 1)
for p in range(1,Y):
self.m.addLConstr(self.M[0,p], GRB.EQUAL, 0)
for t in range(1, T):
bkwd_t = V + t - 1
paths = len(self.si.nodes[bkwd_t].args)
for p in range(paths, Y):
self.m.addLConstr(self.M[t,p], GRB.EQUAL, 0)
# Set failed conv's RF to be 0
if self.si.select_conv_algo:
for i in self.si.conv_list:
for c in range(self.si.num_conv_algos):
if c >= len(self.cpu[i]) or self.cpu[i][c] < 0:
for t in range(T):
self.m.addLConstr(self.RF[t,self.si.conv_list[i],c], GRB.EQUAL, 0)
# create constraints for boolean multiplication linearization
for p in range(Y):
for i in range(V):
for t in range(T):
self.m.addLConstr(self.SM[t,i, p], GRB.GREATER_EQUAL, self.M[t,p] + self.S[t+1,i] - 1)
self.m.addLConstr(self.SM[t,i, p], GRB.LESS_EQUAL, self.M[t,p])
self.m.addLConstr(self.SM[t,i, p], GRB.LESS_EQUAL, self.S[t+1,i])
if self.ablation:
print("Doing ablation")
# Disable all recomputation
for t in range(1,T):
for v in range(V):
self.m.addLConstr(self.R[t,v], GRB.EQUAL, 0)
# Fix all operations to be inplace
for t in range(T):
for v in range(V):
if self.si.do_inplace and v in self.si.inplace_list:
self.m.addLConstr(self.IP[t,v], GRB.EQUAL, self.R[t,v])
# Correctness constraints
# Ensure all checkpoints are in memory
for t in range(T):
for i in range(V):
self.m.addLConstr(self.S[t + 1, i], GRB.LESS_EQUAL, self.S[t, i] + self.P[t, i])
# At least one path should be used
for t in range(T):
self.m.addLConstr(quicksum(self.M[t,p] for p in range(Y)), GRB.GREATER_EQUAL, 1)
# Ensure all computations are possible
for (u, v, p) in self.si.edge_list:
if u < V and v <V:
for t in range(T):
self.m.addLConstr(self.R[t, v], GRB.LESS_EQUAL, self.R[t, u] + self.S[t, u])
if u < V and v >= V:
t = v + 1 - V
self.m.addLConstr(self.M[t,p], GRB.LESS_EQUAL, self.P[t, u] + self.S[t+1, u])
# Ensure that new nodes only computed if main node computed
for t in range(T):
for i in range(V):
if self.si.nodes[i].has_intermediates:
for (_,intid) in self.si.nodes[i].intermediates:
self.m.addLConstr(self.R[t, intid], GRB.LESS_EQUAL, self.R[t,i])
if self.si.do_inplace and i in self.si.inplace_list:
self.m.addLConstr(self.IP[:,intid], GRB.GREATER_EQUAL, self.IP[:,i] + self.R[:,intid] -1)
self.m.addLConstr(self.IP[:,intid], GRB.LESS_EQUAL, self.IP[:,i])
# Constraints for conv selection
if self.si.select_conv_algo:
for i in self.si.conv_list:
if i < self.si.loss:
for t in range(T):
self.m.addLConstr(quicksum(self.RF[t,self.si.conv_list[i],c] for c in range(self.si.num_conv_algos)), GRB.GREATER_EQUAL, self.R[t,i])
else:
t_bwdi = i + 1 - V
self.m.addLConstr(quicksum(self.RF[t_bwdi,self.si.conv_list[i],c] for c in range(self.si.num_conv_algos)), GRB.GREATER_EQUAL, 1)
for c in range(self.si.num_conv_algos):
for t in range(t_bwdi):
self.m.addLConstr(self.RF[t,self.si.conv_list[i],c], GRB.EQUAL, 0)
for t in range(t_bwdi+1, T):
self.m.addLConstr(self.RF[t,self.si.conv_list[i],c], GRB.EQUAL, 0)
# Memory constraints
for t in range(T):
bkwd_t = V + t - 1
bwd_node = self.si.nodes[bkwd_t]
for t in range(T):
bkwd_t = V + t - 1
bwd_node = self.si.nodes[bkwd_t]
bwd_local_memory = self.local_memory(bkwd_t) if t!=0 else 0
gradm = bwd_local_memory + self.fixed_ram(bkwd_t-1)
for j in range(V):
keep_tensors = [[fwdin for fwdin in bwd_node.args[p] if fwdin <= self.si.loss] for p in range(len(bwd_node.args))]
# Forward checkpoint constraint
self.m.addLConstr( gradm +
quicksum(self.S[t,i]*self.ram[i] for i in range(j, V)) +
quicksum(self.S[t+1,i]*self.ram[i] for i in range(j) if t<T-1) +
quicksum( quicksum( self.M[t,p]*self.ram[i] for i in range(j) if i in keep_tensors[p]) for p in range(len(bwd_node.args)) ) -
quicksum( quicksum( self.SM[t,i, p]*self.ram[i] for i in range(j) if i in keep_tensors[p]) for p in range(len(bwd_node.args)) ),
GRB.LESS_EQUAL, budget)
workspace_mem = self.cmem(j) if t==0 else self.cmem(j, recompute=True, inplace=False)
if self.si.select_conv_algo and j in self.si.conv_list:
wm = quicksum(self.RF[t,self.si.conv_list[j],c] * workspace_mem[c] for c in range(len(workspace_mem)))
elif self.si.do_inplace and j in self.si.inplace_list:
inplace_workspace_mem = self.cmem(j, recompute=False, inplace=True)
wm = self.IP[t,j]*(inplace_workspace_mem[0] - workspace_mem[0]) + self.R[t,j] * workspace_mem[0]
else:
wm = self.R[t,j] * workspace_mem[0]
# Forward recomputation constraint
self.m.addLConstr( gradm + wm + self.R[t,j]*self.ram[j] + self.R[t,j]*self.local_memory(j) +
quicksum(self.S[t,i]*self.ram[i] for i in range(j+1, V)) +
quicksum(self.S[t+ 1,i]*self.ram[i] for i in range(j) if i not in self.si.nodes[j].local_tensors) +
quicksum( quicksum( self.M[t,p]*self.ram[i] for i in range(j) if i in keep_tensors[p]) for p in range(len(bwd_node.args)) ) -
quicksum( quicksum( self.SM[t,i, p]*self.ram[i] for i in range(j) if i in keep_tensors[p]) for p in range(len(bwd_node.args)) ),
GRB.LESS_EQUAL, budget)
if t>0:
if self.si.select_conv_algo and bkwd_t in self.si.conv_list:
wmem = self.cmem(bkwd_t)
bwd_wm = quicksum(self.RF[t,self.si.conv_list[bkwd_t],c] * wmem[c] for c in range(len(wmem)))
else:
wmem = self.cmem(bkwd_t)
bwd_wm = quicksum(self.M[t,p]*wmem[p] for p in range(len(bwd_node.args)))
rammem = [sum([self.ram[i] for i in bwd_node.args[p] if i<V]) for p in range(len(bwd_node.args))]
# Backward constraint
self.m.addLConstr( bwd_wm +
bwd_local_memory + self.fixed_ram(bkwd_t) + self.ram[bkwd_t] +
quicksum(self.M[t,p]*rammem[p] for p in range(len(bwd_node.args))) +
quicksum(quicksum(self.SM[t,i,p]*self.ram[i] for i in range(V) if i not in bwd_node.args[p]) for p in range(len(bwd_node.args))),
GRB.LESS_EQUAL, budget)
return None
def solve(self):
from time import time
V = self.si.loss + 1
T = len(self.si.nodes) - self.si.loss
Y = 3
budget = self.budget
self.m.Params.TimeLimit = self.time_limit
t0 = time()
self.m.optimize()
self.solve_time = time() - t0
infeasible = (self.m.status == GRB.INFEASIBLE)
if infeasible:
self.m.computeIIS()
self.m.write("model.ilp")
raise ValueError("Infeasible model, check constraints carefully. Insufficient memory?")
if self.m.solCount < 1:
raise ValueError(f"Model status is {self.m.status} (not infeasible), but solCount is {self.m.solCount}")
Rout = np.zeros((T, V), dtype=bool)
Pout = np.zeros((T, V), dtype=bool)
Sout = np.zeros((T+1, V), dtype=bool)
Mout = np.zeros((T, Y), dtype=bool)
SMout = np.zeros((T, V, Y), dtype=bool)
RFout = np.zeros((T, len(self.si.conv_list), self.si.num_conv_algos), dtype=bool)
IPout = np.zeros((T, V), dtype=bool)
try:
for t in range(T):
for i in range(V):
Rout[t][i] = round(self.R[t, i].X)
Pout[t][i] = round(self.P[t, i].X)
Sout[t][i] = round(self.S[t, i].X)
if self.si.do_inplace:
IPout[t][i] = round(self.IP[t, i].X)
for p in range(Y):
SMout[t][i][p] = round(self.SM[t, i, p].X)
for p in range(Y):
Mout[t][p] = round(self.M[t, p].X)
if self.si.select_conv_algo:
for i in range(len(self.si.conv_list)):
for c in range(self.si.num_conv_algos):
RFout[t][i][c] = round(self.RF[t,i,c].X)
for i in range(V):
Sout[T][i] = round(self.S[T, i].X)
except AttributeError as e:
logging.exception(e)
return None, None, None, None, None, None, None
solution = Solution(Rout, Sout, Mout, RFout, IPout, Pout, self.solve_time, -1, -1, -1)
return solution
def run_ilp(tint, remaining_rids, incomp_rids, ilp_settings, log_prefix):
# Variables directly based on the input ------------------------------------
# I[i,j] = 1 if reads[i]['data'][j]==1 and 0 if reads[i]['data'][j]==0 or 2
# C[i,j] = 1 if exon j is between the first and last exons (inclusively)
# covered by read i and is not in read i but can be turned into a 1
ISOFORM_INDEX_START = 1
M = len(tint['segs'])
MAX_ISOFORM_LG = sum(seg[2] for seg in tint['segs'])
I = tint['ilp_data']['I']
C = tint['ilp_data']['C']
INCOMP_READ_PAIRS = incomp_rids
GARBAGE_COST = tint['ilp_data']['garbage_cost']
informative = informative_segs(tint, remaining_rids)
# ILP model ------------------------------------------------------
ILP_ISOFORMS = Model('isoforms_v8_20210209')
ILP_ISOFORMS.setParam('OutputFlag', 0)
ILP_ISOFORMS.setParam(GRB.Param.Threads, ilp_settings['threads'])
# Decision variables
# R2I[i,k] = 1 if read i assigned to isoform k
R2I = {}
R2I_C1 = {} # Constraint enforcing that each read is assigned to exactly one isoform
for i in remaining_rids:
R2I[i] = {}
for k in range(ilp_settings['K']):
R2I[i][k] = ILP_ISOFORMS.addVar(
vtype=GRB.BINARY,
name='R2I[{i}][{k}]'.format(i=i, k=k)
)
R2I_C1[i] = ILP_ISOFORMS.addLConstr(
lhs=quicksum(R2I[i][k] for k in range(0, ilp_settings['K'])),
sense=GRB.EQUAL,
rhs=1,
name='R2I_C1[{i}]'.format(i=i)
)
# Implied variable: canonical exons presence in isoforms
# E2I[j,k] = 1 if canonical exon j is in isoform k
# E2I_min[j,k] = 1 if canonical exon j is in isoform k and is shared by all reads of that isoform
# E2IR[j,k,i] = 1 if read i assigned to isoform k AND exon j covered by read i
# Auxiliary variable
# E2IR[j,k,i] = R2I[i,k] AND I[i,j]
# E2I[j,k] = max over all reads i of E2IR[j,k,i]
# E2I_min[j,k] = min over all reads i of E2IR[j,k,i]
E2I = {}
E2I_C1 = {}
E2I_min = {}
E2I_min_C1 = {}
E2IR = {}
E2IR_C1 = {}
for j in range(0, M):
if not informative[j]:
continue
E2I[j] = {}
E2I_C1[j] = {}
E2I_min[j] = {}
E2I_min_C1[j] = {}
E2IR[j] = {}
E2IR_C1[j] = {}
# No exon is assignd to the garbage isoform
E2I[j][0] = ILP_ISOFORMS.addVar(
vtype=GRB.BINARY,
name='E2I[{j}][{k}]'.format(j=j, k=0)
)
E2I_C1[j][0] = ILP_ISOFORMS.addLConstr(
lhs=E2I[j][0],
sense=GRB.EQUAL,
rhs=0,
name='E2I_C1[{j}][{k}]'.format(j=j, k=0)
)
# We start assigning exons from the first isoform
for k in range(ISOFORM_INDEX_START, ilp_settings['K']):
E2I[j][k] = ILP_ISOFORMS.addVar(
vtype=GRB.BINARY,
name='E2I[{j}][{k}]'.format(j=j, k=k)
)
E2I_min[j][k] = ILP_ISOFORMS.addVar(
vtype=GRB.BINARY,
name='E2I_min[{j}][{k}]'.format(j=j, k=k)
)
E2IR[j][k] = {}
E2IR_C1[j][k] = {}
for i in remaining_rids:
E2IR[j][k][i] = ILP_ISOFORMS.addVar(
vtype=GRB.BINARY,
name='E2IR[{j}][{k}][{i}]'.format(j=j, k=k, i=i)
)
E2IR_C1[j][k][i] = ILP_ISOFORMS.addLConstr(
lhs=E2IR[j][k][i],
sense=GRB.EQUAL,
rhs=R2I[i][k]*I[i][j],
name='E2IR_C1[{j}][{k}][{i}]'.format(j=j, k=k, i=i)
)
E2I_C1[j][k] = ILP_ISOFORMS.addGenConstrMax(
resvar=E2I[j][k],
vars=[E2IR[j][k][i] for i in remaining_rids],
constant=0.0,
name='E2I_C1[{j}][{k}]'.format(j=j, k=k)
)
E2I_min_C1[j][k] = ILP_ISOFORMS.addGenConstrMin(
resvar=E2I_min[j][k],
vars=[E2IR[j][k][i] for i in remaining_rids],
constant=0.0,
name='E2I_min_C1[{j}][{k}]'.format(j=j, k=k)
)
# Adding constraints for unaligned gaps
# If read i is assigned to isoform k, and reads[i]['gaps'] contains ((j1,j2),l), and
# the sum of the lengths of exons in isoform k between exons j1 and j2 is L
# then (1-EPSILON)L <= l <= (1+EPSILON)L
# GAPI[(j1,j2,k)] = sum of the length of the exons between exons j1 and j2 (inclusively) in isoform k
GAPI = {}
GAPI_C1 = {} # Constraint fixing the value of GAPI
GAPR_C1 = {} # Constraint ensuring that the unaligned gap is not too short for every isoform and gap
GAPR_C2 = {} # Constraint ensuring that the unaligned gap is not too long for every isoform and gap
for i in remaining_rids:
for ((j1, j2), l) in tint['reads'][tint['read_reps'][i][0]]['gaps'].items():
# No such constraint on the garbage isoform if any
for k in range(ISOFORM_INDEX_START, ilp_settings['K']):
if not (j1, j2, k) in GAPI:
assert informative[j1 % M]
assert informative[j2 % M]
assert not any(informative[j+1:j2])
GAPI[(j1, j2, k)] = ILP_ISOFORMS.addVar(
vtype=GRB.INTEGER,
name='GAPI[({j1},{j2},{k})]'.format(j1=j1, j2=j2, k=k)
)
GAPI_C1[(j1, j2, k)] = ILP_ISOFORMS.addLConstr(
lhs=GAPI[(j1, j2, k)],
sense=GRB.EQUAL,
rhs=quicksum(E2I[j][k]*tint['segs'][j][2]
for j in range(j1+1, j2) if informative[j]),
name='GAPI_C1[({j1},{j2},{k})]'.format(
j1=j1, j2=j2, k=k)
)
GAPR_C1[(i, j1, j2, k)] = ILP_ISOFORMS.addLConstr(
lhs=(1.0-ilp_settings['epsilon'])*GAPI[(j1, j2, k)] -
ilp_settings['offset']-((1-R2I[i][k])*MAX_ISOFORM_LG),
sense=GRB.LESS_EQUAL,
rhs=l,
name='GAPR_C1[({i},{j1},{j2},{k})]'.format(
i=i, j1=j1, j2=j2, k=k)
)
GAPR_C2[(i, j1, j2, k)] = ILP_ISOFORMS.addLConstr(
lhs=(1.0+ilp_settings['epsilon'])*GAPI[(j1, j2, k)] +
ilp_settings['offset']+((1-R2I[i][k])*MAX_ISOFORM_LG),
sense=GRB.GREATER_EQUAL,
rhs=l,
name='GAPR_C2[({i},{j1},{j2},{k})]'.format(
i=i, j1=j1, j2=j2, k=k)
)
# Adding constraints for incompatible read pairs
INCOMP_READ_PAIRS_C1 = {}
for (i1, i2) in INCOMP_READ_PAIRS:
if not (i1 in remaining_rids and i2 in remaining_rids):
continue
# Again, no such constraint on the garbage isoform if any
for k in range(ISOFORM_INDEX_START, ilp_settings['K']):
INCOMP_READ_PAIRS_C1[(i1, i2, k)] = ILP_ISOFORMS.addLConstr(
lhs=R2I[i1][k]+R2I[i2][k],
sense=GRB.LESS_EQUAL,
rhs=1,
name='INCOMP_READ_PAIRS_C1[({i1},{i2},{k})]'.format(
i1=i1, i2=i2, k=k)
)
# [OPTIONAL] Labeling non-garbage isoforms by their exon content and forcing them to occur in increasing label order
# LABEL_I = {}
# LABEL_I_C1 = {}
# LABEL_I_C2 = {}
# for k in range(ISOFORM_INDEX_START,ilp_settings['K']):
# LABEL_I[k] = ILP_ISOFORMS.addVar(
# vtype = GRB.INTEGER,
# name = 'LABEL_I[{k}]'.format(k=k)
# )
# LABEL_I_C1[k] = ILP_ISOFORMS.addLConstr(
# lhs = LABEL_I[k],
# sense = GRB.EQUAL,
# rhs = quicksum(E2I[j][k]*(2**j) for j in range(0,M)),
# name = 'LABEL_I_C1[{k}]'.format(k =k)
# )
# if k > ISOFORM_INDEX_START:
# LABEL_I_C2[k] = ILP_ISOFORMS.addLConstr(
# lhs = LABEL_I[k],
# sense = GRB.LESS_EQUAL,
# rhs = LABEL_I[k-1]-0.1,
# name = 'LABEL_I_C2[{k}]'.format(k=k)
# )
# Objective function
# For i,j,k such that i ∈ remaining_rids, C[i,j]=1 (read has a zero that can be
# corrected), and E2I[j,k]=1 (isoform k has exon j), OBJ[i][j][k] = 1
OBJ = {}
OBJ_C1 = {}
OBJ_SUM = LinExpr(0.0)
for i in remaining_rids:
OBJ[i] = {}
OBJ_C1[i] = {}
for j in range(0, M):
if not informative[j]:
continue
if C[i][j] > 0: # 1 if exon j not in read i but can be added to it
OBJ[i][j] = {}
OBJ_C1[i][j] = {}
for k in range(ISOFORM_INDEX_START, ilp_settings['K']):
OBJ[i][j][k] = ILP_ISOFORMS.addVar(
vtype=GRB.BINARY,
name='OBJ[{i}][{j}][{k}]'.format(i=i, j=j, k=k)
)
OBJ_C1[i][j][k] = ILP_ISOFORMS.addGenConstrAnd(
resvar=OBJ[i][j][k],
vars=[R2I[i][k], E2I[j][k]],
name='OBJ_C1[{i}][{j}][{k}]'.format(i=i, j=j, k=k)
)
OBJ_SUM.addTerms(1.0*C[i][j], OBJ[i][j][k])
# coeffs = 1.0,
# vars = OBJ[i][j][k]
# )
# We add the chosen cost for each isoform assigned to the garbage isoform if any
GAR_OBJ = {}
GAR_OBJ_C = {}
for i in remaining_rids:
if ilp_settings['recycle_model'] in ['constant', 'exons', 'introns']:
OBJ_SUM.addTerms(1.0*GARBAGE_COST[i], R2I[i][0])
elif ilp_settings['recycle_model'] == 'relative':
GAR_OBJ[i] = {}
GAR_OBJ_C[i] = {}
for j in range(0, M):
if not informative[j]:
continue
GAR_OBJ[i][j] = {}
GAR_OBJ_C[i][j] = {}
for k in range(ISOFORM_INDEX_START, ilp_settings['K']):
if I[i][j] == 1:
GAR_OBJ[i][j][k] = ILP_ISOFORMS.addVar(
vtype=GRB.BINARY,
name='GAR_OBJ[{i}][{j}][{k}]'.format(i=i, j=j, k=k)
)
GAR_OBJ_C[i][j][k] = ILP_ISOFORMS.addGenConstrAnd(
resvar=GAR_OBJ[i][j][k],
vars=[R2I[i][0], E2I_min[j][k]],
name='GAR_OBJ_C[{i}][{j}][{k}]'.format(
i=i, j=j, k=k)
)
OBJ_SUM.addTerms(1.0, GAR_OBJ[i][j][k])
elif I[i][j] == 0 and C[i][j] == 1:
pass
ILP_ISOFORMS.setObjective(
expr=OBJ_SUM,
sense=GRB.MINIMIZE
)
# Optimization
# ILP_ISOFORMS.Params.PoolSearchMode=2
# ILP_ISOFORMS.Params.PoolSolutions=5
ILP_ISOFORMS.setParam('TuneOutput', 1)
if not log_prefix == None:
ILP_ISOFORMS.setParam('LogFile', '{}.glog'.format(log_prefix))
ILP_ISOFORMS.write('{}.lp'.format(log_prefix))
ILP_ISOFORMS.setParam('TimeLimit', ilp_settings['timeout']*60)
ILP_ISOFORMS.optimize()
ILP_ISOFORMS_STATUS = ILP_ISOFORMS.Status
isoforms = {k: dict()
for k in range(ISOFORM_INDEX_START, ilp_settings['K'])}
# print('STATUS: {}'.format(ILP_ISOFORMS_STATUS))
# if ILP_ISOFORMS_STATUS == GRB.Status.TIME_LIMIT:
# status = 'TIME_LIMIT'
if ILP_ISOFORMS_STATUS != GRB.Status.OPTIMAL:
status = 'NO_SOLUTION'
else:
status = 'OPTIMAL'
# Writing the optimal solution to disk
if not log_prefix == None:
solution_file = open('{}.sol'.format(log_prefix), 'w+')
for v in ILP_ISOFORMS.getVars():
solution_file.write('{}\t{}\n'.format(v.VarName, v.X))
solution_file.close()
# Isoform id to isoform structure
for k in range(ISOFORM_INDEX_START, ilp_settings['K']):
isoforms[k]['exons'] = list()
for j in range(0, M):
if informative[j]:
isoforms[k]['exons'].append(
int(E2I[j][k].getAttr(GRB.Attr.X) > 0.9))
else:
isoforms[k]['exons'].append(
I[next(iter(remaining_rids))][j])
isoforms[k]['rid_to_corrections'] = dict()
# Isoform id to read ids set
for i in remaining_rids:
isoform_id = -1
for k in range(0, ilp_settings['K']):
if R2I[i][k].getAttr(GRB.Attr.X) > 0.9:
assert isoform_id == - \
1, 'Read {} has been assigned to multiple isoforms!'.format(
i)
isoform_id = k
assert isoform_id != - \
1, 'Read {} has not been assigned to any isoform!'.format(i)
if isoform_id == 0:
continue
isoforms[isoform_id]['rid_to_corrections'][i] = -1
# Read id to its exon corrections
for k in range(ISOFORM_INDEX_START, ilp_settings['K']):
for i in isoforms[k]['rid_to_corrections'].keys():
isoforms[k]['rid_to_corrections'][i] = [
str(tint['reads'][tint['read_reps'][i][0]]['data'][j]) for j in range(M)]
for j in range(0, M):
if not informative[j]:
isoforms[k]['rid_to_corrections'][i][j] = '-'
elif C[i][j] == 1 and OBJ[i][j][k].getAttr(GRB.Attr.X) > 0.9:
isoforms[k]['rid_to_corrections'][i][j] = 'X'
return ILP_ISOFORMS_STATUS, status, isoforms
