rappel 0.7.2__py3-none-macosx_15_0_arm64.whl

rappel/serialization.py

@@ -0,0 +1,273 @@
+import dataclasses
+import importlib
+import traceback
+from base64 import b64encode
+from datetime import date, datetime, time, timedelta
+from decimal import Decimal
+from enum import Enum
+from pathlib import PurePath
+from typing import Any
+from uuid import UUID
+
+from google.protobuf import json_format, struct_pb2
+from pydantic import BaseModel
+
+from proto import messages_pb2 as pb2
+
+NULL_VALUE = struct_pb2.NULL_VALUE  # type: ignore[attr-defined]
+
+PRIMITIVE_TYPES = (str, int, float, bool, type(None))
+
+
+def dumps(value: Any) -> pb2.WorkflowArgumentValue:
+    """Serialize a Python value into a WorkflowArgumentValue message."""
+
+    return _to_argument_value(value)
+
+
+def loads(data: Any) -> Any:
+    """Deserialize a workflow argument payload into a Python object."""
+
+    if isinstance(data, pb2.WorkflowArgumentValue):
+        argument = data
+    elif isinstance(data, dict):
+        argument = pb2.WorkflowArgumentValue()
+        json_format.ParseDict(data, argument)
+    else:
+        raise TypeError("argument value payload must be a dict or ArgumentValue message")
+    return _from_argument_value(argument)
+
+
+def build_arguments_from_kwargs(kwargs: dict[str, Any]) -> pb2.WorkflowArguments:
+    arguments = pb2.WorkflowArguments()
+    for key, value in kwargs.items():
+        entry = arguments.arguments.add()
+        entry.key = key
+        entry.value.CopyFrom(dumps(value))
+    return arguments
+
+
+def arguments_to_kwargs(arguments: pb2.WorkflowArguments | None) -> dict[str, Any]:
+    if arguments is None:
+        return {}
+    result: dict[str, Any] = {}
+    for entry in arguments.arguments:
+        result[entry.key] = loads(entry.value)
+    return result
+
+
+def _to_argument_value(value: Any) -> pb2.WorkflowArgumentValue:
+    argument = pb2.WorkflowArgumentValue()
+    if isinstance(value, PRIMITIVE_TYPES):
+        argument.primitive.CopyFrom(_serialize_primitive(value))
+        return argument
+    if isinstance(value, UUID):
+        # Serialize UUID as string primitive
+        argument.primitive.CopyFrom(_serialize_primitive(str(value)))
+        return argument
+    if isinstance(value, datetime):
+        # Serialize datetime as ISO format string
+        argument.primitive.CopyFrom(_serialize_primitive(value.isoformat()))
+        return argument
+    if isinstance(value, date):
+        # Serialize date as ISO format string (must come after datetime check)
+        argument.primitive.CopyFrom(_serialize_primitive(value.isoformat()))
+        return argument
+    if isinstance(value, time):
+        # Serialize time as ISO format string
+        argument.primitive.CopyFrom(_serialize_primitive(value.isoformat()))
+        return argument
+    if isinstance(value, timedelta):
+        # Serialize timedelta as total seconds
+        argument.primitive.CopyFrom(_serialize_primitive(value.total_seconds()))
+        return argument
+    if isinstance(value, Decimal):
+        # Serialize Decimal as string to preserve precision
+        argument.primitive.CopyFrom(_serialize_primitive(str(value)))
+        return argument
+    if isinstance(value, Enum):
+        # Serialize Enum as its value
+        return _to_argument_value(value.value)
+    if isinstance(value, bytes):
+        # Serialize bytes as base64 string
+        argument.primitive.CopyFrom(_serialize_primitive(b64encode(value).decode("ascii")))
+        return argument
+    if isinstance(value, PurePath):
+        # Serialize Path as string
+        argument.primitive.CopyFrom(_serialize_primitive(str(value)))
+        return argument
+    if isinstance(value, (set, frozenset)):
+        # Serialize sets as lists
+        argument.list_value.SetInParent()
+        for item in value:
+            item_value = argument.list_value.items.add()
+            item_value.CopyFrom(_to_argument_value(item))
+        return argument
+    if isinstance(value, BaseException):
+        argument.exception.type = value.__class__.__name__
+        argument.exception.module = value.__class__.__module__
+        argument.exception.message = str(value)
+        tb_text = "".join(traceback.format_exception(type(value), value, value.__traceback__))
+        argument.exception.traceback = tb_text
+        values = _serialize_exception_values(value)
+        for key, item in values.items():
+            entry = argument.exception.values.entries.add()
+            entry.key = key
+            try:
+                entry.value.CopyFrom(_to_argument_value(item))
+            except TypeError:
+                entry.value.CopyFrom(_to_argument_value(str(item)))
+        return argument
+    if _is_base_model(value):
+        model_class = value.__class__
+        model_data = _serialize_model_data(value)
+        argument.basemodel.module = model_class.__module__
+        argument.basemodel.name = model_class.__qualname__
+        # Serialize as dict to preserve types (Struct converts all numbers to float)
+        for key, item in model_data.items():
+            entry = argument.basemodel.data.entries.add()
+            entry.key = key
+            entry.value.CopyFrom(_to_argument_value(item))
+        return argument
+    if _is_dataclass_instance(value):
+        # Dataclasses use the same basemodel serialization path as Pydantic models
+        dc_class = value.__class__
+        dc_data = dataclasses.asdict(value)
+        argument.basemodel.module = dc_class.__module__
+        argument.basemodel.name = dc_class.__qualname__
+        for key, item in dc_data.items():
+            entry = argument.basemodel.data.entries.add()
+            entry.key = key
+            entry.value.CopyFrom(_to_argument_value(item))
+        return argument
+    if isinstance(value, dict):
+        argument.dict_value.SetInParent()
+        for key, item in value.items():
+            if not isinstance(key, str):
+                raise TypeError("workflow dict keys must be strings")
+            entry = argument.dict_value.entries.add()
+            entry.key = key
+            entry.value.CopyFrom(_to_argument_value(item))
+        return argument
+    if isinstance(value, list):
+        argument.list_value.SetInParent()
+        for item in value:
+            item_value = argument.list_value.items.add()
+            item_value.CopyFrom(_to_argument_value(item))
+        return argument
+    if isinstance(value, tuple):
+        argument.tuple_value.SetInParent()
+        for item in value:
+            item_value = argument.tuple_value.items.add()
+            item_value.CopyFrom(_to_argument_value(item))
+        return argument
+    raise TypeError(f"unsupported value type {type(value)!r}")
+
+
+def _from_argument_value(argument: pb2.WorkflowArgumentValue) -> Any:
+    kind = argument.WhichOneof("kind")  # type: ignore[attr-defined]
+    if kind == "primitive":
+        return _primitive_to_python(argument.primitive)
+    if kind == "basemodel":
+        module = argument.basemodel.module
+        name = argument.basemodel.name
+        # Deserialize dict entries to preserve types
+        data: dict[str, Any] = {}
+        for entry in argument.basemodel.data.entries:
+            data[entry.key] = _from_argument_value(entry.value)
+        return _instantiate_serialized_model(module, name, data)
+    if kind == "exception":
+        values: dict[str, Any] = {}
+        if argument.exception.HasField("values"):
+            for entry in argument.exception.values.entries:
+                values[entry.key] = _from_argument_value(entry.value)
+        return {
+            "type": argument.exception.type,
+            "module": argument.exception.module,
+            "message": argument.exception.message,
+            "traceback": argument.exception.traceback,
+            "values": values,
+        }
+    if kind == "list_value":
+        return [_from_argument_value(item) for item in argument.list_value.items]
+    if kind == "tuple_value":
+        return tuple(_from_argument_value(item) for item in argument.tuple_value.items)
+    if kind == "dict_value":
+        result: dict[str, Any] = {}
+        for entry in argument.dict_value.entries:
+            result[entry.key] = _from_argument_value(entry.value)
+        return result
+    raise ValueError("argument value missing kind discriminator")
+
+
+def _serialize_model_data(model: BaseModel) -> dict[str, Any]:
+    if hasattr(model, "model_dump"):
+        return model.model_dump(mode="json")  # type: ignore[attr-defined]
+    if hasattr(model, "dict"):
+        return model.dict()  # type: ignore[attr-defined]
+    return model.__dict__
+
+
+def _serialize_exception_values(exc: BaseException) -> dict[str, Any]:
+    values = dict(vars(exc))
+    if "args" not in values:
+        values["args"] = exc.args
+    return values
+
+
+def _serialize_primitive(value: Any) -> pb2.PrimitiveWorkflowArgument:
+    primitive = pb2.PrimitiveWorkflowArgument()
+    if value is None:
+        primitive.null_value = NULL_VALUE
+    elif isinstance(value, bool):
+        primitive.bool_value = value
+    elif isinstance(value, int) and not isinstance(value, bool):
+        primitive.int_value = value
+    elif isinstance(value, float):
+        primitive.double_value = value
+    elif isinstance(value, str):
+        primitive.string_value = value
+    else:  # pragma: no cover - unreachable given PRIMITIVE_TYPES
+        raise TypeError(f"unsupported primitive type {type(value)!r}")
+    return primitive
+
+
+def _primitive_to_python(primitive: pb2.PrimitiveWorkflowArgument) -> Any:
+    kind = primitive.WhichOneof("kind")  # type: ignore[attr-defined]
+    if kind == "string_value":
+        return primitive.string_value
+    if kind == "double_value":
+        return primitive.double_value
+    if kind == "int_value":
+        return primitive.int_value
+    if kind == "bool_value":
+        return primitive.bool_value
+    if kind == "null_value":
+        return None
+    raise ValueError("primitive argument missing kind discriminator")
+
+
+def _instantiate_serialized_model(module: str, name: str, model_data: dict[str, Any]) -> Any:
+    cls = _import_symbol(module, name)
+    if hasattr(cls, "model_validate"):
+        return cls.model_validate(model_data)  # type: ignore[attr-defined]
+    return cls(**model_data)
+
+
+def _is_base_model(value: Any) -> bool:
+    return isinstance(value, BaseModel)
+
+
+def _is_dataclass_instance(value: Any) -> bool:
+    """Check if value is a dataclass instance (not a class)."""
+    return dataclasses.is_dataclass(value) and not isinstance(value, type)
+
+
+def _import_symbol(module: str, qualname: str) -> Any:
+    module_obj = importlib.import_module(module)
+    attr: Any = module_obj
+    for part in qualname.split("."):
+        attr = getattr(attr, part)
+    if not isinstance(attr, type):
+        raise ValueError(f"{qualname} from {module} is not a class")
+    return attr

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()