oocana 0.15.0__tar.gz

oocana-0.15.0/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.1
+Name: oocana
+Version: 0.15.0
+Summary: python implement of oocana to give a context for oocana block
+License: MIT
+Requires-Python: >=3.9
+Requires-Dist: paho-mqtt>=2
+Requires-Dist: simplejson>=3.19.2
+Requires-Dist: typing-extensions>=4.12.2; python_version < "3.11"
+

oocana-0.15.0/oocana/__init__.py

@@ -0,0 +1,7 @@
+from .data import *  # noqa: F403
+from .context import *  # noqa: F403
+from .service import *  # noqa: F403
+from .handle_data import *  # noqa: F403
+from .preview import *  # noqa: F403
+from .schema import *  # noqa: F403
+from .mainframe import Mainframe as Mainframe  # noqa: F403

oocana-0.15.0/oocana/context.py

@@ -0,0 +1,334 @@
+from dataclasses import asdict
+from json import loads
+from .data import BlockInfo, StoreKey, JobDict, BlockDict, BinValueDict, VarValueDict
+from .handle_data import HandleDef
+from .mainframe import Mainframe
+from typing import Dict, Any, TypedDict, Optional
+from base64 import b64encode
+from io import BytesIO
+from .throtter import throttle
+from .preview import PreviewPayload, TablePreviewData, DataFrame, ShapeDataFrame, PartialDataFrame
+from .data import EXECUTOR_NAME
+import os.path
+import logging
+
+__all__ = ["Context"]
+
+class OnlyEqualSelf:
+    def __eq__(self, value: object) -> bool:
+        return self is value
+
+class OOMOL_LLM_ENV(TypedDict):
+    base_url: str
+    api_key: str
+    models: list[str]
+
+class HostInfo(TypedDict):
+    gpu_vendor: str
+    gpu_renderer: str
+
+class Context:
+    __inputs: Dict[str, Any]
+
+    __block_info: BlockInfo
+    __outputs_def: Dict[str, HandleDef]
+    __store: Any
+    __is_done: bool = False
+    __keep_alive: OnlyEqualSelf = OnlyEqualSelf()
+    __session_dir: str
+    _logger: Optional[logging.Logger] = None
+
+    def __init__(
+        self, inputs: Dict[str, Any], blockInfo: BlockInfo, mainframe: Mainframe, store, outputs, session_dir: str
+    ) -> None:
+
+        self.__block_info = blockInfo
+
+        self.__mainframe = mainframe
+        self.__store = store
+        self.__inputs = inputs
+
+        outputs_defs = {}
+        if outputs is not None:
+            for k, v in outputs.items():
+                outputs_defs[k] = HandleDef(**v)
+        self.__outputs_def = outputs_defs
+        self.__session_dir = session_dir
+
+    @property
+    def logger(self) -> logging.Logger:
+        """a custom logger for the block, you can use it to log the message to the block log. this logger will report the log by context report_logger api.
+        """
+
+        # setup after init, so the logger always exists
+        if self._logger is None:
+            raise ValueError("logger is not setup, please setup the logger in the block init function.")
+        return self._logger
+
+    @property
+    def session_dir(self) -> str:
+        """a temporary directory for the current session, all blocks in the one session will share the same directory.
+        """
+        return self.__session_dir
+
+    @property
+    def keepAlive(self):
+        return self.__keep_alive
+
+    @property
+    def inputs(self):
+        return self.__inputs
+
+    @property
+    def session_id(self):
+        return self.__block_info.session_id
+
+    @property
+    def job_id(self):
+        return self.__block_info.job_id
+    
+    @property
+    def job_info(self) -> JobDict:
+        return self.__block_info.job_info()
+    
+    @property
+    def block_info(self) -> BlockDict:
+        return self.__block_info.block_dict()
+    
+    @property
+    def node_id(self) -> str:
+        return self.__block_info.stacks[-1].get("node_id", None)
+    
+    @property
+    def oomol_llm_env(self) -> OOMOL_LLM_ENV:
+        """this is a dict contains the oomol llm environment variables
+        """
+        return {
+            "base_url": os.getenv("OOMOL_LLM_BASE_URL", ""),
+            "api_key": os.getenv("OOMOL_LLM_API_KEY", ""),
+            "models": os.getenv("OOMOL_LLM_MODELS", "").split(","),
+        }
+    
+    @property
+    def host_info(self) -> HostInfo:
+        """this is a dict contains the host information
+        """
+        return {
+            "gpu_vendor": os.getenv("OOMOL_HOST_GPU_VENDOR", "unknown"),
+            "gpu_renderer": os.getenv("OOMOL_HOST_GPU_RENDERER", "unknown"),
+        }
+
+    def __store_ref(self, handle: str):
+        return StoreKey(
+            executor=EXECUTOR_NAME,
+            handle=handle,
+            job_id=self.job_id,
+            session_id=self.session_id,
+        )
+    
+    def __is_basic_type(self, value: Any) -> bool:
+        return isinstance(value, (int, float, str, bool))
+
+    def output(self, key: str, value: Any, done: bool = False):
+        """
+        output the value to the next block
+
+        key: str, the key of the output, should be defined in the block schema output defs, the field name is handle
+        value: Any, the value of the output
+        """
+
+        v = value
+
+        if self.__outputs_def is not None:
+            output_def = self.__outputs_def.get(key)
+            if (
+                output_def is not None and output_def.is_var_handle() and not self.__is_basic_type(value) # 基础类型即使是变量也不放进 store，直接作为 json 内容传递
+            ):
+                ref = self.__store_ref(key)
+                self.__store[ref] = value
+                d: VarValueDict = {
+                    "__OOMOL_TYPE__": "oomol/var",
+                    "value": asdict(ref)
+                }
+                v = d
+            elif output_def is not None and output_def.is_bin_handle():
+                if not isinstance(value, bytes):
+                    self.send_warning(
+                        f"Output handle key: [{key}] is defined as binary, but the value is not bytes."
+                    )
+                    return
+                
+                bin_file = f"{self.session_dir}/binary/{self.session_id}/{self.job_id}/{key}"
+                os.makedirs(os.path.dirname(bin_file), exist_ok=True)
+                try:
+                    with open(bin_file, "wb") as f:
+                        f.write(value)
+                except IOError as e:
+                    self.send_warning(
+                        f"Output handle key: [{key}] is defined as binary, but an error occurred while writing the file: {e}"
+                    )
+                    return
+
+                if os.path.exists(bin_file):
+                    bin_value: BinValueDict = {
+                        "__OOMOL_TYPE__": "oomol/bin",
+                        "value": bin_file,
+                    }
+                    v = bin_value
+                else:
+                    self.send_warning(
+                        f"Output handle key: [{key}] is defined as binary, but the file is not written."
+                    )
+                    return
+
+        # 如果传入 key 在输出定义中不存在，直接忽略，不发送数据。但是 done 仍然生效。
+        if self.__outputs_def is not None and self.__outputs_def.get(key) is None:
+            self.send_warning(
+                f"Output handle key: [{key}] is not defined in Block outputs schema."
+            )
+            if done:
+                self.done()
+            return
+
+        node_result = {
+            "type": "BlockOutput",
+            "handle": key,
+            "output": v,
+            "done": done,
+        }
+        self.__mainframe.send(self.job_info, node_result)
+
+        if done:
+            self.done()
+
+    def done(self, error: str | None = None):
+        if self.__is_done:
+            self.send_warning("done has been called multiple times, will be ignored.")
+            return
+        self.__is_done = True
+        if error is None:
+            self.__mainframe.send(self.job_info, {"type": "BlockFinished"})
+        else:
+            self.__mainframe.send(
+                self.job_info, {"type": "BlockFinished", "error": error}
+            )
+
+    def send_message(self, payload):
+        self.__mainframe.report(
+            self.block_info,
+            {
+                "type": "BlockMessage",
+                "payload": payload,
+            },
+        )
+    
+    def __dataframe(self, payload: PreviewPayload) -> PreviewPayload:
+        if isinstance(payload, DataFrame):
+            payload = { "type": "table", "data": payload }
+
+        if isinstance(payload, dict) and payload.get("type") == "table":
+            df = payload.get("data")
+            if isinstance(df, ShapeDataFrame):
+                row_count = df.shape[0]
+                if row_count <= 10:
+                    data = df.to_dict(orient='split')
+                    columns = data.get("columns", [])
+                    rows = data.get("data", [])
+                elif isinstance(df, PartialDataFrame):
+                    data_columns = loads(df.head(5).to_json(orient='split'))
+                    columns = data_columns.get("columns", [])
+                    rows_head = data_columns.get("data", [])
+                    data_tail = loads(df.tail(5).to_json(orient='split'))
+                    rows_tail = data_tail.get("data", [])
+                    rows_dots = [["..."] * len(columns)]
+                    rows = rows_head + rows_dots + rows_tail
+                else:
+                    print("dataframe more than 10 rows but not support head and tail is not supported")
+                    return payload
+                data: TablePreviewData = { "rows": rows, "columns": columns, "row_count": row_count }
+                payload = { "type": "table", "data": data }
+            else:
+                print("dataframe is not support shape property")
+        
+        return payload
+
+    def __matplotlib(self, payload: PreviewPayload) -> PreviewPayload:
+        # payload is a matplotlib Figure
+        if hasattr(payload, 'savefig'):
+            fig: Any = payload
+            buffer = BytesIO()
+            fig.savefig(buffer, format='png')
+            buffer.seek(0)
+            png = buffer.getvalue()
+            buffer.close()
+            url = f'data:image/png;base64,{b64encode(png).decode("utf-8")}'
+            payload = { "type": "image", "data": url }
+
+        return payload
+
+    def preview(self, payload: PreviewPayload):
+        payload = self.__dataframe(payload)
+        payload = self.__matplotlib(payload)
+
+        self.__mainframe.report(
+            self.block_info,
+            {
+                "type": "BlockPreview",
+                "payload": payload,
+            },
+        )
+
+    @throttle(0.3)
+    def report_progress(self, progress: float | int):
+        """report progress
+
+        This api is used to report the progress of the block. but it just effect the ui progress not the real progress.
+        This api is throttled. the minimum interval is 0.3s. 
+        When you first call this api, it will report the progress immediately. After it invoked once, it will report the progress at the end of the throttling period.
+
+        |       0.25 s        |   0.2 s  |
+        first call       second call    third call  4 5 6 7's calls
+        |                     |          |          | | | |
+        | -------- 0.3 s -------- | -------- 0.3 s -------- |
+        invoke                  invoke                    invoke
+        :param float | int progress: the progress of the block, the value should be in [0, 100].
+        """
+        self.__mainframe.report(
+            self.block_info,
+            {
+                "type": "BlockProgress",
+                "rate": progress,
+            }
+        )
+
+    def report_log(self, line: str, stdio: str = "stdout"):
+        self.__mainframe.report(
+            self.block_info,
+            {
+                "type": "BlockLog",
+                "log": line,
+                stdio: stdio,
+            },
+        )
+
+    def log_json(self, payload):
+        self.__mainframe.report(
+            self.block_info,
+            {
+                "type": "BlockLogJSON",
+                "json": payload,
+            },
+        )
+
+    def send_warning(self, warning: str):
+        self.__mainframe.report(self.block_info, {"type": "BlockWarning", "warning": warning})
+
+    def send_error(self, error: str):
+        '''
+        deprecated, use error(error) instead.
+        consider to remove in the future.
+        '''
+        self.error(error)
+
+    def error(self, error: str):
+        self.__mainframe.send(self.job_info, {"type": "BlockError", "error": error})

