py-adtools 0.1.2__tar.gz → 0.1.3__tar.gz

{py_adtools-0.1.2 → py_adtools-0.1.3}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: py-adtools
-Version: 0.1.2
-Summary: Useful tools for parsing and evaluating Python programs for algorithm design.
+Version: 0.1.3
+Summary: Useful tools for parsing and evaluating Python programs for LLM-based algorithm design.
 Home-page: https://github.com/RayZhhh/py-adtools
 Author: Rui Zhang
 Author-email: rzhang.cs@gmail.com
@@ -34,7 +34,7 @@ Dynamic: summary
 
 ------
 
-The figure demonstrates how a Python program is parsed into `PyScript`, `PyFunction`, `PyClass,` and `PyProgram` via `adtools`.
+The figure demonstrates how a Python program is parsed into `PyCodeBlock`, `PyFunction`, `PyClass,` and `PyProgram` via `adtools`.
 
 ![pycode](./assets/pycode.png)
 
@@ -68,7 +68,7 @@ Parse your code (in string) into Python code instances, so that you can check ea
 from adtools import PyProgram
 
 code = r'''
-import ast, numba                 # This part will be parsed into PyScript
+import ast, numba                 # This part will be parsed into PyCodeBlock
 import numpy as np
 
 @numba.jit()                      # This part will be parsed into PyFunction
@@ -80,8 +80,9 @@ def function(arg1, arg2=True):
 
 @some.decorators()                # This part will be parsed into PyClass
 class PythonClass(BaseClass):
-    class_var1 = 1                # This part will be parsed into PyScript
-    class_varb = 2                # and placed in PyClass.class_vars_and_code
+    
+    class_var1 = 1                # This part will be parsed into PyCodeBlock
+    class_var2 = 2                # and placed in PyClass.class_vars_and_code
 
     def __init__(self, x):        # This part will be parsed into PyFunction
         self.x = x                # and placed in PyClass.functions
@@ -93,11 +94,11 @@ class PythonClass(BaseClass):
     def method2(self, x, y):
     	return x + y + self.method1(x)
 
-    class InnerClass:             # This part will be parsed into PyScript
+    class InnerClass:             # This part will be parsed into PyCodeBlock
     	def __init__(self):       # and placed in PyClass.class_vars_and_code
     		...
 
-if __name__ == '__main__':        # This part will be parsed into PyScript
+if __name__ == '__main__':        # This part will be parsed into PyCodeBlock
 	res = function(1)
 	print(res)
 	res = PythonClass().method2(1, 2)
@@ -116,7 +117,7 @@ print(p.functions[0].name)
 Evaluate Python programs in a secure process to avoid the abortation of the main process. Two steps:
 
 - Extend the `PyEvaluator` class and override the `evaluate_program` method.
-- Evaluate the program (in str) by calling the `evaluate` (directly evaluate without executing in a sandbox process) or the `secure_evaluate` (evaluate in a sandbox process) methods. 
+- Evaluate the program (in str) by calling the `evaluate` (directly evaluate without executing in a sandbox process) or the `secure_evaluate` (evaluate in a sandbox process) methods.
 
 ```python
 import time
@@ -190,6 +191,7 @@ def merge(left, right):
 
 harmful_code_generated_by_llm = '''
 def merge_sort(arr):
+    print('I am harmful')  # There will be no output since we redirect STDOUT to /dev/null by default.
     while True:
         pass
 '''

{py_adtools-0.1.2 → py_adtools-0.1.3}/README.md

@@ -8,7 +8,7 @@
 
 ------
 
-The figure demonstrates how a Python program is parsed into `PyScript`, `PyFunction`, `PyClass,` and `PyProgram` via `adtools`.
+The figure demonstrates how a Python program is parsed into `PyCodeBlock`, `PyFunction`, `PyClass,` and `PyProgram` via `adtools`.
 
 ![pycode](./assets/pycode.png)
 
@@ -42,7 +42,7 @@ Parse your code (in string) into Python code instances, so that you can check ea
 from adtools import PyProgram
 
 code = r'''
-import ast, numba                 # This part will be parsed into PyScript
+import ast, numba                 # This part will be parsed into PyCodeBlock
 import numpy as np
 
 @numba.jit()                      # This part will be parsed into PyFunction
@@ -54,8 +54,9 @@ def function(arg1, arg2=True):
 
 @some.decorators()                # This part will be parsed into PyClass
 class PythonClass(BaseClass):
-    class_var1 = 1                # This part will be parsed into PyScript
-    class_varb = 2                # and placed in PyClass.class_vars_and_code
+    
+    class_var1 = 1                # This part will be parsed into PyCodeBlock
+    class_var2 = 2                # and placed in PyClass.class_vars_and_code
 
     def __init__(self, x):        # This part will be parsed into PyFunction
         self.x = x                # and placed in PyClass.functions
@@ -67,11 +68,11 @@ class PythonClass(BaseClass):
     def method2(self, x, y):
     	return x + y + self.method1(x)
 
-    class InnerClass:             # This part will be parsed into PyScript
+    class InnerClass:             # This part will be parsed into PyCodeBlock
     	def __init__(self):       # and placed in PyClass.class_vars_and_code
     		...
 
-if __name__ == '__main__':        # This part will be parsed into PyScript
+if __name__ == '__main__':        # This part will be parsed into PyCodeBlock
 	res = function(1)
 	print(res)
 	res = PythonClass().method2(1, 2)
@@ -90,7 +91,7 @@ print(p.functions[0].name)
 Evaluate Python programs in a secure process to avoid the abortation of the main process. Two steps:
 
 - Extend the `PyEvaluator` class and override the `evaluate_program` method.
