upload files

prover/workers/__init__.py

@@ -0,0 +1,6 @@
+from .data_loader import DataLoader
+from .scheduler import Scheduler, ProcessScheduler
+
+from .search import SearchProcess
+
+from .generator import GeneratorProcess

prover/workers/data_loader.py

@@ -0,0 +1,48 @@
+import os
+import copy
+
+import torch
+import torch.multiprocessing as mp
+
+from prover.utils import load_jsonl_objects
+
+
+class DataLoader(object):
+    def __init__(self, data_path, data_split, data_repeat, node_rank, world_size, log_dir):
+        self.manager = mp.Manager()
+        self.queue = self.manager.Queue()
+        self.lock = mp.Lock()
+        self.finished_flag_filename = 'finished_running.txt'
+
+        done_set = set()
+        for dirname in os.listdir(log_dir):
+            run_dir = os.path.join(log_dir, dirname)
+            if os.path.isdir(run_dir):
+                for subdirname in os.listdir(run_dir):
+                    if subdirname.startswith('run') and os.path.exists(os.path.join(run_dir, subdirname, self.finished_flag_filename)):
+                        done_set.add(os.path.join(dirname, subdirname))
+
+        todo_count = 0
+        if isinstance(data_split, str):
+            data_split = [data_split]
+        dataset = load_jsonl_objects(data_path)
+        for _repeat in range(data_repeat):
+            for prob_idx, prob in enumerate(dataset):
+                prob_runname = os.path.join(prob['name'], f'run{_repeat}')
+                if f'{prob_idx}_{prob_runname}' in done_set:
+                    continue
+                if data_split is not None and prob['split'] not in data_split:
+                    continue
+                if todo_count % world_size == node_rank:
+                    self.queue.put((prob_idx, prob_runname, copy.deepcopy(prob)))
+                todo_count += 1
+        print('Number of TODO Problems: {}'.format(self.queue.qsize()))
+    
+    def size(self):
+        return self.queue.qsize()
+    
+    def get(self):
+        with self.lock:
+            if self.queue.qsize() > 0:
+                return self.queue.get()
+        return None, None, None

prover/workers/generator.py

@@ -0,0 +1,52 @@
+import os
+import time
+
+import torch
+import torch.multiprocessing as mp
+from vllm import LLM, SamplingParams
+
+from prover.utils import AttrDict, MODEL_FORMAT
+
+
+class GeneratorProcess(mp.Process):
+    def __init__(self, local_rank, node_rank, model_path, task_queue, request_statuses, lock, args):
+        super().__init__()
+        self.local_rank = local_rank
+        self.node_rank = node_rank
+        self.model_path = model_path
+        self.task_queue = task_queue
+        self.request_statuses = request_statuses
+        self.lock = lock
+        self.sampling_params = SamplingParams(
+            temperature=args.temperature,
+            max_tokens=args.max_tokens,
+            top_p=args.top_p,
+            n=1,
+        )
+        self.prompt_func = MODEL_FORMAT[args.mode]['prompt']
+        self.output_func = MODEL_FORMAT[args.mode]['output']
+
+    def run(self):
+        seed = int(time.time()) % 1000 + (self.node_rank * 8 + self.local_rank) * 1000
+        os.environ['LOCAL_RANK'] = str(self.local_rank)
+        llm = LLM(model=self.model_path, max_num_batched_tokens=8192, seed=seed, trust_remote_code=True)
+        while True:
+            inputs = self.task_queue.get()
+            if inputs is None: # Terminate when receiving None
+                break
+            model_inputs = [
+                ''.join([
+                    item.get('_extra_header', str()),
+                    self.prompt_func(item),
+                    item.get('_extra_prompt', str()),
+                ]) for _, _, item in inputs
+            ]
+            model_outputs = llm.generate(
+                model_inputs,
+                self.sampling_params,
+                use_tqdm=False,
+            )
+            outputs = [self.output_func(_output.outputs[0].text) for _output in model_outputs]
+            with self.lock:
+                for (_, request_id, _), output in zip(inputs, outputs):
+                    self.request_statuses[request_id] = output

