upload files

prover/algorithms/__init__.py

@@ -0,0 +1,2 @@
+from .sampling import Sampling
+from .rmax_tree_search import RMaxTS

prover/algorithms/base.py

@@ -0,0 +1,54 @@
+import os
+
+import numpy as np
+from transformers import AutoTokenizer
+
+from prover.utils import get_datetime, load_jsonl_objects, MODEL_FORMAT
+
+
+class SamplingAlgorithmBase(object):
+    def __init__(self, scheduler, tokenizer_path, process_print, cfg, **kwargs):
+        os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+        self.scheduler = scheduler
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+        self.process_print = process_print
+        self.cfg = cfg
+
+        self.max_tokens = cfg.max_tokens
+        self.few_shot_dataset = cfg.get('few_shot_dataset', None)
+        if self.few_shot_dataset is not None:
+            self.few_shot_dataset = load_jsonl_objects(self.few_shot_dataset)
+        self.few_shot_num = cfg.get('few_shot_num', 3)
+        self.few_shot_func = MODEL_FORMAT[cfg.mode]['few_shot']
+        self.log_interval = cfg.get('log_interval', 32)
+    
+    @property
+    def algorithm_name(self):
+        return self.__class__.__name__
+    
+    def _post_sample_info(self, **kwargs):
+        return dict(
+            algorithm=self.algorithm_name,
+            datetime=get_datetime(),
+            **kwargs,
+        )
+    
+    def _encode_length(self, code):
+        return len(self.tokenizer.encode(code))
+    
+    def _preprocess_data(self, input_data):
+        if self.few_shot_dataset is None or self.few_shot_num == 0:
+            return input_data
+        return {
+            **input_data,
+            '_extra_header': ''.join([
+                self.few_shot_func(self.few_shot_dataset[idx])
+                for idx in np.random.choice([
+                    _idx for _idx, _data in enumerate(self.few_shot_dataset)
+                    if _data['name'] != input_data['name']
+                ], size=self.few_shot_num, replace=False)
+            ] + [input_data.get('_extra_header', str())]),
+        }
+    
+    def sample(self, **kwargs):
+        raise NotImplementedError

prover/algorithms/rmax_tree_search.py