-- Evaluate the program (in str) by calling the `evaluate` (directly evaluate without executing in a sandbox process) or the `secure_evaluate` (evaluate in a sandbox process) methods. 
+- Evaluate the program (in str) by calling the `evaluate` (directly evaluate without executing in a sandbox process) or the `secure_evaluate` (evaluate in a sandbox process) methods.
 
 ```python
 import time
@@ -164,6 +165,7 @@ def merge(left, right):
 
 harmful_code_generated_by_llm = '''
 def merge_sort(arr):
+    print('I am harmful')  # There will be no output since we redirect STDOUT to /dev/null by default.
     while True:
         pass
 '''

py_adtools-0.1.3/adtools/__init__.py

@@ -0,0 +1,3 @@
+from .py_code import PyCodeBlock, PyFunction, PyClass, PyProgram
+from .evaluator import PyEvaluator
+from .evaluator_pool import EvaluatorExecutorPool

{py_adtools-0.1.2 → py_adtools-0.1.3}/adtools/evaluator.py

@@ -1,7 +1,13 @@
+"""
+Copyright (c) 2025 Rui Zhang <rzhang.cs@gmail.com>
+
+NOTICE: This code is under MIT license. This code is intended for academic/research purposes only.
+Commercial use of this software or its derivatives requires prior written permission.
+"""
+
 import multiprocessing
 import os
 import sys
-import time
 from abc import ABC, abstractmethod
 from queue import Empty
 from typing import Any, Literal, Dict, Callable, List
@@ -12,15 +18,26 @@ from .py_code import PyProgram
 
 class PyEvaluator(ABC):
 
-    def __init__(self, debug_mode: bool = False, *, exec_code: bool = True):
-        """Evaluator interface for evaluating the Python algorithm program.
+    def __init__(
+            self,
+            exec_code: bool = True,
+            debug_mode: bool = False,
+            *,
+            join_timeout_seconds: int = 10
+    ):
+        """Evaluator interface for evaluating the Python algorithm program. Override this class and implement
+        'evaluate_program' method, then invoke 'self.evaluate()' or 'self.secure_evaluate()' for evaluation.
         Args:
+            exec_code: Using 'exec()' to execute the program code and obtain the callable functions and classes,
+                which will be passed to 'self.evaluate_program()'. Set this parameter to 'False' if you are going to
+                evaluate a Python scripy. Note that if the parameter is set to 'False', the arguments 'callable_...'
+                in 'self.evaluate_program()' will no longer be affective.
             debug_mode: Debug mode.
-            exec_code : Using 'exec()' to compile the code and provide the callable function.
+            join_timeout_seconds: Timeout in seconds to wait for the process to finish. Kill the process if timeout.
         """
         self._debug_mode = debug_mode
         self._exec_code = exec_code
-        self._JOIN_TIMEOUT_SECONDS = 5
+        self._join_timeout_seconds = join_timeout_seconds
 
     @abstractmethod
     def evaluate_program(
@@ -31,19 +48,21 @@ class PyEvaluator(ABC):
             callable_classes_dict: Dict[str, Callable] | None,
             callable_classes_list: List[Callable] | None,
             **kwargs
-    ) -> Any | None:
+    ) -> Any:
         """Evaluate a given program.
         Args:
-            program_str            : The raw program text.
+            program_str: The raw program text.
             callable_functions_dict: A dict maps function name to callable function.
             callable_functions_list: A list of callable functions.
-            callable_classes_dict  : A dict maps class name to callable class.
-            callable_classes_list  : A list of callable classes.
-        Return:
+            callable_classes_dict: A dict maps class name to callable class.
+            callable_classes_list: A list of callable classes.
+        Returns:
             Returns the evaluation result.
         """
-        raise NotImplementedError('Must provide an evaluator for a python program. '
-                                  'Override this method in a subclass.')
+        raise NotImplementedError(
+            'Must provide an evaluator for a python program. '
+            'Override this method in a subclass.'
+        )
 
     def _kill_process_and_its_children(self, process: multiprocessing.Process):
         # Find all children processes
@@ -54,7 +73,7 @@ class PyEvaluator(ABC):
             children_processes = []
         # Terminate parent process
         process.terminate()
-        process.join(timeout=self._JOIN_TIMEOUT_SECONDS)
+        process.join(timeout=self._join_timeout_seconds)
         if process.is_alive():
             process.kill()
             process.join()
@@ -64,36 +83,42 @@ class PyEvaluator(ABC):
                 print(f"Killing process {process.pid}'s children process {child.pid}")
             child.terminate()
 
-    def evaluate(self, program_str: str, **kwargs):
+    def evaluate(self, program: str | PyProgram, **kwargs):
+        """Evaluate a program.
+        Args:
+            program: the program to be evaluated.
+            **kwargs: additional keyword arguments to pass to 'evaluate_program'.
+        """
         try:
             # Parse to program instance
-            program = PyProgram.from_text(program_str)
+            if isinstance(program, str):
+                program = PyProgram.from_text(program)
             function_names = [f.name for f in program.functions]
             class_names = [c.name for c in program.classes]
+
+            # Execute the code and get callable instances
             if self._exec_code:
-                # Compile the program, and maps the global func/var/class name to its address
                 all_globals_namespace = {}
                 # Execute the program, map func/var/class to global namespace
-                exec(program_str, all_globals_namespace)
+                exec(str(program), all_globals_namespace)
                 # Get callable functions
-                callable_functions_list = [all_globals_namespace[f_name] for f_name in function_names]
-                callable_functions_dict = dict(zip(function_names, callable_functions_list))
+                callable_funcs_list = [all_globals_namespace[f_name] for f_name in function_names]
+                callable_funcs_dict = dict(zip(function_names, callable_funcs_list))
                 # Get callable classes
-                callable_classes_list = [all_globals_namespace[c_name] for c_name in class_names]
-                callable_classes_dict = dict(zip(class_names, callable_classes_list))
+                callable_cls_list = [all_globals_namespace[c_name] for c_name in class_names]
+                callable_cls_dict = dict(zip(class_names, callable_cls_list))
             else:
-                callable_functions_list = None
-                callable_functions_dict = None
-                callable_classes_list = None
-                callable_classes_dict = None
+                callable_funcs_list, callable_funcs_dict, callable_cls_list, callable_cls_dict = (
+                    None, None, None, None
+                )
 
             # Get evaluate result
             res = self.evaluate_program(
-                program_str,
-                callable_functions_dict,
-                callable_functions_list,
-                callable_classes_dict,
-                callable_classes_list,
+                str(program),
+                callable_funcs_dict,
+                callable_funcs_list,
+                callable_cls_dict,
+                callable_cls_list,
                 **kwargs
             )
             return res
@@ -109,10 +134,13 @@ class PyEvaluator(ABC):
             redirect_to_devnull: bool,
             **kwargs
     ):
+        # Redirect STDOUT and STDERR to '/dev/null'
         if redirect_to_devnull:
-            with open('/dev/null', 'w') as devnull:
+            with open(os.devnull, 'w') as devnull:
                 os.dup2(devnull.fileno(), sys.stdout.fileno())
                 os.dup2(devnull.fileno(), sys.stderr.fileno())
+
+        # Evaluate and put the results to the queue
         res = self.evaluate(program_str, **kwargs)
         result_queue.put(res)
 
@@ -121,23 +149,26 @@ class PyEvaluator(ABC):
             program: str | PyProgram,
             timeout_seconds: int | float = None,
             redirect_to_devnull: bool = True,
-            multiprocessing_start_method=Literal['auto', 'fork', 'spawn'],
+            multiprocessing_start_method: Literal['default', 'auto', 'fork', 'spawn'] = 'auto',
             **kwargs
     ):
-        """
+        """Evaluate program in a new process. This enables timeout restriction and output redirection.
         Args:
             program: the program to be evaluated.
             timeout_seconds: return 'None' if the execution time exceeds 'timeout_seconds'.
             redirect_to_devnull: redirect any output to '/dev/null'.
-            multiprocessing_start_method: start a process using 'fork' or 'spawn'.
+            multiprocessing_start_method: start a process using 'fork' or 'spawn'. If set to 'auto',
+                the process will be started using 'fork' with Linux/macOS and 'spawn' with Windows.
+                If set to 'default', there will be no changes to system default.
+            **kwargs: additional keyword arguments to pass to 'evaluate_program'.
         """
         if multiprocessing_start_method == 'auto':
-            # Force MacOS and Linux use 'fork' to generate new process
+            # Force macOS and Linux use 'fork' to generate new process
             if sys.platform.startswith('darwin') or sys.platform.startswith('linux'):
                 multiprocessing.set_start_method('fork', force=True)
         elif multiprocessing_start_method == 'fork':
             multiprocessing.set_start_method('fork', force=True)
-        else:
+        elif multiprocessing_start_method == 'spawn':
             multiprocessing.set_start_method('spawn', force=True)
 
         try:
@@ -156,19 +187,22 @@ class PyEvaluator(ABC):
                     result = result_queue.get(timeout=timeout_seconds)
                     # After getting the result, terminate/kill the process
                     self._kill_process_and_its_children(process)
-                except Empty:
-                    # Timeout
+                except Empty:  # The queue is empty indicates a timeout
                     if self._debug_mode:
                         print(f'DEBUG: the evaluation time exceeds {timeout_seconds}s.')
+                    # Terminate/kill all processes if timeout happens
                     self._kill_process_and_its_children(process)
                     result = None
                 except Exception as e:
                     if self._debug_mode:
                         print(f'DEBUG: evaluation failed with exception:\n{e}')
+                    # Terminate/kill all processes if meet exceptions
                     self._kill_process_and_its_children(process)
                     result = None
             else:
+                # If there is no timeout limit, wait execution to finish
                 result = result_queue.get()
+                # Terminate/kill all processes after evaluation
                 self._kill_process_and_its_children(process)
             return result
         except Exception as e:

py_adtools-0.1.3/adtools/evaluator_pool.py

@@ -0,0 +1,82 @@
+"""
+Copyright (c) 2025 Rui Zhang <rzhang.cs@gmail.com>
+
+NOTICE: This code is under MIT license. This code is intended for academic/research purposes only.
+Commercial use of this software or its derivatives requires prior written permission.
+"""
+
+import time
+from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
+from typing import Literal, Optional
+
+from .evaluator import PyEvaluator
+from .py_code import PyProgram
+
+
+class EvaluatorExecutorPool:
+    def __init__(
+            self,
+            evaluator: PyEvaluator,
+            max_workers: int,
+            pool_type: Literal['thread', 'process'] = 'thread'
+    ):
+        """Multi-thread/process executor pool for parallel evaluation.
+        Args:
+            evaluator: The PyEvaluator instance.
+            max_workers: The maximum number of workers.
+            pool_type: Type of the executor pool.
+        """
+        self._evaluator = evaluator
+        self._max_workers = max_workers
+        if pool_type == 'thread':
+            self._pool = ThreadPoolExecutor(max_workers=self._max_workers)
+        else:
+            self._pool = ProcessPoolExecutor(max_workers=self._max_workers)
+
+    def evaluate(self, program: str | PyProgram, return_time=True, **kwargs):
+        """Evaluate program.
+        Args:
+            program: the program to be evaluated.
+            **kwargs: additional keyword arguments to pass to 'evaluate_program'.
+        """
+        start_time = time.time()
+        future = self._pool.submit(self._evaluator.evaluate, program, **kwargs)
+        res = future.result()
+        duration = time.time() - start_time
+        if return_time:
+            return res, duration
+        else:
+            return res
+
+    def secure_evaluate(
+            self,
+            program: str | PyProgram,
+            timeout_seconds: Optional[float],
+            redirect_to_devnull: bool = True,
+            multiprocessing_start_method: Literal['default', 'auto', 'fork', 'spawn'] = 'auto',
+            return_time=True,
+            **kwargs
+    ):
+        """Evaluate program in a new process. This enables timeout restriction and output redirection.
+        Args:
+            program: the program to be evaluated.
+            timeout_seconds: return 'None' if the execution time exceeds 'timeout_seconds'.
+            redirect_to_devnull: redirect any output to '/dev/null'.
+            multiprocessing_start_method: start a process using 'fork' or 'spawn'.
+            **kwargs: additional keyword arguments to pass to 'evaluate_program'.
+        """
+        start_time = time.time()
+        future = self._pool.submit(
+            self._evaluator.secure_evaluate,
+            program,
+            timeout_seconds,
+            redirect_to_devnull,
+            multiprocessing_start_method,
+            **kwargs
+        )
+        res = future.result()
+        duration = time.time() - start_time
+        if return_time:
+            return res, duration
+        else:
+            return res

py_adtools-0.1.3/adtools/lm_base.py

@@ -0,0 +1,403 @@
+"""
+Copyright (c) 2025 Rui Zhang <rzhang.cs@gmail.com>
+
+NOTICE: This code is under MIT license. This code is intended for academic/research purposes only.
+Commercial use of this software or its derivatives requires prior written permission.
+"""
+
+from abc import abstractmethod
+from typing import Optional, List, Literal, Dict, Any
+import os
+import subprocess
+import sys
+from pathlib import Path
+import psutil
+import requests
+import time
+
+import openai.types.chat
+
+__all__ = ['LanguageModel', 'OpenAIAPI', 'VLLMServer']
+
+
+class LanguageModel:
+    """Base class for language model interface."""
+
+    @abstractmethod
+    def chat_completion(
+            self,
+            message: str | List[openai.types.chat.ChatCompletionMessageParam],
+            max_tokens: int,
+            timeout_seconds: float,
+            *args,
+            **kwargs
+    ):
+        """Send a chat completion query with OpenAI format to the vLLM server. Return the response content.
+        Args:
+            message: The message in str or openai format.
+            max_tokens: The maximum number of tokens to generate.
+            timeout_seconds: The timeout seconds.
+        """
+        pass
+
+    def close(self):
+        """Release resources (if necessary)."""
+        pass
+
+
+class OpenAIAPI(LanguageModel):
+    def __init__(
+            self,
+            model: str,
+            base_url: str = None,
+            api_key: str = None,
+            **openai_init_kwargs
+    ):
+        super().__init__()
+        # If base_url is set to None, find 'OPENAI_BASE_URL' in environment variables
+        if base_url is None:
+            if 'OPENAI_BASE_URL' not in os.environ:
+                raise RuntimeError('If "base_url" is None, the environment variable OPENAI_BASE_URL must be set.')
+            else:
+                base_url = os.environ['OPENAI_BASE_URL']
+
+        # If api_key is set to None, find 'OPENAI_API_KEY' in environment variables
+        if api_key is None:
+            if 'OPENAI_API_KEY' not in os.environ:
+                raise RuntimeError('If "api_key" is None, OPENAI_API_KEY must be set.')
+            else:
+                api_key = os.environ['OPENAI_API_KEY']
+
+        self._model = model
+        self._client = openai.OpenAI(
+            api_key=api_key,
+            base_url=base_url,
+            **openai_init_kwargs
+        )
+
+    def chat_completion(
+            self,
+            message: str | List[openai.types.chat.ChatCompletionMessageParam],
+            max_tokens: int,
+            timeout_seconds: float,
+            *args,
+            **kwargs
+    ):
+        """Send a chat completion query with OpenAI format to the vLLM server. Return the response content.
+        Args:
+            message: The message in str or openai format.
+            max_tokens: The maximum number of tokens to generate.
+            timeout_seconds: The timeout seconds.
+        """
+        if isinstance(message, str):
+            message = [{'role': 'user', 'content': message.strip()}]
+
+        response = self._client.chat.completions.create(
+            model=self._model,
+            messages=message,
+            stream=False,
+            max_tokens=max_tokens,
+            timeout=timeout_seconds,
+            *args,
+            **kwargs,
+        )
+        return response.choices[0].message.content
+
+
+def _print_cmd_list(cmd_list, gpus, host, port):
+    print('\n' + '=' * 80)
+    print(f'[vLLM] Launching vLLM on GPU:{gpus}; URL: https://{host}:{port}')
+    print('=' * 80)
+    cmd = cmd_list[0] + ' \\\n'
+    for c in cmd_list[1:]:
+        cmd += '    ' + c + ' \\\n'
+    print(cmd.strip())
+    print('=' * 80 + '\n', flush=True)
+
+
+class VLLMServer:
+    def __init__(self,
+                 model_path: str,
+                 port: int,
+                 gpus: int | list[int],
+                 tokenizer_path: Optional[str] = None,
+                 max_model_len: int = 16384,
+                 max_lora_rank: Optional[int] = None,
+                 host: str = '0.0.0.0',
+                 mem_util: float = 0.85,
+                 deploy_timeout_seconds: int = 600,
+                 enforce_eager: bool = False,
+                 vllm_log_level: Literal['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'] = 'INFO',
+                 silent_mode: bool = False,
+                 env_variable_dict: Optional[Dict[str, str]] = None,
+                 vllm_serve_args: Optional[List[str]] = None,
+                 vllm_serve_kwargs: Optional[Dict[str, str]] = None,
+                 chat_template_kwargs: Optional[Dict[str, Any]] = None):
+        """Deploy an LLM on specified GPUs.
+        Args:
+            model_path: Path to the model to deploy.
+            tokenizer_path: Path to the tokenizer to use.
+            port: List of ports to deploy.
+            gpus: List of GPUs to deploy.
+            max_lora_rank: Max rank of LoRA adapter. Defaults to `None` which disables LoRA adapter.
+            host: Host address for vLLM server.
+            mem_util: Memory utility for each vLLM deployment.
+            deploy_timeout_seconds: Timeout to deploy (in seconds).
+            enforce_eager: Enforce eager mode.
+            vllm_log_level: Log level of vLLM server.
+            silent_mode: Silent mode.
+            env_variable_dict: Environment variables to use for vLLM server, e.g., {'KEY': 'VALUE'}.
+            vllm_serve_args: Arguments to pass to vLLM server, e.g., ['--enable-reasoning'].
+            vllm_serve_kwargs: Keyword arguments to pass to vLLM server, e.g., {'--reasoning-parser': 'deepseek-r1'}.
+
+        Example:
+            # deploy a model on GPU 0 and 1
+            llm = VLLMServer(
+                model_path='path/to/model',
+                tokenizer_path='path/to/tokenizer',
+                gpus=[0, 1],  # set gpus=0 or gpus=[0] if you only use one GPU
+                port=12001,
+                mem_util=0.8
+            )
+            # draw sample using base model
+            llm.draw_sample('hello')
+
+            # load adapter and draw sample
+            llm.load_lora_adapter('adapter_1', '/path/to/adapter')
+            llm.draw_sample('hello', lora_name='adapter_1')
+
+            # unload adapter
+            llm.unload_lora_adapter('adapter_1')
+
+            # release resources
+            llm.close()
+        """
+        self._model_path = model_path
+        self._port = port
+        self._gpus = gpus
+        self._tokenizer_path = tokenizer_path if tokenizer_path is not None else model_path
+        self._max_model_len = max_model_len
+        self._max_lora_rank = max_lora_rank
+        self._host = host
+        self._mem_util = mem_util
+        self._deploy_timeout_seconds = deploy_timeout_seconds
+        self._enforce_eager = enforce_eager
+        self._vllm_log_level = vllm_log_level
+        self._silent_mode = silent_mode
+        self._env_variable_dict = env_variable_dict
+        self._vllm_serve_args = vllm_serve_args
+        self._vllm_serve_kwargs = vllm_serve_kwargs
+        self._chat_template_kwargs = chat_template_kwargs
+
+        # Deploy vLLMs
+        self._process = self._launch_vllm()
+        self._wait_for_vllm()
+
+    def _launch_vllm(self):
+        """Launch a vLLM server and return the subprocess.
+        """
+        if isinstance(self._gpus, int):
+            gpus = str(self._gpus)
+        else:
+            gpus = ','.join([str(g) for g in self._gpus])
+
+        executable_path = sys.executable
+        cmd = [
+            executable_path, '-m',
+            'vllm.entrypoints.openai.api_server',
+            '--model', self._model_path,
+            '--tokenizer', self._tokenizer_path,
+            '--max_model_len', str(self._max_model_len),
+            '--host', self._host,
+            '--port', str(self._port),
+            '--gpu-memory-utilization', str(self._mem_util),
+            '--tensor-parallel-size', str(len(self._gpus)) if isinstance(self._gpus, list) else '1',
+            '--trust-remote-code',
+            '--chat-template-content-format', 'string',
+        ]
+
+        if self._enforce_eager:
+            cmd.append('--enforce_eager')
+
+        # Other args for vllm serve
+        if self._vllm_serve_args is not None:
+            for arg in self._vllm_serve_args:
+                cmd.append(arg)
+
+        # Other kwargs for vllm serve
+        if self._vllm_serve_kwargs is not None:
+            for kwarg, value in self._vllm_serve_kwargs.items():
+                cmd.extend([kwarg, value])
+
+        # Environmental variables
+        env = os.environ.copy()
+        env['CUDA_VISIBLE_DEVICES'] = gpus
+        env['VLLM_LOGGING_LEVEL'] = self._vllm_log_level
+
+        # FIXME: These code are required for my machine :-(
+        # FIXME: This may due to the bad NCCL environment configuration :-(
+        if isinstance(self._gpus, list) and len(self._gpus) > 1:
+            # set NCCL environment variable
+            env['NCCL_P2P_DISABLE'] = '1'
+            # disable custom all reduce
+            cmd.append('--disable-custom-all-reduce')
+
+        # Enable LoRA dynamic loading
+        if self._max_lora_rank is not None:
+            cmd.extend([
+                '--enable-lora',
+                '--max-lora-rank', str(self._max_lora_rank),
+            ])
+            env['VLLM_ALLOW_RUNTIME_LORA_UPDATING'] = 'True'
+
+        # Other env variables
+        if self._env_variable_dict is not None:
+            for k, v in self._env_variable_dict.items():
+                env[k] = v
+
+        _print_cmd_list(cmd, gpus=self._gpus, host=self._host, port=self._port)
+
+        # Launch vllm using subprocess
+        stdout = Path(os.devnull).open('w') if self._silent_mode else None
+        proc = subprocess.Popen(cmd, env=env, stdout=stdout, stderr=subprocess.STDOUT)
+        return proc
+
+    def _kill_vllm_process(self):
+        try:
+            # Get child processes before terminating parent
+            try:
+                parent = psutil.Process(self._process.pid)
+                children = parent.children(recursive=True)
+            except psutil.NoSuchProcess:
+                children = []
+
+            # Terminate parent process
+            self._process.terminate()
+            self._process.wait(timeout=5)
+            print(f'[vLLM] terminated process: {self._process.pid}')
+
+            # Kill any remaining children
+            for child in children:
+                try:
+                    child.terminate()
+                    child.wait(timeout=2)
+                except (psutil.NoSuchProcess, psutil.TimeoutExpired):
+                    try:
+                        child.kill()
+                    except psutil.NoSuchProcess:
+                        pass
+        except subprocess.TimeoutExpired:
+            self._process.kill()
+            print(f'[vLLM] killed process: {self._process.pid}')
+
+    def _wait_for_vllm(self):
+        """Check each vLLM server's state and check /health. Kill all vLLM server processes if timeout.
+        """
+        for _ in range(self._deploy_timeout_seconds):
+            # check process status
+            if self._process.poll() is not None:
+                sys.exit(f'[vLLM] crashed (exit {self._process.returncode})')
+
+            # check server status
+            health = f'http://{self._host}:{self._port}/health'
+            try:
+                if requests.get(health, timeout=1).status_code == 200:
+                    return
+            except Exception:
+                pass
+            time.sleep(1)
+
+        # Servers fail to initialize
+        print('[vLLM] failed to start within timeout')
+        self._kill_vllm_process()
+        sys.exit('[vLLM] failed to start within timeout')
+
+    def unload_lora_adapter(self, lora_name: str):
+        """Unload lora adapter given the lora name.
+        Args:
+            lora_name: Lora adapter name.
+        """
+        lora_api_url = f'http://{self._host}:{self._port}/v1/unload_lora_adapter'
+        headers = {'Content-Type': 'application/json'}
+        try:
+            payload = {'lora_name': lora_name}
+            requests.post(lora_api_url, json=payload, headers=headers, timeout=10)
+        except requests.exceptions.RequestException:
+            pass
+
+    def load_lora_adapter(self, lora_name: str, new_adapter_path: str, num_trails: int = 5):
+        """Dynamically load a LoRA adapter.
+        Args:
+            lora_name: LoRA adapter name.
+            new_adapter_path: Path to the new LoRA adapter weights.
+        """
+        # First unload lora adapter
+        self.unload_lora_adapter(lora_name)
+
+        if self._max_lora_rank is None:
+            raise ValueError('LoRA is not enabled for this VLLMServer instance, since "max_lora_rank" is not set.')
+
+        # Prepare the payload for LoRA update
+        payload = {'lora_name': lora_name, 'lora_path': new_adapter_path}
+        headers = {'Content-Type': 'application/json'}
+        lora_api_url = f'http://{self._host}:{self._port}/v1/load_lora_adapter'
+
+        # Repeatedly trying to load lora adapters
+        for i in range(num_trails):
+            try:
+                response = requests.post(lora_api_url, json=payload, headers=headers, timeout=60)
+                if response.status_code == 200:
+                    print(f'[vLLM] Successfully load LoRA adapter: {lora_name} from {new_adapter_path}')
+                else:
+                    print(f'[vLLM] Failed to load LoRA adapter. '
+                          f'Status code: {response.status_code}, Response: {response.text}')
+                return True
+            except requests.exceptions.RequestException as e:
+                continue
+
+        print(f'[vLLM] Error loading LoRA adapter: {str(e)}')
+        return False
+
+    def close(self):
+        """Shut down vLLM server and kill all vLLM processes."""
+        self._kill_vllm_process()
+
+    def chat_completion(self,
+                        message: str | List[openai.types.chat.ChatCompletionMessageParam],
+                        max_tokens: Optional[int] = None,
+                        timeout_seconds: Optional[int] = None,
+                        lora_name: Optional[str] = None,
+                        temperature: float = 0.9,
+                        top_p: float = 0.9,
+                        chat_template_kwargs: Optional[Dict[str, Any]] = None) -> str:
+        """Send a chat completion query with OpenAI format to the vLLM server. Return the response content.
+        Args:
+            message: The message in str or openai format.
+            max_tokens: The maximum number of tokens to generate.
+            timeout_seconds: The timeout seconds.
+            lora_name: Lora adapter name. Defaults to None which uses base model.
+            temperature: The temperature parameter.
+            top_p: The top p parameter.
+            chat_template_kwargs: The chat template kwargs, e.g., {'enable_thinking': False}.
+        """
+        data = {
+            'messages': [
+                {'role': 'user', 'content': message.strip()} if isinstance(message, str) else message
+            ],
+            'temperature': temperature,
+            'top_p': top_p,
+            'max_tokens': max_tokens,
+        }
+        # Use the specified lora adapter
+        if lora_name is not None:
+            data['model'] = lora_name
+        # Chat template keyword args
+        if self._chat_template_kwargs is not None:
+            data['chat_template_kwargs'] = self._chat_template_kwargs
+        elif chat_template_kwargs is not None:
+            data['chat_template_kwargs'] = chat_template_kwargs
+        # Request
+        url = f'http://{self._host}:{self._port}/v1/chat/completions'
+        headers = {'Content-Type': 'application/json'}
+        response = requests.post(url, headers=headers, json=data, timeout=timeout_seconds)
+        return response.json()['choices'][0]['message']['content']

{py_adtools-0.1.2 → py_adtools-0.1.3}/adtools/py_code.py

@@ -10,24 +10,27 @@ import dataclasses
 import textwrap
 from typing import List, Optional
 
-__all__ = ['PyScript', 'PyFunction', 'PyClass', 'PyProgram']
+__all__ = ['PyCodeBlock', 'PyFunction', 'PyClass', 'PyProgram']
 
 
 @dataclasses.dataclass
-class PyScript:
-    """A parsed Python script block (top-level code that's not in classes/functions).
+class PyCodeBlock:
+    """A parsed Python code block (e.g., top-level code that's not in classes/functions).
     """
     code: str
 
     def __str__(self) -> str:
-        return self.code + '\n'
+        return self.code
+
+    def __repr__(self) -> str:
+        return self.__str__() + '\n'
 
 
 @dataclasses.dataclass
 class PyFunction:
     """A parsed Python function.
-    > Part of this class is referenced from:
-    > https://github.com/google-deepmind/funsearch/blob/main/implementation/code_manipulation.py
+    Part of this class is referenced from:
+    https://github.com/google-deepmind/funsearch/blob/main/implementation/code_manipulation.py
     """
     decorator: str
     name: str
@@ -38,16 +41,20 @@ class PyFunction:
 
     def __str__(self) -> str:
         return_type = f' -> {self.return_type}' if self.return_type else ''
-        function = f'{self.decorator}\ndef {self.name}({self.args}){return_type}:\n'
+        function = f'{self.decorator}\n' if self.decorator else ''
+        function += f'def {self.name}({self.args}){return_type}:\n'
         if self.docstring:
             # The self.docstring is already indented on every line except the first one.
-            # Here, we assume the indentation is always four spaces.
+            # Here, we assume the indentation is always 4 spaces.
             new_line = '\n' if self.body else ''
             function += f'    """{self.docstring}"""{new_line}'
         # The self.body is already indented.
-        function += self.body + '\n\n'
+        function += self.body
         return function
 
+    def __repr__(self) -> str:
+        return self.__str__() + '\n\n'
+
     def __setattr__(self, name: str, value: str) -> None:
         # Ensure there aren't leading & trailing new lines in `body`
         if name == 'body':
@@ -82,13 +89,14 @@ class PyClass:
     decorator: str
     name: str
     bases: str
-    class_vars_and_code: List[PyScript] = None
+    class_vars_and_code: List[PyCodeBlock] = None
     docstring: str | None = None
     functions: list[PyFunction] = dataclasses.field(default_factory=list)
-    functions_class_vars_and_code: List[PyScript | PyFunction] | None = None
+    functions_class_vars_and_code: List[PyCodeBlock | PyFunction] | None = None
 
     def __str__(self) -> str:
-        class_def = f'{self.decorator}\nclass {self.name}'
+        class_def = f'{self.decorator}\n' if self.decorator else ''
+        class_def += f'class {self.name}'
         if self.bases:
             class_def += f'({self.bases})'
         class_def += ':\n'
@@ -96,15 +104,21 @@ class PyClass:
         if self.docstring:
             class_def += f'    """{self.docstring}"""\n'
 
-        for item in self.functions_class_vars_and_code:
-            if isinstance(item, PyScript):
-                class_def += f'{str(item)}\n'
+        for i, item in enumerate(self.functions_class_vars_and_code):
+            if isinstance(item, PyCodeBlock):
+                # The PyCodeBlock has already indented
+                class_def += f'{str(item)}'
             else:
                 # Add functions with an extra level of indentation
                 class_def += textwrap.indent(str(item).strip(), '    ')
+            # Add '\n\n' if this is not the last element
+            if i != len(self.functions_class_vars_and_code) - 1:
                 class_def += '\n\n'
         return class_def
 
+    def __repr__(self):
+        return self.__str__() + '\n\n'
+
     def __setattr__(self, name: str, value: str) -> None:
         # Ensure there aren't leading & trailing new lines in `body`
         if name == 'body':
@@ -137,19 +151,19 @@ class PyClass:
 class PyProgram:
     """A parsed Python program."""
 
-    scripts: list[PyScript]  # Top-level code that's not in classes/functions
+    scripts: list[PyCodeBlock]  # Top-level code that's not in classes/functions
     functions: list[PyFunction]  # Top-level functions in the code
     classes: list[PyClass]  # Top-level classes in the code
-    classes_functions_scripts: list[PyFunction | PyClass | PyScript]
+    classes_functions_scripts: list[PyFunction | PyClass | PyCodeBlock]
 
     def __str__(self) -> str:
         program = ''
         for class_or_func_or_script in self.classes_functions_scripts:
-            program += str(class_or_func_or_script) + '\n'
+            program += str(class_or_func_or_script) + '\n\n'
         return program
 
     @classmethod
-    def from_text(cls, text: str) -> Optional['PyProgram']:
+    def from_text(cls, text: str) -> 'PyProgram':
         tree = ast.parse(text)
         visitor = _ProgramVisitor(text)
         visitor.visit(tree)
@@ -163,19 +177,31 @@ class _ProgramVisitor(ast.NodeVisitor):
 
     def __init__(self, sourcecode: str):
         self._codelines: list[str] = sourcecode.splitlines()
-        self._scripts: list[PyScript] = []
+        self._scripts: list[PyCodeBlock] = []
         self._functions: list[PyFunction] = []
         self._classes: list[PyClass] = []
-        self._classes_functions_scripts: list[PyFunction | PyClass | PyScript] = []
+        self._classes_functions_scripts: list[PyFunction | PyClass | PyCodeBlock] = []
         self._last_script_end = 0
 
+    def _get_code(self, start_line: int, end_line: int, dedent=False):
+        """Get code between start_line and end_line in 'self._codelines'.
+        """
+        code = []
+        for line in self._codelines[start_line: end_line]:
+            if dedent:
+                code.append(line[4:])
+            else:
+                code.append(line)
+        return '\n'.join(code).rstrip()
+
     def _add_script(self, start_line: int, end_line: int):
-        """Add a script segment from the code."""
+        """Add a script segment from the code.
+        """
         if start_line >= end_line:
             return
-        script_code = '\n'.join(self._codelines[start_line:end_line]).strip()
+        script_code = self._get_code(start_line, end_line).strip()
         if script_code:
-            script = PyScript(code=script_code)
+            script = PyCodeBlock(code=script_code)
             self._scripts.append(script)
             self._classes_functions_scripts.append(script)
 
@@ -189,11 +215,11 @@ class _ProgramVisitor(ast.NodeVisitor):
             if has_decorators:
                 # Find the minimum line number and retain the code above
                 decorator_start_line = min(decorator.lineno for decorator in node.decorator_list)
-                decorator = '\n'.join(self._codelines[decorator_start_line - 1: node.lineno - 1]).strip()
+                decorator = self._get_code(decorator_start_line - 1, node.lineno - 1)
                 # Update script end line
                 script_end_line = decorator_start_line - 1
             else:
-                decorator = ''
+                decorator = None
                 script_end_line = node.lineno - 1
 
             # Add any script code before this function
@@ -204,6 +230,7 @@ class _ProgramVisitor(ast.NodeVisitor):
             body_start_line = node.body[0].lineno - 1
             docstring = None
 
+            # If the first node is ast.Expr, we regard it as a docstring
             if isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant):
                 docstring = ast.literal_eval(ast.unparse(node.body[0])).strip()
                 if len(node.body) > 1:
@@ -211,13 +238,14 @@ class _ProgramVisitor(ast.NodeVisitor):
                 else:
                     body_start_line = function_end_line
 
+            # Return a PyFunction instance
             func = PyFunction(
                 decorator=decorator,
                 name=node.name,
                 args=ast.unparse(node.args),
                 return_type=ast.unparse(node.returns) if node.returns else None,
                 docstring=docstring,
-                body='\n'.join(self._codelines[body_start_line: function_end_line]),
+                body=self._get_code(body_start_line, function_end_line)
             )
 
             self._functions.append(func)
@@ -234,11 +262,11 @@ class _ProgramVisitor(ast.NodeVisitor):
             if has_decorators:
                 # Find the minimum line number and retain the code above
                 decorator_start_line = min(decorator.lineno for decorator in node.decorator_list)
-                class_decorator = '\n'.join(self._codelines[decorator_start_line - 1: node.lineno - 1])
+                class_decorator = self._get_code(decorator_start_line - 1, node.lineno - 1)
                 # Update script end line
                 script_end_line = decorator_start_line - 1
             else:
-                class_decorator = ''
+                class_decorator = None
                 script_end_line = node.lineno - 1
 
             # Add any script code before this class
@@ -250,9 +278,12 @@ class _ProgramVisitor(ast.NodeVisitor):
             if isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant):
                 docstring = ast.literal_eval(ast.unparse(node.body[0]))
 
+            # Record methods
             methods = []
+            # Record class variables or code that are not methods
             class_vars_and_code = []
-            function_class_vars_and_code = []  # record the order of function and class vars and code
+            # Record the order of function and class vars and code
+            function_class_vars_and_code = []
 
             # Traverse each body, if there is a docstring, skip body[0]
             for item in node.body if docstring is None else node.body[1:]:
@@ -263,13 +294,9 @@ class _ProgramVisitor(ast.NodeVisitor):
                         # Find the minimum line number and retain the code above
                         decorator_start_line = min(decorator.lineno for decorator in item.decorator_list)
                         # Dedent decorator code
-                        decorator = []
-                        for line in range(decorator_start_line - 1, item.lineno - 1):
-                            dedented_decorator = self._codelines[line].strip()
-                            decorator.append(dedented_decorator)
-                        decorator = '\n'.join(decorator)
+                        decorator = self._get_code(decorator_start_line - 1, item.lineno - 1, dedent=True)
                     else:
-                        decorator = ''
+                        decorator = None
 
                     method_end_line = item.end_lineno
                     method_body_start_line = item.body[0].lineno - 1
@@ -284,11 +311,7 @@ class _ProgramVisitor(ast.NodeVisitor):
                             method_body_start_line = method_end_line
 
                     # Extract function body and dedent for 4 spaces
-                    body = []
-                    for line in range(method_body_start_line, method_end_line):
-                        dedented_body = self._codelines[line][4:]
-                        body.append(dedented_body)
-                    body = '\n'.join(body)
+                    body = self._get_code(method_body_start_line, method_end_line, dedent=True)
 
                     py_func = PyFunction(
                         decorator=decorator,
@@ -300,17 +323,16 @@ class _ProgramVisitor(ast.NodeVisitor):
                     )
                     methods.append(py_func)
                     function_class_vars_and_code.append(py_func)
-                else:  # If the item is not a function definition,add to class variables and code
-                    code = []
-                    for i in range(item.lineno - 1, item.end_lineno):
-                        code.append(self._codelines[i])
-                    py_script = PyScript(code='\n'.join(code))
+                else:  # If the item is not a function definition, add to class variables and code
+                    code = self._get_code(item.lineno - 1, item.end_lineno)
+                    py_script = PyCodeBlock(code=code)
                     class_vars_and_code.append(py_script)
                     function_class_vars_and_code.append(py_script)
 
             # Get base classes
-            bases = ', '.join([ast.unparse(base) for base in node.bases]) if node.bases else ''
+            bases = ', '.join([ast.unparse(base) for base in node.bases]) if node.bases else None
 
+            # Return a PyClass instance
             class_ = PyClass(
                 decorator=class_decorator,
                 name=node.name,

{py_adtools-0.1.2 → py_adtools-0.1.3}/py_adtools.egg-info/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: py-adtools
-Version: 0.1.2
-Summary: Useful tools for parsing and evaluating Python programs for algorithm design.
+Version: 0.1.3
+Summary: Useful tools for parsing and evaluating Python programs for LLM-based algorithm design.
 Home-page: https://github.com/RayZhhh/py-adtools
 Author: Rui Zhang
 Author-email: rzhang.cs@gmail.com
@@ -34,7 +34,7 @@ Dynamic: summary
 
 ------
 
-The figure demonstrates how a Python program is parsed into `PyScript`, `PyFunction`, `PyClass,` and `PyProgram` via `adtools`.
+The figure demonstrates how a Python program is parsed into `PyCodeBlock`, `PyFunction`, `PyClass,` and `PyProgram` via `adtools`.
 
 ![pycode](./assets/pycode.png)
 
@@ -68,7 +68,7 @@ Parse your code (in string) into Python code instances, so that you can check ea
 from adtools import PyProgram
 
 code = r'''
-import ast, numba                 # This part will be parsed into PyScript
+import ast, numba                 # This part will be parsed into PyCodeBlock
 import numpy as np
 
 @numba.jit()                      # This part will be parsed into PyFunction
@@ -80,8 +80,9 @@ def function(arg1, arg2=True):
 
 @some.decorators()                # This part will be parsed into PyClass
 class PythonClass(BaseClass):
-    class_var1 = 1                # This part will be parsed into PyScript
-    class_varb = 2                # and placed in PyClass.class_vars_and_code
+    
+    class_var1 = 1                # This part will be parsed into PyCodeBlock
+    class_var2 = 2                # and placed in PyClass.class_vars_and_code
 
     def __init__(self, x):        # This part will be parsed into PyFunction
         self.x = x                # and placed in PyClass.functions
@@ -93,11 +94,11 @@ class PythonClass(BaseClass):
     def method2(self, x, y):
     	return x + y + self.method1(x)
 
-    class InnerClass:             # This part will be parsed into PyScript
+    class InnerClass:             # This part will be parsed into PyCodeBlock
     	def __init__(self):       # and placed in PyClass.class_vars_and_code
     		...
 
-if __name__ == '__main__':        # This part will be parsed into PyScript
+if __name__ == '__main__':        # This part will be parsed into PyCodeBlock
 	res = function(1)
 	print(res)
 	res = PythonClass().method2(1, 2)
@@ -116,7 +117,7 @@ print(p.functions[0].name)
 Evaluate Python programs in a secure process to avoid the abortation of the main process. Two steps:
 
 - Extend the `PyEvaluator` class and override the `evaluate_program` method.
-- Evaluate the program (in str) by calling the `evaluate` (directly evaluate without executing in a sandbox process) or the `secure_evaluate` (evaluate in a sandbox process) methods. 
+- Evaluate the program (in str) by calling the `evaluate` (directly evaluate without executing in a sandbox process) or the `secure_evaluate` (evaluate in a sandbox process) methods.
 
 ```python
 import time