prover/workers/scheduler.py

@@ -0,0 +1,121 @@
+import os
+import time
+import ctypes
+import subprocess
+import threading
+import multiprocessing as mp
+
+import numpy as np
+
+from prover.utils import AttrDict
+
+
+class TaskQueue(object):
+    def __init__(self, batch_size=512, name='test'):
+        self.name = name
+        self.batch_size = batch_size
+        self.manager = mp.Manager()
+        self.waiting_list = self.manager.list()
+        self.all_tasks_done = mp.Event()
+        self.lock = mp.Lock()
+
+        self._monitor_log = self.manager.list()
+        self._monitor_thread = threading.Thread(target=self._monitor)
+        self._monitor_thread.start()
+    
+    def _monitor(self):
+        last_log_time = time.time()
+        while not self.all_tasks_done.is_set():
+            if time.time() - last_log_time >= 60.0:
+                with self.lock:
+                    if len(self._monitor_log) > 0:
+                        print('TaskQueue-{}:  {} requests popped with avg batch_size {:.1f} in last period  {} waiting in queue'.format(
+                            self.name, np.sum(self._monitor_log), np.mean(self._monitor_log), len(self.waiting_list),
+                        ))
+                        self._monitor_log[:] = []
+                last_log_time = time.time()
+            time.sleep(1.0)
+    
+    def __len__(self):
+        return len(self.waiting_list)
+    
+    def put(self, item):
+        with self.lock:
+            self.waiting_list.append(item)
+    
+    def get(self, no_wait=False):
+        while not self.all_tasks_done.is_set():
+            with self.lock:
+                if len(self.waiting_list) > 0:
+                    tasks = self.waiting_list[:self.batch_size]
+                    self.waiting_list[:self.batch_size] = []
+                    self._monitor_log.append(len(tasks))
+                    return tasks
+            if no_wait:
+                break
+            time.sleep(0.1)
+        return None
+    
+    def close(self):
+        self.all_tasks_done.set()
+        self._monitor_thread.join()
+
+
+class ProcessScheduler(object):
+    def __init__(self, batch_size=512, name='test'):
+        self.name = name
+        self.manager = mp.Manager()
+        self.batch_size = batch_size
+        self.task_queue = TaskQueue(batch_size=batch_size, name=name)
+        self.request_statuses = self.manager.dict()
+        self.request_counter = mp.Value(ctypes.c_int32, 0)
+        self.lock = mp.Lock()
+
+    def submit_request(self, data):
+        with self.lock:
+            self.request_counter.value += 1
+            request_id = self.request_counter.value
+            self.request_statuses[request_id] = None
+            self.task_queue.put((time.time(), request_id, data))
+        return request_id
+    
+    def submit_all_request(self, data_list):
+        request_id_list = [self.submit_request(data) for data in data_list]
+        return request_id_list
+
+    def get_request_status(self, request_id):
+        with self.lock:
+            response = self.request_statuses.get(request_id, None)
+            if response is not None:
+                self.request_statuses.pop(request_id)
+            return response
+    
+    def get_request_outputs(self, request_id):
+        while True:
+            outputs = self.get_request_status(request_id)
+            if outputs is not None:
+                return outputs
+            time.sleep(1.0)
+    
+    def get_all_request_outputs(self, request_id_list):
+        outputs_list = []
+        for request_id in request_id_list:
+            outputs_list.append(self.get_request_outputs(request_id))
+        return outputs_list
+    
+    def close(self):
+        self.task_queue.close()
+
+
+class Scheduler(object):
+    def __init__(self, scheduler_dict):
+        self._scheduler_dict = scheduler_dict
+        for name, scheduler in scheduler_dict.items():
+            self.__setattr__(name, scheduler)
+            for key in dir(scheduler):
+                if not key.startswith('_'):
+                    self.__setattr__(f'{name}_{key}', scheduler.__getattribute__(key))
+    
+    def close(self):
+        for _, scheduler in self._scheduler_dict.items():
+            scheduler.close()

prover/workers/search.py

