diff --git a/README.md b/README.md index 587a3098..9a160d0f 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 workflow-affine activities 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..9fef4430 --- /dev/null +++ b/activity_sessions/README.md @@ -0,0 +1,81 @@ +# Activity sessions and a worker-local cache + +This sample shows how activity affinity lets a workflow reuse conversation history +cached in its worker's memory, without managing per-worker task queues. + +Each workflow is a conversation: it calls `conversation_step(0)`, +`conversation_step(1)`, ... sequentially (six turns by default). Turn `n` needs +every earlier turn `0..n-1`, then produces turn `n`. Each worker runs in a +separate process with its own dictionary cache of the turns it holds, keyed by +workflow Run ID. Before running a turn, a worker fetches only the earlier turns it +is missing, sleeping `--fetch-seconds` (default 0.5s) per turn to stand in for an +external conversation store. Turn 0 needs no history. The history is +**worker-local, not workflow state**, and can be refetched if absent. + +With affinity **on**, activities use +`workflow.ActivitySessionOptions.workflow_session()`: the workflow's worker +produced every earlier turn itself, so it never fetches. With affinity **off**, +turns spread across the pool. A worker that ran turn 0 and then returns for turn 2 +must fetch turn 1, and a cold worker at turn `n` fetches all `n` earlier turns, so +fetch time grows with how stale each worker is. Turns may still colocate by +chance; ordinary routing does not guarantee either outcome. + +## Prerequisites + +1. Run the server from [temporal#12630](https://github.com/temporalio/temporal/pull/12630), + 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. +3. From the **samples repository root**, select the built SDK's interpreter and + install Rich for the result tables: + + ```bash + export SDK_PYTHON=/absolute/path/to/sdk-python/.venv/bin/python + uv pip install --python "$SDK_PYTHON" 'rich>=13,<15' + ``` + +Use this interpreter for the commands below; 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 + +From the **samples repository root**, run: + +```bash +"$SDK_PYTHON" -m activity_sessions.run +``` + +This starts four session-capable worker processes, runs 20 concurrent workflows +with affinity off and then 20 with affinity on, and shuts down the workers. + +While a batch runs, a live view shows how far each worker holds every +conversation (`03:t4` means workflow 3 through turn 4) and a placement strip per +workflow. In the strip, `○2` means worker 2 ran the turn with every earlier turn +already cached, `●2(3)` means worker 2 first fetched 3 missing turns, and `◌` is +a fetch in progress; color also identifies the worker. The final summary prints +both batches side by side, followed by an impact table: turns fetched, steps that +had to fetch, time spent fetching, cached turns across the fleet, and workflow +latency. +Pass `--no-live` (or redirect output) to print only the summary. + +Useful options (run with `--help` for all): + +```bash +# More workers and longer conversations make ordinary placement fetch even more. +"$SDK_PYTHON" -m activity_sessions.run --session-workers 8 --steps 10 + +# One session-capable worker plus three activity-only workers. +"$SDK_PYTHON" -m activity_sessions.run --session-workers 1 --activity-workers 3 + +# Run only one mode, with a faster simulated fetch. +"$SDK_PYTHON" -m activity_sessions.run --affinity off --fetch-seconds 0.1 +``` + +`--workflows` is the concurrent workflow count **per mode**. Only session-capable +workers register the workflow; activity-only workers never receive session +activities. Workflows start only after every worker is polling, eager activity +execution is disabled, and each invocation uses a unique task queue. The command +exits non-zero if any affinity-on workflow ran on more than one worker or fetched +any turn. diff --git a/activity_sessions/__init__.py b/activity_sessions/__init__.py new file mode 100644 index 00000000..7981b4cb --- /dev/null +++ b/activity_sessions/__init__.py @@ -0,0 +1,15 @@ +from dataclasses import dataclass + +TASK_QUEUE = "activity-session-cache" + + +@dataclass(frozen=True) +class StepEvent: + """Sent by an activity when a step starts and again when it finishes.""" + + workflow_id: str + step: int + worker: str + # Earlier conversation turns this worker had to fetch before running the step. + fetched: int + finished: bool diff --git a/activity_sessions/activities.py b/activity_sessions/activities.py new file mode 100644 index 00000000..57651b85 --- /dev/null +++ b/activity_sessions/activities.py @@ -0,0 +1,44 @@ +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_identity: str, + fetch_seconds: float, + emit: Callable[[StepEvent], None], + ) -> None: + self.worker_identity = worker_identity + self.fetch_seconds = fetch_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) + run_id = cast(str, info.workflow_run_id) + turns = self.cache.setdefault(run_id, []) + missing = range(len(turns), step) + self.emit( + StepEvent(workflow_id, step, self.worker_identity, len(missing), False) + ) + 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}") + turns.append(f"Turn {step} of {workflow_id}") + self.emit( + StepEvent(workflow_id, step, self.worker_identity, len(missing), True) + ) diff --git a/activity_sessions/display.py b/activity_sessions/display.py new file mode 100644 index 00000000..a3bcce34 --- /dev/null +++ b/activity_sessions/display.py @@ -0,0 +1,232 @@ +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field + +from rich import box +from rich.console import Console, Group, RenderableType +from rich.table import Table +from rich.text import Text + +from activity_sessions import StepEvent + +WORKER_COLORS = [ + "cyan", + "magenta", + "green", + "yellow", + "blue", + "red", + "bright_cyan", + "bright_magenta", + "bright_green", + "bright_yellow", +] +LEGEND = ( + "○2 worker 2 had every earlier turn ●2(3) worker 2 fetched 3 missing turns " + "◌ fetching" +) + + +@dataclass +class Batch: + affinity_enabled: bool + workflow_ids: list[str] + steps: dict[str, dict[int, StepEvent]] = field(default_factory=dict) + durations: dict[str, float] = field(default_factory=dict) + wall_seconds: float = 0.0 + + @property + def mode(self) -> str: + return f"Affinity {'ON' if self.affinity_enabled else 'OFF'}" + + def record(self, event: StepEvent) -> None: + if event.workflow_id not in self.workflow_ids: + raise ValueError(f"Event for unknown workflow {event.workflow_id}") + self.steps.setdefault(event.workflow_id, {})[event.step] = event + + def cached_turns(self, workflow_id: str) -> dict[str, int]: + """Turns each worker holds for a workflow: everything up to its latest step.""" + held: dict[str, int] = {} + for event in self.steps.get(workflow_id, {}).values(): + if event.finished: + held[event.worker] = max(held.get(event.worker, 0), event.step + 1) + return held + + def finished_steps(self) -> list[StepEvent]: + return [ + event + for steps in self.steps.values() + for event in steps.values() + if event.finished + ] + + +def worker_number(worker: str) -> int: + return int(worker.rsplit("-", 1)[1]) + + +def worker_color(worker: str) -> str: + return WORKER_COLORS[(worker_number(worker) - 1) % len(WORKER_COLORS)] + + +def placement_strip( + batch: Batch, steps: int, worker_count: int, show_pending: bool = False +) -> Table: + """One row per workflow and one cell per step: glyph = fetches, color = worker.""" + width = len(f"●{worker_count}({steps - 1})") + table = Table(box=None, show_header=False, padding=(0, 1, 0, 0)) + table.add_column(style="dim") + table.add_column() + table.add_column(justify="right") + table.add_column(justify="right", style="dim") + for index, workflow_id in enumerate(batch.workflow_ids, 1): + events = batch.steps.get(workflow_id, {}) + if not events and not show_pending: + continue + cells = Text() + for step in range(steps): + event = events.get(step) + if event is None: + cells.append("·".ljust(width) + " ", style="dim") + continue + number = worker_number(event.worker) + if not event.fetched: + cell = f"○{number}" + else: + cell = f"{'●' if event.finished else '◌'}{number}({event.fetched})" + cells.append(cell.ljust(width) + " ", style=worker_color(event.worker)) + fetched = sum(event.fetched for event in events.values()) + duration = batch.durations.get(workflow_id) + table.add_row( + f"wf-{index:02}", + cells, + f"{duration:.1f}s" if duration is not None else "", + f"{fetched} fetched" if events else "", + ) + return table + + +def print_worker_configuration(workers: Mapping[str, bool], console: Console) -> None: + table = Table(title="Worker configuration", box=box.ROUNDED) + table.add_column("Worker") + table.add_column("Role") + table.add_column("Session support") + for identity, session_enabled in workers.items(): + table.add_row( + Text(identity, style=worker_color(identity)), + "Workflow + activity" if session_enabled else "Activity only", + "[green]ON[/green]" if session_enabled else "[dim]OFF[/dim]", + ) + console.print(table) + + +def percentile(values: Sequence[float], fraction: float) -> float: + ordered = sorted(values) + return ordered[max(0, math.ceil(fraction * len(ordered)) - 1)] + + +def affinity_violations(batch: Batch) -> list[str]: + """Workflows whose turns did not all run on one worker without fetching.""" + return [ + workflow_id + for workflow_id in batch.workflow_ids + if len(batch.cached_turns(workflow_id)) != 1 + or any(event.fetched for event in batch.steps[workflow_id].values()) + ] + + +def comparison_table(batches: Sequence[Batch], fetch_seconds: float) -> Table: + table = Table(title="Impact", box=box.SIMPLE_HEAD) + table.add_column("") + for batch in batches: + table.add_column(batch.mode, justify="right") + compare = len(batches) == 2 + if compare: + table.add_column("Change", justify="right", style="bold green") + + def add( + label: str, values: list[float], fmt: str, change: str | None = None + ) -> None: + row = [label, *(format(value, fmt) for value in values)] + if compare: + row.append(change or "") + table.add_row(*row) + + fetched = [ + float(sum(event.fetched for event in batch.finished_steps())) + for batch in batches + ] + fetching_steps = [ + float(sum(bool(event.fetched) for event in batch.finished_steps())) + for batch in batches + ] + fleet_turns = [ + sum( + sum(batch.cached_turns(workflow_id).values()) + for workflow_id in batch.workflow_ids + ) + / len(batch.workflow_ids) + for batch in batches + ] + p50 = [percentile(list(batch.durations.values()), 0.5) for batch in batches] + p95 = [percentile(list(batch.durations.values()), 0.95) for batch in batches] + walls = [batch.wall_seconds for batch in batches] + + def fewer(values: list[float]) -> str | None: + if not compare or not values[0]: + return None + return f"{1 - values[1] / values[0]:.0%} fewer" + + def faster(values: list[float]) -> str | None: + if not compare or not values[1]: + return None + return f"{values[0] / values[1]:.1f}x faster" + + add("Turns fetched", fetched, ".0f", fewer(fetched)) + add("Steps that had to fetch", fetching_steps, ".0f", fewer(fetching_steps)) + add( + "Time spent fetching (s)", + [count * fetch_seconds for count in fetched], + ".1f", + fewer(fetched), + ) + add("Cached turns across fleet per workflow", fleet_turns, ".1f") + add("Workflow latency p50 (s)", p50, ".1f", faster(p50)) + add("Workflow latency p95 (s)", p95, ".1f", faster(p95)) + add("Batch wall time (s)", walls, ".1f", faster(walls)) + return table + + +def summary( + batches: Sequence[Batch], steps: int, worker_count: int, fetch_seconds: float +) -> RenderableType: + strips = Table.grid(padding=(0, 4)) + for batch in batches: + strips.add_column() + strips.add_row(*(Text(batch.mode, style="bold") for batch in batches)) + strips.add_row(*(placement_strip(batch, steps, worker_count) for batch in batches)) + parts: list[RenderableType] = [ + strips, + Text(LEGEND, style="dim"), + comparison_table(batches, fetch_seconds), + ] + for batch in batches: + if not batch.affinity_enabled: + continue + violations = affinity_violations(batch) + if violations: + parts.append( + Text( + f"Affinity not observed for: {', '.join(violations)}", + style="bold red", + ) + ) + else: + parts.append( + Text( + f"Verified: each of {len(batch.workflow_ids)} workflow(s) ran " + "every turn on one worker without fetching earlier turns.", + style="green", + ) + ) + return Group(*parts) diff --git a/activity_sessions/live.py b/activity_sessions/live.py new file mode 100644 index 00000000..ac3a099b --- /dev/null +++ b/activity_sessions/live.py @@ -0,0 +1,94 @@ +"""Optional live terminal view; the demo runs the same without it (--no-live).""" + +from collections.abc import Mapping +from types import TracebackType + +from rich.columns import Columns +from rich.console import Console, Group, RenderableType +from rich.live import Live +from rich.panel import Panel +from rich.rule import Rule +from rich.text import Text + +from activity_sessions.display import ( + LEGEND, + Batch, + placement_strip, + worker_color, +) + + +class LiveView: + def __init__(self, workers: Mapping[str, bool], steps: int, console: Console): + self.workers = workers + self.steps = steps + self.live = Live(console=console, auto_refresh=False, transient=True) + + def __enter__(self) -> "LiveView": + self.live.start() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self.live.stop() + + def show(self, batch: Batch) -> None: + self.live.update(self.render(batch), refresh=True) + + def render(self, batch: Batch) -> RenderableType: + held = { + workflow_id: batch.cached_turns(workflow_id) + for workflow_id in batch.workflow_ids + } + events = [event for steps in batch.steps.values() for event in steps.values()] + panels = [] + for worker, session_enabled in self.workers.items(): + mine = [event for event in events if event.worker == worker] + fetched = sum(event.fetched for event in mine if event.finished) + fetching = [event for event in mine if event.fetched and not event.finished] + body = Text() + body.append(f"steps {len(mine)} turns fetched {fetched}\n") + body.append( + " ".join(f"◌({event.fetched})" for event in fetching) + "\n", + style="yellow", + ) + # wf:t means this worker holds the workflow's conversation through turn t. + body.append("cache: ", style="dim") + body.append( + " ".join( + f"{index:02}:t{turns[worker] - 1}" + for index, turns in enumerate(held.values(), 1) + if worker in turns + ) + ) + title = worker if session_enabled else f"{worker} (activity only)" + panels.append( + Panel( + body, + title=Text(title, style=f"bold {worker_color(worker)}"), + width=29, + ) + ) + entries = sum(sum(turns.values()) for turns in held.values()) + started = len(batch.steps) + fleet = Text( + f"Fleet: {entries} cached turns for {started} started workflow(s)" + + (f" ({entries / started:.1f} per workflow)" if started else ""), + style="bold", + ) + return Group( + Rule( + Text( + f"{batch.mode} - {len(batch.workflow_ids)} workflow(s)", + style="bold", + ) + ), + Columns(panels), + fleet, + placement_strip(batch, self.steps, len(self.workers), show_pending=True), + Text(LEGEND, style="dim"), + ) diff --git a/activity_sessions/run.py b/activity_sessions/run.py new file mode 100644 index 00000000..7cf0eb0f --- /dev/null +++ b/activity_sessions/run.py @@ -0,0 +1,177 @@ +import argparse +import asyncio +import contextlib +import multiprocessing +import queue +import time +import uuid +from datetime import timedelta +from multiprocessing.queues import Queue +from typing import TYPE_CHECKING + +from rich.console import Console +from temporalio.client import Client +from temporalio.envconfig import ClientConfig + +from activity_sessions import TASK_QUEUE, StepEvent +from activity_sessions.display import ( + Batch, + affinity_violations, + print_worker_configuration, + summary, +) +from activity_sessions.worker import worker_process +from activity_sessions.workflows import ConversationInput, ProcessConversation + +if TYPE_CHECKING: + from activity_sessions.live import LiveView + + +async def run_batch( + client: Client, + args: argparse.Namespace, + affinity_enabled: bool, + events: "Queue[StepEvent]", + view: "LiveView | None", +) -> Batch: + prefix = f"{'on' if affinity_enabled else 'off'}-{uuid.uuid4().hex[:6]}" + batch = Batch( + affinity_enabled, + [f"{prefix}-{index:02}" for index in range(1, args.workflows + 1)], + ) + started = time.monotonic() + + async def run_workflow(workflow_id: str) -> None: + begin = time.monotonic() + await client.execute_workflow( + ProcessConversation.run, + ConversationInput(affinity_enabled, args.steps), + id=workflow_id, + task_queue=args.task_queue, + execution_timeout=timedelta(minutes=5), + ) + batch.durations[workflow_id] = time.monotonic() - begin + + pending = {asyncio.create_task(run_workflow(w)) for w in batch.workflow_ids} + expected_steps = args.workflows * args.steps + # Workflows can complete just before their activities' last events arrive. + while pending or len(batch.finished_steps()) < expected_steps: + if pending: + done, pending = await asyncio.wait(pending, timeout=0.1) + for task in done: + task.result() + else: + await asyncio.sleep(0.05) + while True: + try: + batch.record(events.get_nowait()) + except queue.Empty: + break + if view: + view.show(batch) + batch.wall_seconds = time.monotonic() - started + return batch + + +async def main() -> None: + parser = argparse.ArgumentParser( + description="Run a worker pool and compare activity affinity and cache reuse." + ) + parser.add_argument( + "--affinity", choices=["compare", "on", "off"], default="compare" + ) + parser.add_argument("--task-queue", default=f"{TASK_QUEUE}-{uuid.uuid4().hex[:8]}") + parser.add_argument("--session-workers", type=int, default=4) + parser.add_argument( + "--activity-workers", + type=int, + default=0, + help="Additional activity-only workers without session support", + ) + parser.add_argument( + "--workflows", type=int, default=20, help="Concurrent workflows per mode" + ) + parser.add_argument( + "--steps", type=int, default=6, help="Conversation turns per workflow" + ) + parser.add_argument( + "--fetch-seconds", + type=float, + default=0.5, + help="Simulated latency of fetching one missing conversation turn", + ) + parser.add_argument( + "--no-live", action="store_true", help="Only print the final summary" + ) + args = parser.parse_args() + if args.session_workers < 1: + parser.error("The demo needs at least one session worker to run workflows") + if args.activity_workers < 0: + parser.error("--activity-workers cannot be negative") + if args.workflows < 1 or args.steps < 1: + parser.error("--workflows and --steps must be at least 1") + + workers = { + f"worker-{index}": index <= args.session_workers + for index in range(1, args.session_workers + args.activity_workers + 1) + } + config = ClientConfig.load_client_connect_config() + config.setdefault("target_host", "localhost:7233") + client = await Client.connect(**config) + console = Console() + print_worker_configuration(workers, console) + + # Spawn gives each worker a fresh interpreter, SDK runtime, and private cache. + context = multiprocessing.get_context("spawn") + events: "Queue[StepEvent]" = context.Queue() + ready = context.Barrier(len(workers) + 1) + processes = [ + context.Process( + target=worker_process, + args=( + identity, + session_enabled, + args.task_queue, + args.fetch_seconds, + events, + ready, + ), + daemon=True, + ) + for identity, session_enabled in workers.items() + ] + modes = [False, True] if args.affinity == "compare" else [args.affinity == "on"] + batches = [] + try: + for process in processes: + process.start() + with console.status("Starting workers..."): + # Start workflows only once every worker polls, so none gets a head start. + await asyncio.to_thread(ready.wait, 60) + view = None + if not args.no_live and console.is_terminal: + from activity_sessions.live import LiveView + + view = LiveView(workers, args.steps, console) + with view or contextlib.nullcontext(): + for affinity_enabled in modes: + batches.append( + await run_batch(client, args, affinity_enabled, events, view) + ) + finally: + ready.abort() + for process in processes: + process.terminate() + for process in processes: + process.join() + + console.print(summary(batches, args.steps, len(workers), args.fetch_seconds)) + if any(affinity_violations(batch) for batch in batches if batch.affinity_enabled): + raise SystemExit(1) + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except KeyboardInterrupt: + pass diff --git a/activity_sessions/worker.py b/activity_sessions/worker.py new file mode 100644 index 00000000..0524c05c --- /dev/null +++ b/activity_sessions/worker.py @@ -0,0 +1,54 @@ +import asyncio +import signal +from multiprocessing.queues import Queue +from multiprocessing.synchronize import Barrier + +from temporalio.client import Client +from temporalio.envconfig import ClientConfig +from temporalio.worker import Worker, WorkerSessionOptions + +from activity_sessions import StepEvent +from activity_sessions.activities import ConversationActivities +from activity_sessions.workflows import ProcessConversation + + +async def run_worker( + identity: str, + session_enabled: bool, + task_queue: str, + fetch_seconds: float, + events: "Queue[StepEvent]", + ready: Barrier, +) -> None: + config = ClientConfig.load_client_connect_config() + config.setdefault("target_host", "localhost:7233") + client = await Client.connect(**config, identity=identity) + activities = ConversationActivities(identity, fetch_seconds, events.put) + worker = Worker( + client, + task_queue=task_queue, + workflows=[ProcessConversation] if session_enabled else [], + activities=[activities.conversation_step], + session_worker_options=WorkerSessionOptions() if session_enabled else None, + disable_eager_activity_execution=True, + max_concurrent_activities=100, + ) + async with worker: + await asyncio.to_thread(ready.wait) + # Run until the parent terminates this process. + await asyncio.Event().wait() + + +def worker_process( + identity: str, + session_enabled: bool, + task_queue: str, + fetch_seconds: float, + events: "Queue[StepEvent]", + ready: Barrier, +) -> None: + # The parent handles Ctrl+C and terminates every worker. + signal.signal(signal.SIGINT, signal.SIG_IGN) + asyncio.run( + run_worker(identity, session_enabled, task_queue, fetch_seconds, events, ready) + ) diff --git a/activity_sessions/workflows.py b/activity_sessions/workflows.py new file mode 100644 index 00000000..106b7286 --- /dev/null +++ b/activity_sessions/workflows.py @@ -0,0 +1,30 @@ +from dataclasses import dataclass +from datetime import timedelta + +from temporalio import workflow +from temporalio.common import RetryPolicy + + +@dataclass +class ConversationInput: + affinity_enabled: bool + steps: int + + +@workflow.defn +class ProcessConversation: + @workflow.run + async def run(self, input: ConversationInput) -> None: + session_options = ( + workflow.ActivitySessionOptions.workflow_session() + if input.affinity_enabled + else None + ) + for step in range(input.steps): + await workflow.execute_activity( + "conversation_step", + step, + start_to_close_timeout=timedelta(seconds=30), + retry_policy=RetryPolicy(maximum_attempts=1), + session_options=session_options, + ) 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..8b137891 --- /dev/null +++ b/tests/activity_sessions/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/activity_sessions/cache_test.py b/tests/activity_sessions/cache_test.py new file mode 100644 index 00000000..5bcb7701 --- /dev/null +++ b/tests/activity_sessions/cache_test.py @@ -0,0 +1,66 @@ +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 fetched(events: list[StepEvent]) -> list[int]: + return [event.fetched for event in events if event.finished] + + +async def test_first_turn_needs_no_fetch() -> None: + events: list[StepEvent] = [] + activities = ConversationActivities("worker-1", 0, events.append) + await environment("run-1").run(activities.conversation_step, 0) + assert events == [ + StepEvent("workflow-1", 0, "worker-1", 0, False), + StepEvent("workflow-1", 0, "worker-1", 0, True), + ] + assert activities.cache["run-1"] == ["Turn 0 of workflow-1"] + + +async def test_same_worker_never_fetches_its_own_turns() -> None: + events: list[StepEvent] = [] + activities = ConversationActivities("worker-1", 0, events.append) + env = environment("run-1") + for step in range(4): + await env.run(activities.conversation_step, step) + assert fetched(events) == [0, 0, 0, 0] + assert len(activities.cache["run-1"]) == 4 + + +async def test_stale_worker_fetches_only_missing_turns() -> None: + events: list[StepEvent] = [] + worker_a = ConversationActivities("worker-1", 0, events.append) + worker_b = ConversationActivities("worker-2", 0, events.append) + 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) + # B lacks turn 0; A returns holding turn 0 but not turn 1. + assert fetched(events) == [0, 1, 1] + assert len(worker_a.cache["run-1"]) == 3 + assert len(worker_b.cache["run-1"]) == 2 + + +async def test_cold_worker_fetches_whole_history() -> None: + events: list[StepEvent] = [] + activities = ConversationActivities("worker-1", 0, events.append) + await environment("run-1").run(activities.conversation_step, 4) + assert fetched(events) == [4] + + +async def test_new_run_starts_with_empty_history() -> None: + events: list[StepEvent] = [] + activities = ConversationActivities("worker-1", 0, events.append) + await environment("run-1").run(activities.conversation_step, 0) + await environment("run-2").run(activities.conversation_step, 1) + assert fetched(events) == [0, 1] diff --git a/tests/activity_sessions/integration_test.py b/tests/activity_sessions/integration_test.py new file mode 100644 index 00000000..54fc1378 --- /dev/null +++ b/tests/activity_sessions/integration_test.py @@ -0,0 +1,59 @@ +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 custom server from temporal#12630", +) + + +@pytest.mark.parametrize( + "session_workers,activity_workers", + [(2, 1), (10, 0), (1, 9)], + ids=["three-workers", "ten-session-workers", "one-session-nine-activity-workers"], +) +async def test_demo_cli( + client: Client, + request: pytest.FixtureRequest, + session_workers: int, + activity_workers: int, +) -> 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", + "activity_sessions.run", + "--session-workers", + str(session_workers), + "--activity-workers", + str(activity_workers), + "--workflows", + "2", + "--steps", + "3", + "--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(), 120) + assert demo.returncode == 0, output.decode() + assert "Verified: each of 2 workflow(s)" in output.decode() + print(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/report_test.py b/tests/activity_sessions/report_test.py new file mode 100644 index 00000000..238d5fa6 --- /dev/null +++ b/tests/activity_sessions/report_test.py @@ -0,0 +1,78 @@ +import io + +import pytest +from rich.console import Console + +from activity_sessions import StepEvent +from activity_sessions.display import ( + Batch, + affinity_violations, + print_worker_configuration, + summary, +) + + +def batch(affinity_enabled: bool, placements: dict[str, list[tuple[str, int]]]): + result = Batch(affinity_enabled, list(placements)) + for workflow_id, steps in placements.items(): + for step, (worker, fetched) in enumerate(steps): + result.record(StepEvent(workflow_id, step, worker, fetched, True)) + result.durations[workflow_id] = 1.0 + sum(fetched for _, fetched in steps) + result.wall_seconds = max(result.durations.values()) + return result + + +def render(*batches: Batch) -> str: + output = io.StringIO() + Console(file=output, width=160, color_system=None).print( + summary(batches, 3, 2, 1.0) + ) + return output.getvalue() + + +def test_comparison_reports_fetches_and_speedup() -> None: + off = batch(False, {"off-1": [("worker-1", 0), ("worker-2", 1), ("worker-1", 1)]}) + on = batch(True, {"on-1": [("worker-2", 0), ("worker-2", 0), ("worker-2", 0)]}) + text = render(off, on) + assert "Affinity OFF" in text and "Affinity ON" in text + assert "●2(1)" in text and "○2" in text + assert "2 fetched" in text + assert "100% fewer" in text + assert "3.0x faster" in text + # OFF: worker-1 holds turns 0-2, worker-2 holds turns 0-1. + assert "5.0" in text + assert "Verified: each of 1 workflow(s)" in text + + +def test_cached_turns_track_each_workers_prefix() -> None: + result = batch(False, {"wf": [("worker-1", 0), ("worker-2", 1), ("worker-1", 1)]}) + assert result.cached_turns("wf") == {"worker-1": 3, "worker-2": 2} + + +def test_affinity_violations_flag_split_or_fetching_workflows() -> None: + result = batch( + True, + { + "good": [("worker-1", 0), ("worker-1", 0), ("worker-1", 0)], + "split": [("worker-1", 0), ("worker-2", 1), ("worker-1", 1)], + "fetched": [("worker-1", 0), ("worker-1", 0), ("worker-1", 2)], + }, + ) + assert affinity_violations(result) == ["split", "fetched"] + assert "Affinity not observed for: split, fetched" in render(result) + + +def test_unknown_workflow_event_is_rejected() -> None: + result = Batch(True, ["known"]) + with pytest.raises(ValueError): + result.record(StepEvent("other", 0, "worker-1", 0, True)) + + +def test_worker_configuration_lists_roles() -> None: + output = io.StringIO() + print_worker_configuration( + {"worker-1": True, "worker-2": False}, + Console(file=output, width=140, color_system=None), + ) + assert "Workflow + activity" in output.getvalue() + assert "Activity only" in output.getvalue() 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" },