opt_fuzz.py

import time
import random
import os
import math
import pickle
import heapq
from copy import deepcopy

from typing import List, Tuple
from queue import Queue

from generator import MPSProgram
from mps_mut import DenseMatMut, FlipObjSense, ScaleMut, SparseMatMut, FlipSignMut, NoOpMut, MoreIntMut
from mps_util import load_mps_dict
from optimizer import OPTIMAL, CplexSolver, GurobiSolver

# TODO
# 1. Make this class checkpoint-able
# 2. Use two threads for cplex and gurobi each
# 3. Collect more stats
class OptFuzzer(object):
    MAX_MUTANTS_IN_DIR = 10
    def __init__(self, seed : List, 
                 mut_method : List, 
                 cplex : CplexSolver, 
                 gurobi : GurobiSolver, 
                 save_path : str, 
                 checkpoint_freq : int, 
                 checkpoint_path : str, 
                 diff_path: str, 
                 abs_tol = 1e-6) -> None:
        self.seed = seed # a list of seed mps in dict format
        self.mut_method = mut_method # a list of mutation methods
        self.no_op_mut_method = NoOpMut()
        self.next_mut = 0
        self.n_mut = len(self.mut_method)
        self.cplex = cplex # cplex solver
        self.gurobi = gurobi # gurobi solver
        self.save_path = save_path # path to save the mutated mps files
        os.makedirs(self.save_path, exist_ok=True)
        self.checkpoint_freq = checkpoint_freq
        self.checkpoint_path = checkpoint_path # path to store the checkpoint files
        os.makedirs(self.checkpoint_path, exist_ok=True)
        self.diff_path = diff_path # path to store the mutants that cause differences
        os.makedirs(self.diff_path, exist_ok=True)
        self.abs_tol = abs_tol
        self.diff = [] # mutants that cause differences
        self.all_mutant = []
        # self.queue = Queue(maxsize=0) # unbounded queue for seeds
        self.queue = []
        self.next_id = 0 # id assigned to the next mutant
        print("Initializing seeds")
        for s in self.seed:
            mutant = {
                "id" : self.next_id,
                "root" : self.next_id,
                "parent" : [-1], # root seed, no parent
                "mut_method" : "",
                "mps" : load_mps_dict(s)
            }
            heapq.heappush(self.queue, ((-1, self.next_id, 1), mutant))
            self.next_id += 1
            # self.queue.put_nowait(mutant)
        self.mut_cnt_per_seed = [0] * len(self.seed)
        print("Seeds initialized")

    def get_mutant(self, mut_choice : str, nr_mutations: int = 3) -> Tuple[dict, dict]:
        # parent = self.queue.get()
        parent_ent = heapq.heappop(self.queue)
        parent = parent_ent[1]
        # print(parent_ent[0])
        self.mut_cnt_per_seed[parent["root"]] += 1

        mps = deepcopy(parent["mps"])
        idx = 0
        
        if mut_choice == "fixed":
            idx = 0
        elif mut_choice == "round_robin":
            idx = self.next_mut
            self.next_mut = (self.next_mut + 1) % self.n_mut
        elif mut_choice == "random":
            idx = random.randint(0, self.n_mut - 1)
        else:
            print(f"Unsupported mutant choice {mut_choice}")
            os.abort()
        
        #print("========== GOING TO MUTATE ============")
        if mut_choice == 'random':
            mps_prime = mps
            for _ct in range(nr_mutations):
                idx = random.randint(0, self.n_mut - 1)
                mut_method = self.no_op_mut_method if (random.random() < 0.5) else self.mut_method[idx]
                # print(self.mut_method)
                # print('idx = ', idx, ' Mutating using ', mut_method)
                mps_prime = mut_method(mps_prime)
        else:
            mps_prime = self.mut_method[idx](mps)
        
        ret = {
            "id" : self.next_id,
            "parent" : [parent["id"]], # root seed, no parent
            "mut_method" : self.mut_method[idx].__str__(),
            "mps" : mps_prime,
            "root" : parent["root"],
            "mut_no" : self.mut_cnt_per_seed[parent["root"]]
        }

        self.next_id += 1

        # self.queue.put(parent)
        # heapq.heappush(self.queue, parent_ent)

        return parent_ent, ret
    
    def sol_equal(self, sol1 : dict, sol2 : dict, check_var = False) -> bool:
        if sol1['status'] != sol2['status']:
            return False
        
        if sol1['status'] == OPTIMAL:
            # check if the obj value is the same
            if not math.isclose(sol1['obj_val'], sol2['obj_val'], abs_tol=self.abs_tol):
                return False
            
            # check if the variable values are the same
            # not applicable to MIP where continuous var values can be approx.
            if check_var:
                for v1 in sol1['var_val']:
                    var1 = sol1['var_val'][v1]
                    
                    # Gurobi omits variables if 0
                    var2 = 0
                    try:
                        var2 = sol2['var_val'][v1]
                    except KeyError:
                        var2 = 0
                    
                    if not math.isclose(var1, var2, abs_tol=self.abs_tol):
                        return False
            
        return True

    def get_scaled_time_diff(self, time1, time2) -> float:
        return math.fabs(time1-time2) / max(time1, time2)

    def get_perturbed_time_diff(self, time) -> float:
        return min(1, time + random.uniform(-time,time) * 0.5)

    def fuzz_once(self, mut_choice):
        p, mut = self.get_mutant(mut_choice)

        mps_prog = MPSProgram(**mut["mps"])
        idx = mut["id"] % OptFuzzer.MAX_MUTANTS_IN_DIR # prevent running our of inodes
        fname = os.path.join(self.save_path, "mutant_" + str(idx)+".mps")
        with open(fname, "w") as f:
            f.write(mps_prog.__str__())
        
        c_start = time.time()
        cplex_sol = self.cplex(fname)
        c_time = time.time() - c_start

        g_start = time.time()
        gurobi_sol = self.gurobi(fname)
        g_time = time.time() - g_start

        if cplex_sol is None or gurobi_sol is None:
            print("[Fatal] Cannot load the solution")
            self.checkpoint()
            os.abort()

        are_equal = self.sol_equal(cplex_sol, gurobi_sol)
        
        scaled_time_diff = self.get_scaled_time_diff(c_time, g_time)
        new_key = scaled_time_diff * 0.5 + math.fabs(scaled_time_diff - p[0][2]) * 0.5
        new_key = self.get_perturbed_time_diff(new_key)
        heapq.heappush(self.queue, ((-new_key, p[1]["id"], scaled_time_diff), p[1]))
        # print("time-diff = ", self.get_perturbed_time_diff(scaled_time_diff))
        
        if not are_equal:
            dfname = os.path.join(self.diff_path, f"diff_{len(self.diff)}.mps")
            with open(dfname, "w") as f:
                f.write(mps_prog.__str__())
            self.diff.append(mut)

        # print_ind_stats = False
        # if(print_ind_stats):
        #     #Id, Name, Equal, Cplex time, Gurobi time
        #     print('{},{},{},{},{},{}'.format(mut['id']-203, mut['seed_name'].replace('seed/',''), are_equal, c_time, g_time,scaled_time_diff))

        
        self.all_mutant.append(mut)

    def print_stats(self):
        print(f"Seeds: {self.seed}")
        print("Mutation methods used: ", self.mut_method)
        print("# Mutations: ", self.next_id - len(self.seed))
        print("# Mutants that caused a difference: ", len(self.diff))
        for i, s in enumerate(self.seed):
            print(i, " ", s, "\t", self.mut_cnt_per_seed[i])
        for i, d in enumerate(self.diff):
            print("Difference ", i, " is derived from ", d["root"], " at round ", d["mut_no"])
    
    def checkpoint(self):
        # store all the mutants that have been generated to disks
        fname = os.path.join(self.checkpoint_path, f"ckpt_{self.next_id - 1}")
        with open(fname, "wb") as f:
            pickle.dump(self.all_mutant, f)
        self.all_mutant = []

    def save_states(self):
        self.checkpoint()
        if len(self.diff) > 0:
            fname = os.path.join(self.diff_path, "diff.pickle")
            with open(fname, "wb") as f:
                pickle.dump(self.diff, f)

    def fuzz(self, max_time, mut_choice = "fixed"):
        elapse = 0 # in seconds
        nround = 0
        while elapse < max_time:
            start = time.time()
            self.fuzz_once(mut_choice)
            elapse += time.time() - start
            if nround % self.checkpoint_freq == 0 and nround != 0:
                self.checkpoint()
            nround += 1
        self.print_stats()
        self.cplex.print_solver_stats()
        self.gurobi.print_solver_stats()
        self.save_states()