@@ -0,0 +1,103 @@
+import os
+import time
+import copy
+import json
+import pickle
+from pathlib import Path
+
+import torch
+import torch.multiprocessing as mp
+import numpy as np
+
+from prover.utils import AttrDict, get_datetime
+
+
+class SearchProcess(mp.Process):
+    def __init__(self, idx, log_dir, tokenizer_path, scheduler, data_loader, cfg):
+        self.idx = idx
+        self.log_dir = Path(log_dir)
+        self.scheduler = scheduler
+        self.data_loader = data_loader
+        super().__init__()
+
+        self._current_prob_idx = None
+        sampler_cls = cfg.sampler['algorithm']
+        self.sampler = sampler_cls(
+            scheduler=self.scheduler,
+            tokenizer_path=tokenizer_path,
+            process_print=self.process_print,
+            cfg=AttrDict({
+                **cfg.sampler,
+                'mode': cfg.model_args.mode,
+                'max_tokens': cfg.model_args.max_tokens,
+            })
+        )
+    
+    def _post_process(self, data: dict, proof_code: str):
+        header = data.get('header', str())
+        tailer = data.get('tailer', str())
+        formal_statement = data['formal_statement']
+        return dict(
+            statement_proposal=f'{header}{formal_statement}{proof_code}{tailer}',
+            proof_code=proof_code,
+        )
+    
+    def process_print(self, logs, **kwargs):
+        print('Process ID: {:3d}    Problem ID: {}    {}'.format(self.idx, self._current_prob, logs), **kwargs)
+
+    def run(self):
+        while True:
+            prob_idx, prob_runname, data = self.data_loader.get()
+            if prob_idx is None: break
+            
+            sample_start_time = time.time()
+            # build a yield-iterator object to generate samples
+            self._current_prob = f'{prob_idx}_{prob_runname}'
+            prob_log_dir = self.log_dir / self._current_prob
+            os.makedirs(prob_log_dir, exist_ok=True)
+            sample_generator = self.sampler.sample(
+                data=data,
+                prob_log_dir=prob_log_dir,
+            )
+            # submit requests to the verification server when receiving from the generator
+            candidate_list, info_list, request_id_list = [], [], []
+            for sample, info in sample_generator:
+                candidate = self._post_process(data, sample)
+                candidate_list.append(candidate)
+                info_list.append(copy.deepcopy(info))
+                request_id = self.scheduler.verifier_submit_request(candidate['statement_proposal'])
+                request_id_list.append(request_id)
+            sample_timecost = time.time() - sample_start_time
+
+            verification_start_wait_time = time.time()
+            result_list = self.scheduler.verifier_get_all_request_outputs(request_id_list)
+            verification_timecost = time.time() - verification_start_wait_time
+
+            success_count = sum([int(result['complete']) for result in result_list])
+            self.process_print('Success: {} / {}    Generation: {:.2f} secs    Verfication: {:.2f} secs'.format(
+                success_count, len(candidate_list), sample_timecost, verification_timecost,
+            ))
+            
+
+            summary_dict = dict(success=[], failure=[])
+            for _idx, (candidate, result, info) in enumerate(zip(candidate_list, result_list, info_list)):
+                success_flag = 'success' if result['complete'] else 'failure'
+                summary_dict[success_flag].append(dict(
+                    problem_name=data['name'],
+                    sample_info=info,
+                    formal_statement=data['formal_statement'],
+                    proof_code=candidate['proof_code'],
+                    result=result,
+                ))
+            
+            prob_name, run_id = prob_runname.split('/')
+            prob_log_basedir = self.log_dir / f'{prob_idx}_{data["name"]}'
+            log_tag = f'{self.sampler.algorithm_name}-{run_id}'
+            # separately save success and failure results
+            for success_flag, summary_list in summary_dict.items():
+                if len(summary_list) > 0:
+                    with open(prob_log_basedir / f'{success_flag}-{log_tag}-{get_datetime()}.pkl', 'wb') as pkl_f:
+                        pickle.dump(summary_list, pkl_f)
+            # create a 'finished' placeholder
+            with open(prob_log_dir / self.data_loader.finished_flag_filename, 'w') as f:
+                print('finished', file=f)