waymark 0.30.0.dev2__py3-none-macosx_26_0_arm64.whl

waymark/__init__.py

@@ -0,0 +1,62 @@
+"""Public API for user-defined waymark actions."""
+
+from . import bridge  # noqa: F401
+from . import workflow_runtime as _workflow_runtime  # noqa: F401
+from .actions import (
+    ActionResultPayload,
+    action,
+    deserialize_result_payload,
+    serialize_error_payload,
+    serialize_result_payload,
+)
+from .dependencies import Depend, Depends, provide_dependencies
+from .exceptions import (
+    ExhaustedRetries,
+    ExhaustedRetriesError,
+    ScheduleAlreadyExistsError,
+)
+from .ir_builder import UnsupportedPatternError, build_workflow_ir
+from .registry import registry
+from .schedule import (
+    ScheduleInfo,
+    delete_schedule,
+    list_schedules,
+    pause_schedule,
+    resume_schedule,
+    schedule_workflow,
+)
+from .workflow import (
+    RetryPolicy,
+    Workflow,
+    workflow,
+    workflow_registry,
+)
+
+__all__ = [
+    "action",
+    "registry",
+    "ActionResultPayload",
+    "Workflow",
+    "workflow",
+    "workflow_registry",
+    "RetryPolicy",
+    "build_workflow_ir",
+    "serialize_result_payload",
+    "deserialize_result_payload",
+    "serialize_error_payload",
+    "Depend",
+    "Depends",
+    "provide_dependencies",
+    "bridge",
+    "ExhaustedRetries",
+    "ExhaustedRetriesError",
+    "ScheduleAlreadyExistsError",
+    "UnsupportedPatternError",
+    # Schedule functions
+    "schedule_workflow",
+    "pause_schedule",
+    "resume_schedule",
+    "delete_schedule",
+    "list_schedules",
+    "ScheduleInfo",
+]

waymark/actions.py

@@ -0,0 +1,108 @@
+import inspect
+from dataclasses import dataclass
+from functools import wraps
+from typing import Any, Callable, Optional, TypeVar, overload
+
+from waymark.proto import messages_pb2 as pb2
+
+from .dependencies import provide_dependencies
+from .registry import AsyncAction, registry
+from .serialization import dumps, loads
+
+TAsync = TypeVar("TAsync", bound=AsyncAction)
+
+
+@dataclass
+class ActionResultPayload:
+    result: Any | None
+    error: dict[str, Any] | None
+
+
+def serialize_result_payload(value: Any) -> pb2.WorkflowArguments:
+    """Serialize a successful action result."""
+    arguments = pb2.WorkflowArguments()
+    entry = arguments.arguments.add()
+    entry.key = "result"
+    entry.value.CopyFrom(dumps(value))
+    return arguments
+
+
+def serialize_error_payload(_action: str, exc: BaseException) -> pb2.WorkflowArguments:
+    """Serialize an error raised during action execution."""
+    arguments = pb2.WorkflowArguments()
+    entry = arguments.arguments.add()
+    entry.key = "error"
+    entry.value.CopyFrom(dumps(exc))
+    return arguments
+
+
+def deserialize_result_payload(payload: pb2.WorkflowArguments | None) -> ActionResultPayload:
+    """Deserialize WorkflowArguments produced by serialize_result_payload/error."""
+    if payload is None:
+        return ActionResultPayload(result=None, error=None)
+    values = {entry.key: entry.value for entry in payload.arguments}
+    if "error" in values:
+        error_value = values["error"]
+        data = loads(error_value)
+        if not isinstance(data, dict):
+            raise ValueError("error payload must deserialize to a mapping")
+        return ActionResultPayload(result=None, error=data)
+    result_value = values.get("result")
+    if result_value is None:
+        raise ValueError("result payload missing 'result' field")
+    return ActionResultPayload(result=loads(result_value), error=None)
+
+
+@overload
+def action(func: TAsync, /) -> TAsync: ...
+
+
+@overload
+def action(*, name: Optional[str] = None) -> Callable[[TAsync], TAsync]: ...
+
+
+def action(
+    func: Optional[TAsync] = None,
+    *,
+    name: Optional[str] = None,
+) -> Callable[[TAsync], TAsync] | TAsync:
+    """Decorator for registering async actions.
+
+    Actions decorated with @action will automatically resolve dependency markers
+    when called directly (e.g., during pytest runs where workflows bypass the
+    gRPC bridge).
+    """
+
+    def decorator(target: TAsync) -> TAsync:
+        if not inspect.iscoroutinefunction(target):
+            raise TypeError(f"action '{target.__name__}' must be defined with 'async def'")
+        action_name = name or target.__name__
+        action_module = target.__module__
+
+        @wraps(target)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Convert positional args to kwargs based on the signature
+            sig = inspect.signature(target)
+            params = list(sig.parameters.keys())
+            for i, arg in enumerate(args):
+                if i < len(params):
+                    kwargs[params[i]] = arg
+
+            # Resolve dependencies using the same mechanism as execute_action
+            async with provide_dependencies(target, kwargs) as call_kwargs:
+                return await target(**call_kwargs)
+
+        # Copy over the original function's attributes for introspection
+        wrapper.__wrapped__ = target  # type: ignore[attr-defined]
+        wrapper.__waymark_action_name__ = action_name  # type: ignore[attr-defined]
+        wrapper.__waymark_action_module__ = action_module  # type: ignore[attr-defined]
+
+        # Register the original function (not the wrapper) so execute_action
+        # doesn't double-resolve dependencies
+        registry.register(action_module, action_name, target)
+
+        return wrapper  # type: ignore[return-value]
+
+    if func is not None:
+        return decorator(func)
+    return decorator

