oocana-python-executor 0.15.0__tar.gz

oocana_python_executor-0.15.0/PKG-INFO

@@ -0,0 +1,9 @@
+Metadata-Version: 2.1
+Name: oocana-python-executor
+Version: 0.15.0
+Summary: a client subscribe mqtt topic to execute oocana's block
+Author-Email: l1shen <lishen1635@gmail.com>, yleaf <11785335+leavesster@users.noreply.github.com>
+License: MIT
+Requires-Python: >=3.9
+Requires-Dist: oocana
+

oocana_python_executor-0.15.0/pyproject.toml

@@ -0,0 +1,27 @@
+[project]
+name = "oocana-python-executor"
+version = "0.15.0"
+authors = [
+    { name = "l1shen", email = "lishen1635@gmail.com" },
+    { name = "yleaf", email = "11785335+leavesster@users.noreply.github.com" },
+]
+description = "a client subscribe mqtt topic to execute oocana's block"
+requires-python = ">=3.9"
+dependencies = [
+    "oocana",
+]
+
+[project.license]
+text = "MIT"
+
+[project.scripts]
+python-executor = "python_executor.executor:main"
+
+[build-system]
+requires = [
+    "pdm-backend",
+]
+build-backend = "pdm.backend"
+
+[tool.pdm]
+distribution = true

oocana_python_executor-0.15.0/python_executor/block.py