oocana-0.15.0/oocana/data.py

@@ -0,0 +1,96 @@
+from dataclasses import dataclass, asdict
+from typing import TypedDict, Literal
+from simplejson import JSONEncoder
+import simplejson as json
+
+EXECUTOR_NAME = "python"
+
+__all__ = ["dumps", "BinValueDict", "VarValueDict", "JobDict", "BlockDict", "StoreKey", "BlockInfo", "EXECUTOR_NAME", "JobDict", "BinValueDict", "VarValueDict"]
+
+def dumps(obj, **kwargs):
+    return json.dumps(obj, cls=DataclassJSONEncoder, ignore_nan=True, **kwargs)
+
+class DataclassJSONEncoder(JSONEncoder):
+    def default(self, o): # pyright: ignore[reportIncompatibleMethodOverride]
+        if hasattr(o, '__dataclass_fields__'):
+            return asdict(o)
+        return JSONEncoder.default(self, o)
+
+class BinValueDict(TypedDict):
+    value: str
+    __OOMOL_TYPE__: Literal["oomol/bin"]
+
+class VarValueDict(TypedDict):
+    value: dict
+    __OOMOL_TYPE__: Literal["oomol/var"]
+
+class JobDict(TypedDict):
+    session_id: str
+    job_id: str
+
+class BlockDict(TypedDict):
+
+    try:
+        # NotRequired, Required was added in version 3.11
+        from typing import NotRequired, Required, TypedDict  # type: ignore
+    except ImportError:
+        from typing_extensions import NotRequired, Required, TypedDict
+
+    session_id: str
+    job_id: str
+    stacks: list
+    block_path: NotRequired[str]
+
+# dataclass 默认字段必须一一匹配
+# 如果多一个或者少一个字段，就会报错。
+# 这里想兼容额外多余字段，所以需要自己重写 __init__ 方法，忽略处理多余字段。同时需要自己处理缺少字段的情况。
+@dataclass(frozen=True, kw_only=True)
+class StoreKey:
+    executor: str
+    handle: str
+    job_id: str
+    session_id: str
+
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            object.__setattr__(self, key, value)
+        for key in self.__annotations__.keys():
+            if key not in kwargs:
+                raise ValueError(f"missing key {key}")
+
+
+# 发送 reporter 时，固定需要的 block 信息参数
+@dataclass(frozen=True, kw_only=True)
+class BlockInfo:
+
+    session_id: str
+    job_id: str
+    stacks: list
+    block_path: str | None = None
+
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            object.__setattr__(self, key, value)
+        for key in self.__annotations__.keys():
+            if key not in kwargs and key != "block_path":
+                raise ValueError(f"missing key {key}")
+
+    def job_info(self) -> JobDict:
+        return {"session_id": self.session_id, "job_id": self.job_id}
+
+    def block_dict(self) -> BlockDict:
+        if self.block_path is None:
+            return {
+                "session_id": self.session_id,
+                "job_id": self.job_id,
+                "stacks": self.stacks,
+            }
+
+        return {
+            "session_id": self.session_id,
+            "job_id": self.job_id,
+            "stacks": self.stacks,
+            "block_path": self.block_path,
+        }
+    
+

