rainbow-rb-flow-manager 0.0.9.dev5__py3-none-any.whl

rainbow/rb_flow_manager/context.py

@@ -0,0 +1,610 @@
+import builtins
+import contextlib
+import queue
+import sys
+import time
+from collections.abc import (
+    Callable,
+    MutableMapping,
+)
+from multiprocessing import (
+    Queue,
+)
+from multiprocessing.synchronize import (
+    Event as EventType,
+)
+from typing import Any, Literal, cast
+
+from rainbow.rb_sdk.amr import RBAmrSDK
+from rainbow.rb_sdk.base import RBBaseSDK
+from rainbow.rb_sdk.manipulate import RBManipulateSDK
+from rainbow.rb_sdk.program_sdk.program import RBProgramSDK
+from rainbow.rb_utils.parser import t_to_dict
+
+from .exception import BreakFolder, StopExecution, SubTaskHaltException
+from .schema import ExecutorState, RB_Flow_Manager_ProgramState
+
+
+class ExecutionContext:
+    """스크립트 실행 컨텍스트"""
+
+    def __init__(
+        self,
+        process_id: str,
+        state_dict: dict[str, Any],
+        result_queue: Queue,
+        pause_event: EventType,
+        resume_event: EventType,
+        stop_event: EventType,
+        *,
+        min_step_interval: float | None = None,
+        step_mode: bool = False,
+        shared_state_dicts: dict[str, MutableMapping[str, Any]] | None = None,
+    ):
+        self.process_id = process_id
+        self.state_dict: ExecutorState = cast(ExecutorState, state_dict)
+        self.parent_process_id = state_dict.get("parent_process_id")
+        self.pause_event = pause_event
+        self.resume_event = resume_event
+        self.stop_event = stop_event
+        self.result_queue = result_queue
+        self.variables: dict[str, dict[str, Any]] = {
+            "local": {},
+            "global": {},
+        }
+        self._dirty_local_variables: set[str] = set()
+        self._dirty_global_variables: set[str] = set()
+        self.sdk_functions: dict[str, Callable] = {}
+        self._sdk_roots: dict[str, Any] | None = None
+        self._arg_scope: list[dict[str, Any]] = []
+        self._generation = state_dict.get("generation")
+        self.min_step_interval = min_step_interval
+        self._folder_depth = 0
+        self.step_mode = step_mode
+        self.step_num = 1
+        self.current_depth = 0
+        self._step_ticket = 1 if not step_mode else 0
+        self.is_ui_execution = state_dict.get("is_ui_execution", False)
+        self.data: dict[str, Any] = {}  # 사용자 정의 데이터 저장소
+        self.shared_state_dicts = shared_state_dicts
+
+        # parent_process_id가 있으면 부모 변수 저장소를 owner로 사용
+        self.sync_state_variables()
+        # 서브프로세스에서 첫 스텝 호출 시점의 lazy init 레이스/지연을 줄이기 위해
+        # ExecutionContext 생성 시 SDK 루트를 선초기화한다.
+        self._ensure_sdk_roots()
+
+    def _resolve_variable_owner_state(self) -> ExecutorState:
+        if self.parent_process_id is not None and self.shared_state_dicts is not None:
+            parent_state = self.shared_state_dicts.get(self.parent_process_id)
+            if parent_state is not None:
+                return cast(ExecutorState, parent_state)
+        return self.state_dict
+
+    @staticmethod
+    def _normalize_state_variables(raw: Any) -> dict[str, dict[str, Any]]:
+        if not isinstance(raw, dict):
+            return {}
+        normalized: dict[str, dict[str, Any]] = {}
+        for key, payload in raw.items():
+            if not isinstance(payload, dict) or "value" not in payload:
+                continue
+            var_type = "global" if payload.get("type") == "global" else "local"
+            normalized[key] = {"value": payload.get("value"), "type": var_type}
+        return normalized
+
+    def update_vars_to_state_dict(self):
+        owner_state = self._resolve_variable_owner_state()
+        current = dict(owner_state.get("variables", {}))
+
+        for k in self._dirty_global_variables:
+            v = self.variables.get("global", {}).get(k)
+            current[k] = {"value": v, "type": "global"}
+
+        for k in self._dirty_local_variables:
+            v = self.variables.get("local", {}).get(k)
+            current[k] = {"value": v, "type": "local"}
+
+        owner_state["variables"] = current
+        # 로컬 state에도 미러링해 UI/호환성 유지
+        self.state_dict["variables"] = current
+        self._dirty_global_variables.clear()
+        self._dirty_local_variables.clear()
+
+    def sync_state_variables(self):
+        """owner(state 혹은 parent state)의 variables를 런타임 캐시에 반영"""
+        owner_state = self._resolve_variable_owner_state()
+        normalized = self._normalize_state_variables(owner_state.get("variables", {}))
+
+        self.variables["local"].clear()
+        self.variables["global"].clear()
+
+        for key, payload in normalized.items():
+            if payload["type"] == "global":
+                self.variables["global"][key] = payload["value"]
+            else:
+                self.variables["local"][key] = payload["value"]
+
+        self.state_dict["variables"] = dict(owner_state.get("variables", {}))
+
+    def update_local_variables(self, variables: dict[str, Any]):
+        normalized = t_to_dict(variables)
+        self.variables["local"].update(normalized)
+        self._dirty_local_variables.update(normalized.keys())
+        self.update_vars_to_state_dict()
+
+    def update_global_variables(self, variables: dict[str, Any]):
+        self.variables["global"].update(variables)
+        self._dirty_global_variables.update(variables.keys())
+        self.update_vars_to_state_dict()
+
+    def mark_dirty_variables(
+        self,
+        *,
+        local: set[str] | None = None,
+        global_: set[str] | None = None,
+    ):
+        self._dirty_local_variables.update(local or set())
+        self._dirty_global_variables.update(global_ or set())
+
+    def get_global_variable(self, var_name: str) -> Any:
+        robot_model = self.state_dict.get("robot_model", None)
+        category = self.state_dict.get("category", None)
+
+        if category is None or robot_model is None:
+            return None
+
+        try:
+            if category == "manipulate":
+                fn = self.get_sdk_function("rb_manipulate_sdk.get_data.get_variable")
+
+                res = fn(robot_model, var_name) if fn is not None else None
+
+                if res is None:
+                    raise RuntimeError(f"Failed to get global variable: {var_name}")
+
+                self.update_global_variables({var_name: res})
+
+                return res
+        except Exception as e:
+            raise e
+
+    def _ensure_sdk_roots(self):
+        if self._sdk_roots is not None:
+            return
+
+        try:
+            category = self.state_dict.get("category", None)
+
+            rb_base_sdk = RBBaseSDK()
+            rb_manipulate_sdk: RBManipulateSDK | None = None
+            rb_amr_sdk: RBAmrSDK | None = None
+            rb_program_sdk = RBProgramSDK()
+
+            if category == "manipulate":
+                rb_manipulate_sdk = RBManipulateSDK()
+            elif category == "amr":
+                rb_amr_sdk = RBAmrSDK()
+
+            self._sdk_roots = {
+                "rb_base_sdk": rb_base_sdk,
+                "rb_program_sdk": rb_program_sdk,
+                "rb_manipulate_sdk": rb_manipulate_sdk,
+                "rb_amr_sdk": rb_amr_sdk,
+            }
+        except Exception as e:
+            raise RuntimeError(
+                f"SDK initialization failed in process {self.process_id}: {e}"
+            ) from e
+
+    def get_sdk_function(self, func_name: str) -> Callable | None:
+        if not func_name:
+            return None
+
+        cached = self.sdk_functions.get(func_name)
+        if cached is not None:
+            return cached
+
+        self._ensure_sdk_roots()
+        if self._sdk_roots is None:
+            return None
+
+        parts = func_name.split(".")
+        if len(parts) < 2:
+            return None
+
+        obj = self._sdk_roots.get(parts[0])
+        if obj is None:
+            return None
+
+        for part in parts[1:]:
+            if part.startswith("_") or not hasattr(obj, part):
+                return None
+            obj = getattr(obj, part)
+
+        if not callable(obj):
+            return None
+
+        self.sdk_functions[func_name] = obj
+        return obj
+
+    def step_barrier(self):
+        """step_mode면, 이 스텝 수행 후 일시정지하고
+        resume 신호 올 때까지 대기"""
+        if not self.step_mode or self.step_num < 2:
+            return
+
+        if self.stop_event.is_set():
+            self.state_dict["state"] = RB_Flow_Manager_ProgramState.STOPPED
+            raise StopExecution("Execution stopped by user")
+
+        # pause_event가 이미 set된 상태로 barrier에 들어와도
+        # executor/UI 상태가 RUNNING으로 남지 않도록 WAITING 이벤트를 보낸다.
+        if not self.pause_event.is_set():
+            self.pause_event.set()
+        self.emit_pause(self.state_dict["current_step_id"], is_wait=True)
+        self._wait_for_resume()
+
+    def push_args(self, mapping: dict[str, Any] | None):
+        self._arg_scope.append(mapping or {})
+
+    def pop_args(self):
+        if self._arg_scope:
+            self._arg_scope.pop()
+
+    def lookup(self, key: str) -> Any:
+        idx = len(self._arg_scope) - 2
+
+        while idx >= 0:
+            scope = self._arg_scope[idx]
+            if key in scope:
+                return scope[key]
+            idx -= 1
+        raise KeyError(f"parent pointer not found: {key}")
+
+    def pause(self, is_wait: bool = False):
+        """현재 스크립트를 일시정지"""
+        if not self.pause_event.is_set():
+            self.pause_event.set()
+            self.emit_pause(self.state_dict["current_step_id"], is_wait)
+
+        self._wait_for_resume()
+
+    def resume(self):
+        """현재 스크립트를 재개"""
+        if self.resume_event.is_set():
+            return
+
+        self.resume_event.set()
+        self.pause_event.clear()
+        self.state_dict["state"] = RB_Flow_Manager_ProgramState.RUNNING
+
+    def stop(self):
+        """현재 스크립트를 중지"""
+        if not self.stop_event.is_set():
+            self.stop_event.set()
+            self.emit_stop(self.state_dict["current_step_id"])
+
+        self.pause_event.clear()
+
+        self._wait_for_resume()
+
+    def _wait_for_resume(self):
+        if self.stop_event.is_set() and not self.state_dict.get("ignore_stop", False):
+            self.state_dict["state"] = RB_Flow_Manager_ProgramState.STOPPED
+            raise StopExecution("Execution stopped by user")
+
+        if self.state_dict.get("ignore_pause", False):
+            self.pause_event.clear()
+            self.resume_event.clear()
+            if self.state_dict.get("state") == RB_Flow_Manager_ProgramState.PAUSED:
+                self.state_dict["state"] = RB_Flow_Manager_ProgramState.RUNNING
+            return
+
+        if self.pause_event.is_set():
+            # resume이 이미 와 있으면 바로 통과 (레이스 방지)
+            if self.resume_event.is_set():
+                self.resume_event.clear()
+                self.pause_event.clear()
+                self.state_dict["state"] = RB_Flow_Manager_ProgramState.RUNNING
+                return
+
+            self.state_dict["state"] = RB_Flow_Manager_ProgramState.PAUSED
+
+            while True:
+                if self.stop_event.is_set() and not self.state_dict.get("ignore_stop", False):
+                    self.pause_event.clear()
+                    self.resume_event.clear()
+                    self.state_dict["state"] = RB_Flow_Manager_ProgramState.STOPPED
+                    raise StopExecution("Execution stopped by user")
+
+                if self.state_dict.get("ignore_pause", False):
+                    self.pause_event.clear()
+                    self.resume_event.clear()
+                    self.state_dict["state"] = RB_Flow_Manager_ProgramState.RUNNING
+                    break
+
+                if self.resume_event.wait(timeout=0.01):
+                    self.resume_event.clear()
+                    self.pause_event.clear()
+                    if self.stop_event.is_set() and not self.state_dict.get("ignore_stop", False):
+                        self.state_dict["state"] = RB_Flow_Manager_ProgramState.STOPPED
+                        raise StopExecution("Execution stopped by user")
+                    self.state_dict["state"] = RB_Flow_Manager_ProgramState.RUNNING
+                    break
+
+    def check_stop(self):
+        """stop/pause/resume 명령 처리"""
+        self._wait_for_resume()  # 즉시 멈춤/중지 반영
+
+    def emit_next(self, step_id: str):
+        """next 이벤트 발생"""
+        try:
+            self.result_queue.put_nowait(
+                {
+                    "type": "next",
+                    "process_id": self.process_id,
+                    "step_id": step_id,
+                    "ts": time.time(),
+                    "generation": self._generation,
+                    "error": None,
+                }
+            )
+        except queue.Full:
+            # 고빈도 next 이벤트는 큐 포화 시 드롭해 stop 반응성을 보장한다.
+            return
+
+    def emit_pause(self, step_id: str, is_wait: bool = False):
+        """pause 이벤트 발생"""
+        self.result_queue.put(
+            {
+                "type": "pause",
+                "process_id": self.process_id,
+                "step_id": step_id,
+                "is_wait": is_wait,
+                "parent_process_id": self.parent_process_id,
+                "ts": time.time(),
+                "generation": self._generation,
+                "error": None,
+            }
+        )
+
+    def emit_wait(self, step_id: str):
+        self.result_queue.put(
+            {
+                "type": "wait",
+                "process_id": self.process_id,
+                "step_id": step_id,
+                "parent_process_id": self.parent_process_id,
+                "ts": time.time(),
+                "generation": self._generation,
+                "error": None,
+            }
+        )
+
+    def emit_resume(self, step_id: str):
+        """resume 이벤트 발생"""
+        self.result_queue.put(
+            {
+                "type": "resume",
+                "process_id": self.process_id,
+                "step_id": step_id,
+                "parent_process_id": self.parent_process_id,
+                "ts": time.time(),
+                "generation": self._generation,
+                "error": None,
+            }
+        )
+
+    def emit_stop(self, step_id: str):
+        """stop 이벤트 발생"""
+        self.result_queue.put(
+            {
+                "type": "stop",
+                "process_id": self.process_id,
+                "step_id": step_id,
+                "parent_process_id": self.parent_process_id,
+                "ts": time.time(),
+                "generation": self._generation,
+                "error": None,
+            }
+        )
+
+    def emit_error(self, step_id: str, error: Exception):
+        """error 이벤트 발생"""
+
+        self.result_queue.put(
+            {
+                "type": "error",
+                "process_id": self.process_id,
+                "step_id": step_id,
+                "ts": time.time(),
+                "generation": self._generation,
+                "error": str(error),
+            }
+        )
+
+    def emit_subtask_sync_register(
+        self, sub_task_tree, post_tree=None, subtask_type: Literal["INSERT", "CHANGE"] = "INSERT"
+    ):
+        """서브태스크 sync flags 등록 요청 이벤트"""
+        self.result_queue.put(
+            {
+                "type": "subtask_sync_register",
+                "process_id": self.process_id,
+                "subtask_type": subtask_type,
+                "sub_task_tree": sub_task_tree.to_dict() if sub_task_tree else None,
+                "post_tree": post_tree.to_dict() if post_tree else None,
+                "ts": time.time(),
+                "generation": self._generation,
+            }
+        )
+
+    def emit_subtask_sync_unregister(
+        self, sub_task_tree, post_tree=None, subtask_type: Literal["INSERT", "CHANGE"] = "INSERT"
+    ):
+        """서브태스크 sync flags 해제 요청 이벤트"""
+        self.result_queue.put(
+            {
+                "type": "subtask_sync_unregister",
+                "process_id": self.process_id,
+                "subtask_type": subtask_type,
+                "sub_task_tree": sub_task_tree.to_dict() if sub_task_tree else None,
+                "post_tree": post_tree.to_dict() if post_tree else None,
+                "ts": time.time(),
+                "generation": self._generation,
+            }
+        )
+
+    def emit_main_tree_sync_unregister(self):
+        """메인 트리 sync flags 해제 요청 (CHANGE 타입 전환 시)"""
+        self.result_queue.put(
+            {
+                "type": "main_tree_sync_unregister",
+                "process_id": self.process_id,
+                "ts": time.time(),
+                "generation": self._generation,
+            }
+        )
+
+    def emit_sub_task_start(self, task_id: str, sub_task_type: Literal["INSERT", "CHANGE"]):
+        """sub_task_start 이벤트 발생"""
+        self.result_queue.put(
+            {
+                "type": "sub_task_start",
+                "process_id": self.process_id,
+                "sub_task_id": task_id,
+                "sub_task_type": sub_task_type,
+                "ts": time.time(),
+                "generation": self._generation,
+                "error": None,
+            }
+        )
+
+    def emit_sub_task_done(self, task_id: str, sub_task_type: Literal["INSERT", "CHANGE"]):
+        """sub_task_done 이벤트 발생"""
+        self.result_queue.put(
+            {
+                "type": "sub_task_done",
+                "process_id": self.process_id,
+                "sub_task_id": task_id,
+                "sub_task_type": sub_task_type,
+                "ts": time.time(),
+                "generation": self._generation,
+                "error": None,
+            }
+        )
+
+    def emit_event_sub_task_start(
+        self,
+        *,
+        event_task_id: str,
+        event_tree,
+        step_id: str | None = None,
+        run_mode: Literal["SYNC", "ASYNC"] = "ASYNC",
+        call_seq: int | None = None,
+    ):
+        """EventSubTask 실행 요청 이벤트"""
+        self.result_queue.put(
+            {
+                "type": "event_sub_task_start",
+                "process_id": self.process_id,
+                "event_task_id": event_task_id,
+                "event_tree": event_tree.to_dict() if event_tree is not None else None,
+                "step_id": step_id,
+                "run_mode": run_mode,
+                "call_seq": call_seq,
+                "ts": time.time(),
+                "generation": self._generation,
+                "error": None,
+            }
+        )
+
+    def emit_post_start(self):
+        """post_start 이벤트 발생"""
+        self.result_queue.put(
+            {
+                "type": "post_start",
+                "process_id": self.process_id,
+                "ts": time.time(),
+                "generation": self._generation,
+                "error": None,
+            }
+        )
+
+    def emit_done(self, step_id: str):
+        """done 이벤트 발생"""
+        try:
+            self.result_queue.put_nowait(
+                {
+                    "type": "done",
+                    "process_id": self.process_id,
+                    "step_id": step_id,
+                    "ts": time.time(),
+                    "generation": self._generation,
+                    "error": None,
+                }
+            )
+        except queue.Full:
+            # 고빈도 done 이벤트는 큐 포화 시 드롭해 stop 반응성을 보장한다.
+            return
+
+    def pause_all(self):
+        self.result_queue.put(
+            {
+                "type": "control",
+                "action": "pause_all",
+                "process_id": self.process_id,
+                "ts": time.time(),
+                "generation": self._generation,
+            }
+        )
+
+    def resume_all(self):
+        self.result_queue.put(
+            {
+                "type": "control",
+                "action": "resume_all",
+                "process_id": self.process_id,
+                "ts": time.time(),
+                "generation": self._generation,
+            }
+        )
+
+    def enter_folder(self):
+        """Folder 진입"""
+        self._folder_depth += 1
+
+    def leave_folder(self):
+        """Folder 탈출"""
+        self._folder_depth -= 1
+
+    def break_folder(self):
+        """Folder 실행 중단"""
+        if self._folder_depth > 0:
+            raise BreakFolder()
+
+    def halt_sub_task(self):
+        """서브 태스크 실행 중단"""
+        raise SubTaskHaltException()
+
+    def _safe_close_sdk(self):
+        """SDK 안전 종료"""
+        with contextlib.suppress(Exception):
+            if self._sdk_roots is not None:
+                for sdk in self._sdk_roots.values():
+                    close_fn = getattr(sdk, "close", None)
+                    if callable(close_fn):
+                        close_fn()
+
+    def close(self):
+        """안전하게 close"""
+        self._safe_close_sdk()
+
+    def __del__(self):
+        """소멸자"""
+        if getattr(sys, "is_finalizing", lambda: False)():
+            return
+
+        with contextlib.suppress(builtins.BaseException):
+            self.close()