@@ -0,0 +1,189 @@
+from oocana import Context, Mainframe
+from dataclasses import dataclass
+from typing import Optional, TypedDict
+import inspect
+import traceback
+import logging
+from .data import store, vars, EXECUTOR_NAME
+from .context import createContext
+from .hook import ExitFunctionException
+import os
+import sys
+import importlib
+import importlib.util
+import threading
+
+
+class ExecutorOptionsDict(TypedDict):
+    function: Optional[str]
+    entry: Optional[str]
+    source: Optional[str]
+
+# entry 与 source 是二选一的存在
+class ExecutorDict(TypedDict):
+    options: Optional[ExecutorOptionsDict]
+
+tmp_files = set()
+
+@dataclass
+class ExecutePayload:
+    session_id: str
+    job_id: str
+    dir: str
+    executor: ExecutorDict
+    outputs: Optional[dict] = None
+
+    def __init__(self, *args, **kwargs):
+        if args:
+            self.session_id = args[0]
+            self.job_id = args[1]
+            self.executor = args[2]
+            self.dir = args[3]
+            self.outputs = args[4]
+        if kwargs:
+            for key, value in kwargs.items():
+                setattr(self, key, value)
+
+lock = threading.Lock()
+
+def load_module(file_path: str, source_dir=None):
+
+    if (os.path.isabs(file_path)):
+        file_abs_path = file_path
+    else:
+        dirname = source_dir if source_dir else os.getcwd()
+        file_abs_path = os.path.abspath(os.path.join(dirname, file_path))
+    with lock:
+        if file_abs_path in sys.modules:
+            return sys.modules[file_abs_path]
+
+        module_dir = os.path.dirname(file_abs_path)
+        sys.path.insert(0, module_dir)
+
+        file_spec = importlib.util.spec_from_file_location(file_abs_path, file_abs_path)
+        module = importlib.util.module_from_spec(file_spec)  # type: ignore
+        sys.modules[file_abs_path] = module
+
+        file_spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+def output_return_object(obj, context: Context):
+    if obj is None:
+        context.done()
+    elif obj is context.keepAlive:
+        pass
+    elif isinstance(obj, dict):
+        for k, v in obj.items():
+            context.output(k, v)
+        context.done()
+    else:
+        context.done(f"return object needs to be a dictionary, but get type: {type(obj)}")
+
+logger = logging.getLogger(EXECUTOR_NAME)
+
+async def run_block(message, mainframe: Mainframe, session_dir: str):
+
+    logger.info(f"block {message.get('job_id')} start")
+    try:
+        payload = ExecutePayload(**message)
+        context = createContext(mainframe, payload.session_id, payload.job_id, store, payload.outputs, session_dir)
+    except Exception:
+        traceback_str = traceback.format_exc()
+        # rust 那边会保证传过来的 message 一定是符合格式的，所以这里不应该出现异常。这里主要是防止 rust 修改错误。
+        mainframe.send(
+        {
+            "job_id": message["job_id"],
+            "session_id": message["session_id"],
+        },
+        {
+            "type": "BlockFinished",
+            "job_id": message["job_id"], 
+            "error": traceback_str
+        })
+        return
+
+    vars.set(context)
+
+    load_dir = payload.dir
+
+    options = payload.executor.get("options")
+
+    node_id = context.node_id
+
+    file_path = options["entry"] if options is not None and options.get("entry") is not None else 'index.py'
+
+    source = options.get("source") if options is not None else None
+    if source is not None:
+        if not os.path.exists(load_dir):
+            os.makedirs(load_dir)
+        
+        dir_path = os.path.join(load_dir, ".scriptlets")
+        if not os.path.exists(dir_path):
+            os.makedirs(dir_path)
+        tmp_py = os.path.join(dir_path, f"{node_id}.py")
+        # 记录临时文件，但是现在不再清理
+        tmp_files.add(tmp_py)
+
+        with open(tmp_py, "w") as f:
+            f.write(source)
+        file_path = tmp_py
+
+    try:
+        # TODO: 这里的异常处理，应该跟详细一些，提供语法错误提示。
+        index_module = load_module(file_path, load_dir) # type: ignore
+    except Exception:
+        traceback_str = traceback.format_exc()
+        context.done(traceback_str)
+        return
+    function_name: str = options.get("function") if payload.executor is not None and options.get("function") is not None else 'main' # type: ignore
+    fn = index_module.__dict__.get(function_name)
+
+    if fn is None:
+        context.done(f"function {function_name} not found in {file_path}")
+        return
+    if not callable(fn):
+        context.done(f"{function_name} is not a function in {file_path}")
+        return
+
+    try:
+        signature = inspect.signature(fn)
+        params_count = len(signature.parameters)
+        result = None
+        traceback_str = None
+
+        try:
+            if inspect.iscoroutinefunction(fn):
+                if params_count == 0:
+                    result = await fn()
+                elif params_count == 1:
+                    only_context_param = list(signature.parameters.values())[0].annotation is Context
+                    result = await fn(context) if only_context_param else await fn(context.inputs)
+                else:
+                    result = await fn(context.inputs, context)
+            else:
+                if params_count == 0:
+                    result = fn()
+                elif params_count == 1:
+                    only_context_param = list(signature.parameters.values())[0].annotation is Context
+                    result = fn(context) if only_context_param else fn(context.inputs)
+                else:
+                    result = fn(context.inputs, context)
+        except ExitFunctionException as e:
+            if e.args[0] is not None:
+                context.done("block call exit with message: " + str(e.args[0]))
+            else:
+                context.done()
+        except Exception:
+            traceback_str = traceback.format_exc()
+
+        if traceback_str is not None:
+            context.done(traceback_str)
+        else:
+            output_return_object(result, context)
+    except Exception:
+        traceback_str = traceback.format_exc()
+        context.done(traceback_str)
+    finally:
+        logger.info(f"block {message.get('job_id')} done")
+

oocana_python_executor-0.15.0/python_executor/context.py