waymark/bridge.py

@@ -0,0 +1,432 @@
+import asyncio
+import os
+import shlex
+import shutil
+import subprocess
+import tempfile
+import time
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from threading import Lock, RLock
+from typing import AsyncIterator, NoReturn, Optional
+
+import grpc
+from grpc import aio  # type: ignore[attr-defined]
+
+from waymark.logger import configure as configure_logger
+from waymark.proto import messages_pb2 as pb2
+from waymark.proto import messages_pb2_grpc as pb2_grpc
+
+from .actions import serialize_error_payload, serialize_result_payload
+from .grpc_config import GRPC_CHANNEL_OPTIONS
+from .workflow_runtime import execute_action
+
+DEFAULT_HOST = "127.0.0.1"
+LOGGER = configure_logger("waymark.bridge")
+
+_PORT_LOCK = RLock()
+_CACHED_GRPC_PORT: Optional[int] = None
+_GRPC_TARGET: Optional[str] = None
+_GRPC_CHANNEL: Optional[aio.Channel] = None
+_GRPC_STUB: Optional[pb2_grpc.WorkflowServiceStub] = None
+_GRPC_LOOP: Optional[asyncio.AbstractEventLoop] = None
+_BOOT_MUTEX = Lock()
+_ASYNC_BOOT_LOCK: asyncio.Lock = asyncio.Lock()
+
+
+@dataclass
+class RunInstanceResult:
+    workflow_version_id: str
+    workflow_instance_id: str
+
+
+@dataclass
+class RunBatchResult:
+    workflow_version_id: str
+    workflow_instance_ids: list[str]
+    queued: int
+
+
+def _boot_command() -> list[str]:
+    override = os.environ.get("WAYMARK_BOOT_COMMAND")
+    if override:
+        LOGGER.debug("Using WAYMARK_BOOT_COMMAND=%s", override)
+        return shlex.split(override)
+    binary = os.environ.get("WAYMARK_BOOT_BINARY", "waymark-boot-singleton")
+    LOGGER.debug("Using WAYMARK_BOOT_BINARY=%s", binary)
+    return [binary]
+
+
+def _repo_root() -> Path:
+    return Path(__file__).resolve().parents[3]
+
+
+def _resolve_boot_binary(binary: str) -> str:
+    if Path(binary).is_absolute():
+        return binary
+    resolved = shutil.which(binary)
+    if resolved:
+        return resolved
+    repo_root = _repo_root()
+    # Check multiple target directories:
+    # - CARGO_TARGET_DIR (set by cargo llvm-cov and other tools)
+    # - target/llvm-cov-target (used by cargo llvm-cov by default)
+    # - target (standard cargo target)
+    target_dirs = []
+    if cargo_target_dir := os.environ.get("CARGO_TARGET_DIR"):
+        target_dirs.append(Path(cargo_target_dir))
+    target_dirs.extend(
+        [
+            repo_root / "target" / "llvm-cov-target",
+            repo_root / "target",
+        ]
+    )
+    for target_dir in target_dirs:
+        for profile in ("debug", "release"):
+            candidate = target_dir / profile / binary
+            if candidate.exists():
+                return str(candidate)
+    return binary
+
+
+def _ensure_boot_binary(binary: str) -> str:
+    resolved = _resolve_boot_binary(binary)
+    if Path(resolved).exists():
+        return resolved
+    repo_root = _repo_root()
+    cargo_toml = repo_root / "Cargo.toml"
+    if cargo_toml.exists():
+        LOGGER.info("boot binary %s not found; building via cargo", binary)
+        subprocess.run(
+            [
+                "cargo",
+                "build",
+                "--bin",
+                "waymark-boot-singleton",
+                "--bin",
+                "waymark-bridge",
+            ],
+            cwd=repo_root,
+            check=True,
+        )
+        resolved = _resolve_boot_binary(binary)
+    return resolved
+
+
+def _remember_grpc_port(port: int) -> int:
+    global _CACHED_GRPC_PORT
+    with _PORT_LOCK:
+        _CACHED_GRPC_PORT = port
+    return port
+
+
+def _cached_grpc_port() -> Optional[int]:
+    with _PORT_LOCK:
+        return _CACHED_GRPC_PORT
+
+
+def _env_grpc_port_override() -> Optional[int]:
+    """Check for explicit gRPC port override via environment."""
+    override = os.environ.get("WAYMARK_BRIDGE_GRPC_PORT")
+    if not override:
+        return None
+    try:
+        return int(override)
+    except ValueError as exc:  # pragma: no cover
+        raise RuntimeError(f"invalid WAYMARK_BRIDGE_GRPC_PORT value: {override}") from exc
+
+
+def _boot_singleton_blocking() -> int:
+    """Boot the singleton and return the gRPC port."""
+    command = _boot_command()
+    if os.environ.get("WAYMARK_BOOT_COMMAND") is None:
+        command[0] = _ensure_boot_binary(command[0])
+    with tempfile.NamedTemporaryFile(mode="w+", suffix=".txt") as f:
+        output_file = Path(f.name)
+
+        command.extend(["--output-file", str(output_file)])
+        LOGGER.info("Booting waymark singleton via: %s", " ".join(command))
+
+        try:
+            subprocess.run(
+                command,
+                check=True,
+                timeout=10,
+            )
+        except subprocess.TimeoutExpired as exc:  # pragma: no cover
+            LOGGER.error("boot command timed out after %s seconds", exc.timeout)
+            raise RuntimeError("unable to boot waymark server") from exc
+        except subprocess.CalledProcessError as exc:  # pragma: no cover
+            LOGGER.error("boot command failed: %s", exc)
+            raise RuntimeError("unable to boot waymark server") from exc
+        except OSError as exc:  # pragma: no cover
+            LOGGER.error("unable to spawn boot command: %s", exc)
+            raise RuntimeError("unable to boot waymark server") from exc
+
+        try:
+            # We use a file as a message passer because passing a PIPE to the singleton launcher
+            # will block our code indefinitely
+            # The singleton launches the webserver subprocess to inherit the stdin/stdout that the
+            # singleton launcher receives; which means that in the case of a PIPE it would pass that
+            # pipe to the subprocess and therefore never correctly close the file descriptor and signal
+            # exit process status to Python.
+            port_str = output_file.read_text().strip()
+            grpc_port = int(port_str)
+            LOGGER.info("boot command reported singleton gRPC port %s", grpc_port)
+            return grpc_port
+        except (ValueError, FileNotFoundError) as exc:  # pragma: no cover
+            raise RuntimeError(f"unable to read port from output file: {exc}") from exc
+
+
+def _resolve_grpc_port() -> int:
+    """Resolve the gRPC port, booting singleton if necessary."""
+    cached = _cached_grpc_port()
+    if cached is not None:
+        return cached
+    env_port = _env_grpc_port_override()
+    if env_port is not None:
+        return _remember_grpc_port(env_port)
+    with _BOOT_MUTEX:
+        cached = _cached_grpc_port()
+        if cached is not None:
+            return cached
+        port = _boot_singleton_blocking()
+        return _remember_grpc_port(port)
+
+
+async def _ensure_grpc_port_async() -> int:
+    """Ensure we have a gRPC port, booting singleton if necessary."""
+    cached = _cached_grpc_port()
+    if cached is not None:
+        return cached
+    env_port = _env_grpc_port_override()
+    if env_port is not None:
+        return _remember_grpc_port(env_port)
+    async with _ASYNC_BOOT_LOCK:
+        cached = _cached_grpc_port()
+        if cached is not None:
+            return cached
+        loop = asyncio.get_running_loop()
+        LOGGER.info("No cached singleton found, booting new instance")
+        port = await loop.run_in_executor(None, _boot_singleton_blocking)
+        LOGGER.info("Singleton ready on gRPC port %s", port)
+        return _remember_grpc_port(port)
+
+
+@asynccontextmanager
+async def ensure_singleton() -> AsyncIterator[int]:
+    """Yield the gRPC port for the singleton server, booting it exactly once."""
+    port = await _ensure_grpc_port_async()
+    yield port
+
+
+def _grpc_target() -> str:
+    """Get the gRPC target address for the bridge server."""
+    # Check for explicit full address override
+    explicit = os.environ.get("WAYMARK_BRIDGE_GRPC_ADDR")
+    if explicit:
+        return explicit
+
+    # Otherwise, use host + port
+    host = os.environ.get("WAYMARK_BRIDGE_GRPC_HOST", DEFAULT_HOST)
+    port = _resolve_grpc_port()
+    return f"{host}:{port}"
+
+
+def assert_never(value: object) -> NoReturn:
+    raise AssertionError(f"Unhandled value: {value!r}")
+
+
+async def _workflow_stub() -> pb2_grpc.WorkflowServiceStub:
+    global _GRPC_TARGET, _GRPC_CHANNEL, _GRPC_STUB, _GRPC_LOOP
+    target = _grpc_target()
+    loop = asyncio.get_running_loop()
+    channel_to_wait: Optional[aio.Channel] = None
+    with _PORT_LOCK:
+        if (
+            _GRPC_STUB is not None
+            and _GRPC_TARGET == target
+            and _GRPC_LOOP is loop
+            and not loop.is_closed()
+        ):
+            return _GRPC_STUB
+        channel = aio.insecure_channel(target, options=GRPC_CHANNEL_OPTIONS)
+        stub = pb2_grpc.WorkflowServiceStub(channel)
+        _GRPC_CHANNEL = channel
+        _GRPC_STUB = stub
+        _GRPC_TARGET = target
+        _GRPC_LOOP = loop
+        channel_to_wait = channel
+    if channel_to_wait is not None:
+        await channel_to_wait.channel_ready()
+    return _GRPC_STUB  # type: ignore[return-value]
+
+
+async def run_instance(payload: bytes) -> RunInstanceResult:
+    """Register a workflow definition and start an instance over the gRPC bridge."""
+    async with ensure_singleton():
+        stub = await _workflow_stub()
+    registration = pb2.WorkflowRegistration()
+    registration.ParseFromString(payload)
+    request = pb2.RegisterWorkflowRequest(
+        registration=registration,
+    )
+    try:
+        response = await stub.RegisterWorkflow(request, timeout=30.0)
+    except aio.AioRpcError as exc:  # pragma: no cover
+        raise RuntimeError(f"register_workflow failed: {exc}") from exc
+    return RunInstanceResult(
+        workflow_version_id=response.workflow_version_id,
+        workflow_instance_id=response.workflow_instance_id,
+    )
+
+
+async def run_instances_batch(
+    payload: bytes,
+    *,
+    count: int = 1,
+    inputs: Optional[pb2.WorkflowArguments] = None,
+    inputs_list: Optional[list[pb2.WorkflowArguments]] = None,
+    batch_size: int = 500,
+    include_instance_ids: bool = False,
+) -> RunBatchResult:
+    """Register a workflow definition and start multiple instances over the gRPC bridge."""
+    if count < 1 and not inputs_list:
+        raise ValueError("count must be >= 1 when inputs_list is empty")
+    if batch_size < 1:
+        raise ValueError("batch_size must be >= 1")
+
+    async with ensure_singleton():
+        stub = await _workflow_stub()
+    registration = pb2.WorkflowRegistration()
+    registration.ParseFromString(payload)
+    request = pb2.RegisterWorkflowBatchRequest(
+        registration=registration,
+        count=count,
+        batch_size=batch_size,
+        include_instance_ids=include_instance_ids,
+    )
+    if inputs is not None:
+        request.inputs.CopyFrom(inputs)
+    if inputs_list:
+        request.inputs_list.extend(inputs_list)
+    try:
+        response = await stub.RegisterWorkflowBatch(request, timeout=30.0)
+    except aio.AioRpcError as exc:  # pragma: no cover
+        raise RuntimeError(f"register_workflow_batch failed: {exc}") from exc
+    return RunBatchResult(
+        workflow_version_id=response.workflow_version_id,
+        workflow_instance_ids=list(response.workflow_instance_ids),
+        queued=response.queued,
+    )
+
+
+async def execute_workflow(payload: bytes) -> bytes:
+    """Execute a workflow via the in-memory workflow streaming API."""
+    os.environ.setdefault("WAYMARK_BRIDGE_IN_MEMORY", "1")
+    async with ensure_singleton():
+        stub = await _workflow_stub()
+
+    registration = pb2.WorkflowRegistration()
+    registration.ParseFromString(payload)
+    LOGGER.debug(
+        "pytest stream start: workflow=%s ir_hash=%s",
+        registration.workflow_name,
+        registration.ir_hash,
+    )
+
+    queue: asyncio.Queue[Optional[pb2.WorkflowStreamRequest]] = asyncio.Queue()
+    skip_sleep = bool(os.environ.get("PYTEST_CURRENT_TEST"))
+    await queue.put(pb2.WorkflowStreamRequest(registration=registration, skip_sleep=skip_sleep))
+
+    async def request_stream() -> AsyncIterator[pb2.WorkflowStreamRequest]:
+        while True:
+            item = await queue.get()
+            if item is None:
+                return
+            yield item
+
+    call = stub.ExecuteWorkflow(request_stream(), timeout=300.0)
+    result_payload: Optional[bytes] = None
+
+    async for response in call:
+        kind = response.WhichOneof("kind")
+        match kind:
+            case "action_dispatch":
+                dispatch = response.action_dispatch
+                LOGGER.debug(
+                    "pytest stream dispatch: action_id=%s module=%s action=%s",
+                    dispatch.action_id,
+                    dispatch.module_name,
+                    dispatch.action_name,
+                )
+                start_ns = time.monotonic_ns()
+                execution = await execute_action(dispatch)
+                end_ns = time.monotonic_ns()
+                action_result = pb2.ActionResult(
+                    action_id=dispatch.action_id,
+                    success=execution.exception is None,
+                    payload=(
+                        serialize_result_payload(execution.result)
+                        if execution.exception is None
+                        else serialize_error_payload(dispatch.action_name, execution.exception)
+                    ),
+                    worker_start_ns=start_ns,
+                    worker_end_ns=end_ns,
+                    error_type=(
+                        type(execution.exception).__name__
+                        if execution.exception is not None
+                        else ""
+                    ),
+                    error_message=str(execution.exception)
+                    if execution.exception is not None
+                    else "",
+                )
+                # Echo the opaque server correlation metadata untouched.
+                if dispatch.metadata:
+                    action_result.metadata = dispatch.metadata
+                LOGGER.debug(
+                    "pytest stream result: action_id=%s success=%s",
+                    dispatch.action_id,
+                    execution.exception is None,
+                )
+                await queue.put(pb2.WorkflowStreamRequest(action_result=action_result))
+            case "workflow_result":
+                result_payload = response.workflow_result.payload
+                LOGGER.debug(
+                    "pytest stream complete: workflow=%s payload_bytes=%s",
+                    registration.workflow_name,
+                    len(result_payload),
+                )
+                await queue.put(None)
+                break
+            case None:
+                continue
+            case _:
+                assert_never(kind)
+
+    if result_payload is None:
+        raise RuntimeError("workflow stream ended without a result")
+    return result_payload
+
+
+async def wait_for_instance(
+    instance_id: str,
+    poll_interval_secs: float = 1.0,
+) -> Optional[bytes]:
+    """Block until the workflow daemon produces the requested instance payload."""
+    async with ensure_singleton():
+        stub = await _workflow_stub()
+    request = pb2.WaitForInstanceRequest(
+        instance_id=instance_id,
+        poll_interval_secs=poll_interval_secs,
+    )
+    try:
+        response = await stub.WaitForInstance(request, timeout=None)
+    except aio.AioRpcError as exc:  # pragma: no cover
+        status_fn = exc.code
+        if callable(status_fn) and status_fn() == grpc.StatusCode.NOT_FOUND:
+            return None
+        raise RuntimeError(f"wait_for_instance failed: {exc}") from exc
+    return bytes(response.payload)