oocana-0.15.0/oocana/handle_data.py

@@ -0,0 +1,55 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+from .schema import FieldSchema, ContentMediaType
+
+__all__ = ["HandleDef", "InputHandleDef"]
+
+@dataclass(frozen=True, kw_only=True)
+class HandleDef:
+    """The base handle for output def, can be directly used for output def 
+    """
+    handle: str
+    """The name of the handle. it should be unique in handle list."""
+
+    json_schema: Optional[FieldSchema] = None
+    """The schema of the handle. It is used to validate the handle's content."""
+
+    name: Optional[str] = None
+    """A alias of the handle's type name. It is used to display in the UI and connect to the other handle match"""
+
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            object.__setattr__(self, key, value)
+        if "handle" not in kwargs:
+            raise ValueError("missing attr key: 'handle'")
+        json_schema = self.json_schema
+        if json_schema is not None and not isinstance(json_schema, FieldSchema):
+            object.__setattr__(self, "json_schema", FieldSchema.generate_schema(json_schema))
+
+    def check_handle_type(self, type: ContentMediaType) -> bool:
+        if self.handle is None:
+            return False
+        if self.json_schema is None:
+            return False
+        if self.json_schema.contentMediaType is None:
+            return False
+        return self.json_schema.contentMediaType == type
+
+    def is_var_handle(self) -> bool:
+        return self.check_handle_type("oomol/var")
+    
+    def is_secret_handle(self) -> bool:
+        return self.check_handle_type("oomol/secret")
+
+    def is_bin_handle(self) -> bool:
+        return self.check_handle_type("oomol/bin")
+
+@dataclass(frozen=True, kw_only=True)
+class InputHandleDef(HandleDef):
+
+    value: Optional[Any] = None
+    """default value for input handle, can be None.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)

oocana-0.15.0/oocana/mainframe.py

@@ -0,0 +1,135 @@
+from simplejson import loads
+import paho.mqtt.client as mqtt
+from paho.mqtt.enums import CallbackAPIVersion
+import operator
+from urllib.parse import urlparse
+import uuid
+from .data import BlockDict, JobDict, dumps
+import logging
+from typing import Optional
+
+__all__ = ["Mainframe"]
+
+class Mainframe:
+    address: str
+    client: mqtt.Client
+    client_id: str
+    _subscriptions: set[str] = set()
+    _logger: logging.Logger
+
+    def __init__(self, address: str, client_id: Optional[str] = None, logger = None) -> None:
+        self.address = address
+        self.client_id = client_id or f"python-executor-{uuid.uuid4().hex[:8]}"
+        self._logger = logger or logging.getLogger(__name__)
+
+    def connect(self):
+        connect_address = (
+            self.address
+            if operator.contains(self.address, "://")
+            else f"mqtt://{self.address}"
+        )
+        url = urlparse(connect_address)
+        client = self._setup_client()
+        client.connect(host=url.hostname, port=url.port) # type: ignore
+        client.loop_start()
+    
+    def _setup_client(self):
+        self.client = mqtt.Client(
+            callback_api_version=CallbackAPIVersion.VERSION2,
+            client_id=self.client_id,
+        )
+        self.client.logger = self._logger
+        self.client.on_connect = self.on_connect
+        self.client.on_disconnect = self.on_disconnect
+        self.client.on_connect_fail = self.on_connect_fail # type: ignore
+        return self.client
+
+    # mqtt v5 重连后，订阅和队列信息会丢失(v3 在初始化时，设置 clean_session 后，会保留两者。
+    # 我们的 broker 使用的是 v5，在 on_connect 里订阅，可以保证每次重连都重新订阅上。
+    def on_connect(self, client, userdata, flags, reason_code, properties):
+        if reason_code != 0:
+            self._logger.error("connect to broker failed, reason_code: %s", reason_code)
+            return
+        else:
+            self._logger.info("connect to broker success")
+
+        for topic in self._subscriptions.copy(): # 进程冲突
+            self._logger.info("resubscribe to topic: {}".format(topic))
+            self.client.subscribe(topic, qos=1)
+
+    def on_connect_fail(self) -> None:
+        self._logger.error("connect to broker failed")
+
+    def on_disconnect(self, client, userdata, flags, reason_code, properties):
+        self._logger.warning("disconnect to broker, reason_code: %s", reason_code)
+
+    # 不等待 publish 完成，使用 qos 参数来会保证消息到达。
+    def send(self, job_info: JobDict, msg) -> mqtt.MQTTMessageInfo:
+        return self.client.publish(
+            f'session/{job_info["session_id"]}', dumps({"job_id": job_info["job_id"], "session_id": job_info["session_id"], **msg}), qos=1
+        )
+
+    def report(self, block_info: BlockDict, msg: dict) -> mqtt.MQTTMessageInfo:
+        return self.client.publish("report", dumps({**block_info, **msg}), qos=1)
+    
+    def notify_executor_ready(self, session_id: str, executor_name: str, package: str | None) -> None:
+        self.client.publish(f"session/{session_id}", dumps({
+            "type": "ExecutorReady",
+            "session_id": session_id,
+            "executor_name": executor_name,
+            "package": package,
+        }), qos=1)
+
+    def notify_block_ready(self, session_id: str, job_id: str) -> dict:
+
+        topic = f"inputs/{session_id}/{job_id}"
+        replay = None
+
+        def on_message_once(_client, _userdata, message):
+            nonlocal replay
+            self.client.unsubscribe(topic)
+            replay = loads(message.payload)
+
+        self.client.subscribe(topic, qos=1)
+        self.client.message_callback_add(topic, on_message_once)
+
+        self.client.publish(f"session/{session_id}", dumps({
+            "type": "BlockReady",
+            "session_id": session_id,
+            "job_id": job_id,
+        }), qos=1)
+
+        while True:
+            if replay is not None:
+                self._logger.info("notify ready success in {} {}".format(session_id, job_id))
+                return replay
+            
+    def publish(self, topic, payload):
+        self.client.publish(topic, dumps(payload), qos=1)
+    
+    def subscribe(self, topic: str, callback):
+        def on_message(_client, _userdata, message):
+            self._logger.info("receive topic: {} payload: {}".format(topic, message.payload))
+            payload = loads(message.payload)
+            callback(payload)
+
+        self.client.message_callback_add(topic, on_message)
+        self._subscriptions.add(topic)
+
+        if self.client.is_connected():
+            self.client.subscribe(topic, qos=1)
+            self._logger.info("subscribe to topic: {}".format(topic))
+        else:
+            self._logger.info("wait connected to subscribe to topic: {}".format(topic))
+
+
+    def unsubscribe(self, topic):
+        self.client.message_callback_remove(topic)
+        self.client.unsubscribe(topic)
+        self._subscriptions.remove(topic)
+
+    def loop(self):
+        self.client.loop_forever()
+
+    def disconnect(self):
+        self.client.disconnect()