@@ -0,0 +1,68 @@
+import logging
+from oocana import Mainframe, Context, StoreKey, BlockInfo, BinValueDict, VarValueDict, InputHandleDef, is_bin_value, is_var_value
+from typing import Dict
+from .secret import replace_secret
+import os.path
+from .logger import ContextHandler
+from .data import EXECUTOR_NAME
+
+logger = logging.getLogger(EXECUTOR_NAME)
+
+def createContext(
+    mainframe: Mainframe, session_id: str, job_id: str, store, output, session_dir: str
+) -> Context:
+
+    node_props = mainframe.notify_block_ready(session_id, job_id)
+
+    inputs_def: Dict[str, Dict] | None = node_props.get("inputs_def")
+    inputs = node_props.get("inputs")
+
+    if inputs_def is not None and inputs is not None:
+
+        inputs_def_handles: Dict[str, InputHandleDef] = {}
+        for k, v in inputs_def.items():
+            inputs_def_handles[k] = InputHandleDef(**v)
+
+        inputs = replace_secret(inputs, inputs_def_handles, node_props.get("inputs_def_patch"))
+
+        for k, v in inputs.items():
+            input_def = inputs_def_handles.get(k)
+            if input_def is None:
+                continue
+            if is_var_value(v):
+                wrap_var: VarValueDict = v
+                try:
+                    ref = StoreKey(**wrap_var["value"])
+                except:  # noqa: E722
+                    logger.warning(f"not valid object ref: {wrap_var}")
+                    continue
+                if ref in store:
+                    inputs[k] = store.get(ref)
+                else:
+                    logger.error(f"object {ref} not found in store")
+            elif is_bin_value(v):
+                wrap_bin: BinValueDict = v
+                path = wrap_bin["value"]
+                if isinstance(path, str):
+                    # check file path v is exist
+                    if not os.path.exists(path):
+                        logger.error(f"file {path} for oomol/bin is not found")
+                        continue
+
+                    with open(path, "rb") as f:
+                        inputs[k] = f.read()
+                else:
+                    logger.error(f"not valid bin handle: {v}")
+
+    if inputs is None:
+        inputs = {}
+    
+    blockInfo = BlockInfo(**node_props)
+
+    ctx = Context(inputs, blockInfo, mainframe, store, output, session_dir)
+    # 跟 executor 日志分开，避免有的库在 logger 里面使用 print，导致 hook 出现递归调用。
+    block_logger = logging.getLogger(f"block {job_id}")
+    ctx_handler = ContextHandler(ctx)
+    block_logger.addHandler(ctx_handler)
+    ctx._logger = logger
+    return ctx

oocana_python_executor-0.15.0/python_executor/data.py

@@ -0,0 +1,5 @@
+from contextvars import ContextVar
+from oocana import Context, EXECUTOR_NAME
+
+vars: ContextVar[Context] = ContextVar('context')
+store = {}

oocana_python_executor-0.15.0/python_executor/executor.py

