Fraser/dream-coder

1from dreamcoder.likelihoodModel import AllOrNothingLikelihoodModel2from dreamcoder.grammar import *3from dreamcoder.utilities import get_root_dir4 5import os6import traceback7import subprocess8 9 10def multicoreEnumeration(g, tasks, _=None,11                         enumerationTimeout=None,12                         solver='ocaml',13                         CPUs=1,14                         maximumFrontier=None,15                         verbose=True,16                         evaluationTimeout=None,17                         testing=False):18    '''g: Either a Grammar, or a map from task to grammar.19    Returns (list-of-frontiers, map-from-task-to-search-time)'''20 21    # We don't use actual threads but instead use the multiprocessing22    # library. This is because we need to be able to kill workers.23    #from multiprocess import Process, Queue24 25    from multiprocessing import Queue26 27     # everything that gets sent between processes will be dilled28    import dill29 30    solvers = {"ocaml": solveForTask_ocaml,   31               "pypy": solveForTask_pypy,   32               "python": solveForTask_python}   33    assert solver in solvers, "You must specify a valid solver. options are ocaml, pypy, or python." 34 35    likelihoodModel = None36    if solver == 'pypy' or solver == 'python':37      # Use an all or nothing likelihood model.38      likelihoodModel = AllOrNothingLikelihoodModel(timeout=evaluationTimeout) 39      40    solver = solvers[solver]41 42    if not isinstance(g, dict):43        g = {t: g for t in tasks}44    task2grammar = g45 46    # If we are not evaluating on held out testing tasks:47    # Bin the tasks by request type and grammar48    # If these are the same then we can enumerate for multiple tasks simultaneously49    # If we are evaluating testing tasks:50    # Make sure that each job corresponds to exactly one task51    jobs = {}52    for i, t in enumerate(tasks):53        if testing:54            k = (task2grammar[t], t.request, i)55        else:56            k = (task2grammar[t], t.request)57        jobs[k] = jobs.get(k, []) + [t]58 59    disableParallelism = len(jobs) == 160    parallelCallback = launchParallelProcess if not disableParallelism else lambda f, * \61        a, **k: f(*a, **k)62    if disableParallelism:63        eprint("Disabling parallelism on the Python side because we only have one job.")64        eprint("If you are using ocaml, there could still be parallelism.")65 66    # Map from task to the shortest time to find a program solving it67    bestSearchTime = {t: None for t in task2grammar}68 69    lowerBounds = {k: 0. for k in jobs}70 71    frontiers = {t: Frontier([], task=t) for t in task2grammar}72 73    # For each job we keep track of how long we have been working on it74    stopwatches = {t: Stopwatch() for t in jobs}75 76    # Map from task to how many programs we enumerated for that task77    taskToNumberOfPrograms = {t: 0 for t in tasks }78 79    def numberOfHits(f):80        return sum(e.logLikelihood > -0.01 for e in f)81 82    def budgetIncrement(lb):83        if True:84            return 1.585        # Very heuristic - not sure what to do here86        if lb < 24.:87            return 1.88        elif lb < 27.:89            return 0.590        else:91            return 0.2592 93    def maximumFrontiers(j):94        tasks = jobs[j]95        return {t: maximumFrontier - numberOfHits(frontiers[t]) for t in tasks}96 97    def allocateCPUs(n, tasks):98        allocation = {t: 0 for t in tasks}99        while n > 0:100            for t in tasks:101                # During testing we use exactly one CPU per task102                if testing and allocation[t] > 0:103                    return allocation104                allocation[t] += 1105                n -= 1106                if n == 0:107                    break108        return allocation109 110    def refreshJobs():111        for k in list(jobs.keys()):112            v = [t for t in jobs[k]113                 if numberOfHits(frontiers[t]) < maximumFrontier114                 and stopwatches[k].elapsed <= enumerationTimeout]115            if v:116                jobs[k] = v117            else:118                del jobs[k]119 120    # Workers put their messages in here121    q = Queue()122 123    # How many CPUs are we using?124    activeCPUs = 0125 126    # How many CPUs was each job allocated?127    id2CPUs = {}128    # What job was each ID working on?129    id2job = {}130    nextID = 0131 132    while True:133        refreshJobs()134        # Don't launch a job that we are already working on135        # We run the stopwatch whenever the job is being worked on136        # freeJobs are things that we are not working on but could be137        freeJobs = [j for j in jobs if not stopwatches[j].running138                    and stopwatches[j].elapsed < enumerationTimeout - 0.5]139        if freeJobs and activeCPUs < CPUs:140            # Allocate a CPU to each of the jobs that we have made the least141            # progress on142            freeJobs.sort(key=lambda j: lowerBounds[j])143            # Launch some more jobs until all of the CPUs are being used144            availableCPUs = CPUs - activeCPUs145            allocation = allocateCPUs(availableCPUs, freeJobs)146            for j in freeJobs:147                if allocation[j] == 0:148                    continue149                g, request = j[:2]150                bi = budgetIncrement(lowerBounds[j])151                thisTimeout = enumerationTimeout - stopwatches[j].elapsed152                eprint("(python) Launching %s (%d tasks) w/ %d CPUs. %f <= MDL < %f. Timeout %f." %153                       (request, len(jobs[j]), allocation[j], lowerBounds[j], lowerBounds[j] + bi, thisTimeout))154                stopwatches[j].start()155                parallelCallback(wrapInThread(solver),156                                 q=q, g=g, ID=nextID,157                                 elapsedTime=stopwatches[j].elapsed,158                                 CPUs=allocation[j],159                                 tasks=jobs[j],160                                 lowerBound=lowerBounds[j],161                                 upperBound=lowerBounds[j] + bi,162                                 budgetIncrement=bi,163                                 timeout=thisTimeout,164                                 evaluationTimeout=evaluationTimeout,165                                 maximumFrontiers=maximumFrontiers(j),166                                 testing=testing,167                                 likelihoodModel=likelihoodModel)168                id2CPUs[nextID] = allocation[j]169                id2job[nextID] = j170                nextID += 1171 172                activeCPUs += allocation[j]173                lowerBounds[j] += bi174 175        # If nothing is running, and we just tried to launch jobs,176        # then that means we are finished177        if all(not s.running for s in stopwatches.values()):178            break179 180        # Wait to get a response181        message = Bunch(dill.loads(q.get()))182 183        if message.result == "failure":184            eprint("PANIC! Exception in child worker:", message.exception)185            eprint(message.stacktrace)186            assert False187        elif message.result == "success":188            # Mark the CPUs is no longer being used and pause the stopwatch189            activeCPUs -= id2CPUs[message.ID]190            stopwatches[id2job[message.ID]].stop()191 192            newFrontiers, searchTimes, pc = message.value193            for t, f in newFrontiers.items():194                oldBest = None if len(195                    frontiers[t]) == 0 else frontiers[t].bestPosterior196                frontiers[t] = frontiers[t].combine(f)197                newBest = None if len(198                    frontiers[t]) == 0 else frontiers[t].bestPosterior199 200                taskToNumberOfPrograms[t] += pc201 202                dt = searchTimes[t]203                if dt is not None:204                    if bestSearchTime[t] is None:205                        bestSearchTime[t] = dt206                    else:207                        # newBest & oldBest should both be defined208                        assert oldBest is not None209                        assert newBest is not None210                        newScore = newBest.logPrior + newBest.logLikelihood211                        oldScore = oldBest.logPrior + oldBest.logLikelihood212 213                        if newScore > oldScore:214                            bestSearchTime[t] = dt215                        elif newScore == oldScore:216                            bestSearchTime[t] = min(bestSearchTime[t], dt)217        else:218            eprint("Unknown message result:", message.result)219            assert False220 221    eprint("We enumerated this many programs, for each task:\n\t",222           list(taskToNumberOfPrograms.values()))223 224    return [frontiers[t] for t in tasks], bestSearchTime225 226def wrapInThread(f):227    """228    Returns a function that is designed to be run in a thread/threadlike process.229    Result will be either put into the q230    """231    import dill232 233    def _f(*a, **k):234        q = k.pop("q")235        ID = k.pop("ID")236 237        try:238            r = f(*a, **k)239            q.put(dill.dumps({"result": "success",240                   "ID": ID,241                   "value": r}))242        except Exception as e:243            q.put(dill.dumps({"result": "failure",244                   "exception": e,245                   "stacktrace": traceback.format_exc(),246                   "ID": ID}))247            return248    return _f249 250 251def solveForTask_ocaml(_=None,252                       elapsedTime=0.,253                       CPUs=1,254                       g=None, tasks=None,255                       lowerBound=None, upperBound=None, budgetIncrement=None,256                       timeout=None,257                       testing=None, # FIXME: unused258                       likelihoodModel=None,259                       evaluationTimeout=None, maximumFrontiers=None):260 261    import json262 263    def taskMessage(t):264        m = {265            "examples": [{"inputs": list(xs), "output": y} for xs, y in t.examples],266            "name": t.name,267            "request": t.request.json(),268            "maximumFrontier": maximumFrontiers[t]}269        if hasattr(t, "specialTask"):270            special, extra = t.specialTask271            m["specialTask"] = special272            m["extras"] = extra273        return m274 275 276    message = {"DSL": g.json(),277               "tasks": [taskMessage(t)278                         for t in tasks],279 280               "programTimeout": evaluationTimeout,281               "nc": CPUs,282               "timeout": timeout,283               "lowerBound": lowerBound,284               "upperBound": upperBound,285               "budgetIncrement": budgetIncrement,286               "verbose": False,287               "shatter": 5 if len(tasks) == 1 and "turtle" in str(tasks[0].request) else 10}288 289    if hasattr(tasks[0], 'maxParameters') and tasks[0].maxParameters is not None:290        message["maxParameters"] = tasks[0].maxParameters291 292    message = json.dumps(message)293    # uncomment this if you want to save the messages being sent to the solver294    295 296    try:297        solver_file = os.path.join(get_root_dir(), 'solver')298        process = subprocess.Popen(solver_file,299                                   stdin=subprocess.PIPE,300                                   stdout=subprocess.PIPE)301        response, error = process.communicate(bytes(message, encoding="utf-8"))302        response = json.loads(response.decode("utf-8"))303    except OSError as exc:304        raise exc305 306    except:307        print("response:", response)308        print("error:", error)309        with open("message", "w") as f:310            f.write(message)311        print("message,", message)312        assert False, "MAX RAISE"313 314 315    pc = response.get("number_enumerated",0)  # TODO316    frontiers = {}317    searchTimes = {}318    for t in tasks:319        solutions = response[t.name]320        frontier = Frontier([FrontierEntry(program=p,321                                           logLikelihood=e["logLikelihood"],322                                           logPrior=g.logLikelihood(t.request, p))323                             for e in solutions324                             for p in [Program.parse(e["program"])]],325                            task=t)326        frontiers[t] = frontier327        if frontier.empty:328            searchTimes[t] = None329        # This is subtle:330        # The search time we report is actually not be minimum time to find any solution331        # Rather it is the time to find the MAP solution332        # This is important for regression problems,333        # where we might find something with a good prior but bad likelihood early on,334        # and only later discovered the good high likelihood program335        else:336            searchTimes[t] = min(337                (e["logLikelihood"] + e["logPrior"],338                 e["time"]) for e in solutions)[1] + elapsedTime339 340    return frontiers, searchTimes, pc341 342def solveForTask_pypy(_=None,343                      elapsedTime=0.,344                      g=None, task=None,345                      lowerBound=None, upperBound=None, budgetIncrement=None,346                      timeout=None,347                      likelihoodModel=None,348                      evaluationTimeout=None, maximumFrontier=None, testing=False):349    return callCompiled(enumerateForTasks,350                        g, tasks, likelihoodModel,351                        timeout=timeout,352                        testing=testing,353                        elapsedTime=elapsedTime,354                        evaluationTimeout=evaluationTimeout,355                        maximumFrontiers=maximumFrontiers,356                        budgetIncrement=budgetIncrement,357                        lowerBound=lowerBound, upperBound=upperBound)358 359def solveForTask_python(_=None,360                        elapsedTime=0.,361                        g=None, tasks=None,362                        lowerBound=None, upperBound=None, budgetIncrement=None,363                        timeout=None,364                        CPUs=1,365                        likelihoodModel=None,366                        evaluationTimeout=None, maximumFrontiers=None, testing=False):367    return enumerateForTasks(g, tasks, likelihoodModel,368                             timeout=timeout,369                             testing=testing,370                             elapsedTime=elapsedTime,371                             evaluationTimeout=evaluationTimeout,372                             maximumFrontiers=maximumFrontiers,373                             budgetIncrement=budgetIncrement,374                             lowerBound=lowerBound, upperBound=upperBound)375 376 377class EnumerationTimeout(Exception):378    pass379 380def enumerateForTasks(g, tasks, likelihoodModel, _=None,381                      verbose=False,382                      timeout=None,383                      elapsedTime=0.,384                      CPUs=1,385                      testing=False, #unused386                      evaluationTimeout=None,387                      lowerBound=0.,388                      upperBound=100.,389                      budgetIncrement=1.0, maximumFrontiers=None):390    assert timeout is not None, \391        "enumerateForTasks: You must provide a timeout."392 393    from time import time394 395    request = tasks[0].request396    assert all(t.request == request for t in tasks), \397        "enumerateForTasks: Expected tasks to all have the same type"398 399    maximumFrontiers = [maximumFrontiers[t] for t in tasks]400    # store all of the hits in a priority queue401    # we will never maintain maximumFrontier best solutions402    hits = [PQ() for _ in tasks]403 404    starting = time()405    previousBudget = lowerBound406    budget = lowerBound + budgetIncrement407    try:408        totalNumberOfPrograms = 0409        while time() < starting + timeout and \410                any(len(h) < mf for h, mf in zip(hits, maximumFrontiers)) and \411                budget <= upperBound:412            numberOfPrograms = 0413 414            for prior, _, p in g.enumeration(Context.EMPTY, [], request,415                                             maximumDepth=99,416                                             upperBound=budget,417                                             lowerBound=previousBudget):418                descriptionLength = -prior419                # Shouldn't see it on this iteration420                assert descriptionLength <= budget421                # Should already have seen it422                assert descriptionLength > previousBudget423 424                numberOfPrograms += 1425                totalNumberOfPrograms += 1426 427                for n in range(len(tasks)):428                    task = tasks[n]429 430                    #Warning:changed to max's new likelihood model situation431                    #likelihood = task.logLikelihood(p, evaluationTimeout)432                    #if invalid(likelihood):433                        #continue434                    success, likelihood = likelihoodModel.score(p, task)435                    if not success:436                        continue437                        438                    dt = time() - starting + elapsedTime439                    priority = -(likelihood + prior)440                    hits[n].push(priority,441                                 (dt, FrontierEntry(program=p,442                                                    logLikelihood=likelihood,443                                                    logPrior=prior)))444                    if len(hits[n]) > maximumFrontiers[n]:445                        hits[n].popMaximum()446 447                if timeout is not None and time() - starting > timeout:448                    raise EnumerationTimeout449 450            previousBudget = budget451            budget += budgetIncrement452 453            if budget > upperBound:454                break455    except EnumerationTimeout:456        pass457    frontiers = {tasks[n]: Frontier([e for _, e in hits[n]],458                                    task=tasks[n])459                 for n in range(len(tasks))}460    searchTimes = {461        tasks[n]: None if len(hits[n]) == 0 else \462        min(t for t,_ in hits[n]) for n in range(len(tasks))}463 464    return frontiers, searchTimes, totalNumberOfPrograms465 466 467 468 469 470