@@ -0,0 +1,320 @@
+import os
+import gc
+import time
+import math
+import random
+import pickle
+import subprocess
+
+import numpy as np
+
+from .base import SamplingAlgorithmBase
+from prover.lean.proof import ProofSummarizer
+from prover.utils import ConcurrentJob
+
+
+class TreeNode(object):
+    def __init__(self, parent=None, code=None, **kwargs):
+        self.parent = parent
+        self.children = dict()
+        self._info = {key: val for key, val in kwargs.items()}
+        if code is not None:
+            self.update_code(code)
+        
+        if '_discounted_rewards' not in self._info:
+            self._info['_discounted_rewards'] = 0.0
+        if '_discounted_visitation' not in self._info:
+            self._info['_discounted_visitation'] = 0.0
+        if '_subtree_discounted_rewards' not in self._info:
+            self._info['_subtree_discounted_rewards'] = 0.0
+        if '_subtree_discounted_visitation' not in self._info:
+            self._info['_subtree_discounted_visitation'] = 0.0
+        self._num_running_jobs = 0
+        self._subtree_num_running_jobs = 0
+        self._update_value(gamma=0.0)  # gamma=0.0 is okay for initialization
+    
+    # basic tree supports
+    @property
+    def code(self):
+        return random.choice(self._info['_code_list'])
+
+    def update_code(self, code):
+        if '_code_list' not in self._info:
+            self._info['_code_list'] = []
+        if code not in self._info['_code_list']:
+            self._info['_code_list'].append(code)
+    
+    def __getitem__(self, key):
+        return self._info[key]
+    
+    def to_node_list(self):
+        return sum([child.to_node_list() for _, child in self.children.items()], start=[self])
+    
+    def to_dict(self):
+        return dict(
+            info=self._info,
+            children={
+                edge: child_node.to_dict()
+                for edge, child_node in self.children.items()
+            }
+        )
+    
+    @classmethod
+    def from_dict(cls, dict_data, parent=None):
+        node = cls(
+            parent=parent,
+            **dict_data['info'],
+        )
+        node.children = {
+            edge: cls.from_dict(child_dict, parent=node)
+            for edge, child_dict in dict_data['children'].items()
+        }
+        return node
+    
+    # algorithm supports
+    def update_reward(self, reward, gamma, first_node=True):
+        if first_node:
+            self._info['_discounted_rewards'] = self._info['_discounted_rewards'] * gamma + reward
+            self._info['_discounted_visitation'] = self._info['_discounted_visitation'] * gamma + 1.0
+        self._info['_subtree_discounted_rewards'] = self._info['_subtree_discounted_rewards'] * gamma + reward
+        self._info['_subtree_discounted_visitation'] = self._info['_subtree_discounted_visitation'] * gamma + 1.0
+        self._update_value(gamma)
+        if self.parent is not None:
+            self.parent.update_reward(reward, gamma, first_node=False)
+    
+    def start_new_job(self, gamma, first_node=True):
+        if first_node:
+            self._num_running_jobs += 1
+        self._subtree_num_running_jobs += 1
+        self._update_value(gamma)
+        if self.parent is not None:
+            self.parent.start_new_job(gamma, first_node=False)
+    
+    def complete_job(self, gamma, first_node=True):
+        if first_node:
+            self._num_running_jobs -= 1
+        self._subtree_num_running_jobs -= 1
+        self._update_value(gamma)
+        if self.parent is not None:
+            self.parent.complete_job(gamma, first_node=False)
+    
+    def _update_value(self, gamma):
+        discounted_rewards = self._info['_discounted_rewards'] * (gamma ** self._num_running_jobs)
+        discounted_visitation = \
+            self._info['_discounted_visitation'] * (gamma ** self._num_running_jobs) \
+            + (1.0 - (gamma ** self._num_running_jobs)) / (1.0 - gamma)
+        self.value = discounted_rewards / max(discounted_visitation, 1e-2)
+        self.visitation = discounted_visitation
+
+        subtree_discounted_rewards = self._info['_subtree_discounted_rewards'] * (gamma ** self._subtree_num_running_jobs)
+        subtree_discounted_visitation = \
+            self._info['_subtree_discounted_visitation'] * (gamma ** self._subtree_num_running_jobs) \
+            + (1.0 - (gamma ** self._subtree_num_running_jobs)) / (1.0 - gamma)
+        self.subtree_value = subtree_discounted_rewards / max(subtree_discounted_visitation, 1e-2)
+        self.subtree_visitation = subtree_discounted_visitation
+
+
+class RMaxTS(SamplingAlgorithmBase):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.gamma = self.cfg.get('gamma', 0.99)
+        self.sample_num = self.cfg.get('sample_num', 6400)
+        self.concurrent_num = self.cfg.get('concurrent_num', 32)
+        self.tactic_state_comment = self.cfg.get('tactic_state_comment', True)
+        self.ckpt_interval = self.cfg.get('ckpt_interval', 128)
+
+        self.ckpt_filename = 'checkpoint.pkl'
+        self.node_cls = TreeNode
+        self.algorithm_pipeline = [
+            self._tactic_tree_generate_proof,
+            self._tactic_tree_parse_proof,
+            self._rmax_exploration_summarize_results,
+        ]
+    
+    # basic supports
+    def _save_ckpt(self, ckpt_dict: dict):
+        # save a backup before overwriting the checkpoint file
+        if os.path.exists(self.ckpt_path):
+            subprocess.run(['cp', self.ckpt_path, self.ckpt_path + '.backup'])
+        # overwrite the checkpoint file
+        with open(self.ckpt_path, 'wb') as pkl_f:
+            pickle.dump(ckpt_dict, pkl_f)
+    
+    # tree structure supports
+    def _tree_setup(self, data):
+        # initialize tree
+        ckpt_info = None
+        for _ckpt_path in [self.ckpt_path, self.ckpt_path + '.backup']:
+            if os.path.exists(_ckpt_path):
+                try:
+                    with open(_ckpt_path, 'rb') as pkl_f:
+                        ckpt_info = pickle.load(pkl_f)
+                except:
+                    self.process_print(f'Checkpoint saved at {_ckpt_path} is broken.')
+        if ckpt_info is not None:
+            root = self.node_cls.from_dict(ckpt_info['root'])
+            sample_count = ckpt_info['sample_count']
+            yield_cache = ckpt_info['yield_cache']
+            self.process_print(f'Load checkpoint from sample_count={sample_count}')
+        else:
+            root = self.node_cls(code=dict(tactic_code=str(), state_comment=str()), depth=0)
+            sample_count = 0
+            yield_cache = []
+        
+        # compile the root node with `sorry`
+        self.proof_summarizer = ProofSummarizer(data=data, scheduler=self.scheduler)
+        root_sorry = self.proof_summarizer.analyze('  sorry', require_verification=True)
+        assert root_sorry.result['pass'], "Cannot parse a `sorry` tactic on root."
+        self.root_goal = root_sorry.result['sorries'][-1]['goal']
+
+        # other initialization
+        self._last_selected_node = root
+        
+        return root, sample_count, yield_cache
+    
+    def _tree_new_child(self, parent):
+        return self.node_cls(
+            parent=parent,
+            depth=parent['depth'] + 1,
+        )
+    
+    def _tree_step(self, node, edge, code):
+        if edge not in node.children:
+            new_node = self._tree_new_child(node)
+            node.children[edge] = new_node
+            self.node_list.append(new_node)
+        child_node = node.children[edge]
+        child_node.update_code(code)
+        return child_node
+    
+    def _tree_update(self, proof):
+        node_walk, partial_code = self.root, str()
+
+        # use tactic goals as tree edges
+        segments = proof.segmentation()
+        prev_goal = self.root_goal
+        for info in segments:
+            partial_code += info.tactic_code
+            code = partial_code + info.state_comment if self.tactic_state_comment else partial_code
+            if self._encode_length(code) < self.max_tokens:
+                if info.goal != prev_goal:
+                    node_walk = self._tree_step(
+                        node=node_walk, edge=info.goal,
+                        code=dict(tactic_code=partial_code, state_comment=info.state_comment)
+                    )
+                    prev_goal = info.goal
+        return node_walk
+    
+    # algorithm pipeline
+    def _select_node(self):
+        node = self.root
+        while len(node.children) > 0:
+            num_choice = 1 + len(node.children)
+            total_visitation = node.visitation + np.sum([child.subtree_visitation for _, child in node.children.items()])
+            def hoeffding_ucb(node_visitation):
+                return math.sqrt(2.0 * math.log(max(total_visitation, 2)) / max(node_visitation, 1e-2))
+            choice_list = [(node.value + hoeffding_ucb(node.visitation), None)]
+            for _, child in node.children.items():
+                choice_list.append((child.subtree_value + hoeffding_ucb(child.subtree_visitation), child))
+            choice_list.sort(reverse=True, key=lambda x: x[0])
+            if choice_list[0][1] is None:
+                node.start_new_job(gamma=self.gamma)
+                return node
+            else:
+                node = choice_list[0][1]
+        node.start_new_job(gamma=self.gamma)
+        return node
+    
+    def _tactic_tree_generate_proof(self, data, node):
+        code_prefix = node.code
+        extra_prompt = code_prefix['tactic_code']
+        if self.tactic_state_comment:
+            extra_prompt += code_prefix['state_comment']
+        return dict(
+            node=node,
+            code_prefix=code_prefix,
+            generator_request_id=self.scheduler.generator_submit_request(
+                self._preprocess_data({**data, '_extra_prompt': extra_prompt}),
+            ),
+        )
+    
+    def _tactic_tree_parse_proof(self, node, code_prefix, generator_request_id):
+        code = self.scheduler.generator_get_request_status(generator_request_id)
+        if code is None:
+            return None
+        code = code_prefix['tactic_code'] + code
+        proof = self.proof_summarizer.analyze(code, require_verification=True)
+        return dict(node=node, proof=proof)
+    
+    def _rmax_exploration_summarize_results(self, node, proof):
+        if not proof.is_result_ready():
+            return None
+        
+        num_nodes_before = len(self.node_list)
+        self._tree_update(proof)
+        # RMax reward
+        node.update_reward(int(len(self.node_list) > num_nodes_before), gamma=self.gamma)
+        node.complete_job(gamma=self.gamma)
+        
+        return dict(
+            code=proof.cleaned_code,
+            result=proof.result,
+        )
+    
+    # sampler interface
+    def sample(self, data, prob_log_dir, **kwargs):
+        self.ckpt_path = os.path.join(prob_log_dir, self.ckpt_filename)
+        self.root, sample_count, yield_cache = self._tree_setup(data)
+        self.node_list = self.root.to_node_list()
+        for _proposal, _sample_info in yield_cache:
+            yield _proposal, _sample_info
+        gc.collect()  # release memory
+
+        job_slots = [
+            ConcurrentJob(self.algorithm_pipeline)
+            for _ in range(self.concurrent_num)
+        ]
+
+        sample_budget = self.sample_num - sample_count if len(yield_cache) == 0 else 0
+        while (sample_budget > 0) or any([not job.is_idle() for job in job_slots]):
+            for job in job_slots:
+                if job.is_idle() and sample_budget > 0:
+                    node = self._select_node()
+                    self._last_selected_node = node
+                    job.start(data=data, node=node)
+                    sample_budget -= 1
+                if not job.is_idle():
+                    info = job.get_status()
+                    if info is not None:
+                        # output samples
+                        sample_count += 1
+                        if info['result']['complete']:
+                            _proposal, _sample_info = info['code'], self._post_sample_info(
+                                cost=sample_count, tree_size=len(self.node_list),
+                            )
+                            yield_cache.append((_proposal, _sample_info))
+                            yield _proposal, _sample_info
+                        # logging
+                        if sample_count % self.log_interval == 0:
+                            self.process_print('Progress: {} / {}    Tree Size: {}'.format(
+                                sample_count, self.sample_num, len(self.node_list),
+                            ))
+                        # saving checkpoints
+                        if sample_count % self.ckpt_interval == 0:
+                            self._save_ckpt(dict(
+                                root=self.root.to_dict(),
+                                sample_count=sample_count,
+                                yield_cache=yield_cache,
+                            ))
+                            if len(yield_cache) > 0:
+                                # return after saving the checkpoint
+                                # avoid overestimation caused by interrupt-restart loop
+                                sample_budget = 0
+            time.sleep(0.1)
+        
+        # save the final tree structure
+        self._save_ckpt(dict(
+            root=self.root.to_dict(),
+            sample_count=sample_count,
+            yield_cache=yield_cache,
+        ))