@@ -0,0 +1,247 @@
+#!/usr/bin/env python
+
+import asyncio
+import os
+import queue
+import sys
+import logging
+from . import hook
+from oocana import Mainframe, ServiceExecutePayload
+from .utils import run_in_new_thread, run_async_code, oocana_dir
+from .block import run_block
+from oocana import EXECUTOR_NAME
+from .matplot.oomol_matplot_helper import import_helper, add_matplot_module
+from typing import Literal
+from .topic import prepare_report_topic, service_config_topic, run_action_topic, ServiceTopicParams, ReportStatusPayload, exit_report_topic, status_report_topic
+
+logger = logging.getLogger(EXECUTOR_NAME)
+service_store: dict[str, Literal["launching", "running"]] = {}
+job_set = set()
+
+# 日志目录 ~/.oocana/sessions/{session_id}
+# executor 的日志都会记录在 [python-executor-{suffix}.log | python-executor.log]
+# 全局 logger 会记录在 python-{suffix}.log | python.log
+def config_logger(session_id: str, suffix: str | None, output: Literal["console", "file"]):
+
+
+    format = '%(asctime)s - %(levelname)s - {%(pathname)s:%(lineno)d} - %(message)s'
+    fmt = logging.Formatter(format)
+    logger.setLevel(logging.DEBUG)
+    if output == "file":
+        executor_dir = os.path.join(oocana_dir(), "sessions", session_id)
+        logger_file = os.path.join(executor_dir, f"python-executor-{suffix}.log") if suffix is not None else os.path.join(executor_dir, "python-executor.log")
+
+        if not os.path.exists(logger_file):
+            os.makedirs(os.path.dirname(logger_file), exist_ok=True)
+
+        print(f"setup logging in file {logger_file}")
+        h = logging.FileHandler(logger_file)
+
+        global_logger_file = os.path.join(executor_dir, f"python-{suffix}.log") if suffix is not None else os.path.join(executor_dir, "python.log")
+        logging.basicConfig(filename=global_logger_file, level=logging.DEBUG, format=format)
+    else:
+        logging.basicConfig(level=logging.DEBUG, format=format)
+        h = logging.StreamHandler(sys.stdout)
+
+    h.setFormatter(fmt)
+    logger.addHandler(h)
+    # 跟全局日志分开。避免有的库在全局 logger 里面使用了 print 等 API，导致 hook 出现递归调用
+    logger.propagate = False
+
+
+async def run_executor(address: str, session_id: str, package: str | None, session_dir: str, suffix: str | None = None):
+
+    if suffix is not None:
+        mainframe = Mainframe(address, f"python-executor-{suffix}", logger)
+    else:
+        mainframe = Mainframe(address, f"python-executor-{session_id}", logger)
+
+    mainframe.connect()
+
+    print(f"connecting to broker {address} success")
+    sys.stdout.flush()
+
+    logger.info("executor start") if package is None else logger.info(f"executor start for package {package}")
+
+    add_matplot_module()
+    import_helper(logger)
+
+    # add package to sys.path
+    if package is not None:
+        sys.path.append(package)
+    elif os.path.exists("/app/workspace"):
+        sys.path.append("/app/workspace")
+
+
+    def not_current_session(message):
+        return message.get("session_id") != session_id
+    
+    def not_current_package(message):
+        return message.get("package") != package
+
+    # 目前的 mqtt 库，在 subscribe 回调里 publish 消息会导致死锁无法工作，参考 https://github.com/eclipse/paho.mqtt.python/issues/527 或者 https://stackoverflow.com/a/36964192/4770006
+    # 通过这种方式来绕过，所有需要 callback 后 publish message 的情况，都需要使用 future 类似方式来绕过。
+    fs = queue.Queue()
+    loop = asyncio.get_event_loop()
+
+    def execute_block(message):
+        if not_current_session(message):
+            return
+
+        if not_current_package(message):
+            return
+        
+        # https://github.com/oomol/oocana-rust/issues/310 临时解决方案
+        job_id = message.get("job_id")
+        if job_id in job_set:
+            logger.warning(f"job {job_id} already running, ignore")
+            return
+        job_set.add(job_id)
+
+        nonlocal fs
+        f = loop.create_future()
+        fs.put(f)
+        f.set_result(message)
+    
+    def execute_service_block(message):
+        if not_current_session(message):
+            return
+        
+        if not_current_package(message):
+            return
+
+        nonlocal fs
+        f = loop.create_future()
+        fs.put(f)
+        f.set_result(message)
+
+    def service_exit(message: ReportStatusPayload):
+        service_hash = message.get("service_hash")
+        if service_hash in service_store:
+            del service_store[service_hash]
+
+    def service_status(message: ReportStatusPayload):
+        service_hash = message.get("service_hash")
+        if service_hash in service_store:
+            service_store[service_hash] = "running"
+
+    def report_message(message):
+        type = message.get("type")
+        if type == "SessionFinished":
+            if not_current_session(message):
+                return
+            logger.info(f"session {session_id} finished, exit executor")
+            mainframe.disconnect() # TODO: 即使调用 disconnect，在 broker 上也无法看不到主动断开的信息，有时间再调查。
+            if os.getenv("IS_FORKED"): # fork 进程无法直接使用 sys.exit 退出
+                os._exit(0)
+            else:
+                hook.original_exit(0)
+        
+
+    mainframe.subscribe(f"executor/{EXECUTOR_NAME}/run_block", execute_block)
+    mainframe.subscribe(f"executor/{EXECUTOR_NAME}/run_service_block", execute_service_block)
+    mainframe.subscribe('report', report_message)
+    mainframe.subscribe(exit_report_topic(), service_exit)
+    mainframe.subscribe(status_report_topic(), service_status)
+
+    mainframe.notify_executor_ready(session_id, EXECUTOR_NAME, package)
+
+    async def spawn_service(message: ServiceExecutePayload, service_hash: str):
+        logger.info(f"create new service {message.get('dir')}")
+        service_store[service_hash] = "launching"
+
+        parent_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
+
+        is_global_service = message.get("service_executor").get("stop_at") in ["app_end", "never"]
+
+        if is_global_service:
+            process = await asyncio.create_subprocess_shell(
+                f"python -u -m python_executor.service --address {address}  --service-hash {service_hash} --session-dir {session_dir}",
+                cwd=parent_dir
+            )
+        else:
+            process = await asyncio.create_subprocess_shell(
+                f"python -u -m python_executor.service --address {address} --session-id {session_id}  --service-hash {service_hash} --session-dir {session_dir}",
+                cwd=parent_dir
+            )
+        params: ServiceTopicParams = {
+            "service_hash": service_hash,
+            "session_id": session_id
+        }
+
+        def send_service_config(params: ServiceTopicParams, message: ServiceExecutePayload):
+
+            async def run():
+                mainframe.publish(service_config_topic(params), message)
+                service_store[service_hash] = "running"
+            run_in_new_thread(run)
+
+        # FIXME: mqtt 不能在 subscribe 后立即 publish，需要修复。
+        mainframe.subscribe(prepare_report_topic(params), lambda _: send_service_config(params, message))
+
+        await process.wait()
+        logger.info(f"service {service_hash} exit")
+        del service_store[service_hash]
+    
+
+    def run_service_block(message: ServiceExecutePayload):
+        logger.info(f"service block {message.get('job_id')} start")
+        service_hash = message.get("service_hash")
+        params: ServiceTopicParams = {
+            "service_hash": service_hash,
+            "session_id": session_id
+        }
+        mainframe.publish(run_action_topic(params), message)
+
+    while True:
+        await asyncio.sleep(1)
+        if not fs.empty():
+            f = fs.get()
+            message = await f
+            if message.get("service_executor") is not None:
+                service_hash = message.get("service_hash")
+                status = service_store.get(service_hash)
+                if status is None:
+                    asyncio.create_task(spawn_service(message, service_hash))
+                elif status == "running":
+                    run_service_block(message)
+                elif status == "launching":
+                    logger.info(f"service {service_hash} is launching, set message back to fs to wait next time")
+                    fs.put(f)
+            else:
+                if not_current_session(message):
+                    continue
+                run_block_in_new_thread(message, mainframe, session_dir=session_dir)
+
+def run_block_in_new_thread(message, mainframe: Mainframe, session_dir: str):
+
+    async def run():
+        await run_block(message, mainframe, session_dir=session_dir)
+    run_in_new_thread(run)
+
+def main():
+
+    import argparse
+    parser = argparse.ArgumentParser(description="run executor with address, session-id, tmp-dir")
+    parser.add_argument("--session-id", help="executor subscribe session id", required=True)
+    parser.add_argument("--address", help="mqtt address", default="mqtt://127.0.0.1:47688")
+    parser.add_argument("--session-dir", help="a tmp dir for whole session", required=True)
+    parser.add_argument("--output", help="output log to console or file", default="file", choices=["console", "file"])
+    parser.add_argument("--package", help="package path, if set, executor will only run same package block", default=None)
+    parser.add_argument("--suffix", help="suffix for log file", default=None)
+
+    args = parser.parse_args()
+
+    address: str = args.address
+    session_id: str = str(args.session_id)
+    output: Literal["console", "file"] = args.output
+    package: str | None = args.package
+    suffix: str | None = args.suffix
+    session_dir: str = args.session_dir
+
+    config_logger(session_id, suffix, output)
+
+    run_async_code(run_executor(address=address, session_id=session_id, package=package, session_dir=session_dir, suffix=suffix))
+
+if __name__ == '__main__':
+    main()