rainbow/rb_flow_manager/controller/base_controller.py

@@ -0,0 +1,88 @@
+from abc import (
+    ABC,
+    abstractmethod,
+)
+from collections.abc import (
+    MutableMapping,
+)
+from typing import Any, Literal
+
+
+class BaseController(ABC):
+    """모든 컨트롤러가 따라야 하는 공통 인터페이스"""
+
+    @abstractmethod
+    def on_init(self, state_dicts: dict[str, MutableMapping[str, Any]]) -> None:
+        """컨트롤러 초기화"""
+
+    @abstractmethod
+    def on_start(self, task_id: str) -> None:
+        """스크립트/플로우 시작 시점에 한 번 호출"""
+
+    @abstractmethod
+    def on_stop(self, task_id: str, step_id: str) -> None:
+        """각 Step이 Stop 될때 동작을 수행"""
+
+    @abstractmethod
+    def on_pause(self, task_id: str, step_id: str) -> None:
+        """각 Step이 Pause 될때 실제 동작을 수행"""
+
+    @abstractmethod
+    def on_wait(self, task_id: str, step_id: str) -> None:
+        """각 Step이 Wait 될때 실제 동작을 수행"""
+
+    @abstractmethod
+    def on_resume(self, task_id: str, step_id: str) -> None:
+        """각 Step이 Resume 될때 실제 동작을 수행"""
+
+    @abstractmethod
+    def on_next(self, task_id: str, step_id: str) -> None:
+        """각 Step이 다음 Step으로 이동 할때 동작을 수행"""
+
+    @abstractmethod
+    def on_sub_task_start(self, task_id: str, sub_task_id: str, sub_task_type: Literal["INSERT", "CHANGE"]) -> None:
+        """각 Step이 서브 태스크 시작 시점에 동작을 수행"""
+
+    @abstractmethod
+    def on_sub_task_done(self, task_id: str, sub_task_id: str, sub_task_type: Literal["INSERT", "CHANGE"]) -> None:
+        """각 Step이 서브 태스크 완료 시점에 동작을 수행"""
+
+    @abstractmethod
+    def on_done(self, task_id: str, step_id: str) -> None:
+        """각 Step이 완료 되었을때 동작을 수행"""
+
+    @abstractmethod
+    def on_error(self, task_id: str, step_id: str, error: Exception) -> None:
+        """각 Step에서 Error 발생 시 동작을 수행"""
+
+    @abstractmethod
+    def on_post_start(self, task_id: str) -> None:
+        """스크립트/플로우 Post Tree의 시작 시점에 동작을 수행"""
+
+    @abstractmethod
+    def on_complete(self, task_id: str) -> None:
+        """스크립트/플로우 완료 시점에 동작을 수행"""
+
+    @abstractmethod
+    def on_close(self) -> None:
+        """컨트롤러 종료"""
+
+    @abstractmethod
+    def on_all_complete(self) -> None:
+        """모든 스크립트/플로우 완료 시점에 동작을 수행"""
+
+    @abstractmethod
+    def on_all_stop(self) -> None:
+        """모든 스크립트/플로우 Stop 시점에 동작을 수행"""
+
+    @abstractmethod
+    def on_all_pause(self) -> None:
+        """모든 스크립트/플로우 Pause 시점에 동작을 수행"""
+
+    @abstractmethod
+    def on_all_wait(self) -> None:
+        """모든 스크립트/플로우 Wait 시점에 동작을 수행"""
+
+    @abstractmethod
+    def on_all_resume(self) -> None:
+        """모든 스크립트/플로우 Resume 시점에 동작을 수행"""