prover/algorithms/sampling.py

@@ -0,0 +1,22 @@
+from .base import SamplingAlgorithmBase
+
+
+class Sampling(SamplingAlgorithmBase):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.sample_num = self.cfg.get('sample_num', 32)
+
+    def sample(self, data, **kwargs):
+        request_id_list = [
+            self.scheduler.generator_submit_request(
+                # add few-shot prompts
+                self._preprocess_data(data),
+            ) for _ in range(self.sample_num)
+        ]
+        for _idx, request_id in enumerate(request_id_list):
+            outputs = self.scheduler.generator_get_request_outputs(request_id)
+            yield outputs, self._post_sample_info(cost=_idx+1)
+            if _idx + 1 < self.sample_num and (_idx + 1) % self.log_interval == 0:
+                self.process_print('Progress: {} / {}'.format(
+                    _idx + 1, self.sample_num
+                ))

prover/launch.py

@@ -0,0 +1,95 @@
+import os
+import copy
+import time
+import warnings
+import argparse
+
+import torch
+
+from prover.workers import DataLoader, Scheduler, ProcessScheduler, GeneratorProcess, SearchProcess
+from prover.lean.verifier import Lean4ServerScheduler
+from prover.utils import get_datetime, load_config, AttrDict
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str)
+    parser.add_argument("--log_dir", type=str, default=f'logs/{get_datetime()}')
+    parser.add_argument("--node_rank", type=int, default=0)
+    parser.add_argument("--world_size", type=int, default=1)
+    args = parser.parse_args()
+
+    cfg = load_config(args.config)
+    os.makedirs(args.log_dir, exist_ok=True)
+
+    ngpus = torch.cuda.device_count()
+    assert ngpus >= 1
+    
+    # create data loader
+    data_loader = DataLoader(
+        data_path=cfg.data_path,
+        data_split=cfg.get('data_split', None),
+        data_repeat=cfg.get('data_repeat', 1),
+        node_rank=args.node_rank,
+        world_size=args.world_size,
+        log_dir=args.log_dir,
+    )
+
+    # build Lean verifier
+    verifier_scheduler = Lean4ServerScheduler(
+        max_concurrent_requests=cfg.lean_max_concurrent_requests,
+        memory_limit=cfg.lean_memory_limit,
+        timeout=cfg.lean_timeout,
+        name='verifier',
+    )
+
+    # load LLM models on gpus
+    generator_scheduler = ProcessScheduler(batch_size=cfg.batch_size, name='generator')
+    llm_processes = [
+        GeneratorProcess(
+            local_rank=local_rank,
+            node_rank=args.node_rank,
+            model_path=cfg.model_path,
+            task_queue=generator_scheduler.task_queue,
+            request_statuses=generator_scheduler.request_statuses,
+            lock=generator_scheduler.lock,
+            args=cfg.model_args,
+        )
+        for local_rank in range(ngpus)
+    ]
+
+    # create a unified scheduler interface
+    scheduler = Scheduler(dict(
+        verifier=verifier_scheduler,
+        generator=generator_scheduler,
+    ))
+
+    # launch search processes
+    search_processes = [
+        SearchProcess(
+            idx=i+args.node_rank*cfg.n_search_procs,
+            log_dir=args.log_dir,
+            tokenizer_path=cfg.model_path,
+            scheduler=scheduler,
+            data_loader=data_loader,
+            cfg=cfg,
+        )
+        for i in range(min(cfg.n_search_procs, data_loader.size()))
+    ]
+    for p in search_processes:
+        p.start()
+    print(f'Complete launching {len(search_processes)} SearchProcesses')
+
+    for p in llm_processes:
+        p.start()
+    print(f'Complete launching {len(llm_processes)} LLMProcesses')
+
+    for p in search_processes:
+        p.join()
+    print(f'All {len(search_processes)} SearchProcesses stopped')
+
+    scheduler.close()
+
+    for p in llm_processes:
+        p.join()
+    print(f'All {len(llm_processes)} LLMProcesses stopped')