oocana_python_executor-0.15.0/python_executor/hook.py

@@ -0,0 +1,50 @@
+from sys import exit
+from builtins import exit as global_exit
+from typing import TypeAlias, Any
+import sys
+import builtins
+from .data import vars, EXECUTOR_NAME
+import logging
+
+logger = logging.getLogger(EXECUTOR_NAME)
+
+class ExitFunctionException(Exception):
+    pass
+
+original_exit = exit
+original_global_exit = global_exit
+original_print = print
+
+_ExitCode: TypeAlias = str | int | None
+
+def sys_exit(status: _ExitCode = None) -> None:
+    raise ExitFunctionException(status)
+
+def sys_global_exit(status: _ExitCode = None) -> None:
+    raise ExitFunctionException(status)
+
+def global_print(*values: object, sep: str | None = " ", end: str | None = "\n", file: Any | None = None, flush: bool = False) -> None:
+    
+    context = None  # 初始化 context 变量
+    try:
+        context = vars.get()
+    except LookupError:
+        # 这个 logger 不会上报到 root handle 中，所以即使 root logger 的 Handler 里面有 print 函数，也不会导致递归调用
+        logger.warning("print called outside of block")
+    except Exception as e:
+        logger.error(f"print error: {e}")
+
+    if context is not None:
+        try:
+            msg_sep = sep or " "
+            msg = msg_sep.join(map(str, values))
+            context.report_log(msg)
+        except Exception as e:
+            logger.error(f"transform print message to context log error: {e}")
+
+    original_print(*values, sep=sep, end=end, file=file, flush=flush)
+
+
+sys.exit = sys_exit
+builtins.exit = sys_global_exit
+builtins.print = global_print