waymark/dependencies.py

@@ -0,0 +1,36 @@
+from typing import Any, Callable
+
+from mountaineer_di import (
+    DependencyResolver,
+    Depends,
+    get_function_dependencies,
+    isolate_dependency_only_function,
+    provide_dependencies,
+    strip_depends_from_signature,
+)
+
+DependMarker = type(Depends())
+DependsMarker = DependMarker
+
+
+def Depend(  # noqa: N802
+    dependency: Callable[..., Any] | None = None,
+    *,
+    use_cache: bool = True,
+) -> Any:
+    """Compatibility alias for ``mountaineer_di.Depends``."""
+
+    return Depends(dependency, use_cache=use_cache)
+
+
+__all__ = [
+    "DependencyResolver",
+    "Depend",
+    "DependMarker",
+    "Depends",
+    "DependsMarker",
+    "get_function_dependencies",
+    "isolate_dependency_only_function",
+    "provide_dependencies",
+    "strip_depends_from_signature",
+]

waymark/exceptions.py

@@ -0,0 +1,18 @@
+"""Custom exception types raised by waymark workflows."""
+
+
+class ExhaustedRetriesError(Exception):
+    """Raised when an action exhausts its allotted retry attempts."""
+
+    def __init__(self, message: str | None = None) -> None:
+        super().__init__(message or "action exhausted retries")
+
+
+ExhaustedRetries = ExhaustedRetriesError
+
+
+class ScheduleAlreadyExistsError(Exception):
+    """Raised when a schedule name is already registered."""
+
+    def __init__(self, message: str | None = None) -> None:
+        super().__init__(message or "schedule already exists")