prover/lean/proof.py

@@ -0,0 +1,183 @@
+import re
+import json
+
+import numpy as np
+
+from prover.utils import AttrDict, LEAN4_DEFAULT_HEADER
+
+
+class Proof(object):
+    def __init__(self, full_code, _args, _result_backup=None, **kwargs):
+        self._kwargs_backup = kwargs
+        for key, val in kwargs.items():
+            self.__setattr__(key, val)
+        self._args = _args
+        self._update_full_code(full_code, _result_backup=_result_backup)
+    
+    @property
+    def result(self):
+        if self._verifier_request_id is not None:
+            self._result = self._scheduler.verifier_get_request_outputs(self._verifier_request_id)
+            self._verifier_request_id = None
+        return self._result
+    
+    def is_result_ready(self):
+        if self._verifier_request_id is None:
+            return True
+        status = self._scheduler.verifier_get_request_status(self._verifier_request_id)
+        if status is not None:
+            self._result = status
+            self._verifier_request_id = None
+        return self._result is not None
+    
+    @property
+    def cleaned_code(self):
+        return self.full_code[len(self.header) + len(self.formal_statement): len(self.full_code) - len(self.tailer)]
+    
+    def _update_full_code(self, full_code, _result_backup=None):
+        self.full_code = full_code
+        self._verifier_request_id, self._result = None, None
+        if _result_backup is not None:
+            self._result = _result_backup
+        elif self._args.require_verification:  # need to call verification server
+            self._verifier_request_id = self._scheduler.verifier_submit_request(dict(
+                code=self.full_code,
+                ast=True, tactics=True,
+            ))
+        self._parse_full_code_lines()
+    
+    def _parse_full_code_lines(self):
+        self._full_code_lines = self.full_code.split('\n')
+        self._line_offset, _offset = [], -1
+        for _line in self._full_code_lines:
+            _offset += 1  # '\n'
+            self._line_offset.append(_offset)
+            _offset += len(_line)
+    
+    def _get_idx(self, pos_info):
+        return self._line_offset[pos_info['line'] - 1] + pos_info['column']
+    
+    def segmentation(self, result=None):
+        if result is None:
+            result = self.result
+        if 'errors' not in result:
+            # compiler timeout
+            return []
+        
+        _prefix_len = len(self.header) + len(self.formal_statement)
+        truncate_pos = len(self.full_code) - len(self.tailer)
+        for info in result['sorries'] + result['errors']:
+            info_pos = self._get_idx(info['pos'])
+            if info_pos >= _prefix_len and not info.get('data', str()).lstrip().startswith('unsolved goals'):
+                truncate_pos = min(truncate_pos, info_pos)
+        partial_code = self.full_code[:truncate_pos]
+        if len(partial_code) <= _prefix_len:
+            # all proof lines are invalid
+            return []
+        
+        code_lines = partial_code.split('\n')
+        pos_last, segments = _prefix_len, []
+        for line_idx in range(len(code_lines)):
+            if self._line_offset[line_idx] >= _prefix_len:
+                def compute_last_valid_char_pos(line):
+                    idx, last_non_blank = 0, len(line) + 1
+                    while idx < len(line):
+                        if line[idx: idx+2] == '--':
+                            return last_non_blank
+                        elif line[idx: idx+2] == '/-':
+                            if '-/' not in line[idx+2:]:
+                                # cannot split in this line
+                                return len(line) + 1
+                            idx = line.find('-/', idx+2) + 1
+                        elif line[idx] != ' ':
+                            last_non_blank = idx
+                        idx += 1
+                    return last_non_blank
+                line_lastChar = self._line_offset[line_idx] + compute_last_valid_char_pos(code_lines[line_idx])
+                line_endPos = self._line_offset[line_idx] + len(code_lines[line_idx])
+
+                pos_min, goal = 1e9, None
+                for tactic_info in result['ast']['tactics']:
+                    pos, endPos = tactic_info['pos'], tactic_info['endPos']
+                    if line_lastChar <= endPos and endPos <= line_endPos and pos < pos_min:
+                        pos_min = pos
+                        goal = tactic_info['stateAfter']
+                if goal is not None:
+                    for tactic_info in result['ast']['tactics']:
+                        pos, endPos = tactic_info['pos'], tactic_info['endPos']
+                        if pos_last < endPos and endPos <= line_endPos and pos < pos_min:
+                            pos_min = pos
+
+                    while pos_min > 0 and partial_code[pos_min - 1] != '\n':
+                        pos_min -= 1
+                    indent_len = 0
+                    while partial_code[pos_min + indent_len] == ' ':
+                        indent_len += 1
+                    newline_with_indent = '\n' + ' ' * indent_len
+                    
+                    segments.append(AttrDict(
+                        tactic_code=partial_code[pos_last: line_endPos] + '\n',
+                        state_comment=newline_with_indent.join([
+                            ' ' * indent_len + '/- tactic state:',
+                            '  ' + goal.replace('\n', newline_with_indent + '  '),
+                            '-/\n'
+                        ]),
+                        goal=goal,
+                        indent=indent_len,
+                    ))
+                    pos_last = line_endPos + 1
+        if result['complete'] and (len(segments) == 0 or segments[-1].goal != 'no goals' or segments[-1].indent != segments[0].indent):
+            indent_len = 2 if len(segments) == 0 else segments[0].indent
+            newline_with_indent = '\n' + ' ' * indent_len
+            segments.append(AttrDict(
+                tactic_code=partial_code[pos_last:].rstrip(' \n') + '\n',
+                state_comment=newline_with_indent.join([
+                    ' ' * indent_len + '/- tactic state:',
+                    '  no goals',
+                    '-/\n'
+                ]),
+                goal='no goals',
+                indent=indent_len,
+            ))
+        segments = [seg for seg in segments if len(seg.tactic_code.strip(' \n')) > 0]
+        return segments
+
+
+class ProofSummarizer(object):
+    def __init__(self, data, scheduler=None):
+        """
+        Inputs:
+            data (`dict`): The problem information storing in a `dict` object.
+                formal_statement (`str`): The formal statement of the unproved problem;
+                header (`str`, *optional*, defaults to ''): The code header required by the complier;
+                tailer (`str`, *optional*, defaults to ''): The code tailer required by the complier.
+            scheduler (prover.workers.scheduler.Scheduler, *optional*, defaults to None):
+                An interface to submit requests to models and the verification server.
+                If set to None, the downstream tasks may require the verification result as inputs.
+        """
+        self.formal_statement = data['formal_statement']
+        self.header = data.get('header', LEAN4_DEFAULT_HEADER)
+        self.tailer = data.get('tailer', str())
+        self.scheduler = scheduler
+    
+    def analyze(self, code, require_verification=True):
+        """
+        Inputs:
+            code (`str`): The code of formal proof.
+            require_verification (`bool`, *optional*, defaults to True):
+                Whether to submit a request to the verification server.
+                If set to False, the downstream tasks may require the verification result as inputs.
+        Return:
+            A `Proof` object that summarizes the code.
+        """
+        return Proof(
+            full_code=''.join([self.header, self.formal_statement, code.rstrip(' \n'), self.tailer]),
+            raw_code=code,
+            formal_statement=self.formal_statement,
+            header=self.header,
+            tailer=self.tailer,
+            _scheduler=self.scheduler,
+            _args=AttrDict(
+                require_verification=require_verification,
+            )
+        )

