rappel 0.4.1__py3-none-win_amd64.whl

rappel/worker.py

@@ -0,0 +1,191 @@
+"""gRPC worker client that executes rappel actions."""
+
+import argparse
+import asyncio
+import importlib
+import logging
+import sys
+import time
+from typing import Any, AsyncIterator, cast
+
+import grpc
+
+from proto import messages_pb2 as pb2
+from proto import messages_pb2_grpc as pb2_grpc
+from rappel.actions import serialize_error_payload, serialize_result_payload
+
+from . import workflow_runtime
+from .logger import configure as configure_logger
+
+LOGGER = configure_logger("rappel.worker")
+aio = cast(Any, grpc).aio
+
+
+def _parse_args(argv: list[str] | None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Rappel workflow worker")
+    parser.add_argument("--bridge", required=True, help="gRPC address of the Rust bridge")
+    parser.add_argument("--worker-id", required=True, type=int, help="Logical worker identifier")
+    parser.add_argument(
+        "--user-module",
+        action="append",
+        default=[],
+        help="Optional user module(s) to import eagerly",
+    )
+    return parser.parse_args(argv)
+
+
+async def _outgoing_stream(
+    queue: "asyncio.Queue[pb2.Envelope]", worker_id: int
+) -> AsyncIterator[pb2.Envelope]:
+    hello = pb2.WorkerHello(worker_id=worker_id)
+    envelope = pb2.Envelope(
+        delivery_id=0,
+        partition_id=0,
+        kind=pb2.MessageKind.MESSAGE_KIND_WORKER_HELLO,
+        payload=hello.SerializeToString(),
+    )
+    yield envelope
+    try:
+        while True:
+            message = await queue.get()
+            yield message
+    except asyncio.CancelledError:  # pragma: no cover - best effort shutdown
+        return
+
+
+async def _send_ack(outgoing: "asyncio.Queue[pb2.Envelope]", envelope: pb2.Envelope) -> None:
+    ack = pb2.Ack(acked_delivery_id=envelope.delivery_id)
+    ack_envelope = pb2.Envelope(
+        delivery_id=envelope.delivery_id,
+        partition_id=envelope.partition_id,
+        kind=pb2.MessageKind.MESSAGE_KIND_ACK,
+        payload=ack.SerializeToString(),
+    )
+    await outgoing.put(ack_envelope)
+
+
+async def _handle_dispatch(
+    envelope: pb2.Envelope,
+    outgoing: "asyncio.Queue[pb2.Envelope]",
+) -> None:
+    await _send_ack(outgoing, envelope)
+    dispatch = pb2.ActionDispatch()
+    dispatch.ParseFromString(envelope.payload)
+    timeout_seconds = dispatch.timeout_seconds if dispatch.HasField("timeout_seconds") else 0
+
+    worker_start = time.perf_counter_ns()
+    success = True
+    action_name = dispatch.action_name
+    execution: workflow_runtime.ActionExecutionResult | None = None
+    try:
+        if timeout_seconds > 0:
+            execution = await asyncio.wait_for(
+                workflow_runtime.execute_action(dispatch), timeout=timeout_seconds
+            )
+        else:
+            execution = await workflow_runtime.execute_action(dispatch)
+
+        if execution.exception:
+            success = False
+            response_payload = serialize_error_payload(action_name, execution.exception)
+        else:
+            response_payload = serialize_result_payload(execution.result)
+    except asyncio.TimeoutError:
+        success = False
+        error = TimeoutError(f"action {action_name} timed out after {timeout_seconds} seconds")
+        response_payload = serialize_error_payload(action_name, error)
+        LOGGER.warning(
+            "Action %s timed out after %ss for action_id=%s sequence=%s",
+            action_name,
+            timeout_seconds,
+            dispatch.action_id,
+            dispatch.sequence,
+        )
+    except Exception as exc:  # noqa: BLE001 - propagate structured errors
+        success = False
+        response_payload = serialize_error_payload(action_name, exc)
+        LOGGER.exception(
+            "Action %s failed for action_id=%s sequence=%s",
+            action_name,
+            dispatch.action_id,
+            dispatch.sequence,
+        )
+    worker_end = time.perf_counter_ns()
+    response = pb2.ActionResult(
+        action_id=dispatch.action_id,
+        success=success,
+        worker_start_ns=worker_start,
+        worker_end_ns=worker_end,
+    )
+    response.payload.CopyFrom(response_payload)
+    if dispatch.dispatch_token:
+        response.dispatch_token = dispatch.dispatch_token
+    response_envelope = pb2.Envelope(
+        delivery_id=envelope.delivery_id,
+        partition_id=envelope.partition_id,
+        kind=pb2.MessageKind.MESSAGE_KIND_ACTION_RESULT,
+        payload=response.SerializeToString(),
+    )
+    await outgoing.put(response_envelope)
+    LOGGER.debug("Handled action=%s seq=%s success=%s", action_name, dispatch.sequence, success)
+
+
+async def _handle_incoming_stream(
+    stub: pb2_grpc.WorkerBridgeStub,
+    worker_id: int,
+    outgoing: "asyncio.Queue[pb2.Envelope]",
+) -> None:
+    """Process incoming messages, running action dispatches concurrently."""
+    pending_tasks: set[asyncio.Task[None]] = set()
+
+    async for envelope in stub.Attach(_outgoing_stream(outgoing, worker_id)):
+        kind = envelope.kind
+        if kind == pb2.MessageKind.MESSAGE_KIND_ACTION_DISPATCH:
+            # Spawn task to handle dispatch concurrently
+            task = asyncio.create_task(_handle_dispatch(envelope, outgoing))
+            pending_tasks.add(task)
+            task.add_done_callback(pending_tasks.discard)
+        elif kind == pb2.MessageKind.MESSAGE_KIND_HEARTBEAT:
+            LOGGER.debug("Received heartbeat delivery=%s", envelope.delivery_id)
+            await _send_ack(outgoing, envelope)
+        else:
+            LOGGER.warning("Unhandled message kind: %s", kind)
+            await _send_ack(outgoing, envelope)
+
+    # Wait for any remaining tasks on stream close
+    if pending_tasks:
+        await asyncio.gather(*pending_tasks, return_exceptions=True)
+
+
+async def _run_worker(args: argparse.Namespace) -> None:
+    outgoing: "asyncio.Queue[pb2.Envelope]" = asyncio.Queue()
+    for module_name in args.user_module:
+        if not module_name:
+            continue
+        LOGGER.info("Preloading user module %s", module_name)
+        importlib.import_module(module_name)
+
+    async with aio.insecure_channel(args.bridge) as channel:
+        stub = pb2_grpc.WorkerBridgeStub(channel)
+        LOGGER.info("Worker %s connected to %s", args.worker_id, args.bridge)
+        try:
+            await _handle_incoming_stream(stub, args.worker_id, outgoing)
+        except aio.AioRpcError as exc:  # pragma: no cover
+            status = exc.code()
+            LOGGER.error("Worker stream closed: %s", status)
+            raise
+
+
+def main(argv: list[str] | None = None) -> None:
+    args = _parse_args(argv)
+    logging.basicConfig(level=logging.INFO, format="[worker] %(message)s", stream=sys.stderr)
+    try:
+        asyncio.run(_run_worker(args))
+    except KeyboardInterrupt:  # pragma: no cover - exit quietly on Ctrl+C
+        return
+    except grpc.RpcError:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

rappel/workflow.py

@@ -0,0 +1,236 @@
+"""
+Workflow base class and registration decorator.
+
+This module provides the foundation for defining workflows that can be
+compiled to IR and executed by the Rappel runtime.
+"""
+
+import hashlib
+import inspect
+import os
+from dataclasses import dataclass
+from datetime import timedelta
+from functools import wraps
+from threading import RLock
+from typing import Any, Awaitable, ClassVar, Optional, TypeVar
+
+from proto import ast_pb2 as ir
+from proto import messages_pb2 as pb2
+
+from . import bridge
+from .actions import deserialize_result_payload
+from .ir_builder import build_workflow_ir
+from .logger import configure as configure_logger
+from .serialization import build_arguments_from_kwargs
+from .workflow_runtime import WorkflowNodeResult
+
+logger = configure_logger("rappel.workflow")
+
+TWorkflow = TypeVar("TWorkflow", bound="Workflow")
+TResult = TypeVar("TResult")
+
+
+@dataclass(frozen=True)
+class RetryPolicy:
+    """Retry policy for action execution.
+
+    Maps to IR RetryPolicy: [ExceptionType -> retry: N, backoff: Xs]
+
+    Args:
+        attempts: Maximum number of retry attempts.
+        exception_types: List of exception type names to retry on. Empty = catch all.
+        backoff_seconds: Constant backoff duration between retries in seconds.
+    """
+
+    attempts: Optional[int] = None
+    exception_types: Optional[list[str]] = None
+    backoff_seconds: Optional[float] = None
+
+
+class Workflow:
+    """Base class for workflow definitions."""
+
+    name: ClassVar[Optional[str]] = None
+    """Human-friendly identifier. Override to pin the registry key; defaults to lowercase class name."""
+
+    concurrent: ClassVar[bool] = False
+    """When True, downstream engines may respect DAG-parallel execution; False preserves sequential semantics."""
+
+    _workflow_ir: ClassVar[Optional[ir.Program]] = None
+    _ir_lock: ClassVar[RLock] = RLock()
+    _workflow_version_id: ClassVar[Optional[str]] = None
+
+    async def run(self) -> Any:
+        raise NotImplementedError
+
+    async def run_action(
+        self,
+        awaitable: Awaitable[TResult],
+        *,
+        retry: Optional[RetryPolicy] = None,
+        timeout: Optional[float | int | timedelta] = None,
+    ) -> TResult:
+        """Helper that simply awaits the provided action coroutine.
+
+        The retry and timeout arguments are consumed by the workflow compiler
+        (IR builder) rather than the runtime execution path.
+
+        Args:
+            awaitable: The action coroutine to execute.
+            retry: Retry policy including max attempts, exception types, and backoff.
+            timeout: Timeout duration in seconds (or timedelta).
+        """
+        # Parameters are intentionally unused at runtime; the workflow compiler
+        # inspects the AST to record them.
+        del retry, timeout
+        return await awaitable
+
+    @classmethod
+    def short_name(cls) -> str:
+        if cls.name:
+            return cls.name
+        return cls.__name__.lower()
+
+    @classmethod
+    def workflow_ir(cls) -> ir.Program:
+        """Build and cache the IR program for this workflow."""
+        if cls._workflow_ir is None:
+            with cls._ir_lock:
+                if cls._workflow_ir is None:
+                    cls._workflow_ir = build_workflow_ir(cls)
+        return cls._workflow_ir
+
+    @classmethod
+    def _build_registration_payload(
+        cls, initial_context: Optional[pb2.WorkflowArguments] = None
+    ) -> pb2.WorkflowRegistration:
+        """Build a registration payload with the serialized IR."""
+        program = cls.workflow_ir()
+
+        # Serialize IR to bytes
+        ir_bytes = program.SerializeToString()
+        ir_hash = hashlib.sha256(ir_bytes).hexdigest()
+
+        message = pb2.WorkflowRegistration(
+            workflow_name=cls.short_name(),
+            ir=ir_bytes,
+            ir_hash=ir_hash,
+            concurrent=cls.concurrent,
+        )
+
+        if initial_context:
+            message.initial_context.CopyFrom(initial_context)
+
+        return message
+
+
+class WorkflowRegistry:
+    """Registry of workflow definitions keyed by workflow name."""
+
+    def __init__(self) -> None:
+        self._workflows: dict[str, type[Workflow]] = {}
+        self._lock = RLock()
+
+    def register(self, name: str, workflow_cls: type[Workflow]) -> None:
+        with self._lock:
+            if name in self._workflows:
+                raise ValueError(f"workflow '{name}' already registered")
+            self._workflows[name] = workflow_cls
+
+    def get(self, name: str) -> Optional[type[Workflow]]:
+        with self._lock:
+            return self._workflows.get(name)
+
+    def names(self) -> list[str]:
+        with self._lock:
+            return sorted(self._workflows.keys())
+
+    def reset(self) -> None:
+        with self._lock:
+            self._workflows.clear()
+
+
+workflow_registry = WorkflowRegistry()
+
+
+def workflow(cls: type[TWorkflow]) -> type[TWorkflow]:
+    """Decorator that registers workflow classes and caches their IR."""
+
+    if not issubclass(cls, Workflow):
+        raise TypeError("workflow decorator requires Workflow subclasses")
+    run_impl = cls.run
+    if not inspect.iscoroutinefunction(run_impl):
+        raise TypeError("workflow run() must be defined with 'async def'")
+
+    @wraps(run_impl)
+    async def run_public(self: Workflow, *args: Any, **kwargs: Any) -> Any:
+        if _running_under_pytest():
+            cls.workflow_ir()
+            return await run_impl(self, *args, **kwargs)
+
+        # Get the signature of run() to map positional args to parameter names
+        sig = inspect.signature(run_impl)
+        params = list(sig.parameters.keys())[1:]  # Skip 'self'
+
+        # Convert positional args to kwargs
+        for i, arg in enumerate(args):
+            if i < len(params):
+                kwargs[params[i]] = arg
+
+        # Serialize kwargs using common logic
+        initial_context = build_arguments_from_kwargs(kwargs)
+
+        payload = cls._build_registration_payload(initial_context)
+        run_result = await bridge.run_instance(payload.SerializeToString())
+        cls._workflow_version_id = run_result.workflow_version_id
+        if _skip_wait_for_instance():
+            logger.info(
+                "Skipping wait_for_instance for workflow %s due to RAPPEL_SKIP_WAIT_FOR_INSTANCE",
+                cls.short_name(),
+            )
+            return None
+        result_bytes = await bridge.wait_for_instance(
+            instance_id=run_result.workflow_instance_id,
+            poll_interval_secs=1.0,
+        )
+        if result_bytes is None:
+            raise TimeoutError(
+                f"workflow instance {run_result.workflow_instance_id} did not complete"
+            )
+        arguments = pb2.WorkflowArguments()
+        arguments.ParseFromString(result_bytes)
+        result = deserialize_result_payload(arguments)
+        if result.error:
+            raise RuntimeError(f"workflow failed: {result.error}")
+
+        # Unwrap WorkflowNodeResult if present (internal worker representation)
+        if isinstance(result.result, WorkflowNodeResult):
+            # Extract the actual result from the variables dict
+            variables = result.result.variables
+            program = cls.workflow_ir()
+            # Get the return variable from the IR if available
+            if program.functions:
+                outputs = list(program.functions[0].io.outputs)
+                if outputs:
+                    return_var = outputs[0]
+                    if return_var in variables:
+                        return variables[return_var]
+            return None
+
+        return result.result
+
+    cls.__workflow_run_impl__ = run_impl
+    cls.run = run_public  # type: ignore[assignment]
+    workflow_registry.register(cls.short_name(), cls)
+    return cls
+
+
+def _running_under_pytest() -> bool:
+    return bool(os.environ.get("PYTEST_CURRENT_TEST"))
+
+
+def _skip_wait_for_instance() -> bool:
+    value = os.environ.get("RAPPEL_SKIP_WAIT_FOR_INSTANCE")
+    if not value:
+        return False
+    return value.strip().lower() not in {"0", "false", "no"}

rappel/workflow_runtime.py

@@ -0,0 +1,137 @@
+"""Runtime helpers for executing actions inside the worker.
+
+This module provides the execution layer for Python workers that receive
+action dispatch commands from the Rust scheduler.
+"""
+
+import asyncio
+import dataclasses
+from dataclasses import dataclass
+from typing import Any, Dict, get_type_hints
+
+from pydantic import BaseModel
+
+from proto import messages_pb2 as pb2
+
+from .dependencies import provide_dependencies
+from .registry import registry
+from .serialization import arguments_to_kwargs
+
+
+class WorkflowNodeResult(BaseModel):
+    """Result from a workflow node execution containing variable bindings."""
+
+    variables: Dict[str, Any]
+
+
+@dataclass
+class ActionExecutionResult:
+    """Result of an action execution."""
+
+    result: Any
+    exception: BaseException | None = None
+
+
+def _is_pydantic_model(cls: type) -> bool:
+    """Check if a class is a Pydantic BaseModel subclass."""
+    try:
+        return isinstance(cls, type) and issubclass(cls, BaseModel)
+    except TypeError:
+        return False
+
+
+def _is_dataclass_type(cls: type) -> bool:
+    """Check if a class is a dataclass."""
+    return dataclasses.is_dataclass(cls) and isinstance(cls, type)
+
+
+def _coerce_dict_to_model(value: Any, target_type: type) -> Any:
+    """Convert a dict to a Pydantic model or dataclass if needed.
+
+    If value is a dict and target_type is a Pydantic model or dataclass,
+    instantiate the model with the dict values. Otherwise, return value unchanged.
+    """
+    if not isinstance(value, dict):
+        return value
+
+    if _is_pydantic_model(target_type):
+        # Use model_validate for Pydantic v2, fall back to direct instantiation
+        model_validate = getattr(target_type, "model_validate", None)
+        if model_validate is not None:
+            return model_validate(value)
+        return target_type(**value)
+
+    if _is_dataclass_type(target_type):
+        return target_type(**value)
+
+    return value
+
+
+def _coerce_kwargs_to_type_hints(handler: Any, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    """Coerce dict kwargs to Pydantic models or dataclasses based on type hints.
+
+    When the IR converts a Pydantic model or dataclass constructor call to a dict,
+    the action runner needs to convert that dict back to the expected type based
+    on the handler's type annotations.
+    """
+    try:
+        type_hints = get_type_hints(handler)
+    except Exception:
+        # If we can't get type hints (e.g., forward references), return as-is
+        return kwargs
+
+    coerced = {}
+    for key, value in kwargs.items():
+        if key in type_hints:
+            target_type = type_hints[key]
+            coerced[key] = _coerce_dict_to_model(value, target_type)
+        else:
+            coerced[key] = value
+
+    return coerced
+
+
+async def execute_action(dispatch: pb2.ActionDispatch) -> ActionExecutionResult:
+    """Execute an action based on the dispatch command.
+
+    Args:
+        dispatch: The action dispatch command from the Rust scheduler.
+
+    Returns:
+        The result of executing the action.
+    """
+    action_name = dispatch.action_name
+    module_name = dispatch.module_name
+
+    # Import the module if specified (this registers actions via @action decorator)
+    if module_name:
+        import importlib
+
+        importlib.import_module(module_name)
+
+    # Get the action handler using both module and name
+    handler = registry.get(module_name, action_name)
+    if handler is None:
+        return ActionExecutionResult(
+            result=None,
+            exception=KeyError(f"action '{module_name}:{action_name}' not registered"),
+        )
+
+    # Deserialize kwargs
+    kwargs = arguments_to_kwargs(dispatch.kwargs)
+
+    # Coerce dict arguments to Pydantic models or dataclasses based on type hints
+    # This is needed because the IR converts model constructor calls to dicts
+    kwargs = _coerce_kwargs_to_type_hints(handler, kwargs)
+
+    try:
+        async with provide_dependencies(handler, kwargs) as call_kwargs:
+            value = handler(**call_kwargs)
+            if asyncio.iscoroutine(value):
+                value = await value
+        return ActionExecutionResult(result=value)
+    except Exception as e:
+        return ActionExecutionResult(
+            result=None,
+            exception=e,
+        )