Add activity sessions sample (affinity comparison and single-session lifecycle)

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 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.

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

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

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

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