def main():
    from config import parser
    args = parser.parse_args()
    print(args)

    mut_method = []
    for m in args.mut_methods:
        c = None
        if m == "ScaleMut":
            c = ScaleMut(0.3,2.5)
        elif m == "SparseMatMut":
            c = SparseMatMut(0.5)
        elif m == "DenseMatMut":
            c = DenseMatMut(0.5)
        elif m == "FlipSignMut":
            c = FlipSignMut(0.5)
        elif m == "MoreIntMut":
            c = MoreIntMut(0.3)
        elif m == "FlipObjSense":
            c = FlipObjSense()
        elif m == "NoOpMut":
            c = NoOpMut()
        mut_method.append(c)

    import glob
    seed = glob.glob(os.path.join(args.seed_path, "*.mps"))
    
    cpx = CplexSolver()
    grb = GurobiSolver()
    
    ts = str(time.time())
    print("Timestamp = ", ts)
    save_path = os.path.join(args.save_path, ts)
    checkpoint_freq = args.ckpt_freq
    checkpoint_path = os.path.join(args.ckpt_path, ts)
    diff_path = os.path.join(args.diff_path, ts)
    max_time = args.max_time
    mut_choice = args.mut_choice
    
    fuzzer = OptFuzzer(seed, mut_method, cpx, grb, save_path, checkpoint_freq, checkpoint_path, diff_path)
    fuzzer.fuzz(max_time, mut_choice)

if __name__ == "__main__":
    main()