prover/lean/verifier.py

@@ -0,0 +1,156 @@
+import os
+import time
+import json
+import ctypes
+import resource
+import tempfile
+import traceback
+import threading
+import subprocess
+import multiprocessing as mp
+from pprint import pprint
+
+import numpy as np
+
+from prover.lean.ast_parser import lean4_parser
+from prover.workers import ProcessScheduler
+from prover.utils import AttrDict
+
+
+HOME_DIR = os.path.expanduser('~')
+DEFAULT_LAKE_PATH = f'{HOME_DIR}/.elan/bin/lake'
+DEFAULT_LEAN_WORKSPACE = 'mathlib4/'
+
+
+def verify_lean4_file(code, lake_path=DEFAULT_LAKE_PATH, lean_workspace=DEFAULT_LEAN_WORKSPACE, last_env=None, verbose=False, timeout=300, allTactics=False, ast=False, premises=False, tactics=False):
+    command = dict(cmd=code, allTactics=allTactics, ast=ast, tactics=tactics, premises=premises)
+    if last_env is not None:
+        command.update(env=last_env)
+    message_str = json.dumps(command, ensure_ascii=False)
+    if verbose:
+        print(message_str)
+    start_time = time.time()
+    system_messages = ''
+    try:
+        with tempfile.TemporaryFile(mode='w+', encoding='utf-8') as temp_file:
+            temp_file.write(message_str + "\r\n\r\n")
+            temp_file.seek(0)
+            outputs = subprocess.run([lake_path, "exe", 'repl'], stdin=temp_file, capture_output=True, text=True, cwd=lean_workspace, timeout=timeout)
+        result = json.loads(outputs.stdout)
+        ast_results = lean4_parser(code, result['ast']) if 'ast' in result and result['ast'] else {}
+        result = {
+            "sorries" : result.get('sorries', []), 
+            "tactics" : result.get('tactics', []),
+            "errors" : [m for m in result.get('messages', []) if m['severity'] == 'error'],
+            "warnings" : [m for m in result.get('messages', []) if m['severity'] == 'warning'],
+            "infos" : [m for m in result.get('messages', []) if m['severity'] == 'info'],
+            "system_messages" : system_messages,
+            "system_errors" : None,
+            "ast" : ast_results,
+            "verified_code" : code,
+        }
+        result['pass'] = not result['errors']
+        result['complete'] = result['pass'] and not result['sorries'] and not any("declaration uses 'sorry'" in warning['data'] or 'failed' in warning['data'] for warning in result['warnings'])
+    except:
+        result = {
+            "pass": False,
+            "complete": False,
+            "system_errors": traceback.format_exc(),
+            "system_messages": system_messages
+        }
+    result['verify_time'] = time.time() - start_time
+    return result
+
+
+class Lean4ServerProcess(mp.Process):
+    def __init__(self, idx, task_queue, request_statuses, lock, extra_args=AttrDict()):
+        super().__init__()
+        self.idx = idx
+        self.task_queue = task_queue
+        self.request_statuses = request_statuses
+        self.lock = lock
+        self.extra_args = extra_args
+
+        self.timeout = extra_args.get('timeout', 300)
+        self.memory_limit = extra_args.get('memory_limit', -1)
+        self.last_output_time = mp.Value(ctypes.c_double, time.time())
+        self.complete_count = mp.Value(ctypes.c_int, 0)
+    
+    def run(self):
+        if self.memory_limit > 0:
+            resource.setrlimit(
+                resource.RLIMIT_AS,
+                (self.memory_limit * (1000 ** 3), self.memory_limit * (1000 ** 3))
+            )
+        while True:
+            inputs = self.task_queue.get()
+            if inputs is None: # Terminate when receiving None
+                break
+            for _, request_id, task in inputs:
+                if isinstance(task, str):
+                    task = dict(code=task)
+                if 'timeout' not in task:
+                    task['timeout'] = self.timeout
+                result = verify_lean4_file(**task)
+                if len(result['system_messages']) > 0:
+                    retry_start_time = time.time()
+                    while ('lean::exception: failed to create thread' in result['system_messages'] or
+                           'std::bad_alloc: std::bad_alloc' in result['system_messages'] or
+                           'Cannot allocate memory' in result['system_messages']) \
+                          and time.time() - retry_start_time < self.timeout:
+                        time.sleep(0.1)
+                        result = verify_lean4_file(**task)
+                with self.lock:
+                    self.request_statuses[request_id] = result
+                    self.last_output_time.value = time.time()
+                    self.complete_count.value += 1
+
+
+class Lean4ServerScheduler(ProcessScheduler):
+    def __init__(self, max_concurrent_requests=64, timeout=300, memory_limit=-1, name='verifier'):
+        super().__init__(batch_size=1, name=name)
+        
+        self.processes = [
+            Lean4ServerProcess(
+                idx=idx,
+                task_queue=self.task_queue,
+                request_statuses=self.request_statuses,
+                lock=self.lock,
+                extra_args=AttrDict(
+                    timeout=timeout,
+                    memory_limit=memory_limit,
+                )
+            )
+            for idx in range(max_concurrent_requests)
+        ]
+        for p in self.processes:
+            p.start()
+        print(f'Complete launching {len(self.processes)} LeanServerProcesses')
+
+        self.timeout = timeout
+        self._running_monitor = mp.Value(ctypes.c_bool, True)
+        self._last_complete_count = mp.Value(ctypes.c_int, 0)
+        self._monitor_process = mp.Process(target=self._monitor)
+        self._monitor_process.start()
+    
+    def _monitor(self):
+        while self._running_monitor.value:
+            time.sleep(1.0)
+            subprocess.run(['killall', 'repl', f'--older-than={int(self.timeout) + 10}s'], capture_output=True)
+    
+    def close(self):
+        super().close()
+        for p in self.processes:
+            p.join()
+        self._running_monitor.value = False
+        self._monitor_process.join()
+        print(f'All {len(self.processes)} LeanServerProcesses stopped')
+
+
+if __name__ == '__main__':
+    code = open('mathlib4/.lake/packages/REPL/test/aime_1983_p9.code.in').read()
+    lean4_scheduler = Lean4ServerScheduler(max_concurrent_requests=1, timeout=300, memory_limit=10, name='verifier')
+    request_id_list = lean4_scheduler.submit_all_request([dict(code=code, ast=True, tactics=True)])
+    outputs_list = lean4_scheduler.get_all_request_outputs(request_id_list)
+    lean4_scheduler.close()
+    pprint(outputs_list)