oocana_python_executor-0.15.0/python_executor/logger.py

@@ -0,0 +1,19 @@
+import logging
+from oocana import Context
+import weakref
+
+class ContextHandler(logging.Handler):
+
+    @property
+    def context(self):
+        return self._context
+
+    def __init__(self, context: Context):
+        super().__init__()
+        self._context = weakref.ref(context)
+
+    def emit(self, record):
+        msg = self.format(record)
+        ctx = self.context()
+        if ctx:
+            ctx.report_log(msg)

oocana_python_executor-0.15.0/python_executor/matplot/matplotlib_oomol/__init__.py

@@ -0,0 +1 @@
+from .oomol import show, FigureCanvas

oocana_python_executor-0.15.0/python_executor/matplot/matplotlib_oomol/oomol.py

@@ -0,0 +1,28 @@
+"""matplotlib.use('module://matplotlib_oomol'), remember to add this file to PYTHONPATH"""
+
+from matplotlib.backend_bases import Gcf # type: ignore
+from matplotlib.backends.backend_agg import FigureCanvasAgg # type: ignore
+from python_executor.data import vars
+
+FigureCanvas = FigureCanvasAgg
+
+def show(*args, **kwargs):
+    import sys
+    from io import BytesIO
+    from base64 import b64encode
+    if vars is not None:
+        context = vars.get()
+        images = []
+        for figmanager in Gcf.get_all_fig_managers():
+                buffer = BytesIO()
+                figmanager.canvas.figure.savefig(buffer, format='png')
+                buffer.seek(0)
+                png = buffer.getvalue()
+                buffer.close()
+                base64Data = b64encode(png).decode('utf-8')
+                url = f'data:image/png;base64,{base64Data}'
+                images.append(url)
+        if images:
+            context.preview({ "type": "image", "data": images })
+    else:
+        print('matplotlib_oomol: no sys.modules["oomol"]', file=sys.stderr)