diff --git a/CHANGELOG.md b/CHANGELOG.md index f6606d6..77022c9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,59 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ### Added +- **`WorkflowEngine.workflow_scope(workflow_id=None)`** — new async context + manager that groups all tasks registered inside under a shared + `asyncflow.workflow_id`. Internally sets `_workflow_id_ctx` (a module-level + `ContextVar`), which asyncio copies into every `create_task` / `gather` + branch automatically so concurrent workflow instances remain isolated. When + telemetry is active, also opens an OTel `"workflow"` span, making all task + spans inside structural children in the trace hierarchy. Auto-generates a + short UUID if `workflow_id` is `None`. + +- **`_workflow_id_ctx` ContextVar** — per-instance `ContextVar` (default + `None`) that carries the active workflow ID across asyncio task boundaries + without explicit argument passing. Each `WorkflowEngine` instance owns its + own uniquely-named ContextVar (`asyncflow_workflow_id.`), preventing + context leakage when multiple engines are used within the same coroutine. + `@flow.block` execution sets it to the block's UID so tasks inside a block + inherit the workflow ID automatically. + +- **`WorkflowEngine.start_telemetry()` new parameters** — `span_processors`, + `metric_readers`, `resource` — forwarded to `TelemetryManager.__init__()`, + mirroring `Session.start_telemetry()`. Enables passing pre-built OTel + exporters (OTLP, Prometheus, Jaeger) without any RHAPSODY code changes. + +- **Span enricher for `asyncflow.workflow_id`** — registered automatically by + `start_telemetry()` via `register_span_enricher()`. Stamps + `asyncflow.workflow_id` onto every task OTel span when the task was created + inside a `workflow_scope()` or `@flow.block`. Enables per-workflow Gantt + views and span filtering in Jaeger / Grafana Tempo. + +- **`_emit()` `workflow_id` kwarg** — when `workflow_id` is set, injects + `asyncflow.workflow_id` into the event's `attributes` dict. All task lifecycle + events emitted by `WorkflowEngine` (TaskCreated, asyncflow.TaskResolved, + TaskSubmitted) carry the active workflow ID. + +- **Example `01-workflow_grouping.py`** — complete rewrite as an HPC Campaign + Manager simulation. Models 4 workflow types with resource tracking and + dependency chains: + - `simulate` (4 tasks, GPU, no deps) — molecular dynamics runs + - `analyze` (4 tasks, GPU, deps=simulate) — post-processing per simulation + - `train` (8 tasks, GPU, no deps) — distributed ML training + - `evaluate` (8 tasks, CPU, deps=train) — lightweight model evaluation + Uses `ResourcePool` (asyncio-queue-based GPU/CPU slot tracking) and emits + `campaign.ResourceAssigned` custom events to record per-instance resource + assignments in the JSONL checkpoint. + +- **`plot_campaign.py`** — new plotting script producing a two-panel Campaign + Manager timeline figure: + - **Top panel**: Gantt chart with rows per workflow instance, bars coloured by + workflow type and labelled with the assigned resource (`gpu:N` / `cpu:N`). + Right-margin annotations show priority / cpu / gpu per type. Right column + renders a config table, scheduler description box, and dependency graph. + - **Bottom panel**: GPU (left axis) and CPU (right axis) utilisation as step + functions over elapsed time, with total-capacity dashed reference lines. + - **`capture_stdio` decorator parameter** — `@flow.executable_task(capture_stdio=True)` redirects stdout/stderr from executable tasks directly to files instead of collecting them in memory. The awaited future resolves to the stdout **file path** (`{work_dir}/{uid}/{task_uid}.stdout`) @@ -25,6 +78,43 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). - Unit tests for `capture_stdio` field placement, default value, `_work_dir` authority, and end-to-end file I/O with `ConcurrentExecutionBackend`. +### Fixed + +- **Block spans parented correctly** — `execute_block()` now benefits from the + `span_scope()` session-span fallback added in RHAPSODY: block spans that run + in a context with no active OTel span (e.g. spawned from `run()` before + `start_telemetry()` was awaited) now correctly nest under the session root + span instead of floating as unrooted traces. + +- **`execute_block()` dead branch removed** — the `run_in_executor` code path + (used when the wrapped function was sync) is removed. `WorkflowEngine` + enforces a strict async API; all block functions must be `async def`. + +### Changed + +- **`execute_block()` uses `nullcontext`** — the no-telemetry code path uses + stdlib `nullcontext` (Python >= 3.7) instead of the former custom + `_null_context()` helper, which is removed. + +- **`execute_block()` sets `_workflow_id_ctx`** — the block's UID is set as the + active `_workflow_id_ctx` for the duration of block execution so every task + registered inside the block inherits it as `asyncflow.workflow_id` without + requiring an explicit `workflow_scope()` call. + +### Docs + +- **`docs/telemetry.md`**: + - Added "Forwarding to an external backend" section with corrected + `span_processors` / `metric_readers` code examples (replaces the broken + `set_tracer_provider()` pattern). + - Added "`workflow_scope()` context manager" section with usage examples and + auto-ID generation. + - Added "OTel span hierarchy" section showing the four-level + `session -> workflow -> block -> task` tree and explaining how + `asyncflow.workflow_id` propagates to every span attribute and JSONL event. + - Updated `start_telemetry()` signature to include `span_processors`, + `metric_readers`, and `resource` parameters. + ## [0.3.1] - 2026-03-09 ### Fixed diff --git a/docs/assets/workflow-telemetry-dashboard.png b/docs/assets/workflow-telemetry-dashboard.png new file mode 100644 index 0000000..fd38025 Binary files /dev/null and b/docs/assets/workflow-telemetry-dashboard.png differ diff --git a/docs/telemetry.md b/docs/telemetry.md index c36581a..6d6d0eb 100644 --- a/docs/telemetry.md +++ b/docs/telemetry.md @@ -21,8 +21,8 @@ backend = await ConcurrentExecutionBackend(ProcessPoolExecutor()) flow = await WorkflowEngine.create(backend) telemetry = await flow.start_telemetry( - resource_poll_interval=5.0, # node CPU/memory/GPU every 5 s - checkpoint_path="./telemetry/", # write a JSONL file (optional) + resource_poll_interval=5.0, # node CPU/memory/GPU every 5 s + checkpoint_path="./telemetry/", # write a JSONL file (optional) ) ``` @@ -34,6 +34,99 @@ await flow.shutdown() # also stops telemetry await telemetry.stop() ``` +### Forwarding to an external backend + +Pass pre-built OTel `SpanProcessor` and/or `MetricReader` instances to forward data to Jaeger, Grafana Tempo, Honeycomb, Prometheus, or any other OTel-compatible backend: + +```python +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader +from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter + +telemetry = await flow.start_telemetry( + span_processors=[BatchSpanProcessor(OTLPSpanExporter())], + metric_readers=[PeriodicExportingMetricReader(OTLPMetricExporter())], +) +``` + +Exporters read `OTEL_EXPORTER_OTLP_ENDPOINT`, `OTEL_EXPORTER_OTLP_HEADERS`, and `OTEL_SERVICE_NAME` from the environment. See [RHAPSODY Integrations](https://radical-cybertools.github.io/rhapsody/telemetry/integrations/) for the full parameter reference. + +--- + +## Workflow grouping with `workflow_scope()` + +By default all tasks share the same `session_id`. Use `workflow_scope()` to tag every task submitted inside the scope with a `asyncflow.workflow_id` attribute — enabling per-workflow filtering in Jaeger, Tempo, and the JSONL checkpoint: + +```python +async with flow.workflow_scope("etl-run-42"): + raw = ingest("source.csv") + val = validate(raw) + result = await report(val) +``` + +Every span and JSONL event emitted inside the scope carries `asyncflow.workflow_id = "etl-run-42"`. The scope maps directly onto a named OTel span (`name="workflow"`), so the flame graph in Tempo shows `session → workflow(etl-run-42) → task` hierarchy. + +If no `workflow_id` is passed, one is auto-generated (`wf-`). + +```python +# Auto-ID +async with flow.workflow_scope() as wid: + print(wid) # e.g. "wf-3a1b9c2d" +``` + +`@flow.block`-decorated functions are automatically tagged with the block's UID as the `workflow_id` — no explicit `workflow_scope()` needed inside a block body. + +### Visualising workflow telemetry + +The example in `examples/telemetry/01-workflow_grouping.py` +runs six parallel ML training pipelines (load → preprocess → train → evaluate), each +wrapped in a `workflow_scope()`. Running the companion plotting script against the +JSONL checkpoint: + +```bash +python examples/telemetry/01-workflow_grouping.py --out results/ +python examples/telemetry/plot_workflow_gantt.py results/ +``` + +produces a six-panel dashboard: + +![AsyncFlow workflow telemetry dashboard](assets/workflow-telemetry-dashboard.png) + +| Panel | What it shows | +|---|---| +| **Workflow Gantt** | One row per workflow instance; coloured bars = pipeline stages; light band = dependency wait | +| **Stage Execution Time** | Stacked bars in ms per stage per workflow — reveals which stage dominates each run | +| **Workflow Concurrency** | In-flight workflows over time — shows scheduler saturation | +| **Workflow Throughput** | Completion histogram — when pipelines finish relative to session start | +| **E2E Latency Distribution** | Histogram of total elapsed time per workflow with mean/median lines | +| **Stage × Time Heatmap** | 2-D heatmap (stages × time bins): task occupancy per cell — reveals bottleneck stages | + +--- + +## OTel span hierarchy + +When telemetry is enabled, AsyncFlow produces a four-level span tree: + +``` +Trace (one per session) +└── session span [SessionStarted … SessionEnded] + │ + ├── workflow span [workflow_scope("etl-run-42")] + │ ├── task span task_id=ingest-001 + │ ├── task span task_id=validate-001 + │ └── task span task_id=report-001 + │ + ├── block span [execute_block("pipeline_block")] + │ ├── task span task_id=preprocess-001 + │ ├── task span task_id=compute-001 + │ └── task span task_id=store-001 + │ + └── task span (ungrouped — submitted outside any scope) +``` + +The `asyncflow.workflow_id` attribute is stamped on every task span and every JSONL lifecycle event inside a scope. This makes per-workflow Gantt views, performance comparison across workflow types, and span filtering all possible without any post-processing join. + --- ## AsyncFlow-specific events diff --git a/examples/telemetry/01-workflow_grouping.py b/examples/telemetry/01-workflow_grouping.py new file mode 100644 index 0000000..2d2c8bd --- /dev/null +++ b/examples/telemetry/01-workflow_grouping.py @@ -0,0 +1,173 @@ +"""01-workflow_grouping.py — Parallel ML training campaign with workflow telemetry. + +Demonstrates ``workflow_scope()``: six independent training pipelines run +concurrently, each with four sequential stages: + + load_data → preprocess → train → evaluate + +Every task inside a ``workflow_scope()`` automatically carries +``asyncflow.workflow_id`` in every telemetry event, enabling per-pipeline +Gantt charts, stage breakdowns, and cross-pipeline comparisons — with zero +changes to the task functions themselves. + +Usage +----- + python 01-workflow_grouping.py + python 01-workflow_grouping.py --out results/ + python plot_workflow_gantt.py results/ +""" + +import argparse +import asyncio +import logging +import time +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path + +from rhapsody.backends import ConcurrentExecutionBackend +from rhapsody.telemetry import define_event +from rhapsody.telemetry.events import make_event + +from radical.asyncflow import WorkflowEngine +from radical.asyncflow.logging import init_default_logger + +logger = logging.getLogger(__name__) + +# ── Custom event ────────────────────────────────────────────────────────────── + +ExperimentResult = define_event( + "experiment.Result", + experiment_id=str, + dataset_size=int, + learning_rate=float, + loss=float, + accuracy=float, + elapsed_s=float, +) + +# ── Experiment matrix: (dataset_size, learning_rate) ───────────────────────── + +EXPERIMENTS = [ + (100, 0.010), + (200, 0.005), + (150, 0.020), + (80, 0.050), + (250, 0.001), + (120, 0.010), +] + +STAGES = ["load", "preprocess", "train", "evaluate"] + + +# ── Main ────────────────────────────────────────────────────────────────────── + + +async def main(out_dir: str = "telemetry-output") -> None: + Path(out_dir).mkdir(parents=True, exist_ok=True) + + init_default_logger(logging.INFO) + + backend = await ConcurrentExecutionBackend(ProcessPoolExecutor(max_workers=4)) + flow = await WorkflowEngine.create(backend) + telemetry = await flow.start_telemetry( + resource_poll_interval=0.5, + checkpoint_path=out_dir, + ) + + # ── Task definitions ────────────────────────────────────────────────────── + # Task functions are plain async coroutines — no telemetry imports needed. + + @flow.function_task + async def load_data(experiment_id: str, size: int) -> dict: + await asyncio.sleep(0.05 + size * 0.0008) + return {"experiment_id": experiment_id, "size": size, "data": list(range(size))} + + @flow.function_task + async def preprocess(raw: dict) -> dict: + await asyncio.sleep(0.08 + raw["size"] * 0.0006) + return {**raw, "features": [x / raw["size"] for x in raw["data"]]} + + @flow.function_task + async def train(processed: dict, lr: float) -> dict: + await asyncio.sleep(0.20 + processed["size"] * 0.0010) + loss = 0.5 / (1.0 + processed["size"] * lr * 10) + return {**processed, "loss": loss, "lr": lr} + + @flow.function_task + async def evaluate(model: dict) -> dict: + await asyncio.sleep(0.04 + model["size"] * 0.0003) + accuracy = min(0.99, 0.75 + (1.0 - model["loss"]) * 0.24) + return {**model, "accuracy": accuracy} + + # ── Per-experiment runner ───────────────────────────────────────────────── + + async def run_experiment(i: int, size: int, lr: float) -> dict: + t0 = time.time() + exp_id = f"exp-{i}" + + async with flow.workflow_scope(exp_id) as wid: + # asyncflow.workflow_id propagates to all four tasks automatically. + raw = load_data(exp_id, size) + prep = preprocess(raw) + model = train(prep, lr) + result = await evaluate(model) + + elapsed = time.time() - t0 + + # Emit application-level result from the orchestration layer. + telemetry.emit( + make_event( + ExperimentResult, + session_id=telemetry.session_id, + backend="pipeline", + task_id=wid, + experiment_id=exp_id, + dataset_size=size, + learning_rate=lr, + loss=result["loss"], + accuracy=result["accuracy"], + elapsed_s=elapsed, + ) + ) + print( + f" {exp_id} size={size:3d} lr={lr:.3f} " + f"loss={result['loss']:.4f} acc={result['accuracy']:.4f} " + f"({elapsed * 1000:.0f} ms)" + ) + return result + + # ── Run all experiments concurrently ────────────────────────────────────── + + print(f"Running {len(EXPERIMENTS)} experiments concurrently …\n") + t_start = time.time() + await asyncio.gather( + *(run_experiment(i, size, lr) for i, (size, lr) in enumerate(EXPERIMENTS)) + ) + total = time.time() - t_start + + summary = telemetry.summary() + print( + f"\n {len(EXPERIMENTS)} experiments · " + f"{summary['tasks']['completed']} tasks completed in {total * 1000:.0f} ms" + ) + + await flow.shutdown() + await telemetry.stop() + + jsonl = next(Path(out_dir).glob("*.jsonl"), None) + print(f"\nCheckpoint : {jsonl}") + print(f"Plot : python plot_workflow_gantt.py {out_dir}/") + + +if __name__ == "__main__": + ap = argparse.ArgumentParser( + description="Parallel ML training campaign with workflow_scope() telemetry." + ) + ap.add_argument( + "--out", + default="telemetry-output", + metavar="DIR", + help="Output directory for JSONL checkpoint (default: telemetry-output/)", + ) + args = ap.parse_args() + asyncio.run(main(out_dir=args.out)) diff --git a/examples/telemetry/README.md b/examples/telemetry/README.md index ba2792b..1272f38 100644 --- a/examples/telemetry/README.md +++ b/examples/telemetry/README.md @@ -2,9 +2,11 @@ > **Requires:** `pip install rhapsody-py[telemetry]` +Two self-contained examples, each paired with a plotting script. + --- -## Files +## Example 00 — Basic telemetry & custom events ### `00-telemetry_subscribe_basic.py` @@ -20,33 +22,80 @@ Demonstrates: - Emitting **custom application events** (`etl.StageTimer`) per stage using `define_event` + `telemetry.emit()` - Full task lifecycle: `TaskCreated` → `asyncflow.TaskResolved` → `TaskSubmitted` → `TaskStarted` → `TaskCompleted` -Writes a JSONL checkpoint to `telemetry-output/` for use with `plot_workflow_dashboard.py`. - ```bash python 00-telemetry_subscribe_basic.py ``` +Writes a JSONL checkpoint to `telemetry-output/` for `plot_workflow_dashboard.py`. + --- ### `plot_workflow_dashboard.py` -Workflow-focused visualization of the JSONL checkpoint. Produces a PNG with 7 panels: +Comprehensive 8-panel dashboard for the ETL workflow checkpoint: ```bash python plot_workflow_dashboard.py telemetry-output/run.jsonl python plot_workflow_dashboard.py telemetry-output/run.jsonl --out report.png -# Split each figure in the dashboard to be a standalone plot -python plot_workflow_dashboard.py --split telemetry-output/run.jsonl +# Save each panel as a separate file +python plot_workflow_dashboard.py telemetry-output/run.jsonl --split ``` | Panel | What it shows | |---|---| | **Task Gantt** | One bar per task coloured by lifecycle phase (Created → Resolved → Submitted → Running → Completed) | -| **Phase Duration Distribution** | Boxplot of time spent in each phase across all tasks — reveals dep wait vs routing vs execution | +| **Phase Duration Distribution** | Boxplot of time spent in each phase — reveals dep wait vs routing vs execution | | **Concurrency** | Tasks running simultaneously — shows DAG parallelism | | **Task Waterfall** | Stacked bar sorted by e2e time — identifies the slowest tasks and which phase dominates | | **Dependency Wait** | Histogram of `asyncflow.TaskResolved → TaskStarted` — the core AsyncFlow scheduling metric | | **Stage Timer** | Mean ± std per ETL stage from `etl.StageTimer` custom events — application-level breakdown | | **Task Lifecycle Swim-lanes** *(full row)* | One row per task: purple = dep wait · amber = pre-exec · green = execution | | **Concurrency + Node Resources** *(full row)* | Concurrency step (left axis) + CPU % and Memory % per node (right axis) | + +--- + +## Example 01 — Workflow grouping with `workflow_scope()` + +### `01-workflow_grouping.py` + +Six parallel ML training pipelines, each grouped with `workflow_scope()`: + +``` +load_data → preprocess → train → evaluate +``` + +Every task inside a `workflow_scope()` automatically carries +`asyncflow.workflow_id` in every telemetry event — with zero changes to the +task functions themselves. Emits a custom `experiment.Result` event at the +end of each pipeline with loss, accuracy, and elapsed time. + +```bash +python 01-workflow_grouping.py +python 01-workflow_grouping.py --out results/ +``` + +Writes a JSONL checkpoint to `telemetry-output/` (or `--out` dir) for `plot_workflow_gantt.py`. + +--- + +### `plot_workflow_gantt.py` + +Workflow-grouping dashboard — accepts the output directory or JSONL file as a CLI argument: + +```bash +python plot_workflow_gantt.py telemetry-output/ +python plot_workflow_gantt.py telemetry-output/ --out dashboard.png +python plot_workflow_gantt.py telemetry-output/session.jsonl --dpi 180 +``` + +Produces a six-panel PNG: + +| Panel | What it shows | +|---|---| +| **Workflow Gantt** *(full row)* | One row per workflow instance; bars coloured by stage (load / preprocess / train / evaluate); light band = dependency wait | +| **Stage Execution Time** | Stacked horizontal bars: ms per stage per workflow — reveals which stage dominates each run | +| **Workflow Concurrency** | Step function of in-flight workflows over time — shows scheduler saturation | +| **Workflow Throughput** | Histogram of workflow completion times — when pipelines finish relative to session start | +| **E2E Latency Distribution** | Histogram of total elapsed time per workflow with mean/median lines | +| **Stage × Time Heatmap** *(full row)* | 2-D `YlOrRd` heatmap (stages × time bins): fractional task occupancy per cell — reveals bottleneck stages | diff --git a/examples/telemetry/plot_workflow_gantt.py b/examples/telemetry/plot_workflow_gantt.py new file mode 100644 index 0000000..c29ae49 --- /dev/null +++ b/examples/telemetry/plot_workflow_gantt.py @@ -0,0 +1,589 @@ +#!/usr/bin/env python3 +"""plot_workflow_gantt.py — Workflow-grouped telemetry dashboard. + +Reads a RHAPSODY JSONL checkpoint produced by ``01-workflow_grouping.py`` and +renders a six-panel PNG: + + 1. Gantt — one row per workflow instance, bars coloured by stage. + Lighter band shows dependency-wait before each stage. + 2. Stage Breakdown— stacked bars: execution time per stage per workflow. + 3. Workflow Concurrency + — how many workflows are in-flight simultaneously. + 4. Workflow Throughput + — workflows completed per time bin (completion histogram). + 5. E2E Latency Distribution + — histogram of total elapsed time per workflow instance. + 6. Stage × Time Heatmap + — active tasks per stage per time slice; reveals pipeline + saturation and bottleneck stages. + +The ``asyncflow.workflow_id`` attribute stamped on every task event by +``workflow_scope()`` drives all grouping — no changes to task functions needed. + +Usage +----- + python plot_workflow_gantt.py / + python plot_workflow_gantt.py / --out dashboard.png + python plot_workflow_gantt.py session.jsonl --dpi 180 +""" + +from __future__ import annotations + +import argparse +import collections +import json +import sys +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.gridspec as gridspec +import matplotlib.patches as mpatches +import matplotlib.pyplot as plt +import numpy as np + +# ── Theme ───────────────────────────────────────────────────────────────────── + +plt.rcParams.update( + { + "font.family": "DejaVu Sans", + "font.size": 11, + "axes.titlesize": 12, + "axes.titleweight": "bold", + "axes.labelsize": 10, + "axes.labelweight": "bold", + "xtick.labelsize": 8.5, + "ytick.labelsize": 8.5, + "legend.fontsize": 8.5, + "legend.framealpha": 0.92, + "lines.linewidth": 2.0, + "figure.dpi": 140, + } +) + +BG = "#f4f6f9" +PANEL_BG = "#ffffff" +BORDER = "#c8d0da" +GRID = "#e2e6ea" + +STAGE_NAMES = ["load", "preprocess", "train", "evaluate"] +STAGE_COLORS = ["#2e86de", "#ee5a24", "#009432", "#8854d0"] +DEP_WAIT_COLOR = "#dfe6e9" + +# Maps RHAPSODY event_type → task lifecycle field name +EVENT_TYPE_MAP = { + "TaskCreated": "created", + "asyncflow.TaskResolved": "resolved", + "TaskSubmitted": "submitted", + "TaskStarted": "started", + "TaskCompleted": "completed", + "TaskFailed": "failed", + "TaskCanceled": "canceled", +} + + +# ── I/O helpers ─────────────────────────────────────────────────────────────── + + +def _find_jsonl(path: str) -> Path: + p = Path(path) + if p.is_file(): + return p + matches = sorted(p.glob("*.jsonl")) + if not matches: + print(f"error: no .jsonl file found in {p}", file=sys.stderr) + sys.exit(1) + return matches[-1] + + +def load_events(jsonl: Path) -> list[dict]: + events = [] + with open(jsonl) as f: + for line in f: + try: + rec = json.loads(line) + except json.JSONDecodeError: + continue + if rec.get("section") in ("event", None): + events.append(rec) + if not events: + print(f"error: no events in {jsonl}", file=sys.stderr) + sys.exit(1) + return events + + +def session_t0(events: list[dict]) -> float: + ss = [e for e in events if e.get("event_type") == "SessionStarted"] + return ss[0]["event_time"] if ss else min(e["event_time"] for e in events) + + +# ── Data model ──────────────────────────────────────────────────────────────── + + +def build_workflow_tasks(events: list[dict], t0: float) -> dict[str, list[dict]]: + """Group task lifecycle dicts by asyncflow.workflow_id. + + Returns {wf_id: [task_dict, ...]}, tasks sorted by TaskCreated time. + Each task_dict: task_id, workflow_id, created, resolved, submitted, + started, completed, failed, canceled (seconds rel. to t0). + """ + by_task: dict[str, dict] = collections.defaultdict( + lambda: dict( + task_id=None, + workflow_id=None, + created=None, + resolved=None, + submitted=None, + started=None, + completed=None, + failed=None, + canceled=None, + ) + ) + + for e in events: + et = e.get("event_type", "") + tid = e.get("task_id") + if not tid or et not in EVENT_TYPE_MAP: + continue + td = by_task[tid] + td["task_id"] = tid + td[EVENT_TYPE_MAP[et]] = e["event_time"] - t0 + wid = e.get("attributes", {}).get("asyncflow.workflow_id") + if wid: + td["workflow_id"] = wid + + by_wf: dict[str, list[dict]] = collections.defaultdict(list) + for td in by_task.values(): + if td["workflow_id"]: + by_wf[td["workflow_id"]].append(td) + + for wid in by_wf: + by_wf[wid].sort(key=lambda x: x["created"] or 0) + + return dict(by_wf) + + +def _wf_span(tasks: list[dict]) -> tuple[float | None, float | None]: + """(earliest created, latest completed) for a workflow's task list.""" + starts = [t["created"] for t in tasks if t["created"] is not None] + ends = [ + t["completed"] or t["failed"] for t in tasks if (t["completed"] or t["failed"]) + ] + return (min(starts) if starts else None, max(ends) if ends else None) + + +def build_wf_concurrency(by_wf: dict[str, list[dict]]): + """Step function: number of in-flight workflows over time.""" + deltas = [] + for tasks in by_wf.values(): + t_s, t_e = _wf_span(tasks) + if t_s is not None: + deltas.append((t_s, +1)) + if t_e is not None: + deltas.append((t_e, -1)) + if not deltas: + return [], [] + deltas.sort(key=lambda x: (x[0], -x[1])) + t_out, y_out, running = [], [], 0 + for tc, d in deltas: + t_out += [tc, tc] + y_out += [running, running + d] + running += d + t_out.append(t_out[-1]) + y_out.append(0) + return t_out, y_out + + +def build_stage_heatmap( + by_wf: dict[str, list[dict]], n_bins: int = 40 +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """2-D array [stages × time_bins] counting tasks active per cell. + + Returns: (matrix, time_edges, stage_names) + """ + all_t: list[float] = [] + for tasks in by_wf.values(): + for td in tasks: + if td["started"] is not None: + all_t.append(td["started"]) + t_e = td.get("completed") or td.get("failed") + if t_e is not None: + all_t.append(t_e) + if not all_t: + return ( + np.zeros((len(STAGE_NAMES), n_bins)), + np.linspace(0, 1, n_bins + 1), + STAGE_NAMES, + ) + + t_min, t_max = min(all_t), max(all_t) + edges = np.linspace(t_min, t_max, n_bins + 1) + bin_w = edges[1] - edges[0] + mat = np.zeros((len(STAGE_NAMES), n_bins)) + + for tasks in by_wf.values(): + for stage_idx, td in enumerate(tasks): + if stage_idx >= len(STAGE_NAMES): + break + t_s = td.get("started") + t_e = td.get("completed") or td.get("failed") + if t_s is None or t_e is None: + continue + # Distribute task's execution uniformly across the bins it overlaps + for b in range(n_bins): + b_lo, b_hi = edges[b], edges[b + 1] + overlap = max(0.0, min(t_e, b_hi) - max(t_s, b_lo)) + if overlap > 0: + mat[stage_idx, b] += overlap / bin_w # fractional occupancy + + return mat, edges, STAGE_NAMES + + +# ── Style helpers ───────────────────────────────────────────────────────────── + + +def _decorate(ax, title, xlabel="Time since start (s)", ylabel="", grid=True): + ax.set_facecolor(PANEL_BG) + ax.set_title(title, pad=7, color="#2d3436") + ax.set_xlabel(xlabel, labelpad=4) + ax.set_ylabel(ylabel, labelpad=4) + for spine in ax.spines.values(): + spine.set_visible(True) + spine.set_linewidth(1.1) + spine.set_edgecolor(BORDER) + if grid: + ax.grid(color=GRID, linewidth=0.6, linestyle="-", zorder=0) + ax.set_axisbelow(True) + + +# ── Panels ──────────────────────────────────────────────────────────────────── + + +def _panel_gantt(ax, by_wf: dict[str, list[dict]]): + wf_ids = sorted(by_wf.keys()) + stage_patches = [ + mpatches.Patch(color=STAGE_COLORS[i], label=STAGE_NAMES[i]) + for i in range(len(STAGE_NAMES)) + ] + dep_patch = mpatches.Patch(color=DEP_WAIT_COLOR, label="dep wait") + + for row, wid in enumerate(wf_ids): + for stage_idx, td in enumerate(by_wf[wid]): + color = STAGE_COLORS[stage_idx % len(STAGE_COLORS)] + t_s, t_e = td.get("started"), td.get("completed") or td.get("failed") + t_c = td.get("created") + + if t_s is None: + continue + if t_c is not None and t_s > t_c: + ax.barh( + row, + t_s - t_c, + left=t_c, + height=0.55, + color=DEP_WAIT_COLOR, + alpha=0.85, + edgecolor="none", + zorder=2, + ) + if t_e is not None and t_e > t_s: + ax.barh( + row, + t_e - t_s, + left=t_s, + height=0.65, + color=color, + edgecolor="white", + linewidth=0.4, + zorder=3, + ) + + ax.set_yticks(range(len(wf_ids))) + ax.set_yticklabels(wf_ids, fontsize=9) + ax.set_ylim(-0.6, len(wf_ids) - 0.4) + ax.legend( + handles=stage_patches + [dep_patch], + ncol=len(stage_patches) + 1, + fontsize=8, + loc="upper right", + framealpha=0.92, + ) + _decorate( + ax, "Workflow Gantt — tasks grouped by workflow_scope()", ylabel="Workflow" + ) + + +def _panel_stage_breakdown(ax, by_wf: dict[str, list[dict]]): + wf_ids = sorted(by_wf.keys()) + n = len(wf_ids) + lefts = np.zeros(n) + + for stage_idx, stage_name in enumerate(STAGE_NAMES): + widths = [] + for wid in wf_ids: + tasks = by_wf[wid] + td = tasks[stage_idx] if stage_idx < len(tasks) else {} + t_s, t_e = td.get("started"), td.get("completed") or td.get("failed") + widths.append((t_e - t_s) * 1000 if (t_s and t_e) else 0.0) + + ax.barh( + range(n), + widths, + left=lefts, + color=STAGE_COLORS[stage_idx], + edgecolor="white", + linewidth=0.5, + label=stage_name, + height=0.55, + zorder=3, + ) + for row, (left, w) in enumerate(zip(lefts, widths)): + if w > 8: + ax.text( + left + w / 2, + row, + f"{w:.0f}", + ha="center", + va="center", + fontsize=7, + color="white", + fontweight="bold", + zorder=4, + ) + lefts += np.array(widths) + + ax.set_yticks(range(n)) + ax.set_yticklabels(wf_ids, fontsize=9) + ax.set_ylim(-0.6, n - 0.4) + ax.legend(ncol=len(STAGE_NAMES), fontsize=8, loc="lower right") + _decorate( + ax, + "Stage Execution Time per Workflow (ms)", + xlabel="Execution time (ms)", + ylabel="Workflow", + ) + + +def _panel_wf_concurrency(ax, t_wf, y_wf): + if t_wf: + ax.step(t_wf, y_wf, color="#6c5ce7", linewidth=2.2, where="post", zorder=3) + ax.fill_between(t_wf, y_wf, step="post", alpha=0.22, color="#6c5ce7", zorder=2) + ax.yaxis.set_major_locator(plt.MaxNLocator(integer=True)) + ax.set_ylim(bottom=0) + _decorate( + ax, + "Workflow Concurrency\n(in-flight simultaneously)", + ylabel="Workflows running", + ) + + +def _panel_wf_throughput(ax, by_wf: dict[str, list[dict]], n_bins: int = 12): + """Histogram of workflow completion times.""" + completions = [] + for tasks in by_wf.values(): + _, t_e = _wf_span(tasks) + if t_e is not None: + completions.append(t_e) + + if completions: + ax.hist( + completions, + bins=n_bins, + color="#00b894", + edgecolor="white", + linewidth=0.7, + zorder=3, + ) + ax.yaxis.set_major_locator(plt.MaxNLocator(integer=True)) + _decorate( + ax, + "Workflow Throughput\n(completions per time bin)", + ylabel="Workflows completed", + ) + + +def _panel_e2e_latency(ax, by_wf: dict[str, list[dict]]): + """Histogram of end-to-end elapsed time per workflow.""" + latencies_ms = [] + for tasks in by_wf.values(): + t_s, t_e = _wf_span(tasks) + if t_s is not None and t_e is not None: + latencies_ms.append((t_e - t_s) * 1000) + + if latencies_ms: + n_bins = min(max(6, len(latencies_ms)), 20) + ax.hist( + latencies_ms, + bins=n_bins, + color="#fdcb6e", + edgecolor="white", + linewidth=0.7, + zorder=3, + ) + mu = np.mean(latencies_ms) + ax.axvline( + mu, + color="#d63031", + linestyle="--", + linewidth=1.8, + label=f"mean {mu:.0f} ms", + zorder=4, + ) + if len(latencies_ms) > 1: + ax.axvline( + np.median(latencies_ms), + color="#2d3436", + linestyle=":", + linewidth=1.6, + label=f"median {np.median(latencies_ms):.0f} ms", + zorder=4, + ) + ax.legend(fontsize=8) + ax.yaxis.set_major_locator(plt.MaxNLocator(integer=True)) + _decorate( + ax, + "E2E Workflow Latency Distribution", + xlabel="Elapsed time (ms)", + ylabel="Workflows", + ) + + +def _panel_heatmap(ax, mat: np.ndarray, edges: np.ndarray, stage_names: list[str]): + """Stage × time heatmap: fractional task occupancy per cell.""" + if mat.max() == 0: + ax.text( + 0.5, + 0.5, + "no data", + ha="center", + va="center", + transform=ax.transAxes, + fontsize=12, + color="#636e72", + ) + _decorate(ax, "Stage × Time Throughput Heatmap", grid=False) + return + + im = ax.imshow( + mat, + aspect="auto", + origin="lower", + cmap="YlOrRd", + vmin=0, + extent=[edges[0], edges[-1], -0.5, len(stage_names) - 0.5], + interpolation="nearest", + ) + ax.set_yticks(range(len(stage_names))) + ax.set_yticklabels(stage_names, fontsize=9) + ax.set_ylim(-0.5, len(stage_names) - 0.5) + + cbar = ax.get_figure().colorbar(im, ax=ax, pad=0.01, fraction=0.025) + cbar.set_label("Tasks active (fractional)", fontsize=8) + cbar.ax.tick_params(labelsize=7) + + for spine in ax.spines.values(): + spine.set_visible(True) + spine.set_linewidth(1.1) + spine.set_edgecolor(BORDER) + ax.set_facecolor(PANEL_BG) + ax.set_title("Stage × Time Throughput Heatmap", pad=7, color="#2d3436") + ax.set_xlabel("Time since start (s)", labelpad=4) + ax.set_ylabel("Pipeline stage", labelpad=4) + + +# ── Main ────────────────────────────────────────────────────────────────────── + + +def plot(jsonl: Path, out: str | None, dpi: int) -> None: + events = load_events(jsonl) + t0 = session_t0(events) + name = jsonl.stem + + by_wf = build_workflow_tasks(events, t0) + if not by_wf: + print( + "warning: no asyncflow.workflow_id attributes found — " + "did you run 01-workflow_grouping.py?", + file=sys.stderr, + ) + + t_wf, y_wf = build_wf_concurrency(by_wf) + mat, edges, snames = build_stage_heatmap(by_wf) + n_wf = len(by_wf) + n_tasks = sum(len(t) for t in by_wf.values()) + + gantt_h = max(1.6, n_wf * 0.52 + 0.8) + fig = plt.figure(figsize=(20, gantt_h * 2 + 14), facecolor=BG) + fig.suptitle( + f"AsyncFlow Workflow Telemetry · {name}" + f" ({n_wf} workflows · {n_tasks} tasks)", + fontsize=15, + fontweight="bold", + y=0.999, + color="#2d3436", + ) + + gs = gridspec.GridSpec( + 3, + 2, + figure=fig, + hspace=0.62, + wspace=0.38, + top=0.965, + bottom=0.05, + left=0.10, + right=0.97, + height_ratios=[gantt_h, gantt_h, 2.2], + ) + + _panel_gantt(fig.add_subplot(gs[0, :]), by_wf) + _panel_stage_breakdown(fig.add_subplot(gs[1, 0]), by_wf) + + # Right column: three stacked panels + gs_right = gridspec.GridSpecFromSubplotSpec( + 3, 1, subplot_spec=gs[1, 1], hspace=0.75, height_ratios=[1, 1, 1] + ) + _panel_wf_concurrency(fig.add_subplot(gs_right[0]), t_wf, y_wf) + _panel_wf_throughput(fig.add_subplot(gs_right[1]), by_wf) + _panel_e2e_latency(fig.add_subplot(gs_right[2]), by_wf) + + _panel_heatmap(fig.add_subplot(gs[2, :]), mat, edges, snames) + + out_path = out or str(jsonl.with_suffix(".png")) + fig.savefig(out_path, dpi=dpi, bbox_inches="tight", facecolor=BG) + print(f"Saved → {out_path}") + plt.close(fig) + + +# ── CLI ─────────────────────────────────────────────────────────────────────── + + +def main() -> None: + ap = argparse.ArgumentParser( + description="Plot workflow-grouped AsyncFlow telemetry from a JSONL checkpoint." + ) + ap.add_argument( + "path", + metavar="DIR_OR_FILE", + help=( + "Output directory produced by 01-workflow_grouping.py " + "(the .jsonl file inside is auto-detected), or a .jsonl file directly." + ), + ) + ap.add_argument( + "--out", + metavar="FILE", + default=None, + help="Output PNG path (default: .png next to the input).", + ) + ap.add_argument("--dpi", type=int, default=140, help="Image DPI (default: 140).") + args = ap.parse_args() + + jsonl = _find_jsonl(args.path) + plot(jsonl, out=args.out, dpi=args.dpi) + + +if __name__ == "__main__": + main() diff --git a/src/radical/asyncflow/workflow_manager.py b/src/radical/asyncflow/workflow_manager.py index 8ccf956..439e75f 100644 --- a/src/radical/asyncflow/workflow_manager.py +++ b/src/radical/asyncflow/workflow_manager.py @@ -10,10 +10,12 @@ import time import uuid from collections import defaultdict, deque +from contextlib import asynccontextmanager, nullcontext +from contextvars import ContextVar from functools import wraps from itertools import count from pathlib import Path -from typing import Any, Callable, Optional, Union +from typing import Any, AsyncIterator, Callable, Optional, Union import typeguard @@ -76,6 +78,11 @@ def __init__( self.uid = uid or f"asyncflow.session.{uuid.uuid4().hex[:8]}" self.work_dir = work_dir or os.getcwd() + # Per-engine ContextVar avoids cross-engine context leakage in shared coroutines, + # while asyncio still propagates it across create_task/gather branches. + self._workflow_id_ctx: ContextVar[str | None] = ContextVar( + f"asyncflow_workflow_id.{self.uid}", default=None + ) # Normalize backend: accept a single backend or a list of backends if not isinstance(backend, list): @@ -148,6 +155,9 @@ async def start_telemetry( resource_poll_interval: float = 5.0, checkpoint_interval: float | None = None, checkpoint_path: str | None = None, + span_processors: list | None = None, + metric_readers: list | None = None, + resource: object | None = None, ) -> Any: """Create and start a RHAPSODY TelemetryManager for this workflow engine. @@ -156,6 +166,21 @@ async def start_telemetry( NoopExecutionBackend which have no adapter), starts the async dispatch loop, and returns the manager. + Args: + resource_poll_interval: Seconds between resource metric polls (default: 5.0). + checkpoint_interval: Seconds between periodic metric+span flushes to disk. + None = no periodic flush (file still written at stop). + checkpoint_path: Directory for the JSONL checkpoint file. + None = no file output. + span_processors: Optional list of OTel SpanProcessor instances added to + RHAPSODY's TracerProvider alongside the internal + SpanBuffer. Callers own exporter construction. + metric_readers: Optional list of OTel MetricReader instances added to + RHAPSODY's MeterProvider alongside InMemoryMetricReader. + resource: Optional ``opentelemetry.sdk.resources.Resource``. + When None, ``Resource.create()`` reads + ``OTEL_SERVICE_NAME`` from the environment automatically. + Returns: The active TelemetryManager instance. """ @@ -165,6 +190,9 @@ async def start_telemetry( session_id=self.uid, checkpoint_interval=checkpoint_interval, checkpoint_path=checkpoint_path, + span_processors=span_processors, + metric_readers=metric_readers, + resource=resource, ) for backend in self._backends.values(): @@ -177,6 +205,15 @@ async def start_telemetry( await self._telemetry.start() + self._telemetry.register_span_enricher( + lambda event: ( + {"asyncflow.workflow_id": event.attributes["asyncflow.workflow_id"]} + if isinstance(event.attributes, dict) + and "asyncflow.workflow_id" in event.attributes + else {} + ) + ) + from rhapsody.telemetry.events import ( # noqa: PLC0415 TaskCanceled, TaskCompleted, @@ -208,7 +245,13 @@ async def start_telemetry( return self._telemetry def _emit( - self, event_name: str, *, task_id: str, backend: str = None, **kwargs + self, + event_name: str, + *, + task_id: str, + backend: str = None, + workflow_id: str | None = None, + **kwargs, ) -> None: """Emit a telemetry event by name. @@ -218,6 +261,8 @@ def _emit( return if "event_time" not in kwargs: kwargs["event_time"] = time.time() + if workflow_id is not None: + kwargs.setdefault("attributes", {})["asyncflow.workflow_id"] = workflow_id self._telemetry.emit( self._tel_make_event( self._tel_events[event_name], @@ -228,6 +273,35 @@ def _emit( ) ) + @asynccontextmanager + async def workflow_scope( + self, workflow_id: str | None = None + ) -> AsyncIterator[str]: + """Tag all tasks registered inside this scope with a shared workflow_id. + + Also creates an OTel workflow span (when telemetry is active) so that task spans + registered inside become structural children in the trace hierarchy. + + Args: + workflow_id: Explicit id for this workflow instance. + If omitted, a short UUID is generated automatically. + + Yields: + The workflow_id string in use. + """ + wid = workflow_id or f"wf-{uuid.uuid4().hex[:8]}" + token = self._workflow_id_ctx.set(wid) + try: + scope = ( + self._telemetry.span_scope("workflow", {"asyncflow.workflow_id": wid}) + if self._telemetry is not None + else nullcontext() + ) + with scope: + yield wid + finally: + self._workflow_id_ctx.reset(token) + def _setup_signal_handlers(self): """Register signal handlers for graceful shutdown on SIGHUP, SIGTERM, and SIGINT.""" @@ -629,6 +703,7 @@ def _register_component( # make sure not to specify both func and executable at the same time comp_desc["name"] = comp_desc["function"].__name__ comp_desc["uid"] = self._assign_uid(prefix=comp_type) + comp_desc["workflow_id"] = self._workflow_id_ctx.get() if task_type == EXECUTABLE: comp_desc[FUNCTION] = None # Clear function since we're using executable @@ -691,6 +766,7 @@ def _register_component( self._emit( "TaskCreated", task_id=comp_desc["uid"], + workflow_id=comp_desc.get("workflow_id"), attributes={ "executable": comp_desc["_tel_executable"], "task_type": comp_desc["_tel_task_type"], @@ -704,6 +780,7 @@ def _register_component( self._emit( "asyncflow.TaskResolved", task_id=comp_desc["uid"], + workflow_id=comp_desc.get("workflow_id"), attributes={ "executable": comp_desc["_tel_executable"], "task_type": comp_desc["_tel_task_type"], @@ -865,6 +942,7 @@ def _notify_dependents(self, comp_uid: str): self._emit( "asyncflow.TaskResolved", task_id=dependent_uid, + workflow_id=dep_desc.get("workflow_id"), attributes={ "executable": dep_desc["_tel_executable"], "task_type": dep_desc["_tel_task_type"], @@ -1108,6 +1186,7 @@ async def run(self): "TaskQueued", task_id=comp_uid, backend=comp_desc.get("target_backend"), + workflow_id=comp_desc.get("workflow_id"), attributes={ "executable": comp_desc["_tel_executable"], "task_type": comp_desc["_tel_task_type"], @@ -1216,6 +1295,7 @@ async def submit(self, objects: list) -> None: "TaskSubmitted", task_id=t["uid"], backend=t.get("target_backend"), + workflow_id=t.get("workflow_id"), event_time=now, attributes={ "executable": t["_tel_executable"], @@ -1304,18 +1384,25 @@ async def execute_block( None """ + block_uid = block_fut.block["uid"] + token = self._workflow_id_ctx.set(block_uid) try: - if asyncio.iscoroutinefunction(func): + scope = ( + self._telemetry.span_scope( + "block", {"asyncflow.workflow_id": block_uid} + ) + if self._telemetry is not None + else nullcontext() + ) + with scope: result = await func(*args, **kwargs) - else: - # Run sync function in executor - result = await self.loop.run_in_executor(None, func, *args, **kwargs) - if not block_fut.done(): block_fut.set_result(result) except Exception as e: if not block_fut.done(): block_fut.set_exception(e) + finally: + self._workflow_id_ctx.reset(token) def handle_task_success(self, task: dict, task_fut: asyncio.Future) -> None: """Handles successful task completion and updates the associated future. @@ -1499,6 +1586,7 @@ def wait_and_set(): self._emit( "TaskCompleted", task_id=uid, + workflow_id=task_dct.get("workflow_id"), event_time=now, duration_seconds=now - start, attributes=tel_attrs, @@ -1514,6 +1602,7 @@ def wait_and_set(): self._emit( "TaskStarted", task_id=uid, + workflow_id=task_dct.get("workflow_id"), event_time=now, attributes=tel_attrs, ) @@ -1528,6 +1617,7 @@ def wait_and_set(): self._emit( "TaskCanceled", task_id=uid, + workflow_id=task_dct.get("workflow_id"), event_time=now, duration_seconds=now - start, attributes=tel_attrs, @@ -1543,6 +1633,7 @@ def wait_and_set(): self._emit( "TaskFailed", task_id=uid, + workflow_id=task_dct.get("workflow_id"), event_time=now, duration_seconds=now - start, error_type="unknown", diff --git a/tests/unit/test_workflow_scope_isolation.py b/tests/unit/test_workflow_scope_isolation.py new file mode 100644 index 0000000..802028f --- /dev/null +++ b/tests/unit/test_workflow_scope_isolation.py @@ -0,0 +1,87 @@ +"""Unit tests for workflow_scope() ContextVar isolation (per-instance fix).""" + +import asyncio + +from radical.asyncflow import NoopExecutionBackend, WorkflowEngine + + +async def test_workflow_scope_does_not_leak_to_sibling_engine(): + """engine1.workflow_scope() must not stamp workflow_id onto engine2's tasks.""" + engine1 = await WorkflowEngine.create(backend=NoopExecutionBackend()) + engine2 = await WorkflowEngine.create(backend=NoopExecutionBackend()) + + async with engine1.workflow_scope("wf-A"): + + @engine1.function_task + async def task1(x): + return x + + @engine2.function_task + async def task2(x): + return x + + await task1(1) + await task2(2) + await asyncio.sleep(0.05) + + desc1 = next(iter(engine1.components.values()))["description"] + desc2 = next(iter(engine2.components.values()))["description"] + + assert desc1["workflow_id"] == "wf-A", "engine1 task must carry the workflow_id" + assert desc2["workflow_id"] is None, "engine2 task must NOT inherit engine1's scope" + + await engine1.shutdown() + await engine2.shutdown() + + +async def test_workflow_scope_stamps_all_tasks_on_same_engine(): + """All tasks registered inside workflow_scope() on the same engine get the id.""" + engine = await WorkflowEngine.create(backend=NoopExecutionBackend()) + + async with engine.workflow_scope("wf-B"): + + @engine.function_task + async def t1(x): + return x + + @engine.function_task + async def t2(x): + return x + + await t1(1) + await t2(2) + await asyncio.sleep(0.05) + + wids = [c["description"]["workflow_id"] for c in engine.components.values()] + assert all(w == "wf-B" for w in wids), f"Expected all 'wf-B', got {wids}" + + await engine.shutdown() + + +async def test_workflow_scope_resets_after_exit(): + """After exiting workflow_scope(), tasks on the same engine get workflow_id=None.""" + engine = await WorkflowEngine.create(backend=NoopExecutionBackend()) + + async with engine.workflow_scope("wf-C"): + + @engine.function_task + async def inside(x): + return x + + await inside(1) + await asyncio.sleep(0.05) + + @engine.function_task + async def outside(x): + return x + + await outside(2) + await asyncio.sleep(0.05) + + comps = list(engine.components.values()) + assert len(comps) == 2 + wids = {c["description"]["name"]: c["description"]["workflow_id"] for c in comps} + assert wids["inside"] == "wf-C" + assert wids["outside"] is None + + await engine.shutdown()