prover/summarize.py

@@ -0,0 +1,56 @@
+import os
+import argparse
+
+import pandas as pd
+from termcolor import colored
+
+from prover.utils import get_datetime, load_config, load_jsonl_objects
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str)
+    parser.add_argument("--log_dir", type=str)
+    args = parser.parse_args()
+
+    cfg = load_config(args.config)
+    dataset = load_jsonl_objects(cfg.data_path)
+    log_dir_dict = {
+        os.path.basename(args.log_dir): args.log_dir,
+    }
+
+    for data in dataset:
+        data['success'] = dict()
+    for runname, log_dir in log_dir_dict.items():
+        for prob_idx, data in enumerate(dataset):
+            res_dir = os.path.join(log_dir, f'{prob_idx}_{dataset[prob_idx]["name"]}')
+            _success_flag = False
+            if os.path.exists(res_dir):
+                for filename in os.listdir(res_dir):
+                    if filename[:7] == 'success':
+                        _success_flag = True
+            data['success'][runname] = _success_flag
+    
+    def make_inner_list(info):
+        return {key: [val] for key, val in info.items()}
+    
+    def add_color(info):
+        return {key: colored(val, 'cyan', attrs=['bold']) for key, val in info.items()} if info['prob_type'] == '<all>' else info
+
+    def aggregate(split, prob_type):
+        info = dict(split=split, prob_type=prob_type)
+        for runname in log_dir_dict:
+            success_count, total_count = 0, 0
+            for prob_idx, data in enumerate(dataset):
+                if data['split'] == split and (data['name'].startswith(prob_type) or prob_type == '<all>'):
+                    total_count += 1
+                    success_count += int(data['success'][runname])
+            info[runname] = '{:3d} / {:3d} = {:.3f}'.format(success_count, total_count, success_count / total_count)
+        return pd.DataFrame(make_inner_list(add_color(info)))
+    
+    summary = pd.concat([
+        aggregate(split, '<all>')
+        for split in set([data['split'] for data in dataset])
+    ])
+    print('DateTime:', get_datetime(readable=True))
+    print(summary.to_markdown(index=False, tablefmt="github", colalign=["left"] * 2 + ["right"] * len(log_dir_dict)))

