diff --git a/CHANGELOG.md b/CHANGELOG.md index e9dde1828..3492ce2ec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ ### Improvements - databricks-openai: `DatabricksOpenAI` and `AsyncDatabricksOpenAI` clients now follow HTTP redirects by default, configurable via the new `follow_redirects` parameter (#445) +- databricks-ai-bridge: Add a transport-neutral `DatabricksDurableRuntime` with Lakebase request, response, event persistence, and stale-attempt recovery ### Bug Fixes - databricks-ai-bridge: Genie now returns the full answer text aggregated from all text attachments (#432) diff --git a/README.md b/README.md index 615d57c52..1786cffa1 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,20 @@ For frameworks without dedicated integration packages: pip install databricks-ai-bridge ``` +## Durable Runtime + +[`DatabricksDurableRuntime`](./src/databricks_ai_bridge/durable_runtime/README.md) +adds Lakebase-backed request, result, and event persistence with heartbeat-based +stale-attempt recovery around a caller-owned async handler. The handler remains +responsible for agent sessions and checkpoints. + +```sh +pip install 'databricks-ai-bridge[memory]' +``` + +See the [transport-neutral library cookbook](./cookbooks/durable-runtime-library/README.md) +for the developer-owned FastAPI wiring. + ### Install from source With https: diff --git a/cookbooks/durable-runtime-library/README.md b/cookbooks/durable-runtime-library/README.md new file mode 100644 index 000000000..7933b44a7 --- /dev/null +++ b/cookbooks/durable-runtime-library/README.md @@ -0,0 +1,136 @@ +# Durable Runtime Library + +This cookbook shows the transport-neutral library option. + +`openai_agent.py` contains an OpenAI Agents SDK approval flow. It uses the +stable `session_id` for an `AsyncDatabricksSession` and emits SDK events into +the durable event log. `agent.py` decides how recovered attempts should invoke +the SDK by checking `context.is_recovery`. + +The developer-owned `server.py` must: + +- start and stop `DatabricksDurableRuntime`; +- map the application's request into a stable `run_id`, `session_id`, and JSON payload; +- expose submission and polling endpoints; and +- convert persisted events into the application's streaming protocol. + +## Background, streaming, and client impact + +| Capability | Library contract | Developer and client contract | +| --- | --- | --- | +| Background | `runtime.submit()` persists and schedules work; `runtime.get()` returns stored status/result. | Developer chooses how the server returns `202`, exposes the run ID, and implements polling. | +| Durable streaming | `context.emit()` stores ordered JSON events; `runtime.events()` reads them by cursor. | Developer defines the SSE/event format and reconnect route; the client must retain the chosen cursor. | + +The library does not mandate any HTTP contract. The cookbook chooses +`POST /runs`, `GET /runs/{run_id}`, and `GET /runs/{run_id}/events?after=N` to +make the missing server work visible. + +### OpenAI Agents SDK: before and after + +The OpenAI Agents SDK has no remote deployment client. A developer may preserve +an existing client exactly, but must map that route to the runtime and decide +how it expresses background and replay semantics. The cookbook instead chooses: + +```python +async with http.stream( + "POST", + "/runs", + json={ + "run_id": "run-1", + "session_id": "conversation-1", + "background": True, + "stream": True, + "payload": {"message": "hello"}, + }, +) as response: + async for line in response.aiter_lines(): + last_event_id = remember_sse_id(line, last_event_id) + +final = (await http.get("/runs/run-1")).json() +``` + +The server-side agent still calls `Runner.run_streamed()` and forwards each SDK +event to `context.emit()`. + +### LangGraph SDK: before and after + +The native client starts and observes work through LangGraph's own protocol: + +```python +thread = await langgraph.threads.create() +run = await langgraph.runs.create( + thread["thread_id"], "agent", input={"messages": messages} +) +result = await langgraph.runs.join(thread["thread_id"], run["run_id"]) + +async for event in langgraph.runs.stream( + thread["thread_id"], "agent", input={"messages": messages} +): + consume(event) +``` + +The library is flexible enough for a developer to recreate those routes and +keep `langgraph_sdk`, but the library does not provide that integration. With +the cookbook server, the client changes to `/runs` and passes `thread_id` as +`session_id`. This is the option with the least mandated client change and the +most developer-owned protocol work. + +References: [OpenAI Agents SDK streaming](https://openai.github.io/openai-agents-python/streaming/), +[LangGraph background runs](https://docs.langchain.com/langsmith/runs), and +[LangGraph resumable streaming](https://docs.langchain.com/langsmith/streaming). + +Run it locally with: + +```bash +pip install -r requirements.txt +export OPENAI_API_KEY=... +export LAKEBASE_AUTOSCALING_ENDPOINT=projects/.../endpoints/... +uvicorn server:app --reload +``` + +## Durable HITL flow + +HITL is two durable runs. The first run returns `requires_action`; the second +run carries the human decision and reuses the same SDK session. + +Start the proposal as a background stream: + +```bash +curl -N -X POST localhost:8000/runs \ + -H 'content-type: application/json' \ + -d '{ + "run_id": "proposal-1", + "session_id": "approval-session-1", + "background": true, + "stream": true, + "payload": {"action": "publish the release notes"} + }' +``` + +Poll `/runs/proposal-1`, review the persisted result, and submit approval: + +```bash +curl -N -X POST localhost:8000/runs \ + -H 'content-type: application/json' \ + -d '{ + "run_id": "approval-1", + "session_id": "approval-session-1", + "background": true, + "stream": true, + "payload": { + "action": "publish the release notes", + "decision": "approve", + "wait_seconds": 60 + } + }' +``` + +Stop the process during the wait. On restart, the library reclaims the run and +calls the same developer-owned executor with `context.is_recovery=True`. The +developer is responsible for reconnecting the SDK session and choosing the +recovery instruction. Replay from +`/runs/approval-1/events?after=` and poll `/runs/approval-1` for +the final result. + +This option makes the ownership difference explicit: the library stores runs +and events, while the developer implements all HTTP and recovery wiring. diff --git a/cookbooks/durable-runtime-library/agent.py b/cookbooks/durable-runtime-library/agent.py new file mode 100644 index 000000000..4b300238a --- /dev/null +++ b/cookbooks/durable-runtime-library/agent.py @@ -0,0 +1,24 @@ +"""OpenAI Agents SDK adapter for the transport-neutral durable runtime.""" + +from openai_agent import run_openai_agent + +from databricks_ai_bridge.durable_runtime import DurableExecutionContext, JsonObject + + +async def run_agent( + request: JsonObject, context: DurableExecutionContext +) -> JsonObject: + payload = request["payload"] + session_id = str(request["session_id"]) + result = await run_openai_agent( + payload=payload, + session_id=session_id, + emit=context.emit, + is_recovery=context.is_recovery, + ) + + return { + "result": result, + "session_id": session_id, + "attempt": context.attempt, + } diff --git a/cookbooks/durable-runtime-library/app.yaml b/cookbooks/durable-runtime-library/app.yaml new file mode 100644 index 000000000..21fc2d5e2 --- /dev/null +++ b/cookbooks/durable-runtime-library/app.yaml @@ -0,0 +1,3 @@ +command: + - python + - server.py diff --git a/cookbooks/durable-runtime-library/openai_agent.py b/cookbooks/durable-runtime-library/openai_agent.py new file mode 100644 index 000000000..bee733c3c --- /dev/null +++ b/cookbooks/durable-runtime-library/openai_agent.py @@ -0,0 +1,85 @@ +"""OpenAI Agents SDK loop for durable streaming and human approval.""" + +import asyncio +import os +from collections.abc import Awaitable, Callable +from typing import Any + +from agents import Agent, Runner, function_tool +from databricks_openai.agents import AsyncDatabricksSession + +EventEmitter = Callable[[dict[str, Any]], Awaitable[int]] + + +@function_tool +async def complete_approved_action(action: str, wait_seconds: int) -> str: + """Simulate an approved side effect after a bounded delay.""" + if wait_seconds < 0 or wait_seconds > 300: + raise ValueError("wait_seconds must be between 0 and 300") + await asyncio.sleep(wait_seconds) + return f"Completed approved action: {action}" + + +def _create_agent() -> Agent: + return Agent( + name="Durable approval assistant", + model=os.getenv("OPENAI_MODEL", "gpt-4.1-mini"), + instructions=( + "For a PROPOSAL request, produce a short plan and do not call tools. " + "For an APPROVED request, always call complete_approved_action. " + "For a REJECTED request, acknowledge the rejection without calling tools." + ), + tools=[complete_approved_action], + ) + + +async def run_openai_agent( + payload: dict[str, Any], + session_id: str, + emit: EventEmitter, + *, + is_recovery: bool = False, +) -> dict[str, Any]: + session = AsyncDatabricksSession( + session_id=session_id, + autoscaling_endpoint=os.environ["LAKEBASE_AUTOSCALING_ENDPOINT"], + schema=os.getenv( + "OPENAI_AGENT_SESSION_SCHEMA", "durable_openai_agent_sessions" + ), + ) + action = str(payload["action"]) + decision = payload.get("decision") + wait_seconds = int(payload.get("wait_seconds", 0)) + if decision is None: + instruction = f"PROPOSAL: Draft a plan for this action: {action}" + elif decision == "approve": + instruction = ( + f"APPROVED: Complete this action: {action}. " + f"Call complete_approved_action with wait_seconds={wait_seconds}." + ) + elif decision == "reject": + instruction = f"REJECTED: Do not perform this action: {action}" + else: + raise ValueError("decision must be approve, reject, or omitted") + + if is_recovery: + instruction = ( + "RECOVERY: Continue the interrupted durable run using the existing session. " + + instruction + ) + + result = Runner.run_streamed(_create_agent(), input=instruction, session=session) + + async for event in result.stream_events(): + if event.type == "raw_response_event": + await emit(event.data.model_dump(mode="json")) + + output = str(result.final_output or "") + if decision is None: + response = {"status": "requires_action", "action": action, "proposal": output} + await emit({"type": "agent.approval_required", "action": action}) + return response + + status = "completed" if decision == "approve" else "rejected" + await emit({"type": f"agent.{status}", "action": action}) + return {"status": status, "action": action, "output": output} diff --git a/cookbooks/durable-runtime-library/requirements.txt b/cookbooks/durable-runtime-library/requirements.txt new file mode 100644 index 000000000..c8fb8a711 --- /dev/null +++ b/cookbooks/durable-runtime-library/requirements.txt @@ -0,0 +1,3 @@ +databricks-ai-bridge[agent-server] +databricks-openai[memory]>=0.17.0 +openai-agents>=0.19.4,<0.20 diff --git a/cookbooks/durable-runtime-library/server.py b/cookbooks/durable-runtime-library/server.py new file mode 100644 index 000000000..f64a89251 --- /dev/null +++ b/cookbooks/durable-runtime-library/server.py @@ -0,0 +1,120 @@ +"""Developer-owned FastAPI wiring for DatabricksDurableRuntime.""" + +import asyncio +import json +import os +from contextlib import asynccontextmanager +from typing import Any + +import uvicorn +from fastapi import FastAPI, HTTPException, Response +from fastapi.responses import JSONResponse, StreamingResponse +from pydantic import BaseModel, Field + +from agent import run_agent +from databricks_ai_bridge.durable_runtime import ( + DatabricksDurableRuntime, + DurableExecution, + DurableExecutionStatus, +) + +runtime = DatabricksDurableRuntime(run_agent) + + +@asynccontextmanager +async def lifespan(_: FastAPI): + await runtime.start() + try: + yield + finally: + await runtime.stop() + + +app = FastAPI(lifespan=lifespan) + + +class RunSubmission(BaseModel): + run_id: str + session_id: str + payload: dict[str, Any] = Field(default_factory=dict) + background: bool = True + stream: bool = False + + +def _state_payload(state: DurableExecution) -> dict[str, Any]: + return { + "run_id": state.execution_id, + "status": state.status.value, + "attempt": state.attempt, + "result": state.response, + } + + +@app.post("/runs") +async def submit_run(submission: RunSubmission) -> Response: + request = {"session_id": submission.session_id, "payload": submission.payload} + if submission.stream: + await runtime.submit(submission.run_id, request) + return StreamingResponse( + _event_stream(submission.run_id, 0), + media_type="text/event-stream", + headers={"Databricks-Run-Id": submission.run_id}, + ) + + if submission.background: + state = await runtime.submit(submission.run_id, request) + status_code = 200 if state.is_terminal else 202 + return JSONResponse(_state_payload(state), status_code=status_code) + + await runtime.invoke(submission.run_id, request) + state = await runtime.get(submission.run_id) + if state is None: + raise RuntimeError(f"run {submission.run_id!r} disappeared after completion") + return JSONResponse(_state_payload(state)) + + +@app.get("/runs/{run_id}") +async def get_run(run_id: str) -> dict[str, Any]: + state = await runtime.get(run_id) + if state is None: + raise HTTPException(404, "run not found") + return _state_payload(state) + + +@app.get("/runs/{run_id}/events") +async def stream_run_events(run_id: str, after: int = 0) -> StreamingResponse: + if await runtime.get(run_id) is None: + raise HTTPException(404, "run not found") + + return StreamingResponse( + _event_stream(run_id, after), media_type="text/event-stream" + ) + + +async def _event_stream(run_id: str, after: int): + cursor = after + while True: + events = await runtime.events(run_id, after_sequence=cursor) + for event in events: + cursor = event.sequence_number + yield ( + f"id: {cursor}\n" + f"event: {event.event.get('type', 'message')}\n" + f"data: {json.dumps(event.event)}\n\n" + ) + + state = await runtime.get(run_id) + if state is None or state.status in { + DurableExecutionStatus.COMPLETED, + DurableExecutionStatus.FAILED, + }: + return + await asyncio.sleep(0.25) + + +if __name__ == "__main__": + uvicorn.run( + app, + host="0.0.0.0", + port=int(os.getenv("DATABRICKS_APP_PORT", "8000")), + ) diff --git a/src/databricks_ai_bridge/durable_runtime/README.md b/src/databricks_ai_bridge/durable_runtime/README.md new file mode 100644 index 000000000..82a839f85 --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/README.md @@ -0,0 +1,136 @@ +# Databricks Durable Runtime + +`DatabricksDurableRuntime` is a transport-neutral durability layer for an +idempotent JSON request/response handler. It stores the request, response, and +heartbeat state in Lakebase and re-executes stale work after a process exits. + +It does not own agent session history. An OpenAI Agents SDK session, LangGraph +checkpointer, or other harness-managed store remains an executor concern. + +## Contract + +The caller provides: + +- a stable `execution_id`, used as the idempotency and recovery key; +- a JSON-object request; and +- an async executor that returns a JSON-object response. + +The executor receives `DurableExecutionContext`. On recovery, `attempt > 1` +and `is_recovery` is true. The runtime passes the exact persisted request on +every attempt and does not add a recovery message or reconstruct history. + +```python +from databricks_ai_bridge.durable_runtime import ( + DatabricksDurableRuntime, + DurableExecutionContext, +) + + +async def execute(request: dict, context: DurableExecutionContext) -> dict: + session = create_agent_session(request["session_id"]) + await context.emit({"type": "progress", "message": "agent started"}) + if context.is_recovery: + result = await resume_from_session(session) + else: + result = await start_from_request(request, session) + return {"output": result} + + +runtime = DatabricksDurableRuntime( + execute, + autoscaling_endpoint="projects/project/branches/branch/endpoints/primary", +) +``` + +`execution_id` identifies one durable request, not an SDK session. A response +ID or client idempotency key is normally the right value. A session ID can be +used only when the application intentionally allows one durable request per +session. For multi-turn agents, keep the harness session or conversation ID in +the persisted request and assign each invocation its own `execution_id`. + +Call `await runtime.start()` during process startup and `await runtime.stop()` +during shutdown. Shutdown cancels local tasks without marking them failed, so a +different process can claim them after the heartbeat becomes stale. + +## Input and Output Wiring + +- `submit(execution_id, request)` accepts background work and returns persisted + state immediately. +- `invoke(execution_id, request)` accepts work and waits for the persisted + response, including when another process owns the attempt. +- `get(execution_id)` returns status, attempt, heartbeat, request, and response. +- `wait(execution_id)` waits for a previously submitted response. +- `events(execution_id, after_sequence=N)` returns persisted events for replay. + +Executors may call `await context.emit(event)` to append an ordered JSON event. +Events are stored before delivery, survive worker replacement, and use the returned +sequence number as a replay cursor. HTTP and SSE adapters decide how those generic +events are presented to clients. + +Submitting the same ID and same JSON request is idempotent. If the response is +already complete, `invoke` returns the cached response. Reusing the ID with a +different request raises `DurableRequestConflictError`. + +An HTTP adapter remains small and transport-specific: + +```python +@app.post("/responses") +async def responses(request: RequestModel): + execution_id = stable_id_from(request) + payload = executor_payload(request) + if request.background: + state = await runtime.submit(execution_id, payload) + return status_response(state, status_code=202) + return await runtime.invoke(execution_id, payload) + + +@app.get("/responses/{execution_id}") +async def retrieve(execution_id: str): + state = await runtime.get(execution_id) + return status_or_response(state) +``` + +The adapter decides how IDs are supplied, which HTTP status shape to return, +and how JSON is converted to framework-specific request or response models. +`executor_payload` should remove transport-only fields such as `background`, +`stream`, polling cursors, or trace-return flags. Changing only the transport +mode must not create a request conflict for the same durable operation. + +For background execution, an adapter may generate the ID and return it in the +initial `202` response. For a blocking request that must survive a lost HTTP +connection, the client must supply a stable ID (for example an +`Idempotency-Key`, session ID, or response ID) so it can retry or retrieve the +same operation. The runtime does not hide this client/server recovery contract. + +Recommended HTTP error mappings are: + +- different request for an existing ID: `409 Conflict`; +- unknown ID on retrieval: `404 Not Found`; +- failed execution: a terminal failed-status payload; and +- blocking wait timeout: a gateway timeout while execution continues in the + background and remains retrievable by ID. + +## Durability Store + +The default schema is `databricks_durable_runtime`. It contains two tables: + +```text +executions + execution_id TEXT PRIMARY KEY + status TEXT + attempt INTEGER + heartbeat_at TIMESTAMPTZ + request JSONB + response JSONB + +execution_events + sequence_number BIGSERIAL PRIMARY KEY + execution_id TEXT + attempt INTEGER + event JSONB +``` + +Recovery is at-least-once. A process can exit after an external side effect but +before persisting its response, so executor tools must tolerate retries where +needed. A compare-and-swap claim on the stale row prevents multiple recovery +attempts from acquiring the same durability ownership. diff --git a/src/databricks_ai_bridge/durable_runtime/__init__.py b/src/databricks_ai_bridge/durable_runtime/__init__.py new file mode 100644 index 000000000..efba50012 --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/__init__.py @@ -0,0 +1,49 @@ +"""Lakebase-backed durable execution for JSON request/response handlers. + +Install the memory extra before importing this module:: + + pip install databricks-ai-bridge[memory] +""" + +try: + import psycopg # noqa: F401 + import sqlalchemy # noqa: F401 +except ImportError as exc: + raise ImportError( + "DatabricksDurableRuntime requires databricks-ai-bridge[memory]. " + "Install it with: pip install databricks-ai-bridge[memory]" + ) from exc + +from databricks_ai_bridge.durable_runtime.runtime import DatabricksDurableRuntime +from databricks_ai_bridge.durable_runtime.store import ( + DEFAULT_DURABILITY_SCHEMA, + LakebaseDurabilityStore, +) +from databricks_ai_bridge.durable_runtime.types import ( + DurabilityStore, + DurableEvent, + DurableExecution, + DurableExecutionContext, + DurableExecutionFailedError, + DurableExecutionNotFoundError, + DurableExecutionStatus, + DurableExecutor, + DurableRequestConflictError, + JsonObject, +) + +__all__ = [ + "DEFAULT_DURABILITY_SCHEMA", + "DatabricksDurableRuntime", + "DurabilityStore", + "DurableEvent", + "DurableExecution", + "DurableExecutionContext", + "DurableExecutionFailedError", + "DurableExecutionNotFoundError", + "DurableExecutionStatus", + "DurableExecutor", + "DurableRequestConflictError", + "JsonObject", + "LakebaseDurabilityStore", +] diff --git a/src/databricks_ai_bridge/durable_runtime/runtime.py b/src/databricks_ai_bridge/durable_runtime/runtime.py new file mode 100644 index 000000000..a782a66ff --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/runtime.py @@ -0,0 +1,334 @@ +"""Transport-neutral durable request runtime.""" + +from __future__ import annotations + +import asyncio +import copy +import json +import logging +from datetime import datetime, timezone +from typing import TYPE_CHECKING + +from databricks_ai_bridge.durable_runtime.store import LakebaseDurabilityStore +from databricks_ai_bridge.durable_runtime.types import ( + DurabilityStore, + DurableEvent, + DurableExecution, + DurableExecutionContext, + DurableExecutionFailedError, + DurableExecutionNotFoundError, + DurableExecutionStatus, + DurableExecutor, + JsonObject, +) + +if TYPE_CHECKING: + from databricks.sdk import WorkspaceClient + +logger = logging.getLogger(__name__) + + +def _copy_json_object(value: JsonObject, name: str) -> JsonObject: + if not isinstance(value, dict): + raise TypeError(f"{name} must be a JSON object, got {type(value).__name__}") + try: + return json.loads(json.dumps(value, allow_nan=False)) + except (TypeError, ValueError) as exc: + raise TypeError(f"{name} must be JSON serializable") from exc + + +class DatabricksDurableRuntime: + """Execute idempotent JSON requests with Lakebase-backed crash recovery. + + The runtime persists only request execution state. The executor owns agent + sessions, checkpoints, tools, and any recovery-specific prompt or behavior. + """ + + def __init__( + self, + executor: DurableExecutor | None = None, + *, + durability_store: DurabilityStore | None = None, + autoscaling_endpoint: str | None = None, + project: str | None = None, + branch: str | None = None, + workspace_client: WorkspaceClient | None = None, + schema: str = "databricks_durable_runtime", + heartbeat_seconds: float = 3.0, + stale_seconds: float = 10.0, + scan_seconds: float = 3.0, + poll_seconds: float = 1.0, + ) -> None: + if heartbeat_seconds <= 0: + raise ValueError("heartbeat_seconds must be positive") + if stale_seconds <= heartbeat_seconds: + raise ValueError("stale_seconds must be greater than heartbeat_seconds") + if scan_seconds <= 0: + raise ValueError("scan_seconds must be positive") + if poll_seconds <= 0: + raise ValueError("poll_seconds must be positive") + + self._executor = executor + self.durability_store = ( + durability_store + if durability_store is not None + else LakebaseDurabilityStore( + autoscaling_endpoint=autoscaling_endpoint, + project=project, + branch=branch, + workspace_client=workspace_client, + schema=schema, + ) + ) + self.heartbeat_seconds = heartbeat_seconds + self.stale_seconds = stale_seconds + self.scan_seconds = scan_seconds + self.poll_seconds = poll_seconds + self._tasks: dict[str, asyncio.Task[None]] = {} + self._scanner: asyncio.Task[None] | None = None + self._started = False + + async def execute( + self, + request: JsonObject, + context: DurableExecutionContext, + ) -> JsonObject: + """Run one attempt; subclasses may override this method.""" + if self._executor is None: + raise NotImplementedError("provide an executor or override execute()") + return await self._executor(request, context) + + async def start(self) -> None: + """Initialize storage and start proactive recovery scanning.""" + if self._started: + return + await self.durability_store.initialize() + self._started = True + self._scanner = asyncio.create_task( + self._scan_loop(), + name="databricks-durable-runtime-scanner", + ) + + async def stop(self) -> None: + """Stop local work, leaving active rows recoverable by another process.""" + if not self._started: + return + tasks = list(self._tasks.values()) + if self._scanner is not None: + self._scanner.cancel() + for task in tasks: + task.cancel() + await asyncio.gather( + *tasks, + *([self._scanner] if self._scanner is not None else []), + return_exceptions=True, + ) + self._tasks.clear() + self._scanner = None + self._started = False + await self.durability_store.close() + + async def submit(self, execution_id: str, request: JsonObject) -> DurableExecution: + """Accept an idempotent request and ensure recoverable work is scheduled.""" + self._require_started() + if not execution_id: + raise ValueError("execution_id must not be empty") + state = await self.durability_store.accept( + execution_id, + _copy_json_object(request, "request"), + ) + self._ensure_scheduled(state) + return state + + async def invoke( + self, + execution_id: str, + request: JsonObject, + *, + timeout: float | None = None, + ) -> JsonObject: + """Accept a request and wait for its persisted terminal response.""" + await self.submit(execution_id, request) + return await self.wait(execution_id, timeout=timeout) + + async def get(self, execution_id: str) -> DurableExecution | None: + """Return persisted state and schedule recovery if it is currently eligible.""" + self._require_started() + state = await self.durability_store.get(execution_id) + if state is not None: + self._ensure_scheduled(state) + return state + + async def wait( + self, + execution_id: str, + *, + timeout: float | None = None, + ) -> JsonObject: + """Wait for a completed response, including work owned by another process.""" + self._require_started() + + async def poll() -> JsonObject: + while True: + state = await self.get(execution_id) + if state is None: + raise DurableExecutionNotFoundError(execution_id) + if state.status == DurableExecutionStatus.COMPLETED: + if state.response is None: + raise RuntimeError( + f"execution {execution_id!r} completed without a response" + ) + return copy.deepcopy(state.response) + if state.status == DurableExecutionStatus.FAILED: + raise DurableExecutionFailedError(execution_id) + await asyncio.sleep(self.poll_seconds) + + if timeout is None: + return await poll() + return await asyncio.wait_for(poll(), timeout=timeout) + + async def events( + self, + execution_id: str, + *, + after_sequence: int | None = None, + ) -> list[DurableEvent]: + """Return persisted events after an optional replay cursor.""" + self._require_started() + return await self.durability_store.events(execution_id, after_sequence) + + def _ensure_scheduled(self, state: DurableExecution) -> None: + if not self._is_recoverable(state): + return + current = self._tasks.get(state.execution_id) + if current is not None and not current.done(): + return + task = asyncio.create_task( + self._execute_attempt(state.execution_id), + name=f"durable-execution-{state.execution_id}", + ) + self._tasks[state.execution_id] = task + task.add_done_callback(lambda completed: self._discard_task(state.execution_id, completed)) + + def _is_recoverable(self, state: DurableExecution) -> bool: + if state.status == DurableExecutionStatus.QUEUED: + return True + if state.status != DurableExecutionStatus.ACTIVE: + return False + if state.heartbeat_at is None: + return True + heartbeat_at = state.heartbeat_at + if heartbeat_at.tzinfo is None: + heartbeat_at = heartbeat_at.replace(tzinfo=timezone.utc) + age = (datetime.now(timezone.utc) - heartbeat_at).total_seconds() + return age >= self.stale_seconds + + async def _execute_attempt(self, execution_id: str) -> None: + try: + claimed = await self.durability_store.claim(execution_id, self.stale_seconds) + except Exception: + logger.exception("Failed to claim durable execution: %s", execution_id) + return + if claimed is None: + return + + heartbeat = asyncio.create_task( + self._heartbeat_loop(execution_id, claimed.attempt), + name=f"durable-heartbeat-{execution_id}-{claimed.attempt}", + ) + try: + + async def emit(event: JsonObject) -> int: + sequence_number = await self.durability_store.append_event( + execution_id, + claimed.attempt, + _copy_json_object(event, "event"), + ) + if sequence_number is None: + raise RuntimeError( + f"execution {execution_id!r} no longer owns attempt {claimed.attempt}" + ) + return sequence_number + + response = await self.execute( + copy.deepcopy(claimed.request), + DurableExecutionContext( + execution_id=execution_id, + attempt=claimed.attempt, + _emit=emit, + ), + ) + response = _copy_json_object(response, "executor response") + completed = await self.durability_store.complete( + execution_id, + claimed.attempt, + response, + ) + if not completed: + logger.info( + "Skipped completion after durability ownership changed: %s attempt=%d", + execution_id, + claimed.attempt, + ) + except asyncio.CancelledError: + raise + except Exception: + logger.exception( + "Durable execution failed: %s attempt=%d", + execution_id, + claimed.attempt, + ) + try: + await self.durability_store.fail(execution_id, claimed.attempt) + except Exception: + logger.exception( + "Failed to persist durable failure: %s attempt=%d", + execution_id, + claimed.attempt, + ) + finally: + heartbeat.cancel() + await asyncio.gather(heartbeat, return_exceptions=True) + + async def _heartbeat_loop(self, execution_id: str, attempt: int) -> None: + while True: + try: + owns_attempt = await self.durability_store.heartbeat(execution_id, attempt) + except Exception: + logger.warning( + "Durable heartbeat failed: %s attempt=%d", + execution_id, + attempt, + exc_info=True, + ) + await asyncio.sleep(self.heartbeat_seconds) + continue + if not owns_attempt: + return + await asyncio.sleep(self.heartbeat_seconds) + + async def _scan_loop(self) -> None: + while True: + try: + execution_ids = await self.durability_store.recoverable_execution_ids( + self.stale_seconds + ) + for execution_id in execution_ids: + state = await self.durability_store.get(execution_id) + if state is not None: + self._ensure_scheduled(state) + except asyncio.CancelledError: + raise + except Exception: + logger.exception("Databricks durable runtime recovery scan failed") + await asyncio.sleep(self.scan_seconds) + + def _discard_task(self, execution_id: str, completed: asyncio.Task[None]) -> None: + if self._tasks.get(execution_id) is completed: + self._tasks.pop(execution_id, None) + if not completed.cancelled(): + completed.exception() + + def _require_started(self) -> None: + if not self._started: + raise RuntimeError("DatabricksDurableRuntime.start() must be called first") diff --git a/src/databricks_ai_bridge/durable_runtime/store.py b/src/databricks_ai_bridge/durable_runtime/store.py new file mode 100644 index 000000000..5969904b9 --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/store.py @@ -0,0 +1,394 @@ +"""Lakebase persistence for durable request execution.""" + +from __future__ import annotations + +import json +import os +import re +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any + +from sqlalchemy import text + +from databricks_ai_bridge.durable_runtime.types import ( + DurableEvent, + DurableExecution, + DurableExecutionStatus, + DurableRequestConflictError, + JsonObject, +) +from databricks_ai_bridge.lakebase import AsyncLakebaseSQLAlchemy + +if TYPE_CHECKING: + from databricks.sdk import WorkspaceClient + +DEFAULT_DURABILITY_SCHEMA = "databricks_durable_runtime" +_SCHEMA_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def _serialize_json_object(value: JsonObject) -> str: + if not isinstance(value, dict): + raise TypeError(f"expected a JSON object, got {type(value).__name__}") + return json.dumps(value, allow_nan=False) + + +def _validate_execution_id(execution_id: str) -> None: + if not execution_id: + raise ValueError("execution_id must not be empty") + + +class LakebaseDurabilityStore: + """Store durable execution state in one Lakebase table.""" + + def __init__( + self, + *, + autoscaling_endpoint: str | None = None, + project: str | None = None, + branch: str | None = None, + workspace_client: WorkspaceClient | None = None, + schema: str = DEFAULT_DURABILITY_SCHEMA, + lakebase: AsyncLakebaseSQLAlchemy | None = None, + ) -> None: + if not _SCHEMA_NAME.fullmatch(schema): + raise ValueError(f"invalid durability schema name: {schema!r}") + + if lakebase is None: + autoscaling_endpoint = autoscaling_endpoint or os.getenv( + "LAKEBASE_AUTOSCALING_ENDPOINT" + ) + if autoscaling_endpoint is None: + project = project or os.getenv("LAKEBASE_AUTOSCALING_PROJECT") + branch = branch or os.getenv("LAKEBASE_AUTOSCALING_BRANCH") + lakebase = AsyncLakebaseSQLAlchemy( + autoscaling_endpoint=autoscaling_endpoint, + project=project, + branch=branch, + workspace_client=workspace_client, + schema=schema, + pool_pre_ping=True, + ) + + self._lakebase = lakebase + self._engine = lakebase.engine + self._table = f"{schema}.executions" + self._events_table = f"{schema}.execution_events" + + async def initialize(self) -> None: + await self._lakebase.create_schema() + async with self._engine.begin() as connection: + await connection.execute( + text( + f""" + CREATE TABLE IF NOT EXISTS {self._table} ( + execution_id TEXT PRIMARY KEY, + status TEXT NOT NULL, + attempt INTEGER NOT NULL DEFAULT 0, + heartbeat_at TIMESTAMPTZ, + request JSONB NOT NULL, + response JSONB, + CHECK (status IN ('QUEUED', 'ACTIVE', 'COMPLETED', 'FAILED')), + CHECK (jsonb_typeof(request) = 'object'), + CHECK (response IS NULL OR jsonb_typeof(response) = 'object') + ) + """ + ) + ) + await connection.execute( + text( + f""" + CREATE TABLE IF NOT EXISTS {self._events_table} ( + sequence_number BIGSERIAL PRIMARY KEY, + execution_id TEXT NOT NULL + REFERENCES {self._table}(execution_id) ON DELETE CASCADE, + attempt INTEGER NOT NULL, + event JSONB NOT NULL, + CHECK (jsonb_typeof(event) = 'object') + ) + """ + ) + ) + await connection.execute( + text( + f""" + CREATE INDEX IF NOT EXISTS execution_events_replay_idx + ON {self._events_table} (execution_id, sequence_number) + """ + ) + ) + await connection.execute( + text( + f""" + CREATE INDEX IF NOT EXISTS executions_recovery_idx + ON {self._table} (status, heartbeat_at) + WHERE status IN ('QUEUED', 'ACTIVE') + """ + ) + ) + + async def close(self) -> None: + await self._engine.dispose() + + async def accept(self, execution_id: str, request: JsonObject) -> DurableExecution: + _validate_execution_id(execution_id) + serialized_request = _serialize_json_object(request) + async with self._engine.begin() as connection: + await connection.execute( + text( + f""" + INSERT INTO {self._table} (execution_id, status, request) + VALUES (:execution_id, 'QUEUED', CAST(:request AS JSONB)) + ON CONFLICT (execution_id) DO NOTHING + """ + ), + {"execution_id": execution_id, "request": serialized_request}, + ) + row = ( + ( + await connection.execute( + text( + f""" + SELECT execution_id, status, attempt, heartbeat_at, + request::TEXT AS request_json, + response::TEXT AS response_json + FROM {self._table} + WHERE execution_id=:execution_id + """ + ), + {"execution_id": execution_id}, + ) + ) + .mappings() + .one() + ) + + state = self._to_execution(row) + if state.request != request: + raise DurableRequestConflictError( + f"execution {execution_id!r} was already accepted with a different request" + ) + return state + + async def get(self, execution_id: str) -> DurableExecution | None: + _validate_execution_id(execution_id) + async with self._engine.connect() as connection: + row = ( + ( + await connection.execute( + text( + f""" + SELECT execution_id, status, attempt, heartbeat_at, + request::TEXT AS request_json, + response::TEXT AS response_json + FROM {self._table} + WHERE execution_id=:execution_id + """ + ), + {"execution_id": execution_id}, + ) + ) + .mappings() + .one_or_none() + ) + return self._to_execution(row) if row is not None else None + + async def recoverable_execution_ids(self, stale_seconds: float) -> list[str]: + async with self._engine.connect() as connection: + rows = ( + ( + await connection.execute( + text( + f""" + SELECT execution_id + FROM {self._table} + WHERE status='QUEUED' + OR (status='ACTIVE' AND ( + heartbeat_at IS NULL + OR heartbeat_at < NOW() - (:stale * INTERVAL '1 second') + )) + ORDER BY heartbeat_at NULLS FIRST + """ + ), + {"stale": stale_seconds}, + ) + ) + .scalars() + .all() + ) + return list(rows) + + async def claim( + self, + execution_id: str, + stale_seconds: float, + ) -> DurableExecution | None: + _validate_execution_id(execution_id) + async with self._engine.begin() as connection: + row = ( + ( + await connection.execute( + text( + f""" + UPDATE {self._table} + SET status='ACTIVE', attempt=attempt+1, heartbeat_at=NOW() + WHERE execution_id=:execution_id + AND ( + status='QUEUED' + OR (status='ACTIVE' AND ( + heartbeat_at IS NULL + OR heartbeat_at < NOW() - (:stale * INTERVAL '1 second') + )) + ) + RETURNING execution_id, status, attempt, heartbeat_at, + request::TEXT AS request_json, + response::TEXT AS response_json + """ + ), + {"execution_id": execution_id, "stale": stale_seconds}, + ) + ) + .mappings() + .one_or_none() + ) + return self._to_execution(row) if row is not None else None + + async def heartbeat(self, execution_id: str, attempt: int) -> bool: + _validate_execution_id(execution_id) + async with self._engine.begin() as connection: + result = await connection.execute( + text( + f""" + UPDATE {self._table} + SET heartbeat_at=NOW() + WHERE execution_id=:execution_id + AND attempt=:attempt + AND status='ACTIVE' + """ + ), + {"execution_id": execution_id, "attempt": attempt}, + ) + return result.rowcount == 1 + + async def complete( + self, + execution_id: str, + attempt: int, + response: JsonObject, + ) -> bool: + _validate_execution_id(execution_id) + serialized_response = _serialize_json_object(response) + async with self._engine.begin() as connection: + result = await connection.execute( + text( + f""" + UPDATE {self._table} + SET status='COMPLETED', response=CAST(:response AS JSONB) + WHERE execution_id=:execution_id + AND attempt=:attempt + AND status='ACTIVE' + """ + ), + { + "execution_id": execution_id, + "attempt": attempt, + "response": serialized_response, + }, + ) + return result.rowcount == 1 + + async def fail(self, execution_id: str, attempt: int) -> bool: + _validate_execution_id(execution_id) + async with self._engine.begin() as connection: + result = await connection.execute( + text( + f""" + UPDATE {self._table} + SET status='FAILED' + WHERE execution_id=:execution_id + AND attempt=:attempt + AND status='ACTIVE' + """ + ), + {"execution_id": execution_id, "attempt": attempt}, + ) + return result.rowcount == 1 + + async def append_event( + self, + execution_id: str, + attempt: int, + event: JsonObject, + ) -> int | None: + """Append an event only while the caller owns the active attempt.""" + _validate_execution_id(execution_id) + serialized_event = _serialize_json_object(event) + async with self._engine.begin() as connection: + result = await connection.execute( + text( + f""" + INSERT INTO {self._events_table} (execution_id, attempt, event) + SELECT :execution_id, :attempt, CAST(:event AS JSONB) + WHERE EXISTS ( + SELECT 1 + FROM {self._table} + WHERE execution_id=:execution_id + AND attempt=:attempt + AND status='ACTIVE' + ) + RETURNING sequence_number + """ + ), + { + "execution_id": execution_id, + "attempt": attempt, + "event": serialized_event, + }, + ) + sequence_number = result.scalar_one_or_none() + return int(sequence_number) if sequence_number is not None else None + + async def events( + self, + execution_id: str, + after_sequence: int | None = None, + ) -> list[DurableEvent]: + """Return ordered events for one execution after an optional cursor.""" + _validate_execution_id(execution_id) + async with self._engine.connect() as connection: + result = await connection.execute( + text( + f""" + SELECT sequence_number, execution_id, attempt, + event::TEXT AS event_json + FROM {self._events_table} + WHERE execution_id=:execution_id + AND (:after_sequence IS NULL OR sequence_number > :after_sequence) + ORDER BY sequence_number + """ + ), + { + "execution_id": execution_id, + "after_sequence": after_sequence, + }, + ) + rows = result.mappings().all() + return [ + DurableEvent( + sequence_number=int(row["sequence_number"]), + execution_id=str(row["execution_id"]), + attempt=int(row["attempt"]), + event=json.loads(row["event_json"]), + ) + for row in rows + ] + + @staticmethod + def _to_execution(row: Mapping[str, Any]) -> DurableExecution: + return DurableExecution( + execution_id=str(row["execution_id"]), + status=DurableExecutionStatus(str(row["status"])), + attempt=int(row["attempt"]), + heartbeat_at=row["heartbeat_at"], + request=json.loads(row["request_json"]), + response=json.loads(row["response_json"]) if row["response_json"] else None, + ) diff --git a/src/databricks_ai_bridge/durable_runtime/types.py b/src/databricks_ai_bridge/durable_runtime/types.py new file mode 100644 index 000000000..a03b73f64 --- /dev/null +++ b/src/databricks_ai_bridge/durable_runtime/types.py @@ -0,0 +1,126 @@ +"""Public types for durable request execution.""" + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from datetime import datetime +from enum import Enum +from typing import Any, Protocol + +JsonObject = dict[str, Any] +DurableEventEmitter = Callable[[JsonObject], Awaitable[int]] + + +class DurableExecutionStatus(str, Enum): + """Lifecycle states stored by the durability layer.""" + + QUEUED = "QUEUED" + ACTIVE = "ACTIVE" + COMPLETED = "COMPLETED" + FAILED = "FAILED" + + +@dataclass(frozen=True) +class DurableExecution: + """Persisted state for one idempotent request.""" + + execution_id: str + status: DurableExecutionStatus + attempt: int + heartbeat_at: datetime | None + request: JsonObject + response: JsonObject | None + + @property + def is_terminal(self) -> bool: + return self.status in { + DurableExecutionStatus.COMPLETED, + DurableExecutionStatus.FAILED, + } + + +@dataclass(frozen=True) +class DurableEvent: + """One persisted event emitted by a durable execution attempt.""" + + sequence_number: int + execution_id: str + attempt: int + event: JsonObject + + +@dataclass(frozen=True) +class DurableExecutionContext: + """Attempt metadata passed to the caller-owned executor.""" + + execution_id: str + attempt: int + _emit: DurableEventEmitter | None = field(default=None, repr=False, compare=False) + + @property + def is_recovery(self) -> bool: + return self.attempt > 1 + + async def emit(self, event: JsonObject) -> int: + """Persist an ordered event and return its replay cursor.""" + if self._emit is None: + raise RuntimeError("event emission is not available for this execution context") + return await self._emit(event) + + +DurableExecutor = Callable[[JsonObject, DurableExecutionContext], Awaitable[JsonObject]] + + +class DurabilityStore(Protocol): + """Persistence contract used by :class:`DatabricksDurableRuntime`.""" + + async def initialize(self) -> None: ... + + async def close(self) -> None: ... + + async def accept(self, execution_id: str, request: JsonObject) -> DurableExecution: ... + + async def get(self, execution_id: str) -> DurableExecution | None: ... + + async def recoverable_execution_ids(self, stale_seconds: float) -> list[str]: ... + + async def claim( + self, + execution_id: str, + stale_seconds: float, + ) -> DurableExecution | None: ... + + async def heartbeat(self, execution_id: str, attempt: int) -> bool: ... + + async def complete( + self, + execution_id: str, + attempt: int, + response: JsonObject, + ) -> bool: ... + + async def fail(self, execution_id: str, attempt: int) -> bool: ... + + async def append_event( + self, + execution_id: str, + attempt: int, + event: JsonObject, + ) -> int | None: ... + + async def events( + self, + execution_id: str, + after_sequence: int | None = None, + ) -> list[DurableEvent]: ... + + +class DurableRequestConflictError(ValueError): + """Raised when an execution ID is reused with a different request.""" + + +class DurableExecutionNotFoundError(LookupError): + """Raised when waiting for an unknown execution ID.""" + + +class DurableExecutionFailedError(RuntimeError): + """Raised when a durable execution reaches the failed state.""" diff --git a/tests/databricks_ai_bridge/test_durable_runtime.py b/tests/databricks_ai_bridge/test_durable_runtime.py new file mode 100644 index 000000000..debfe1614 --- /dev/null +++ b/tests/databricks_ai_bridge/test_durable_runtime.py @@ -0,0 +1,456 @@ +"""Tests for DatabricksDurableRuntime orchestration.""" + +import asyncio +import copy +from datetime import datetime, timedelta, timezone + +import pytest + +pytest.importorskip("psycopg") +pytest.importorskip("sqlalchemy") + +from databricks_ai_bridge.durable_runtime import ( + DatabricksDurableRuntime, + DurableEvent, + DurableExecution, + DurableExecutionContext, + DurableExecutionFailedError, + DurableExecutionStatus, + DurableRequestConflictError, +) + + +class MemoryDurabilityStore: + def __init__(self) -> None: + self.states: dict[str, DurableExecution] = {} + self.initialized = False + self.closed = False + self.heartbeats: list[tuple[str, int]] = [] + self.persisted_events: list[DurableEvent] = [] + + async def initialize(self) -> None: + self.initialized = True + + async def close(self) -> None: + self.closed = True + + async def accept(self, execution_id: str, request: dict) -> DurableExecution: + existing = self.states.get(execution_id) + if existing is not None: + if existing.request != request: + raise DurableRequestConflictError(execution_id) + return existing + state = DurableExecution( + execution_id=execution_id, + status=DurableExecutionStatus.QUEUED, + attempt=0, + heartbeat_at=None, + request=copy.deepcopy(request), + response=None, + ) + self.states[execution_id] = state + return state + + async def get(self, execution_id: str) -> DurableExecution | None: + return self.states.get(execution_id) + + async def recoverable_execution_ids(self, stale_seconds: float) -> list[str]: + return [ + execution_id + for execution_id, state in self.states.items() + if state.status == DurableExecutionStatus.QUEUED + or ( + state.status == DurableExecutionStatus.ACTIVE + and ( + state.heartbeat_at is None + or datetime.now(timezone.utc) - state.heartbeat_at + >= timedelta(seconds=stale_seconds) + ) + ) + ] + + async def claim( + self, + execution_id: str, + stale_seconds: float, + ) -> DurableExecution | None: + state = self.states[execution_id] + recoverable = state.status == DurableExecutionStatus.QUEUED or ( + state.status == DurableExecutionStatus.ACTIVE + and ( + state.heartbeat_at is None + or datetime.now(timezone.utc) - state.heartbeat_at + >= timedelta(seconds=stale_seconds) + ) + ) + if not recoverable: + return None + state = DurableExecution( + execution_id=state.execution_id, + status=DurableExecutionStatus.ACTIVE, + attempt=state.attempt + 1, + heartbeat_at=datetime.now(timezone.utc), + request=state.request, + response=None, + ) + self.states[execution_id] = state + return state + + async def heartbeat(self, execution_id: str, attempt: int) -> bool: + state = self.states[execution_id] + if state.status != DurableExecutionStatus.ACTIVE or state.attempt != attempt: + return False + self.heartbeats.append((execution_id, attempt)) + self.states[execution_id] = DurableExecution( + execution_id=state.execution_id, + status=state.status, + attempt=state.attempt, + heartbeat_at=datetime.now(timezone.utc), + request=state.request, + response=state.response, + ) + return True + + async def complete(self, execution_id: str, attempt: int, response: dict) -> bool: + state = self.states[execution_id] + if state.status != DurableExecutionStatus.ACTIVE or state.attempt != attempt: + return False + self.states[execution_id] = DurableExecution( + execution_id=state.execution_id, + status=DurableExecutionStatus.COMPLETED, + attempt=state.attempt, + heartbeat_at=state.heartbeat_at, + request=state.request, + response=copy.deepcopy(response), + ) + return True + + async def fail(self, execution_id: str, attempt: int) -> bool: + state = self.states[execution_id] + if state.status != DurableExecutionStatus.ACTIVE or state.attempt != attempt: + return False + self.states[execution_id] = DurableExecution( + execution_id=state.execution_id, + status=DurableExecutionStatus.FAILED, + attempt=state.attempt, + heartbeat_at=state.heartbeat_at, + request=state.request, + response=None, + ) + return True + + async def append_event( + self, + execution_id: str, + attempt: int, + event: dict, + ) -> int | None: + state = self.states[execution_id] + if state.status != DurableExecutionStatus.ACTIVE or state.attempt != attempt: + return None + persisted = DurableEvent( + sequence_number=len(self.persisted_events) + 1, + execution_id=execution_id, + attempt=attempt, + event=copy.deepcopy(event), + ) + self.persisted_events.append(persisted) + return persisted.sequence_number + + async def events( + self, + execution_id: str, + after_sequence: int | None = None, + ) -> list[DurableEvent]: + return [ + copy.deepcopy(event) + for event in self.persisted_events + if event.execution_id == execution_id + and (after_sequence is None or event.sequence_number > after_sequence) + ] + + +def make_runtime(executor, store=None, **kwargs): + return DatabricksDurableRuntime( + executor, + durability_store=store or MemoryDurabilityStore(), + heartbeat_seconds=0.01, + stale_seconds=0.05, + scan_seconds=0.01, + poll_seconds=0.005, + **kwargs, + ) + + +@pytest.mark.asyncio +async def test_invoke_persists_request_and_response(): + calls = [] + + async def execute(request: dict, context: DurableExecutionContext) -> dict: + calls.append((request, context)) + return {"output": request["input"]} + + store = MemoryDurabilityStore() + runtime = make_runtime(execute, store) + await runtime.start() + try: + response = await runtime.invoke("session-1", {"input": "hello"}) + state = await runtime.get("session-1") + finally: + await runtime.stop() + + assert response == {"output": "hello"} + assert state is not None + assert state.request == {"input": "hello"} + assert state.response == {"output": "hello"} + assert calls[0][1].attempt == 1 + assert calls[0][1].is_recovery is False + + +@pytest.mark.asyncio +async def test_completed_request_returns_cached_response_without_reexecution(): + call_count = 0 + + async def execute(request: dict, context: DurableExecutionContext) -> dict: + nonlocal call_count + call_count += 1 + return {"output": "done"} + + runtime = make_runtime(execute) + await runtime.start() + try: + first = await runtime.invoke("session-1", {"input": "hello"}) + second = await runtime.invoke("session-1", {"input": "hello"}) + finally: + await runtime.stop() + + assert first == second == {"output": "done"} + assert call_count == 1 + + +@pytest.mark.asyncio +async def test_same_id_with_different_request_is_rejected(): + async def execute(request: dict, context: DurableExecutionContext) -> dict: + return {"output": "done"} + + runtime = make_runtime(execute) + await runtime.start() + try: + await runtime.invoke("session-1", {"input": "first"}) + with pytest.raises(DurableRequestConflictError): + await runtime.submit("session-1", {"input": "second"}) + finally: + await runtime.stop() + + +@pytest.mark.asyncio +async def test_stale_attempt_reuses_request_and_marks_recovery(): + contexts = [] + + async def execute(request: dict, context: DurableExecutionContext) -> dict: + contexts.append((request, context)) + return {"output": "recovered"} + + store = MemoryDurabilityStore() + store.states["session-1"] = DurableExecution( + execution_id="session-1", + status=DurableExecutionStatus.ACTIVE, + attempt=1, + heartbeat_at=datetime.now(timezone.utc) - timedelta(minutes=1), + request={"input": "original"}, + response=None, + ) + runtime = make_runtime(execute, store) + await runtime.start() + try: + response = await runtime.wait("session-1") + finally: + await runtime.stop() + + assert response == {"output": "recovered"} + assert contexts == [ + ( + {"input": "original"}, + DurableExecutionContext(execution_id="session-1", attempt=2), + ) + ] + assert contexts[0][1].is_recovery is True + + +@pytest.mark.asyncio +async def test_executor_failure_is_persisted_as_terminal_state(): + async def execute(request: dict, context: DurableExecutionContext) -> dict: + raise RuntimeError("boom") + + runtime = make_runtime(execute) + await runtime.start() + try: + with pytest.raises(DurableExecutionFailedError): + await runtime.invoke("session-1", {"input": "hello"}) + state = await runtime.get("session-1") + finally: + await runtime.stop() + + assert state is not None + assert state.status == DurableExecutionStatus.FAILED + assert state.response is None + + +@pytest.mark.asyncio +async def test_submit_returns_before_background_execution_finishes(): + release = asyncio.Event() + + async def execute(request: dict, context: DurableExecutionContext) -> dict: + await release.wait() + return {"output": "done"} + + runtime = make_runtime(execute) + await runtime.start() + try: + state = await runtime.submit("session-1", {"input": "hello"}) + assert state.status == DurableExecutionStatus.QUEUED + release.set() + assert await runtime.wait("session-1") == {"output": "done"} + finally: + await runtime.stop() + + +@pytest.mark.asyncio +async def test_wait_observes_response_completed_by_another_process(): + executor_called = False + + async def execute(request: dict, context: DurableExecutionContext) -> dict: + nonlocal executor_called + executor_called = True + return {} + + store = MemoryDurabilityStore() + store.states["session-1"] = DurableExecution( + execution_id="session-1", + status=DurableExecutionStatus.ACTIVE, + attempt=1, + heartbeat_at=datetime.now(timezone.utc), + request={"input": "hello"}, + response=None, + ) + runtime = make_runtime(execute, store) + await runtime.start() + + async def complete_elsewhere(): + await asyncio.sleep(0.02) + await store.complete("session-1", 1, {"output": "remote"}) + + completion = asyncio.create_task(complete_elsewhere()) + try: + assert await runtime.wait("session-1") == {"output": "remote"} + finally: + await completion + await runtime.stop() + + assert executor_called is False + + +@pytest.mark.asyncio +async def test_blocking_timeout_does_not_cancel_execution(): + release = asyncio.Event() + + async def execute(request: dict, context: DurableExecutionContext) -> dict: + await release.wait() + return {"output": "done"} + + runtime = make_runtime(execute) + await runtime.start() + try: + with pytest.raises(TimeoutError): + await runtime.invoke("session-1", {"input": "hello"}, timeout=0.02) + state = await runtime.get("session-1") + assert state is not None + assert state.status == DurableExecutionStatus.ACTIVE + release.set() + assert await runtime.wait("session-1") == {"output": "done"} + finally: + await runtime.stop() + + +@pytest.mark.asyncio +async def test_executor_can_persist_replayable_events(): + async def execute(request: dict, context: DurableExecutionContext) -> dict: + sequence_number = await context.emit({"type": "progress", "step": 1}) + return {"last_sequence_number": sequence_number} + + runtime = make_runtime(execute) + await runtime.start() + try: + assert await runtime.invoke("session-1", {"input": "hello"}) == {"last_sequence_number": 1} + events = await runtime.events("session-1") + assert [(event.sequence_number, event.event) for event in events] == [ + (1, {"type": "progress", "step": 1}) + ] + assert await runtime.events("session-1", after_sequence=1) == [] + finally: + await runtime.stop() + + +@pytest.mark.asyncio +async def test_request_must_be_json_object(): + async def execute(request: dict, context: DurableExecutionContext) -> dict: + return {} + + runtime = make_runtime(execute) + await runtime.start() + try: + with pytest.raises(TypeError, match="request must be a JSON object"): + await runtime.submit("session-1", ["not", "an", "object"]) + finally: + await runtime.stop() + + +@pytest.mark.asyncio +async def test_subclass_can_own_execution_wiring(): + class Runtime(DatabricksDurableRuntime): + async def execute( + self, + request: dict, + context: DurableExecutionContext, + ) -> dict: + return {"attempt": context.attempt, "input": request["input"]} + + runtime = Runtime( + durability_store=MemoryDurabilityStore(), + heartbeat_seconds=0.01, + stale_seconds=0.05, + scan_seconds=0.01, + poll_seconds=0.005, + ) + await runtime.start() + try: + assert await runtime.invoke("session-1", {"input": "hello"}) == { + "attempt": 1, + "input": "hello", + } + finally: + await runtime.stop() + + +@pytest.mark.asyncio +async def test_start_and_stop_manage_store_lifecycle(): + async def execute(request: dict, context: DurableExecutionContext) -> dict: + return {} + + store = MemoryDurabilityStore() + runtime = make_runtime(execute, store) + + await runtime.start() + assert store.initialized is True + await runtime.stop() + assert store.closed is True + + +@pytest.mark.asyncio +async def test_runtime_requires_start_before_use(): + async def execute(request: dict, context: DurableExecutionContext) -> dict: + return {} + + runtime = make_runtime(execute) + with pytest.raises(RuntimeError, match=r"start\(\)"): + await runtime.submit("session-1", {}) diff --git a/tests/databricks_ai_bridge/test_durable_runtime_store.py b/tests/databricks_ai_bridge/test_durable_runtime_store.py new file mode 100644 index 000000000..174340e3a --- /dev/null +++ b/tests/databricks_ai_bridge/test_durable_runtime_store.py @@ -0,0 +1,207 @@ +"""Tests for the Lakebase durability store.""" + +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +pytest.importorskip("psycopg") +pytest.importorskip("sqlalchemy") + +from databricks_ai_bridge.durable_runtime import ( + DurableExecutionStatus, + DurableRequestConflictError, + LakebaseDurabilityStore, +) + + +def mock_lakebase(): + connection = AsyncMock() + engine = MagicMock() + + @asynccontextmanager + async def begin(): + yield connection + + @asynccontextmanager + async def connect(): + yield connection + + engine.begin = begin + engine.connect = connect + engine.dispose = AsyncMock() + lakebase = MagicMock(engine=engine) + lakebase.create_schema = AsyncMock() + return lakebase, connection + + +def mapping_result(value): + result = MagicMock() + result.mappings.return_value.one.return_value = value + result.mappings.return_value.one_or_none.return_value = value + result.mappings.return_value.all.return_value = value if isinstance(value, list) else [value] + return result + + +def execution_row(**overrides): + row = { + "execution_id": "session-1", + "status": "QUEUED", + "attempt": 0, + "heartbeat_at": None, + "request_json": '{"input": "hello"}', + "response_json": None, + } + row.update(overrides) + return row + + +@pytest.mark.asyncio +async def test_initialize_creates_execution_and_event_tables(): + lakebase, connection = mock_lakebase() + store = LakebaseDurabilityStore(lakebase=lakebase) + + await store.initialize() + + sql = " ".join(str(call.args[0]) for call in connection.execute.await_args_list) + assert "databricks_durable_runtime.executions" in sql + assert "execution_id TEXT PRIMARY KEY" in sql + assert "request JSONB NOT NULL" in sql + assert "response JSONB" in sql + assert "databricks_durable_runtime.execution_events" in sql + assert "sequence_number BIGSERIAL PRIMARY KEY" in sql + lakebase.create_schema.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_accept_returns_existing_request_when_it_matches(): + lakebase, connection = mock_lakebase() + connection.execute.side_effect = [MagicMock(), mapping_result(execution_row())] + store = LakebaseDurabilityStore(lakebase=lakebase) + + state = await store.accept("session-1", {"input": "hello"}) + + assert state.execution_id == "session-1" + assert state.status == DurableExecutionStatus.QUEUED + assert state.request == {"input": "hello"} + + +@pytest.mark.asyncio +async def test_accept_rejects_same_id_with_different_request(): + lakebase, connection = mock_lakebase() + connection.execute.side_effect = [MagicMock(), mapping_result(execution_row())] + store = LakebaseDurabilityStore(lakebase=lakebase) + + with pytest.raises(DurableRequestConflictError): + await store.accept("session-1", {"input": "different"}) + + +@pytest.mark.asyncio +async def test_claim_returns_request_and_incremented_attempt(): + lakebase, connection = mock_lakebase() + heartbeat = datetime.now(timezone.utc) + connection.execute.return_value = mapping_result( + execution_row(status="ACTIVE", attempt=2, heartbeat_at=heartbeat) + ) + store = LakebaseDurabilityStore(lakebase=lakebase) + + state = await store.claim("session-1", 10) + + assert state is not None + assert state.attempt == 2 + assert state.heartbeat_at == heartbeat + assert state.request == {"input": "hello"} + + +@pytest.mark.asyncio +async def test_get_decodes_cached_response(): + lakebase, connection = mock_lakebase() + connection.execute.return_value = mapping_result( + execution_row( + status="COMPLETED", + attempt=1, + response_json='{"output": "done"}', + ) + ) + store = LakebaseDurabilityStore(lakebase=lakebase) + + state = await store.get("session-1") + + assert state is not None + assert state.status == DurableExecutionStatus.COMPLETED + assert state.response == {"output": "done"} + + +@pytest.mark.asyncio +async def test_complete_persists_response_for_owned_attempt(): + lakebase, connection = mock_lakebase() + connection.execute.return_value = MagicMock(rowcount=1) + store = LakebaseDurabilityStore(lakebase=lakebase) + + assert await store.complete("session-1", 2, {"output": "done"}) is True + + parameters = connection.execute.await_args.args[1] + assert parameters["execution_id"] == "session-1" + assert parameters["attempt"] == 2 + assert parameters["response"] == '{"output": "done"}' + + +@pytest.mark.asyncio +async def test_append_event_returns_replay_cursor_for_owned_attempt(): + lakebase, connection = mock_lakebase() + result = MagicMock() + result.scalar_one_or_none.return_value = 7 + connection.execute.return_value = result + store = LakebaseDurabilityStore(lakebase=lakebase) + + sequence_number = await store.append_event( + "session-1", + 2, + {"type": "progress", "step": 1}, + ) + + assert sequence_number == 7 + parameters = connection.execute.await_args.args[1] + assert parameters == { + "execution_id": "session-1", + "attempt": 2, + "event": '{"type": "progress", "step": 1}', + } + + +@pytest.mark.asyncio +async def test_events_returns_ordered_replay_data(): + lakebase, connection = mock_lakebase() + connection.execute.return_value = mapping_result( + [ + { + "sequence_number": 8, + "execution_id": "session-1", + "attempt": 2, + "event_json": '{"type": "progress", "step": 2}', + } + ] + ) + store = LakebaseDurabilityStore(lakebase=lakebase) + + events = await store.events("session-1", after_sequence=7) + + assert len(events) == 1 + assert events[0].sequence_number == 8 + assert events[0].attempt == 2 + assert events[0].event == {"type": "progress", "step": 2} + + +def test_schema_name_is_validated(): + lakebase, _ = mock_lakebase() + with pytest.raises(ValueError, match="invalid durability schema"): + LakebaseDurabilityStore(lakebase=lakebase, schema="bad-schema;drop") + + +@pytest.mark.asyncio +async def test_store_rejects_empty_execution_id(): + lakebase, _ = mock_lakebase() + store = LakebaseDurabilityStore(lakebase=lakebase) + with pytest.raises(ValueError, match="must not be empty"): + await store.accept("", {})