itp-interface 1.0.0__py3-none-any.whl

itp_interface/tools/lean_local_data_generation_transform.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+
+import sys
+root_dir = f"{__file__.split('itp_interface')[0]}"
+if root_dir not in sys.path:
+    sys.path.append(root_dir)
+import typing
+import uuid
+from itp_interface.tools.lean_cmd_executor import Lean3Executor
+from itp_interface.tools.lean_context_helper import Lean3ContextHelper
+from itp_interface.tools.coq_training_data_generator import GenericTrainingDataGenerationTransform, TrainingDataGenerationType
+from itp_interface.tools.training_data_format import Goal, MergableCollection, TrainingDataMetadataFormat, TrainingDataCollection, TrainingDataFormat
+from itp_interface.tools.training_data import TrainingData
+
+class LocalDataGenerationTransform(GenericTrainingDataGenerationTransform):
+    def __init__(self,
+                depth = None,
+                max_search_results = None,
+                buffer_size : int = 10000,
+                logger = None,
+                max_parallelism : int = 4):
+        super().__init__(TrainingDataGenerationType.LOCAL, buffer_size, logger)
+        self.depth = depth
+        self.max_search_results = max_search_results
+        self.max_parallelism = max_parallelism
+
+    def get_meta_object(self) -> MergableCollection:
+        return TrainingDataMetadataFormat(training_data_buffer_size=self.buffer_size)
+
+    def get_data_collection_object(self) -> MergableCollection:
+        return TrainingDataCollection()
+    
+    def load_meta_from_file(self, file_path) -> MergableCollection:
+        return TrainingDataMetadataFormat.load_from_file(file_path)
+    
+    def load_data_from_file(self, file_path) -> MergableCollection:
+        return TrainingDataCollection.load_from_file(file_path, self.logger)
+
+    def __call__(self, training_data: TrainingData, project_id : str, lean_executor: Lean3Executor, print_coq_executor_callback: typing.Callable[[], Lean3Executor], theorems: typing.List[str] = None, other_args: dict = {}) -> TrainingData:
+        print_lean_executor = print_coq_executor_callback()
+        lean_context_helper = Lean3ContextHelper(print_lean_executor, self.depth, self.logger)
+        lean_context_helper.__enter__()
+        file_namespace = lean_executor.main_file.replace('/', '.')
+        self.logger.info(f"=========================Processing {file_namespace}=========================")
+        proof_running = False
+        cmd_ran = lean_executor.run_next()
+        cmd_exec = lean_executor.current_stmt
+        prev_goal : typing.List[Goal] = lean_context_helper.get_focussed_goals(lean_executor) if lean_executor.is_in_proof_mode() else []
+        line_number = lean_executor.line_num
+        lemma_name = lean_executor.get_lemma_name_if_running()
+        if lemma_name is None:
+            lemma_name = "__NONE__"
+        theorem_id = str(uuid.uuid4())
+        proof_id = theorem_id # self.get_proof_id(theorem_id, file_namespace, line_number, lemma_name)
+        local_lemma_refs_cnt = 0
+        external_lemma_refs_cnt = 0
+        theorems = set(theorems) if theorems is not None else None
+        while cmd_ran:
+            if lean_executor.is_in_proof_mode() and lemma_name != "__NONE__" and (theorems is None or lemma_name in theorems):
+                proof_running = True
+                prev_goal : typing.List[Goal] = [Goal(goal.hypotheses, goal.goal) for goal in prev_goal]
+                next_goal : typing.List[Goal] = lean_context_helper.get_focussed_goals(lean_executor)
+                if len(prev_goal) > 0:
+                    training_data_format = TrainingDataFormat(
+                        proof_id=proof_id,
+                        all_useful_defns_theorems=[],
+                        start_goals=prev_goal,
+                        end_goals=next_goal,
+                        proof_steps=[cmd_exec],
+                        simplified_goals=[], 
+                        addition_state_info={},
+                        file_path=lean_executor.main_file,
+                        theorem_name=lemma_name,
+                        project_id=project_id)
+                    assert len(training_data_format.proof_steps) > 0, f"Proof steps cannot be empty for {proof_id}"
+                    for goal in training_data_format.start_goals:
+                        lemma_cnt = len(training_data_format.all_useful_defns_theorems)
+                        assert all([0 <= lemma_ref.lemma_idx < lemma_cnt for lemma_ref in goal.used_theorems_local]), f"Invalid lemma idx in {proof_id}"
+                        assert all([0 <= lemma_ref.lemma_idx < lemma_cnt for lemma_ref in goal.used_theorems_external]), f"Invalid lemma idx in {proof_id}"
+                        assert all([0 <= lemma_ref.lemma_idx < lemma_cnt for lemma_ref in goal.possible_useful_theorems_external]), f"Invalid lemma idx in {proof_id}"
+                        assert all([0 <= lemma_ref.lemma_idx < lemma_cnt for lemma_ref in goal.possible_useful_theorems_local]), f"Invalid lemma idx in {proof_id}"
+                    training_data.merge(training_data_format)
+                    local_lemma_refs_cnt += sum([len(goal.used_theorems_local) for goal in training_data_format.start_goals])
+                    external_lemma_refs_cnt += sum([len(goal.used_theorems_external) for goal in training_data_format.start_goals])
+                prev_goal = next_goal
+            else:
+                prev_goal = []
+            cmd_ran = lean_executor.run_next()
+            cmd_exec = lean_executor.current_stmt
+            line_number = lean_executor.line_num
+            if proof_running and not lean_executor.is_in_proof_mode():
+                proof_running = False
+                self.logger.info(f"Finished processing lemma [{theorem_id}] {lemma_name}")
+                theorem_id = str(uuid.uuid4())
+            lemma_name = lean_executor.get_lemma_name_if_running()
+            if lemma_name is None:
+                lemma_name = "__NONE__"
+            proof_id = theorem_id # self.get_proof_id(theorem_id, file_namespace, line_number, lemma_name)
+            
+        self.logger.info(f"===============Finished processing {file_namespace}=====================")
+        try:
+            lean_context_helper.__exit__(None, None, None)
+        except:
+            pass
+
+
+if __name__ == "__main__":
+    import os
+    import logging
+    import time
+    os.chdir(root_dir)
+    project_dir = "data/test/lean_proj"
+    file_name = "data/test/lean_proj/src/simple_solved.lean"
+    project_id = project_dir.replace('/', '.')
+    time_str = time.strftime("%Y%m%d-%H%M%S")
+    output_path = f".log/local_data_generation_transform/data/{time_str}"
+    log_path = f".log/local_data_generation_transform/log/{time_str}"
+    log_file = f"{log_path}/local_data_generation_transform-{time_str}.log"
+    os.makedirs(output_path, exist_ok=True)
+    os.makedirs(log_path, exist_ok=True)
+    logging.basicConfig(filename=log_file, level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+    logger = logging.getLogger(__name__)
+    def _print_lean_executor_callback():
+        search_lean_exec = Lean3Executor(project_root=project_dir, main_file=file_name, use_human_readable_proof_context=True, suppress_error_log=True)
+        search_lean_exec.__enter__()
+        return search_lean_exec
+    transform = LocalDataGenerationTransform(0, buffer_size=1000)
+    training_data = TrainingData(
+        output_path, 
+        "training_metadata.json",
+        training_meta=transform.get_meta_object(), 
+        logger=logger)
+    with Lean3Executor(project_root=project_dir, main_file=file_name, use_human_readable_proof_context=True, suppress_error_log=True) as coq_exec:
+        transform(training_data, project_id, coq_exec, _print_lean_executor_callback)
+    save_info = training_data.save()
+    logger.info(f"Saved training data to {save_info}")

itp_interface/tools/lean_parse_utils.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+
+import sys
+
+
+root_dir = f"{__file__.split('itp_interface')[0]}"
+if root_dir not in sys.path:
+    sys.path.append(root_dir)
+import typing
+from itp_interface.lean_server.lean_utils import Lean3Utils
+
+class LeanLineByLineReader(object):
+    def __init__(self, file_name: str = None, file_content: str = None, remove_comments: bool = False, no_strip: bool = False):
+        assert file_name is not None or file_content is not None, "Either file_name or file_content must be provided"
+        assert file_name is None or file_content is None, "Only one of file_name or file_content must be provided"
+        self.file_name : str = file_name
+        self.file_content : str = file_content
+        self.no_strip = no_strip
+        if self.file_name is not None:
+            with open(file_name, 'r') as fd:
+                self.file_content : str = fd.read()
+        if remove_comments:
+            self.file_content = Lean3Utils.remove_comments(self.file_content)
+    
+    def instruction_step_generator(self) -> typing.Iterator[str]:
+        lines = self.file_content.split('\n')
+        for line in lines:
+            if not self.no_strip:
+                line = line.strip()
+            else:
+                line = line
+            yield line

itp_interface/tools/log_utils.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+
+import sys
+
+root_dir = f"{__file__.split('itp_interface')[0]}"
+if root_dir not in sys.path:
+    sys.path.append(root_dir)
+import logging
+
+def setup_logger(name, log_file, level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'):
+    """To setup as many loggers as you want"""
+
+    handler = logging.FileHandler(log_file)
+    formatter = logging.Formatter(format)
+    handler.setFormatter(formatter)
+
+    logger = logging.getLogger(name)
+    logger.setLevel(level)
+    logger.addHandler(handler)
+    return logger

itp_interface/tools/proof_exec_callback.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+
+import sys
+
+root_dir = f"{__file__.split('itp_interface')[0]}"
+if root_dir not in sys.path:
+    sys.path.append(root_dir)
+import logging
+import typing
+from itp_interface.rl.proof_action import ProofAction
+from itp_interface.tools.coq_context_helper import CoqContextHelper
+from itp_interface.tools.lean_context_helper import Lean3ContextHelper
+from itp_interface.tools.lean4_context_helper import Lean4ContextHelper
+from itp_interface.tools.isabelle_context_helper import IsabelleContextHelper
+from itp_interface.tools.coq_executor import CoqExecutor
+from itp_interface.tools.lean_cmd_executor import Lean3Executor
+from itp_interface.tools.lean4_sync_executor import Lean4SyncExecutor
+from itp_interface.tools.isabelle_executor import IsabelleExecutor
+from itp_interface.tools.dynamic_coq_proof_exec import DynamicProofExecutor as DynamicCoqProofExecutor
+from itp_interface.tools.dynamic_lean_proof_exec import DynamicProofExecutor as DynamicLeanProofExecutor
+from itp_interface.tools.dynamic_lean4_proof_exec import DynamicProofExecutor as DynamicLean4ProofExecutor
+from itp_interface.tools.dynamic_isabelle_proof_exec import DynamicProofExecutor as DynamicIsabelleProofExecutor
+
+class ProofExecutorCallback(object):
+    def __init__(self,
+                project_folder: str,
+                file_path: str,
+                language: ProofAction.Language = ProofAction.Language.COQ,
+                prefix: str = None,
+                use_hammer: bool = False,
+                timeout_in_secs: int = 60,
+                use_human_readable_proof_context: bool = True,
+                suppress_error_log: bool = True,
+                search_depth: int = 0,
+                logger: logging.Logger = None,
+                always_use_retrieval: bool = False,
+                keep_local_context: bool = False,
+                setup_cmds: typing.List[str] = [],
+                port: typing.Optional[int] = None,
+                enable_search: bool = True):
+        self.project_folder = project_folder
+        self.file_path = file_path
+        self.language = language
+        self.use_hammer = use_hammer
+        self.timeout_in_secs = timeout_in_secs
+        self.use_human_readable_proof_context = use_human_readable_proof_context
+        self.suppress_error_log = suppress_error_log
+        self.search_depth = search_depth
+        self.logger = logger
+        self.prefix = prefix
+        self.always_use_retrieval = always_use_retrieval
+        self.keep_local_context = keep_local_context
+        self.setup_cmds = setup_cmds
+        self.port = port
+        self.enable_search = enable_search
+        pass
+
+    def get_proof_executor(self) -> typing.Union[DynamicCoqProofExecutor, DynamicLeanProofExecutor, DynamicLean4ProofExecutor, DynamicIsabelleProofExecutor]:
+        if self.language == ProofAction.Language.COQ:
+            search_exec = CoqExecutor(self.project_folder, self.file_path, use_hammer=self.use_hammer, timeout_in_sec=self.timeout_in_secs, suppress_error_log=self.suppress_error_log, use_human_readable_proof_context=self.use_human_readable_proof_context, setup_cmds=self.setup_cmds)
+            coq_context_helper = CoqContextHelper(search_exec, self.search_depth, logger=self.logger, enable_search=self.enable_search)
+            return DynamicCoqProofExecutor(coq_context_helper, self.project_folder, self.file_path, context_type=DynamicCoqProofExecutor.ContextType.BestContext, use_hammer=self.use_hammer, timeout_in_seconds=self.timeout_in_secs, suppress_error_log=self.suppress_error_log, use_human_readable_proof_context=self.use_human_readable_proof_context, setup_cmds=self.setup_cmds)
+        elif self.language == ProofAction.Language.LEAN:
+            search_exec = Lean3Executor(self.project_folder, self.prefix, self.file_path, use_hammer=self.use_hammer, timeout_in_sec=self.timeout_in_secs, suppress_error_log=self.suppress_error_log, use_human_readable_proof_context=self.use_human_readable_proof_context, enable_search=self.always_use_retrieval, keep_local_context=self.keep_local_context)
+            lean_context_helper = Lean3ContextHelper(search_exec, self.search_depth, logger=self.logger)
+            return DynamicLeanProofExecutor(lean_context_helper, self.project_folder, self.file_path, context_type=DynamicLeanProofExecutor.ContextType.NoContext, use_hammer=self.use_hammer, timeout_in_seconds=self.timeout_in_secs, suppress_error_log=self.suppress_error_log, use_human_readable_proof_context=self.use_human_readable_proof_context, keep_local_context=self.keep_local_context)
+        elif self.language == ProofAction.Language.LEAN4:
+            search_exec = Lean4SyncExecutor(self.project_folder, self.prefix, self.file_path, use_hammer=self.use_hammer, timeout_in_sec=self.timeout_in_secs, suppress_error_log=self.suppress_error_log, use_human_readable_proof_context=self.use_human_readable_proof_context, enable_search=self.always_use_retrieval, keep_local_context=self.keep_local_context)
+            lean4_context_helper = Lean4ContextHelper(search_exec, self.search_depth, logger=self.logger)
+            return DynamicLean4ProofExecutor(lean4_context_helper, self.project_folder, self.file_path, context_type=DynamicLeanProofExecutor.ContextType.NoContext, use_hammer=self.use_hammer, timeout_in_seconds=self.timeout_in_secs, suppress_error_log=self.suppress_error_log, use_human_readable_proof_context=self.use_human_readable_proof_context, keep_local_context=self.keep_local_context)
+        elif self.language == ProofAction.Language.ISABELLE:
+            search_exec = IsabelleExecutor(self.project_folder, self.file_path, use_hammer=self.use_hammer, timeout_in_sec=self.timeout_in_secs, suppress_error_log=self.suppress_error_log, use_human_readable_proof_context=self.use_human_readable_proof_context, port=self.port)
+            isabelle_context_helper = IsabelleContextHelper(search_exec, self.search_depth, logger=self.logger)
+            return DynamicIsabelleProofExecutor(isabelle_context_helper, self.project_folder, self.file_path, context_type=DynamicIsabelleProofExecutor.ContextType.BestContext, use_hammer=self.use_hammer, timeout_in_seconds=self.timeout_in_secs, suppress_error_log=self.suppress_error_log, use_human_readable_proof_context=self.use_human_readable_proof_context, port=self.port)
+        else:
+            raise Exception(f"Unknown context type: {self.context_type}")

itp_interface/tools/ray_utils.py

@@ -0,0 +1,265 @@
+#!/usr/bin/env python3
+
+import sys
+root_dir = f"{__file__.split('itp_interface')[0]}"
+if root_dir not in sys.path:
+    sys.path.append(root_dir)
+import asyncio
+import time
+import ray
+import typing
+import psutil
+import logging
+import gc
+
+class RayUtils(object):
+
+    @staticmethod
+    def init_ray(num_of_cpus: int = 10, object_store_memory_in_gb: float = 25, memory_in_gb: float = 0.5, runtime_env: typing.Dict[str, str] = None):
+        gb = 2**30
+        object_store_memory = int(object_store_memory_in_gb * gb)
+        memory = int(memory_in_gb * gb)
+        return ray.init(num_cpus=num_of_cpus, object_store_memory=object_store_memory, _memory=memory, ignore_reinit_error=True, runtime_env=runtime_env)
+
+    @staticmethod
+    def ray_run_within_parallel_limits(
+        max_parallel: int,
+        num_objects: int,
+        transform_outputs: typing.Callable[[typing.List[typing.Any]], None],
+        prepare_next: typing.Callable[[int], typing.List[typing.Any]],
+        create_remotes: typing.Callable[[typing.List[typing.Any]], typing.List[typing.Any]],
+        logger: logging.Logger = None,
+        turn_off_logging: bool = False
+    ):
+        logger = logger or logging.getLogger(__name__)
+        idx = 0
+        next_batch = prepare_next(max_parallel)
+        if not turn_off_logging:
+            logger.info(f"Loading next_batch: {len(next_batch)}, max_parallel: {max_parallel}")
+        assert len(next_batch) <= max_parallel, f"next_batch: {len(next_batch)}, max_parallel: {max_parallel}"
+        process = psutil.Process()
+        if not turn_off_logging:
+            logger.info(f"[Process Id = {process.pid}] [After Next Batch] Memory used: {process.memory_info().rss/2**30} GiB")
+        remotes = create_remotes(next_batch)
+        process = psutil.Process()
+        if not turn_off_logging:
+            logger.info(f"[Process Id = {process.pid}] [After Create] Memory used: {process.memory_info().rss/2**30} GiB")
+            logger.info(f"Created remotes: {len(remotes)}")
+        diff_remotes = len(remotes)
+        while idx < num_objects or len(remotes) > 0:
+            idx += diff_remotes
+            idx = min(idx, num_objects)
+            if not turn_off_logging:
+                logger.info(f"Waiting for idx: {idx}, num_objects: {num_objects}, len(remotes): {len(remotes)}")
+            ready, remotes = ray.wait(remotes)
+            if len(ready) > 0:
+                if not turn_off_logging:
+                    logger.info(f"Got ready: {len(ready)}")
+                process = psutil.Process()
+                if not turn_off_logging:
+                    logger.info(f"[Process Id = {process.pid}] [After Ready] Memory used: {process.memory_info().rss/2**30} GiB")
+                results = ray.get(ready)
+                transform_outputs(results)
+                process = psutil.Process()
+                if not turn_off_logging:
+                    logger.info(f"[Process Id = {process.pid}] [After Transform] Memory used: {process.memory_info().rss/2**30} GiB")
+                next_batch = prepare_next(len(results))
+                process = psutil.Process()
+                if not turn_off_logging:
+                    logger.info(f"[Process Id = {process.pid}] [After Next Batch] Memory used: {process.memory_info().rss/2**30} GiB")
+                assert len(next_batch) <= len(results), f"next_batch: {len(next_batch)}, ready: {len(results)}"
+                new_remotes = create_remotes(next_batch)
+                process = psutil.Process()
+                if not turn_off_logging:
+                    logger.info(f"[Process Id = {process.pid}] [After Create] Memory used: {process.memory_info().rss/2**30} GiB")
+                remotes.extend(new_remotes)
+                diff_remotes = len(new_remotes)
+                # Delete results to free up memory
+                del results
+                process = psutil.Process()
+                if not turn_off_logging:
+                    logger.info(f"[Process Id = {process.pid}] [After Delete] Memory used: {process.memory_info().rss/2**30} GiB")
+                    logger.info(f"Running GC collect")
+                gc.collect()
+                process = psutil.Process()
+                if not turn_off_logging:
+                    logger.info(f"[Process Id = {process.pid}] [After GC] Memory used: {process.memory_info().rss/2**30} GiB")
+            else:
+                diff_remotes = 0
+
+@ray.remote(num_cpus=1)
+class RayResourcePoolActor(object):
+    def __init__(self, resources: list):
+        self.resources = resources
+        self.available = list(resources)
+        self.lock_event = asyncio.Lock()
+        self.acquired = []
+
+    async def acquire(self, num: int):
+        async with self.lock_event:
+            if len(self.available) < num:
+                return None
+            acquired = self.available[:num]
+            self.available = self.available[num:]
+            self.acquired.extend(acquired)
+            return acquired
+    
+    async def release(self, resources: list):
+        async with self.lock_event:
+            for resource in resources:
+                try:
+                    self.acquired.remove(resource)
+                    self.available.append(resource)
+                except ValueError:
+                    pass
+            return True
+    
+    def get_acquired(self):
+        return list(self.acquired)
+    
+    def get_available(self):
+        return list(self.available)
+    
+    async def wait_and_acquire(self, num: int, timeout: typing.Optional[float] = None):
+        polling_time = 0.1
+        start_time = time.time()
+        while True:
+            acquired = await self.acquire(num)
+            if acquired is not None:
+                return acquired
+            if timeout is not None and time.time() - start_time > timeout:
+                return None
+            await asyncio.sleep(polling_time)
+
+
+class RayTimedException(Exception):
+    pass
+
+@ray.remote
+class TimedRayExec(object):
+    def __init__(self, func, args=None, kwargs=None):
+        # Check that func is a remote function
+        assert hasattr(func, "remote"), f"func: {func} is not a remote function"
+        self.func = func
+        self.args = args if args else []
+        self.kwargs = kwargs if kwargs else {}
+        self._is_cancelled = False
+
+    def execute_with_timeout(self, timeout: float = 60):
+        ray_id = self.func.remote(*self.args, **self.kwargs)
+        finished, unfinished = ray.wait([ray_id], timeout=timeout)
+        is_cancelled = False
+        if len(unfinished) > 0:
+            ray.cancel(ray_id, force=True)
+            is_cancelled = True
+        self._is_cancelled = is_cancelled
+        if not is_cancelled:
+            assert len(finished) == 1, f"len(finished): {len(finished)}"
+            return_typ = ray.get(finished[0])
+            return return_typ
+        else:
+            return None
+    
+    def is_cancelled(self):
+        return self._is_cancelled
+
+if __name__ == "__main__":
+    import os
+    import time
+    import random
+    import unittest
+    class TestRayResourcePoolActor(unittest.TestCase):
+        @classmethod
+        def setUpClass(cls):
+            ray.init(ignore_reinit_error=True)
+
+        @classmethod
+        def tearDownClass(cls):
+            ray.shutdown()
+
+        def setUp(self):
+            resources = ["res1", "res2", "res3", "res4"]
+            self.pool = RayResourcePoolActor.remote(resources)
+
+        def test_acquire(self):
+            result = ray.get(self.pool.acquire.remote(2))
+            self.assertEqual(result, ["res1", "res2"])
+            self.assertEqual(ray.get(self.pool.get_available.remote()), ["res3", "res4"])
+            self.assertEqual(ray.get(self.pool.get_acquired.remote()), ["res1", "res2"])
+
+        def test_acquire_more_than_available(self):
+            result = ray.get(self.pool.acquire.remote(5))
+            self.assertIsNone(result)
+
+        def test_release(self):
+            ray.get(self.pool.acquire.remote(2))
+            ray.get(self.pool.release.remote(["res1"]))
+            self.assertEqual(ray.get(self.pool.get_available.remote()), ["res3", "res4", "res1"])
+            self.assertEqual(ray.get(self.pool.get_acquired.remote()), ["res2"])
+
+        def test_release_non_acquired_resource(self):
+            ray.get(self.pool.acquire.remote(2))
+            result = ray.get(self.pool.release.remote(["res5"]))
+            self.assertTrue(result)
+            self.assertEqual(ray.get(self.pool.get_available.remote()), ["res3", "res4"])
+            self.assertEqual(ray.get(self.pool.get_acquired.remote()), ["res1", "res2"])
+
+        def test_wait_and_acquire_success(self):
+            print("Available: " + str(ray.get(self.pool.get_available.remote())))
+            print("Acquired: " + str(ray.get(self.pool.get_acquired.remote())))
+            result = ray.get(self.pool.wait_and_acquire.remote(2))
+            self.assertEqual(result, ["res1", "res2"])
+
+        def test_wait_and_acquire_timeout(self):
+            result = ray.get(self.pool.wait_and_acquire.remote(5, timeout=1))
+            self.assertIsNone(result)
+    
+    unittest.main()
+
+
+    log_folder = f".log/ray_utils"
+    os.makedirs(log_folder, exist_ok=True)
+    log_file = f"{log_folder}/ray_utils-{time.strftime('%Y-%m-%d-%H-%M-%S')}.log"
+    logging.basicConfig(filename=log_file, filemode='w', level=logging.INFO)
+    logger = logging.getLogger(__name__)
+    size = 1000
+    example_cnt = 100000
+    last_job_idx = 0
+    total_sum = 0
+    total_sum_serial = 0
+    job_spec = [[random.random() for _ in range(example_cnt)] for _ in range(size)]
+
+    @ray.remote
+    def _do_job(job):
+        idx, arr = job
+        for i in range(size*10):
+            # This is just to stress the CPUs
+            sum_val = sum(arr)
+        return sum_val, arr
+
+    def _prepare_remotes(num: int):
+        global last_job_idx
+        job_list = job_spec[last_job_idx:last_job_idx+num]
+        job_list = [(last_job_idx + idx, job) for idx, job in enumerate(job_list)]
+        last_job_idx += len(job_list)
+        return job_list
+
+    def _create_remotes(job_list: typing.List[typing.Tuple[int, typing.List[float]]]):
+        remotes = []
+        for job in job_list:
+            logger.info(f"Queuing job {job[0]}")
+            job_ref = ray.put(job)
+            remotes.append(_do_job.remote(job_ref))
+        return remotes
+    
+    def _transform_output(results):
+        global total_sum, total_sum_serial
+        for sum_val, arr in results:
+            total_sum += sum_val
+            total_sum_serial += sum(arr)
+        del results # This is important to free up memory
+    parallel = 30
+    RayUtils.init_ray(num_of_cpus=parallel)
+    RayUtils.ray_run_within_parallel_limits(parallel, size, _transform_output, _prepare_remotes, _create_remotes, logger=logger)
+    assert total_sum == total_sum_serial, f"total_sum: {total_sum}, total_sum_serial: {total_sum_serial}"
+    logger.info(f"total_sum: {total_sum}, total_sum_serial: {total_sum_serial}")

itp_interface/tools/repl/.git

@@ -0,0 +1 @@
+gitdir: ../../../../.git/modules/src/itp_interface/tools/repl

itp_interface/tools/repl/.github/workflows/ci.yml

@@ -0,0 +1,24 @@
+name: Run Tests
+
+on: [push, pull_request]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v2
+
+    - name: install elan
+      run: |
+        set -o pipefail
+        curl -sSfL https://github.com/leanprover/elan/releases/download/v3.0.0/elan-x86_64-unknown-linux-gnu.tar.gz | tar xz
+        ./elan-init -y --default-toolchain none
+        echo "$HOME/.elan/bin" >> $GITHUB_PATH
+
+    - name: build
+      run: lake build
+
+    - name: Run tests
+      run: ./test.sh

itp_interface/tools/repl/.gitignore

@@ -0,0 +1,7 @@
+/build
+/lake-packages/*
+/lakefile.olean
+/.lake
+/test/Mathlib/.lake
+/test/*.olean
+/test/*.olean.tmp

itp_interface/tools/repl/.vscode/copyright.code-snippets

@@ -0,0 +1,13 @@
+{
+	"Copyright header for mathlib": {
+		"scope": "lean4",
+		"prefix": "copyright",
+		"body": [
+			"/-",
+			"Copyright (c) ${CURRENT_YEAR} $1. All rights reserved.",
+			"Released under Apache 2.0 license as described in the file LICENSE.",
+			"Authors: $1",
+			"-/"
+		]
+	}
+}

itp_interface/tools/repl/.vscode/extensions.json

@@ -0,0 +1,13 @@
+{
+	// See https://go.microsoft.com/fwlink/?LinkId=827846 to learn about workspace recommendations.
+	// Extension identifier format: ${publisher}.${name}. Example: vscode.csharp
+
+	// List of extensions which should be recommended for users of this workspace.
+	"recommendations": [
+		"leanprover.lean4"
+	],
+	// List of extensions recommended by VS Code that should not be recommended for users of this workspace.
+	"unwantedRecommendations": [
+		"ms-vscode-remote.remote-containers"
+	]
+}

itp_interface/tools/repl/.vscode/module-docstring.code-snippets

@@ -0,0 +1,35 @@
+{
+	"Module docstring for mathlib": {
+		"scope": "lean4",
+		"prefix": "module docstring",
+		"body": [
+			"/-!",
+      "# ${TM_FILENAME_BASE/([^_]*)(_?)/${1:/capitalize}${2:+ }/g}",
+			"",
+			"## Main definitions",
+			"",
+			"* `FooBar`",
+			"",
+			"## Main statements",
+			"",
+			"* `fooBar_unique`",
+			"",
+			"## Notation",
+			"",
+			"",
+			"",
+			"## Implementation details",
+			"",
+			"",
+			"",
+			"## References",
+			"",
+			"* [F. Bar, *Quuxes*][bibkey]",
+			"",
+			"## Tags",
+			"",
+			"Foobars, barfoos",
+			"-/",
+			"",
+		]},
+}

itp_interface/tools/repl/.vscode/settings.json

@@ -0,0 +1,11 @@
+{
+  "editor.insertSpaces": true,
+  "editor.tabSize": 2,
+  "editor.rulers" : [100],
+  "files.encoding": "utf8",
+  "files.eol": "\n",
+  "files.insertFinalNewline": true,
+  // We don't use this: it messes up our test files!
+  // "files.trimFinalNewlines": true,
+  "files.trimTrailingWhitespace": true,
+}