prover/utils.py

@@ -0,0 +1,105 @@
+import os
+import json
+import pytz
+from pathlib import Path
+from datetime import datetime
+from collections import UserDict
+from importlib.machinery import SourceFileLoader
+from easydict import EasyDict as AttrDict
+
+
+LEAN4_DEFAULT_HEADER = "import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 0\n\nopen BigOperators Real Nat Topology Rat\n\n"
+
+def non_cot_prompt(data):
+    return "Complete the following Lean 4 code:\n\n```lean4\n{header}{informal_prefix}{formal_statement}".format(
+        header=data.get('header', LEAN4_DEFAULT_HEADER),
+        informal_prefix=data.get('informal_prefix', str()),
+        formal_statement=data['formal_statement'],
+    )
+
+def non_cot_few_shot_prompt(data):
+    return "Complete the following Lean 4 code:\n\n```lean4\n{header}{informal_prefix}{formal_statement}{formal_proof}\n```\n\n\n".format(
+        header=data.get('header', LEAN4_DEFAULT_HEADER),
+        informal_prefix=data.get('informal_prefix', str()),
+        formal_statement=data['formal_statement'],
+        formal_proof=data['formal_proof'],
+    )
+
+def cot_prompt(data):
+    return "Complete the following Lean 4 code with explanatory comments preceding each line of code:\n\n```lean4\n{header}{informal_prefix}{formal_statement}".format(
+        header=data.get('header', LEAN4_DEFAULT_HEADER),
+        informal_prefix=data.get('informal_prefix', str()),
+        formal_statement=data['formal_statement'],
+    )
+
+def cot_few_shot_prompt(data):
+    return "Complete the following Lean 4 code with explanatory comments preceding each line of code:\n\n```lean4\n{header}{informal_prefix}{formal_statement}{formal_proof}\n```\n\n\n".format(
+        header=data.get('header', LEAN4_DEFAULT_HEADER),
+        informal_prefix=data.get('informal_prefix', str()),
+        formal_statement=data['formal_statement'],
+        formal_proof=data['formal_proof'],
+    )
+
+def post_process_output(output):
+    _find_idx = output.find("```")
+    return output[:_find_idx] if _find_idx >= 0 else output
+
+MODEL_FORMAT = dict(
+    non_cot=dict(prompt=non_cot_prompt, output=post_process_output, few_shot=non_cot_few_shot_prompt),
+    cot=dict(prompt=cot_prompt, output=post_process_output, few_shot=cot_few_shot_prompt),
+)
+
+
+def get_datetime(readable=False):
+    if readable:
+        return datetime.now(pytz.timezone('Asia/Shanghai')).strftime("%Y/%m/%d %H:%M:%S")
+    return datetime.now(pytz.timezone('Asia/Shanghai')).strftime("%Y%m%d_%H%M%S")
+
+def load_config(fname):
+    name = Path(fname).stem
+    mod = SourceFileLoader(name, fname).load_module()
+
+    config = {}
+    for n in dir(mod):
+        if not n.startswith("__"):
+            config[n] = getattr(mod, n)
+    config = AttrDict(config)
+
+    return config
+
+def load_jsonl_objects(input_path):
+    objects = []
+    with open(input_path, 'r', encoding='utf-8') as fr:
+        for line in fr:
+            objects.append(json.loads(line))
+    return objects
+
+
+class ConcurrentJob(object):
+    def __init__(self, stage_list):
+        assert len(stage_list) > 1
+        self.stage_list = stage_list
+        self.reset()
+    
+    def is_idle(self):
+        return self._stage_idx is None
+    
+    def reset(self):
+        self._stage_idx = None
+        self._stage_cache = None
+    
+    def start(self, **kwargs):
+        self._stage_idx = 1
+        self._stage_cache = self.stage_list[0](**kwargs)
+    
+    def get_status(self):
+        assert not self.is_idle()
+        while True:
+            status = self.stage_list[self._stage_idx](**self._stage_cache)
+            if status is None:
+                return None
+            self._stage_idx += 1
+            if self._stage_idx == len(self.stage_list):
+                self.reset()
+                return status
+            self._stage_cache = status

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)