Add workflow-session activity cache sample

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.
 <!-- Keep this list in alphabetical order -->
+* [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.

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.

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

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

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)