diff --git a/README.md b/README.md index 587a3098..99b69229 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,7 @@ Some examples require extra dependencies. See each sample's directory for specif * [hello_signal](hello/hello_signal.py) - Send signals to a workflow. * [hello update](hello/hello_update.py) - Send a request to and a response from a client to a workflow execution. +* [activity_sessions](activity_sessions) - Experimental activity sessions reuse worker-local conversation context. * [activity_worker](activity_worker) - Use Python activities from a workflow in another language. * [batch_sliding_window](batch_sliding_window) - Batch processing with a sliding window of child workflows. * [bedrock](bedrock) - Orchestrate a chatbot with Amazon Bedrock. diff --git a/activity_sessions/README.md b/activity_sessions/README.md new file mode 100644 index 00000000..5406b04c --- /dev/null +++ b/activity_sessions/README.md @@ -0,0 +1,106 @@ +# Activity sessions and a worker-local cache + +This sample shows how activity sessions keep a workflow's activities on one worker, +so they can reuse state cached in that worker's memory without per-worker task +queues. + +Each workflow is a conversation: it calls `conversation_step(0)`, +`conversation_step(1)`, ... sequentially. Turn `n` needs every earlier turn. Each +worker runs in a separate process with its own cache of the turns it holds; before +running a turn, a worker fetches the earlier turns it is missing, sleeping +`--fetch-seconds` per turn to stand in for an external conversation store. + +The workflow, activity, and worker code is shared; only whether the workflow opens +a session and how workers are configured changes: + +| Mode | Workflow | Workers | +| --- | --- | --- | +| `none` | No session: ordinary task-queue routing | `WorkerSessionConfig()` | +| `affinity` | Turns run inside `workflow.activity_session(ActivitySessionOptions(placement=UseWorkflowSession()))` | `WorkerSessionConfig()` | +| `single-session` | Same as `affinity` | `WorkerSessionConfig(admission=SingleSessionAdmission(stable_identity=...))` | + +## Prerequisites + +1. Run a server built from the `shivam/agentcore-single-session` branch of + [temporalio/temporal](https://github.com/temporalio/temporal), with the + `default` namespace and `system.enableStickyActivityExecution` enabled. +2. Build the SDK from [sdk-python#1979](https://github.com/temporalio/sdk-python/pull/1979) + by running `poe build-develop` in that SDK checkout, then select its + interpreter: + + ```bash + export SDK_PYTHON=/absolute/path/to/sdk-python/.venv/bin/python + ``` + +3. Install Rich, used for the live view and the summary: + + ```bash + uv pip install --python "$SDK_PYTHON" 'rich>=13,<15' + ``` + +The released SDK does not provide the experimental session API. The sample +connects to `localhost:7233`, namespace `default`, unless overridden by +`TEMPORAL_ADDRESS` / `TEMPORAL_NAMESPACE` or SDK profile configuration. Run the +commands below from the **samples repository root**. + +## Compare modes + +```bash +"$SDK_PYTHON" -m activity_sessions.affinity # none and affinity +"$SDK_PYTHON" -m activity_sessions.affinity --modes none affinity single-session +``` + +Each mode runs `--workflows` (default 8) workflows on a fresh worker fleet: +`--workers` (default 4) workers for `none` and `affinity`, one worker per workflow +for `single-session`. It then compares turns fetched, time spent fetching, and the +most workflows any one worker served: + +- `none`: turns spread across the fleet, so workers keep fetching turns they missed. +- `affinity`: every turn runs on the workflow's worker, which already holds the + earlier turns, so nothing is fetched. Workers serve many workflows. +- `single-session`: the same reuse, with each worker serving exactly one workflow. + +A live view shows each running batch as a grid: one row per workflow, one cell per +turn. `w2` means worker `w2` ran the turn; `w2+3` means it first fetched 3 missing +turns, and the cell's shade deepens with the number fetched. The summary shows a +tile of headline numbers per mode, each mode's grid, and the checks. Pass `--no-live` to print only the summary. + +The command exits non-zero if `affinity` or `single-session` workflows run on more +than one worker or fetch, or if a `single-session` worker serves more than one +workflow. + +## Single session across a fleet loss + +```bash +"$SDK_PYTHON" -m activity_sessions.single_session +``` + +This shows both sides of single-session binding with `--workflows` N (default 6) +workflows and `--workers` M (default 3) single-session workers, in four phases: + +1. **Bind.** Start N workflows and M workers. Each worker binds to one workflow; + the other N - M workflows get no worker, because a bound worker takes no + further work. +2. **Kill the fleet.** Kill every worker (SIGKILL). Every workflow stalls. +3. **New identities.** Start N - M workers with new stable identities. They run + the N - M unbound workflows; the server does not hand them the bound ones, + which stay stalled. +4. **Rebind.** Restart the original M workers with their stable identities. Each + resumes its own workflow. Its new process starts with an empty cache, so its + first turn fetches the history once. + +Each phase advances once its condition is observed, and records a check; the +command exits non-zero if any check fails. A live view shows the phase, each +workflow's state (`… waiting`, `▶ running`, `⏸ stalled`, `✓ done`) and the +identity it is bound to, each worker's process (`● alive`, `✖ killed`) and the +workflow it serves, and the turn grid. `--turn-seconds` (default 0.3) paces +the turns so the phases are easy to follow. + +## Tests + +```bash +"$SDK_PYTHON" -m pytest tests/activity_sessions +# Also run both entrypoints end to end against a session-capable server: +TEMPORAL_TEST_ACTIVITY_SESSION_SAMPLE=1 "$SDK_PYTHON" -m pytest tests/activity_sessions \ + --workflow-environment localhost:7233 +``` diff --git a/activity_sessions/__init__.py b/activity_sessions/__init__.py new file mode 100644 index 00000000..c1e1c49e --- /dev/null +++ b/activity_sessions/__init__.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class StepEvent: + """Sent by an activity when it finishes a conversation turn.""" + + workflow_id: str + step: int + worker: str + # Earlier turns this worker had to fetch before running the step. + fetched: int diff --git a/activity_sessions/activities.py b/activity_sessions/activities.py new file mode 100644 index 00000000..47131baf --- /dev/null +++ b/activity_sessions/activities.py @@ -0,0 +1,42 @@ +import asyncio +from collections.abc import Callable +from typing import cast + +from temporalio import activity + +from activity_sessions import StepEvent + + +class ConversationActivities: + """Activities sharing one worker-local cache; each worker process has its own.""" + + def __init__( + self, + worker: str, + fetch_seconds: float, + turn_seconds: float, + emit: Callable[[StepEvent], None], + ) -> None: + self.worker = worker + self.fetch_seconds = fetch_seconds + self.turn_seconds = turn_seconds + self.emit = emit + # Workflow Run ID -> conversation turns this worker holds, always a prefix. + self.cache: dict[str, list[str]] = {} + + @activity.defn + async def conversation_step(self, step: int) -> None: + """Run turn `step`, which needs every earlier turn of the conversation.""" + info = activity.info() + # Only standalone activities lack workflow IDs; this sample calls it from a workflow. + workflow_id = cast(str, info.workflow_id) + turns = self.cache.setdefault(cast(str, info.workflow_run_id), []) + missing = range(len(turns), step) + for turn in missing: + # Stands in for a slow fetch of one turn from an external conversation store. + await asyncio.sleep(self.fetch_seconds) + turns.append(f"Turn {turn} of {workflow_id}") + # Stands in for the turn's own work. + await asyncio.sleep(self.turn_seconds) + turns.append(f"Turn {step} of {workflow_id}") + self.emit(StepEvent(workflow_id, step, self.worker, len(missing))) diff --git a/activity_sessions/affinity.py b/activity_sessions/affinity.py new file mode 100644 index 00000000..7ca62356 --- /dev/null +++ b/activity_sessions/affinity.py @@ -0,0 +1,181 @@ +"""Compare how much fetching each mode causes and how many workflows share a worker.""" + +import argparse +import asyncio +import contextlib +import uuid +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from multiprocessing.queues import Queue +from typing import Literal, get_args + +from rich.console import Console, Group, RenderableType +from rich.text import Text +from temporalio.client import Client + +from activity_sessions import StepEvent +from activity_sessions.display import GRID_LEGEND, check_list, stat_tiles, turn_grid +from activity_sessions.fleet import ( + Fleet, + Turns, + add_common_args, + connect, + new_event_queue, + start_conversations, +) +from activity_sessions.live import run_live +from activity_sessions.worker import WorkerSpec + +# none: ordinary task-queue routing. affinity: activities join the workflow's +# session on shared session workers. single-session: the same, on workers that +# each serve exactly one workflow. +Mode = Literal["none", "affinity", "single-session"] +LABELS: dict[Mode, str] = { + "none": "No session", + "affinity": "Affinity", + "single-session": "Single session", +} + + +@dataclass +class ModeRun: + mode: Mode + workflow_ids: list[str] + turns: Turns = field(default_factory=Turns) + + +def checks(run: ModeRun) -> list[tuple[bool, str]]: + """What the mode promises; ordinary routing promises nothing.""" + if run.mode == "none": + return [] + label = LABELS[run.mode] + one_worker = all( + len({event.worker for event in run.turns.of(w)}) == 1 for w in run.workflow_ids + ) + no_fetch = not any(event.fetched for event in run.turns.events.values()) + results = [ + ( + one_worker and no_fetch, + f"{label}: each workflow ran every turn on one worker without fetching", + ) + ] + if run.mode == "single-session": + results.append( + ( + all(len(w) == 1 for w in run.turns.workflows_by_worker().values()), + f"{label}: each worker served exactly one workflow", + ) + ) + return results + + +def metrics(run: ModeRun, fetch_seconds: float) -> dict[str, str]: + fetched = sum(event.fetched for event in run.turns.events.values()) + most = max(map(len, run.turns.workflows_by_worker().values())) + return { + "Turns fetched": str(fetched), + "Time spent fetching (s)": f"{fetched * fetch_seconds:.1f}", + "Most workflows on one worker": str(most), + } + + +async def run_mode( + client: Client, + args: argparse.Namespace, + mode: Mode, + events: "Queue[StepEvent]", + on_progress: Callable[[ModeRun], None], +) -> ModeRun: + """Run one batch of workflows on a fresh worker fleet for this mode.""" + prefix = f"{mode}-{uuid.uuid4().hex[:6]}" + single_session = mode == "single-session" + # Single session: one worker per workflow, each bound to it for its lifetime. + workers = [ + WorkerSpec( + f"w{index}", + stable_identity=f"{prefix}-w{index}" if single_session else None, + ) + for index in range(1, (args.workflows if single_session else args.workers) + 1) + ] + run = ModeRun(mode, [f"{prefix}-{i:02}" for i in range(1, args.workflows + 1)]) + task_queue = f"activity-sessions-{prefix}" + fleet = Fleet(task_queue, args, events) + try: + # Start workflows only once every worker polls, so none gets a head start. + await fleet.start(*workers) + handles = await start_conversations( + client, run.workflow_ids, task_queue, mode != "none", args.steps + ) + results = asyncio.gather(*(handle.result() for handle in handles)) + # Workflows can complete just before their last turn's event arrives. + while not results.done() or len(run.turns.events) < args.workflows * args.steps: + await asyncio.sleep(0.1) + run.turns.drain(events) + on_progress(run) + await results + finally: + fleet.stop() + return run + + +def grid(run: ModeRun, steps: int) -> Group: + return Group( + Text(LABELS[run.mode], style="bold"), + turn_grid(run.workflow_ids, run.turns, steps), + ) + + +def report(runs: Sequence[ModeRun], steps: int, fetch_seconds: float) -> Group: + return Group( + stat_tiles({LABELS[run.mode]: metrics(run, fetch_seconds) for run in runs}), + *(grid(run, steps) for run in runs), + Text(GRID_LEGEND, style="bright_black"), + check_list([check for run in runs for check in checks(run)]), + ) + + +async def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--modes", + nargs="+", + choices=get_args(Mode), + default=["none", "affinity"], + help="Modes to run and compare (default: none affinity)", + ) + parser.add_argument("--workflows", type=int, default=8, help="Workflows per mode") + parser.add_argument( + "--workers", + type=int, + default=4, + help="Workers for none and affinity; single-session starts one per workflow", + ) + add_common_args(parser, fetch_seconds=0.5, turn_seconds=0.1) + args = parser.parse_args() + + client = await connect() + events = new_event_queue() + console = Console() + current: list[ModeRun] = [] + + async def work(refresh: Callable[[], None]) -> list[ModeRun]: + def on_progress(run: ModeRun) -> None: + current[:] = [run] + refresh() + + return [ + await run_mode(client, args, mode, events, on_progress) + for mode in args.modes + ] + + runs = await run_live( + console, args.no_live, lambda: grid(current[0], args.steps), work + ) + console.print(report(runs, args.steps, args.fetch_seconds)) + if not all(ok for run in runs for ok, _ in checks(run)): + raise SystemExit(1) + + +if __name__ == "__main__": + with contextlib.suppress(KeyboardInterrupt): + asyncio.run(main()) diff --git a/activity_sessions/display.py b/activity_sessions/display.py new file mode 100644 index 00000000..7e50f1b0 --- /dev/null +++ b/activity_sessions/display.py @@ -0,0 +1,126 @@ +"""Rich building blocks shared by the sample's entrypoints.""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Literal + +from rich import box +from rich.columns import Columns +from rich.console import Group, RenderableType +from rich.panel import Panel +from rich.progress_bar import ProgressBar +from rich.rule import Rule +from rich.table import Table +from rich.text import Text + +from activity_sessions.fleet import Turns + +# Turns fetched is a magnitude: one hue, from muted (none) to strong (many). +FETCH_SHADES = [ + "bright_black", + "color(223)", + "color(215)", + "color(208)", + "bold color(202)", +] +GRID_LEGEND = ( + "w2: worker w2 ran the turn w2+3: it first fetched 3 missing turns (shaded)" +) + +WorkflowState = Literal["waiting", "running", "stalled", "done"] +# Status colors always come with an icon and a label. +WORKFLOW_STATES: dict[WorkflowState, tuple[str, str]] = { + "waiting": ("… waiting", "bright_black"), + "running": ("▶ running", "green"), + "stalled": ("⏸ stalled", "yellow"), + "done": ("✓ done", "bright_black"), +} + + +def turn_grid(workflow_ids: Sequence[str], turns: Turns, steps: int) -> Table: + """One row per workflow, one cell per turn: who ran it, and what it fetched.""" + table = Table(box=None, show_header=False, padding=(0, 1, 0, 0)) + for index, workflow_id in enumerate(workflow_ids, 1): + cells = Text() + for step in range(steps): + event = turns.events.get((workflow_id, step)) + if event is None: + cells.append(f"{'·':<6}", style="bright_black") + continue + label = event.worker + (f"+{event.fetched}" if event.fetched else "") + shade = FETCH_SHADES[min(event.fetched, len(FETCH_SHADES) - 1)] + cells.append(f"{label:<6}", style=shade) + table.add_row(Text(f"wf-{index:02}", style="bright_black"), cells) + return table + + +def stat_tiles(columns: Mapping[str, Mapping[str, str]]) -> Columns: + """A headline tile per column: each metric's value above its label.""" + tiles = [] + for title, metrics in columns.items(): + body = Text() + for label, value in metrics.items(): + body.append(f"{value}\n", style="bold") + body.append(f"{label}\n", style="bright_black") + body.rstrip() + tiles.append(Panel(body, title=title, width=32, box=box.ROUNDED)) + return Columns(tiles) + + +def check_list(checks: Sequence[tuple[bool, str]]) -> Group: + return Group( + *( + Text(f"{'✓' if ok else '✗'} {message}", style="green" if ok else "bold red") + for ok, message in checks + ) + ) + + +def phase_banner(phase: int, phases: Sequence[str]) -> Rule: + return Rule(Text(f"Phase {phase + 1}/{len(phases)}: {phases[phase]}", style="bold")) + + +@dataclass(frozen=True) +class WorkflowRow: + workflow_id: str + state: WorkflowState + worker: str | None + turns: int + steps: int + + +def workflow_table(rows: Sequence[WorkflowRow]) -> Table: + table = Table(title="Workflows", box=box.SIMPLE_HEAD, title_justify="left") + for column in ["Workflow", "State", "Bound to", "Turns"]: + table.add_column(column) + for index, row in enumerate(rows, 1): + label, style = WORKFLOW_STATES[row.state] + table.add_row( + f"wf-{index:02}", + Text(label, style=style), + row.worker or Text("unbound", style="bright_black"), + ProgressBar(total=row.steps, completed=row.turns, width=row.steps * 2), + ) + return table + + +@dataclass(frozen=True) +class WorkerRow: + identity: str + alive: bool + serving: str | None + + +def worker_table(rows: Sequence[WorkerRow]) -> Table: + table = Table(title="Workers", box=box.SIMPLE_HEAD, title_justify="left") + for column in ["Worker", "Process", "Serving"]: + table.add_column(column) + for row in rows: + table.add_row( + row.identity, + Text("● alive", style="green") + if row.alive + else Text("✖ killed", style="red"), + row.serving or Text("—", style="bright_black"), + ) + return table diff --git a/activity_sessions/fleet.py b/activity_sessions/fleet.py new file mode 100644 index 00000000..7443afb9 --- /dev/null +++ b/activity_sessions/fleet.py @@ -0,0 +1,125 @@ +"""Worker processes, the turn events they send, and starting conversations.""" + +import argparse +import asyncio +import contextlib +import multiprocessing +import queue +from dataclasses import dataclass, field +from datetime import timedelta +from multiprocessing.process import BaseProcess +from multiprocessing.queues import Queue + +from temporalio.client import Client, WorkflowHandle +from temporalio.envconfig import ClientConfig + +from activity_sessions import StepEvent +from activity_sessions.worker import WorkerSpec, worker_process +from activity_sessions.workflows import ConversationInput, ProcessConversation + +CONTEXT = multiprocessing.get_context("spawn") + + +@dataclass +class Turns: + """Finished turns, built from the activities' events.""" + + # Latest event per (workflow ID, step); a retried turn replaces its event. + events: dict[tuple[str, int], StepEvent] = field(default_factory=dict) + + def drain(self, events: "Queue[StepEvent]") -> None: + with contextlib.suppress(queue.Empty): + while True: + event = events.get_nowait() + self.events[event.workflow_id, event.step] = event + + def of(self, workflow_id: str) -> list[StepEvent]: + return [e for (w, _), e in sorted(self.events.items()) if w == workflow_id] + + def workflows_by_worker(self) -> dict[str, set[str]]: + result: dict[str, set[str]] = {} + for event in self.events.values(): + result.setdefault(event.worker, set()).add(event.workflow_id) + return result + + +class Fleet: + """Worker processes, each with a fresh interpreter, SDK runtime, and cache.""" + + def __init__( + self, task_queue: str, args: argparse.Namespace, events: "Queue[StepEvent]" + ) -> None: + self.worker_args = (task_queue, args.fetch_seconds, args.turn_seconds, events) + self.processes: dict[str, BaseProcess] = {} + + async def start(self, *specs: WorkerSpec) -> None: + """Start workers and return once every one of them is polling.""" + readies = [] + for spec in specs: + ready = CONTEXT.Event() + process = CONTEXT.Process( + target=worker_process, + args=(spec, *self.worker_args, ready), + daemon=True, + ) + process.start() + self.processes[spec.identity] = process + readies.append(ready) + for ready in readies: + await asyncio.to_thread(ready.wait) + + def kill(self, *identities: str) -> None: + # SIGKILL: no graceful shutdown, like a crash. + for identity in identities: + process = self.processes.pop(identity) + process.kill() + process.join() + + def stop(self) -> None: + self.kill(*self.processes) + + +async def connect() -> Client: + config = ClientConfig.load_client_connect_config() + config.setdefault("target_host", "localhost:7233") + return await Client.connect(**config) + + +async def start_conversations( + client: Client, ids: list[str], task_queue: str, use_session: bool, steps: int +) -> list[WorkflowHandle]: + return [ + await client.start_workflow( + ProcessConversation.run, + ConversationInput(use_session=use_session, steps=steps), + id=workflow_id, + task_queue=task_queue, + execution_timeout=timedelta(minutes=5), + ) + for workflow_id in ids + ] + + +def new_event_queue() -> "Queue[StepEvent]": + return CONTEXT.Queue() + + +def add_common_args( + parser: argparse.ArgumentParser, fetch_seconds: float, turn_seconds: float +) -> None: + parser.add_argument("--steps", type=int, default=6, help="Turns per workflow") + parser.add_argument( + "--fetch-seconds", + type=float, + default=fetch_seconds, + help="Simulated latency of fetching one missing turn", + ) + parser.add_argument( + "--turn-seconds", + type=float, + default=turn_seconds, + help="Simulated work per turn", + ) + parser.add_argument( + "--no-live", action="store_true", help="Only print the final summary" + ) diff --git a/activity_sessions/live.py b/activity_sessions/live.py new file mode 100644 index 00000000..1f97b23d --- /dev/null +++ b/activity_sessions/live.py @@ -0,0 +1,22 @@ +"""Live terminal view; the entrypoints run the same without it (--no-live).""" + +from collections.abc import Awaitable, Callable +from typing import TypeVar + +from rich.console import Console, RenderableType +from rich.live import Live + +T = TypeVar("T") + + +async def run_live( + console: Console, + no_live: bool, + render: Callable[[], RenderableType], + work: Callable[[Callable[[], None]], Awaitable[T]], +) -> T: + """Run `work`, re-rendering `render` each time it calls its refresh argument.""" + if no_live or not console.is_terminal: + return await work(lambda: None) + with Live(console=console, auto_refresh=False, transient=True) as live: + return await work(lambda: live.update(render(), refresh=True)) diff --git a/activity_sessions/single_session.py b/activity_sessions/single_session.py new file mode 100644 index 00000000..cfedacfc --- /dev/null +++ b/activity_sessions/single_session.py @@ -0,0 +1,218 @@ +"""Show single-session workers binding workflows to stable identities, across the +loss of the whole fleet.""" + +import argparse +import asyncio +import contextlib +import uuid +from collections.abc import Callable +from dataclasses import dataclass, field + +from rich.console import Console, Group, RenderableType +from rich.text import Text + +from activity_sessions.display import ( + GRID_LEGEND, + WorkerRow, + WorkflowRow, + WorkflowState, + check_list, + phase_banner, + turn_grid, + worker_table, + workflow_table, +) +from activity_sessions.fleet import ( + Fleet, + Turns, + add_common_args, + connect, + new_event_queue, + start_conversations, +) +from activity_sessions.live import run_live +from activity_sessions.worker import WorkerSpec + +PHASES = ["Bind", "Kill the fleet", "New identities", "Rebind"] + + +@dataclass +class Lifecycle: + """What the run has done so far, for the checks and the display.""" + + workflow_ids: list[str] + steps: int + phase: int = 0 + turns: Turns = field(default_factory=Turns) + # Worker identity -> whether its process is running. + alive: dict[str, bool] = field(default_factory=dict) + checks: list[tuple[bool, str]] = field(default_factory=list) + + def progress(self) -> dict[str, int]: + return {w: len(self.turns.of(w)) for w in self.workflow_ids} + + def worker_of(self, workflow_id: str) -> str | None: + """The worker identity a workflow is bound to, once it has run a turn.""" + turns = self.turns.of(workflow_id) + return turns[-1].worker if turns else None + + def state(self, workflow_id: str) -> WorkflowState: + worker = self.worker_of(workflow_id) + if worker is None: + return "waiting" + if self.progress()[workflow_id] == self.steps: + return "done" + return "running" if self.alive[worker] else "stalled" + + def render(self) -> RenderableType: + labels = {w: f"wf-{i:02}" for i, w in enumerate(self.workflow_ids, 1)} + serving = { + worker: labels[min(ids)] + for worker, ids in self.turns.workflows_by_worker().items() + } + return Group( + phase_banner(self.phase, PHASES), + workflow_table( + [ + WorkflowRow(w, self.state(w), self.worker_of(w), n, self.steps) + for w, n in self.progress().items() + ] + ), + worker_table( + [WorkerRow(w, alive, serving.get(w)) for w, alive in self.alive.items()] + ), + turn_grid(self.workflow_ids, self.turns, self.steps), + Text(GRID_LEGEND, style="bright_black"), + ) + + +async def run_lifecycle( + args: argparse.Namespace, life: Lifecycle, refresh: Callable[[], None] +) -> None: + prefix = life.workflow_ids[0].rsplit("-", 1)[0] + task_queue = f"activity-sessions-{prefix}" + events = new_event_queue() + fleet = Fleet(task_queue, args, events) + + def workers(group: str, count: int) -> list[WorkerSpec]: + return [ + WorkerSpec(f"{group}{i}", stable_identity=f"{prefix}-{group}{i}") + for i in range(1, count + 1) + ] + + originals = workers("a", args.workers) + replacements = workers("b", args.workflows - args.workers) + + async def start(specs: list[WorkerSpec]) -> None: + await fleet.start(*specs) + life.alive.update({spec.identity: True for spec in specs}) + + def kill(specs: list[WorkerSpec]) -> None: + fleet.kill(*(spec.identity for spec in specs)) + life.alive.update({spec.identity: False for spec in specs}) + + async def watch(until: Callable[[], bool]) -> None: + """Record turns until the condition holds.""" + while True: + life.turns.drain(events) + refresh() + if until(): + return + await asyncio.sleep(0.1) + + async def watch_for(seconds: float) -> None: + loop = asyncio.get_running_loop() + deadline = loop.time() + seconds + await watch(lambda: loop.time() >= deadline) + + def done(workflow_ids: list[str]) -> bool: + return all(life.progress()[w] == args.steps for w in workflow_ids) + + try: + # 1. Each original worker binds to one workflow; the rest get none. + await start(originals) + await start_conversations( + await connect(), life.workflow_ids, task_queue, True, args.steps + ) + await watch( + lambda: sum(n >= 2 for n in life.progress().values()) >= args.workers + ) + bound = [w for w, n in life.progress().items() if n] + unbound = [w for w in life.workflow_ids if w not in bound] + life.checks.append( + ( + len(bound) == args.workers, + f"{len(bound)} workflows bound to the {args.workers} workers; " + f"{len(unbound)} got none", + ) + ) + + # 2. With the whole fleet gone, nothing runs. + life.phase = 1 + kill(originals) + # Let turns that finished just before the kill arrive. + await watch_for(0.5) + stalled = life.progress() + await watch_for(2) + life.checks.append((life.progress() == stalled, "Every workflow stalled")) + + # 3. Workers with new identities take only the unbound workflows. + life.phase = 2 + await start(replacements) + await watch(lambda: done(unbound)) + life.checks.append( + ( + all(life.progress()[w] == stalled[w] for w in bound), + "New identities ran the unbound workflows; the bound ones stayed stalled", + ) + ) + + # 4. Restarted with their stable identities, the original workers resume + # their own workflows. + life.phase = 3 + await start(originals) + await watch(lambda: done(life.workflow_ids)) + life.checks.append( + ( + all(len(w) == 1 for w in life.turns.workflows_by_worker().values()) + and all( + len({e.worker for e in life.turns.of(w)}) == 1 + for w in life.workflow_ids + ), + "Each workflow ran every turn on its own worker identity", + ) + ) + finally: + fleet.stop() + + +async def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--workflows", type=int, default=6, help="Workflows (N)") + parser.add_argument( + "--workers", type=int, default=3, help="Original workers (M, fewer than N)" + ) + add_common_args(parser, fetch_seconds=0.2, turn_seconds=0.3) + args = parser.parse_args() + if not 0 < args.workers < args.workflows: + parser.error("--workers must be at least 1 and fewer than --workflows") + + prefix = f"single-session-{uuid.uuid4().hex[:6]}" + life = Lifecycle( + [f"{prefix}-{i:02}" for i in range(1, args.workflows + 1)], args.steps + ) + console = Console() + await run_live( + console, + args.no_live, + life.render, + lambda refresh: run_lifecycle(args, life, refresh), + ) + console.print(life.render(), check_list(life.checks)) + if not all(ok for ok, _ in life.checks): + raise SystemExit(1) + + +if __name__ == "__main__": + with contextlib.suppress(KeyboardInterrupt): + asyncio.run(main()) diff --git a/activity_sessions/worker.py b/activity_sessions/worker.py new file mode 100644 index 00000000..a1f20a42 --- /dev/null +++ b/activity_sessions/worker.py @@ -0,0 +1,58 @@ +import asyncio +import signal +from dataclasses import dataclass +from multiprocessing.queues import Queue +from multiprocessing.synchronize import Event + +from temporalio.client import Client +from temporalio.envconfig import ClientConfig +from temporalio.worker import SingleSessionAdmission, Worker, WorkerSessionConfig + +from activity_sessions import StepEvent +from activity_sessions.activities import ConversationActivities +from activity_sessions.workflows import ProcessConversation + + +@dataclass(frozen=True) +class WorkerSpec: + identity: str + # Dedicates the worker to one workflow for its lifetime, including across + # restarts that reuse this identity. + stable_identity: str | None = None + + +def worker_process( + spec: WorkerSpec, + task_queue: str, + fetch_seconds: float, + turn_seconds: float, + events: "Queue[StepEvent]", + ready: Event, +) -> None: + async def run() -> None: + config = ClientConfig.load_client_connect_config() + config.setdefault("target_host", "localhost:7233") + client = await Client.connect(**config, identity=spec.identity) + activities = ConversationActivities( + spec.identity, fetch_seconds, turn_seconds, events.put + ) + admission = ( + SingleSessionAdmission(stable_identity=spec.stable_identity) + if spec.stable_identity + else None + ) + async with Worker( + client, + task_queue=task_queue, + workflows=[ProcessConversation], + activities=[activities.conversation_step], + session_config=WorkerSessionConfig(admission=admission), + disable_eager_activity_execution=True, + ): + ready.set() + # Run until the parent kills this process. + await asyncio.Event().wait() + + # The parent handles Ctrl+C and stops every worker. + signal.signal(signal.SIGINT, signal.SIG_IGN) + asyncio.run(run()) diff --git a/activity_sessions/workflows.py b/activity_sessions/workflows.py new file mode 100644 index 00000000..b7bc4359 --- /dev/null +++ b/activity_sessions/workflows.py @@ -0,0 +1,34 @@ +import contextlib +from dataclasses import dataclass +from datetime import timedelta + +from temporalio import workflow +from temporalio.common import ActivitySessionOptions, UseWorkflowSession + + +@dataclass +class ConversationInput: + use_session: bool + steps: int + + +@workflow.defn +class ProcessConversation: + @workflow.run + async def run(self, input: ConversationInput) -> None: + # Activities started in the scope run on this workflow's worker. + session = ( + workflow.activity_session( + ActivitySessionOptions(placement=UseWorkflowSession()) + ) + if input.use_session + else contextlib.nullcontext() + ) + with session: + for step in range(input.steps): + await workflow.execute_activity( + "conversation_step", + step, + # Short, so a turn lost with a killed worker is retried soon. + start_to_close_timeout=timedelta(seconds=10), + ) diff --git a/pyproject.toml b/pyproject.toml index 9fff6239..93e386d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,7 @@ Documentation = "https://docs.temporal.io/docs/python" "Bug Tracker" = "https://github.com/temporalio/samples-python/issues" [dependency-groups] +activity-sessions = ["rich>=13,<15"] dev = [ "ruff>=0.5.0,<0.6", "mypy>=1.4.1,<2", diff --git a/tests/activity_sessions/__init__.py b/tests/activity_sessions/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/activity_sessions/affinity_test.py b/tests/activity_sessions/affinity_test.py new file mode 100644 index 00000000..a8daa0d1 --- /dev/null +++ b/tests/activity_sessions/affinity_test.py @@ -0,0 +1,42 @@ +from activity_sessions import StepEvent +from activity_sessions.affinity import Mode, ModeRun, checks, metrics + + +def run(mode: Mode, placements: dict[str, list[tuple[str, int]]]) -> ModeRun: + """Workflow ID -> (worker, turns fetched) per step.""" + result = ModeRun(mode, list(placements)) + for workflow_id, steps in placements.items(): + for step, (worker, fetched) in enumerate(steps): + result.turns.events[workflow_id, step] = StepEvent( + workflow_id, step, worker, fetched + ) + return result + + +def passed(result: ModeRun) -> list[bool]: + return [ok for ok, _ in checks(result)] + + +def test_no_session_has_no_checks() -> None: + assert checks(run("none", {"wf": [("w1", 0), ("w2", 1)]})) == [] + + +def test_affinity_requires_one_worker_without_fetches() -> None: + assert passed(run("affinity", {"wf": [("w1", 0), ("w1", 0)]})) == [True] + assert passed(run("affinity", {"wf": [("w1", 0), ("w2", 1)]})) == [False] + + +def test_single_session_requires_one_workflow_per_worker() -> None: + shared = run("single-session", {"a": [("w1", 0)], "b": [("w1", 0)]}) + assert passed(shared) == [True, False] + dedicated = run("single-session", {"a": [("w1", 0)], "b": [("w2", 0)]}) + assert passed(dedicated) == [True, True] + + +def test_metrics() -> None: + result = run("none", {"a": [("w1", 0), ("w1", 1)], "b": [("w1", 2)]}) + assert metrics(result, fetch_seconds=0.5) == { + "Turns fetched": "3", + "Time spent fetching (s)": "1.5", + "Most workflows on one worker": "2", + } diff --git a/tests/activity_sessions/cache_test.py b/tests/activity_sessions/cache_test.py new file mode 100644 index 00000000..74918556 --- /dev/null +++ b/tests/activity_sessions/cache_test.py @@ -0,0 +1,37 @@ +from dataclasses import replace + +from temporalio.testing import ActivityEnvironment + +from activity_sessions import StepEvent +from activity_sessions.activities import ConversationActivities + + +def environment(run_id: str) -> ActivityEnvironment: + env = ActivityEnvironment() + env.info = replace(env.info, workflow_id="workflow-1", workflow_run_id=run_id) + return env + + +def worker(name: str, events: list[StepEvent]) -> ConversationActivities: + return ConversationActivities(name, 0, 0, events.append) + + +async def test_same_worker_never_fetches_its_own_turns() -> None: + events: list[StepEvent] = [] + activities = worker("w1", events) + env = environment("run-1") + for step in range(3): + await env.run(activities.conversation_step, step) + assert [event.fetched for event in events] == [0, 0, 0] + assert len(activities.cache["run-1"]) == 3 + + +async def test_stale_or_cold_worker_fetches_only_missing_turns() -> None: + events: list[StepEvent] = [] + worker_a, worker_b = worker("w1", events), worker("w2", events) + env = environment("run-1") + await env.run(worker_a.conversation_step, 0) + await env.run(worker_b.conversation_step, 1) + await env.run(worker_a.conversation_step, 2) + # w2 lacks turn 0; w1 returns holding turn 0 but not turn 1. + assert [event.fetched for event in events] == [0, 1, 1] diff --git a/tests/activity_sessions/display_test.py b/tests/activity_sessions/display_test.py new file mode 100644 index 00000000..0f87cd0e --- /dev/null +++ b/tests/activity_sessions/display_test.py @@ -0,0 +1,48 @@ +import io + +from rich.console import Console, RenderableType + +from activity_sessions import StepEvent +from activity_sessions.display import ( + WorkerRow, + WorkflowRow, + check_list, + stat_tiles, + turn_grid, + worker_table, + workflow_table, +) +from activity_sessions.fleet import Turns + + +def render(renderable: RenderableType) -> str: + output = io.StringIO() + Console(file=output, width=120, color_system=None).print(renderable) + return output.getvalue() + + +def test_turn_grid_labels_worker_and_fetches() -> None: + turns = Turns() + turns.events["wf", 0] = StepEvent("wf", 0, "w1", 0) + turns.events["wf", 1] = StepEvent("wf", 1, "w2", 1) + text = render(turn_grid(["wf"], turns, 3)) + assert "wf-01 w1 w2+1 ·" in text + + +def test_status_tables_and_checks_use_labels() -> None: + text = render( + workflow_table( + [ + WorkflowRow("a", "stalled", "a1", 2, 4), + WorkflowRow("b", "waiting", None, 0, 4), + ] + ) + ) + assert "⏸ stalled" in text and "… waiting" in text and "unbound" in text + assert "✖ killed" in render(worker_table([WorkerRow("a1", False, "wf-01")])) + assert "✗ bad" in render(check_list([(True, "good"), (False, "bad")])) + + +def test_stat_tiles_show_each_metric() -> None: + text = render(stat_tiles({"Affinity": {"Turns fetched": "0"}})) + assert "Affinity" in text and "Turns fetched" in text diff --git a/tests/activity_sessions/fleet_test.py b/tests/activity_sessions/fleet_test.py new file mode 100644 index 00000000..1a8e9093 --- /dev/null +++ b/tests/activity_sessions/fleet_test.py @@ -0,0 +1,20 @@ +import queue + +from activity_sessions import StepEvent +from activity_sessions.fleet import Turns + + +def test_turns_keep_the_latest_event_per_step() -> None: + events: "queue.Queue[StepEvent]" = queue.Queue() + for event in [ + StepEvent("a", 1, "w1", 0), + StepEvent("a", 0, "w1", 0), + StepEvent("b", 0, "w2", 0), + # A retried turn replaces its earlier event. + StepEvent("a", 1, "w2", 1), + ]: + events.put(event) + turns = Turns() + turns.drain(events) # type: ignore[arg-type] + assert [event.worker for event in turns.of("a")] == ["w1", "w2"] + assert turns.workflows_by_worker() == {"w1": {"a"}, "w2": {"a", "b"}} diff --git a/tests/activity_sessions/integration_test.py b/tests/activity_sessions/integration_test.py new file mode 100644 index 00000000..16ce45ae --- /dev/null +++ b/tests/activity_sessions/integration_test.py @@ -0,0 +1,52 @@ +import asyncio +import os +import signal +import sys + +import pytest +from temporalio.client import Client + +pytestmark = pytest.mark.skipif( + os.getenv("TEMPORAL_TEST_ACTIVITY_SESSION_SAMPLE") != "1", + reason="requires the session SDK and a session-capable server", +) + + +@pytest.mark.parametrize( + "entrypoint", + [ + ["activity_sessions.affinity", "--modes", "none", "affinity", "single-session"], + ["activity_sessions.single_session", "--turn-seconds", "0.2"], + ], + ids=["affinity", "single-session"], +) +async def test_entrypoint( + client: Client, request: pytest.FixtureRequest, entrypoint: list[str] +) -> None: + if request.config.getoption("--workflow-environment") in ("local", "time-skipping"): + pytest.fail("Use --workflow-environment host:port for the custom server") + demo = await asyncio.create_subprocess_exec( + sys.executable, + "-m", + *entrypoint, + "--steps", + "4", + "--fetch-seconds", + "0.1", + "--no-live", + env={ + **os.environ, + "TEMPORAL_ADDRESS": client.service_client.config.target_host, + "TEMPORAL_NAMESPACE": client.namespace, + }, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + ) + try: + output, _ = await asyncio.wait_for(demo.communicate(), 180) + assert demo.returncode == 0, output.decode() + assert "✗" not in output.decode() + finally: + if demo.returncode is None: + demo.send_signal(signal.SIGINT) + await asyncio.wait_for(demo.wait(), 25) diff --git a/tests/activity_sessions/single_session_test.py b/tests/activity_sessions/single_session_test.py new file mode 100644 index 00000000..7d51509d --- /dev/null +++ b/tests/activity_sessions/single_session_test.py @@ -0,0 +1,17 @@ +from activity_sessions import StepEvent +from activity_sessions.single_session import Lifecycle + + +def test_lifecycle_states() -> None: + life = Lifecycle(["a", "b", "c"], steps=2) + for event in [ + StepEvent("a", 0, "a1", 0), + StepEvent("a", 1, "a1", 0), + StepEvent("b", 0, "a2", 0), + ]: + life.turns.events[event.workflow_id, event.step] = event + life.alive = {"a1": False, "a2": False} + assert life.progress() == {"a": 2, "b": 1, "c": 0} + assert [life.state(w) for w in life.workflow_ids] == ["done", "stalled", "waiting"] + life.alive["a2"] = True + assert life.state("b") == "running" and life.worker_of("b") == "a2" diff --git a/uv.lock b/uv.lock index 74d3488d..5515a4f7 100644 --- a/uv.lock +++ b/uv.lock @@ -4742,6 +4742,9 @@ dependencies = [ ] [package.dev-dependencies] +activity-sessions = [ + { name = "rich" }, +] bedrock = [ { name = "boto3" }, ] @@ -4866,6 +4869,7 @@ requires-dist = [ ] [package.metadata.requires-dev] +activity-sessions = [{ name = "rich", specifier = ">=13,<15" }] bedrock = [{ name = "boto3", specifier = ">=1.34.92,<2" }] cloud-export-to-parquet = [ { name = "boto3", specifier = ">=1.34.89,<2" },