@@ -190,6 +191,7 @@ def merge(left, right):
 
 harmful_code_generated_by_llm = '''
 def merge_sort(arr):
+    print('I am harmful')  # There will be no output since we redirect STDOUT to /dev/null by default.
     while True:
         pass
 '''

{py_adtools-0.1.2 → py_adtools-0.1.3}/py_adtools.egg-info/SOURCES.txt

@@ -3,6 +3,8 @@ README.md
 setup.py
 adtools/__init__.py
 adtools/evaluator.py
+adtools/evaluator_pool.py
+adtools/lm_base.py
 adtools/py_code.py
 py_adtools.egg-info/PKG-INFO
 py_adtools.egg-info/SOURCES.txt

{py_adtools-0.1.2 → py_adtools-0.1.3}/setup.py

@@ -5,10 +5,10 @@ with open('README.md', 'r', encoding='utf-8') as fh:
 
 setup(
     name='py-adtools',
-    version='0.1.2',
+    version='0.1.3',
     author='Rui Zhang',
     author_email='rzhang.cs@gmail.com',
-    description='Useful tools for parsing and evaluating Python programs for algorithm design.',
+    description='Useful tools for parsing and evaluating Python programs for LLM-based algorithm design.',
     long_description=long_description,
     long_description_content_type='text/markdown',
     url='https://github.com/RayZhhh/py-adtools',

py_adtools-0.1.2/adtools/__init__.py

@@ -1,2 +0,0 @@
-from .py_code import PyScript, PyFunction, PyClass, PyProgram
-from .evaluator import PyEvaluator